fix(workbuddy): stream TTFB timeout failover (from #336)

internal/providers/workbuddy/client.go

@@ -537,7 +537,9 @@ func (c *Client) sendChat(ctx context.Context, accountID string, credential Cred
 	client.Timeout = 0
 	refreshed := false
 	for {
-		resp, err := client.Do(httpReq)
+		// Fail over when the upstream accepts but never sends a response head
+		// (see ttfb.go). Kept inside sendChat so the 401 refresh loop still runs.
+		resp, err := doWithFirstByteTimeout(client, httpReq, resolveFirstByteTimeout())
 		if err != nil {
 			return nil, wrapTransportError(err)
 		}

internal/providers/workbuddy/ttfb.go

@@ -0,0 +1,93 @@
+package workbuddy
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"os"
+	"strconv"
+	"strings"
+	"time"
+
+	"github.com/caigee-cmd/cli2api/internal/providers"
+)
+
+// Streaming first-byte timeout.
+//
+// WorkBuddy's CN gateway sometimes accepts a chat request and then sends nothing
+// at all for minutes — an upstream hang, not a rejection. Without a deadline the
+// relay waits until the client gives up, the turn is spent, and a retry lands on
+// the same account. Qoder already fails over on this (QODER_STREAM_TTFB_MS); this
+// is the WorkBuddy counterpart.
+const defaultFirstByteTimeout = 30 * time.Second
+
+// workbuddyStreamTTFBEnv overrides the first-byte deadline.
+const workbuddyStreamTTFBEnv = "WORKBUDDY_STREAM_TTFB_MS"
+
+// resolveFirstByteTimeout reads WORKBUDDY_STREAM_TTFB_MS. Absent or invalid
+// values fall back to the default; sub-5s values clamp up so the knob cannot be
+// set so low that normal upstream latency looks like a hang.
+func resolveFirstByteTimeout() time.Duration {
+	raw := strings.TrimSpace(os.Getenv(workbuddyStreamTTFBEnv))
+	if raw == "" {
+		return defaultFirstByteTimeout
+	}
+	ms, err := strconv.Atoi(raw)
+	if err != nil || ms <= 0 {
+		return defaultFirstByteTimeout
+	}
+	d := time.Duration(ms) * time.Millisecond
+	if d < 5*time.Second {
+		return 5 * time.Second
+	}
+	return d
+}
+
+// doWithFirstByteTimeout sends req and returns once the response head arrives.
+// If the head does not arrive within timeout the error wraps
+// providers.ErrFirstByteTimeout, which the executor classifies as an upstream
+// failure and fails over to another account.
+//
+// The deadline context must NOT be cancelled after a successful head: Go ties
+// the response body to the request context, so cancelling would truncate the
+// stream that has already started. Only the error path cancels; on success the
+// timer stops and the derived context is released by Body.Close.
+func doWithFirstByteTimeout(client *http.Client, req *http.Request, timeout time.Duration) (*http.Response, error) {
+	if timeout <= 0 {
+		timeout = defaultFirstByteTimeout
+	}
+	parent := req.Context()
+	ttfbCtx, cancelTTFB := context.WithCancel(parent)
+	timer := time.AfterFunc(timeout, cancelTTFB)
+	resp, err := client.Do(req.WithContext(ttfbCtx))
+	if err == nil {
+		timer.Stop()
+		resp.Body = &cancelOnCloseBody{ReadCloser: resp.Body, cancel: cancelTTFB}
+		return resp, nil
+	}
+	timer.Stop()
+	// Decide the cause before cancelTTFB, which would otherwise make
+	// ttfbCtx.Err() non-nil regardless of what actually happened. A parent
+	// cancel stays a plain cancel so a client disconnect is not cooled down.
+	timedOut := parent.Err() == nil && ttfbCtx.Err() != nil && !errors.Is(parent.Err(), context.Canceled)
+	cancelTTFB()
+	if timedOut {
+		return nil, fmt.Errorf("%w: no response head within %s", providers.ErrFirstByteTimeout, timeout)
+	}
+	return nil, err
+}
+
+// cancelOnCloseBody releases the TTFB-derived context when the caller is done
+// with the stream, without cutting the body short the way an eager cancel would.
+type cancelOnCloseBody struct {
+	io.ReadCloser
+	cancel context.CancelFunc
+}
+
+func (b *cancelOnCloseBody) Close() error {
+	err := b.ReadCloser.Close()
+	b.cancel()
+	return err
+}

internal/providers/workbuddy/ttfb_test.go

@@ -0,0 +1,98 @@
+package workbuddy
+
+import (
+	"context"
+	"errors"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/caigee-cmd/cli2api/internal/providers"
+)
+
+func TestResolveFirstByteTimeout(t *testing.T) {
+	t.Setenv(workbuddyStreamTTFBEnv, "")
+	if got := resolveFirstByteTimeout(); got != defaultFirstByteTimeout {
+		t.Fatalf("default = %s, want %s", got, defaultFirstByteTimeout)
+	}
+
+	t.Setenv(workbuddyStreamTTFBEnv, "45000")
+	if got := resolveFirstByteTimeout(); got != 45*time.Second {
+		t.Fatalf("override = %s, want 45s", got)
+	}
+
+	// Below the floor must clamp up, otherwise normal latency looks like a hang.
+	t.Setenv(workbuddyStreamTTFBEnv, "100")
+	if got := resolveFirstByteTimeout(); got != 5*time.Second {
+		t.Fatalf("floor = %s, want 5s", got)
+	}
+
+	// Invalid values fall back to the default rather than disabling the guard.
+	for _, bad := range []string{"0", "-1", "soon", "  "} {
+		t.Setenv(workbuddyStreamTTFBEnv, bad)
+		if got := resolveFirstByteTimeout(); got != defaultFirstByteTimeout {
+			t.Fatalf("value %q = %s, want default %s", bad, got, defaultFirstByteTimeout)
+		}
+	}
+}
+
+func TestChatStreamFailsOverWhenHeadNeverArrives(t *testing.T) {
+	release := make(chan struct{})
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// Accept the request, then never write a head — the upstream hang this
+		// guard exists for.
+		<-release
+	}))
+	defer server.Close()
+	defer close(release)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, server.URL, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	start := time.Now()
+	_, err = doWithFirstByteTimeout(server.Client(), req, 150*time.Millisecond)
+	if err == nil {
+		t.Fatal("expected a first-byte timeout")
+	}
+	if !errors.Is(err, providers.ErrFirstByteTimeout) {
+		t.Fatalf("err = %v, want ErrFirstByteTimeout", err)
+	}
+	if elapsed := time.Since(start); elapsed > 3*time.Second {
+		t.Fatalf("waited %s, the deadline did not fire", elapsed)
+	}
+}
+
+func TestChatStreamKeepsBodyAliveAfterHead(t *testing.T) {
+	// A head that arrives in time must leave the body readable: cancelling the
+	// deadline context after success would truncate the stream.
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+		w.(http.Flusher).Flush()
+		time.Sleep(300 * time.Millisecond)
+		_, _ = w.Write([]byte("data: ok\n\n"))
+	}))
+	defer server.Close()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, server.URL, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	resp, err := doWithFirstByteTimeout(server.Client(), req, 2*time.Second)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	defer resp.Body.Close()
+	buf := make([]byte, 64)
+	n, _ := resp.Body.Read(buf)
+	if n == 0 {
+		t.Fatal("body was truncated after a successful head")
+	}
+}