From 5dbd92a0e0dfcf1d98adf58e793f3ce6c95ff082 Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Sat, 10 Oct 2026 20:26:36 +0800 Subject: [PATCH] fix(auth): preserve upstream error cooldown on canceled stream contexts - Record credential and model cooldowns when stream execution encounters an upstream error despite context cancellation - Prioritize reading ready stream chunks before handling context cancellation in conductor - Defer Codex websocket disconnect notifications until error frames are delivered to the stream Closes: #6503 --- .../codex_stream_bootstrap_buffering_test.go | 66 ++++++ .../executor/codex_websockets_stream.go | 13 +- sdk/cliproxy/auth/conductor_issue6503_test.go | 194 ++++++++++++++++++ sdk/cliproxy/auth/conductor_stream.go | 30 ++- 4 files changed, 299 insertions(+), 4 deletions(-) create mode 100644 sdk/cliproxy/auth/conductor_issue6503_test.go diff --git a/internal/runtime/executor/codex_stream_bootstrap_buffering_test.go b/internal/runtime/executor/codex_stream_bootstrap_buffering_test.go index 8b8cba1b0..f2ff79d9e 100644 --- a/internal/runtime/executor/codex_stream_bootstrap_buffering_test.go +++ b/internal/runtime/executor/codex_stream_bootstrap_buffering_test.go @@ -1746,3 +1746,69 @@ func TestCodexWebsocketsExecutor_BootstrapBuffering_StatusBearingErrorAfterTimeo t.Fatal("expected status-bearing error to be delivered in-stream") } } + +// The downstream handler closes the client connection as soon as the executor publishes an +// upstream disconnect, which cancels the request context. A usage-limit error must reach the +// conductor before that happens, otherwise the cancellation drops the error and the credential +// cooldown it carries (issue 6503). +func TestCodexWebsocketsExecutor_UnbufferedUpstreamErrorIsDeliveredBeforeDisconnectNotify(t *testing.T) { + frames := map[string]string{ + "error frame": `{"type":"error","status":429,"error":{"type":"usage_limit_reached","message":"The usage limit has been reached","resets_in_seconds":600}}`, + "response.failed": `{"type":"response.failed","response":{"error":{"type":"usage_limit_reached","status_code":429,"message":"The usage limit has been reached","resets_in_seconds":600}}}`, + } + for name, frame := range frames { + t.Run(name, func(t *testing.T) { + server := codexWebsocketServerHoldingConnection(t, frame) + defer server.Close() + + exec := NewCodexWebsocketsExecutor(codexBufferingConfig(false)) + exec.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} + + const sessionID = "issue-6503-session" + disconnectCh := exec.UpstreamDisconnectChan(sessionID) + if disconnectCh == nil { + t.Fatal("expected a disconnect channel") + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + canceled := make(chan struct{}) + go func() { + <-disconnectCh + cancel() + close(canceled) + }() + + req, opts := codexWebsocketRequest() + opts.Metadata = map[string]any{cliproxyexecutor.ExecutionSessionMetadataKey: sessionID} + result, err := exec.ExecuteStream(ctx, codexTestAuth(server.URL), req, opts) + if err != nil { + t.Fatalf("ExecuteStream() error = %v", err) + } + + // Give a premature disconnect notification time to cancel the context before the error is read. + select { + case <-canceled: + case <-time.After(200 * time.Millisecond): + } + + var streamErr error + for chunk := range result.Chunks { + if chunk.Err != nil { + streamErr = chunk.Err + } + } + if streamErr == nil { + t.Fatal("the upstream usage-limit error was dropped by the disconnect-triggered cancellation") + } + if got := statusCodeFromTestError(t, streamErr); got != http.StatusTooManyRequests { + t.Fatalf("status code = %d, want %d", got, http.StatusTooManyRequests) + } + select { + case <-canceled: + case <-time.After(5 * time.Second): + t.Fatal("downstream disconnect was never signalled after the error was delivered") + } + }) + } +} diff --git a/internal/runtime/executor/codex_websockets_stream.go b/internal/runtime/executor/codex_websockets_stream.go index 3c924248c..6269e51e7 100644 --- a/internal/runtime/executor/codex_websockets_stream.go +++ b/internal/runtime/executor/codex_websockets_stream.go @@ -623,38 +623,47 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr if wsErr, ok := parseCodexWebsocketErrorWithCooling(payload, e.modelLevelCooling()); ok { terminateReason = "upstream_error" terminateErr = wsErr + // The disconnect notification makes the downstream handler close the client and + // cancel the request context, so it must follow delivery of the error to the + // conductor; otherwise the cancellation can drop the error and its cooldown. if sess != nil { - e.invalidateUpstreamConn(sess, conn, "upstream_error", wsErr) + e.invalidateUpstreamConnWithoutDisconnectNotify(sess, conn, "upstream_error", wsErr) } if errClearReplay := clearCodexReasoningReplayOnWebsocketError(ctx, replayScope, payload); errClearReplay != nil { terminateErr = errClearReplay helps.RecordAPIWebsocketError(ctx, e.cfg, "replay_clear_error", errClearReplay) reporter.PublishFailure(ctx, errClearReplay) _ = send(cliproxyexecutor.StreamChunk{Err: errClearReplay}) + sess.notifyUpstreamDisconnect(wsErr) return } helps.RecordAPIWebsocketError(ctx, e.cfg, "upstream_error", wsErr) reporter.PublishFailure(ctx, wsErr) _ = send(cliproxyexecutor.StreamChunk{Err: wsErr}) + sess.notifyUpstreamDisconnect(wsErr) return } if streamErr, terminalBody, ok := codexTerminalFailureErrWithCooling(payload, e.modelLevelCooling()); ok { terminateReason = "upstream_error" terminateErr = streamErr + // Publish the downstream disconnect only after the error was delivered, see the + // upstream error branch above. if sess != nil { unlockStreamSession() - e.invalidateUpstreamConn(sess, conn, "terminal_failure", streamErr) + e.invalidateUpstreamConnWithoutDisconnectNotify(sess, conn, "terminal_failure", streamErr) } if errClearReplay := clearCodexReasoningReplayOnInvalidSignature(ctx, replayScope, streamErr.StatusCode(), terminalBody); errClearReplay != nil { terminateErr = errClearReplay helps.RecordAPIWebsocketError(ctx, e.cfg, "replay_clear_error", errClearReplay) reporter.PublishFailure(ctx, errClearReplay) _ = send(cliproxyexecutor.StreamChunk{Err: errClearReplay}) + sess.notifyUpstreamDisconnect(streamErr) return } helps.RecordAPIWebsocketError(ctx, e.cfg, "upstream_error", streamErr) reporter.PublishFailure(ctx, streamErr) _ = send(cliproxyexecutor.StreamChunk{Err: streamErr}) + sess.notifyUpstreamDisconnect(streamErr) return } diff --git a/sdk/cliproxy/auth/conductor_issue6503_test.go b/sdk/cliproxy/auth/conductor_issue6503_test.go new file mode 100644 index 000000000..fa30ad37b --- /dev/null +++ b/sdk/cliproxy/auth/conductor_issue6503_test.go @@ -0,0 +1,194 @@ +package auth + +import ( + "context" + "errors" + "net/http" + "testing" + "time" + + "github.com/google/uuid" + "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" +) + +type issue6503QuotaError struct { + retryAfter time.Duration +} + +func (issue6503QuotaError) Error() string { return "usage_limit_reached" } +func (issue6503QuotaError) StatusCode() int { return http.StatusTooManyRequests } +func (issue6503QuotaError) IsCredentialScoped() bool { return true } +func (e issue6503QuotaError) RetryAfter() *time.Duration { return &e.retryAfter } + +// issue6503StreamExecutor cancels the request context before returning, which mirrors +// the Codex websocket executor tearing down the downstream connection on an error frame. +type issue6503StreamExecutor struct { + cancel context.CancelFunc + err error +} + +func (*issue6503StreamExecutor) Identifier() string { return "issue6503" } + +func (*issue6503StreamExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, errors.New("not implemented") +} + +func (e *issue6503StreamExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + e.cancel() + return nil, e.err +} + +func (*issue6503StreamExecutor) Refresh(_ context.Context, auth *Auth) (*Auth, error) { + return auth, nil +} + +func (*issue6503StreamExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, errors.New("not implemented") +} + +func (*issue6503StreamExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, errors.New("not implemented") +} + +func runIssue6503CanceledStream(t *testing.T, executorErr error) (*Manager, string, string, error) { + t.Helper() + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + model := "issue6503-model-" + uuid.NewString() + auth := &Auth{ID: "issue6503-auth-" + uuid.NewString(), Provider: "issue6503", Status: StatusActive} + manager := NewManager(nil, nil, nil) + manager.SetRetryConfig(0, 0, 0) + manager.RegisterExecutor(&issue6503StreamExecutor{cancel: cancel, err: executorErr}) + registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth.ID) }) + if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { + t.Fatalf("Register() error = %v", errRegister) + } + _, errStream := manager.ExecuteStream(ctx, []string{"issue6503"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}) + return manager, auth.ID, model, errStream +} + +func TestManagerExecuteStreamRecordsQuotaCooldownWhenContextCanceled(t *testing.T) { + manager, authID, _, errStream := runIssue6503CanceledStream(t, issue6503QuotaError{retryAfter: 10 * time.Minute}) + if !errors.Is(errStream, context.Canceled) { + t.Fatalf("ExecuteStream() error = %v, want context.Canceled", errStream) + } + auth, ok := manager.GetByID(authID) + if !ok || auth == nil { + t.Fatalf("GetByID(%q) did not return auth", authID) + } + if !auth.Unavailable || !auth.NextRetryAfter.After(time.Now().Add(5*time.Minute)) { + t.Fatalf("auth was not cooled: unavailable=%t next=%v", auth.Unavailable, auth.NextRetryAfter) + } +} + +func TestManagerExecuteStreamPlainCancellationDoesNotCoolCredential(t *testing.T) { + manager, authID, model, errStream := runIssue6503CanceledStream(t, context.Canceled) + if !errors.Is(errStream, context.Canceled) { + t.Fatalf("ExecuteStream() error = %v, want context.Canceled", errStream) + } + auth, ok := manager.GetByID(authID) + if !ok || auth == nil { + t.Fatalf("GetByID(%q) did not return auth", authID) + } + if auth.Unavailable || !auth.NextRetryAfter.IsZero() { + t.Fatalf("auth was cooled: unavailable=%t next=%v", auth.Unavailable, auth.NextRetryAfter) + } + if state := auth.ModelStates[model]; state != nil && (state.Unavailable || !state.NextRetryAfter.IsZero() || state.Quota.Exceeded) { + t.Fatalf("model was cooled: %#v", state) + } +} + +// issue6503BootstrapExecutor returns a stream whose first chunk is an upstream quota error +// while the request context is already canceled, which is what the unbuffered Codex websocket +// path produces after the downstream connection was closed for the same upstream error. +type issue6503BootstrapExecutor struct { + issue6503StreamExecutor +} + +func (e *issue6503BootstrapExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + e.cancel() + chunks := make(chan cliproxyexecutor.StreamChunk, 1) + chunks <- cliproxyexecutor.StreamChunk{Err: e.err} + close(chunks) + return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil +} + +func TestManagerExecuteStreamRecordsQuotaCooldownFromCanceledBootstrapChunk(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + model := "issue6503-bootstrap-model-" + uuid.NewString() + auth := &Auth{ID: "issue6503-bootstrap-auth-" + uuid.NewString(), Provider: "issue6503", Status: StatusActive} + manager := NewManager(nil, nil, nil) + manager.SetRetryConfig(0, 0, 0) + manager.RegisterExecutor(&issue6503BootstrapExecutor{issue6503StreamExecutor{cancel: cancel, err: issue6503QuotaError{retryAfter: 10 * time.Minute}}}) + registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth.ID) }) + if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { + t.Fatalf("Register() error = %v", errRegister) + } + + _, errStream := manager.ExecuteStream(ctx, []string{"issue6503"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}) + if !errors.Is(errStream, context.Canceled) { + t.Fatalf("ExecuteStream() error = %v, want context.Canceled", errStream) + } + current, ok := manager.GetByID(auth.ID) + if !ok || current == nil { + t.Fatalf("GetByID(%q) did not return auth", auth.ID) + } + if !current.Unavailable || !current.NextRetryAfter.After(time.Now().Add(5*time.Minute)) { + t.Fatalf("auth was not cooled: unavailable=%t next=%v", current.Unavailable, current.NextRetryAfter) + } +} + +// issue6503RefreshExecutor fails the first stream with an unauthorized bootstrap chunk and, after +// the credential is refreshed, cancels the request context while returning a quota error. +type issue6503RefreshExecutor struct { + issue6503StreamExecutor + calls int +} + +func (e *issue6503RefreshExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + e.calls++ + if e.calls > 1 { + e.cancel() + return nil, e.err + } + chunks := make(chan cliproxyexecutor.StreamChunk, 1) + chunks <- cliproxyexecutor.StreamChunk{Err: &Error{HTTPStatus: http.StatusUnauthorized, Message: "token invalidated"}} + close(chunks) + return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil +} + +func TestManagerExecuteStreamRecordsQuotaCooldownFromCanceledRefreshRetry(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + model := "issue6503-refresh-model-" + uuid.NewString() + auth := &Auth{ + ID: "issue6503-refresh-auth-" + uuid.NewString(), + Provider: "issue6503", + Status: StatusActive, + Metadata: map[string]any{"access_token": "stale", "refresh_token": "refresh"}, + } + manager := NewManager(nil, nil, nil) + manager.SetRetryConfig(0, 0, 0) + manager.RegisterExecutor(&issue6503RefreshExecutor{issue6503StreamExecutor: issue6503StreamExecutor{cancel: cancel, err: issue6503QuotaError{retryAfter: 10 * time.Minute}}}) + registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth.ID) }) + if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { + t.Fatalf("Register() error = %v", errRegister) + } + + _, errStream := manager.ExecuteStream(ctx, []string{"issue6503"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}) + if !errors.Is(errStream, context.Canceled) { + t.Fatalf("ExecuteStream() error = %v, want context.Canceled", errStream) + } + current, ok := manager.GetByID(auth.ID) + if !ok || current == nil { + t.Fatalf("GetByID(%q) did not return auth", auth.ID) + } + if !current.Unavailable || !current.NextRetryAfter.After(time.Now().Add(5*time.Minute)) { + t.Fatalf("auth was not cooled: unavailable=%t next=%v", current.Unavailable, current.NextRetryAfter) + } +} diff --git a/sdk/cliproxy/auth/conductor_stream.go b/sdk/cliproxy/auth/conductor_stream.go index 565c8d2f7..9558f813e 100644 --- a/sdk/cliproxy/auth/conductor_stream.go +++ b/sdk/cliproxy/auth/conductor_stream.go @@ -99,10 +99,16 @@ func readStreamBootstrap(ctx context.Context, ch <-chan cliproxyexecutor.StreamC ok bool ) if ctx != nil { + // Prefer a chunk that is already available so an upstream error delivered together + // with a context cancellation is not discarded. select { - case <-ctx.Done(): - return nil, false, ctx.Err() case chunk, ok = <-ch: + default: + select { + case <-ctx.Done(): + return nil, false, ctx.Err() + case chunk, ok = <-ch: + } } } else { chunk, ok = <-ch @@ -120,6 +126,22 @@ func readStreamBootstrap(ctx context.Context, ch <-chan cliproxyexecutor.StreamC } } +// recordCanceledUpstreamCooldown keeps an upstream-confirmed cooldown (an error carrying +// RetryAfter) when the request context was canceled while the executor was returning, for +// example because the downstream connection was closed in response to the same upstream error. +// Errors without RetryAfter, such as plain client cancellations, are not recorded. +func (m *Manager) recordCanceledUpstreamCooldown(ctx context.Context, auth *Auth, provider, resultModel, routeModel string, opts cliproxyexecutor.Options, err error, ephemeral bool) { + retryAfter := retryAfterFromError(err) + if m == nil || auth == nil || retryAfter == nil { + return + } + action, okAction := matchRequestScopedErrorAction(auth, err, m.runtimeConfigSnapshot()) + result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, RouteModel: routeModel, Success: false, Error: resultErrorFromError(err), Options: opts, RetryAfter: retryAfter, CredentialVersion: auth.CredentialVersion, RegistrationEpoch: auth.RegistrationEpoch} + result.CredentialScope = isCredentialScopedError(err) + applyRequestScopedActionToResult(action, okAction, &result) + m.recordExecutionResult(context.WithoutCancel(ctx), result, auth, ephemeral) +} + func (m *Manager) wrapStreamResult(ctx context.Context, auth *Auth, provider, resultModel, routeModel string, headers http.Header, buffered []cliproxyexecutor.StreamChunk, remaining <-chan cliproxyexecutor.StreamChunk, aliasResult OAuthModelAliasResult, ephemeralResult bool, opts cliproxyexecutor.Options) *cliproxyexecutor.StreamResult { out := make(chan cliproxyexecutor.StreamChunk) streamStart := time.Now() @@ -248,6 +270,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi durationStream := time.Since(startStream) if errStream != nil { if errCtx := ctx.Err(); errCtx != nil { + m.recordCanceledUpstreamCooldown(ctx, auth, provider, resultModel, routeModel, execOpts, errStream, ephemeralResult) return nil, errCtx } if allowRetry && !ephemeralResult { @@ -269,6 +292,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi if errStream != nil { warnLogUpstreamFailure(ctx, entry, provider, execModel, auth, durationRetry, errStream) if errCtx := ctx.Err(); errCtx != nil { + m.recordCanceledUpstreamCooldown(ctx, auth, provider, resultModel, routeModel, execOpts, errStream, ephemeralResult) return nil, errCtx } } @@ -323,6 +347,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi } if bootstrapErr != nil { if errCtx := ctx.Err(); errCtx != nil { + m.recordCanceledUpstreamCooldown(ctx, auth, provider, resultModel, routeModel, execOpts, bootstrapErr, ephemeralResult) discardStreamChunks(streamResult.Chunks) return nil, errCtx } @@ -342,6 +367,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi retryErr = markUpstreamExecutionAttemptFromContext(ctx, retryErr) if retryErr != nil { if errCtx := ctx.Err(); errCtx != nil { + m.recordCanceledUpstreamCooldown(ctx, auth, provider, resultModel, routeModel, execOpts, retryErr, ephemeralResult) return nil, errCtx } bootstrapErr = retryErr