From d2a31556789d431b3af58f80c8f3d94384200237 Mon Sep 17 00:00:00 2001 From: caigee-cmd <210803857+caigee-cmd@users.noreply.github.com> Date: Sat, 24 Oct 2026 20:13:34 +0800 Subject: [PATCH] fix(workbuddy): fail over when the stream head never arrives Port WORKBUDDY_STREAM_TTFB_MS from #336, but wire doWithFirstByteTimeout inside the existing sendChat 401-refresh loop instead of replacing it. The author's branch removed sendChat entirely; that would regress token refresh / TLS classify / KindUnavailable handling that landed on main after the PR branched. --- internal/providers/workbuddy/client.go | 4 +- internal/providers/workbuddy/ttfb.go | 93 +++++++++++++++++++++ internal/providers/workbuddy/ttfb_test.go | 98 +++++++++++++++++++++++ 3 files changed, 194 insertions(+), 1 deletion(-) create mode 100644 internal/providers/workbuddy/ttfb.go create mode 100644 internal/providers/workbuddy/ttfb_test.go diff --git a/internal/providers/workbuddy/client.go b/internal/providers/workbuddy/client.go index 617f1ef4..1cafea66 100644 --- a/internal/providers/workbuddy/client.go +++ b/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) } diff --git a/internal/providers/workbuddy/ttfb.go b/internal/providers/workbuddy/ttfb.go new file mode 100644 index 00000000..d03595bf --- /dev/null +++ b/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 +} diff --git a/internal/providers/workbuddy/ttfb_test.go b/internal/providers/workbuddy/ttfb_test.go new file mode 100644 index 00000000..80226128 --- /dev/null +++ b/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") + } +}