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

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")
+			}
+		})
+	}
+}

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
 			}
 

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)
+	}
+}

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