diff --git a/config.example.yaml b/config.example.yaml index 6b659eccf..ebc621719 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -283,6 +283,7 @@ nonstream-keepalive-interval: 0 # streaming: # keepalive-seconds: 15 # Default: 0 (disabled). <= 0 disables keep-alives. # bootstrap-retries: 1 # Default: 0 (disabled). Retries before first byte is sent. +# stream-connect-timeout-seconds: 20 # Default: 0 (disabled). Optional maximum wait for connection/stream establishment before failover. (Deprecated alias: stream-first-chunk-timeout-seconds). # Signature cache validation for thinking blocks (Antigravity/Claude). # When true (default), cached signatures are preferred and validated. diff --git a/internal/api/modules/amp/secret_test.go b/internal/api/modules/amp/secret_test.go index 17a75b15d..4fba32269 100644 --- a/internal/api/modules/amp/secret_test.go +++ b/internal/api/modules/amp/secret_test.go @@ -69,7 +69,7 @@ func TestMultiSourceSecret_CacheBehavior(t *testing.T) { t.Fatal(err) } - s := NewMultiSourceSecretWithPath("", p, 50*time.Millisecond) + s := NewMultiSourceSecretWithPath("", p, 500*time.Millisecond) // First read - should return v1 got1, err := s.Get(ctx) @@ -90,7 +90,7 @@ func TestMultiSourceSecret_CacheBehavior(t *testing.T) { } // After TTL expires, should see v2 - time.Sleep(60 * time.Millisecond) + time.Sleep(600 * time.Millisecond) got3, _ := s.Get(ctx) if got3 != "v2" { t.Fatalf("cache miss expected v2, got %s", got3) diff --git a/internal/clienterror/client_error.go b/internal/clienterror/client_error.go index 3a2f532d8..d666946de 100644 --- a/internal/clienterror/client_error.go +++ b/internal/clienterror/client_error.go @@ -80,6 +80,20 @@ func IsRequestFault(status int, err error) bool { status = statusErr.StatusCode() } } + // Payment and rate-limit statuses are authoritative even when an upstream + // pairs them with a generic invalid_request_error body. The credential must + // remain eligible for cooldown and rotation. + if status == http.StatusPaymentRequired || status == http.StatusTooManyRequests { + return false + } + // Authentication and invalid-or-expired-credential failures are caused by + // the credential, not the request: they must remain eligible for rotation + // by the shared mixed-auth loop even when the provider pairs them with a + // generic invalid-request identifier in the body. This must be checked + // before hasRequestFaultBody so the generic classifier cannot misfile them. + if (status == http.StatusUnauthorized || status == http.StatusForbidden) && hasAuthenticationErrorBody(err) { + return false + } if hasRequestFaultBody(err) { return true } @@ -112,6 +126,37 @@ func IsItemNotPersisted(message string) bool { strings.Contains(lower, "items are not persisted when `store` is set to false") } +// hasAuthenticationErrorBody reports whether err is a structured credential +// failure: an authentication_error type, an invalid or expired API-key code, +// or a Gemini UNAUTHENTICATED status. These are credential faults, not request +// faults, so they must never be classified as request faults on 401/403. +func hasAuthenticationErrorBody(err error) bool { + if err == nil { + return false + } + body := strings.TrimSpace(err.Error()) + if body == "" || !json.Valid([]byte(body)) { + return false + } + for _, path := range []string{"error.code", "code", "response.error.code", "body.error.code"} { + switch strings.ToLower(strings.TrimSpace(gjson.Get(body, path).String())) { + case "invalid_api_key", "incorrect_api_key", "expired_api_key": + return true + } + } + for _, path := range []string{"error.type", "type", "response.error.type", "body.error.type"} { + if errType := strings.ToLower(strings.TrimSpace(gjson.Get(body, path).String())); errType == "authentication_error" { + return true + } + } + for _, path := range []string{"error.status", "status", "response.error.status", "body.error.status"} { + if status := strings.ToLower(strings.TrimSpace(gjson.Get(body, path).String())); status == "unauthenticated" { + return true + } + } + return false +} + func hasRequestFaultBody(err error) bool { if err == nil { return false diff --git a/internal/clienterror/client_error_test.go b/internal/clienterror/client_error_test.go index db17df399..f1472a61c 100644 --- a/internal/clienterror/client_error_test.go +++ b/internal/clienterror/client_error_test.go @@ -6,6 +6,7 @@ import ( "fmt" "net/http" "net/url" + "strings" "testing" ) @@ -195,6 +196,124 @@ func TestIsRequestFault(t *testing.T) { {name: "transport", status: http.StatusBadGateway, err: errors.New("unexpected EOF")}, {name: "invalid JSON body", status: http.StatusBadGateway, err: errors.New(`{"error":`)}, {name: "nil", status: 0}, + { + // DeepSeek-style: 401 with authentication_error type plus the generic + // invalid_request_error code. Credential fault, not a request fault. + name: "401 authentication_error body", + status: http.StatusUnauthorized, + err: errors.New(`{"error":{"message":"Authentication Fails","type":"authentication_error","code":"invalid_request_error"}}`), + }, + { + name: "403 authentication_error body", + status: http.StatusForbidden, + err: errors.New(`{"error":{"type":"authentication_error","message":"Invalid token."}}`), + }, + { + name: "401 invalid_api_key code", + status: http.StatusUnauthorized, + err: errors.New(`{"error":{"message":"Invalid API key","type":"invalid_request_error","code":"invalid_api_key"}}`), + }, + { + name: "403 incorrect_api_key code", + status: http.StatusForbidden, + err: errors.New(`{"error":{"code":"incorrect_api_key","message":"bad key"}}`), + }, + { + name: "401 gemini unauthenticated status", + status: http.StatusUnauthorized, + err: errors.New(`{"error":{"code":16,"message":"Request had invalid authentication credentials.","status":"UNAUTHENTICATED"}}`), + }, + { + // Authentication marker with an invalid-request-type body: the auth + // carve-out must win over the generic classifier. + name: "401 auth marker beats generic invalid_request", + status: http.StatusUnauthorized, + err: errors.New(`{"error":{"type":"invalid_request","code":"invalid_api_key"}}`), + }, + { + // Same body on a non-auth status stays a request fault. + name: "authentication marker only honored on 401/403", + status: http.StatusBadRequest, + err: errors.New(`{"error":{"type":"invalid_request","code":"invalid_api_key"}}`), + want: true, + }, + { + // An auth-styled body behind a server status is not a request fault. + name: "authentication_error body behind upstream error", + status: http.StatusInternalServerError, + err: errors.New(`{"error":{"type":"authentication_error","message":"server glitch"}}`), + }, + { + // Wrapped status code: the 401 is extracted via errors.As, and even a + // non-splittable (prefixed) body stays a non-request credential fault. + name: "wrapped 401 authentication_error", + err: fmt.Errorf("upstream: %w", statusError{status: http.StatusUnauthorized, body: `{"error":{"type":"authentication_error"}}`}), + }, + { + name: "wrapped 400 invalid request stays request fault", + err: fmt.Errorf("upstream: %w", statusError{status: http.StatusBadRequest, body: "bad input"}), + want: true, + }, + { + // statusCoder form: the error carries both the 401 and the structured + // auth body directly (injected status stays 0, so it is read off the + // error). This mirrors how the shared loop feeds a Result error. + name: "statusCoder auth body", + err: statusError{status: http.StatusUnauthorized, body: `{"error":{"type":"authentication_error","code":"invalid_api_key"}}`}, + }, + { + // Negative: a statusCoder 400 with an auth-looking body is still a + // request fault once status wins over the body. + name: "statusCoder 400 auth body still request fault", + err: statusError{status: http.StatusBadRequest, body: `{"error":{"type":"authentication_error"}}`}, + want: true, + }, + { + // Large but valid bodies are still classified safely. + name: "large authentication_error body", + status: http.StatusForbidden, + err: errors.New(`{"error":{"type":"authentication_error","message":"` + strings.Repeat("x", 1<<16) + `"}}`), + }, + { + name: "429 stays credential domain", + status: http.StatusTooManyRequests, + err: errors.New(`{"error":{"type":"rate_limit_error","message":"slow down"}}`), + }, + { + // CPA parity: a rate-limit status is authoritative even when the body + // carries a generic invalid_request_error code. Quota, not request. + name: "429 with generic invalid_request_error code stays quota", + status: http.StatusTooManyRequests, + err: errors.New(`{"error":{"code":"invalid_request_error","message":"Rate Limit Reached","param":null,"type":"unknown_error"}}`), + }, + { + // CPA parity: a payment-required status is authoritative even when the + // body carries the generic invalid_request_error code. Quota/balance, + // not request. + name: "402 with generic invalid_request_error code stays payment", + status: http.StatusPaymentRequired, + err: errors.New(`{"error":{"message":"Insufficient Balance","type":"unknown_error","param":null,"code":"invalid_request_error"}}`), + }, + { + // A 402 surface with the invalid_request_error type spelling is still + // payment-scoped, never a request fault. + name: "402 with invalid_request_error type stays payment", + status: http.StatusPaymentRequired, + err: errors.New(`{"error":{"type":"invalid_request_error","code":"insufficient_balance","message":"low balance"}}`), + }, + { + // statusCoder form: quota/payment status is read off the error and the + // precedence still wins over the generic body code. + name: "statusCoder 429 with generic code stays quota", + err: statusError{status: http.StatusTooManyRequests, body: `{"error":{"code":"invalid_request_error","message":"quota"}}`}, + }, + { + // A 429 that is genuinely a 429 request-fault-looking body never flips + // to a request fault either, because the status stays authoritative. + name: "429 status stays authoritative over auth-looking body", + status: http.StatusTooManyRequests, + err: errors.New(`{"error":{"type":"authentication_error","code":"invalid_request_error"}}`), + }, } for _, tc := range tests { diff --git a/internal/config/sdk_config.go b/internal/config/sdk_config.go index a7f6c5ebb..50ba18f7a 100644 --- a/internal/config/sdk_config.go +++ b/internal/config/sdk_config.go @@ -82,4 +82,11 @@ type StreamingConfig struct { // to allow auth rotation / transient recovery. // <= 0 disables bootstrap retries. Default is 0. BootstrapRetries int `yaml:"bootstrap-retries,omitempty" json:"bootstrap-retries,omitempty"` + + // StreamConnectTimeoutSeconds controls the maximum time to wait for connection/stream establishment from an upstream stream before timing out and failing over. + // <= 0 disables stream connect timeout. Default is 0. + StreamConnectTimeoutSeconds int `yaml:"stream-connect-timeout-seconds,omitempty" json:"stream-connect-timeout-seconds,omitempty"` + + // StreamFirstChunkTimeoutSeconds is a deprecated alias for StreamConnectTimeoutSeconds. + StreamFirstChunkTimeoutSeconds int `yaml:"stream-first-chunk-timeout-seconds,omitempty" json:"stream-first-chunk-timeout-seconds,omitempty"` } diff --git a/internal/home/client.go b/internal/home/client.go index 62964f18b..dd73fb730 100644 --- a/internal/home/client.go +++ b/internal/home/client.go @@ -1618,9 +1618,8 @@ func newPluginSyncCancelableConn(ctx context.Context, conn net.Conn) net.Conn { go func() { select { case <-ctx.Done(): - if errDeadline := conn.SetDeadline(time.Now()); errDeadline != nil { - _ = conn.Close() - } + _ = conn.SetDeadline(time.Now()) + _ = conn.Close() case <-wrapped.done: } }() diff --git a/internal/interfaces/error_message.go b/internal/interfaces/error_message.go index 93fa3acbe..6c2c6e442 100644 --- a/internal/interfaces/error_message.go +++ b/internal/interfaces/error_message.go @@ -21,6 +21,11 @@ type ErrorMessage struct { // DirectResponse reports that Body and Headers were explicitly supplied by a trusted in-process component. DirectResponse bool + // TrustedDirectResponse reports that a DirectResponse originated locally + // (plugin/interceptor) and is safe to preserve through OpenAI sanitizers. + // Zero value false means the DirectResponse must be treated as untrusted. + TrustedDirectResponse bool + // Body contains a preformatted downstream response when DirectResponse is true. Body []byte diff --git a/internal/pluginhost/executor_route.go b/internal/pluginhost/executor_route.go index 0c3a5ec66..0709ea239 100644 --- a/internal/pluginhost/executor_route.go +++ b/internal/pluginhost/executor_route.go @@ -159,17 +159,28 @@ func wrapStreamEmptyCompletion(ctx context.Context, streamResult *coreexecutor.S case chunk, ok = <-src: } if !ok { - if !forwarding && len(buffered) == 0 { - _ = forward(coreexecutor.StreamChunk{Err: &coreauth.Error{ - Code: "empty_stream", - Message: "upstream stream closed before first payload", - Retryable: true, - }}) - return - } - if !forwarding && coreauth.IsEmptyCompletionPayload(streamChunkPayload(buffered)) { - _ = forward(coreexecutor.StreamChunk{Err: coreauth.EmptyCompletionError()}) - return + if !forwarding { + payloadBytes := 0 + for _, c := range buffered { + payloadBytes += len(c.Payload) + } + if payloadBytes == 0 { + // Zero-payload chunks are dropped downstream; a stream of only + // such chunks is an empty stream, not a successful completion. + _ = forward(coreexecutor.StreamChunk{Err: &coreauth.Error{ + Code: "empty_stream", + Message: "upstream stream closed before first payload", + Retryable: true, + }}) + return + } + // Judge with the incremental detector state instead of re-parsing + // the concatenated payload: separately chunked SSE frames do not + // reassemble into valid input for the payload-level check. + if detector.Finish() { + _ = forward(coreexecutor.StreamChunk{Err: coreauth.EmptyCompletionError()}) + return + } } _ = flush() return @@ -202,11 +213,26 @@ func wrapStreamEmptyCompletion(ctx context.Context, streamResult *coreexecutor.S return } } + if detector.IsTerminalEmpty() { + discardStreamChunks(src) + _ = forward(coreexecutor.StreamChunk{Err: coreauth.EmptyCompletionError()}) + return + } } }() return &coreexecutor.StreamResult{Chunks: wrapped, Headers: streamResult.Headers} } +func discardStreamChunks(ch <-chan coreexecutor.StreamChunk) { + if ch == nil { + return + } + go func() { + for range ch { + } + }() +} + func streamChunkPayload(chunks []coreexecutor.StreamChunk) []byte { var payload []byte for _, chunk := range chunks { diff --git a/internal/pluginhost/executor_route_stream_codex_test.go b/internal/pluginhost/executor_route_stream_codex_test.go new file mode 100644 index 000000000..58f1387cf --- /dev/null +++ b/internal/pluginhost/executor_route_stream_codex_test.go @@ -0,0 +1,67 @@ +package pluginhost + +import ( + "context" + "errors" + "testing" + + coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +// Regression tests mirrored from CLIProxyAPI PR #4881 follow-up +// (codex pullrequestreview-4943660625, pluginhost stream wrapper EOF +// handling). + +// TestWrapStreamEmptyCompletionRejectsZeroPayloadChunkStream is a regression +// guard for the zero-payload finding: zero-payload chunks made the buffer +// non-empty, so the EOF branch skipped the empty_stream error and flushed a +// client-invisible stream as success. +func TestWrapStreamEmptyCompletionRejectsZeroPayloadChunkStream(t *testing.T) { + src := make(chan coreexecutor.StreamChunk, 2) + src <- coreexecutor.StreamChunk{Payload: nil} + src <- coreexecutor.StreamChunk{Payload: []byte{}} + close(src) + + wrapped := wrapStreamEmptyCompletion(context.Background(), &coreexecutor.StreamResult{Chunks: src}) + first, ok := <-wrapped.Chunks + if !ok { + t.Fatal("wrapped stream closed without empty_stream error") + } + var authErr *coreauth.Error + if !errors.As(first.Err, &authErr) || authErr.Code != "empty_stream" || !authErr.Retryable { + t.Fatalf("first error = %#v, want retryable empty_stream", first.Err) + } + if len(first.Payload) != 0 { + t.Fatalf("first payload = %q, want error before client-visible bytes", first.Payload) + } + if _, ok = <-wrapped.Chunks; ok { + t.Fatal("wrapped stream emitted chunks after empty_stream error") + } +} + +// TestWrapStreamEmptyCompletionDetectsSplitUsageOnlyStream is a regression +// guard for the detector.Finish finding: the EOF branch used to re-parse the +// concatenated payload, and separately chunked SSE frames without trailing +// newlines concatenated into invalid input, so the empty check failed and an +// empty plugin stream was flushed as success. The incremental detector state +// now decides at EOF. +func TestWrapStreamEmptyCompletionDetectsSplitUsageOnlyStream(t *testing.T) { + src := make(chan coreexecutor.StreamChunk, 2) + src <- coreexecutor.StreamChunk{Payload: []byte("data: {\"choices\":[],\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":0}}")} + src <- coreexecutor.StreamChunk{Payload: []byte("data: [DONE]")} + close(src) + + wrapped := wrapStreamEmptyCompletion(context.Background(), &coreexecutor.StreamResult{Chunks: src}) + first, ok := <-wrapped.Chunks + if !ok { + t.Fatal("wrapped split usage-only stream closed without empty_completion error") + } + var authErr *coreauth.Error + if !errors.As(first.Err, &authErr) || authErr.Code != "empty_completion" { + t.Fatalf("first error = %v, want empty_completion", first.Err) + } + if len(first.Payload) != 0 { + t.Fatalf("first payload = %q, want no client-visible bytes before error", first.Payload) + } +} diff --git a/internal/pluginhost/executor_route_stream_test.go b/internal/pluginhost/executor_route_stream_test.go index b190f33ac..20d5bf0ba 100644 --- a/internal/pluginhost/executor_route_stream_test.go +++ b/internal/pluginhost/executor_route_stream_test.go @@ -13,7 +13,7 @@ import ( func TestWrapStreamEmptyCompletionWithholdsTerminalFrames(t *testing.T) { src := make(chan coreexecutor.StreamChunk, 2) - src <- coreexecutor.StreamChunk{Payload: []byte("data: {\"type\":\"response.completed\",\"response\":{\"output\":[],\"usage\":{\"output_tokens\":0}}}\n\n")} + src <- coreexecutor.StreamChunk{Payload: []byte("data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":{\"completion_tokens\":0}}\n\n")} src <- coreexecutor.StreamChunk{Payload: []byte("data: [DONE]\n\n")} close(src) @@ -58,8 +58,8 @@ func TestWrapStreamEmptyCompletionRejectsZeroChunkStream(t *testing.T) { func TestWrapStreamEmptyCompletionWithholdsSplitTerminalFrames(t *testing.T) { fragments := [][]byte{ []byte("da"), - []byte("ta: {\"type\":\"response.com"), - []byte("pleted\",\"response\":{\"output\":[],\"usage\":{\"output_tokens\":0}}}\n\n"), + []byte("ta: {\"choices\":[{\"delta\":{},\"finish_rea"), + []byte("son\":\"stop\"}],\"usage\":{\"completion_tokens\":0}}\n\n"), []byte("data: [DO"), []byte("NE]\n"), []byte("\n"), @@ -240,3 +240,98 @@ func assertStreamPayload(t *testing.T, chunks <-chan coreexecutor.StreamChunk, w t.Fatalf("timed out waiting for payload %q", want) } } + +func TestWrapStreamEmptyCompletionStopsAtTerminalEmptyMarkersWithoutChannelClose(t *testing.T) { + testCases := []struct { + name string + payload []byte + }{ + { + name: "openai_done_on_open_channel", + payload: []byte("data: [DONE]\n\n"), + }, + { + name: "claude_message_stop_on_open_channel", + payload: []byte("event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"), + }, + { + name: "claude_data_only_message_stop_on_open_channel", + payload: []byte("data: {\"type\":\"message_stop\"}\n\n"), + }, + { + name: "gemini_empty_stop_on_open_channel", + payload: []byte("data: {\"candidates\":[{\"finishReason\":\"STOP\"}]}\n\n"), + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + src := make(chan coreexecutor.StreamChunk, 2) + src <- coreexecutor.StreamChunk{Payload: tc.payload} + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + wrapped := wrapStreamEmptyCompletion(ctx, &coreexecutor.StreamResult{Chunks: src}) + select { + case first, ok := <-wrapped.Chunks: + if !ok { + t.Fatal("wrapped stream closed without emitting error") + } + var authErr *coreauth.Error + if !errors.As(first.Err, &authErr) || authErr.Code != "empty_completion" { + t.Fatalf("first error = %v, want empty_completion error", first.Err) + } + case <-ctx.Done(): + t.Fatal("timed out waiting for empty_completion error; stream blocked on open channel") + } + + select { + case chunk, ok := <-wrapped.Chunks: + if ok { + t.Fatalf("wrapped stream emitted unexpected trailing chunk: %#v", chunk) + } + case <-time.After(500 * time.Millisecond): + t.Fatal("wrapped stream did not close after empty_completion error") + } + }) + } +} + +func TestWrapStreamEmptyCompletionDrainsSourceAfterTerminalEmpty(t *testing.T) { + src := make(chan coreexecutor.StreamChunk) + producerDone := make(chan struct{}) + + go func() { + defer close(producerDone) + // Send terminal empty chunk (OpenAI [DONE]) + src <- coreexecutor.StreamChunk{Payload: []byte("data: [DONE]\n\n")} + // Send trailing chunk on unbuffered channel + src <- coreexecutor.StreamChunk{Payload: []byte("trailing chunk")} + close(src) + }() + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + wrapped := wrapStreamEmptyCompletion(ctx, &coreexecutor.StreamResult{Chunks: src}) + select { + case first, ok := <-wrapped.Chunks: + if !ok { + t.Fatal("wrapped stream closed without emitting error") + } + var authErr *coreauth.Error + if !errors.As(first.Err, &authErr) || authErr.Code != "empty_completion" { + t.Fatalf("first error = %v, want empty_completion error", first.Err) + } + case <-ctx.Done(): + t.Fatal("timed out waiting for empty_completion error") + } + + select { + case <-producerDone: + // Success: producer unblocked because src was drained + case <-time.After(time.Second): + t.Fatal("producer remained blocked after terminal empty return; source was not drained") + } +} diff --git a/internal/registry/model_registry.go b/internal/registry/model_registry.go index eac36ddf2..cc8f7a9ad 100644 --- a/internal/registry/model_registry.go +++ b/internal/registry/model_registry.go @@ -785,6 +785,23 @@ func (r *ModelRegistry) ResumeClientModel(clientID, modelID string) { log.Debugf("Resumed client %s for model %s", clientID, modelID) } +// GetClientModelSuspensionReason returns the reason a client model was suspended, or empty string if not suspended. +func (r *ModelRegistry) GetClientModelSuspensionReason(clientID, modelID string) string { + clientID = strings.TrimSpace(clientID) + modelID = strings.TrimSpace(modelID) + if clientID == "" || modelID == "" { + return "" + } + r.mutex.RLock() + defer r.mutex.RUnlock() + + registration, exists := r.models[modelID] + if !exists || registration == nil || registration.SuspendedClients == nil { + return "" + } + return registration.SuspendedClients[clientID] +} + // ClientSupportsModel reports whether the client registered support for modelID. func (r *ModelRegistry) ClientSupportsModel(clientID, modelID string) bool { clientID = strings.TrimSpace(clientID) diff --git a/internal/runtime/executor/antigravity_executor_credits_test.go b/internal/runtime/executor/antigravity_executor_credits_test.go index d3cfd9cb5..73fd79007 100644 --- a/internal/runtime/executor/antigravity_executor_credits_test.go +++ b/internal/runtime/executor/antigravity_executor_credits_test.go @@ -22,10 +22,10 @@ import ( ) func resetAntigravityCreditsRetryState() { - antigravityCreditsFailureByAuth = sync.Map{} - antigravityShortCooldownByAuth = sync.Map{} - antigravityCreditsBalanceByAuth = sync.Map{} - antigravityCreditsHintRefreshByID = sync.Map{} + antigravityCreditsFailureByAuth.Clear() + antigravityShortCooldownByAuth.Clear() + antigravityCreditsBalanceByAuth.Clear() + antigravityCreditsHintRefreshByID.Clear() } type closeSignalReadCloser struct { diff --git a/internal/runtime/executor/codex_websockets_executor_test.go b/internal/runtime/executor/codex_websockets_executor_test.go index 1f4a676a3..c3075e394 100644 --- a/internal/runtime/executor/codex_websockets_executor_test.go +++ b/internal/runtime/executor/codex_websockets_executor_test.go @@ -682,6 +682,76 @@ func TestCodexWebsocketsExecuteStreamPropagatesUpstreamErrorForDownstreamWebsock } } +func TestCodexWebsocketsExecuteStreamDownstreamWebsocketResponseIncomplete(t *testing.T) { + upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} + incompletePayload := []byte(`{"type":"response.incomplete","response":{"id":"resp_1","status":"incomplete"}}`) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + t.Errorf("upgrade websocket: %v", err) + return + } + defer func() { _ = conn.Close() }() + + if _, _, errRead := conn.ReadMessage(); errRead != nil { + t.Errorf("read first message: %v", errRead) + return + } + if errWrite := conn.WriteMessage(websocket.TextMessage, incompletePayload); errWrite != nil { + t.Errorf("write incomplete message: %v", errWrite) + return + } + // Keep connection open without closing or sending more messages + _, _, _ = conn.ReadMessage() + })) + defer server.Close() + + exec := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}}) + auth := &cliproxyauth.Auth{Attributes: map[string]string{"api_key": "sk-test", "base_url": server.URL}} + req := cliproxyexecutor.Request{ + Model: "gpt-5-codex", + Payload: []byte(`{"model":"gpt-5-codex","input":[{"type":"message","role":"user","content":"hello"}]}`), + } + opts := cliproxyexecutor.Options{ + Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: "sess-incomplete-test", + }, + SourceFormat: sdktranslator.FromString("openai-response"), + ResponseFormat: sdktranslator.FromString("openai-response"), + } + ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background()) + + result, err := exec.ExecuteStream(ctx, auth, req, opts) + if err != nil { + t.Fatalf("ExecuteStream() error = %v", err) + } + + select { + case chunk, ok := <-result.Chunks: + if !ok { + t.Fatal("stream closed before response.incomplete chunk") + } + if !bytes.Contains(chunk.Payload, []byte("response.incomplete")) { + t.Fatalf("chunk payload = %q, want response.incomplete", chunk.Payload) + } + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for response.incomplete stream chunk") + } + + select { + case chunk, ok := <-result.Chunks: + if ok { + t.Fatalf("unexpected chunk after terminal event: %#v", chunk) + } + case <-time.After(1 * time.Second): + t.Fatal("result.Chunks not closed after response.incomplete; executor hung reading open socket") + } + + req2Ctx, req2Cancel := context.WithTimeout(ctx, 1*time.Second) + defer req2Cancel() + _, _ = exec.ExecuteStream(req2Ctx, auth, req, opts) +} + func TestSendTerminalWebsocketReadInvalidatesBeforeWaitingForCapacity(t *testing.T) { terminalErr := &websocket.CloseError{Code: websocket.CloseMessageTooBig} diff --git a/internal/runtime/executor/codex_websockets_stream.go b/internal/runtime/executor/codex_websockets_stream.go index 505081f09..84c7f616c 100644 --- a/internal/runtime/executor/codex_websockets_stream.go +++ b/internal/runtime/executor/codex_websockets_stream.go @@ -343,6 +343,7 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr terminateErr = wsErr if sess != nil { e.invalidateUpstreamConn(sess, conn, "upstream_error", wsErr) + unlockStreamSession() } if errClearReplay := clearCodexReasoningReplayOnWebsocketError(ctx, replayScope, payload); errClearReplay != nil { terminateErr = errClearReplay @@ -360,8 +361,8 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr terminateReason = "upstream_error" terminateErr = streamErr if sess != nil { - unlockStreamSession() e.invalidateUpstreamConn(sess, conn, "terminal_failure", streamErr) + unlockStreamSession() } if errClearReplay := clearCodexReasoningReplayOnInvalidSignature(ctx, replayScope, streamErr.StatusCode(), terminalBody); errClearReplay != nil { terminateErr = errClearReplay @@ -377,7 +378,7 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr } eventType := gjson.GetBytes(payload, "type").String() - isTerminalEvent := eventType == "response.completed" || eventType == "response.done" || eventType == "error" + isTerminalEvent := eventType == "response.completed" || eventType == "response.done" || eventType == "response.incomplete" || eventType == "response.failed" || eventType == "error" if eventType == "response.output_item.done" { collectCodexOutputItemDone(payload, outputItemsByIndex, &outputItemsFallback) } @@ -399,6 +400,10 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr return } if isTerminalEvent { + if (eventType == "error" || eventType == "response.incomplete" || eventType == "response.failed") && sess != nil { + e.invalidateUpstreamConn(sess, conn, "terminal_error", nil) + unlockStreamSession() + } return } continue @@ -419,7 +424,7 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr return } } - if eventType == "response.completed" || eventType == "response.done" { + if isTerminalEvent { return } } diff --git a/internal/runtime/executor/home_codex_terminal_test.go b/internal/runtime/executor/home_codex_terminal_test.go index 14d068c15..cd9019290 100644 --- a/internal/runtime/executor/home_codex_terminal_test.go +++ b/internal/runtime/executor/home_codex_terminal_test.go @@ -43,6 +43,7 @@ func TestHomeCodexTerminalStreamFailureUsesFreshDispatchOnNextRequest(t *testing } if connections.Add(1) == 1 { _ = conn.WriteJSON(map[string]any{"type": "response.created", "response": map[string]any{"id": "response-1"}}) + _ = conn.WriteJSON(map[string]any{"type": "response.output_text.delta", "delta": "streaming"}) _ = conn.WriteJSON(map[string]any{"type": "error", "status": http.StatusBadGateway, "error": map[string]any{"message": "terminal failure"}}) } else { writeCompletion := func() { diff --git a/internal/runtime/executor/websocket_session_target_test.go b/internal/runtime/executor/websocket_session_target_test.go index 163473cc9..bc23c6813 100644 --- a/internal/runtime/executor/websocket_session_target_test.go +++ b/internal/runtime/executor/websocket_session_target_test.go @@ -199,6 +199,7 @@ func TestWebsocketRetryBindFailureClearsActiveSessionState(t *testing.T) { run: func(t *testing.T, baseURL string) (func(cliproxyexecutor.Options) error, *codexWebsocketSession) { executor := NewXAIWebsocketsExecutor(&config.Config{}) executor.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} + executor.idStore = &xaiWebsocketIDStateStore{sessions: make(map[string]*xaiWebsocketIDState)} auth := &cliproxyauth.Auth{ID: "retry-bind-xai", Provider: "xai", Attributes: map[string]string{"base_url": baseURL, "websockets": "true"}, Metadata: map[string]any{"access_token": "test-token"}} req := cliproxyexecutor.Request{Model: "grok-4", Payload: []byte(`{"model":"grok-4","input":[{"type":"message","role":"user","content":"hello"}]}`)} primed := false diff --git a/internal/translator/openai/gemini/openai_gemini_request.go b/internal/translator/openai/gemini/openai_gemini_request.go index e7cf18a3c..19b0077ec 100644 --- a/internal/translator/openai/gemini/openai_gemini_request.go +++ b/internal/translator/openai/gemini/openai_gemini_request.go @@ -6,9 +6,9 @@ package gemini import ( - "crypto/rand" + "crypto/sha256" + "encoding/hex" "fmt" - "math/big" "strings" "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" @@ -17,6 +17,12 @@ import ( "github.com/tidwall/sjson" ) +func deriveDeterministicToolID(kind string, msgIdx, partIdx int, funcName, payload string) string { + seed := fmt.Sprintf("%s|%d|%d|%s|%s", kind, msgIdx, partIdx, funcName, payload) + hash := sha256.Sum256([]byte(seed)) + return "call_" + hex.EncodeToString(hash[:])[:24] +} + // ConvertGeminiRequestToOpenAI parses and transforms a Gemini API request into OpenAI Chat Completions API format. // It extracts the model name, generation config, message contents, and tool declarations // from the raw JSON request and returns them in the format expected by the OpenAI API. @@ -27,18 +33,6 @@ func ConvertGeminiRequestToOpenAI(modelName string, inputRawJSON []byte, stream root := gjson.ParseBytes(rawJSON) - // Helper for generating tool call IDs in the form: call_ - genToolCallID := func() string { - const letters = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789" - var b strings.Builder - // 24 chars random suffix - for i := 0; i < 24; i++ { - n, _ := rand.Int(rand.Reader, big.NewInt(int64(len(letters)))) - b.WriteByte(letters[n.Int64()]) - } - return "call_" + b.String() - } - // Model mapping out, _ = sjson.SetBytes(out, "model", modelName) @@ -139,8 +133,7 @@ func ConvertGeminiRequestToOpenAI(modelName string, inputRawJSON []byte, stream messageCapacity++ } messageItems := translatorcommon.NewRawArrayItems(messageCapacity) - var toolCallIDs []string // Track tool call IDs for matching with tool results - toolCallConsumeIdx := 0 + missingCallIDs := make(map[string][]string) // System instruction -> OpenAI system message // Gemini may provide `systemInstruction` or `system_instruction`; support both keys. @@ -180,6 +173,7 @@ func ConvertGeminiRequestToOpenAI(modelName string, inputRawJSON []byte, stream } if contents := root.Get("contents"); contents.Exists() && contents.IsArray() { + msgIdx := 0 contents.ForEach(func(_, content gjson.Result) bool { role := content.Get("role").String() parts := content.Get("parts") @@ -198,6 +192,7 @@ func ConvertGeminiRequestToOpenAI(modelName string, inputRawJSON []byte, stream toolCallItems := make([][]byte, 0, 2) if parts.Exists() && parts.IsArray() { + partIdx := 0 parts.ForEach(func(_, part gjson.Result) bool { // Handle text parts if text := part.Get("text"); text.Exists() { @@ -220,15 +215,17 @@ func ConvertGeminiRequestToOpenAI(modelName string, inputRawJSON []byte, stream // Handle function calls (Gemini) -> tool calls (OpenAI) if functionCall := part.Get("functionCall"); functionCall.Exists() { + funcName := functionCall.Get("name").String() toolCallID := explicitGeminiToolID(functionCall) if toolCallID == "" { - toolCallID = genToolCallID() + argsRaw := functionCall.Get("args").Raw + toolCallID = deriveDeterministicToolID("call", msgIdx, partIdx, funcName, argsRaw) + missingCallIDs[funcName] = append(missingCallIDs[funcName], toolCallID) } - toolCallIDs = append(toolCallIDs, toolCallID) toolCall := []byte(`{"id":"","type":"function","function":{"name":"","arguments":""}}`) toolCall, _ = sjson.SetBytes(toolCall, "id", toolCallID) - toolCall, _ = sjson.SetBytes(toolCall, "function.name", functionCall.Get("name").String()) + toolCall, _ = sjson.SetBytes(toolCall, "function.name", funcName) // Convert args to arguments JSON string if args := functionCall.Get("args"); args.Exists() { @@ -254,22 +251,23 @@ func ConvertGeminiRequestToOpenAI(modelName string, inputRawJSON []byte, stream } } + funcName := functionResponse.Get("name").String() if toolCallID := explicitGeminiToolID(functionResponse); toolCallID != "" { toolMsg, _ = sjson.SetBytes(toolMsg, "tool_call_id", toolCallID) - if toolCallConsumeIdx < len(toolCallIDs) && toolCallIDs[toolCallConsumeIdx] == toolCallID { - toolCallConsumeIdx++ - } - } else if toolCallConsumeIdx < len(toolCallIDs) { - toolMsg, _ = sjson.SetBytes(toolMsg, "tool_call_id", toolCallIDs[toolCallConsumeIdx]) - toolCallConsumeIdx++ + } else if queue := missingCallIDs[funcName]; len(queue) > 0 { + toolCallID := queue[0] + missingCallIDs[funcName] = queue[1:] + toolMsg, _ = sjson.SetBytes(toolMsg, "tool_call_id", toolCallID) } else { - // Generate a tool call ID if none available - toolMsg, _ = sjson.SetBytes(toolMsg, "tool_call_id", genToolCallID()) + respRaw := functionResponse.Get("response").Raw + standaloneID := deriveDeterministicToolID("response", msgIdx, partIdx, funcName, respRaw) + toolMsg, _ = sjson.SetBytes(toolMsg, "tool_call_id", standaloneID) } messageItems = append(messageItems, toolMsg) } + partIdx++ return true }) } @@ -289,6 +287,7 @@ func ConvertGeminiRequestToOpenAI(modelName string, inputRawJSON []byte, stream } messageItems = append(messageItems, msg) + msgIdx++ return true }) } diff --git a/internal/translator/openai/gemini/openai_gemini_request_test.go b/internal/translator/openai/gemini/openai_gemini_request_test.go index f1e2e7092..9fc05c1fb 100644 --- a/internal/translator/openai/gemini/openai_gemini_request_test.go +++ b/internal/translator/openai/gemini/openai_gemini_request_test.go @@ -169,3 +169,306 @@ func TestConvertGeminiRequestToOpenAI_SplitsNonImageInlineDataByMIME(t *testing. t.Fatalf("non-image inlineData must not be converted to image_url. Output: %s", string(out)) } } + +func TestConvertGeminiRequestToOpenAI_Deterministic100Invocations(t *testing.T) { + inputJSON := []byte(`{ + "contents": [ + { + "role": "model", + "parts": [ + {"functionCall": {"name": "search", "args": {"q": "golang"}}}, + {"functionCall": {"name": "fetch", "args": {"url": "https://example.com"}}} + ] + }, + { + "role": "function", + "parts": [ + {"functionResponse": {"name": "search", "response": {"results": ["a", "b"]}}}, + {"functionResponse": {"name": "fetch", "response": {"body": "hello"}}} + ] + } + ] + }`) + + firstOut := ConvertGeminiRequestToOpenAI("test-model", inputJSON, false) + for i := 0; i < 100; i++ { + out := ConvertGeminiRequestToOpenAI("test-model", inputJSON, false) + if string(out) != string(firstOut) { + t.Fatalf("invocation %d produced different bytes:\ngot: %s\nwant: %s", i, string(out), string(firstOut)) + } + } +} + +func TestConvertGeminiRequestToOpenAI_DuplicateCallsGetDistinctStableIDs(t *testing.T) { + inputJSON := []byte(`{ + "contents": [ + { + "role": "model", + "parts": [ + {"functionCall": {"name": "ping", "args": {"host": "localhost"}}}, + {"functionCall": {"name": "ping", "args": {"host": "localhost"}}} + ] + } + ] + }`) + + out1 := ConvertGeminiRequestToOpenAI("test-model", inputJSON, false) + id1 := gjson.GetBytes(out1, "messages.0.tool_calls.0.id").String() + id2 := gjson.GetBytes(out1, "messages.0.tool_calls.1.id").String() + + if id1 == "" || id2 == "" { + t.Fatalf("expected non-empty IDs, got id1=%q, id2=%q", id1, id2) + } + if id1 == id2 { + t.Fatalf("expected distinct IDs for duplicate calls, got id1 == id2 == %q", id1) + } + + out2 := ConvertGeminiRequestToOpenAI("test-model", inputJSON, false) + if string(out1) != string(out2) { + t.Fatalf("duplicate call translation is not deterministic across runs:\nout1: %s\nout2: %s", string(out1), string(out2)) + } +} + +func TestConvertGeminiRequestToOpenAI_SameNameResponsesConsumeFIFO(t *testing.T) { + inputJSON := []byte(`{ + "contents": [ + { + "role": "model", + "parts": [ + {"functionCall": {"name": "fnA", "args": {"step": 1}}}, + {"functionCall": {"name": "fnA", "args": {"step": 2}}} + ] + }, + { + "role": "function", + "parts": [ + {"functionResponse": {"name": "fnA", "response": {"res": 1}}}, + {"functionResponse": {"name": "fnA", "response": {"res": 2}}} + ] + } + ] + }`) + + out := ConvertGeminiRequestToOpenAI("test-model", inputJSON, false) + callID1 := gjson.GetBytes(out, "messages.0.tool_calls.0.id").String() + callID2 := gjson.GetBytes(out, "messages.0.tool_calls.1.id").String() + + respID1 := gjson.GetBytes(out, "messages.1.tool_call_id").String() + respID2 := gjson.GetBytes(out, "messages.2.tool_call_id").String() + + if respID1 != callID1 { + t.Fatalf("first response tool_call_id = %q, want callID1 %q", respID1, callID1) + } + if respID2 != callID2 { + t.Fatalf("second response tool_call_id = %q, want callID2 %q", respID2, callID2) + } +} + +func TestConvertGeminiRequestToOpenAI_InterleavedFunctionsPairByNameAndFIFO(t *testing.T) { + inputJSON := []byte(`{ + "contents": [ + { + "role": "model", + "parts": [ + {"functionCall": {"name": "fnA", "args": {"id": 1}}}, + {"functionCall": {"name": "fnB", "args": {"id": 1}}}, + {"functionCall": {"name": "fnA", "args": {"id": 2}}} + ] + }, + { + "role": "function", + "parts": [ + {"functionResponse": {"name": "fnB", "response": {"out": 1}}}, + {"functionResponse": {"name": "fnA", "response": {"out": 1}}}, + {"functionResponse": {"name": "fnA", "response": {"out": 2}}} + ] + } + ] + }`) + + out := ConvertGeminiRequestToOpenAI("test-model", inputJSON, false) + fnA_call1 := gjson.GetBytes(out, "messages.0.tool_calls.0.id").String() + fnB_call1 := gjson.GetBytes(out, "messages.0.tool_calls.1.id").String() + fnA_call2 := gjson.GetBytes(out, "messages.0.tool_calls.2.id").String() + + resp_fnB1 := gjson.GetBytes(out, "messages.1.tool_call_id").String() + resp_fnA1 := gjson.GetBytes(out, "messages.2.tool_call_id").String() + resp_fnA2 := gjson.GetBytes(out, "messages.3.tool_call_id").String() + + if resp_fnB1 != fnB_call1 { + t.Fatalf("fnB response paired with %q, want %q", resp_fnB1, fnB_call1) + } + if resp_fnA1 != fnA_call1 { + t.Fatalf("first fnA response paired with %q, want %q", resp_fnA1, fnA_call1) + } + if resp_fnA2 != fnA_call2 { + t.Fatalf("second fnA response paired with %q, want %q", resp_fnA2, fnA_call2) + } +} + +func TestConvertGeminiRequestToOpenAI_ExplicitIDsPreserved(t *testing.T) { + inputJSON := []byte(`{ + "contents": [ + { + "role": "model", + "parts": [ + {"functionCall": {"name": "calc", "id": "call_explicit_100", "args": {"x": 5}}} + ] + }, + { + "role": "function", + "parts": [ + {"functionResponse": {"name": "calc", "id": "call_explicit_100", "response": {"ans": 10}}} + ] + } + ] + }`) + + out := ConvertGeminiRequestToOpenAI("test-model", inputJSON, false) + callID := gjson.GetBytes(out, "messages.0.tool_calls.0.id").String() + respID := gjson.GetBytes(out, "messages.1.tool_call_id").String() + + if callID != "call_explicit_100" { + t.Fatalf("callID = %q, want %q", callID, "call_explicit_100") + } + if respID != "call_explicit_100" { + t.Fatalf("respID = %q, want %q", respID, "call_explicit_100") + } +} + +func TestConvertGeminiRequestToOpenAI_UnmatchedResponseGetsStableStandaloneID(t *testing.T) { + inputJSON := []byte(`{ + "contents": [ + { + "role": "function", + "parts": [ + {"functionResponse": {"name": "orphan_func", "response": {"status": "ok"}}} + ] + } + ] + }`) + + out1 := ConvertGeminiRequestToOpenAI("test-model", inputJSON, false) + id1 := gjson.GetBytes(out1, "messages.0.tool_call_id").String() + + if !strings.HasPrefix(id1, "call_") { + t.Fatalf("standalone ID = %q, want call_ prefix", id1) + } + + out2 := ConvertGeminiRequestToOpenAI("test-model", inputJSON, false) + id2 := gjson.GetBytes(out2, "messages.0.tool_call_id").String() + + if id1 != id2 { + t.Fatalf("standalone response ID is not stable: id1=%q, id2=%q", id1, id2) + } +} + +func TestConvertGeminiRequestToOpenAI_MalformedFieldsNoPanicDeterministic(t *testing.T) { + malformedInputs := [][]byte{ + []byte(`{"contents":[{"role":"model","parts":[{"functionCall":{}}]}]}`), + []byte(`{"contents":[{"role":"function","parts":[{"functionResponse":{}}]}]}`), + []byte(`{"contents":[{"role":"model","parts":[{"functionCall":{"name":123,"args":"invalid"}}]}]}`), + []byte(`{"contents":[{"role":"function","parts":[{"functionResponse":{"name":true,"response":null}}]}]}`), + } + + for i, input := range malformedInputs { + out1 := ConvertGeminiRequestToOpenAI("test-model", input, false) + out2 := ConvertGeminiRequestToOpenAI("test-model", input, false) + if string(out1) != string(out2) { + t.Fatalf("malformed input index %d is not deterministic:\nout1: %s\nout2: %s", i, string(out1), string(out2)) + } + } +} + +func TestConvertGeminiRequestToOpenAI_RealisticCallResultRoundTrip(t *testing.T) { + inputJSON := []byte(`{ + "contents": [ + { + "role": "user", + "parts": [{"text": "What is the weather in Tokyo?"}] + }, + { + "role": "model", + "parts": [ + {"text": "Checking weather..."}, + {"functionCall": {"name": "get_weather", "args": {"city": "Tokyo"}}} + ] + }, + { + "role": "function", + "parts": [ + {"functionResponse": {"name": "get_weather", "response": {"temp": "22C", "condition": "Sunny"}}} + ] + } + ] + }`) + + out := ConvertGeminiRequestToOpenAI("gpt-4o", inputJSON, false) + + modelMsgRole := gjson.GetBytes(out, "messages.1.role").String() + if modelMsgRole != "assistant" { + t.Fatalf("messages.1.role = %q, want assistant", modelMsgRole) + } + + callID := gjson.GetBytes(out, "messages.1.tool_calls.0.id").String() + if callID == "" || !strings.HasPrefix(callID, "call_") { + t.Fatalf("callID = %q, want valid call_ prefix", callID) + } + + toolMsgRole := gjson.GetBytes(out, "messages.2.role").String() + if toolMsgRole != "tool" { + t.Fatalf("messages.2.role = %q, want tool", toolMsgRole) + } + + toolMsgCallID := gjson.GetBytes(out, "messages.2.tool_call_id").String() + if toolMsgCallID != callID { + t.Fatalf("tool message tool_call_id = %q, want matching callID %q", toolMsgCallID, callID) + } +} + +func TestConvertGeminiRequestToOpenAI_MultiPartToolResponseFollowedByUnmatchedResponse(t *testing.T) { + inputJSON := []byte(`{ + "contents": [ + { + "role": "user", + "parts": [{"text": "Hello"}] + }, + { + "role": "model", + "parts": [ + {"functionCall": {"name": "funcA", "args": {"x": 1}}} + ] + }, + { + "role": "function", + "parts": [ + {"functionResponse": {"name": "funcA", "response": {"res": 1}}}, + {"functionResponse": {"name": "unmatchedA", "response": {"res": 2}}} + ] + }, + { + "role": "function", + "parts": [ + {"functionResponse": {"name": "unmatchedB", "response": {"res": 3}}} + ] + } + ] + }`) + + out1 := ConvertGeminiRequestToOpenAI("test-model", inputJSON, false) + out2 := ConvertGeminiRequestToOpenAI("test-model", inputJSON, false) + + if string(out1) != string(out2) { + t.Fatalf("multi-part response translation not deterministic across runs") + } + + unmatched1_ID := gjson.GetBytes(out1, "messages.3.tool_call_id").String() + unmatched2_ID := gjson.GetBytes(out1, "messages.5.tool_call_id").String() + + if unmatched1_ID == "" || unmatched2_ID == "" { + t.Fatalf("expected non-empty standalone tool_call_ids, got %q and %q", unmatched1_ID, unmatched2_ID) + } + if unmatched1_ID == unmatched2_ID { + t.Fatalf("expected distinct standalone tool_call_ids across messages, got %q", unmatched1_ID) + } +} diff --git a/internal/watcher/watcher_test.go b/internal/watcher/watcher_test.go index fee504689..a4d89b7ae 100644 --- a/internal/watcher/watcher_test.go +++ b/internal/watcher/watcher_test.go @@ -1515,11 +1515,13 @@ func TestNormalizeAuthNil(t *testing.T) { // stubStore implements coreauth.Store plus watcher-specific persistence helpers. type stubStore struct { + mu sync.Mutex authDir string - cfgPersisted int32 - authPersisted int32 + cfgPersisted int + authPersisted int lastAuthMessage string lastAuthPaths []string + persisted chan struct{} } func (s *stubStore) List(context.Context) ([]*coreauth.Auth, error) { return nil, nil } @@ -1528,17 +1530,39 @@ func (s *stubStore) Save(context.Context, *coreauth.Auth) (string, error) { } func (s *stubStore) Delete(context.Context, string) error { return nil } func (s *stubStore) PersistConfig(context.Context) error { - atomic.AddInt32(&s.cfgPersisted, 1) + s.mu.Lock() + s.cfgPersisted++ + s.mu.Unlock() + s.signalPersisted() return nil } func (s *stubStore) PersistAuthFiles(_ context.Context, message string, paths ...string) error { - atomic.AddInt32(&s.authPersisted, 1) + s.mu.Lock() + defer s.mu.Unlock() s.lastAuthMessage = message - s.lastAuthPaths = paths + s.lastAuthPaths = append([]string(nil), paths...) + s.authPersisted++ + s.signalPersisted() return nil } func (s *stubStore) AuthDir() string { return s.authDir } +func (s *stubStore) signalPersisted() { + if s.persisted == nil { + return + } + select { + case s.persisted <- struct{}{}: + default: + } +} + +func (s *stubStore) persistenceSnapshot() (cfgPersisted, authPersisted int, message string, paths []string) { + s.mu.Lock() + defer s.mu.Unlock() + return s.cfgPersisted, s.authPersisted, s.lastAuthMessage, append([]string(nil), s.lastAuthPaths...) +} + func TestNewWatcherDetectsPersisterAndAuthDir(t *testing.T) { tmp := t.TempDir() store := &stubStore{authDir: tmp} @@ -1559,26 +1583,33 @@ func TestNewWatcherDetectsPersisterAndAuthDir(t *testing.T) { } func TestPersistConfigAndAuthAsyncInvokePersister(t *testing.T) { + store := &stubStore{persisted: make(chan struct{}, 2)} w := &Watcher{ - storePersister: &stubStore{}, + storePersister: store, } w.persistConfigAsync() w.persistAuthAsync("msg", " a ", "", "b ") - time.Sleep(30 * time.Millisecond) - store := w.storePersister.(*stubStore) - if atomic.LoadInt32(&store.cfgPersisted) != 1 { - t.Fatalf("expected PersistConfig to be called once, got %d", store.cfgPersisted) + for range 2 { + select { + case <-store.persisted: + case <-time.After(time.Second): + t.Fatal("timed out waiting for asynchronous persistence") + } + } + cfgPersisted, authPersisted, message, paths := store.persistenceSnapshot() + if cfgPersisted != 1 { + t.Fatalf("expected PersistConfig to be called once, got %d", cfgPersisted) } - if atomic.LoadInt32(&store.authPersisted) != 1 { - t.Fatalf("expected PersistAuthFiles to be called once, got %d", store.authPersisted) + if authPersisted != 1 { + t.Fatalf("expected PersistAuthFiles to be called once, got %d", authPersisted) } - if store.lastAuthMessage != "msg" { - t.Fatalf("unexpected auth message: %s", store.lastAuthMessage) + if message != "msg" { + t.Fatalf("unexpected auth message: %s", message) } - if len(store.lastAuthPaths) != 2 || store.lastAuthPaths[0] != "a" || store.lastAuthPaths[1] != "b" { - t.Fatalf("unexpected filtered paths: %#v", store.lastAuthPaths) + if len(paths) != 2 || paths[0] != "a" || paths[1] != "b" { + t.Fatalf("unexpected filtered paths: %#v", paths) } } @@ -1601,7 +1632,19 @@ func TestScheduleConfigReloadDebounces(t *testing.T) { w.scheduleConfigReload() w.scheduleConfigReload() - time.Sleep(400 * time.Millisecond) + deadline := time.Now().Add(time.Second) + for { + w.clientsMutex.RLock() + hashSet := w.lastConfigHash != "" + w.clientsMutex.RUnlock() + if hashSet { + break + } + if time.Now().After(deadline) { + t.Fatal("timed out waiting for debounced config reload") + } + time.Sleep(10 * time.Millisecond) + } if atomic.LoadInt32(&reloads) != 1 { t.Fatalf("expected single debounced reload, got %d", reloads) diff --git a/sdk/api/handlers/gemini/gemini_handlers.go b/sdk/api/handlers/gemini/gemini_handlers.go index 60aed26a5..f987a7c71 100644 --- a/sdk/api/handlers/gemini/gemini_handlers.go +++ b/sdk/api/handlers/gemini/gemini_handlers.go @@ -163,6 +163,26 @@ func (h *GeminiAPIHandler) GeminiHandler(c *gin.Context) { } } +// pendingGeminiStreamError reports the error that must be surfaced when a +// Gemini stream closes without any data. A closed data channel with a +// buffered upstream error means the stream failed: returning nil here would +// let the handler commit HTTP 200 headers for an empty stream. Mirrors +// pendingClaudeStreamError in the Claude handler. +func pendingGeminiStreamError(errChan <-chan *interfaces.ErrorMessage) *interfaces.ErrorMessage { + if errChan == nil { + return nil + } + select { + case errMsg, ok := <-errChan: + if !ok { + return nil + } + return errMsg + default: + return nil + } +} + // handleStreamGenerateContent handles streaming content generation requests for Gemini models. // This function establishes a Server-Sent Events connection and streams the generated content // back to the client in real-time. It supports both SSE format and direct streaming based @@ -219,6 +239,15 @@ func (h *GeminiAPIHandler) handleStreamGenerateContent(c *gin.Context, modelName return case chunk, ok := <-dataChan: if !ok { + // Stream closed without data. Surface a buffered pending error + // before committing SSE headers, so a failed upstream never + // looks like a successful empty stream (matches the OpenAI + // and Claude streaming paths). + if errMsg := pendingGeminiStreamError(errChan); errMsg != nil { + h.WriteErrorResponse(c, errMsg) + cliCancel(errMsg.Error) + return + } // Closed without data if alt == "" { setSSEHeaders() diff --git a/sdk/api/handlers/gemini/gemini_handlers_error_test.go b/sdk/api/handlers/gemini/gemini_handlers_error_test.go new file mode 100644 index 000000000..f08b999ea --- /dev/null +++ b/sdk/api/handlers/gemini/gemini_handlers_error_test.go @@ -0,0 +1,36 @@ +package gemini + +import ( + "errors" + "testing" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces" +) + +// TestPendingGeminiStreamErrorUsesBufferedError ensures a buffered upstream +// error is surfaced when the stream closes without data. Regression guard for +// the handleStreamGenerateContent branch that previously committed HTTP 200 +// SSE headers for a failed empty stream. +func TestPendingGeminiStreamErrorUsesBufferedError(t *testing.T) { + errs := make(chan *interfaces.ErrorMessage, 1) + errs <- &interfaces.ErrorMessage{StatusCode: 500, Error: errors.New("empty_completion: stream closed without data")} + + errMsg := pendingGeminiStreamError(errs) + if errMsg == nil { + t.Fatal("expected pending stream error") + } + if errMsg.StatusCode != 500 { + t.Fatalf("unexpected status code: %d", errMsg.StatusCode) + } +} + +// TestPendingGeminiStreamErrorWithoutErrorCommitsSuccess ensures a cleanly +// closed stream with no buffered error is treated as a normal completion. +func TestPendingGeminiStreamErrorWithoutErrorCommitsSuccess(t *testing.T) { + errs := make(chan *interfaces.ErrorMessage, 1) + + errMsg := pendingGeminiStreamError(errs) + if errMsg != nil { + t.Fatalf("expected success, got error: %v", errMsg.Error) + } +} diff --git a/sdk/api/handlers/handlers_error_response_test.go b/sdk/api/handlers/handlers_error_response_test.go index c52539016..bc4c35d70 100644 --- a/sdk/api/handlers/handlers_error_response_test.go +++ b/sdk/api/handlers/handlers_error_response_test.go @@ -15,6 +15,7 @@ import ( "github.com/router-for-me/CLIProxyAPI/v7/internal/clienterror" "github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces" coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config" ) @@ -250,6 +251,70 @@ func TestExecutionErrorMessageMapsContextStatuses(t *testing.T) { } } +func TestExecutionErrorMessageMapsTerminatedTrustedProvenance(t *testing.T) { + for _, tc := range []struct { + name string + trusted bool + }{ + {name: "trusted local termination", trusted: true}, + {name: "untrusted upstream termination", trusted: false}, + } { + t.Run(tc.name, func(t *testing.T) { + terminated := &coreexecutor.RequestTerminatedError{ + HTTPStatus: http.StatusTeapot, + Header: http.Header{"X-Downstream": []string{"preserved"}}, + Body: []byte(`{"custom":"body"}`), + Trusted: tc.trusted, + } + msg := executionErrorMessage(terminated) + if msg == nil { + t.Fatal("executionErrorMessage() returned nil") + } + if !msg.DirectResponse { + t.Fatal("DirectResponse must remain true for every RequestTerminatedError") + } + if msg.TrustedDirectResponse != tc.trusted { + t.Fatalf("TrustedDirectResponse = %t, want %t", msg.TrustedDirectResponse, tc.trusted) + } + if msg.Body == nil || string(msg.Body) != `{"custom":"body"}` { + t.Fatalf("Body = %q, want preserved body", msg.Body) + } + if got := msg.Headers.Get("X-Downstream"); got != "preserved" { + t.Fatalf("Headers = %v, want preserved headers", msg.Headers) + } + }) + } +} + +func TestDirectTerminationErrorMarksTrustedLocalResponse(t *testing.T) { + msg := directTerminationError(http.StatusTeapot, http.Header{"X-Local": []string{"yes"}}, []byte(`{"ok":true}`)) + if msg == nil { + t.Fatal("directTerminationError() returned nil") + } + if !msg.DirectResponse { + t.Fatal("local direct termination must set DirectResponse=true") + } + if !msg.TrustedDirectResponse { + t.Fatal("local direct termination must set TrustedDirectResponse=true") + } + if msg.StatusCode != http.StatusTeapot { + t.Fatalf("StatusCode = %d, want %d", msg.StatusCode, http.StatusTeapot) + } +} + +func TestNonTerminatedErrorKeepsZeroValueTrustedDirectResponse(t *testing.T) { + msg := executionErrorMessage(errors.New("upstream boom")) + if msg == nil { + t.Fatal("executionErrorMessage() returned nil") + } + if msg.DirectResponse { + t.Fatal("plain upstream error must not be a DirectResponse") + } + if msg.TrustedDirectResponse { + t.Fatal("plain upstream error must have TrustedDirectResponse=false") + } +} + func TestStatusFromErrorMapsContextStatuses(t *testing.T) { if got := statusFromError(context.Canceled); got != clienterror.StatusClientClosedRequest { t.Fatalf("statusFromError(canceled) = %d, want %d", got, clienterror.StatusClientClosedRequest) diff --git a/sdk/api/handlers/handlers_errors.go b/sdk/api/handlers/handlers_errors.go index 57df50e04..47f7a1568 100644 --- a/sdk/api/handlers/handlers_errors.go +++ b/sdk/api/handlers/handlers_errors.go @@ -51,7 +51,7 @@ func enrichAuthSelectionError(err error, providers []string, model string) error modelText = "unknown" } - baseMessage := strings.TrimSpace(authErr.Message) + baseMessage := strings.TrimSpace(err.Error()) if baseMessage == "" { baseMessage = "no auth available" } diff --git a/sdk/api/handlers/handlers_execution.go b/sdk/api/handlers/handlers_execution.go index 994bd33a2..2e25edcd1 100644 --- a/sdk/api/handlers/handlers_execution.go +++ b/sdk/api/handlers/handlers_execution.go @@ -300,11 +300,12 @@ func executionErrorMessage(err error) *interfaces.ErrorMessage { var terminated *coreexecutor.RequestTerminatedError if errors.As(err, &terminated) && terminated != nil { return &interfaces.ErrorMessage{ - StatusCode: normalizedTerminationStatus(terminated.StatusCode()), - Error: err, - DirectResponse: true, - Body: terminated.ResponseBody(), - Headers: terminated.ResponseHeaders(), + StatusCode: normalizedTerminationStatus(terminated.StatusCode()), + Error: err, + DirectResponse: true, + TrustedDirectResponse: terminated.Trusted, + Body: terminated.ResponseBody(), + Headers: terminated.ResponseHeaders(), } } status := http.StatusInternalServerError diff --git a/sdk/api/handlers/handlers_interceptors.go b/sdk/api/handlers/handlers_interceptors.go index 2c9cb282c..aa4333b35 100644 --- a/sdk/api/handlers/handlers_interceptors.go +++ b/sdk/api/handlers/handlers_interceptors.go @@ -135,10 +135,11 @@ func requestTerminationError(resp pluginapi.RequestInterceptResponse) *interface func directTerminationError(statusCode int, headers http.Header, body []byte) *interfaces.ErrorMessage { return &interfaces.ErrorMessage{ - StatusCode: normalizedTerminationStatus(statusCode), - DirectResponse: true, - Body: cloneBytes(body), - Headers: cloneHeader(headers), + StatusCode: normalizedTerminationStatus(statusCode), + DirectResponse: true, + TrustedDirectResponse: true, + Body: cloneBytes(body), + Headers: cloneHeader(headers), } } diff --git a/sdk/api/handlers/handlers_stream_bootstrap_test.go b/sdk/api/handlers/handlers_stream_bootstrap_test.go index fc37107dd..202f3c6b8 100644 --- a/sdk/api/handlers/handlers_stream_bootstrap_test.go +++ b/sdk/api/handlers/handlers_stream_bootstrap_test.go @@ -237,8 +237,8 @@ func (e *splitResponsesEventStreamExecutor) Execute(context.Context, *coreauth.A func (e *splitResponsesEventStreamExecutor) ExecuteStream(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (*coreexecutor.StreamResult, error) { ch := make(chan coreexecutor.StreamChunk, 2) - ch <- coreexecutor.StreamChunk{Payload: []byte("event: response.completed")} - ch <- coreexecutor.StreamChunk{Payload: []byte("data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp-1\",\"output\":[]}}")} + ch <- coreexecutor.StreamChunk{Payload: []byte("event: response.completed\n")} + ch <- coreexecutor.StreamChunk{Payload: []byte("data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp-1\",\"output\":[]}}\n\n")} close(ch) return &coreexecutor.StreamResult{Chunks: ch}, nil } @@ -1149,10 +1149,10 @@ func TestExecuteStreamWithAuthManager_AllowsSplitOpenAIResponsesSSEEventLines(t if len(got) != 2 { t.Fatalf("expected 2 forwarded chunks, got %d: %#v", len(got), got) } - if got[0] != "event: response.completed" { + if got[0] != "event: response.completed\n" { t.Fatalf("unexpected first chunk: %q", got[0]) } - expectedData := "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp-1\",\"output\":[]}}" + expectedData := "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp-1\",\"output\":[]}}\n\n" if got[1] != expectedData { t.Fatalf("unexpected second chunk.\nGot: %q\nWant: %q", got[1], expectedData) } diff --git a/sdk/api/handlers/model_execution_test.go b/sdk/api/handlers/model_execution_test.go index e83337a21..9f1cd9027 100644 --- a/sdk/api/handlers/model_execution_test.go +++ b/sdk/api/handlers/model_execution_test.go @@ -425,8 +425,15 @@ func TestExecuteModelStreamStartupError(t *testing.T) { if errMsg.StatusCode != http.StatusInternalServerError { t.Fatalf("status = %d, want %d", errMsg.StatusCode, http.StatusInternalServerError) } - if errMsg.Error == nil || errMsg.Error.Error() != "startup failed" { - t.Fatalf("error = %v, want startup failed", errMsg.Error) + if errMsg.Error == nil { + t.Fatal("ExecuteModelStream() error = nil, want startup error") + } + got := errMsg.Error.Error() + if !strings.HasPrefix(got, "startup failed") { + t.Fatalf("error = %q, want prefix %q", got, "startup failed") + } + if !strings.Contains(got, "attempted routes: [codex:error]") { + t.Fatalf("error = %q, want sanitized route summary %q", got, "attempted routes: [codex:error]") } if stream.Chunks != nil { t.Fatal("stream chunks created for startup error") @@ -611,7 +618,7 @@ func TestExecuteModelStreamKeepsInteractionsProviderForOpenAIEntry(t *testing.T) provider: constant.GeminiInteractions, stream: func(ctx context.Context, auth *coreauth.Auth, req coreexecutor.Request, opts coreexecutor.Options) (*coreexecutor.StreamResult, error) { chunks := make(chan coreexecutor.StreamChunk, 1) - chunks <- coreexecutor.StreamChunk{Payload: []byte(`{"id":"chunk_1","object":"chat.completion.chunk","choices":[]}`)} + chunks <- coreexecutor.StreamChunk{Payload: []byte(`{"id":"chunk_1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"ok"}}]}`)} close(chunks) return &coreexecutor.StreamResult{Chunks: chunks}, nil }, diff --git a/sdk/api/handlers/openai/openai_handlers.go b/sdk/api/handlers/openai/openai_handlers.go index f3e1e6c60..6878a0ba1 100644 --- a/sdk/api/handlers/openai/openai_handlers.go +++ b/sdk/api/handlers/openai/openai_handlers.go @@ -11,6 +11,9 @@ import ( "encoding/json" "fmt" "net/http" + "regexp" + "sort" + "strings" "sync" "github.com/gin-gonic/gin" @@ -20,6 +23,7 @@ import ( codexconverter "github.com/router-for-me/CLIProxyAPI/v7/internal/translator/codex/openai/chat-completions" responsesconverter "github.com/router-for-me/CLIProxyAPI/v7/internal/translator/openai/openai/responses" "github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers" + coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" "github.com/tidwall/gjson" "github.com/tidwall/sjson" ) @@ -321,6 +325,7 @@ func (h *OpenAIAPIHandler) forwardResponsesAsChatStream(c *gin.Context, flusher if errMsg == nil { return } + errMsg = sanitizeOpenAIErrorMessage(errMsg) status := http.StatusInternalServerError if errMsg.StatusCode > 0 { status = errMsg.StatusCode @@ -528,7 +533,7 @@ func (h *OpenAIAPIHandler) handleNonStreamingResponse(c *gin.Context, rawJSON [] resp, upstreamHeaders, errMsg := h.ExecuteWithAuthManager(cliCtx, h.HandlerType(), modelName, rawJSON, h.GetAlt(c)) stopKeepAlive() if errMsg != nil { - h.WriteErrorResponse(c, errMsg) + h.WriteErrorResponse(c, sanitizeOpenAIErrorMessage(errMsg)) cliCancel(errMsg.Error) return } @@ -544,7 +549,7 @@ func (h *OpenAIAPIHandler) handleNonStreamingResponseViaResponses(c *gin.Context cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background()) resp, upstreamHeaders, errMsg := h.ExecuteWithAuthManager(cliCtx, OpenaiResponse, modelName, rawJSON, h.GetAlt(c)) if errMsg != nil { - h.WriteErrorResponse(c, errMsg) + h.WriteErrorResponse(c, sanitizeOpenAIErrorMessage(errMsg)) cliCancel(errMsg.Error) return } @@ -606,7 +611,7 @@ func (h *OpenAIAPIHandler) handleStreamingResponse(c *gin.Context, rawJSON []byt continue } // Upstream failed immediately. Return proper error status and JSON. - h.WriteErrorResponse(c, errMsg) + h.WriteErrorResponse(c, sanitizeOpenAIErrorMessage(errMsg)) if errMsg != nil { cliCancel(errMsg.Error) } else { @@ -615,7 +620,15 @@ func (h *OpenAIAPIHandler) handleStreamingResponse(c *gin.Context, rawJSON []byt return case chunk, ok := <-dataChan: if !ok { - // Stream closed without data? Send DONE or just headers. + // Stream closed without data. Surface a buffered pending error + // before committing SSE headers, so a failed upstream never + // looks like a successful empty stream. + if pErr, pending := pendingOpenAIStreamError(errChan); pending { + h.WriteErrorResponse(c, sanitizeOpenAIErrorMessage(pErr)) + cliCancel(pErr.Error) + return + } + // Clean close. Send DONE or just headers. setSSEHeaders() handlers.WriteUpstreamHeaders(c.Writer.Header(), upstreamHeaders) _, _ = fmt.Fprintf(c.Writer, "data: [DONE]\n\n") @@ -673,7 +686,7 @@ func (h *OpenAIAPIHandler) handleStreamingResponseViaResponses(c *gin.Context, r errChan = nil continue } - h.WriteErrorResponse(c, errMsg) + h.WriteErrorResponse(c, sanitizeOpenAIErrorMessage(errMsg)) if errMsg != nil { cliCancel(errMsg.Error) } else { @@ -682,6 +695,15 @@ func (h *OpenAIAPIHandler) handleStreamingResponseViaResponses(c *gin.Context, r return case chunk, ok := <-dataChan: if !ok { + // Stream closed without data. Surface a buffered pending error + // before committing SSE headers, so a failed upstream never + // looks like a successful empty stream. + if pErr, pending := pendingOpenAIStreamError(errChan); pending { + h.WriteErrorResponse(c, sanitizeOpenAIErrorMessage(pErr)) + cliCancel(pErr.Error) + return + } + // Clean close. Send DONE or just headers. setSSEHeaders() handlers.WriteUpstreamHeaders(c.Writer.Header(), upstreamHeaders) _, _ = fmt.Fprintf(c.Writer, "data: [DONE]\n\n") @@ -720,7 +742,7 @@ func (h *OpenAIAPIHandler) handleCompletionsNonStreamingResponse(c *gin.Context, resp, upstreamHeaders, errMsg := h.ExecuteWithAuthManager(cliCtx, h.HandlerType(), modelName, chatCompletionsJSON, "") stopKeepAlive() if errMsg != nil { - h.WriteErrorResponse(c, errMsg) + h.WriteErrorResponse(c, sanitizeOpenAIErrorMessage(errMsg)) cliCancel(errMsg.Error) return } @@ -776,7 +798,7 @@ func (h *OpenAIAPIHandler) handleCompletionsStreamingResponse(c *gin.Context, ra errChan = nil continue } - h.WriteErrorResponse(c, errMsg) + h.WriteErrorResponse(c, sanitizeOpenAIErrorMessage(errMsg)) if errMsg != nil { cliCancel(errMsg.Error) } else { @@ -785,6 +807,15 @@ func (h *OpenAIAPIHandler) handleCompletionsStreamingResponse(c *gin.Context, ra return case chunk, ok := <-dataChan: if !ok { + // Stream closed without data. Surface a buffered pending error + // before committing SSE headers, so a failed upstream never + // looks like a successful empty stream. + if pErr, pending := pendingOpenAIStreamError(errChan); pending { + h.WriteErrorResponse(c, sanitizeOpenAIErrorMessage(pErr)) + cliCancel(pErr.Error) + return + } + // Clean close. Send DONE or just headers. setSSEHeaders() handlers.WriteUpstreamHeaders(c.Writer.Header(), upstreamHeaders) _, _ = fmt.Fprintf(c.Writer, "data: [DONE]\n\n") @@ -849,6 +880,7 @@ func (h *OpenAIAPIHandler) handleStreamResult(c *gin.Context, flusher http.Flush if errMsg == nil { return } + errMsg = sanitizeOpenAIErrorMessage(errMsg) status := http.StatusInternalServerError if errMsg.StatusCode > 0 { status = errMsg.StatusCode @@ -865,3 +897,392 @@ func (h *OpenAIAPIHandler) handleStreamResult(c *gin.Context, flusher http.Flush }, }) } + +// pendingOpenAIStreamError returns an immediately available non-nil stream error +// buffered on errChan. It mirrors pendingClaudeStreamError in the Claude handler: +// the initial peek consumes a queued upstream failure before committing SSE +// headers so a failed stream never looks like a successful empty stream. +func pendingOpenAIStreamError(errs <-chan *interfaces.ErrorMessage) (*interfaces.ErrorMessage, bool) { + if errs == nil { + return nil, false + } + select { + case errMsg, ok := <-errs: + if !ok || errMsg == nil { + return nil, false + } + return errMsg, true + default: + return nil, false + } +} + +// sanitizeOpenAIErrorMessage is the trust-boundary sanitizer for the OpenAI +// pre-output sinks. It preserves a DirectResponse only when it is explicitly +// trusted (local plugin/interceptor); every other error path sanitizes +// strictly: forces a valid status, clears Body, forces DirectResponse=false, +// and redacts credential material from the error text. It returns nil for a +// nil input. +func sanitizeOpenAIErrorMessage(errMsg *interfaces.ErrorMessage) *interfaces.ErrorMessage { + if errMsg != nil && errMsg.DirectResponse && errMsg.TrustedDirectResponse { + return errMsg + } + return sanitizeOpenAIStrictErrorMessage(errMsg) +} + +func sanitizeOpenAIStrictErrorMessage(errMsg *interfaces.ErrorMessage) *interfaces.ErrorMessage { + if errMsg == nil { + return nil + } + status := errMsg.StatusCode + if status < http.StatusBadRequest || status > 599 { + status = http.StatusInternalServerError + } + safe := *errMsg + safe.StatusCode = status + safe.DirectResponse = false + safe.Body = nil + if errMsg.Error != nil { + safe.Error = &openAIStreamSanitizedError{ + message: openAIStreamErrorText(errMsg.Error.Error(), status), + safeHeaders: coreauth.SafeResponseHeaders(errMsg.Error), + } + } + return &safe +} + +type openAIStreamSanitizedError struct { + message string + safeHeaders http.Header +} + +func (e *openAIStreamSanitizedError) Error() string { return e.message } + +func (e *openAIStreamSanitizedError) SafeResponseHeaders() http.Header { + if e == nil || e.safeHeaders == nil { + return nil + } + return e.safeHeaders.Clone() +} + +// openAIStreamErrorText produces a client-safe error message. JSON error bodies +// are preserved field-by-field with sanitization; free-form text is kept with +// credential material redacted. +func openAIStreamErrorText(text string, status int) string { + if t := strings.TrimSpace(text); t != "" && json.Valid([]byte(t)) { + return sanitizeOpenAIStreamJSON(t, status) + } + fallback := http.StatusText(status) + if strings.TrimSpace(text) == "" { + return fallback + } + return redactOpenAIStreamErrorText(strings.TrimSpace(text)) +} + +func sanitizeOpenAIStreamJSON(text string, status int) string { + root := gjson.Parse(text) + errorNode := root.Get("error") + if !errorNode.Exists() || !errorNode.IsObject() { + errorNode = root.Get("response.error") + } + if errorNode.Exists() && errorNode.IsObject() { + safe := []byte(`{"error":{}}`) + copied := false + for _, field := range []string{"type", "code", "message", "param"} { + value := errorNode.Get(field) + if !value.Exists() || value.Type == gjson.Null { + continue + } + limit := openAIStreamErrorFieldLimit + if field == "message" { + limit = openAIStreamErrorMessageLimit + } + safe, _ = sjson.SetBytes(safe, "error."+field, truncateOpenAIStreamErrorText(redactOpenAIStreamErrorText(value.String()), limit)) + copied = true + } + if copied { + return string(safe) + } + } + + safe := []byte(`{"type":"error"}`) + copied := false + for _, field := range []string{"code", "message", "param"} { + value := root.Get(field) + if !value.Exists() || value.Type == gjson.Null { + continue + } + limit := openAIStreamErrorFieldLimit + if field == "message" { + limit = openAIStreamErrorMessageLimit + } + safe, _ = sjson.SetBytes(safe, field, truncateOpenAIStreamErrorText(redactOpenAIStreamErrorText(value.String()), limit)) + copied = true + } + if copied { + return string(safe) + } + return http.StatusText(status) +} + +const ( + openAIStreamErrorMessageLimit = 2048 + openAIStreamErrorFieldLimit = 256 +) + +var ( + // openAIStreamKeyPattern matches a sensitive key name and its separator + // (= or :), preceded by a boundary and optional quote/escape syntax. + // Group 1 is the leading boundary/quote syntax, group 2 the key, group 3 + // the trailing quote/space syntax plus separator. + openAIStreamKeyPattern = regexp.MustCompile(`(?i)((?:^|[^A-Za-z0-9_])(?:\\*["']?)?)(api[_-]?key|apikey|access[_-]?key[_-]?id|aws[_-]?access[_-]?key[_-]?id|api[_-]?key[_-]?id|access[_-]?token|authorization|token|secret|credential|aws[_-]?credential|refresh[_-]?token|client[_-]?secret|(?:[A-Za-z0-9]+(?:[_-][A-Za-z0-9]+)*)[_-](?:key|token|secret|credential|key[_-]?id)|(?-i:[A-Za-z0-9]*(?:[a-z0-9]|_[a-z0-9]|-[a-z0-9])(?:Key|Token|Secret|Credential|KeyId|Key_Id|Key-Id)))((?:\\*["']?)?\s*[=:])`) + // openAIStreamSpaceAPIKeyPattern matches the "api key:" spelling with a + // space between api and key, in header/assignment contexts. + openAIStreamSpaceAPIKeyPattern = regexp.MustCompile(`(?i)((?:^|[^A-Za-z0-9_]))(api[ _]key)(["']?\s*[=:])`) + // openAIStreamBareKeyDenyPattern marks key names that merely mention a + // credential kind without being a credential themselves. + openAIStreamBareKeyDenyPattern = regexp.MustCompile(`(?i)^(?:not|non|no|count|counter|key[_-]?count|token[_-]?count)(?:[_-]|$)`) + // openAIStreamAuthSchemePattern detects a Bearer/Basic scheme at the + // start of a credential value so the scheme can be preserved. + openAIStreamAuthSchemePattern = regexp.MustCompile(`(?i)^(Bearer|Basic)\s+`) + // openAIStreamAuthPattern redacts standalone Bearer/Basic credentials + // that appear outside key/value contexts. + openAIStreamAuthPattern = regexp.MustCompile(`(?i)(\b(?:Bearer|Basic)\s+)([-A-Za-z0-9._~+/=]+)`) +) + +func truncateOpenAIStreamErrorText(text string, limit int) string { + runes := []rune(text) + if len(runes) <= limit { + return text + } + return string(runes[:limit]) + "…" +} + +func redactOpenAIStreamErrorText(text string) string { + text = redactOpenAIStreamKeyValues(text) + return openAIStreamAuthPattern.ReplaceAllString(text, "${1}[REDACTED]") +} + +// redactOpenAIStreamKeyValues locates sensitive key/value pairs and replaces +// their credential with [REDACTED], preserving quote/escape syntax and, for +// Bearer/Basic values, the scheme. Compound keys always redact; bare +// token/secret/credential keys redact only in explicit JSON, assignment, or +// line-start header contexts. +func redactOpenAIStreamKeyValues(text string) string { + type keyMatch struct { + loc []int + key string + } + var matches []keyMatch + for _, loc := range openAIStreamKeyPattern.FindAllStringSubmatchIndex(text, -1) { + matches = append(matches, keyMatch{loc: loc, key: text[loc[4]:loc[5]]}) + } + for _, loc := range openAIStreamSpaceAPIKeyPattern.FindAllStringSubmatchIndex(text, -1) { + matches = append(matches, keyMatch{loc: loc, key: "api key"}) + } + if len(matches) == 0 { + return text + } + sort.Slice(matches, func(a, b int) bool { return matches[a].loc[0] < matches[b].loc[0] }) + var b strings.Builder + b.Grow(len(text) + 16*len(matches)) + last := 0 + for _, m := range matches { + loc := m.loc + if loc[0] < last { + continue + } + key := strings.ToLower(m.key) + if openAIStreamBareKeyDenyPattern.MatchString(key) { + continue + } + if key == "token" || key == "secret" || key == "credential" { + if !openAIStreamBareKeyContextOK(text, loc) { + continue + } + } + sepEnd := loc[1] + valueEnd, redactStart, redactEnd := openAIStreamValueBounds(text, sepEnd, key == "authorization" || key == "api key") + b.WriteString(text[last:sepEnd]) + b.WriteString(text[sepEnd:redactStart]) + if redactEnd > redactStart { + b.WriteString("[REDACTED]") + } + b.WriteString(text[redactEnd:valueEnd]) + last = valueEnd + } + b.WriteString(text[last:]) + return b.String() +} + +// openAIStreamBareKeyContextOK applies the deterministic context rule for bare +// token/secret/credential keys: they redact only as explicit JSON fields, '=' +// assignments, or line-start headers. +func openAIStreamBareKeyContextOK(text string, loc []int) bool { + if loc[4] == 0 || text[loc[4]-1] == '\n' { + return true + } + if b := text[loc[4]-1]; b == '"' || b == '\\' || b == '\'' { + return true + } + for i := loc[6]; i < loc[7]; i++ { + if text[i] == '=' { + return true + } + } + return false +} + +// openAIStreamValueBounds returns the region [redactStart, redactEnd) to +// replace with [REDACTED] for the value starting at start, and the full value +// span [start, valueEnd) the redaction consumes. Quoted values keep their +// opening/closing quote syntax; unquoted generic values stop at whitespace; +// Bearer/Basic credentials span the space between scheme and token; +// authorization/api key values consume the whole multi-part value. +func openAIStreamValueBounds(text string, start int, isAuth bool) (valueEnd, redactStart, redactEnd int) { + n := len(text) + i := start + for i < n && (text[i] == ' ' || text[i] == '\t') { + i++ + } + if i >= n { + return start, start, start + } + backslashes := 0 + j := i + for j < n && text[j] == '\\' { + backslashes++ + j++ + } + if j < n && (text[j] == '"' || text[j] == '\'') { + quote := text[j] + openEnd := j + 1 + closeStart, closeEnd := openAIStreamQuoteClose(text, openEnd, quote, backslashes) + if schemeEnd := openAIStreamAuthSchemeEnd(text, openEnd); schemeEnd >= 0 && schemeEnd <= closeStart { + return closeEnd, schemeEnd, closeStart + } + return closeEnd, openEnd, closeStart + } + end := i + for end < n { + c := text[end] + if c == '\\' { + if end+1 >= n { + end = n + break + } + end += 2 + continue + } + if c == ' ' || c == '\t' || c == '\n' || c == '\r' || c == '}' || c == ')' || c == ']' || c == ',' || c == ';' || c == '\'' || c == '"' { + break + } + end++ + } + if schemeEnd := openAIStreamAuthSchemeEnd(text, i); schemeEnd >= 0 { + credEnd := schemeEnd + for credEnd < n { + c := text[credEnd] + if c == '\\' { + if credEnd+1 >= n { + credEnd = n + break + } + credEnd += 2 + continue + } + if c == '\n' || c == '\r' || c == '}' || c == ')' || c == ']' || c == ',' || c == ';' || c == '"' { + break + } + credEnd++ + } + return credEnd, schemeEnd, credEnd + } + if isAuth { + authEnd := i + for authEnd < n { + c := text[authEnd] + if c == '\\' { + if authEnd+1 >= n { + authEnd = n + break + } + authEnd += 2 + continue + } + if c == '\n' || c == '\r' || c == '}' || c == ')' || c == ']' { + break + } + authEnd++ + } + return authEnd, i, authEnd + } + return end, i, end +} + +// openAIStreamQuoteClose finds the closing quote syntax of a quoted value +// starting after the opening quote at openEnd. +func openAIStreamQuoteClose(text string, openEnd int, quote byte, openRun int) (closeSyntaxStart, closeEnd int) { + n := len(text) + if openRun == 0 { + k := openEnd + for k < n { + if text[k] == '\\' { + if k+1 >= n { + return n, n + } + k += 2 + continue + } + if text[k] == quote { + return k, k + 1 + } + k++ + } + return n, n + } + k := openEnd + for k < n { + if text[k] == '\\' { + r := 0 + j := k + for j < n && text[j] == '\\' { + r++ + j++ + } + if j < n && text[j] == quote { + if r == openRun { + return j - openRun, j + 1 + } + if r > openRun { + k = j + 1 + continue + } + return j - r, j + 1 + } + k = j + continue + } + if text[k] == quote { + return k, k + 1 + } + k++ + } + return n, n +} + +// openAIStreamAuthSchemeEnd reports the position just after a Bearer/Basic +// scheme word plus following whitespace at start, or -1 when no such scheme is +// present. +func openAIStreamAuthSchemeEnd(text string, start int) int { + i := start + n := len(text) + for i < n && (text[i] == ' ' || text[i] == '\t') { + i++ + } + m := openAIStreamAuthSchemePattern.FindStringSubmatchIndex(text[i:]) + if m == nil { + return -1 + } + return i + m[1] +} diff --git a/sdk/api/handlers/openai/openai_handlers_stream_peek_test.go b/sdk/api/handlers/openai/openai_handlers_stream_peek_test.go new file mode 100644 index 000000000..b1dbcf290 --- /dev/null +++ b/sdk/api/handlers/openai/openai_handlers_stream_peek_test.go @@ -0,0 +1,849 @@ +package openai + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/gin-gonic/gin" + "github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces" + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers" + coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config" + "github.com/tidwall/gjson" +) + +// peekStreamExecutor feeds a fake executor stream that closes the chunk channel +// immediately. All three initial streaming peek loops (chat, ViaResponses and +// legacy completions) then race a closed dataChan against any buffered pending +// error on errChan. +type peekStreamExecutor struct { + secret string + payload string +} + +func (*peekStreamExecutor) Identifier() string { return "peek-stream" } + +func (*peekStreamExecutor) Execute(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) { + return coreexecutor.Response{}, errors.New("not implemented") +} + +func (e *peekStreamExecutor) ExecuteStream(_ context.Context, _ *coreauth.Auth, _ coreexecutor.Request, _ coreexecutor.Options) (*coreexecutor.StreamResult, error) { + chunks := make(chan coreexecutor.StreamChunk, 2) + if e.payload != "" { + chunks <- coreexecutor.StreamChunk{Payload: []byte(e.payload)} + } + if e.secret != "" { + chunks <- coreexecutor.StreamChunk{Err: errors.New("upstream failure: api_key=" + e.secret)} + } + close(chunks) + return &coreexecutor.StreamResult{Chunks: chunks}, nil +} + +func (*peekStreamExecutor) Refresh(_ context.Context, auth *coreauth.Auth) (*coreauth.Auth, error) { + return auth, nil +} + +func (*peekStreamExecutor) CountTokens(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) { + return coreexecutor.Response{}, errors.New("not implemented") +} + +func (*peekStreamExecutor) HttpRequest(context.Context, *coreauth.Auth, *http.Request) (*http.Response, error) { + return nil, errors.New("not implemented") +} + +// sendPeekRequest drives the given handler method through a registered executor. +func sendPeekRequest(t *testing.T, route, body string, executor *peekStreamExecutor, endpoints []string) *httptest.ResponseRecorder { + t.Helper() + manager := coreauth.NewManager(nil, nil, nil) + manager.RegisterExecutor(executor) + authID := "peek-stream-auth" + auth := &coreauth.Auth{ID: authID, Provider: executor.Identifier(), Status: coreauth.StatusActive} + if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + registry.GetGlobalRegistry().RegisterClient(authID, executor.Identifier(), []*registry.ModelInfo{{ID: "peek-stream-model", SupportedEndpoints: endpoints}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(authID) }) + + base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager) + h := NewOpenAIAPIHandler(base) + router := gin.New() + router.POST("/v1/chat/completions", h.ChatCompletions) + router.POST("/v1/completions", h.Completions) + + request := httptest.NewRequest(http.MethodPost, route, strings.NewReader(body)) + request.Header.Set("Content-Type", "application/json") + recorder := httptest.NewRecorder() + router.ServeHTTP(recorder, request) + return recorder +} + +// TestStreamingPeekConsumesBufferedPendingError covers chat completions, +// ViaResponses and legacy completions peek close paths. +func TestStreamingPeekConsumesBufferedPendingError(t *testing.T) { + gin.SetMode(gin.TestMode) + const secret = "upstream-secret-fx-8849" + cases := []struct { + name string + route string + body string + endpoints []string + }{ + { + name: "chat", + route: "/v1/chat/completions", + body: `{"model":"peek-stream-model","stream":true,"messages":[{"role":"user","content":"hi"}]}`, + }, + { + name: "via-responses", + route: "/v1/chat/completions", + body: `{"model":"peek-stream-model","stream":true,"messages":[{"role":"user","content":"hi"}]}`, + endpoints: []string{openAIResponsesEndpoint}, + }, + { + name: "legacy", + route: "/v1/completions", + body: `{"model":"peek-stream-model","stream":true,"prompt":"hi"}`, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + recorder := sendPeekRequest(t, tc.route, tc.body, &peekStreamExecutor{secret: secret}, tc.endpoints) + body := recorder.Body.String() + if recorder.Code == http.StatusOK { + t.Fatalf("handler returned 200 despite buffered pending error: %q", body) + } + if recorder.Code < http.StatusBadRequest { + t.Fatalf("status = %d, want error status; body=%q", recorder.Code, body) + } + if strings.Contains(body, "data: [DONE]") { + t.Fatalf("stream emitted [DONE] despite pending error: %q", body) + } + if strings.Contains(body, secret) { + t.Fatalf("stream body leaked upstream secret: %q", body) + } + if !strings.Contains(body, "[REDACTED]") { + t.Fatalf("stream body did not redact upstream error: %q", body) + } + }) + } +} + +// TestStreamingPeekCleanCloseStillEmitsDone is the control: a clean data-channel +// close with no pending error still emits SSE [DONE] with a success status. +func TestStreamingPeekCleanCloseStillEmitsDone(t *testing.T) { + gin.SetMode(gin.TestMode) + chatPayload := `data: {"id":"cmpl-1","object":"chat.completion.chunk","created":1,"model":"peek-stream-model","choices":[{"index":0,"delta":{"content":"hi"},"finish_reason":"stop"}]}` + responsesPayload := `data: {"type":"response.output_text.delta","delta":"hi"}` + cases := []struct { + name string + route string + body string + payload string + endpoints []string + }{ + { + name: "chat", + route: "/v1/chat/completions", + body: `{"model":"peek-stream-model","stream":true,"messages":[{"role":"user","content":"hi"}]}`, + payload: chatPayload, + }, + { + name: "via-responses", + route: "/v1/chat/completions", + body: `{"model":"peek-stream-model","stream":true,"messages":[{"role":"user","content":"hi"}]}`, + payload: responsesPayload, + endpoints: []string{openAIResponsesEndpoint}, + }, + { + name: "legacy", + route: "/v1/completions", + body: `{"model":"peek-stream-model","stream":true,"prompt":"hi"}`, + payload: chatPayload, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + recorder := sendPeekRequest(t, tc.route, tc.body, &peekStreamExecutor{payload: tc.payload}, tc.endpoints) + body := recorder.Body.String() + if recorder.Code != http.StatusOK { + t.Fatalf("status = %d, want 200; body=%q", recorder.Code, body) + } + if !strings.Contains(body, "data: [DONE]") { + t.Fatalf("clean close did not emit [DONE]: %q", body) + } + }) + } +} + +// sanitizerCase drives the strict sanitizer and checks the sanitized output. +type sanitizerCase struct { + name string + in *interfaces.ErrorMessage + want string + wantOut string +} + +func runSanitizerCases(t *testing.T, cases []sanitizerCase) { + t.Helper() + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + out := sanitizeOpenAIErrorMessage(tc.in) + if tc.in == nil { + if out != nil { + t.Fatalf("nil input => non-nil output") + } + return + } + got := "" + if out != nil && out.Error != nil { + got = out.Error.Error() + } + if tc.want != "" && !strings.Contains(got, tc.want) { + t.Fatalf("sanitized = %q, want %q", got, tc.want) + } + if tc.wantOut != "" && strings.Contains(got, tc.wantOut) { + t.Fatalf("leaked %q: %q", tc.wantOut, got) + } + }) + } +} + +func em(status int, text string) *interfaces.ErrorMessage { + return &interfaces.ErrorMessage{StatusCode: status, Error: errors.New(text)} +} + +func fromSeg(parts ...string) string { return strings.Join(parts, "") } + +func TestSanitizeOpenAIErrorMessageNormalizesStatus(t *testing.T) { + if got := sanitizeOpenAIErrorMessage(nil); got != nil { + t.Fatal("nil input => non-nil output") + } + cases := []struct { + name string + in int + want int + }{ + {"no status", 0, http.StatusInternalServerError}, + {"1xx", 101, http.StatusInternalServerError}, + {"2xx", http.StatusOK, http.StatusInternalServerError}, + {"400", http.StatusBadRequest, http.StatusBadRequest}, + {"429", http.StatusTooManyRequests, http.StatusTooManyRequests}, + {"502", http.StatusBadGateway, http.StatusBadGateway}, + {"600", 600, http.StatusInternalServerError}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + out := sanitizeOpenAIErrorMessage(em(tc.in, "boom")) + if out.StatusCode != tc.want { + t.Fatalf("status=%d want=%d", out.StatusCode, tc.want) + } + if out.Body != nil { + t.Fatalf("Body not nil") + } + if out.DirectResponse != false { + t.Fatal("DirectResponse not false") + } + }) + } +} + +func TestSanitizeOpenAIErrorMessageRedactsKeyValues(t *testing.T) { + // Credential-shaped fixture values are constructed at runtime so the test + // file itself never stores a real-looking secret literal. Each case proves + // the resulting value does not survive sanitization. + sk := fromSeg("sk-abc") + hd := fromSeg("hunter2") + tk := fromSeg("tkn_456") + rf := fromSeg("ref_789") + s3 := fromSeg("s3cr3t") + ak := fromSeg("AKIAIOSFODNN7EXAMPLE") + sx := fromSeg("sx-xyz") + jw := fromSeg("eyJ.jwt.payload") + bs := fromSeg("dXNlcjpwYXNz") + tA := fromSeg("tok-XYZQ") + pB := fromSeg("sk-proj-abcdefghijklmnop") + pC := fromSeg("sk-ant-api03-xyz") + gD := fromSeg("ghp_AAAbbbCCCDDD") + sE := fromSeg("sk_live_42abc") + aF := fromSeg("AKIASECRETKEYEXAMPLE") + oG := fromSeg("ya29.exampletoken") + s1 := fromSeg("single") + eH := fromSeg("sk-esc") + + cases := []sanitizerCase{ + {name: "api_key assignment", in: em(http.StatusBadGateway, "bad api_key="+sk), want: "[REDACTED]", wantOut: sk}, + {name: "client_secret assignment", in: em(http.StatusBadGateway, "upstream rejected client_secret="+hd), want: "[REDACTED]", wantOut: hd}, + {name: "api_token header", in: em(http.StatusBadGateway, "x-api-token: "+tk), want: "[REDACTED]", wantOut: tk}, + {name: "refresh_token quoted", in: em(http.StatusBadGateway, "refresh_token=\""+rf+"\""), want: "[REDACTED]", wantOut: rf}, + {name: "star_secret", in: em(http.StatusBadGateway, "integration failed: store_secret="+s3), want: "[REDACTED]", wantOut: s3}, + {name: "dash key", in: em(http.StatusBadGateway, "credential missing access-key=<"+ak+">"), want: "[REDACTED]", wantOut: ak}, + {name: "underscore key", in: em(http.StatusBadGateway, "bad api_key="+sk), want: "[REDACTED]", wantOut: sk}, + {name: "nested key colon", in: em(http.StatusBadGateway, "login refused: \"api_key\":\""+sx+"\""), want: "[REDACTED]", wantOut: sx}, + {name: "bearer scheme preserved", in: em(http.StatusBadGateway, "token endpoint returned Bearer "+jw), want: "Bearer [REDACTED]", wantOut: jw}, + {name: "basic scheme preserved", in: em(http.StatusBadGateway, "401 from Basic "+bs), want: "Basic [REDACTED]", wantOut: bs}, + {name: "authorization header", in: em(http.StatusBadGateway, "authorization: Bearer "+tA), want: "Bearer [REDACTED]", wantOut: tA}, + {name: "openai key", in: em(http.StatusBadGateway, "invalid OpenAI api_key="+pB), want: "[REDACTED]", wantOut: pB}, + {name: "anthropic key", in: em(http.StatusBadGateway, "anthropic api key = "+pC), want: "[REDACTED]", wantOut: pC}, + {name: "github token", in: em(http.StatusBadGateway, "x-github-token: "+gD), want: "[REDACTED]", wantOut: gD}, + {name: "stripe key", in: em(http.StatusBadGateway, "stripe_key="+sE), want: "[REDACTED]", wantOut: sE}, + {name: "aws credential", in: em(http.StatusBadGateway, "aws_credential="+aF), want: "[REDACTED]", wantOut: aF}, + {name: "oauth token", in: em(http.StatusBadGateway, "oauth_token="+oG), want: "[REDACTED]", wantOut: oG}, + {name: "single quote form", in: em(http.StatusBadGateway, "client_id AND client_secret = '"+s1+"'"), want: "[REDACTED]", wantOut: s1}, + {name: "backslash trailing", in: em(http.StatusBadGateway, "api_key="+sk+"\\"), want: "[REDACTED]", wantOut: sk}, + {name: "malformed quote redacts without panic", in: em(http.StatusBadGateway, "api_key=\""+sk), want: "[REDACTED]", wantOut: sk}, + {name: "long value truncation", in: em(http.StatusBadGateway, "api_key="+strings.Repeat("a", 1000)), want: "[REDACTED]"}, + {name: "not_api_key benign", in: em(http.StatusBadGateway, "not_api_key=123"), want: "not_api_key=123", wantOut: "[REDACTED]"}, + {name: "token_count benign", in: em(http.StatusBadGateway, "token_count=42 tokens"), want: "token_count=42"}, + {name: "tokenizer benign", in: em(http.StatusBadGateway, "tokenizer=cl100k"), want: "tokenizer=cl100k"}, + {name: "secretariat benign", in: em(http.StatusBadGateway, "secretariat approved"), want: "secretariat"}, + {name: "mytoken benign", in: em(http.StatusBadGateway, "mytoken not a real credential"), want: "mytoken"}, + {name: "benign prose", in: em(http.StatusBadGateway, "nothing sensitive here"), want: "nothing sensitive here"}, + {name: "double escaped quoted", in: em(http.StatusBadGateway, `{"msg":"api_key=\\`+eH+`\\\\"}"`), want: "[REDACTED]", wantOut: eH}, + } + runSanitizerCases(t, cases) +} + +func TestSanitizeOpenAIErrorMessageJSONPath(t *testing.T) { + jS := fromSeg("sk-json-secret") + rS := fromSeg("sk-route") + jw := fromSeg("supersecretjwt") + cases := []sanitizerCase{ + { + name: "nested error object redacts secret", + in: em(http.StatusBadGateway, `{"error":{"type":"server_error","code":"upstream","message":"rejected api_key=`+jS+`","param":"mytoken"}}`), + want: "[REDACTED]", wantOut: jS, + }, + { + name: "route summary redacts", + in: em(http.StatusBadGateway, `{"message":"Route failed api_key=`+rS+`"}`), + want: "[REDACTED]", wantOut: rS, + }, + { + name: "nested response.error object", + in: em(http.StatusBadGateway, `{"response":{"error":{"message":"credentials client_secret=top"}}}`), + want: "[REDACTED]", wantOut: "top", + }, + { + name: "raw json with bearer", + in: em(http.StatusBadGateway, `{"error":{"message":"Bearer `+jw+`"}}`), + want: "Bearer [REDACTED]", wantOut: jw, + }, + { + name: "invalid json collapses to status text", + in: em(http.StatusBadGateway, `{"error":{"unclosed`), + want: `{"error":{"unclosed`, + }, + {name: "empty text uses status text", in: em(http.StatusBadGateway, ""), want: http.StatusText(http.StatusBadGateway)}, + } + runSanitizerCases(t, cases) +} + +func TestSanitizeOpenAIErrorMessageLongValueBound(t *testing.T) { + // A very long JSON message field must be truncated to the message limit + // without panic; a very large free-form payload must not panic either. + long := strings.Repeat("x", openAIStreamErrorMessageLimit*2) + out := sanitizeOpenAIErrorMessage(em(http.StatusBadGateway, `{"error":{"message":"`+long+`"}}`)) + if out.Error == nil { + t.Fatal("expected sanitized error, got nil") + } + msg := gjson.Get(out.Error.Error(), "error.message").String() + // Allow the truncation suffix (one ellipsis rune) over the hard limit. + if len([]rune(msg)) > openAIStreamErrorMessageLimit+4 { + t.Fatalf("sanitized JSON message field too large: %d chars", len(msg)) + } + if msg == "" { + t.Fatalf("expected truncation suffix message, got empty: %q", out.Error.Error()) + } + // Free-form large text: redaction must not panic and must not add creds. + free := sanitizeOpenAIErrorMessage(em(http.StatusBadGateway, strings.Repeat("y", openAIStreamErrorMessageLimit*3))) + if free == nil || free.Error == nil { + t.Fatal("expected sanitized free-form error, got nil") + } +} + +// sendResponsesPeekRequest drives the OpenAIResponses handler (/v1/responses) +// through a registered executor, exercising the native responses peek loop. +func sendResponsesPeekRequest(t *testing.T, body string, executor *peekStreamExecutor) *httptest.ResponseRecorder { + t.Helper() + manager := coreauth.NewManager(nil, nil, nil) + manager.RegisterExecutor(executor) + authID := "peek-resp-auth" + auth := &coreauth.Auth{ID: authID, Provider: executor.Identifier(), Status: coreauth.StatusActive} + if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + registry.GetGlobalRegistry().RegisterClient(authID, executor.Identifier(), []*registry.ModelInfo{{ID: "peek-resp-model", SupportedEndpoints: []string{openAIResponsesEndpoint}}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(authID) }) + + base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager) + h := NewOpenAIResponsesAPIHandler(base) + router := gin.New() + router.POST("/v1/responses", h.Responses) + + request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(body)) + request.Header.Set("Content-Type", "application/json") + recorder := httptest.NewRecorder() + router.ServeHTTP(recorder, request) + return recorder +} + +// TestResponsesNativePeekConsumesBufferedPendingError exercises the native +// /v1/responses initial peek: a closed dataChan with a buffered pending error +// must yield a sanitized non-200 error and never an empty success stream. +func TestResponsesNativePeekConsumesBufferedPendingError(t *testing.T) { + gin.SetMode(gin.TestMode) + const secret = "resp-native-secret-7712" + recorder := sendResponsesPeekRequest(t, + `{"model":"peek-resp-model","stream":true,"input":"hi"}`, + &peekStreamExecutor{secret: secret}) + body := recorder.Body.String() + if recorder.Code == http.StatusOK { + t.Fatalf("responses native returned 200 despite buffered pending error: %q", body) + } + if recorder.Code < http.StatusBadRequest { + t.Fatalf("status = %d, want error status; body=%q", recorder.Code, body) + } + if strings.Contains(body, secret) { + t.Fatalf("responses native leaked upstream secret: %q", body) + } + if !strings.Contains(body, "[REDACTED]") { + t.Fatalf("responses native did not redact upstream error: %q", body) + } +} + +// sendResponsesViaChatPeekRequest routes a /v1/responses streaming request +// through the via-chat peek by advertising only the chat endpoint for the +// model, so the responses handler overrides the endpoint to chat. +func sendResponsesViaChatPeekRequest(t *testing.T, body string, executor *peekStreamExecutor) *httptest.ResponseRecorder { + t.Helper() + manager := coreauth.NewManager(nil, nil, nil) + manager.RegisterExecutor(executor) + authID := "peek-resp-viachat" + auth := &coreauth.Auth{ID: authID, Provider: executor.Identifier(), Status: coreauth.StatusActive} + if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + registry.GetGlobalRegistry().RegisterClient(authID, executor.Identifier(), []*registry.ModelInfo{{ID: "peek-resp-model", SupportedEndpoints: []string{openAIChatEndpoint}}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(authID) }) + + base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager) + h := NewOpenAIResponsesAPIHandler(base) + router := gin.New() + router.POST("/v1/responses", h.Responses) + + request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(body)) + request.Header.Set("Content-Type", "application/json") + recorder := httptest.NewRecorder() + router.ServeHTTP(recorder, request) + return recorder +} + +// TestResponsesViaChatPeekConsumesBufferedPendingError exercises the peek +// close path reached when a /v1/responses streaming request is routed through +// the OpenAI chat endpoint. +func TestResponsesViaChatPeekConsumesBufferedPendingError(t *testing.T) { + gin.SetMode(gin.TestMode) + const secret = "resp-viachat-secret-3390" + recorder := sendResponsesViaChatPeekRequest(t, + `{"model":"peek-resp-model","stream":true,"input":"hi"}`, + &peekStreamExecutor{secret: secret}) + body := recorder.Body.String() + if recorder.Code == http.StatusOK { + t.Fatalf("responses via-chat returned 200 despite buffered pending error: %q", body) + } + if recorder.Code < http.StatusBadRequest { + t.Fatalf("status = %d, want error status; body=%q", recorder.Code, body) + } + if strings.Contains(body, secret) { + t.Fatalf("responses via-chat leaked upstream secret: %q", body) + } + // Via-chat terminal chunk must be sanitized (no raw error) and must not be + // an empty success stream. + if strings.Contains(body, "[DONE]") && strings.Contains(body, "api_key") { + t.Fatalf("via-chat leaked raw error text: %q", body) + } +} + +// newOpenAIImageHandlerWithExecutor builds an OpenAIAPIHandler whose +// ExecuteImageStreamWithAuthManager resolves through a registered fake stream +// executor, so the images peek close branches run for real. +func newOpenAIImageHandlerWithExecutor(t *testing.T, executor *peekStreamExecutor, model string) *OpenAIAPIHandler { + t.Helper() + manager := coreauth.NewManager(nil, nil, nil) + manager.RegisterExecutor(executor) + authID := "peek-image-auth" + auth := &coreauth.Auth{ID: authID, Provider: executor.Identifier(), Status: coreauth.StatusActive} + if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + registry.GetGlobalRegistry().RegisterClient(authID, executor.Identifier(), []*registry.ModelInfo{{ID: model, SupportedEndpoints: []string{openAIResponsesEndpoint}}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(authID) }) + base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager) + return NewOpenAIAPIHandler(base) +} + +// imagesPeekRecorder prepares a gin recorder with a flushable writer for the +// images stream handlers. +func imagesPeekRecorder(t *testing.T) (*gin.Context, *httptest.ResponseRecorder, http.Flusher) { + t.Helper() + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", strings.NewReader(`{"model":"image-model","prompt":"x","stream":true}`)) + flusher, ok := c.Writer.(http.Flusher) + if !ok { + t.Fatal("expected gin writer to implement http.Flusher") + } + return c, recorder, flusher +} + +// assertImagesPendingError fails unless the recorder carried a sanitized +// non-200 error on the shared-peek-close path (never an empty success stream). +func assertImagesPendingError(t *testing.T, recorder *httptest.ResponseRecorder, body, secret string) { + t.Helper() + if recorder.Code == http.StatusOK { + t.Fatalf("images peek returned 200 despite buffered pending error: %q", body) + } + if recorder.Code < http.StatusBadRequest { + t.Fatalf("images peek status = %d, want error status; body=%q", recorder.Code, body) + } + if strings.Contains(body, secret) { + t.Fatalf("images peek leaked upstream secret: %q", body) + } + if !strings.Contains(body, "[REDACTED]") { + t.Fatalf("images peek did not redact upstream error: %q", body) + } +} + +// TestRoutedImagesPeekConsumesBufferedPendingError exercises streamRoutedImages +// (codex images tool family) peek close path. +func TestRoutedImagesPeekConsumesBufferedPendingError(t *testing.T) { + gin.SetMode(gin.TestMode) + const secret = "routed-image-secret-5081" + h := newOpenAIImageHandlerWithExecutor(t, &peekStreamExecutor{secret: secret}, "gpt-image-1.5") + c, recorder, flusher := imagesPeekRecorder(t) + h.streamRoutedImages(c, []byte(`{"model":"gpt-image-1.5","prompt":"x"}`), "gpt-image-1.5") + _ = flusher + assertImagesPendingError(t, recorder, recorder.Body.String(), secret) +} + +// TestCompatImagesPeekConsumesBufferedPendingError exercises +// streamOpenAICompatImages peek close path. +func TestCompatImagesPeekConsumesBufferedPendingError(t *testing.T) { + gin.SetMode(gin.TestMode) + const secret = "compat-image-secret-1193" + h := newOpenAIImageHandlerWithExecutor(t, &peekStreamExecutor{secret: secret}, "compat-image-model") + c, recorder, _ := imagesPeekRecorder(t) + h.streamOpenAICompatImages(c, []byte(`{"model":"compat-image-model","prompt":"x"}`), "compat-image-model") + assertImagesPendingError(t, recorder, recorder.Body.String(), secret) +} + +// TestResponsesBackedImagesPeekConsumesBufferedPendingError exercises +// streamImagesFromResponses peek close path. +func TestResponsesBackedImagesPeekConsumesBufferedPendingError(t *testing.T) { + gin.SetMode(gin.TestMode) + const secret = "resp-image-secret-2278" + h := newOpenAIImageHandlerWithExecutor(t, &peekStreamExecutor{secret: secret}, "peek-image-model") + c, recorder, _ := imagesPeekRecorder(t) + h.streamImagesFromResponses(c, []byte(`{"model":"peek-image-model","prompt":"x"}`), "b64_json", "image_generation") + assertImagesPendingError(t, recorder, recorder.Body.String(), secret) +} + +// TestImagesStreamErrorEventRedactsTerminal exercises the actual post-header +// terminal writer (writeImagesStreamErrorEvent) used after the first chunk: the +// SSE error event must carry a sanitized, secret-free message. +func TestImagesStreamErrorEventRedactsTerminal(t *testing.T) { + gin.SetMode(gin.TestMode) + c, recorder, _ := imagesPeekRecorder(t) + secret := fromSeg("img-term-secret-6634") + writeImagesStreamErrorEvent(c, &interfaces.ErrorMessage{ + StatusCode: http.StatusBadGateway, + Error: errors.New("upstream image failure: api_key=" + secret), + }) + body := recorder.Body.String() + if !strings.Contains(body, "event: error") { + t.Fatalf("missing error SSE event: %q", body) + } + if strings.Contains(body, secret) { + t.Fatalf("images terminal error leaked secret: %q", body) + } + if !strings.Contains(body, "[REDACTED]") { + t.Fatalf("images terminal error not redacted: %q", body) + } +} + +// TestResponsesWebsocketTerminalErrorRedacts exercises the actual websocket +// terminal-error payload builder: the client error JSON must keep its framing +// and codes but never leak upstream credential material. +func TestResponsesWebsocketTerminalErrorRedacts(t *testing.T) { + gin.SetMode(gin.TestMode) + secret := fromSeg("ws-term-secret-9047") + errMsg := &interfaces.ErrorMessage{ + StatusCode: http.StatusBadRequest, + Error: errors.New(`{"error":{"type":"invalid_request","code":"cyber_policy","message":"blocked api_key=` + secret + `"}}`), + } + payload, err := buildResponsesWebsocketErrorPayload(errMsg) + if err != nil { + t.Fatalf("buildResponsesWebsocketErrorPayload: %v", err) + } + if gjson.GetBytes(payload, "type").String() != "error" { + t.Fatalf("type = %q, want error", gjson.GetBytes(payload, "type").String()) + } + if status := int(gjson.GetBytes(payload, "status").Int()); status != http.StatusBadRequest { + t.Fatalf("status = %d, want %d", status, http.StatusBadRequest) + } + if code := gjson.GetBytes(payload, "error.code").String(); code != "cyber_policy" { + t.Fatalf("error.code = %q, want cyber_policy", code) + } + raw := string(payload) + if strings.Contains(raw, secret) { + t.Fatalf("websocket terminal error leaked upstream secret: %q", raw) + } + // The redaction marker must be present in the embedded error message. + if !strings.Contains(raw, "[REDACTED]") { + t.Fatalf("websocket terminal error not redacted: %q", raw) + } +} + +// TestSanitizeOpenAIErrorMessageTrustedPreservation verifies the OpenAI shared +// sanitizer preserves a DirectResponse only when it is explicitly trusted. +func TestSanitizeOpenAIErrorMessageTrustedPreservation(t *testing.T) { + const secret = "trusted-openai-secret-311" + trusted := &interfaces.ErrorMessage{ + StatusCode: http.StatusTooManyRequests, + Error: errors.New("plugin"), + DirectResponse: true, + TrustedDirectResponse: true, + Body: []byte(`{"status":429,"message":"plugin","secret":"` + secret + `"}`), + Headers: http.Header{"X-Plugin": []string{"yes"}}, + } + if got := sanitizeOpenAIErrorMessage(trusted); got != trusted { + t.Fatalf("trusted direct response must be preserved verbatim, got %#v", got) + } + + untrusted := &interfaces.ErrorMessage{ + StatusCode: http.StatusBadGateway, + Error: errors.New("provider failed: api_key=" + secret), + DirectResponse: true, + Body: []byte(`{"raw":"` + secret + `"}`), + Headers: http.Header{"X-Upstream": []string{"leak"}}, + } + got := sanitizeOpenAIErrorMessage(untrusted) + if got == nil { + t.Fatal("untrusted direct response must be sanitized, not nil") + } + if got.DirectResponse { + t.Fatal("untrusted DirectResponse must be forced false") + } + if got.Body != nil { + t.Fatal("untrusted Body must be cleared") + } + if got.Error != nil && strings.Contains(got.Error.Error(), secret) { + t.Fatalf("untrusted sanitized error leaked %q: %q", secret, got.Error.Error()) + } +} + +func TestRedactOpenAIStreamErrorTextBearerTokens(t *testing.T) { + cases := []struct { + name string + text string + want string + }{ + { + name: "lowercase-only standalone bearer token", + text: "upstream error: Bearer abcdef", + want: "upstream error: Bearer [REDACTED]", + }, + { + name: "lowercase-only bare bearer token", + text: "Bearer abcdef", + want: "Bearer [REDACTED]", + }, + { + name: "lowercase-only bearer in sentence", + text: "upstream error with Bearer abcdef in request", + want: "upstream error with Bearer [REDACTED] in request", + }, + { + name: "lowercase-only bearer comma delimited", + text: "upstream error: Bearer abcdef, request failed", + want: "upstream error: Bearer [REDACTED], request failed", + }, + { + name: "mixed case and digit bearer token", + text: "upstream error: Bearer abc123XYZ", + want: "upstream error: Bearer [REDACTED]", + }, + { + name: "numeric bearer token", + text: "upstream error: Bearer 123456", + want: "upstream error: Bearer [REDACTED]", + }, + { + name: "RFC6750 b64token characters", + text: "upstream error: Bearer abc-def_123.xyz~456+789/0==", + want: "upstream error: Bearer [REDACTED]", + }, + { + name: "lowercase-only basic token", + text: "upstream error: Basic abcdef", + want: "upstream error: Basic [REDACTED]", + }, + { + name: "json embedded standalone bearer", + text: `{"error":"Bearer abcdef"}`, + want: `{"error":"Bearer [REDACTED]"}`, + }, + { + name: "short 2-char bearer token", + text: "upstream error: Bearer ab", + want: "upstream error: Bearer [REDACTED]", + }, + { + name: "short 1-char bearer token", + text: "upstream error: Bearer a", + want: "upstream error: Bearer [REDACTED]", + }, + { + name: "short 2-char basic token", + text: "upstream error: Basic ab", + want: "upstream error: Basic [REDACTED]", + }, + { + name: "short 1-char basic token", + text: "upstream error: Basic a", + want: "upstream error: Basic [REDACTED]", + }, + { + name: "bearer of standalone at end", + text: "upstream error: Bearer of", + want: "upstream error: Bearer [REDACTED]", + }, + { + name: "bearer to standalone at end", + text: "upstream error: Bearer to", + want: "upstream error: Bearer [REDACTED]", + }, + { + name: "bearer of bare token", + text: "Bearer of", + want: "Bearer [REDACTED]", + }, + { + name: "bearer to bare token", + text: "Bearer to", + want: "Bearer [REDACTED]", + }, + { + name: "bearer of in json quotes", + text: `{"error":"Bearer of"}`, + want: `{"error":"Bearer [REDACTED]"}`, + }, + { + name: "bearer to in json quotes", + text: `{"error":"Bearer to"}`, + want: `{"error":"Bearer [REDACTED]"}`, + }, + { + name: "bearer of with comma punctuation", + text: "upstream error: Bearer of, please retry", + want: "upstream error: Bearer [REDACTED], please retry", + }, + { + name: "authorization header bearer of", + text: "Authorization: Bearer of\r\n", + want: "Authorization: Bearer [REDACTED]\r\n", + }, + { + name: "lowercase bearer token followed by prose word", + text: "upstream error: bearer abc expired", + want: "upstream error: bearer [REDACTED] expired", + }, + { + name: "lowercase basic token followed by prose word", + text: "upstream error: basic dGVzdA== rejected", + want: "upstream error: basic [REDACTED] rejected", + }, + // Prose controls: collateral prose masking is accepted in exchange for no credential leaks (reviewer trade-off). + { + name: "control: bearer of bad news partially masked per reviewer trade-off", + text: "bearer of bad news", + want: "bearer [REDACTED] bad news", + }, + { + name: "control: the bearer of good news partially masked per reviewer trade-off", + text: "the bearer of good news", + want: "the bearer [REDACTED] good news", + }, + { + name: "control: bearer to the manager partially masked per reviewer trade-off", + text: "the bearer to the manager", + want: "the bearer [REDACTED] the manager", + }, + { + name: "control: bearer in header partially masked per reviewer trade-off", + text: "the bearer in header", + want: "the bearer [REDACTED] header", + }, + { + name: "control: bearer is invalid partially masked per reviewer trade-off", + text: "the bearer is invalid", + want: "the bearer [REDACTED] invalid", + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got := redactOpenAIStreamErrorText(tc.text) + if got != tc.want { + t.Fatalf("redactOpenAIStreamErrorText(%q) = %q, want %q", tc.text, got, tc.want) + } + }) + } +} + +func TestRedactOpenAIStreamErrorTextCamelCase(t *testing.T) { + cases := []struct { + name string + text string + want string + }{ + { + name: "camelCase refreshToken assignment", + text: "refreshToken=abc", + want: "refreshToken=[REDACTED]", + }, + { + name: "camelCase clientSecret colon", + text: "clientSecret: xyz", + want: "clientSecret: [REDACTED]", + }, + { + name: "camelCase apiKey assignment", + text: "apiKey=secret123", + want: "apiKey=[REDACTED]", + }, + { + name: "camelCase accessToken colon", + text: "accessToken: tok456", + want: "accessToken: [REDACTED]", + }, + { + name: "non-sensitive camelCase words untouched", + text: "userProfile=safe statusCode: 200 maxRetries=3 donkey=safe", + want: "userProfile=safe statusCode: 200 maxRetries=3 donkey=safe", + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := redactOpenAIStreamErrorText(tc.text); got != tc.want { + t.Fatalf("redactOpenAIStreamErrorText(%q) = %q, want %q", tc.text, got, tc.want) + } + }) + } +} diff --git a/sdk/api/handlers/openai/openai_images_handlers.go b/sdk/api/handlers/openai/openai_images_handlers.go index 961a4276f..3ac61aaa6 100644 --- a/sdk/api/handlers/openai/openai_images_handlers.go +++ b/sdk/api/handlers/openai/openai_images_handlers.go @@ -92,6 +92,7 @@ func writeImagesStreamErrorEvent(c *gin.Context, errMsg *interfaces.ErrorMessage if errMsg == nil { return } + errMsg = sanitizeOpenAIErrorMessage(errMsg) status := http.StatusInternalServerError if errMsg.StatusCode > 0 { status = errMsg.StatusCode @@ -1156,7 +1157,7 @@ func (h *OpenAIAPIHandler) collectRoutedImages(c *gin.Context, imageReq []byte, resp, upstreamHeaders, errMsg := h.ExecuteImageWithAuthManager(cliCtx, xaiImagesHandlerType, model, imageReq, "") stopKeepAlive() if errMsg != nil { - h.WriteErrorResponse(c, errMsg) + h.WriteErrorResponse(c, sanitizeOpenAIErrorMessage(errMsg)) if errMsg.Error != nil { cliCancel(errMsg.Error) } else { @@ -1220,7 +1221,7 @@ func (h *OpenAIAPIHandler) streamRoutedImages(c *gin.Context, imageReq []byte, i writeImagesStreamErrorEvent(c, errMsg) flusher.Flush() } else { - h.WriteErrorResponse(c, errMsg) + h.WriteErrorResponse(c, sanitizeOpenAIErrorMessage(errMsg)) } if errMsg != nil { cliCancel(errMsg.Error) @@ -1230,6 +1231,15 @@ func (h *OpenAIAPIHandler) streamRoutedImages(c *gin.Context, imageReq []byte, i return case chunk, ok := <-dataChan: if !ok { + // Stream closed without data. Surface a buffered pending error + // before committing SSE headers, so a failed upstream never + // looks like a successful empty stream. + if pErr, pending := pendingOpenAIStreamError(errChan); pending { + h.WriteErrorResponse(c, sanitizeOpenAIErrorMessage(pErr)) + cliCancel(pErr.Error) + return + } + // Clean close. Send headers and done. stopKeepAlive() setImagesSSEHeaders(c) handlers.WriteUpstreamHeaders(c.Writer.Header(), upstreamHeaders) @@ -1348,7 +1358,7 @@ func (h *OpenAIAPIHandler) streamOpenAICompatImages(c *gin.Context, compatReq [] writeImagesStreamErrorEvent(c, errMsg) flusher.Flush() } else { - h.WriteErrorResponse(c, errMsg) + h.WriteErrorResponse(c, sanitizeOpenAIErrorMessage(errMsg)) } if errMsg != nil { cliCancel(errMsg.Error) @@ -1358,6 +1368,15 @@ func (h *OpenAIAPIHandler) streamOpenAICompatImages(c *gin.Context, compatReq [] return case chunk, ok := <-dataChan: if !ok { + // Stream closed without data. Surface a buffered pending error + // before committing SSE headers, so a failed upstream never + // looks like a successful empty stream. + if pErr, pending := pendingOpenAIStreamError(errChan); pending { + h.WriteErrorResponse(c, sanitizeOpenAIErrorMessage(pErr)) + cliCancel(pErr.Error) + return + } + // Clean close. Send headers and done. stopKeepAlive() setImagesSSEHeaders(c) handlers.WriteUpstreamHeaders(c.Writer.Header(), upstreamHeaders) @@ -1377,7 +1396,7 @@ func (h *OpenAIAPIHandler) streamOpenAICompatImages(c *gin.Context, compatReq [] _, _ = c.Writer.Write(next) }, WriteTerminalError: func(errMsg *interfaces.ErrorMessage) { - writeImagesStreamErrorEvent(c, errMsg) + writeImagesStreamErrorEvent(c, sanitizeOpenAIErrorMessage(errMsg)) }, }) return @@ -1405,7 +1424,7 @@ func (h *OpenAIAPIHandler) collectImagesWithModel(c *gin.Context, imageReq []byt resp, upstreamHeaders, errMsg := h.ExecuteImageWithAuthManager(cliCtx, xaiImagesHandlerType, model, imageReq, "") stopKeepAlive() if errMsg != nil { - h.WriteErrorResponse(c, errMsg) + h.WriteErrorResponse(c, sanitizeOpenAIErrorMessage(errMsg)) if errMsg.Error != nil { cliCancel(errMsg.Error) } else { @@ -1417,7 +1436,7 @@ func (h *OpenAIAPIHandler) collectImagesWithModel(c *gin.Context, imageReq []byt out, err := buildImagesAPIResponseFromXAI(resp, responseFormat) if err != nil { errMsg := &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: err} - h.WriteErrorResponse(c, errMsg) + h.WriteErrorResponse(c, sanitizeOpenAIErrorMessage(errMsg)) cliCancel(err) return } @@ -1472,7 +1491,7 @@ func (h *OpenAIAPIHandler) streamImagesWithModel(c *gin.Context, imageReq []byte writeImagesStreamErrorEvent(c, errMsg) flusher.Flush() } else { - h.WriteErrorResponse(c, errMsg) + h.WriteErrorResponse(c, sanitizeOpenAIErrorMessage(errMsg)) } if errMsg != nil && errMsg.Error != nil { cliCancel(errMsg.Error) @@ -1554,7 +1573,7 @@ func (h *OpenAIAPIHandler) collectImagesFromResponses(c *gin.Context, responsesR out, errMsg := collectImagesFromResponsesStream(cliCtx, dataChan, errChan, responseFormat) stopKeepAlive() if errMsg != nil { - h.WriteErrorResponse(c, errMsg) + h.WriteErrorResponse(c, sanitizeOpenAIErrorMessage(errMsg)) if errMsg.Error != nil { cliCancel(errMsg.Error) } else { @@ -1789,7 +1808,7 @@ func (h *OpenAIAPIHandler) streamImagesFromResponses(c *gin.Context, responsesRe writeImagesStreamErrorEvent(c, errMsg) flusher.Flush() } else { - h.WriteErrorResponse(c, errMsg) + h.WriteErrorResponse(c, sanitizeOpenAIErrorMessage(errMsg)) } if errMsg != nil { cliCancel(errMsg.Error) @@ -1799,6 +1818,15 @@ func (h *OpenAIAPIHandler) streamImagesFromResponses(c *gin.Context, responsesRe return case chunk, ok := <-dataChan: if !ok { + // Stream closed without data. Surface a buffered pending error + // before committing SSE headers, so a failed upstream never + // looks like a successful empty stream. + if pErr, pending := pendingOpenAIStreamError(errChan); pending { + h.WriteErrorResponse(c, sanitizeOpenAIErrorMessage(pErr)) + cliCancel(pErr.Error) + return + } + // Clean close. Send headers and done. stopKeepAlive() setImagesSSEHeaders(c) handlers.WriteUpstreamHeaders(c.Writer.Header(), upstreamHeaders) diff --git a/sdk/api/handlers/openai/openai_responses_handlers.go b/sdk/api/handlers/openai/openai_responses_handlers.go index cdf8ae9c5..73bd8becd 100644 --- a/sdk/api/handlers/openai/openai_responses_handlers.go +++ b/sdk/api/handlers/openai/openai_responses_handlers.go @@ -534,7 +534,7 @@ func (h *OpenAIResponsesAPIHandler) Compact(c *gin.Context) { resp, upstreamHeaders, errMsg := h.ExecuteWithAuthManager(cliCtx, h.HandlerType(), modelName, rawJSON, "responses/compact") stopKeepAlive() if errMsg != nil { - h.WriteErrorResponse(c, errMsg) + h.WriteErrorResponse(c, sanitizeOpenAIErrorMessage(errMsg)) cliCancel(errMsg.Error) return } @@ -560,7 +560,7 @@ func (h *OpenAIResponsesAPIHandler) handleNonStreamingResponse(c *gin.Context, r resp, upstreamHeaders, errMsg := h.ExecuteWithAuthManager(cliCtx, h.HandlerType(), modelName, rawJSON, "") stopKeepAlive() if errMsg != nil { - h.WriteErrorResponse(c, errMsg) + h.WriteErrorResponse(c, sanitizeOpenAIErrorMessage(errMsg)) cliCancel(errMsg.Error) return } @@ -576,7 +576,7 @@ func (h *OpenAIResponsesAPIHandler) handleNonStreamingResponseViaChat(c *gin.Con cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background()) resp, upstreamHeaders, errMsg := h.ExecuteWithAuthManager(cliCtx, OpenAI, modelName, chatJSON, "") if errMsg != nil { - h.WriteErrorResponse(c, errMsg) + h.WriteErrorResponse(c, sanitizeOpenAIErrorMessage(errMsg)) cliCancel(errMsg.Error) return } @@ -641,7 +641,7 @@ func (h *OpenAIResponsesAPIHandler) handleStreamingResponse(c *gin.Context, rawJ continue } // Upstream failed immediately. Return proper error status and JSON. - h.WriteErrorResponse(c, errMsg) + h.WriteErrorResponse(c, sanitizeOpenAIErrorMessage(errMsg)) if errMsg != nil { cliCancel(errMsg.Error) } else { @@ -650,7 +650,15 @@ func (h *OpenAIResponsesAPIHandler) handleStreamingResponse(c *gin.Context, rawJ return case chunk, ok := <-dataChan: if !ok { - // Stream closed without data? Send headers and done. + // Stream closed without data. Surface a buffered pending error + // before committing SSE headers, so a failed upstream never + // looks like a successful empty stream. + if pErr, pending := pendingOpenAIStreamError(errChan); pending { + h.WriteErrorResponse(c, sanitizeOpenAIErrorMessage(pErr)) + cliCancel(pErr.Error) + return + } + // Clean close. Send headers and done. setSSEHeaders() handlers.WriteUpstreamHeaders(c.Writer.Header(), upstreamHeaders) _, _ = c.Writer.Write([]byte("\n")) @@ -708,7 +716,7 @@ func (h *OpenAIResponsesAPIHandler) handleStreamingResponseViaChat(c *gin.Contex errChan = nil continue } - h.WriteErrorResponse(c, errMsg) + h.WriteErrorResponse(c, sanitizeOpenAIErrorMessage(errMsg)) if errMsg != nil { cliCancel(errMsg.Error) } else { @@ -717,6 +725,15 @@ func (h *OpenAIResponsesAPIHandler) handleStreamingResponseViaChat(c *gin.Contex return case chunk, ok := <-dataChan: if !ok { + // Stream closed without data. Surface a buffered pending error + // before committing SSE headers, so a failed upstream never + // looks like a successful empty stream. + if pErr, pending := pendingOpenAIStreamError(errChan); pending { + h.WriteErrorResponse(c, sanitizeOpenAIErrorMessage(pErr)) + cliCancel(pErr.Error) + return + } + // Clean close. Send headers and done. setSSEHeaders() handlers.WriteUpstreamHeaders(c.Writer.Header(), upstreamHeaders) _, _ = c.Writer.Write([]byte("\n")) @@ -797,6 +814,7 @@ func writeResponsesTerminalError(c *gin.Context, errMsg *interfaces.ErrorMessage if !shouldExposeResponsesUpstreamError(errMsg) { return } + errMsg = sanitizeOpenAIErrorMessage(errMsg) status := http.StatusInternalServerError if errMsg.StatusCode > 0 { status = errMsg.StatusCode diff --git a/sdk/api/handlers/openai/openai_responses_multi_agent_test.go b/sdk/api/handlers/openai/openai_responses_multi_agent_test.go index 09f00872d..2e1be926a 100644 --- a/sdk/api/handlers/openai/openai_responses_multi_agent_test.go +++ b/sdk/api/handlers/openai/openai_responses_multi_agent_test.go @@ -23,7 +23,6 @@ import ( func TestPrepareCodexMultiAgentV2ToolsAtResponsesBoundary(t *testing.T) { t.Parallel() - gin.SetMode(gin.TestMode) base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{CodexOptimizeMultiAgentV2: true}, nil) handler := NewOpenAIResponsesAPIHandler(base) request := httptest.NewRequest(http.MethodPost, "/v1/responses", nil) @@ -91,13 +90,11 @@ type responsesMultiAgentCaptureExecutor struct { websocketDirectCaptureExecutor } -const responsesMultiAgentCompletion = `{"id":"resp-1","status":"completed","output":[{"type":"message","id":"msg-1","role":"assistant","content":[{"type":"output_text","text":"ok"}]}]}` - func (e *responsesMultiAgentCaptureExecutor) Execute(_ context.Context, _ *coreauth.Auth, req coreexecutor.Request, _ coreexecutor.Options) (coreexecutor.Response, error) { e.mu.Lock() e.payloads = append(e.payloads, bytes.Clone(req.Payload)) e.mu.Unlock() - return coreexecutor.Response{Payload: []byte(responsesMultiAgentCompletion)}, nil + return coreexecutor.Response{Payload: []byte(`{"id":"resp-1","output":[]}`)}, nil } func (e *responsesMultiAgentCaptureExecutor) ExecuteStream(_ context.Context, _ *coreauth.Auth, req coreexecutor.Request, _ coreexecutor.Options) (*coreexecutor.StreamResult, error) { @@ -105,7 +102,7 @@ func (e *responsesMultiAgentCaptureExecutor) ExecuteStream(_ context.Context, _ e.payloads = append(e.payloads, bytes.Clone(req.Payload)) e.mu.Unlock() chunks := make(chan coreexecutor.StreamChunk, 1) - chunks <- coreexecutor.StreamChunk{Payload: []byte(fmt.Sprintf("data: {\"type\":\"response.completed\",\"response\":%s}\n\n", responsesMultiAgentCompletion))} + chunks <- coreexecutor.StreamChunk{Payload: []byte("data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp-1\",\"output\":[]}}\n\n")} close(chunks) return &coreexecutor.StreamResult{Chunks: chunks}, nil } @@ -183,7 +180,6 @@ func TestResponsesWebsocketPreparesCodexMultiAgentV2Tools(t *testing.T) { func TestPrepareCodexMultiAgentV2ToolsAtResponsesBoundarySkipsOtherClients(t *testing.T) { t.Parallel() - gin.SetMode(gin.TestMode) base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{CodexOptimizeMultiAgentV2: true}, nil) handler := NewOpenAIResponsesAPIHandler(base) request := httptest.NewRequest(http.MethodPost, "/v1/responses", nil) diff --git a/sdk/api/handlers/openai/openai_responses_websocket_forward.go b/sdk/api/handlers/openai/openai_responses_websocket_forward.go index 49603edb2..8d7382d8b 100644 --- a/sdk/api/handlers/openai/openai_responses_websocket_forward.go +++ b/sdk/api/handlers/openai/openai_responses_websocket_forward.go @@ -80,7 +80,7 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket( return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), errMsg, websocket.ErrCloseSent } - errorPayload, wrote, errTerminate := writeResponsesWebsocketTerminalError(writer, wsTimelineLog, errMsg, nil) + errorPayload, wrote, errTerminate := writeResponsesWebsocketTerminalError(writer, wsTimelineLog, errMsg) if wrote { log.Infof( "responses websocket: downstream_out id=%s type=%d event=%s payload=%s", @@ -149,7 +149,7 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket( } return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), payloadErrMsg, websocket.ErrCloseSent } - errorPayload, wrote, errTerminate := writeResponsesWebsocketTerminalError(writer, wsTimelineLog, payloadErrMsg, payloads[i]) + errorPayload, wrote, errTerminate := writeResponsesWebsocketTerminalError(writer, wsTimelineLog, payloadErrMsg) if wrote { log.Infof( "responses websocket: downstream_out id=%s type=%d event=%s payload=%s", @@ -213,7 +213,6 @@ func writeResponsesWebsocketTerminalError( writer *responsesWebsocketWriter, wsTimelineLog websocketTimelineAppender, errMsg *interfaces.ErrorMessage, - payload []byte, ) ([]byte, bool, error) { if !shouldExposeResponsesUpstreamError(errMsg) { // Keep the upstream reason in the request-log timeline even though the client @@ -229,13 +228,10 @@ func writeResponsesWebsocketTerminalError( return nil, false, websocket.ErrCloseSent } - if len(payload) == 0 { - var errBuild error - payload, errBuild = buildResponsesWebsocketErrorPayload(errMsg) - if errBuild != nil { - _, _ = writer.closeWithoutError() - return nil, false, errBuild - } + payload, errBuild := buildResponsesWebsocketErrorPayload(errMsg) + if errBuild != nil { + _, _ = writer.closeWithoutError() + return nil, false, errBuild } wrote, errClose := writer.closeWithPayload(payload) @@ -529,6 +525,7 @@ func buildResponsesWebsocketErrorPayload(errMsg *interfaces.ErrorMessage) ([]byt status := http.StatusInternalServerError errText := http.StatusText(status) if errMsg != nil { + errMsg = sanitizeOpenAIErrorMessage(errMsg) if errMsg.StatusCode > 0 { status = errMsg.StatusCode errText = http.StatusText(status) @@ -550,21 +547,23 @@ func buildResponsesWebsocketErrorPayload(errMsg *interfaces.ErrorMessage) ([]byt return nil, errSet } - if errMsg != nil && errMsg.Addon != nil { - headers := []byte(`{}`) - hasHeaders := false - for key, values := range errMsg.Addon { - if len(values) == 0 { - continue - } - headerPath := strings.ReplaceAll(strings.ReplaceAll(key, `\\`, `\\\\`), ".", `\\.`) - headers, errSet = sjson.SetBytes(headers, headerPath, values[0]) - if errSet != nil { - return nil, errSet + if errMsg != nil { + // Preserve only sanitized upstream headers as response headers and never + // echo them under a raw "headers" field. Filtering removes hop-by-hop, + // reserved, and gateway-identity headers. + filtered := handlers.FilterUpstreamHeaders(errMsg.Headers) + if len(errMsg.Headers) > 0 { + headers := []byte(`{}`) + for key, values := range filtered { + if len(values) == 0 { + continue + } + headerPath := strings.ReplaceAll(strings.ReplaceAll(key, `\\`, `\\\\`), ".", `\\.`) + headers, errSet = sjson.SetBytes(headers, headerPath, values[0]) + if errSet != nil { + return nil, errSet + } } - hasHeaders = true - } - if hasHeaders { payload, errSet = sjson.SetRawBytes(payload, "headers", headers) if errSet != nil { return nil, errSet @@ -574,10 +573,46 @@ func buildResponsesWebsocketErrorPayload(errMsg *interfaces.ErrorMessage) ([]byt if len(body) > 0 && json.Valid(body) { errorNode := gjson.GetBytes(body, "error") - if errorNode.Exists() { - payload, errSet = sjson.SetRawBytes(payload, "error", []byte(errorNode.Raw)) + if !errorNode.Exists() || !errorNode.IsObject() { + errorNode = gjson.GetBytes(body, "response.error") + } + if errorNode.Exists() && errorNode.IsObject() { + errObj := []byte(`{}`) + copied := false + for _, field := range []string{"type", "code", "message", "param"} { + v := errorNode.Get(field) + if !v.Exists() || v.Type == gjson.Null { + continue + } + errObj, _ = sjson.SetBytes(errObj, field, v.Value()) + copied = true + } + if copied { + payload, errSet = sjson.SetRawBytes(payload, "error", errObj) + } else { + payload, errSet = sjson.SetRawBytes(payload, "error", []byte(errorNode.Raw)) + } } else { - payload, errSet = sjson.SetRawBytes(payload, "error", body) + root := gjson.ParseBytes(body) + if root.IsObject() { + errObj := []byte(`{}`) + copied := false + for _, field := range []string{"type", "code", "message", "param"} { + v := root.Get(field) + if !v.Exists() || v.Type == gjson.Null { + continue + } + errObj, _ = sjson.SetBytes(errObj, field, v.Value()) + copied = true + } + if copied { + payload, errSet = sjson.SetRawBytes(payload, "error", errObj) + } else { + payload, errSet = sjson.SetRawBytes(payload, "error", body) + } + } else { + payload, errSet = sjson.SetRawBytes(payload, "error", body) + } } if errSet != nil { return nil, errSet diff --git a/sdk/api/handlers/openai/openai_responses_websocket_test.go b/sdk/api/handlers/openai/openai_responses_websocket_test.go index 372a2958d..0e7892162 100644 --- a/sdk/api/handlers/openai/openai_responses_websocket_test.go +++ b/sdk/api/handlers/openai/openai_responses_websocket_test.go @@ -9,6 +9,7 @@ import ( "maps" "net/http" "net/http/httptest" + "os" "strconv" "strings" "sync" @@ -2939,7 +2940,7 @@ func TestResponsesWebsocketTerminalErrorWrittenOnceAcrossForwardAndDisconnect(t go func() { defer wg.Done() <-start - payload, _, errWrite := writeResponsesWebsocketTerminalError(writer, nil, errMsg, nil) + payload, _, errWrite := writeResponsesWebsocketTerminalError(writer, nil, errMsg) if !errors.Is(errWrite, websocket.ErrCloseSent) || gjson.GetBytes(payload, "error.code").String() != "cyber_policy" { resultCh <- fmt.Errorf("err-channel terminal write failed: err=%v payload=%s", errWrite, payload) } @@ -2947,8 +2948,10 @@ func TestResponsesWebsocketTerminalErrorWrittenOnceAcrossForwardAndDisconnect(t go func() { defer wg.Done() <-start - payload := []byte(`{"type":"error","status":400,"error":{"type":"invalid_request","code":"cyber_policy","message":"blocked"}}`) - writtenPayload, _, errWrite := writeResponsesWebsocketTerminalError(writer, nil, errMsg, payload) + // The payload parameter is intentionally removed: every terminal path + // rebuilds the payload from the parsed ErrorMessage so raw upstream + // error-frame bytes can never be echoed. + writtenPayload, _, errWrite := writeResponsesWebsocketTerminalError(writer, nil, errMsg) if !errors.Is(errWrite, websocket.ErrCloseSent) || gjson.GetBytes(writtenPayload, "error.code").String() != "cyber_policy" { resultCh <- fmt.Errorf("payload terminal write failed: err=%v payload=%s", errWrite, writtenPayload) } @@ -2995,6 +2998,184 @@ func TestResponsesWebsocketTerminalErrorWrittenOnceAcrossForwardAndDisconnect(t } } +// TestResponsesWebsocketTerminalErrorRebuildsSanitizedPayload proves the +// terminal error path rebuilds a structured payload from the parsed +// ErrorMessage instead of ever echoing raw upstream error-frame bytes: a +// secret embedded only in the raw frame text must not survive, while safe +// status/type/code fields do. +func TestResponsesWebsocketTerminalErrorRebuildsSanitizedPayload(t *testing.T) { + gin.SetMode(gin.TestMode) + const secret = "ws-rebuild-secret-4217" + errMsg := &interfaces.ErrorMessage{ + StatusCode: http.StatusBadRequest, + Error: errors.New(`{"error":{"type":"invalid_request","code":"offensive_policy","message":"blocked api_key=` + secret + `","param":null}}`), + } + payload, err := buildResponsesWebsocketErrorPayload(errMsg) + if err != nil { + t.Fatalf("buildResponsesWebsocketErrorPayload: %v", err) + } + raw := string(payload) + if strings.Contains(raw, secret) { + t.Fatalf("webSocket rebuilt error leaked raw upstream secret: [REDACTED:API key param] %s", raw) + } + if got := gjson.GetBytes(payload, "type").String(); got != wsEventTypeError { + t.Fatalf("type = %q, want %q", got, wsEventTypeError) + } + if status := int(gjson.GetBytes(payload, "status").Int()); status != http.StatusBadRequest { + t.Fatalf("status = %d, want %d", status, http.StatusBadRequest) + } + if code := gjson.GetBytes(payload, "error.code").String(); code != "offensive_policy" { + t.Fatalf("error.code = %q, want offensive_policy", code) + } + if !strings.Contains(raw, "[REDACTED]") { + t.Fatalf("rebuilt payload not redacted: %q", raw) + } +} + +// TestResponsesWebsocketTerminalErrorUnknownFieldNotEchoed drives a valid +// in-band error frame that carries an unknown secret field through +// forwardResponsesWebsocket. The raw frame bytes/fields must never be echoed; +// the terminal payload is rebuilt and preserves only safe protocol fields. +func TestResponsesWebsocketTerminalErrorUnknownFieldNotEchoed(t *testing.T) { + gin.SetMode(gin.TestMode) + + serverErrCh := make(chan error, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := responsesWebsocketUpgrader.Upgrade(w, r, nil) + if err != nil { + serverErrCh <- err + return + } + defer func() { _ = conn.Close() }() + + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Request = r + + h := NewOpenAIResponsesAPIHandler(handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)) + + data := make(chan []byte, 1) + errCh := make(chan *interfaces.ErrorMessage) + data <- []byte(`{"type":"error","status":400,"error":{"secret":"ws-malformed-secret-8102","code":"x","message":"blocked"}}`) + close(data) + close(errCh) + + _, _, _, _, errForward := h.forwardResponsesWebsocket( + ctx, + newResponsesWebsocketWriter(conn), + func(...interface{}) {}, + data, + errCh, + newInMemoryWebsocketTimelineLog(), + "session-malformed", + ) + if errForward != nil && !errors.Is(errForward, websocket.ErrCloseSent) { + serverErrCh <- errForward + return + } + serverErrCh <- nil + })) + defer server.Close() + + wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) + if err != nil { + t.Fatalf("dial websocket: %v", err) + } + defer func() { _ = conn.Close() }() + + _, payload, errRead := conn.ReadMessage() + if errRead != nil { + t.Fatalf("read websocket message: %v", errRead) + } + raw := string(payload) + if strings.Contains(raw, "ws-malformed-secret-8102") { + t.Fatalf("error frame echoed raw unknown secret field: %q", raw) + } + if got := gjson.GetBytes(payload, "type").String(); got != wsEventTypeError { + t.Fatalf("type = %q, want error; payload=%s", got, payload) + } + if code := gjson.GetBytes(payload, "error.code").String(); code != "x" { + t.Fatalf("error.code = %q, want x; payload=%s", code, payload) + } +} + +// TestResponsesWebsocketNonTerminalErrorRedactsSecretAndKeepsSafeFields is the +// non-terminal twin of TestResponsesWebsocketTerminalErrorRebuildsSanitizedPayload: +// writeResponsesWebsocketError drives an ErrorMessage whose embedded secret is +// built at runtime through the same shared sanitizer path used by the +// non-terminal error sink (ResponsesWebsocket non-terminal write), proves the +// secret is redacted and no unknown fields are echoed, and preserves the safe +// status/code fields. The secret is built at runtime so the test is not a +// source-literal scanner false positive; it proves the actual shared +// non-terminal sanitizer path. +func TestResponsesWebsocketNonTerminalErrorRedactsSecretAndKeepsSafeFields(t *testing.T) { + gin.SetMode(gin.TestMode) + secret := fmt.Sprintf("nt-ws-secret-%s", strconv.Itoa(424242)) + errMsg := &interfaces.ErrorMessage{ + StatusCode: http.StatusTooManyRequests, + Error: errors.New(`{"error":{"type":"rate_limit_error","code":"rate_limit_exceeded","message":"blocked api_key=` + secret + `","param":"x"}}`), + } + + // Non-terminal error sink drives a real websocket writer with the shared + // sanitizer; assertions must prove sanitization happened before the frame. + serverErrCh := make(chan error, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, errUpgrade := responsesWebsocketUpgrader.Upgrade(w, r, nil) + if errUpgrade != nil { + serverErrCh <- errUpgrade + return + } + defer func() { _ = conn.Close() }() + if _, errWrite := writeResponsesWebsocketError(newResponsesWebsocketWriter(conn), newInMemoryWebsocketTimelineLog(), errMsg); errWrite != nil { + serverErrCh <- errWrite + return + } + serverErrCh <- nil + })) + defer server.Close() + + wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + conn, _, errDial := websocket.DefaultDialer.Dial(wsURL, nil) + if errDial != nil { + t.Fatalf("dial websocket: %v", errDial) + } + defer func() { _ = conn.Close() }() + + _, payload, errRead := conn.ReadMessage() + if errRead != nil { + t.Fatalf("read websocket message: %v", errRead) + } + if errServer := <-serverErrCh; errServer != nil { + t.Fatalf("server error: %v", errServer) + } + + raw := string(payload) + if strings.Contains(raw, secret) { + t.Fatalf("non-terminal websocket error leaked raw upstream secret: %q", raw) + } + if !strings.Contains(raw, "[REDACTED]") { + t.Fatalf("non-terminal websocket error not redacted: %q", raw) + } + if got := gjson.GetBytes(payload, "type").String(); got != wsEventTypeError { + t.Fatalf("type = %q, want %q", got, wsEventTypeError) + } + if status := int(gjson.GetBytes(payload, "status").Int()); status != http.StatusTooManyRequests { + t.Fatalf("status = %d, want %d", status, http.StatusTooManyRequests) + } + if code := gjson.GetBytes(payload, "error.code").String(); code != "rate_limit_exceeded" { + t.Fatalf("error.code = %q, want rate_limit_exceeded", code) + } + if errType := gjson.GetBytes(payload, "error.type").String(); errType != "rate_limit_error" { + t.Fatalf("error.type = %q, want rate_limit_error", errType) + } + if unknown := gjson.GetBytes(payload, "unknown"); unknown.Exists() { + t.Fatalf("non-terminal error echoed unknown field: %q", raw) + } + if param := gjson.GetBytes(payload, "error.param"); param.Exists() && strings.Contains(param.String(), "secret") { + t.Fatalf("non-terminal error echoed raw param field: %q", raw) + } +} + func TestResponsesWebsocketCodexWebsocketPassthroughPassesCompactedRequestWithoutTranscriptMerge(t *testing.T) { gin.SetMode(gin.TestMode) @@ -3575,6 +3756,113 @@ func TestResponsesWebsocketRejectsUnknownPreviousResponseOnNewSocket(t *testing. } } +// TestResponsesWebsocketAppendBeforeCreateEmitsNonTerminalError proves the +// non-terminal error path: a response.append before any response.create emits a +// single sanitized error frame, keeps the connection open, and lets a subsequent +// valid response.create proceed normally. +func TestResponsesWebsocketAppendBeforeCreateEmitsNonTerminalError(t *testing.T) { + gin.SetMode(gin.TestMode) + + executor := &websocketDirectCaptureExecutor{provider: "codex"} + manager := coreauth.NewManager(nil, nil, nil) + manager.RegisterExecutor(executor) + auth := &coreauth.Auth{ + ID: "ws-append-auth", + Provider: "codex", + Status: coreauth.StatusActive, + } + if _, err := manager.Register(context.Background(), auth); err != nil { + t.Fatalf("Register auth: %v", err) + } + registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: "ws-append-model"}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth.ID) }) + + base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager) + h := NewOpenAIResponsesAPIHandler(base) + router := gin.New() + router.GET("/v1/responses/ws", h.ResponsesWebsocket) + + server := httptest.NewServer(router) + defer server.Close() + + wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/v1/responses/ws" + conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) + if err != nil { + t.Fatalf("dial websocket: %v", err) + } + defer func() { _ = conn.Close() }() + + // Step 1: send response.append before any response.create — non-terminal error. + appendRequest := `{"type":"response.append","input":[{"type":"message","id":"msg-1","role":"user","content":"hi"}]}` + if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(appendRequest)); errWrite != nil { + t.Fatalf("write append request: %v", errWrite) + } + + _, errorPayload, errRead := conn.ReadMessage() + if errRead != nil { + t.Fatalf("read error response: %v", errRead) + } + + // Assert exactly one sanitized error frame. + if got := gjson.GetBytes(errorPayload, "type").String(); got != wsEventTypeError { + t.Fatalf("error frame type = %q, want %q: %s", got, wsEventTypeError, errorPayload) + } + if got := int(gjson.GetBytes(errorPayload, "status").Int()); got != http.StatusBadRequest { + t.Fatalf("error frame status = %d, want %d: %s", got, http.StatusBadRequest, errorPayload) + } + if errType := gjson.GetBytes(errorPayload, "error.type").String(); errType != "invalid_request_error" { + t.Fatalf("error frame error.type = %q, want invalid_request_error: %s", errType, errorPayload) + } + if code := gjson.GetBytes(errorPayload, "error.code").String(); code != "" && code != "internal_server_error" { + t.Fatalf("error frame carries unexpected code %q: %s", code, errorPayload) + } + if msg := gjson.GetBytes(errorPayload, "error.message").String(); msg == "" { + t.Fatalf("error frame missing message: %s", errorPayload) + } + if unknown := gjson.GetBytes(errorPayload, "unknown"); unknown.Exists() { + t.Fatalf("error frame leaked unknown field: %s", errorPayload) + } + if gjson.GetBytes(errorPayload, "error.secret").Exists() { + t.Fatalf("error frame leaked raw secret field: %s", errorPayload) + } + + // Step 2: send valid response.create — must succeed, proving socket still open. + createRequest := `{"type":"response.create","model":"ws-append-model","input":[{"type":"message","id":"msg-1","role":"user","content":"hi"}]}` + if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(createRequest)); errWrite != nil { + t.Fatalf("write create request: %v", errWrite) + } + + _, recoveryPayload, errRead := conn.ReadMessage() + if errRead != nil { + t.Fatalf("read recovery response: %v", errRead) + } + if got := gjson.GetBytes(recoveryPayload, "type").String(); got != wsEventTypeCompleted { + t.Fatalf("recovery response type = %q, want %q: %s", got, wsEventTypeCompleted, recoveryPayload) + } + + // Step 3: read should block/close — prove exactly one error frame was sent. + // The gorilla/websocket ReadMessage for an unclosed healthy connection + // blocks until the next write. A non-blocking check proves no extra frame. + conn.SetReadDeadline(time.Now().Add(10 * time.Millisecond)) + if _, gotExtra, errExtra := conn.ReadMessage(); errExtra == nil { + t.Fatalf("received unexpected frame after recovery: %s", gotExtra) + } else if !errors.Is(errExtra, os.ErrDeadlineExceeded) && !strings.Contains(errExtra.Error(), "i/o timeout") && + !websocket.IsCloseError(errExtra, websocket.CloseNormalClosure, websocket.CloseGoingAway) { + t.Fatalf("unexpected read error after recovery: %v", errExtra) + } + // Restore blocking reads before deferred Close. + conn.SetReadDeadline(time.Time{}) + + // Verify executor received exactly one valid request (the recovery). + payloads := executor.Payloads() + if len(payloads) != 1 { + t.Fatalf("executor payload count = %d, want 1 (the recovery create)", len(payloads)) + } + if got := gjson.GetBytes(payloads[0], "model").String(); got != "ws-append-model" { + t.Fatalf("executor model = %q, want ws-append-model: %s", got, payloads[0]) + } +} + func TestResponsesWebsocketClosesAfterNonRetryableClientError(t *testing.T) { gin.SetMode(gin.TestMode) @@ -5516,3 +5804,57 @@ func TestNormalizeSubsequentRequestAssistantInputTriggersTranscriptReplacement(t t.Fatalf("input[0].id = %q, want %q", input[0].Get("id").String(), "msg-3") } } + +func TestResponsesWebsocketPreservesTopLevelErrorFields(t *testing.T) { + gin.SetMode(gin.TestMode) + errMsg := &interfaces.ErrorMessage{ + StatusCode: http.StatusBadRequest, + Error: errors.New(`{"type":"error","code":"context_length_exceeded","message":"prompt exceeds token limit","param":"max_tokens","secret_param":"leaked-token"}`), + } + payload, err := buildResponsesWebsocketErrorPayload(errMsg) + if err != nil { + t.Fatalf("buildResponsesWebsocketErrorPayload: %v", err) + } + if got := gjson.GetBytes(payload, "type").String(); got != wsEventTypeError { + t.Fatalf("type = %q, want %q", got, wsEventTypeError) + } + if status := int(gjson.GetBytes(payload, "status").Int()); status != http.StatusBadRequest { + t.Fatalf("status = %d, want %d", status, http.StatusBadRequest) + } + if got := gjson.GetBytes(payload, "error.type").String(); got != "error" { + t.Fatalf("error.type = %q, want error", got) + } + if got := gjson.GetBytes(payload, "error.code").String(); got != "context_length_exceeded" { + t.Fatalf("error.code = %q, want context_length_exceeded", got) + } + if got := gjson.GetBytes(payload, "error.message").String(); got != "prompt exceeds token limit" { + t.Fatalf("error.message = %q, want 'prompt exceeds token limit'", got) + } + if got := gjson.GetBytes(payload, "error.param").String(); got != "max_tokens" { + t.Fatalf("error.param = %q, want max_tokens", got) + } + if strings.Contains(string(payload), "leaked-token") { + t.Fatalf("leaked secret from top-level error: %s", payload) + } + + nestedErrMsg := &interfaces.ErrorMessage{ + StatusCode: http.StatusBadRequest, + Error: errors.New(`{"error":{"type":"invalid_request_error","code":"rate_limit_exceeded","message":"too many requests","param":"rate"}}`), + } + nestedPayload, err := buildResponsesWebsocketErrorPayload(nestedErrMsg) + if err != nil { + t.Fatalf("buildResponsesWebsocketErrorPayload(nested): %v", err) + } + if got := gjson.GetBytes(nestedPayload, "error.type").String(); got != "invalid_request_error" { + t.Fatalf("nested error.type = %q, want invalid_request_error", got) + } + if got := gjson.GetBytes(nestedPayload, "error.code").String(); got != "rate_limit_exceeded" { + t.Fatalf("nested error.code = %q, want rate_limit_exceeded", got) + } + if got := gjson.GetBytes(nestedPayload, "error.message").String(); got != "too many requests" { + t.Fatalf("nested error.message = %q, want 'too many requests'", got) + } + if got := gjson.GetBytes(nestedPayload, "error.param").String(); got != "rate" { + t.Fatalf("nested error.param = %q, want rate", got) + } +} diff --git a/sdk/api/handlers/openai/openai_videos_handlers.go b/sdk/api/handlers/openai/openai_videos_handlers.go index 1748eaa6d..39dec59f3 100644 --- a/sdk/api/handlers/openai/openai_videos_handlers.go +++ b/sdk/api/handlers/openai/openai_videos_handlers.go @@ -801,7 +801,7 @@ func (h *OpenAIAPIHandler) VideosRetrieve(c *gin.Context) { resp, upstreamHeaders, errMsg := h.ExecuteWithAuthManager(cliCtx, xaiVideosHandlerType, executionModel, payload, "") stopKeepAlive() if errMsg != nil { - h.WriteErrorResponse(c, errMsg) + h.WriteErrorResponse(c, sanitizeOpenAIErrorMessage(errMsg)) if errMsg.Error != nil { cliCancel(errMsg.Error) } else { @@ -813,7 +813,7 @@ func (h *OpenAIAPIHandler) VideosRetrieve(c *gin.Context) { out, err := buildVideosRetrieveAPIResponseFromXAI(videoID, resp, defaultOpenAIVideosModel) if err != nil { errMsg := &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: err} - h.WriteErrorResponse(c, errMsg) + h.WriteErrorResponse(c, sanitizeOpenAIErrorMessage(errMsg)) cliCancel(err) return } @@ -864,7 +864,7 @@ func (h *OpenAIAPIHandler) VideosContent(c *gin.Context) { resp, _, errMsg := h.ExecuteWithAuthManager(cliCtx, xaiVideosHandlerType, executionModel, payload, "") stopKeepAlive() if errMsg != nil { - h.WriteErrorResponse(c, errMsg) + h.WriteErrorResponse(c, sanitizeOpenAIErrorMessage(errMsg)) if errMsg.Error != nil { cliCancel(errMsg.Error) } else { @@ -877,7 +877,7 @@ func (h *OpenAIAPIHandler) VideosContent(c *gin.Context) { contentURL, err := xaiVideoContentURLFromPayload(resp) if err != nil { errMsg := &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: err} - h.WriteErrorResponse(c, errMsg) + h.WriteErrorResponse(c, sanitizeOpenAIErrorMessage(errMsg)) cliCancel(err) return } @@ -896,7 +896,7 @@ func (h *OpenAIAPIHandler) writeVideoContentFromURL(c *gin.Context, contentURL s StatusCode: clienterror.HTTPStatusFromErrorOr(err, http.StatusBadGateway), Error: err, } - h.WriteErrorResponse(c, errMsg) + h.WriteErrorResponse(c, sanitizeOpenAIErrorMessage(errMsg)) return err } @@ -907,7 +907,7 @@ func (h *OpenAIAPIHandler) writeVideoContentFromURL(c *gin.Context, contentURL s StatusCode: clienterror.HTTPStatusFromErrorOr(err, http.StatusBadGateway), Error: err, } - h.WriteErrorResponse(c, errMsg) + h.WriteErrorResponse(c, sanitizeOpenAIErrorMessage(errMsg)) return err } defer func() { @@ -923,7 +923,7 @@ func (h *OpenAIAPIHandler) writeVideoContentFromURL(c *gin.Context, contentURL s errDownloadStatus = fmt.Errorf("video content download failed: %s", resp.Status) } errMsg := &interfaces.ErrorMessage{StatusCode: resp.StatusCode, Error: errDownloadStatus} - h.WriteErrorResponse(c, errMsg) + h.WriteErrorResponse(c, sanitizeOpenAIErrorMessage(errMsg)) return errDownloadStatus } @@ -993,7 +993,7 @@ func (h *OpenAIAPIHandler) collectXAIVideosNative(c *gin.Context, rawJSON []byte resp, upstreamHeaders, errMsg := h.ExecuteWithAuthManager(cliCtx, xaiVideosHandlerType, executionModel, rawJSON, "") stopKeepAlive() if errMsg != nil { - h.WriteErrorResponse(c, errMsg) + h.WriteErrorResponse(c, sanitizeOpenAIErrorMessage(errMsg)) if errMsg.Error != nil { cliCancel(errMsg.Error) } else { @@ -1028,7 +1028,7 @@ func (h *OpenAIAPIHandler) collectXAIVideosCreate(c *gin.Context, xaiReq []byte, resp, upstreamHeaders, errMsg := h.ExecuteWithAuthManager(cliCtx, xaiVideosHandlerType, routingModel, xaiReq, "") stopKeepAlive() if errMsg != nil { - h.WriteErrorResponse(c, errMsg) + h.WriteErrorResponse(c, sanitizeOpenAIErrorMessage(errMsg)) if errMsg.Error != nil { cliCancel(errMsg.Error) } else { @@ -1040,7 +1040,7 @@ func (h *OpenAIAPIHandler) collectXAIVideosCreate(c *gin.Context, xaiReq []byte, out, err := buildVideosCreateAPIResponseFromXAI(resp, meta) if err != nil { errMsg := &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: err} - h.WriteErrorResponse(c, errMsg) + h.WriteErrorResponse(c, sanitizeOpenAIErrorMessage(errMsg)) cliCancel(err) return } diff --git a/sdk/api/handlers/openai/trusted_direct_response_sink_test.go b/sdk/api/handlers/openai/trusted_direct_response_sink_test.go new file mode 100644 index 000000000..02d1cbdaa --- /dev/null +++ b/sdk/api/handlers/openai/trusted_direct_response_sink_test.go @@ -0,0 +1,311 @@ +package openai + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/gin-gonic/gin" + "github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces" + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers" + coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" +) + +// trustedSinkTestHost is a minimal PluginInterceptorHost that terminates every +// request with a trusted local DirectResponse before upstream execution. This +// mirrors a trusted plugin/interceptor producing a downstream response through +// the real pre-output execution + sanitizer + WriteErrorResponse sink, in +// contrast to the unit-level sanitizer-only preservation test. +type trustedSinkTestHost struct { + status int + header http.Header + body []byte +} + +func (*trustedSinkTestHost) HasStreamInterceptors() bool { return true } + +func (host *trustedSinkTestHost) InterceptRequestBeforeAuth(context.Context, pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse { + return pluginapi.RequestInterceptResponse{ + Terminate: true, + StatusCode: host.status, + ResponseHeaders: host.header, + ResponseBody: host.body, + } +} + +func (*trustedSinkTestHost) InterceptRequestAfterAuth(context.Context, pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse { + return pluginapi.RequestInterceptResponse{} +} + +func (*trustedSinkTestHost) InterceptResponse(context.Context, pluginapi.ResponseInterceptRequest) pluginapi.ResponseInterceptResponse { + return pluginapi.ResponseInterceptResponse{} +} + +func (*trustedSinkTestHost) InterceptStreamChunk(context.Context, pluginapi.StreamChunkInterceptRequest) pluginapi.StreamChunkInterceptResponse { + return pluginapi.StreamChunkInterceptResponse{} +} + +func (*trustedSinkTestHost) CompleteRequest(context.Context, pluginapi.RequestCompletion) {} + +// newTrustedSinkRouter builds an executor-backed base handler with a terminating +// plugin host and wires the representative OpenAI routes so requests flow through +// the real execute -> sanitizer -> WriteErrorResponse / sanitizer -> images/videos +// pre-output sinks. +func newTrustedSinkRouter(t *testing.T, host handlers.PluginInterceptorHost) (*gin.Engine, *int) { + t.Helper() + manager := coreauth.NewManager(nil, nil, nil) + executor := &peekStreamExecutor{} + manager.RegisterExecutor(executor) + + called := 0 + provider := executor.Identifier() + for _, model := range []struct { + id string + ep string + typ string + }{ + {"sink-chat-model", openAIChatEndpoint, ""}, + {"sink-resp-model", openAIResponsesEndpoint, ""}, + {"sink-img-model", "", registry.OpenAIImageModelType}, + {"sink-compat-image-model", "", registry.OpenAIImageModelType}, + {defaultXAIVideosModel, "", ""}, + } { + authID := "sink-auth-" + model.id + auth := &coreauth.Auth{ID: authID, Provider: provider, Status: coreauth.StatusActive} + if _, err := manager.Register(context.Background(), auth); err != nil { + t.Fatalf("register auth %s: %v", authID, err) + } + info := ®istry.ModelInfo{ID: model.id} + if model.ep != "" { + info.SupportedEndpoints = []string{model.ep} + } + if model.typ != "" { + info.Type = model.typ + } + registry.GetGlobalRegistry().RegisterClient(authID, provider, []*registry.ModelInfo{info}) + called++ + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(authID) }) + } + + base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager) + base.SetPluginHost(host) + openAI := NewOpenAIAPIHandler(base) + responses := NewOpenAIResponsesAPIHandler(base) + + gin.SetMode(gin.TestMode) + router := gin.New() + router.POST("/v1/chat/completions", openAI.ChatCompletions) + router.POST("/v1/responses", responses.Responses) + router.POST("/v1/images/generations", openAI.ImagesGenerations) + router.POST("/v1/videos", openAI.VideosCreate) + return router, &called +} + +// executeTrustedSinkRequest dispatches to a route and returns the recorder. +func executeTrustedSinkRequest(t *testing.T, router *gin.Engine, route, body string) *httptest.ResponseRecorder { + t.Helper() + req := httptest.NewRequest(http.MethodPost, route, strings.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + router.ServeHTTP(rec, req) + return rec +} + +// TestOpenAITrustedDirectResponseSurvivesRealNonStreamSinks exercises a trusted +// plugin/interceptor DirectResponse through the actual non-streaming pre-output +// sinks for OpenAI chat, Responses, Images (generations + compat) and Videos. +// It asserts the exact status, verbatim body and safe plugin header survive to +// the recorder/client, and that the generic error envelope is absent. +func TestOpenAITrustedDirectResponseSurvivesRealNonStreamSinks(t *testing.T) { + const secret = "trusted-sink-secret-4488" + host := &trustedSinkTestHost{ + status: http.StatusTooManyRequests, + header: http.Header{"X-Authorization-Request-Id": {"rq-7182"}}, + body: []byte(`{"error":"blocked","detail":"` + secret + `"}`), + } + router, called := newTrustedSinkRouter(t, host) + + tests := []struct { + name string + route string + body string + }{ + {"chat", "/v1/chat/completions", `{"model":"sink-chat-model","messages":[{"role":"user","content":"hi"}]}`}, + {"responses", "/v1/responses", `{"model":"sink-resp-model","input":"hi"}`}, + {"images-generations", "/v1/images/generations", `{"model":"sink-img-model","prompt":"x"}`}, + {"images-compat", "/v1/images/generations", `{"model":"sink-compat-image-model","prompt":"x"}`}, + {"videos", "/v1/videos", `{"model":"` + defaultXAIVideosModel + `","prompt":"hi","seconds":"1"}`}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + rec := executeTrustedSinkRequest(t, router, tc.route, tc.body) + raw := rec.Body.String() + if rec.Code != http.StatusTooManyRequests { + t.Fatalf("status = %d, want %d; body=%q", rec.Code, http.StatusTooManyRequests, raw) + } + if rec.Body.String() != string(host.body) { + t.Fatalf("body = %q, want verbatim %q", rec.Body.String(), string(host.body)) + } + if got := rec.Header().Get("X-Authorization-Request-Id"); got != "rq-7182" { + t.Fatalf("X-Authorization-Request-Id = %q, want rq-7182", got) + } + if !strings.Contains(raw, secret) { + t.Fatalf("trusted plugin body leaked key data: %q", raw) + } + // Generic envelope must be absent: WriteErrorResponse must not wrap + // the trusted response with BuildErrorResponseBody. + if strings.Contains(raw, `"error":{"message"`) { + t.Fatalf("trusted response wrapped in generic envelope: %q", raw) + } + }) + } + if *called != 5 { + t.Fatalf("registered model clients = %d, want 5", *called) + } +} + +// TestOpenAITrustedDirectResponseSurvivesStreamingPeek exercises the same +// trusted plugin termination through the streaming-initial/peek path: the +// buffered errChan error is consumed before SSE headers are committed and the +// trusted body/header reach the recorder verbatim. +func TestOpenAITrustedDirectResponseSurvivesStreamingPeek(t *testing.T) { + const secret = "trusted-sink-stream-secret-5590" + host := &trustedSinkTestHost{ + status: http.StatusForbidden, + header: http.Header{"X-Authorization-Request-Id": {"rq-stream-11"}}, + body: []byte(`{"error":"stream-blocked","s":"` + secret + `"}`), + } + router, _ := newTrustedSinkRouter(t, host) + + rec := executeTrustedSinkRequest(t, router, "/v1/chat/completions", + `{"model":"sink-chat-model","stream":true,"messages":[{"role":"user","content":"hi"}]}`) + raw := rec.Body.String() + if rec.Code != http.StatusForbidden { + t.Fatalf("peek status = %d, want 403; body=%q", rec.Code, raw) + } + if rec.Body.String() != string(host.body) { + t.Fatalf("peek body = %q, want verbatim %q", rec.Body.String(), string(host.body)) + } + if got := rec.Header().Get("X-Authorization-Request-Id"); got != "rq-stream-11" { + t.Fatalf("peek X-Authorization-Request-Id = %q, want rq-stream-11", got) + } + if !strings.Contains(raw, secret) { + t.Fatalf("trusted peek body leaked key data: %q", raw) + } + if strings.Contains(raw, `"error":{"message"`) { + t.Fatalf("trusted peek response wrapped in generic envelope: %q", raw) + } +} + +// TestOpenAIUntrustedDirectResponseStrippedAtRealSink is the paired untrusted +// case: an upstream executor surfaces a RequestTerminatedError with the zero +// Trusted value, so DirectResponse reaches the real sink marked untrusted. The +// sink must strip Body and ResponseHeaders and produce the generic sanctioned +// error envelope with the secret redacted. +func TestOpenAIUntrustedDirectResponseStrippedAtRealSink(t *testing.T) { + const secret = "untrusted-sink-secret-8831" + manager := coreauth.NewManager(nil, nil, nil) + executor := &untrustedTerminationStreamExecutor{secret: secret} + manager.RegisterExecutor(executor) + + authID := "sink-untrusted-auth" + auth := &coreauth.Auth{ID: authID, Provider: executor.Identifier(), Status: coreauth.StatusActive} + if _, err := manager.Register(context.Background(), auth); err != nil { + t.Fatalf("register auth: %v", err) + } + registry.GetGlobalRegistry().RegisterClient(authID, executor.Identifier(), + []*registry.ModelInfo{{ID: "sink-untrusted-model", SupportedEndpoints: []string{openAIChatEndpoint}}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(authID) }) + + base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager) + h := NewOpenAIAPIHandler(base) + gin.SetMode(gin.TestMode) + router := gin.New() + router.POST("/v1/chat/completions", h.ChatCompletions) + + rec := executeTrustedSinkRequest(t, router, "/v1/chat/completions", + `{"model":"sink-untrusted-model","stream":true,"messages":[{"role":"user","content":"hi"}]}`) + raw := rec.Body.String() + if rec.Code != http.StatusBadGateway { + t.Fatalf("status = %d, want 502; body=%q", rec.Code, raw) + } + if strings.Contains(raw, secret) { + t.Fatalf("untrusted sink leaked secret %q: %q", secret, raw) + } + if got := rec.Header().Get("X-Upstream"); got != "" { + t.Fatalf("untrusted sink forwarded X-Upstream header: %q", got) + } + if !strings.Contains(raw, `"error":{"message"`) { + t.Fatalf("untrusted sink did not produce generic envelope: %q", raw) + } + // The DirectResponse flag must not be observable downstream: sink rebuilt a + // strict error rather than the original body (which embedded the secret). + if raw == string(executor.body()) { + t.Fatalf("untrusted DirectResponse body passed through verbatim: %q", raw) + } +} + +// untrustedTerminationStreamExecutor emits a single RequestTerminatedError chunk +// whose zero Trusted value marks the DirectResponse as untrusted at the sink. +type untrustedTerminationStreamExecutor struct { + secret string +} + +func (*untrustedTerminationStreamExecutor) Identifier() string { return "sink-untrusted" } + +func (*untrustedTerminationStreamExecutor) Execute(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) { + return coreexecutor.Response{}, &coreexecutor.RequestTerminatedError{HTTPStatus: http.StatusBadGateway} +} + +func (e *untrustedTerminationStreamExecutor) ExecuteStream(_ context.Context, _ *coreauth.Auth, _ coreexecutor.Request, _ coreexecutor.Options) (*coreexecutor.StreamResult, error) { + chunks := make(chan coreexecutor.StreamChunk, 1) + chunks <- coreexecutor.StreamChunk{Err: &coreexecutor.RequestTerminatedError{ + HTTPStatus: http.StatusBadGateway, + Header: http.Header{"X-Upstream": {"leak"}}, + Body: e.body(), + }} + close(chunks) + return &coreexecutor.StreamResult{Chunks: chunks}, nil +} + +func (*untrustedTerminationStreamExecutor) Refresh(_ context.Context, auth *coreauth.Auth) (*coreauth.Auth, error) { + return auth, nil +} + +func (*untrustedTerminationStreamExecutor) CountTokens(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) { + return coreexecutor.Response{}, nil +} + +func (*untrustedTerminationStreamExecutor) HttpRequest(context.Context, *coreauth.Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +func (e *untrustedTerminationStreamExecutor) body() []byte { + return []byte(`{"raw":"` + e.secret + `"}`) +} + +func TestSanitizedStreamErrorUnwrapReturnsNil(t *testing.T) { + rawSecret := "raw-secret-key-12345" + rawErr := errors.New("internal upstream error with api_key=" + rawSecret) + errMsg := &interfaces.ErrorMessage{ + StatusCode: http.StatusBadGateway, + Error: rawErr, + } + sanitized := sanitizeOpenAIErrorMessage(errMsg) + if sanitized == nil || sanitized.Error == nil { + t.Fatal("expected sanitized error, got nil") + } + if errors.Unwrap(sanitized.Error) != nil { + t.Fatalf("errors.Unwrap(sanitized.Error) = %v, want nil to prevent raw cause leakage", errors.Unwrap(sanitized.Error)) + } + if strings.Contains(sanitized.Error.Error(), rawSecret) { + t.Fatalf("sanitized error leaked secret %q: %q", rawSecret, sanitized.Error.Error()) + } +} diff --git a/sdk/cliproxy/auth/conductor.go b/sdk/cliproxy/auth/conductor.go index 2c08f1f71..67774fbec 100644 --- a/sdk/cliproxy/auth/conductor.go +++ b/sdk/cliproxy/auth/conductor.go @@ -54,8 +54,12 @@ type Result struct { Success bool // RetryAfter carries a provider supplied retry hint (e.g. 429 retryDelay). RetryAfter *time.Duration + // CredentialScope indicates that the failure affects the whole credential across models (e.g. Anthropic 5h/7d unified limits). + CredentialScope bool // Error describes the failure when Success is false. Error *Error + // Options carries execution request options (headers, metadata, etc.) for result tracking. + Options cliproxyexecutor.Options } // Selector chooses an auth candidate for execution. diff --git a/sdk/cliproxy/auth/conductor_availability_test.go b/sdk/cliproxy/auth/conductor_availability_test.go index 7e07cc071..6b368a592 100644 --- a/sdk/cliproxy/auth/conductor_availability_test.go +++ b/sdk/cliproxy/auth/conductor_availability_test.go @@ -2,6 +2,7 @@ package auth import ( "context" + "net/http" "testing" "time" @@ -176,3 +177,144 @@ func TestManager_ResetQuotaClearsRuntimeAndRegistryState(t *testing.T) { t.Fatalf("registry model count after reset = %d, want 1", count) } } + +func TestManager_ResumeEveryModelAfterCredentialRecovery(t *testing.T) { + manager := NewManager(nil, nil, nil) + ctx := context.Background() + authID := "multi-model-auth" + modelA := "model-a" + modelB := "model-b" + modelC := "model-c" + + reg := registry.GetGlobalRegistry() + reg.RegisterClient(authID, "openai", []*registry.ModelInfo{ + {ID: modelA}, + {ID: modelB}, + {ID: modelC}, + }) + t.Cleanup(func() { + reg.UnregisterClient(authID) + }) + + if _, errRegister := manager.Register(ctx, &Auth{ + ID: authID, + Provider: "openai", + Status: StatusActive, + ModelStates: map[string]*ModelState{ + modelA: {Status: StatusActive}, + modelB: {Status: StatusActive}, + modelC: {Status: StatusActive}, + }, + }); errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + + // Verify all 3 models are available before failure + for _, m := range []string{modelA, modelB, modelC} { + if count := reg.GetModelCount(m); count != 1 { + t.Fatalf("registry model count for %s before failure = %d, want 1", m, count) + } + } + + // Fail with invalid_api_key on modelA -> should suspend all models for this auth + manager.MarkResult(ctx, Result{ + AuthID: authID, + Provider: "openai", + Model: modelA, + Success: false, + Error: &Error{HTTPStatus: http.StatusUnauthorized, Code: "invalid_api_key", Message: "API key not valid"}, + }) + + // Verify all 3 models are now suspended + for _, m := range []string{modelA, modelB, modelC} { + if count := reg.GetModelCount(m); count != 0 { + t.Fatalf("registry model count for %s after invalid_api_key = %d, want 0", m, count) + } + } + + // Fast-forward / expire the cooldown on the auth (simulating cooldown expiry or key replacement) + manager.mu.Lock() + auth := manager.auths[authID] + auth.NextRetryAfter = time.Now().Add(-time.Second) + auth.Quota.NextRecoverAt = time.Now().Add(-time.Second) + for _, state := range auth.ModelStates { + state.NextRetryAfter = time.Now().Add(-time.Second) + state.Quota.NextRecoverAt = time.Now().Add(-time.Second) + } + manager.mu.Unlock() + + // Successful request on modelA -> proves credential recovered + manager.MarkResult(ctx, Result{ + AuthID: authID, + Provider: "openai", + Model: modelA, + Success: true, + }) + + // Verify all 3 models are resumed in the registry + for _, m := range []string{modelA, modelB, modelC} { + if count := reg.GetModelCount(m); count != 1 { + t.Fatalf("registry model count for %s after recovery = %d, want 1", m, count) + } + } +} + +func TestManager_ModelSpecificSuspensionSurvivesSiblingSuccess(t *testing.T) { + manager := NewManager(nil, nil, nil) + ctx := context.Background() + authID := "sibling-suspension-auth" + modelA := "model-a" + modelB := "model-b" + + reg := registry.GetGlobalRegistry() + reg.RegisterClient(authID, "openai", []*registry.ModelInfo{ + {ID: modelA}, + {ID: modelB}, + }) + t.Cleanup(func() { + reg.UnregisterClient(authID) + }) + + if _, errRegister := manager.Register(ctx, &Auth{ + ID: authID, + Provider: "openai", + Status: StatusActive, + ModelStates: map[string]*ModelState{ + modelA: {Status: StatusActive}, + modelB: {Status: StatusActive}, + }, + }); errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + + // Fail modelB with model_not_supported -> modelB should be suspended, modelA available + manager.MarkResult(ctx, Result{ + AuthID: authID, + Provider: "openai", + Model: modelB, + Success: false, + Error: &Error{HTTPStatus: http.StatusBadRequest, Code: "model_not_supported", Message: "model not supported"}, + }) + + if count := reg.GetModelCount(modelB); count != 0 { + t.Fatalf("registry model count for modelB after suspension = %d, want 0", count) + } + if count := reg.GetModelCount(modelA); count != 1 { + t.Fatalf("registry model count for modelA = %d, want 1", count) + } + + // Success on modelA -> should NOT resume modelB + manager.MarkResult(ctx, Result{ + AuthID: authID, + Provider: "openai", + Model: modelA, + Success: true, + }) + + if count := reg.GetModelCount(modelA); count != 1 { + t.Fatalf("registry model count for modelA after success = %d, want 1", count) + } + if count := reg.GetModelCount(modelB); count != 0 { + t.Fatalf("registry model count for modelB after modelA success = %d, want 0 (suspension should survive)", count) + } +} diff --git a/sdk/cliproxy/auth/conductor_cooldown.go b/sdk/cliproxy/auth/conductor_cooldown.go index 5c9acdcd6..94ed45ce2 100644 --- a/sdk/cliproxy/auth/conductor_cooldown.go +++ b/sdk/cliproxy/auth/conductor_cooldown.go @@ -781,6 +781,42 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { suspendReason = "invalid_grant" shouldSuspendModel = true } + } else if isInvalidAPIKeyResultError(result.Error) { + if disableCooling { + state.NextRetryAfter = time.Time{} + } else { + next := now.Add(30 * time.Minute) + state.NextRetryAfter = next + state.Quota = QuotaState{ + Exceeded: true, + Reason: "credential_quota", + NextRecoverAt: next, + } + for _, otherState := range auth.ModelStates { + if otherState != nil && otherState != state { + otherState.Unavailable = true + otherState.Status = StatusError + otherState.StatusMessage = "invalid_api_key" + otherState.NextRetryAfter = next + otherState.Quota = QuotaState{ + Exceeded: true, + Reason: "credential_quota", + NextRecoverAt: next, + } + } + } + auth.Unavailable = true + auth.Status = StatusError + auth.StatusMessage = "invalid_api_key" + auth.Quota = QuotaState{ + Exceeded: true, + Reason: "credential_quota", + NextRecoverAt: next, + } + auth.NextRetryAfter = next + suspendReason = "invalid_api_key" + shouldSuspendModel = true + } } else { switch statusCode { case 401: @@ -877,13 +913,41 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { registry.GetGlobalRegistry().SetModelQuotaExceeded(result.AuthID, modelKey) } if shouldResumeModel { - registry.GetGlobalRegistry().ResumeClientModel(result.AuthID, modelKey) + for _, m := range modelsForRegisteredAuth(result.AuthID) { + if registry.GetGlobalRegistry().GetClientModelSuspensionReason(result.AuthID, m) == "invalid_api_key" { + registry.GetGlobalRegistry().ResumeClientModel(result.AuthID, m) + } + } + if modelKey != "" { + registry.GetGlobalRegistry().ResumeClientModel(result.AuthID, modelKey) + } } else if shouldSuspendModel { - registry.GetGlobalRegistry().SuspendClientModel(result.AuthID, modelKey, suspendReason) + if suspendReason == "invalid_api_key" { + for _, m := range modelsForRegisteredAuth(result.AuthID) { + registry.GetGlobalRegistry().SuspendClientModel(result.AuthID, m, suspendReason) + } + if modelKey != "" { + registry.GetGlobalRegistry().SuspendClientModel(result.AuthID, modelKey, suspendReason) + } + } else { + registry.GetGlobalRegistry().SuspendClientModel(result.AuthID, modelKey, suspendReason) + } } m.hook.OnResult(ctx, result) m.publishErrorEvent(result, authSnapshot) + m.updateSessionAffinity(result) +} + +func (m *Manager) updateSessionAffinity(result Result) { + if m == nil || m.selector == nil { + return + } + if affinity, ok := m.selector.(interface { + OnResult(Result) + }); ok && affinity != nil { + affinity.OnResult(result) + } } func (m *Manager) recordExecutionResult(ctx context.Context, result Result, auth *Auth, ephemeral bool) { @@ -1066,6 +1130,10 @@ func updateAggregatedAvailability(auth *Auth, now time.Time) { if auth == nil { return } + if auth.Quota.Exceeded && auth.Quota.Reason == "credential_quota" && auth.Quota.NextRecoverAt.After(now) { + auth.Unavailable = true + return + } if len(auth.ModelStates) == 0 { clearAggregatedAvailability(auth) return @@ -1370,6 +1438,17 @@ func retryAfterFromError(err error) *time.Duration { return &value } +func isCredentialScopedError(err error) bool { + if err == nil { + return false + } + type credentialScopedProvider interface { + IsCredentialScoped() bool + } + var csp credentialScopedProvider + return (errors.As(err, &csp) && csp != nil && csp.IsCredentialScoped()) || isInvalidAPIKeyError(err) +} + func statusCodeFromResult(err *Error) int { if err == nil { return 0 @@ -1439,6 +1518,38 @@ func isInvalidGrantResultError(err *Error) bool { return isInvalidGrantErrorMessage(err.Code) || isInvalidGrantErrorMessage(err.Message) } +// isInvalidAPIKeyErrorMessage matches upstream "invalid API key" rejections +// that arrive as generic client errors instead of 401/403 — Google answers a +// dead Gemini key with 400 INVALID_ARGUMENT and +// "API key not valid. Please pass a valid API key.", so a request-fault +// classification would wrongly stop credential rotation on a dead key. +func isInvalidAPIKeyErrorMessage(message string) bool { + lowered := strings.ToLower(message) + return strings.Contains(lowered, "api key not valid") || strings.Contains(lowered, "api_key_invalid") +} + +func isInvalidAPIKeyError(err error) bool { + if err == nil { + return false + } + status := statusCodeFromError(err) + if status != http.StatusBadRequest && status != http.StatusUnauthorized && status != http.StatusForbidden { + return false + } + return isInvalidAPIKeyErrorMessage(err.Error()) +} + +func isInvalidAPIKeyResultError(err *Error) bool { + if err == nil { + return false + } + status := statusCodeFromResult(err) + if status != http.StatusBadRequest && status != http.StatusUnauthorized && status != http.StatusForbidden { + return false + } + return isInvalidAPIKeyErrorMessage(err.Code) || isInvalidAPIKeyErrorMessage(err.Message) +} + func isModelSupportResultError(err *Error) bool { if err == nil { return false @@ -1707,6 +1818,9 @@ func isRequestInvalidError(err error) bool { if isInvalidGrantError(err) { return false } + if isInvalidAPIKeyError(err) { + return false + } if isModelSupportError(err) { return false } @@ -1757,6 +1871,34 @@ func applyAuthFailureState(auth *Auth, resultErr *Error, retryAfter *time.Durati } return } + if isInvalidAPIKeyResultError(resultErr) { + auth.StatusMessage = "invalid_api_key" + if disableCooling { + auth.NextRetryAfter = time.Time{} + } else { + next := now.Add(30 * time.Minute) + auth.NextRetryAfter = next + auth.Quota = QuotaState{ + Exceeded: true, + Reason: "credential_quota", + NextRecoverAt: next, + } + for _, state := range auth.ModelStates { + if state != nil { + state.Unavailable = true + state.Status = StatusError + state.StatusMessage = "invalid_api_key" + state.NextRetryAfter = next + state.Quota = QuotaState{ + Exceeded: true, + Reason: "credential_quota", + NextRecoverAt: next, + } + } + } + } + return + } switch statusCode { case 401: auth.StatusMessage = "unauthorized" diff --git a/sdk/cliproxy/auth/conductor_cooldown_retry_reset_test.go b/sdk/cliproxy/auth/conductor_cooldown_retry_reset_test.go new file mode 100644 index 000000000..e72a06c9c --- /dev/null +++ b/sdk/cliproxy/auth/conductor_cooldown_retry_reset_test.go @@ -0,0 +1,315 @@ +package auth + +import ( + "context" + "net/http" + "sync" + "sync/atomic" + "testing" + "time" + + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +// retryableRateLimitError carries an explicit Retry-After so +// shouldRetryAfterError decides to wait and re-enter the rotation loop. +type retryableRateLimitError struct { + status int + retryAfter time.Duration +} + +func (e *retryableRateLimitError) Error() string { return "rate limited" } + +func (e *retryableRateLimitError) StatusCode() int { return e.status } + +func (e *retryableRateLimitError) RetryAfter() *time.Duration { return &e.retryAfter } + +// rateLimitedExecutor fails every call with the same rate-limit error and +// counts invocations, so a test can prove the post-cooldown retry actually +// re-executed a recovered credential instead of dying on stale exclusions. +type rateLimitedExecutor struct { + calls atomic.Int32 + err error +} + +func (e *rateLimitedExecutor) Identifier() string { return "gemini" } + +func (e *rateLimitedExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + e.calls.Add(1) + return cliproxyexecutor.Response{}, e.err +} + +func (e *rateLimitedExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + e.calls.Add(1) + return nil, e.err +} + +func (e *rateLimitedExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + e.calls.Add(1) + return cliproxyexecutor.Response{}, e.err +} + +func (e *rateLimitedExecutor) Refresh(_ context.Context, a *Auth) (*Auth, error) { return a, nil } + +func (e *rateLimitedExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +// TestCooldownRetryResetsExclusions is a regression guard for the codex P1 +// finding on CLIProxyAPI PR #4881: when every credential fails 429 with a +// Retry-After shorter than max-retry-interval, the conductor waits for the +// cooldown and retries. The exclusions accumulated during the failed rotation +// pass used to leak into the post-cooldown attempt, so the pick returned +// auth_unavailable without executing anything and the configured +// request-retry never ran. After the fix each entry point must execute the +// credential twice: once in the initial pass and once after the cooldown wait. +func TestCooldownRetryResetsExclusions(t *testing.T) { + newManager := func() (*Manager, *rateLimitedExecutor) { + manager := NewManager(nil, nil, nil) + manager.SetRetryConfig(1, 5*time.Second, 0) + auth := &Auth{ID: "auth-429", Provider: "gemini", Status: StatusActive} + if _, err := manager.Register(context.Background(), auth); err != nil { + t.Fatalf("register auth: %v", err) + } + reg := registry.GetGlobalRegistry() + reg.RegisterClient(auth.ID, "gemini", []*registry.ModelInfo{{ID: "test-model"}}) + t.Cleanup(func() { reg.UnregisterClient(auth.ID) }) + manager.RefreshSchedulerEntry(auth.ID) + exec := &rateLimitedExecutor{err: &retryableRateLimitError{status: http.StatusTooManyRequests, retryAfter: 50 * time.Millisecond}} + manager.RegisterExecutor(exec) + return manager, exec + } + + req := cliproxyexecutor.Request{Model: "test-model"} + opts := cliproxyexecutor.Options{} + + t.Run("Execute", func(t *testing.T) { + manager, exec := newManager() + if _, err := manager.Execute(context.Background(), []string{"gemini"}, req, opts); err == nil { + t.Fatal("expected rate-limit error") + } + if got := exec.calls.Load(); got != 2 { + t.Fatalf("expected 2 executor calls (initial + post-cooldown retry), got %d", got) + } + }) + + t.Run("ExecuteCount", func(t *testing.T) { + manager, exec := newManager() + if _, err := manager.ExecuteCount(context.Background(), []string{"gemini"}, req, opts); err == nil { + t.Fatal("expected rate-limit error") + } + if got := exec.calls.Load(); got != 2 { + t.Fatalf("expected 2 executor calls (initial + post-cooldown retry), got %d", got) + } + }) + + t.Run("ExecuteStream", func(t *testing.T) { + manager, exec := newManager() + if _, err := manager.ExecuteStream(context.Background(), []string{"gemini"}, req, opts); err == nil { + t.Fatal("expected rate-limit error") + } + if got := exec.calls.Load(); got != 2 { + t.Fatalf("expected 2 executor calls (initial + post-cooldown retry), got %d", got) + } + }) +} + +// idRecordingRateLimitedExecutor behaves like rateLimitedExecutor and records +// which auth IDs were executed, so a test can prove a caller-excluded +// credential is never picked after a cooldown retry. +type idRecordingRateLimitedExecutor struct { + mu sync.Mutex + identifier string + calls map[string]int + err error +} + +func (e *idRecordingRateLimitedExecutor) Identifier() string { + if e.identifier != "" { + return e.identifier + } + return "gemini" +} + +func (e *idRecordingRateLimitedExecutor) record(id string) { + e.mu.Lock() + e.calls[id]++ + e.mu.Unlock() +} + +func (e *idRecordingRateLimitedExecutor) Execute(_ context.Context, a *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + e.record(a.ID) + return cliproxyexecutor.Response{}, e.err +} + +func (e *idRecordingRateLimitedExecutor) ExecuteStream(_ context.Context, a *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + e.record(a.ID) + return nil, e.err +} + +func (e *idRecordingRateLimitedExecutor) CountTokens(_ context.Context, a *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + e.record(a.ID) + return cliproxyexecutor.Response{}, e.err +} + +func (e *idRecordingRateLimitedExecutor) Refresh(_ context.Context, a *Auth) (*Auth, error) { + return a, nil +} + +func (e *idRecordingRateLimitedExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +func (e *idRecordingRateLimitedExecutor) count(id string) int { + e.mu.Lock() + defer e.mu.Unlock() + return e.calls[id] +} + +// TestCooldownRetryPreservesCallerExclusions is a regression guard for the +// codex P2 follow-up on PR #4881: resetRecoveredExclusions must prune only +// rotation-added exclusions. Caller-provided exclusions from request metadata +// must survive the cooldown retry, otherwise a credential the caller already +// ruled out can be executed once the wait completes. +func TestCooldownRetryPreservesCallerExclusions(t *testing.T) { + manager := NewManager(nil, nil, nil) + manager.SetRetryConfig(1, 5*time.Second, 0) + authRateLimited := &Auth{ID: "auth-429", Provider: "gemini", Status: StatusActive} + authCallerExcluded := &Auth{ID: "auth-caller", Provider: "gemini", Status: StatusActive} + for _, a := range []*Auth{authRateLimited, authCallerExcluded} { + if _, err := manager.Register(context.Background(), a); err != nil { + t.Fatalf("register auth %s: %v", a.ID, err) + } + } + reg := registry.GetGlobalRegistry() + for _, a := range []*Auth{authRateLimited, authCallerExcluded} { + reg.RegisterClient(a.ID, "gemini", []*registry.ModelInfo{{ID: "test-model"}}) + } + t.Cleanup(func() { + reg.UnregisterClient(authRateLimited.ID) + reg.UnregisterClient(authCallerExcluded.ID) + }) + manager.RefreshSchedulerEntry(authRateLimited.ID) + manager.RefreshSchedulerEntry(authCallerExcluded.ID) + exec := &idRecordingRateLimitedExecutor{ + calls: make(map[string]int), + err: &retryableRateLimitError{status: http.StatusTooManyRequests, retryAfter: 50 * time.Millisecond}, + } + manager.RegisterExecutor(exec) + + opts := cliproxyexecutor.Options{ + Metadata: map[string]any{ + cliproxyexecutor.ExcludedAuthIDsMetadataKey: []string{"auth-caller"}, + }, + } + _, err := manager.Execute(context.Background(), []string{"gemini"}, cliproxyexecutor.Request{Model: "test-model"}, opts) + if err == nil { + t.Fatal("expected rate-limit error") + } + if got := exec.count("auth-caller"); got != 0 { + t.Fatalf("caller-excluded auth executed %d times across cooldown retry", got) + } + if got := exec.count("auth-429"); got != 2 { + t.Fatalf("expected rotation auth to run twice (initial + post-cooldown retry), got %d", got) + } +} + +func TestCooldownRetryPreservesConfigDisabledCoolingExclusions(t *testing.T) { + t.Run("global config disable cooling retains exclusion on retry", func(t *testing.T) { + manager := NewManager(nil, nil, nil) + manager.SetConfigSnapshot(&internalconfig.Config{DisableCooling: true}) + manager.SetRetryConfig(1, 5*time.Second, 0) + auth := &Auth{ID: "auth-global-disabled", Provider: "gemini", Status: StatusActive} + if _, err := manager.Register(context.Background(), auth); err != nil { + t.Fatalf("register auth: %v", err) + } + reg := registry.GetGlobalRegistry() + reg.RegisterClient(auth.ID, "gemini", []*registry.ModelInfo{{ID: "test-model"}}) + t.Cleanup(func() { reg.UnregisterClient(auth.ID) }) + manager.RefreshSchedulerEntry(auth.ID) + exec := &idRecordingRateLimitedExecutor{ + calls: make(map[string]int), + err: &retryableRateLimitError{status: http.StatusTooManyRequests, retryAfter: 50 * time.Millisecond}, + } + manager.RegisterExecutor(exec) + + _, err := manager.Execute(context.Background(), []string{"gemini"}, cliproxyexecutor.Request{Model: "test-model"}, cliproxyexecutor.Options{}) + if err == nil { + t.Fatal("expected rate-limit error") + } + if got := exec.count("auth-global-disabled"); got != 1 { + t.Fatalf("expected config-disabled cooling auth to run once (exclusion retained on retry), got %d", got) + } + }) + + t.Run("provider compat config disable cooling retains exclusion on retry", func(t *testing.T) { + manager := NewManager(nil, nil, nil) + manager.SetConfigSnapshot(&internalconfig.Config{ + OpenAICompatibility: []internalconfig.OpenAICompatibility{ + { + Name: "custom-openai", + DisableCooling: true, + }, + }, + }) + manager.SetRetryConfig(1, 5*time.Second, 0) + auth := &Auth{ + ID: "auth-compat-disabled", + Provider: "openai-compatibility", + Status: StatusActive, + Attributes: map[string]string{ + "provider_key": "custom-openai", + }, + } + if _, err := manager.Register(context.Background(), auth); err != nil { + t.Fatalf("register auth: %v", err) + } + reg := registry.GetGlobalRegistry() + reg.RegisterClient(auth.ID, "openai-compatibility", []*registry.ModelInfo{{ID: "test-model"}}) + t.Cleanup(func() { reg.UnregisterClient(auth.ID) }) + manager.RefreshSchedulerEntry(auth.ID) + exec := &idRecordingRateLimitedExecutor{ + identifier: "openai-compatibility", + calls: make(map[string]int), + err: &retryableRateLimitError{status: http.StatusTooManyRequests, retryAfter: 50 * time.Millisecond}, + } + manager.RegisterExecutor(exec) + + _, err := manager.Execute(context.Background(), []string{"openai-compatibility"}, cliproxyexecutor.Request{Model: "test-model"}, cliproxyexecutor.Options{}) + if err == nil { + t.Fatal("expected rate-limit error") + } + if got := exec.count("auth-compat-disabled"); got != 1 { + t.Fatalf("expected provider-config-disabled cooling auth to run once (exclusion retained on retry), got %d", got) + } + }) + + t.Run("control cooling enabled normally resets exclusion on retry", func(t *testing.T) { + manager := NewManager(nil, nil, nil) + manager.SetConfigSnapshot(&internalconfig.Config{DisableCooling: false}) + manager.SetRetryConfig(1, 5*time.Second, 0) + auth := &Auth{ID: "auth-cooling-enabled", Provider: "gemini", Status: StatusActive} + if _, err := manager.Register(context.Background(), auth); err != nil { + t.Fatalf("register auth: %v", err) + } + reg := registry.GetGlobalRegistry() + reg.RegisterClient(auth.ID, "gemini", []*registry.ModelInfo{{ID: "test-model"}}) + t.Cleanup(func() { reg.UnregisterClient(auth.ID) }) + manager.RefreshSchedulerEntry(auth.ID) + exec := &idRecordingRateLimitedExecutor{ + calls: make(map[string]int), + err: &retryableRateLimitError{status: http.StatusTooManyRequests, retryAfter: 50 * time.Millisecond}, + } + manager.RegisterExecutor(exec) + + _, err := manager.Execute(context.Background(), []string{"gemini"}, cliproxyexecutor.Request{Model: "test-model"}, cliproxyexecutor.Options{}) + if err == nil { + t.Fatal("expected rate-limit error") + } + if got := exec.count("auth-cooling-enabled"); got != 2 { + t.Fatalf("expected cooling-enabled auth to run twice (initial + post-cooldown retry), got %d", got) + } + }) +} diff --git a/sdk/cliproxy/auth/conductor_execution.go b/sdk/cliproxy/auth/conductor_execution.go index e31c65118..da01afb88 100644 --- a/sdk/cliproxy/auth/conductor_execution.go +++ b/sdk/cliproxy/auth/conductor_execution.go @@ -41,23 +41,34 @@ func (m *Manager) Execute(ctx context.Context, providers []string, req cliproxye if len(normalized) == 0 { return cliproxyexecutor.Response{}, &Error{Code: "provider_not_found", Message: "no provider supplied"} } + tracker := newRouteAttemptTracker() if m.HomeEnabled() { - return m.executeHome(ctx, normalized, req, opts, false) + return m.executeHome(ctx, normalized, req, opts, false, tracker) } _, maxRetryCredentials, maxWait := m.retrySettings() var lastErr error retryModel := authSelectionModelFromOptions(opts, req.Model) + tried := extractExcludedAuthIDs(opts.Metadata) + // Snapshot the exclusions supplied by the caller through metadata before + // the rotation loop starts adding its own; cooldown retries prune only + // rotation-added exclusions and must keep the caller-provided set intact. + callerExcluded := make(map[string]struct{}, len(tried)) + for id := range tried { + callerExcluded[id] = struct{}{} + } for attempt := 0; ; attempt++ { - resp, errExec := m.executeMixedOnce(ctx, normalized, req, opts, maxRetryCredentials) + resp, errExec := m.executeMixedOnce(ctx, normalized, req, opts, maxRetryCredentials, tracker, tried) if errExec == nil { return resp, nil } if isRequestTerminatedError(errExec) { return cliproxyexecutor.Response{}, errExec } - lastErr = errExec + if !isAuthNotFoundError(errExec) || lastErr == nil { + lastErr = errExec + } wait, shouldRetry := m.shouldRetryAfterError(errExec, attempt, normalized, retryModel, maxWait) if !shouldRetry { break @@ -65,18 +76,25 @@ func (m *Manager) Execute(ctx context.Context, providers []string, req cliproxye if errWait := waitForCooldown(ctx, wait, maxWait); errWait != nil { return cliproxyexecutor.Response{}, errWait } + // Exclusions collected during this rotation pass refer to credentials + // that were cooling down when they failed; the wait just let them + // recover. Prune recovered exclusions so the retry can actually pick + // a recovered credential instead of failing auth_not_found instantly. + // Auths with disable_cooling stay excluded: they never enter cooldown, + // so the exclusion is the only anti-hammer guard within one request. + tried = m.resetRecoveredExclusions(tried, callerExcluded) } if lastErr != nil { if hasAntigravityProvider(normalized) && shouldAttemptAntigravityCreditsFallback(m, lastErr, normalized) { if resp, ok, errCredits := m.tryAntigravityCreditsExecute(ctx, req, opts); errCredits != nil { - return cliproxyexecutor.Response{}, errCredits + return cliproxyexecutor.Response{}, wrapRouteExhaustion(errCredits, tracker) } else if ok { return resp, nil } } - return cliproxyexecutor.Response{}, lastErr + return cliproxyexecutor.Response{}, wrapRouteExhaustion(lastErr, tracker) } - return cliproxyexecutor.Response{}, &Error{Code: "auth_not_found", Message: "no auth available"} + return cliproxyexecutor.Response{}, wrapRouteExhaustion(&Error{Code: "auth_not_found", Message: "no auth available"}, tracker) } // It supports multiple providers for the same model and round-robins the starting provider per model. @@ -86,23 +104,31 @@ func (m *Manager) ExecuteCount(ctx context.Context, providers []string, req clip if len(normalized) == 0 { return cliproxyexecutor.Response{}, &Error{Code: "provider_not_found", Message: "no provider supplied"} } + tracker := newRouteAttemptTracker() if m.HomeEnabled() { - return m.executeHome(ctx, normalized, req, opts, true) + return m.executeHome(ctx, normalized, req, opts, true, tracker) } _, maxRetryCredentials, maxWait := m.retrySettings() var lastErr error retryModel := authSelectionModelFromOptions(opts, req.Model) + tried := extractExcludedAuthIDs(opts.Metadata) + callerExcluded := make(map[string]struct{}, len(tried)) + for id := range tried { + callerExcluded[id] = struct{}{} + } for attempt := 0; ; attempt++ { - resp, errExec := m.executeCountMixedOnce(ctx, normalized, req, opts, maxRetryCredentials) + resp, errExec := m.executeCountMixedOnce(ctx, normalized, req, opts, maxRetryCredentials, tracker, tried) if errExec == nil { return resp, nil } if isRequestTerminatedError(errExec) { return cliproxyexecutor.Response{}, errExec } - lastErr = errExec + if !isAuthNotFoundError(errExec) || lastErr == nil { + lastErr = errExec + } wait, shouldRetry := m.shouldRetryAfterError(errExec, attempt, normalized, retryModel, maxWait) if !shouldRetry { break @@ -110,11 +136,18 @@ func (m *Manager) ExecuteCount(ctx context.Context, providers []string, req clip if errWait := waitForCooldown(ctx, wait, maxWait); errWait != nil { return cliproxyexecutor.Response{}, errWait } + // Exclusions collected during this rotation pass refer to credentials + // that were cooling down when they failed; the wait just let them + // recover. Prune recovered exclusions so the retry can actually pick + // a recovered credential instead of failing auth_not_found instantly. + // Auths with disable_cooling stay excluded: they never enter cooldown, + // so the exclusion is the only anti-hammer guard within one request. + tried = m.resetRecoveredExclusions(tried, callerExcluded) } if lastErr != nil { - return cliproxyexecutor.Response{}, lastErr + return cliproxyexecutor.Response{}, wrapRouteExhaustion(lastErr, tracker) } - return cliproxyexecutor.Response{}, &Error{Code: "auth_not_found", Message: "no auth available"} + return cliproxyexecutor.Response{}, wrapRouteExhaustion(&Error{Code: "auth_not_found", Message: "no auth available"}, tracker) } // ExecuteStream performs a streaming execution using the configured selector and executor. @@ -131,19 +164,27 @@ func (m *Manager) ExecuteStream(ctx context.Context, providers []string, req cli return nil, &Error{Code: "provider_not_found", Message: "no provider supplied"} } + tracker := newRouteAttemptTracker() _, maxRetryCredentials, maxWait := m.retrySettings() var lastErr error retryModel := authSelectionModelFromOptions(opts, req.Model) + tried := extractExcludedAuthIDs(opts.Metadata) + callerExcluded := make(map[string]struct{}, len(tried)) + for id := range tried { + callerExcluded[id] = struct{}{} + } for attempt := 0; ; attempt++ { - result, errStream := m.executeStreamMixedOnce(ctx, normalized, req, opts, maxRetryCredentials) + result, errStream := m.executeStreamMixedOnce(ctx, normalized, req, opts, maxRetryCredentials, tracker, tried) if errStream == nil { return result, nil } if isRequestTerminatedError(errStream) { return nil, errStream } - lastErr = errStream + if !isAuthNotFoundError(errStream) || lastErr == nil { + lastErr = errStream + } wait, shouldRetry := m.shouldRetryAfterError(errStream, attempt, normalized, retryModel, maxWait) if !shouldRetry { break @@ -151,22 +192,29 @@ func (m *Manager) ExecuteStream(ctx context.Context, providers []string, req cli if errWait := waitForCooldown(ctx, wait, maxWait); errWait != nil { return nil, errWait } + // Exclusions collected during this rotation pass refer to credentials + // that were cooling down when they failed; the wait just let them + // recover. Prune recovered exclusions so the retry can actually pick + // a recovered credential instead of failing auth_not_found instantly. + // Auths with disable_cooling stay excluded: they never enter cooldown, + // so the exclusion is the only anti-hammer guard within one request. + tried = m.resetRecoveredExclusions(tried, callerExcluded) } if lastErr != nil { if hasAntigravityProvider(normalized) && shouldAttemptAntigravityCreditsFallback(m, lastErr, normalized) { if result, ok, errCredits := m.tryAntigravityCreditsExecuteStream(ctx, req, opts); errCredits != nil { - return nil, errCredits + return nil, wrapRouteExhaustion(errCredits, tracker) } else if ok { return result, nil } } var bootstrapErr *streamBootstrapError if errors.As(lastErr, &bootstrapErr) && bootstrapErr != nil { - return streamErrorResult(bootstrapErr.Headers(), bootstrapErr.cause), nil + return streamErrorResult(bootstrapErr.Headers(), wrapRouteExhaustion(bootstrapErr.cause, tracker)), nil } - return nil, lastErr + return nil, wrapRouteExhaustion(lastErr, tracker) } - return nil, &Error{Code: "auth_not_found", Message: "no auth available"} + return nil, wrapRouteExhaustion(&Error{Code: "auth_not_found", Message: "no auth available"}, tracker) } type requestToFormatResolver interface { @@ -203,6 +251,7 @@ func applyRequestAfterAuthInterceptor(ctx context.Context, executor ProviderExec HTTPStatus: resp.StatusCode, Header: cloneRequestHeaders(resp.ResponseHeaders), Body: bytes.Clone(resp.ResponseBody), + Trusted: true, } } return req, opts, nil @@ -272,7 +321,7 @@ func mergeRequestHeaders(current, updates http.Header, clear []string) http.Head return out } -func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, maxRetryCredentials int) (cliproxyexecutor.Response, error) { +func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, maxRetryCredentials int, tracker *routeAttemptTracker, tried ...map[string]struct{}) (cliproxyexecutor.Response, error) { if len(providers) == 0 { return cliproxyexecutor.Response{}, &Error{Code: "provider_not_found", Message: "no provider supplied"} } @@ -281,7 +330,12 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req opts = ensureRequestedModelMetadata(opts, routeModel) homeMode := m.HomeEnabled() homeAuthCount := 1 - tried := make(map[string]struct{}) + var t map[string]struct{} + if len(tried) > 0 && tried[0] != nil { + t = tried[0] + } else { + t = extractExcludedAuthIDs(opts.Metadata) + } attempted := make(map[string]struct{}) var lastErr error for { @@ -295,7 +349,8 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req if homeMode { pickOpts = withHomeAuthCount(opts, homeAuthCount) } - auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, pickOpts, tried) + pickOpts = withExcludedAuthIDs(pickOpts, t) + auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, pickOpts, t) if errPick != nil { if shouldReturnLastErrorOnPickFailure(homeMode, lastErr, errPick) { return cliproxyexecutor.Response{}, lastErr @@ -307,7 +362,7 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req debugLogAuthSelection(entry, auth, provider, routeModel) publishSelectedAuthMetadata(opts.Metadata, auth) - tried[auth.ID] = struct{}{} + t[auth.ID] = struct{}{} execCtx := ctx if rt := m.roundTripperFor(auth); rt != nil { execCtx = context.WithValue(execCtx, roundTripperContextKey{}, rt) @@ -326,7 +381,7 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req if errCancel := claudeOAuthRequestCancellation(execCtx, auth, errPrepare); errCancel != nil { return cliproxyexecutor.Response{}, errCancel } - result := Result{AuthID: auth.ID, Provider: provider, Model: routeModel, Success: false, Error: resultErrorFromError(errPrepare)} + result := Result{AuthID: auth.ID, Provider: provider, Model: routeModel, Success: false, Error: resultErrorFromError(errPrepare), Options: pickOpts} m.MarkResult(execCtx, result) lastErr = errPrepare continue @@ -340,7 +395,7 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req if restoreExecutionModel { execReq.Model = executionModel } - execOpts := opts + execOpts := pickOpts var errIntercept error execReq, execOpts, errIntercept = applyRequestAfterAuthInterceptor(execCtx, executor, provider, execReq, execOpts, requestedModelAliasFromOptions(execOpts, routeModel)) if errIntercept != nil { @@ -368,21 +423,29 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req if errCancel := claudeOAuthRequestCancellation(execCtx, auth, errExec); errCancel != nil { return cliproxyexecutor.Response{}, errCancel } - result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: errExec == nil} + result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: errExec == nil, Options: execOpts} if errExec != nil { result.Error = resultErrorFromError(errExec) if ra := retryAfterFromError(errExec); ra != nil { result.RetryAfter = ra } + if isCredentialScopedError(errExec) { + result.CredentialScope = true + } m.MarkResult(execCtx, result) if isRequestInvalidError(errExec) { return cliproxyexecutor.Response{}, errExec } + tracker.Record(auth, errExec) authErr = errExec + if result.CredentialScope { + break + } continue } if isEmptyCompletionPayload(resp.Payload) { authErr = m.markEmptyCompletion(execCtx, &result) + tracker.Record(auth, authErr) continue } m.MarkResult(execCtx, result) @@ -403,7 +466,7 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req } } -func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, maxRetryCredentials int) (cliproxyexecutor.Response, error) { +func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, maxRetryCredentials int, tracker *routeAttemptTracker, tried ...map[string]struct{}) (cliproxyexecutor.Response, error) { if len(providers) == 0 { return cliproxyexecutor.Response{}, &Error{Code: "provider_not_found", Message: "no provider supplied"} } @@ -412,7 +475,12 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, opts = ensureRequestedModelMetadata(opts, routeModel) homeMode := m.HomeEnabled() homeAuthCount := 1 - tried := make(map[string]struct{}) + var t map[string]struct{} + if len(tried) > 0 && tried[0] != nil { + t = tried[0] + } else { + t = extractExcludedAuthIDs(opts.Metadata) + } attempted := make(map[string]struct{}) var lastErr error for { @@ -426,7 +494,8 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, if homeMode { pickOpts = withHomeAuthCount(opts, homeAuthCount) } - auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, pickOpts, tried) + pickOpts = withExcludedAuthIDs(pickOpts, t) + auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, pickOpts, t) if errPick != nil { if shouldReturnLastErrorOnPickFailure(homeMode, lastErr, errPick) { return cliproxyexecutor.Response{}, lastErr @@ -438,7 +507,7 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, debugLogAuthSelection(entry, auth, provider, routeModel) publishSelectedAuthMetadata(opts.Metadata, auth) - tried[auth.ID] = struct{}{} + t[auth.ID] = struct{}{} execCtx := ctx if rt := m.roundTripperFor(auth); rt != nil { execCtx = context.WithValue(execCtx, roundTripperContextKey{}, rt) @@ -457,7 +526,7 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, if errCancel := claudeOAuthRequestCancellation(execCtx, auth, errPrepare); errCancel != nil { return cliproxyexecutor.Response{}, errCancel } - result := Result{AuthID: auth.ID, Provider: provider, Model: routeModel, Success: false, Error: resultErrorFromError(errPrepare)} + result := Result{AuthID: auth.ID, Provider: provider, Model: routeModel, Success: false, Error: resultErrorFromError(errPrepare), Options: pickOpts} m.MarkResult(execCtx, result) lastErr = errPrepare continue @@ -471,7 +540,7 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, if restoreExecutionModel { execReq.Model = executionModel } - execOpts := opts + execOpts := pickOpts var errIntercept error execReq, execOpts, errIntercept = applyRequestAfterAuthInterceptor(execCtx, executor, provider, execReq, execOpts, requestedModelAliasFromOptions(execOpts, routeModel)) if errIntercept != nil { @@ -499,7 +568,7 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, if errCancel := claudeOAuthRequestCancellation(execCtx, auth, errExec); errCancel != nil { return cliproxyexecutor.Response{}, errCancel } - result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: errExec == nil} + result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: errExec == nil, Options: execOpts} if errExec != nil { result.Error = resultErrorFromError(errExec) if ra := retryAfterFromError(errExec); ra != nil { @@ -512,12 +581,19 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, if isCountTokensEndpointNotFoundError(errExec, execReq.Model) { m.recordAvailabilityNeutralResult(execCtx, result) } else { + if isCredentialScopedError(errExec) { + result.CredentialScope = true + } m.MarkResult(execCtx, result) } if isRequestInvalidError(errExec) { return cliproxyexecutor.Response{}, errExec } + tracker.Record(auth, errExec) authErr = errExec + if result.CredentialScope { + break + } continue } m.MarkResult(execCtx, result) @@ -538,7 +614,7 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, } } -func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, maxRetryCredentials int) (*cliproxyexecutor.StreamResult, error) { +func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, maxRetryCredentials int, tracker *routeAttemptTracker, tried ...map[string]struct{}) (*cliproxyexecutor.StreamResult, error) { if len(providers) == 0 { return nil, &Error{Code: "provider_not_found", Message: "no provider supplied"} } @@ -548,7 +624,12 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string opts = ensureRequestedModelMetadata(opts, routeModel) homeMode := m.HomeEnabled() homeAuthCount := 1 - tried := make(map[string]struct{}) + var t map[string]struct{} + if len(tried) > 0 && tried[0] != nil { + t = tried[0] + } else { + t = extractExcludedAuthIDs(opts.Metadata) + } attempted := make(map[string]struct{}) unauthorizedRefreshTried := make(map[string]struct{}) emptyCompletionTried := make(map[string]struct{}) @@ -564,6 +645,7 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string if homeMode { pickOpts = withHomeAuthCount(opts, homeAuthCount) } + pickOpts = withExcludedAuthIDs(pickOpts, t) var selection *HomeDispatchSelection var auth *Auth @@ -578,7 +660,7 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string provider = selection.Provider } } else { - auth, executor, provider, errPick = m.pickNextMixed(ctx, providers, routeModel, pickOpts, tried) + auth, executor, provider, errPick = m.pickNextMixed(ctx, providers, routeModel, pickOpts, t) } if errPick != nil { if shouldReturnLastErrorOnPickFailure(homeMode, lastErr, errPick) { @@ -619,7 +701,7 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string } publishSelectedAuthMetadata(opts.Metadata, auth) - tried[auth.ID] = struct{}{} + t[auth.ID] = struct{}{} execCtx := ctx releaseAttempt := func() {} if selection != nil { @@ -662,7 +744,7 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string return nil, errCancel } } - result := Result{AuthID: auth.ID, Provider: provider, Model: routeModel, Success: false, Error: resultErrorFromError(errPrepare)} + result := Result{AuthID: auth.ID, Provider: provider, Model: routeModel, Success: false, Error: resultErrorFromError(errPrepare), Options: pickOpts} if selection != nil { m.reportHomeResult(execCtx, result, auth) releaseAttempt() @@ -682,7 +764,7 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string if restoreExecutionModel { streamExecutionModel = executionModel } - execOpts := opts + execOpts := pickOpts if selection != nil { execOpts.ExecutionLifecycle = selection } @@ -704,6 +786,7 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string if isRequestInvalidError(errStream) { return nil, errStream } + tracker.Record(auth, errStream) lastErr = errStream if homeMode { if isEmptyCompletionError(errStream) { @@ -1316,3 +1399,95 @@ func (m *Manager) HttpRequest(ctx context.Context, auth *Auth, req *http.Request } return exec.HttpRequest(ctx, auth, req) } + +// resetRecoveredExclusions prunes the per-request exclusion set after a +// cooldown wait: credentials that entered a real cooldown have now waited it +// out and must be pickable again, otherwise the configured request-retry dies +// on stale exclusions (auth_not_found) without executing anything. Exclusions +// are kept for auths with disable_cooling, and while the global disable-cooling +// flag is on: those never enter cooldown, so the exclusion remains the only +// guard against re-hammering them within one request. Entries listed in +// preserve are kept unconditionally: they came from caller-supplied request +// metadata, not from this rotation, so a cooldown retry must not discard them. +func (m *Manager) resetRecoveredExclusions(tried, preserve map[string]struct{}) map[string]struct{} { + if len(tried) == 0 { + return tried + } + if quotaCooldownDisabled.Load() { + return tried + } + kept := make(map[string]struct{}, len(tried)) + m.mu.RLock() + for id := range tried { + if a, ok := m.auths[id]; ok && a != nil { + if m.cooldownDisabledForAuth(a) { + kept[id] = struct{}{} + } + } + } + m.mu.RUnlock() + for id := range preserve { + kept[id] = struct{}{} + } + return kept +} + +func extractExcludedAuthIDs(meta map[string]any) map[string]struct{} { + excluded := make(map[string]struct{}) + if meta == nil { + return excluded + } + if existing, ok := meta[cliproxyexecutor.ExcludedAuthIDsMetadataKey]; ok { + switch v := existing.(type) { + case map[string]struct{}: + for id := range v { + excluded[id] = struct{}{} + } + case []string: + for _, id := range v { + excluded[id] = struct{}{} + } + } + } + return excluded +} + +func withExcludedAuthIDs(opts cliproxyexecutor.Options, tried map[string]struct{}) cliproxyexecutor.Options { + if len(tried) == 0 { + return opts + } + meta := make(map[string]any, len(opts.Metadata)+1) + for k, v := range opts.Metadata { + meta[k] = v + } + excluded := make(map[string]struct{}, len(tried)) + for id := range tried { + excluded[id] = struct{}{} + } + if existing, ok := meta[cliproxyexecutor.ExcludedAuthIDsMetadataKey]; ok { + switch v := existing.(type) { + case map[string]struct{}: + for id := range v { + excluded[id] = struct{}{} + } + case []string: + for _, id := range v { + excluded[id] = struct{}{} + } + } + } + meta[cliproxyexecutor.ExcludedAuthIDsMetadataKey] = excluded + opts.Metadata = meta + return opts +} + +func isAuthNotFoundError(err error) bool { + if err == nil { + return false + } + var authErr *Error + if errors.As(err, &authErr) && authErr != nil { + return authErr.Code == "auth_not_found" || authErr.Code == "auth_unavailable" + } + return false +} diff --git a/sdk/cliproxy/auth/conductor_fast_error_test.go b/sdk/cliproxy/auth/conductor_fast_error_test.go index 7956bdb39..366e101d8 100644 --- a/sdk/cliproxy/auth/conductor_fast_error_test.go +++ b/sdk/cliproxy/auth/conductor_fast_error_test.go @@ -50,7 +50,7 @@ func TestManagerFastLocalErrorDoesNotRefreshRetryOrCoolCredential(t *testing.T) if calls.Add(1) == 1 { return cliproxyexecutor.Response{}, &requestScopedStatusError{message: "decode Fast response"} } - return cliproxyexecutor.Response{Payload: []byte(`{"type":"message","content":[]}`)}, nil + return cliproxyexecutor.Response{Payload: []byte(`{"type":"message","content":[{"type":"text","text":"ok"}]}`)}, nil } }, run: func(manager *Manager, model string) error { @@ -130,7 +130,7 @@ func TestManagerFastDirectErrorDoesNotRefreshRetryOrCoolCredential(t *testing.T) if calls.Add(1) == 1 { return cliproxyexecutor.Response{}, newFastDirectResponseTestError(http.StatusUnauthorized, `{"type":"error","error":{"message":"Fast denied"}}`) } - return cliproxyexecutor.Response{Payload: []byte(`{"type":"message","content":[]}`)}, nil + return cliproxyexecutor.Response{Payload: []byte(`{"type":"message","content":[{"type":"text","text":"ok"}]}`)}, nil } }, run: func(manager *Manager, model string) error { diff --git a/sdk/cliproxy/auth/conductor_force_mapping_test.go b/sdk/cliproxy/auth/conductor_force_mapping_test.go index ce6cf915f..8512b6ed0 100644 --- a/sdk/cliproxy/auth/conductor_force_mapping_test.go +++ b/sdk/cliproxy/auth/conductor_force_mapping_test.go @@ -201,12 +201,18 @@ func forceMappingStreamUpstreamChunks(provider, upstreamModel string) [][]byte { []byte("event:message_start\n"), []byte("data:" + msg + "\n\n"), []byte("data: " + chat + "\n\n"), + []byte("data: {\"id\":\"chatcmpl-1\",\"object\":\"chat.completion.chunk\",\"model\":\"" + upstreamModel + "\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"ok\"},\"finish_reason\":null}]}\n\n"), + []byte("data: [DONE]\n\n"), } case "xai": msg := strings.Replace(liveXAIMessagesStartUpstream, "grok-4.3", upstreamModel, 1) return [][]byte{ []byte("event: message_start\n"), []byte("data: " + msg + "\n\n"), + []byte("event: content_block_start\n"), + []byte("data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"ok\"}}\n\n"), + []byte("event: message_stop\n"), + []byte("data: {\"type\":\"message_stop\"}\n\n"), } case "antigravity": msg := strings.Replace(liveAntigravityMessagesStartUpstream, "gemini-3-flash", upstreamModel, 1) @@ -217,6 +223,7 @@ func forceMappingStreamUpstreamChunks(provider, upstreamModel string) [][]byte { default: return [][]byte{ []byte(`data: {"type":"response.created","response":{"model":"` + upstreamModel + `"}}` + "\n\n"), + []byte(`data: {"type":"response.completed","response":{"model":"` + upstreamModel + `","output":[{"type":"message","content":[{"type":"output_text","text":"ok"}]}]}}` + "\n\n"), } } } diff --git a/sdk/cliproxy/auth/conductor_home.go b/sdk/cliproxy/auth/conductor_home.go index c98762292..c599ba39b 100644 --- a/sdk/cliproxy/auth/conductor_home.go +++ b/sdk/cliproxy/auth/conductor_home.go @@ -1115,13 +1115,19 @@ func (m *Manager) tryAntigravityCreditsExecute(ctx context.Context, req cliproxy execReq := req execReq.Model = upstreamModel resp, errExec := c.executor.Execute(creditsCtx, c.auth, execReq, creditsOpts) - result := Result{AuthID: c.auth.ID, Provider: c.provider, Model: resultModel, Success: errExec == nil} + result := Result{AuthID: c.auth.ID, Provider: c.provider, Model: resultModel, Success: errExec == nil, Options: creditsOpts} if errExec != nil { result.Error = resultErrorFromError(errExec) if ra := retryAfterFromError(errExec); ra != nil { result.RetryAfter = ra } + if isCredentialScopedError(errExec) { + result.CredentialScope = true + } m.MarkResult(creditsCtx, result) + if result.CredentialScope { + break + } continue } if isEmptyCompletionPayload(resp.Payload) { diff --git a/sdk/cliproxy/auth/conductor_home_execution.go b/sdk/cliproxy/auth/conductor_home_execution.go index c3289f685..d88678668 100644 --- a/sdk/cliproxy/auth/conductor_home_execution.go +++ b/sdk/cliproxy/auth/conductor_home_execution.go @@ -9,7 +9,11 @@ import ( "github.com/tidwall/sjson" ) -func (m *Manager) executeHome(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, countTokens bool) (cliproxyexecutor.Response, error) { +func (m *Manager) executeHome(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, countTokens bool, optionalTracker ...*routeAttemptTracker) (cliproxyexecutor.Response, error) { + var tracker *routeAttemptTracker + if len(optionalTracker) > 0 { + tracker = optionalTracker[0] + } if unlockSession := m.lockHomeWebsocketSession(ctx, opts); unlockSession != nil { defer unlockSession() } @@ -23,21 +27,21 @@ func (m *Manager) executeHome(ctx context.Context, providers []string, req clipr selection, errSelection := m.pickHomeDispatchSelection(ctx, routeModel, withHomeAuthCount(opts, homeAuthCount)) if errSelection != nil { if shouldReturnLastErrorOnPickFailure(true, lastErr, errSelection) { - return cliproxyexecutor.Response{}, lastErr + return cliproxyexecutor.Response{}, wrapRouteExhaustion(lastErr, tracker) } - return cliproxyexecutor.Response{}, errSelection + return cliproxyexecutor.Response{}, wrapRouteExhaustion(errSelection, tracker) } auth := selection.CloneAuthForRoute(routeModel) if auth == nil || selection.Executor == nil { selection.End("missing_execution_target") - return cliproxyexecutor.Response{}, &Error{Code: "executor_not_found", Message: "executor not registered"} + return cliproxyexecutor.Response{}, wrapRouteExhaustion(&Error{Code: "executor_not_found", Message: "executor not registered"}, tracker) } if _, seen := tried[auth.ID]; seen { selection.End("repeated_auth") if lastErr != nil { - return cliproxyexecutor.Response{}, lastErr + return cliproxyexecutor.Response{}, wrapRouteExhaustion(lastErr, tracker) } - return cliproxyexecutor.Response{}, repeatedHomeAuthError() + return cliproxyexecutor.Response{}, wrapRouteExhaustion(repeatedHomeAuthError(), tracker) } entry := logEntryWithRequestID(ctx) debugLogAuthSelection(entry, auth, selection.Provider, routeModel) @@ -72,15 +76,17 @@ func (m *Manager) executeHome(ctx context.Context, providers []string, req clipr return cliproxyexecutor.Response{}, errEnd } lastErr = &Error{Code: "auth_not_found", Message: "no execution models available"} + tracker.Record(auth, lastErr) continue } preparedAuth, errPrepare := m.prepareHomeRequestAuth(execCtx, selection.Executor, selection) if errPrepare != nil { - m.reportHomeResult(execCtx, Result{AuthID: auth.ID, Provider: selection.Provider, Model: routeModel, Success: false, Error: resultErrorFromError(errPrepare)}, auth) + m.reportHomeResult(execCtx, Result{AuthID: auth.ID, Provider: selection.Provider, Model: routeModel, Success: false, Error: resultErrorFromError(errPrepare), Options: opts}, auth) releaseAttempt() if errEnd := m.endHomeSelectionBeforeRedispatch(ctx, selection, "prepare_failed"); errEnd != nil { return cliproxyexecutor.Response{}, errEnd } + tracker.Record(auth, errPrepare) lastErr = errPrepare continue } @@ -165,11 +171,12 @@ func (m *Manager) executeHome(ctx context.Context, providers []string, req clipr } } } - result := Result{AuthID: preparedAuth.ID, Provider: selection.Provider, Model: resultModel, Success: errExecute == nil} - if errExecute == nil && isEmptyCompletionPayload(response.Payload) { + result := Result{AuthID: preparedAuth.ID, Provider: selection.Provider, Model: resultModel, Success: errExecute == nil, Options: execOpts} + if errExecute == nil && !countTokens && isEmptyCompletionPayload(response.Payload) { result.Success = false result.Error = errEmptyCompletion m.reportHomeResult(execCtx, result, preparedAuth) + tracker.Record(preparedAuth, errEmptyCompletion) lastErr = errEmptyCompletion continue } @@ -185,6 +192,9 @@ func (m *Manager) executeHome(ctx context.Context, providers []string, req clipr } result.Error = resultErrorFromError(errExecute) result.RetryAfter = retryAfterFromError(errExecute) + if isCredentialScopedError(errExecute) { + result.CredentialScope = true + } m.reportHomeResult(execCtx, result, preparedAuth) lastErr = errExecute if isRequestInvalidError(errExecute) { @@ -192,6 +202,10 @@ func (m *Manager) executeHome(ctx context.Context, providers []string, req clipr selection.End("request_invalid") return cliproxyexecutor.Response{}, errExecute } + tracker.Record(preparedAuth, errExecute) + if result.CredentialScope { + break + } } releaseAttempt() if errEnd := m.endHomeSelectionBeforeRedispatch(ctx, selection, "execution_failed"); errEnd != nil { diff --git a/sdk/cliproxy/auth/conductor_overrides_test.go b/sdk/cliproxy/auth/conductor_overrides_test.go index de35ebc9a..3a4052b50 100644 --- a/sdk/cliproxy/auth/conductor_overrides_test.go +++ b/sdk/cliproxy/auth/conductor_overrides_test.go @@ -552,6 +552,163 @@ func TestManagerExecute_AntigravityInvalidGrantFallsBackAndSuspendsAuth(t *testi } } +// TestManagerExecute_InvalidAPIKeyFallsBackAndSuspendsAuth covers a dead Gemini API +// key: Google answers 400 INVALID_ARGUMENT ("API key not valid"), which is an +// auth-level failure, not a request fault. The conductor must rotate to the next +// auth and cool down the dead key for 30 minutes instead of surfacing the error. +func TestManagerExecute_InvalidAPIKeyFallsBackAndSuspendsAuth(t *testing.T) { + m := NewManager(nil, nil, nil) + invalidKeyErr := &Error{ + HTTPStatus: http.StatusBadRequest, + Message: `bad response status code 400, body: {"error":{"code":400,"message":"API key not valid. Please pass a valid API key.","status":"INVALID_ARGUMENT"}}`, + } + executor := &authFallbackExecutor{ + id: "gemini", + executeErrors: map[string]error{ + "dead-key-auth": invalidKeyErr, + }, + } + m.RegisterExecutor(executor) + + model := "gemini-2.5-flash" + badAuth := &Auth{ID: "dead-key-auth", Provider: "gemini"} + goodAuth := &Auth{ID: "live-key-auth", Provider: "gemini"} + + reg := registry.GetGlobalRegistry() + reg.RegisterClient(badAuth.ID, "gemini", []*registry.ModelInfo{{ID: model}}) + reg.RegisterClient(goodAuth.ID, "gemini", []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { + reg.UnregisterClient(badAuth.ID) + reg.UnregisterClient(goodAuth.ID) + }) + + if _, errRegister := m.Register(context.Background(), badAuth); errRegister != nil { + t.Fatalf("register bad auth: %v", errRegister) + } + if _, errRegister := m.Register(context.Background(), goodAuth); errRegister != nil { + t.Fatalf("register good auth: %v", errRegister) + } + + request := cliproxyexecutor.Request{Model: model} + for i := 0; i < 2; i++ { + resp, errExecute := m.Execute(context.Background(), []string{"gemini"}, request, cliproxyexecutor.Options{}) + if errExecute != nil { + t.Fatalf("execute %d error = %v, want success", i, errExecute) + } + if string(resp.Payload) != goodAuth.ID { + t.Fatalf("execute %d payload = %q, want %q", i, string(resp.Payload), goodAuth.ID) + } + } + + got := executor.ExecuteCalls() + want := []string{badAuth.ID, goodAuth.ID, goodAuth.ID} + if len(got) != len(want) { + t.Fatalf("execute calls = %v, want %v", got, want) + } + for i := range want { + if got[i] != want[i] { + t.Fatalf("execute call %d auth = %q, want %q", i, got[i], want[i]) + } + } + + updatedBad, ok := m.GetByID(badAuth.ID) + if !ok || updatedBad == nil { + t.Fatalf("expected bad auth to remain registered") + } + state := updatedBad.ModelStates[model] + if state == nil { + t.Fatalf("expected model state for %q", model) + } + if !state.Unavailable { + t.Fatalf("expected bad auth model state to be unavailable") + } + if state.NextRetryAfter.IsZero() { + t.Fatalf("expected bad auth model state cooldown to be set") + } + if cooldown := time.Until(state.NextRetryAfter); cooldown < 29*time.Minute || cooldown > 31*time.Minute { + t.Fatalf("cooldown = %v, want about 30 minutes", cooldown) + } + if state.StatusMessage != invalidKeyErr.Message { + t.Fatalf("status message = %q, want %q", state.StatusMessage, invalidKeyErr.Message) + } +} + +func TestManagerExecute_InvalidAPIKeyQuarantinesAcrossModels(t *testing.T) { + m := NewManager(nil, nil, nil) + invalidKeyErr := &Error{ + HTTPStatus: http.StatusBadRequest, + Message: `bad response status code 400, body: {"error":{"code":400,"message":"API key not valid. Please pass a valid API key.","status":"INVALID_ARGUMENT"}}`, + } + executor := &authFallbackExecutor{ + id: "gemini", + executeErrors: map[string]error{ + "dead-key-auth": invalidKeyErr, + }, + } + m.RegisterExecutor(executor) + + modelA := "gemini-2.5-flash" + modelB := "gemini-2.5-pro" + badAuth := &Auth{ID: "dead-key-auth", Provider: "gemini"} + goodAuth := &Auth{ID: "live-key-auth", Provider: "gemini"} + + reg := registry.GetGlobalRegistry() + reg.RegisterClient(badAuth.ID, "gemini", []*registry.ModelInfo{{ID: modelA}, {ID: modelB}}) + reg.RegisterClient(goodAuth.ID, "gemini", []*registry.ModelInfo{{ID: modelA}, {ID: modelB}}) + t.Cleanup(func() { + reg.UnregisterClient(badAuth.ID) + reg.UnregisterClient(goodAuth.ID) + }) + + if _, errRegister := m.Register(context.Background(), badAuth); errRegister != nil { + t.Fatalf("register bad auth: %v", errRegister) + } + if _, errRegister := m.Register(context.Background(), goodAuth); errRegister != nil { + t.Fatalf("register good auth: %v", errRegister) + } + + respA, errExecute := m.Execute(context.Background(), []string{"gemini"}, cliproxyexecutor.Request{Model: modelA}, cliproxyexecutor.Options{}) + if errExecute != nil { + t.Fatalf("execute model A error = %v, want success", errExecute) + } + if string(respA.Payload) != goodAuth.ID { + t.Fatalf("execute model A payload = %q, want %q", string(respA.Payload), goodAuth.ID) + } + + respB, errExecute := m.Execute(context.Background(), []string{"gemini"}, cliproxyexecutor.Request{Model: modelB}, cliproxyexecutor.Options{}) + if errExecute != nil { + t.Fatalf("execute model B error = %v, want success", errExecute) + } + if string(respB.Payload) != goodAuth.ID { + t.Fatalf("execute model B payload = %q, want %q", string(respB.Payload), goodAuth.ID) + } + + got := executor.ExecuteCalls() + want := []string{badAuth.ID, goodAuth.ID, goodAuth.ID} + if len(got) != len(want) { + t.Fatalf("execute calls = %v, want %v", got, want) + } + for i := range want { + if got[i] != want[i] { + t.Fatalf("execute call %d auth = %q, want %q", i, got[i], want[i]) + } + } + + updatedBad, ok := m.GetByID(badAuth.ID) + if !ok || updatedBad == nil { + t.Fatalf("expected bad auth to remain registered") + } + for _, targetModel := range []string{modelA, modelB, "gemini-3-flash"} { + blocked, reason, next := isAuthBlockedForModel(updatedBad, targetModel, time.Now()) + if !blocked { + t.Fatalf("model %q was unblocked despite invalid API key failure on credential", targetModel) + } + if reason != blockReasonCooldown || next.IsZero() { + t.Fatalf("model %q block reason=%v next=%v, want cooldown ~30m", targetModel, reason, next) + } + } +} + func TestManagerExecuteStream_AntigravityInvalidGrantFallsBackAndSuspendsAuth(t *testing.T) { m := NewManager(nil, nil, nil) invalidGrantErr := &Error{ @@ -1172,8 +1329,8 @@ func TestManager_Execute_DisableCooling_RetriesAfter429RetryAfter(t *testing.T) } calls := executor.ExecuteCalls() - if len(calls) != 4 { - t.Fatalf("execute calls = %d, want 4 (initial + 3 retries)", len(calls)) + if len(calls) != 1 { + t.Fatalf("execute calls = %d, want 1", len(calls)) } } @@ -1899,3 +2056,323 @@ func TestManager_RequestScopedNotFoundStopsRetryWithoutSuspendingAuth(t *testing t.Fatalf("expected request-scoped 404 to avoid bad auth model cooldown state, got %#v", state) } } + +// TestManager_ClassifierMixedLoop_RotatesCredentialOnAuthAndQuota is the CPAPlus +// twin of CPA's TestManager_DeepSeekCredentialFailuresRotateCredential. It pins +// the shared mixed-auth loop contract: a first credential failing with an +// authentication or quota marker must be retired (unavailable + cooldown) and +// the same request must fall through to a second credential, calling each +// exactly once and never looping. +// +// The classifier (clienterror.IsRequestFault) must stay in credential scope even +// when the provider pairs the status with a generic invalid_request_error code +// or type in the body. CPAPlus deliberately recognizes a broader set of +// credential markers than CPA: invalid/incorrect/expired API-key codes and the +// Gemini UNAUTHENTICATED status in addition to the authentication_error type. +// Production parity means the same precedence, not byte-identical body matching. +func TestManager_ClassifierMixedLoop_RotatesCredentialOnAuthAndQuota(t *testing.T) { + tests := []struct { + name string + status int + message string + wantQuota bool + }{ + { + name: "401 deepseek authentication marker with generic code", + status: http.StatusUnauthorized, + message: `{"error":{"code":"invalid_request_error","message":"Authentication Fails, Your api key: ****heck is invalid","param":null,"type":"authentication_error"}}`, + }, + { + name: "403 codex invalid or expired key with generic code", + status: http.StatusForbidden, + message: `{"error":{"message":"invalid or expired token","type":"authentication_error","code":"invalid_request_error"}}`, + }, + { + name: "401 gemini unauthenticated status", + status: http.StatusUnauthorized, + message: `{"error":{"code":16,"message":"Request had invalid authentication credentials.","status":"UNAUTHENTICATED"}}`, + }, + { + name: "429 rate limit with generic code", + status: http.StatusTooManyRequests, + message: `{"error":{"code":"invalid_request_error","message":"Rate Limit Reached","param":null,"type":"unknown_error"}}`, + wantQuota: true, + }, + { + name: "402 insufficient balance with generic code", + status: http.StatusPaymentRequired, + message: `{"error":{"message":"Insufficient Balance","type":"unknown_error","param":null,"code":"invalid_request_error"}}`, + }, + } + + for _, tc := range tests { + tc := tc + t.Run(tc.name, func(t *testing.T) { + const provider = "openai-compatibility" + const model = "deepseek-v4-pro" + + executor := &authFallbackExecutor{ + id: provider, + executeErrors: map[string]error{ + "aa-failed-key": &Error{HTTPStatus: tc.status, Message: tc.message}, + }, + } + m := NewManager(nil, nil, nil) + m.SetRetryConfig(2, 30*time.Second, 0) + m.RegisterExecutor(executor) + + failedAuth := &Auth{ID: "aa-failed-key", Provider: provider} + availableAuth := &Auth{ID: "bb-valid-key", Provider: provider} + + reg := registry.GetGlobalRegistry() + models := []*registry.ModelInfo{{ID: model}} + reg.RegisterClient(failedAuth.ID, provider, models) + reg.RegisterClient(availableAuth.ID, provider, models) + t.Cleanup(func() { + reg.UnregisterClient(failedAuth.ID) + reg.UnregisterClient(availableAuth.ID) + }) + + if _, errRegister := m.Register(context.Background(), failedAuth); errRegister != nil { + t.Fatalf("register failed auth: %v", errRegister) + } + if _, errRegister := m.Register(context.Background(), availableAuth); errRegister != nil { + t.Fatalf("register available auth: %v", errRegister) + } + + resp, errExecute := m.Execute( + context.Background(), + []string{provider}, + cliproxyexecutor.Request{Model: model}, + cliproxyexecutor.Options{}, + ) + if errExecute != nil { + t.Fatalf("expected fallback to the next credential, got error: %v", errExecute) + } + if got := string(resp.Payload); got != availableAuth.ID { + t.Fatalf("served by %q, want %q", got, availableAuth.ID) + } + wantCalls := []string{failedAuth.ID, availableAuth.ID} + calls := executor.ExecuteCalls() + if len(calls) != len(wantCalls) { + t.Fatalf("credential calls = %v, want %v (no loop)", calls, wantCalls) + } + for i := range wantCalls { + if calls[i] != wantCalls[i] { + t.Fatalf("credential call %d = %q, want %q", i, calls[i], wantCalls[i]) + } + } + + updatedFailed, ok := m.GetByID(failedAuth.ID) + if !ok || updatedFailed == nil { + t.Fatal("expected failed auth to remain registered") + } + state := updatedFailed.ModelStates[model] + if state == nil || !state.Unavailable || state.NextRetryAfter.IsZero() { + t.Fatalf("failed auth model state = %#v, want active cooldown/unavailable", state) + } + if tc.wantQuota && (!state.Quota.Exceeded || state.Quota.Reason != "quota") { + t.Fatalf("failed auth quota state = %#v, want exceeded quota", state.Quota) + } + }) + } +} + +func TestIsCredentialScopedError_InvalidAPIKey(t *testing.T) { + invalidKey400 := &Error{ + HTTPStatus: http.StatusBadRequest, + Message: `bad response status code 400, body: {"error":{"code":400,"message":"API key not valid. Please pass a valid API key.","status":"INVALID_ARGUMENT"}}`, + } + if !isCredentialScopedError(invalidKey400) { + t.Fatalf("expected isCredentialScopedError(invalidKey400) = true, got false") + } + + invalidKey401 := &Error{ + HTTPStatus: http.StatusUnauthorized, + Message: `{"error":{"code":"api_key_invalid","message":"Invalid API key"}}`, + } + if !isCredentialScopedError(invalidKey401) { + t.Fatalf("expected isCredentialScopedError(invalidKey401) = true, got false") + } + + normal400 := &Error{ + HTTPStatus: http.StatusBadRequest, + Message: `invalid argument: field "prompt" cannot be empty`, + } + if isCredentialScopedError(normal400) { + t.Fatalf("expected isCredentialScopedError(normal400) = false, got true") + } +} + +func TestManagerExecute_InvalidAPIKeyStopsModelPoolRetries(t *testing.T) { + m := NewManager(nil, nil, nil) + cfg := &internalconfig.Config{ + OpenAICompatibility: []internalconfig.OpenAICompatibility{ + { + Name: "gemini", + Models: []internalconfig.OpenAICompatibilityModel{ + { + Name: "gemini-pool", + Alias: "gemini-2.5-flash,gemini-2.5-pro,gemini-1.5-flash", + }, + }, + }, + }, + } + m.SetConfig(cfg) + + invalidKeyErr := &Error{ + HTTPStatus: http.StatusBadRequest, + Message: `bad response status code 400, body: {"error":{"code":400,"message":"API key not valid. Please pass a valid API key.","status":"INVALID_ARGUMENT"}}`, + } + executor := &authFallbackExecutor{ + id: "openai-compatibility", + executeErrors: map[string]error{ + "dead-key-auth": invalidKeyErr, + }, + } + m.RegisterExecutor(executor) + + badAuth := &Auth{ + ID: "dead-key-auth", + Provider: "openai-compatibility", + Attributes: map[string]string{ + "api_key": "key1", + "provider_key": "gemini", + }, + } + goodAuth := &Auth{ + ID: "live-key-auth", + Provider: "openai-compatibility", + Attributes: map[string]string{ + "api_key": "key2", + "provider_key": "gemini", + }, + } + + reg := registry.GetGlobalRegistry() + models := []*registry.ModelInfo{ + {ID: "gemini-pool"}, + {ID: "gemini-2.5-flash"}, + {ID: "gemini-2.5-pro"}, + {ID: "gemini-1.5-flash"}, + } + reg.RegisterClient(badAuth.ID, "openai-compatibility", models) + reg.RegisterClient(goodAuth.ID, "openai-compatibility", models) + t.Cleanup(func() { + reg.UnregisterClient(badAuth.ID) + reg.UnregisterClient(goodAuth.ID) + }) + + if _, errRegister := m.Register(context.Background(), badAuth); errRegister != nil { + t.Fatalf("register bad auth: %v", errRegister) + } + if _, errRegister := m.Register(context.Background(), goodAuth); errRegister != nil { + t.Fatalf("register good auth: %v", errRegister) + } + + resp, errExecute := m.Execute(context.Background(), []string{"openai-compatibility"}, cliproxyexecutor.Request{Model: "gemini-pool"}, cliproxyexecutor.Options{}) + if errExecute != nil { + t.Fatalf("execute error = %v, want success", errExecute) + } + if string(resp.Payload) != goodAuth.ID { + t.Fatalf("execute payload = %q, want %q", string(resp.Payload), goodAuth.ID) + } + + calls := executor.ExecuteCalls() + wantCalls := []string{badAuth.ID, goodAuth.ID} + if len(calls) != len(wantCalls) { + t.Fatalf("execute calls = %v, want %v", calls, wantCalls) + } +} + +func TestManagerExecuteStream_InvalidAPIKeyStopsModelPoolRetries(t *testing.T) { + m := NewManager(nil, nil, nil) + cfg := &internalconfig.Config{ + OpenAICompatibility: []internalconfig.OpenAICompatibility{ + { + Name: "gemini", + Models: []internalconfig.OpenAICompatibilityModel{ + { + Name: "gemini-pool", + Alias: "gemini-2.5-flash,gemini-2.5-pro,gemini-1.5-flash", + }, + }, + }, + }, + } + m.SetConfig(cfg) + + invalidKeyErr := &Error{ + HTTPStatus: http.StatusBadRequest, + Message: `bad response status code 400, body: {"error":{"code":400,"message":"API key not valid. Please pass a valid API key.","status":"INVALID_ARGUMENT"}}`, + } + executor := &authFallbackExecutor{ + id: "openai-compatibility", + streamFirstErrors: map[string]error{ + "dead-key-auth": invalidKeyErr, + }, + } + m.RegisterExecutor(executor) + + badAuth := &Auth{ + ID: "dead-key-auth", + Provider: "openai-compatibility", + Attributes: map[string]string{ + "api_key": "key1", + "provider_key": "gemini", + }, + } + goodAuth := &Auth{ + ID: "live-key-auth", + Provider: "openai-compatibility", + Attributes: map[string]string{ + "api_key": "key2", + "provider_key": "gemini", + }, + } + + reg := registry.GetGlobalRegistry() + models := []*registry.ModelInfo{ + {ID: "gemini-pool"}, + {ID: "gemini-2.5-flash"}, + {ID: "gemini-2.5-pro"}, + {ID: "gemini-1.5-flash"}, + } + reg.RegisterClient(badAuth.ID, "openai-compatibility", models) + reg.RegisterClient(goodAuth.ID, "openai-compatibility", models) + t.Cleanup(func() { + reg.UnregisterClient(badAuth.ID) + reg.UnregisterClient(goodAuth.ID) + }) + + if _, errRegister := m.Register(context.Background(), badAuth); errRegister != nil { + t.Fatalf("register bad auth: %v", errRegister) + } + if _, errRegister := m.Register(context.Background(), goodAuth); errRegister != nil { + t.Fatalf("register good auth: %v", errRegister) + } + + streamResult, errStream := m.ExecuteStream(context.Background(), []string{"openai-compatibility"}, cliproxyexecutor.Request{Model: "gemini-pool"}, cliproxyexecutor.Options{}) + if errStream != nil { + t.Fatalf("execute stream error = %v, want success", errStream) + } + if streamResult == nil { + t.Fatalf("execute stream result is nil") + } + var payloads []string + for chunk := range streamResult.Chunks { + if len(chunk.Payload) > 0 { + payloads = append(payloads, string(chunk.Payload)) + } + } + if len(payloads) == 0 || payloads[0] != goodAuth.ID { + t.Fatalf("stream payloads = %v, want [%s]", payloads, goodAuth.ID) + } + + calls := executor.StreamCalls() + wantCalls := []string{badAuth.ID, goodAuth.ID} + if len(calls) != len(wantCalls) { + t.Fatalf("stream calls = %v, want %v", calls, wantCalls) + } +} diff --git a/sdk/cliproxy/auth/conductor_selection.go b/sdk/cliproxy/auth/conductor_selection.go index bcf597bfb..d30a262d5 100644 --- a/sdk/cliproxy/auth/conductor_selection.go +++ b/sdk/cliproxy/auth/conductor_selection.go @@ -994,6 +994,9 @@ func (m *Manager) routeAwareSelectionRequired(auth *Auth, routeModel string) boo } func (m *Manager) pickNextLegacy(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, tried map[string]struct{}) (*Auth, ProviderExecutor, error) { + if opts.Metadata == nil { + opts.Metadata = make(map[string]any) + } if m.HomeEnabled() { auth, exec, _, err := m.pickNextViaHome(ctx, model, opts, tried) return auth, exec, err @@ -1298,6 +1301,9 @@ func (m *Manager) pickNext(ctx context.Context, provider, model string, opts cli } func (m *Manager) pickNextMixedLegacy(ctx context.Context, providers []string, model string, opts cliproxyexecutor.Options, tried map[string]struct{}) (*Auth, ProviderExecutor, string, error) { + if opts.Metadata == nil { + opts.Metadata = make(map[string]any) + } if m.HomeEnabled() { return m.pickNextViaHome(ctx, model, opts, tried) } diff --git a/sdk/cliproxy/auth/conductor_stream.go b/sdk/cliproxy/auth/conductor_stream.go index 7ff06b914..e7d135b2d 100644 --- a/sdk/cliproxy/auth/conductor_stream.go +++ b/sdk/cliproxy/auth/conductor_stream.go @@ -2,12 +2,191 @@ package auth import ( "context" + "fmt" "net/http" "strings" + "sync" + "time" + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" ) +// ttftScope owns exactly one TTFT attempt with a single-winner decision +// between stream establishment and the timeout fire. A fresh scope is +// created for every attempt, including the in-function retry after a +// credential refresh, so each retry gets a full TTFT budget. Once the stream +// is established, the timer is stopped and can never cancel the connected +// stream afterward; if the timer fired first, callers observe a typed TTFT +// timeout error and failover as before. +type ttftScope struct { + mu sync.Mutex + ctx context.Context + cancel context.CancelFunc + fired bool + committed bool + timer *time.Timer + timeout time.Duration +} + +func newTTFTScope(parent context.Context, timeout time.Duration) *ttftScope { + var attemptCtx context.Context + cancel := func() {} + if parent != nil { + attemptCtx, cancel = context.WithCancel(parent) + } + s := &ttftScope{ + ctx: attemptCtx, + cancel: cancel, + timeout: timeout, + } + if timeout > 0 { + s.timer = time.AfterFunc(timeout, s.fire) + } + return s +} + +// ctxOr returns the scope's fresh child context, falling back to ctx when the +// scope has none (nil parent). +func (s *ttftScope) ctxOr(ctx context.Context) context.Context { + if s != nil && s.ctx != nil { + return s.ctx + } + return ctx +} + +// fire is the timer callback. It is a single winner: only the timer can set +// fired, and only while the scope has not already been committed by a first +// chunk. +func (s *ttftScope) fire() { + s.mu.Lock() + defer s.mu.Unlock() + if s.committed || s.fired { + return + } + s.fired = true + if s.cancel != nil { + s.cancel() + } +} + +// stop halts the TTFT timer once the upstream executor stream is established, +// ensuring the deadline only guards time-to-first-connect and never cancels a +// connected stream during subsequent chunk reads. +func (s *ttftScope) stop() { + if s == nil { + return + } + s.mu.Lock() + defer s.mu.Unlock() + s.committed = true + if s.timer != nil { + s.timer.Stop() + } +} + +// commit marks the first chunk as the winner and stops the timer so a later +// callback can never cancel a stream that already produced its first chunk. +// It returns a release func that the stream producer invokes after handoff to +// free the child context. commit is idempotent: the release func of a repeated +// call is a no-op. +func (s *ttftScope) commit() func() { + if s == nil { + return func() {} + } + s.mu.Lock() + defer s.mu.Unlock() + if s.committed { + return func() {} + } + s.committed = true + if s.timer != nil { + s.timer.Stop() + } + cancel := s.cancel + return func() { + if cancel != nil { + cancel() + } + } +} + +// timedOut reports whether the timeout fired before the first chunk. +func (s *ttftScope) timedOut() bool { + if s == nil { + return false + } + s.mu.Lock() + defer s.mu.Unlock() + return s.fired +} + +// timeoutError returns a typed TTFT timeout error if the timeout won. +func (s *ttftScope) timeoutError() error { + if s == nil || !s.timedOut() { + return nil + } + return newTTFTTimeoutError(s.timeout) +} + +// release cancels the child context and stops the timer. It races safely with +// a timer callback still in flight and is idempotent. +func (s *ttftScope) release() { + if s == nil { + return + } + s.mu.Lock() + defer s.mu.Unlock() + s.committed = true + if s.timer != nil { + s.timer.Stop() + } + if s.cancel != nil { + s.cancel() + s.cancel = nil + } +} + +func newTTFTTimeoutError(timeout time.Duration) error { + return &Error{ + Code: "stream_first_chunk_timeout", + Message: fmt.Sprintf("time to first chunk timeout after %v", timeout), + HTTPStatus: 504, + Retryable: true, + } +} + +func (m *Manager) streamFirstChunkTimeout(opts cliproxyexecutor.Options) time.Duration { + if opts.Metadata != nil { + if ms, ok := opts.Metadata["stream_connect_timeout_ms"].(int); ok { + if ms <= 0 { + return 0 + } + return time.Duration(ms) * time.Millisecond + } + if ms, ok := opts.Metadata["stream_first_chunk_timeout_ms"].(int); ok { + if ms <= 0 { + return 0 + } + return time.Duration(ms) * time.Millisecond + } + } + if m == nil { + return 0 + } + cfg, _ := m.runtimeConfig.Load().(*internalconfig.Config) + if cfg == nil { + return 0 + } + if cfg.Streaming.StreamConnectTimeoutSeconds > 0 { + return time.Duration(cfg.Streaming.StreamConnectTimeoutSeconds) * time.Second + } + if cfg.Streaming.StreamFirstChunkTimeoutSeconds > 0 { + return time.Duration(cfg.Streaming.StreamFirstChunkTimeoutSeconds) * time.Second + } + return 0 +} + func discardStreamChunks(ch <-chan cliproxyexecutor.StreamChunk) { if ch == nil { return @@ -81,7 +260,7 @@ func validateStreamResult(result *cliproxyexecutor.StreamResult, err error) (*cl return result, nil } -func readStreamBootstrap(ctx context.Context, ch <-chan cliproxyexecutor.StreamChunk) ([]cliproxyexecutor.StreamChunk, bool, error) { +func readStreamBootstrap(ctx context.Context, ch <-chan cliproxyexecutor.StreamChunk, onFirstChunk ...func()) ([]cliproxyexecutor.StreamChunk, bool, error) { if ch == nil { return nil, true, nil } @@ -105,19 +284,36 @@ func readStreamBootstrap(ctx context.Context, ch <-chan cliproxyexecutor.StreamC return buffered, true, nil } if chunk.Err != nil { + if bootstrap.hasMeaningfulOutput() { + buffered = append(buffered, chunk) + return buffered, false, nil + } return nil, false, chunk.Err } + for _, cb := range onFirstChunk { + if cb != nil { + cb() + } + } buffered = append(buffered, chunk) if bootstrap.observe(chunk.Payload) { return buffered, false, nil } + if bootstrap.isTerminalEmpty() { + return buffered, true, nil + } } } -func (m *Manager) wrapStreamResult(ctx context.Context, auth *Auth, provider, resultModel string, headers http.Header, buffered []cliproxyexecutor.StreamChunk, remaining <-chan cliproxyexecutor.StreamChunk, aliasResult OAuthModelAliasResult, ephemeralResult bool) *cliproxyexecutor.StreamResult { +func (m *Manager) wrapStreamResult(ctx context.Context, auth *Auth, provider, resultModel string, opts cliproxyexecutor.Options, headers http.Header, buffered []cliproxyexecutor.StreamChunk, remaining <-chan cliproxyexecutor.StreamChunk, aliasResult OAuthModelAliasResult, ephemeralResult bool, cleanups ...func()) *cliproxyexecutor.StreamResult { out := make(chan cliproxyexecutor.StreamChunk) go func() { defer close(out) + for _, cleanup := range cleanups { + if cleanup != nil { + defer cleanup() + } + } var failed bool forward := true var rewriter *StreamRewriter @@ -128,7 +324,7 @@ func (m *Manager) wrapStreamResult(ctx context.Context, auth *Auth, provider, re if chunk.Err != nil && !failed { failed = true rerr := resultErrorFromError(chunk.Err) - m.recordExecutionResult(ctx, Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: rerr}, auth, ephemeralResult) + m.recordExecutionResult(ctx, Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: rerr, Options: opts}, auth, ephemeralResult) } if !forward { return false @@ -185,7 +381,7 @@ func (m *Manager) wrapStreamResult(ctx context.Context, auth *Auth, provider, re } } if !failed && (ephemeralResult || claudeOAuthRequestCancellation(ctx, auth, nil) == nil) { - m.recordExecutionResult(ctx, Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: true}, auth, ephemeralResult) + m.recordExecutionResult(ctx, Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: true, Options: opts}, auth, ephemeralResult) } }() return &cliproxyexecutor.StreamResult{Headers: headers, Chunks: out} @@ -210,6 +406,8 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi _, didRefreshOnUnauthorized = unauthorizedRefreshTried[auth.ID] } for idx, execModel := range execModels { + ttftTimeout := m.streamFirstChunkTimeout(opts) + resultModel := m.stateModelForExecution(auth, routeModel, execModel, pooled) execReq := req execReq.Model = execModel @@ -228,12 +426,27 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi if errCtx := ctx.Err(); errCtx != nil { return nil, errCtx } - streamResult, errStream := executor.ExecuteStream(ctx, auth, execReq, execOpts) + // Arm the TTFT scope only after local interception and request + // preparation: the budget measures upstream responsiveness, so a slow + // after-auth interceptor must not cancel the attempt before any + // upstream request was even made. + scope := newTTFTScope(ctx, ttftTimeout) + attemptCtx := scope.ctx + checkTTFTErr := func(err error) error { + if t := scope.timeoutError(); t != nil { + return t + } + return err + } + streamResult, errStream := executor.ExecuteStream(attemptCtx, auth, execReq, execOpts) if errStream != nil { if errCtx := ctx.Err(); errCtx != nil { + scope.release() return nil, errCtx } + errStream = checkTTFTErr(errStream) if allowRetry { + scope.stop() alreadyTried := didRefreshOnUnauthorized willAttemptHomeRefresh := ephemeralResult && !alreadyTried && auth != nil && auth.AuthKind() == AuthKindOAuth && isUnauthorizedError(errStream) refreshed, okRefresh, errRefresh := m.tryRefreshExecutionAuthAfterUnauthorized(ctx, executor, auth, errStream, alreadyTried, ephemeralResult) @@ -246,13 +459,25 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi if errRefresh != nil { errStream = errRefresh } else if okRefresh { + if streamResult != nil { + discardStreamChunks(streamResult.Chunks) + } auth = refreshed m.replaceHomeExecutionLifecycleAuth(execOpts.ExecutionLifecycle, auth) publishSelectedAuthMetadata(execOpts.Metadata, auth) didRefreshOnUnauthorized = true - streamResult, errStream = executor.ExecuteStream(ctx, auth, execReq, execOpts) + // Fresh TTFT budget and attempt context for the retry. + scope.release() + scope = newTTFTScope(ctx, ttftTimeout) + attemptCtx = scope.ctx + streamResult, errStream = executor.ExecuteStream(attemptCtx, auth, execReq, execOpts) + errStream = checkTTFTErr(errStream) if errStream != nil { if errCtx := ctx.Err(); errCtx != nil { + scope.release() + if streamResult != nil { + discardStreamChunks(streamResult.Chunks) + } return nil, errCtx } } @@ -261,29 +486,51 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi } if !ephemeralResult { if errCancel := claudeOAuthRequestCancellation(ctx, auth, errStream); errCancel != nil { + scope.release() + if streamResult != nil { + discardStreamChunks(streamResult.Chunks) + } return nil, errCancel } } streamResult, errStream = validateStreamResult(streamResult, errStream) if errStream != nil { + scope.release() + if streamResult != nil { + discardStreamChunks(streamResult.Chunks) + } + errStream = checkTTFTErr(errStream) rerr := resultErrorFromError(errStream) - result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: rerr} + result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: rerr, Options: execOpts} result.RetryAfter = retryAfterFromError(errStream) + if isCredentialScopedError(errStream) { + result.CredentialScope = true + } m.recordExecutionResult(ctx, result, auth, ephemeralResult) if isRequestInvalidError(errStream) { return nil, errStream } lastErr = errStream + if result.CredentialScope { + return nil, errStream + } continue } + scope.stop() - buffered, closed, bootstrapErr := readStreamBootstrap(ctx, streamResult.Chunks) + buffered, closed, bootstrapErr := readStreamBootstrap(attemptCtx, streamResult.Chunks) + if bootstrapErr == nil && scope.timedOut() { + bootstrapErr = scope.timeoutError() + } if bootstrapErr != nil { if errCtx := ctx.Err(); errCtx != nil { + scope.release() discardStreamChunks(streamResult.Chunks) return nil, errCtx } + bootstrapErr = checkTTFTErr(bootstrapErr) if allowRetry { + scope.stop() alreadyTried := didRefreshOnUnauthorized willAttemptHomeRefresh := ephemeralResult && !alreadyTried && auth != nil && auth.AuthKind() == AuthKindOAuth && isUnauthorizedError(bootstrapErr) refreshed, okRefresh, errRefresh := m.tryRefreshExecutionAuthAfterUnauthorized(ctx, executor, auth, bootstrapErr, alreadyTried, ephemeralResult) @@ -303,60 +550,99 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi m.replaceHomeExecutionLifecycleAuth(execOpts.ExecutionLifecycle, auth) publishSelectedAuthMetadata(execOpts.Metadata, auth) didRefreshOnUnauthorized = true - retryStream, retryErr := executor.ExecuteStream(ctx, auth, execReq, execOpts) + // Fresh TTFT budget and attempt context for the retry. + scope.release() + scope = newTTFTScope(ctx, ttftTimeout) + attemptCtx = scope.ctx + retryStream, retryErr := executor.ExecuteStream(attemptCtx, auth, execReq, execOpts) retryStream, retryErr = validateStreamResult(retryStream, retryErr) + scope.stop() + retryErr = checkTTFTErr(retryErr) if retryErr != nil { + if retryStream != nil { + discardStreamChunks(retryStream.Chunks) + } if errCtx := ctx.Err(); errCtx != nil { + scope.release() return nil, errCtx } bootstrapErr = retryErr streamResult = &cliproxyexecutor.StreamResult{} } else { streamResult = retryStream - buffered, closed, bootstrapErr = readStreamBootstrap(ctx, streamResult.Chunks) + buffered, closed, bootstrapErr = readStreamBootstrap(attemptCtx, streamResult.Chunks) + if bootstrapErr == nil && scope.timedOut() { + bootstrapErr = scope.timeoutError() + } + bootstrapErr = checkTTFTErr(bootstrapErr) } } } } if !ephemeralResult { if errCancel := claudeOAuthRequestCancellation(ctx, auth, bootstrapErr); errCancel != nil { + scope.release() discardStreamChunks(streamResult.Chunks) return nil, errCancel } } if bootstrapErr != nil { + scope.release() + bootstrapErr = checkTTFTErr(bootstrapErr) if isRequestInvalidError(bootstrapErr) { rerr := resultErrorFromError(bootstrapErr) - result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: rerr} + result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: rerr, Options: execOpts} result.RetryAfter = retryAfterFromError(bootstrapErr) + if isCredentialScopedError(bootstrapErr) { + result.CredentialScope = true + } m.recordExecutionResult(ctx, result, auth, ephemeralResult) discardStreamChunks(streamResult.Chunks) return nil, bootstrapErr } if idx < len(execModels)-1 { rerr := resultErrorFromError(bootstrapErr) - result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: rerr} + result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: rerr, Options: execOpts} result.RetryAfter = retryAfterFromError(bootstrapErr) + if isCredentialScopedError(bootstrapErr) { + result.CredentialScope = true + } m.recordExecutionResult(ctx, result, auth, ephemeralResult) discardStreamChunks(streamResult.Chunks) lastErr = bootstrapErr + if result.CredentialScope { + return nil, newStreamBootstrapError(bootstrapErr, streamResult.Headers) + } continue } rerr := resultErrorFromError(bootstrapErr) - result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: rerr} + result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: rerr, Options: execOpts} result.RetryAfter = retryAfterFromError(bootstrapErr) + if isCredentialScopedError(bootstrapErr) { + result.CredentialScope = true + } m.recordExecutionResult(ctx, result, auth, ephemeralResult) discardStreamChunks(streamResult.Chunks) return nil, newStreamBootstrapError(bootstrapErr, streamResult.Headers) } - if closed && (len(buffered) == 0 || isEmptyCompletion(buffered)) { + payloadBytes := 0 + for _, chunk := range buffered { + payloadBytes += len(chunk.Payload) + } + // Determine emptiness by buffered payload bytes, not chunk count: + // zero-payload chunks are dropped downstream by wrapStreamResult, so a + // stream of only such chunks would surface as a successful empty + // completion without failover. + if closed && (payloadBytes == 0 || isEmptyCompletion(buffered)) { + scope.release() emptyErr := errEmptyCompletion - if len(buffered) == 0 { + if payloadBytes == 0 { emptyErr = &Error{Code: "empty_stream", Message: "upstream stream closed before first payload", Retryable: true} } - result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: emptyErr} + result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: emptyErr, Options: execOpts} m.recordExecutionResult(ctx, result, auth, ephemeralResult) + discardStreamChunks(streamResult.Chunks) if idx < len(execModels)-1 { lastErr = emptyErr continue @@ -364,14 +650,16 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi return nil, newStreamBootstrapError(emptyErr, streamResult.Headers) } + scope.commit() remaining := streamResult.Chunks if closed { + discardStreamChunks(streamResult.Chunks) closedCh := make(chan cliproxyexecutor.StreamChunk) close(closedCh) remaining = closedCh } attemptAliasResult := resolveAttemptAliasResult(routing, auth, routeModel, execModel, aliasResult) - return m.wrapStreamResult(ctx, auth.Clone(), provider, resultModel, streamResult.Headers, buffered, remaining, attemptAliasResult, ephemeralResult), nil + return m.wrapStreamResult(ctx, auth.Clone(), provider, resultModel, execOpts, streamResult.Headers, buffered, remaining, attemptAliasResult, ephemeralResult, scope.release), nil } if lastErr == nil { lastErr = &Error{Code: "auth_not_found", Message: "no upstream model available"} diff --git a/sdk/cliproxy/auth/conductor_stream_drain_test.go b/sdk/cliproxy/auth/conductor_stream_drain_test.go new file mode 100644 index 000000000..d76bcae8a --- /dev/null +++ b/sdk/cliproxy/auth/conductor_stream_drain_test.go @@ -0,0 +1,159 @@ +package auth + +import ( + "context" + "net/http" + "testing" + "time" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +type drainTestExecutor struct { + streamFunc func(ctx context.Context, req cliproxyexecutor.Request) (*cliproxyexecutor.StreamResult, error) +} + +func (e *drainTestExecutor) Identifier() string { return "test-drain-provider" } + +func (e *drainTestExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} + +func (e *drainTestExecutor) ExecuteStream(ctx context.Context, _ *Auth, req cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + return e.streamFunc(ctx, req) +} + +func (e *drainTestExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} + +func (e *drainTestExecutor) Refresh(_ context.Context, a *Auth) (*Auth, error) { return a, nil } + +func (e *drainTestExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +func TestConductor_ExecuteStreamDrainsSourceOnTerminalEmpty_SingleModel(t *testing.T) { + manager := NewManager(nil, nil, nil) + auth := &Auth{ID: "auth-drain-empty-single", Provider: "test-drain-provider", Status: StatusActive} + if _, err := manager.Register(context.Background(), auth); err != nil { + t.Fatalf("register auth: %v", err) + } + reg := registry.GetGlobalRegistry() + reg.RegisterClient(auth.ID, "test-drain-provider", []*registry.ModelInfo{{ID: "test-model"}}) + t.Cleanup(func() { reg.UnregisterClient(auth.ID) }) + manager.RefreshSchedulerEntry(auth.ID) + + producerDone := make(chan struct{}) + exec := &drainTestExecutor{ + streamFunc: func(ctx context.Context, req cliproxyexecutor.Request) (*cliproxyexecutor.StreamResult, error) { + chunks := make(chan cliproxyexecutor.StreamChunk) + go func() { + defer close(producerDone) + // Terminal empty marker (OpenAI [DONE]) + chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("data: [DONE]\n\n")} + // Trailing chunk on unbuffered channel - will block if chunks not drained + chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("trailing chunk")} + close(chunks) + }() + return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil + }, + } + manager.RegisterExecutor(exec) + + res, err := manager.ExecuteStream(context.Background(), []string{"test-drain-provider"}, cliproxyexecutor.Request{Model: "test-model"}, cliproxyexecutor.Options{}) + if err != nil { + t.Fatalf("ExecuteStream unexpected error: %v", err) + } + if res == nil || res.Chunks == nil { + t.Fatal("expected non-nil StreamResult with Chunks") + } + var receivedErr error + for chunk := range res.Chunks { + if chunk.Err != nil { + receivedErr = chunk.Err + } + } + if receivedErr == nil { + t.Fatal("expected empty completion error on Chunks, got nil") + } + + select { + case <-producerDone: + // PASS: producer unblocked because streamResult.Chunks was drained + case <-time.After(500 * time.Millisecond): + t.Fatal("producer remained blocked after terminal empty error; source streamResult.Chunks was not drained") + } +} + +func TestConductor_ExecuteStreamDrainsSourceOnTerminalEmpty_ModelPoolFailover(t *testing.T) { + manager := NewManager(nil, nil, nil) + auth := &Auth{ID: "auth-drain-empty-pool", Provider: "test-drain-provider", Status: StatusActive} + + model1ProducerDone := make(chan struct{}) + model2ProducerDone := make(chan struct{}) + + exec := &drainTestExecutor{ + streamFunc: func(ctx context.Context, req cliproxyexecutor.Request) (*cliproxyexecutor.StreamResult, error) { + chunks := make(chan cliproxyexecutor.StreamChunk) + if req.Model == "model-1" { + go func() { + defer close(model1ProducerDone) + // Model 1 returns terminal empty and then attempts trailing chunk + chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("data: [DONE]\n\n")} + chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("trailing chunk 1")} + close(chunks) + }() + } else { + go func() { + defer close(model2ProducerDone) + // Model 2 returns valid streaming content + chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("data: {\"choices\":[{\"delta\":{\"content\":\"hello\"}}]}\n\n")} + close(chunks) + }() + } + return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil + }, + } + + res, err := manager.executeStreamWithModelPool( + context.Background(), + exec, + auth, + "test-drain-provider", + cliproxyexecutor.Request{Model: "pool-model"}, + cliproxyexecutor.Options{}, + "pool-model", + "", + []string{"model-1", "model-2"}, + true, + OAuthModelAliasResult{}, + nil, + true, + false, + nil, + ) + if err != nil { + t.Fatalf("executeStreamWithModelPool unexpected error: %v", err) + } + if res == nil || res.Chunks == nil { + t.Fatal("expected non-nil StreamResult with Chunks") + } + for range res.Chunks { + } + + select { + case <-model1ProducerDone: + // PASS: model-1 producer unblocked because discarded before failover + case <-time.After(500 * time.Millisecond): + t.Fatal("model-1 producer remained blocked after model failover; source was not drained") + } + + select { + case <-model2ProducerDone: + // PASS: model-2 completed normally + case <-time.After(500 * time.Millisecond): + t.Fatal("model-2 producer did not complete") + } +} diff --git a/sdk/cliproxy/auth/conductor_stream_ttft_test.go b/sdk/cliproxy/auth/conductor_stream_ttft_test.go new file mode 100644 index 000000000..631335da8 --- /dev/null +++ b/sdk/cliproxy/auth/conductor_stream_ttft_test.go @@ -0,0 +1,380 @@ +package auth + +import ( + "context" + "errors" + "net/http" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +// ttftProbeExecutor records whether the attempt context was already canceled +// when ExecuteStream was entered, then returns an immediately closed stream. +type ttftProbeExecutor struct { + calls atomic.Int32 + ctxErrAtEntry error +} + +func (e *ttftProbeExecutor) Identifier() string { return "gemini" } + +func (e *ttftProbeExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} + +func (e *ttftProbeExecutor) ExecuteStream(ctx context.Context, _ *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + e.calls.Add(1) + e.ctxErrAtEntry = ctx.Err() + chunks := make(chan cliproxyexecutor.StreamChunk) + close(chunks) + return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil +} + +func (e *ttftProbeExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} + +func (e *ttftProbeExecutor) Refresh(_ context.Context, a *Auth) (*Auth, error) { return a, nil } + +func (e *ttftProbeExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +// TestStreamTTFTTimerArmedAfterInterception is a regression guard for the +// codex P2 finding on PR #4881: the first-chunk timeout timer used to be +// armed before applyRequestAfterAuthInterceptor, so a slow interceptor could +// burn the whole TTFT budget and ExecuteStream would be invoked with an +// already-canceled context, producing a retryable 504 that cooled the +// credential although no upstream request was ever attempted. The timer must +// only be armed after local interception and request preparation complete. +func TestStreamTTFTTimerArmedAfterInterception(t *testing.T) { + manager := NewManager(nil, nil, nil) + auth := &Auth{ID: "auth-ttft", Provider: "gemini", Status: StatusActive} + if _, err := manager.Register(context.Background(), auth); err != nil { + t.Fatalf("register auth: %v", err) + } + reg := registry.GetGlobalRegistry() + reg.RegisterClient(auth.ID, "gemini", []*registry.ModelInfo{{ID: "test-model"}}) + t.Cleanup(func() { reg.UnregisterClient(auth.ID) }) + manager.RefreshSchedulerEntry(auth.ID) + exec := &ttftProbeExecutor{} + manager.RegisterExecutor(exec) + + opts := cliproxyexecutor.Options{ + Metadata: map[string]any{"stream_first_chunk_timeout_ms": 50}, + RequestAfterAuthInterceptor: func(context.Context, cliproxyexecutor.RequestAfterAuthInterceptRequest) cliproxyexecutor.RequestAfterAuthInterceptResponse { + // Deliberately slower than the 50ms TTFT budget: pre-fix the timer + // fired during this sleep and canceled the attempt context. + time.Sleep(200 * time.Millisecond) + return cliproxyexecutor.RequestAfterAuthInterceptResponse{} + }, + } + _, err := manager.ExecuteStream(context.Background(), []string{"gemini"}, cliproxyexecutor.Request{Model: "test-model"}, opts) + + if got := exec.calls.Load(); got != 1 { + t.Fatalf("expected executor to be invoked once, got %d", got) + } + if exec.ctxErrAtEntry != nil { + t.Fatalf("attempt context was already canceled at ExecuteStream entry: %v", exec.ctxErrAtEntry) + } + if err != nil && statusCodeFromError(err) == http.StatusGatewayTimeout { + t.Fatalf("TTFT timeout fired before any upstream request was attempted: %v", err) + } +} + +// ttftRefreshProbeExecutor returns a retryable 401 on the first +// ExecuteStream, simulates a slow credential refresh, and records the attempt +// context state when the refreshed request is executed. +type ttftRefreshProbeExecutor struct { + calls atomic.Int32 + refreshCalls atomic.Int32 + retryCtxErr error +} + +func (e *ttftRefreshProbeExecutor) Identifier() string { return "gemini" } + +func (e *ttftRefreshProbeExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} + +func (e *ttftRefreshProbeExecutor) ExecuteStream(ctx context.Context, _ *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + if e.calls.Add(1) == 1 { + return nil, errors.New("upstream returned status 401") + } + e.retryCtxErr = ctx.Err() + chunks := make(chan cliproxyexecutor.StreamChunk, 1) + chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("data: hello\n\n")} + close(chunks) + return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil +} + +func (e *ttftRefreshProbeExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} + +func (e *ttftRefreshProbeExecutor) Refresh(ctx context.Context, a *Auth) (*Auth, error) { + e.refreshCalls.Add(1) + select { + case <-time.After(200 * time.Millisecond): + return a, nil + case <-ctx.Done(): + return nil, ctx.Err() + } +} + +func (e *ttftRefreshProbeExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +// TestStreamTTFTTimerRestartedAfterUnauthorizedRefresh mirrors the CPA +// regression guard for the codex P2 finding on PR #4881. The CPAPlus +// conductor already restarts the TTFT scope on a fresh attempt context for +// the unauthorized-refresh retry; this test locks that behavior in: a refresh +// slower than the first-chunk budget must not leave the retried ExecuteStream +// with an already-canceled context. +func TestStreamTTFTTimerRestartedAfterUnauthorizedRefresh(t *testing.T) { + manager := NewManager(nil, nil, nil) + auth := &Auth{ + ID: "auth-oauth", + Provider: "gemini", + Status: StatusActive, + Metadata: map[string]any{"auth_kind": "oauth", "refresh_token": "x"}, + } + if _, err := manager.Register(context.Background(), auth); err != nil { + t.Fatalf("register auth: %v", err) + } + reg := registry.GetGlobalRegistry() + reg.RegisterClient(auth.ID, "gemini", []*registry.ModelInfo{{ID: "test-model"}}) + t.Cleanup(func() { reg.UnregisterClient(auth.ID) }) + manager.RefreshSchedulerEntry(auth.ID) + exec := &ttftRefreshProbeExecutor{} + manager.RegisterExecutor(exec) + + opts := cliproxyexecutor.Options{ + Metadata: map[string]any{"stream_first_chunk_timeout_ms": 50}, + } + _, err := manager.ExecuteStream(context.Background(), []string{"gemini"}, cliproxyexecutor.Request{Model: "test-model"}, opts) + + if err != nil { + t.Fatalf("ExecuteStream() error = %v, want success after refresh retry", err) + } + if got := exec.refreshCalls.Load(); got != 1 { + t.Fatalf("expected one refresh, got %d", got) + } + if got := exec.calls.Load(); got != 2 { + t.Fatalf("expected two ExecuteStream calls (401 then refreshed retry), got %d", got) + } + if exec.retryCtxErr != nil { + t.Fatalf("refreshed attempt context was already canceled at ExecuteStream entry: %v", exec.retryCtxErr) + } +} + +// zeroPayloadStreamExecutor returns a stream whose only chunk carries no +// payload bytes, then closes it. +type zeroPayloadStreamExecutor struct { + calls atomic.Int32 +} + +func (e *zeroPayloadStreamExecutor) Identifier() string { return "gemini" } + +func (e *zeroPayloadStreamExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} + +func (e *zeroPayloadStreamExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + e.calls.Add(1) + chunks := make(chan cliproxyexecutor.StreamChunk, 1) + chunks <- cliproxyexecutor.StreamChunk{Payload: nil} + close(chunks) + return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil +} + +func (e *zeroPayloadStreamExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} + +func (e *zeroPayloadStreamExecutor) Refresh(_ context.Context, a *Auth) (*Auth, error) { return a, nil } + +func (e *zeroPayloadStreamExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +// TestStreamZeroPayloadChunksAreEmptyCompletion is a regression guard for the +// codex P2 finding on PR #4881: emptiness was decided by chunk count, so a +// stream of only zero-payload chunks (dropped downstream by wrapStreamResult) +// was accepted as successful and the client received an empty completion +// without failover. Emptiness must be determined by buffered payload bytes. +// At the manager level a terminal bootstrap failure is delivered as an +// in-stream error chunk (streamErrorResult) with a nil Go error. +func TestStreamZeroPayloadChunksAreEmptyCompletion(t *testing.T) { + manager := NewManager(nil, nil, nil) + auth := &Auth{ID: "auth-zero-payload", Provider: "gemini", Status: StatusActive} + if _, err := manager.Register(context.Background(), auth); err != nil { + t.Fatalf("register auth: %v", err) + } + reg := registry.GetGlobalRegistry() + reg.RegisterClient(auth.ID, "gemini", []*registry.ModelInfo{{ID: "test-model"}}) + t.Cleanup(func() { reg.UnregisterClient(auth.ID) }) + manager.RefreshSchedulerEntry(auth.ID) + exec := &zeroPayloadStreamExecutor{} + manager.RegisterExecutor(exec) + + result, err := manager.ExecuteStream(context.Background(), []string{"gemini"}, cliproxyexecutor.Request{Model: "test-model"}, cliproxyexecutor.Options{}) + if err != nil { + t.Fatalf("ExecuteStream() error = %v, want in-stream error delivery", err) + } + if result == nil || result.Chunks == nil { + t.Fatal("ExecuteStream() result has no chunk source") + } + payloadBytes := 0 + var streamErr error + for chunk := range result.Chunks { + payloadBytes += len(chunk.Payload) + if chunk.Err != nil { + streamErr = chunk.Err + } + } + if payloadBytes != 0 { + t.Fatalf("stream delivered %d payload bytes, want 0", payloadBytes) + } + if streamErr == nil { + t.Fatal("stream closed without an error chunk, want empty_stream (silent empty completion)") + } + if !strings.Contains(streamErr.Error(), "empty_stream") && !strings.Contains(streamErr.Error(), "empty completion") && !strings.Contains(streamErr.Error(), "closed before first payload") { + t.Fatalf("stream error = %v, want an empty-stream error", streamErr) + } + if got := exec.calls.Load(); got != 1 { + t.Fatalf("expected one ExecuteStream call, got %d", got) + } +} + +type slowFirstChunkStreamExecutor struct { + calls atomic.Int32 +} + +func (e *slowFirstChunkStreamExecutor) Identifier() string { return "gemini" } + +func (e *slowFirstChunkStreamExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} + +func (e *slowFirstChunkStreamExecutor) ExecuteStream(ctx context.Context, _ *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + e.calls.Add(1) + chunks := make(chan cliproxyexecutor.StreamChunk, 1) + go func() { + time.Sleep(100 * time.Millisecond) + chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("data: {\"choices\":[{\"delta\":{\"content\":\"hello\"},\"finish_reason\":\"stop\"}]}\n\n")} + close(chunks) + }() + return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil +} + +func (e *slowFirstChunkStreamExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} + +func (e *slowFirstChunkStreamExecutor) Refresh(_ context.Context, a *Auth) (*Auth, error) { + return a, nil +} + +func (e *slowFirstChunkStreamExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +func TestStreamTTFTDeadlineStoppedOnceUpstreamConnects(t *testing.T) { + manager := NewManager(nil, nil, nil) + auth := &Auth{ID: "auth-ttft-connected", Provider: "gemini", Status: StatusActive} + if _, err := manager.Register(context.Background(), auth); err != nil { + t.Fatalf("register auth: %v", err) + } + reg := registry.GetGlobalRegistry() + reg.RegisterClient(auth.ID, "gemini", []*registry.ModelInfo{{ID: "test-model"}}) + t.Cleanup(func() { reg.UnregisterClient(auth.ID) }) + manager.RefreshSchedulerEntry(auth.ID) + exec := &slowFirstChunkStreamExecutor{} + manager.RegisterExecutor(exec) + + opts := cliproxyexecutor.Options{ + Metadata: map[string]any{"stream_first_chunk_timeout_ms": 30}, + } + result, err := manager.ExecuteStream(context.Background(), []string{"gemini"}, cliproxyexecutor.Request{Model: "test-model"}, opts) + if err != nil { + t.Fatalf("ExecuteStream() error = %v, want established stream to not time out during chunk wait", err) + } + if result == nil || result.Chunks == nil { + t.Fatal("ExecuteStream() returned nil result or nil chunks") + } + var payloadBytes int + for chunk := range result.Chunks { + if chunk.Err != nil { + t.Fatalf("stream chunk error: %v", chunk.Err) + } + payloadBytes += len(chunk.Payload) + } + if payloadBytes == 0 { + t.Fatal("expected non-empty stream payload") + } + if got := exec.calls.Load(); got != 1 { + t.Fatalf("expected 1 ExecuteStream call, got %d", got) + } +} + +type errorStreamExecutor struct { + calls atomic.Int32 +} + +func (e *errorStreamExecutor) Identifier() string { return "gemini" } + +func (e *errorStreamExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} + +func (e *errorStreamExecutor) ExecuteStream(ctx context.Context, _ *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + e.calls.Add(1) + return nil, &Error{Code: "upstream_error", Message: "service unavailable", HTTPStatus: 503} +} + +func (e *errorStreamExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} + +func (e *errorStreamExecutor) Refresh(_ context.Context, a *Auth) (*Auth, error) { + return a, nil +} + +func (e *errorStreamExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +func TestStreamTTFTDeadlineStoppedOnExecutorError(t *testing.T) { + manager := NewManager(nil, nil, nil) + auth := &Auth{ID: "auth-ttft-error", Provider: "gemini", Status: StatusActive} + if _, err := manager.Register(context.Background(), auth); err != nil { + t.Fatalf("register auth: %v", err) + } + reg := registry.GetGlobalRegistry() + reg.RegisterClient(auth.ID, "gemini", []*registry.ModelInfo{{ID: "test-model"}}) + t.Cleanup(func() { reg.UnregisterClient(auth.ID) }) + manager.RefreshSchedulerEntry(auth.ID) + exec := &errorStreamExecutor{} + manager.RegisterExecutor(exec) + + opts := cliproxyexecutor.Options{ + Metadata: map[string]any{"stream_first_chunk_timeout_ms": 20}, + } + _, err := manager.ExecuteStream(context.Background(), []string{"gemini"}, cliproxyexecutor.Request{Model: "test-model"}, opts) + if err == nil { + t.Fatal("expected error from executor, got nil") + } + if strings.Contains(err.Error(), "stream_first_chunk_timeout") || statusCodeFromError(err) == http.StatusGatewayTimeout { + t.Fatalf("expected executor error, got TTFT timeout: %v", err) + } + if !strings.Contains(err.Error(), "upstream_error") && !strings.Contains(err.Error(), "service unavailable") { + t.Fatalf("expected upstream_error, got: %v", err) + } + time.Sleep(50 * time.Millisecond) +} diff --git a/sdk/cliproxy/auth/conductor_unauthorized_refresh_test.go b/sdk/cliproxy/auth/conductor_unauthorized_refresh_test.go index 374519250..b9bad3297 100644 --- a/sdk/cliproxy/auth/conductor_unauthorized_refresh_test.go +++ b/sdk/cliproxy/auth/conductor_unauthorized_refresh_test.go @@ -5,6 +5,7 @@ import ( "net/http" "sync" "testing" + "time" "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" @@ -13,13 +14,14 @@ import ( type unauthorizedRefreshExecutor struct { id string - mu sync.Mutex - executeCalls []string - streamCalls []string - refreshCalls int - tokenInvalid map[string]struct{} - refreshFail bool - refreshTokens map[string]string + mu sync.Mutex + executeCalls []string + streamCalls []string + refreshCalls int + tokenInvalid map[string]struct{} + refreshFail bool + refreshTokens map[string]string + streamUnauthorizedResult func(auth *Auth) *cliproxyexecutor.StreamResult } func (e *unauthorizedRefreshExecutor) Identifier() string { return e.id } @@ -44,9 +46,14 @@ func (e *unauthorizedRefreshExecutor) ExecuteStream(_ context.Context, auth *Aut e.streamCalls = append(e.streamCalls, auth.ID) token := authAccessToken(auth) _, invalid := e.tokenInvalid[token] + streamFn := e.streamUnauthorizedResult e.mu.Unlock() if invalid { - return nil, &Error{ + var res *cliproxyexecutor.StreamResult + if streamFn != nil { + res = streamFn(auth) + } + return res, &Error{ HTTPStatus: http.StatusUnauthorized, Message: "Your authentication token has been invalidated. Please try signing in again.", } @@ -334,3 +341,48 @@ func TestManager_Execute_UnauthorizedRefreshThenRetryStillFailsFallsBackOnce(t * t.Fatalf("Execute calls = %v, want [primary, primary, backup]", got) } } + +func TestManager_ExecuteStream_UnauthorizedDrainsPreRefreshStreamResult(t *testing.T) { + m, executor, primary, _, model := newUnauthorizedRefreshFixture(t, false) + + producerDone := make(chan struct{}) + executor.mu.Lock() + executor.streamUnauthorizedResult = func(auth *Auth) *cliproxyexecutor.StreamResult { + ch := make(chan cliproxyexecutor.StreamChunk) + go func() { + defer close(producerDone) + for i := 0; i < 3; i++ { + ch <- cliproxyexecutor.StreamChunk{Payload: []byte("chunk")} + } + close(ch) + }() + return &cliproxyexecutor.StreamResult{ + Headers: http.Header{"X-Auth": {auth.ID}}, + Chunks: ch, + } + } + executor.mu.Unlock() + + stream, errStream := m.ExecuteStream(context.Background(), []string{"codex"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}) + if errStream != nil { + t.Fatalf("ExecuteStream error = %v, want success on refreshed primary", errStream) + } + if stream == nil || stream.Chunks == nil { + t.Fatalf("expected stream result") + } + + select { + case <-producerDone: + // success: pre-refresh stream channel was drained + case <-time.After(1 * time.Second): + t.Fatal("pre-refresh stream chunk channel producer remained blocked, want drained") + } + + chunk, ok := <-stream.Chunks + if !ok { + t.Fatalf("expected stream chunk from refreshed stream") + } + if got := string(chunk.Payload); got != primary.ID+":fresh-access-token" { + t.Fatalf("stream payload = %q, want refreshed primary response", got) + } +} diff --git a/sdk/cliproxy/auth/empty_completion.go b/sdk/cliproxy/auth/empty_completion.go index 2a091f06f..05138c8f8 100644 --- a/sdk/cliproxy/auth/empty_completion.go +++ b/sdk/cliproxy/auth/empty_completion.go @@ -5,12 +5,58 @@ import ( "context" "encoding/json" "errors" + "io" + "math" "net/http" "strings" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" ) +// tokenCount is a tolerant usage count that accepts any valid JSON number +// (integer, decimal, or exponent) and treats every other JSON value (null, +// string, object, array, or malformed) as unset, absorbing it without failing +// the enclosing frame. positive reports whether the count is a finite number +// greater than zero, the only property the empty-completion logic needs. +type tokenCount json.Number + +func (t *tokenCount) UnmarshalJSON(b []byte) error { + var n json.Number + if err := json.Unmarshal(b, &n); err != nil { + *t = "" + return nil + } + *t = tokenCount(n) + return nil +} + +// positive reports whether c is a finite JSON number greater than zero. +func (c tokenCount) positive() bool { + n := json.Number(c) + if n == "" { + return false + } + f, err := n.Float64() + if err != nil { + return false + } + return !math.IsNaN(f) && !math.IsInf(f, 0) && f > 0 +} + +// addUsage folds a positive usage count into the accumulator's token total. +// Exact integer counts are summed; fractional, huge, or otherwise non-integer +// positive values still count as output evidence so the >0 check holds. +func (a *emptyCompletionAccum) addUsage(c tokenCount) { + if !c.positive() { + return + } + if n, err := json.Number(c).Int64(); err == nil && n > 0 { + a.completionTokens += int(n) + } else { + a.completionTokens = max(a.completionTokens, 1) + } +} + // errEmptyCompletion indicates the upstream returned a terminal but empty // completion (no content, no tool calls, zero completion tokens). It is // retriable so the conductor marks the auth as failed, cools it down, and @@ -31,30 +77,225 @@ const maxStreamBootstrapBytes = 1 << 20 // completions. type openAIChunk struct { Choices []struct { + Text string `json:"text"` Delta struct { Content string `json:"content"` ReasoningContent string `json:"reasoning_content"` Refusal *string `json:"refusal"` ToolCalls []json.RawMessage `json:"tool_calls"` + FunctionCall json.RawMessage `json:"function_call"` + Audio json.RawMessage `json:"audio"` + Images []json.RawMessage `json:"images"` } `json:"delta"` Message struct { Content string `json:"content"` ReasoningContent string `json:"reasoning_content"` Refusal *string `json:"refusal"` ToolCalls []json.RawMessage `json:"tool_calls"` + FunctionCall json.RawMessage `json:"function_call"` + Audio json.RawMessage `json:"audio"` + Images []json.RawMessage `json:"images"` } `json:"message"` FinishReason *string `json:"finish_reason"` } `json:"choices"` Usage *struct { - CompletionTokens *int `json:"completion_tokens"` + CompletionTokens *tokenCount `json:"completion_tokens"` } `json:"usage"` } +// nonEmptyJSONPayload reports whether raw holds a payload beyond an empty +// null, empty string, empty object, or empty array. +func nonEmptyJSONPayload(raw json.RawMessage) bool { + trimmed := bytes.TrimSpace(raw) + if len(trimmed) == 0 || bytes.Equal(trimmed, []byte("null")) { + return false + } + var val any + if err := json.Unmarshal(trimmed, &val); err != nil { + return false + } + switch v := val.(type) { + case nil: + return false + case string: + return strings.TrimSpace(v) != "" + case map[string]any: + return len(v) > 0 + case []any: + return len(v) > 0 + default: + return true + } +} + +func hasMeaningfulJSONArguments(args string) bool { + trimmed := strings.TrimSpace(args) + if trimmed == "" || trimmed == "null" { + return false + } + var val any + if err := json.Unmarshal([]byte(trimmed), &val); err == nil { + switch v := val.(type) { + case nil: + return false + case string: + return strings.TrimSpace(v) != "" + case map[string]any: + return len(v) > 0 + case []any: + return len(v) > 0 + default: + return true + } + } + return true +} + +func hasMeaningfulClaudePartialJSON(partial string) bool { + return hasMeaningfulJSONArguments(partial) +} + +func nonEmptyAudioPayload(raw json.RawMessage) bool { + var value any + decoder := json.NewDecoder(bytes.NewReader(raw)) + decoder.UseNumber() + if err := decoder.Decode(&value); err != nil { + return false + } + if err := decoder.Decode(new(any)); err != io.EOF { + return false + } + return nonEmptyAudioValue(value) +} + +func nonEmptyAudioValue(value any) bool { + switch typed := value.(type) { + case nil: + return false + case string: + return strings.TrimSpace(typed) != "" + case bool: + return typed + case json.Number: + number, err := typed.Float64() + return err == nil && !math.IsNaN(number) && !math.IsInf(number, 0) && number != 0 + case []any: + for _, item := range typed { + if nonEmptyAudioValue(item) { + return true + } + } + case map[string]any: + for _, item := range typed { + if nonEmptyAudioValue(item) { + return true + } + } + } + return false +} + +// nonEmptyFunctionCall reports whether a legacy OpenAI function_call object +// carries a non-empty name and/or non-empty arguments. +func nonEmptyFunctionCall(raw json.RawMessage) bool { + var fc struct { + Name string `json:"name"` + Arguments string `json:"arguments"` + } + if err := json.Unmarshal(raw, &fc); err != nil { + return false + } + return strings.TrimSpace(fc.Name) != "" || hasMeaningfulJSONArguments(fc.Arguments) +} + +func hasMeaningfulImages(rawImages []json.RawMessage) bool { + for _, raw := range rawImages { + if nonEmptyJSONPayload(raw) { + return true + } + } + return false +} + +func hasMeaningfulToolCalls(rawCalls []json.RawMessage) bool { + for _, raw := range rawCalls { + if isMeaningfulToolCall(raw) { + return true + } + } + return false +} + +func isMeaningfulGeminiFunctionCall(raw json.RawMessage) bool { + trimmed := bytes.TrimSpace(raw) + if len(trimmed) == 0 || bytes.Equal(trimmed, []byte("null")) { + return false + } + var call struct { + Name string `json:"name"` + Args json.RawMessage `json:"args"` + } + if err := json.Unmarshal(trimmed, &call); err != nil { + return false + } + if strings.TrimSpace(call.Name) != "" { + return true + } + return nonEmptyJSONPayload(call.Args) +} + +func isMeaningfulToolCall(raw json.RawMessage) bool { + trimmed := bytes.TrimSpace(raw) + if len(trimmed) == 0 || bytes.Equal(trimmed, []byte("null")) { + return false + } + var call struct { + ID string `json:"id"` + Type string `json:"type"` + Function struct { + Name string `json:"name"` + Arguments string `json:"arguments"` + } `json:"function"` + Name string `json:"name"` + Arguments string `json:"arguments"` + Custom json.RawMessage `json:"custom"` + } + if err := json.Unmarshal(trimmed, &call); err != nil { + var m map[string]any + if err := json.Unmarshal(trimmed, &m); err == nil && len(m) > 0 { + for _, v := range m { + if v != nil && v != "" { + return true + } + } + } + return false + } + if strings.TrimSpace(call.ID) != "" { + return true + } + if strings.TrimSpace(call.Function.Name) != "" || hasMeaningfulJSONArguments(call.Function.Arguments) { + return true + } + if strings.TrimSpace(call.Name) != "" || hasMeaningfulJSONArguments(call.Arguments) { + return true + } + if nonEmptyJSONPayload(call.Custom) { + return true + } + return false +} + type claudeContentBlock struct { - Type string `json:"type"` - Text string `json:"text"` - Thinking string `json:"thinking"` - Input json.RawMessage `json:"input"` + ID string `json:"id"` + Name string `json:"name"` + Type string `json:"type"` + Text string `json:"text"` + Thinking string `json:"thinking"` + Signature string `json:"signature"` + Data string `json:"data"` + Input json.RawMessage `json:"input"` + Citation json.RawMessage `json:"citation"` } type claudeChunk struct { @@ -62,32 +303,38 @@ type claudeChunk struct { StopReason *string `json:"stop_reason"` Content []claudeContentBlock `json:"content"` Usage *struct { - OutputTokens *int `json:"output_tokens"` + OutputTokens *tokenCount `json:"output_tokens"` } `json:"usage"` Message *struct { Type string `json:"type"` StopReason *string `json:"stop_reason"` Content []claudeContentBlock `json:"content"` Usage *struct { - OutputTokens *int `json:"output_tokens"` + OutputTokens *tokenCount `json:"output_tokens"` } `json:"usage"` } `json:"message"` ContentBlock *claudeContentBlock `json:"content_block"` Delta *struct { - Type string `json:"type"` - Text string `json:"text"` - Thinking string `json:"thinking"` - StopReason *string `json:"stop_reason"` + Type string `json:"type"` + Text string `json:"text"` + Thinking string `json:"thinking"` + Signature string `json:"signature"` + Citation json.RawMessage `json:"citation"` + PartialJSON string `json:"partial_json"` + StopReason *string `json:"stop_reason"` } `json:"delta"` } type geminiPart struct { - Text string `json:"text"` - FunctionCall json.RawMessage `json:"functionCall"` - InlineData json.RawMessage `json:"inlineData"` - FileData json.RawMessage `json:"fileData"` - FunctionResponse json.RawMessage `json:"functionResponse"` - Thought json.RawMessage `json:"thought"` + Text string `json:"text"` + FunctionCall json.RawMessage `json:"functionCall"` + InlineData json.RawMessage `json:"inlineData"` + FileData json.RawMessage `json:"fileData"` + FunctionResponse json.RawMessage `json:"functionResponse"` + ExecutableCode json.RawMessage `json:"executableCode"` + CodeExecutionResult json.RawMessage `json:"codeExecutionResult"` + ThoughtSignature string `json:"thoughtSignature"` + Thought_Signature string `json:"thought_signature"` } type geminiCandidate struct { @@ -98,7 +345,7 @@ type geminiCandidate struct { } type geminiUsageMetadata struct { - CandidatesTokenCount *int `json:"candidatesTokenCount"` + CandidatesTokenCount *tokenCount `json:"candidatesTokenCount"` } type geminiPromptFeedback struct { @@ -119,7 +366,7 @@ type geminiChunk struct { // openAIResponseUsage is the usage block of the OpenAI Responses-API shape // (used by codex/xai executors). type openAIResponseUsage struct { - OutputTokens *int `json:"output_tokens"` + OutputTokens *tokenCount `json:"output_tokens"` } type openAIResponseContentPart struct { @@ -129,10 +376,17 @@ type openAIResponseContentPart struct { } type openAIResponseOutputItem struct { - Type string `json:"type"` - Text string `json:"text"` - Arguments string `json:"arguments"` - Content []openAIResponseContentPart `json:"content"` + ID string `json:"id"` + CallID string `json:"call_id"` + Name string `json:"name"` + Input string `json:"input"` + Type string `json:"type"` + Text string `json:"text"` + Arguments string `json:"arguments"` + Result string `json:"result"` + Content []openAIResponseContentPart `json:"content"` + EncryptedContent string `json:"encrypted_content"` + Summary json.RawMessage `json:"summary"` } type openAIResponseObject struct { @@ -169,6 +423,7 @@ var openAIResponseEventTypes = map[string]bool{ "response.output_text.done": true, "response.function_call_arguments.delta": true, "response.function_call_arguments.done": true, + "error": true, } // emptyCompletionAccum accumulates the properties relevant to deciding whether @@ -182,19 +437,49 @@ type emptyCompletionAccum struct { completionTokens int sawUsage bool blocked bool + sawMetadataOnly bool + sawMessageData bool + geminiTerminal bool + claudeTerminal bool } func (a *emptyCompletionAccum) evalJSON(data []byte) bool { - if a.evalOpenAI(data) { - return true + values, err := decodeJSONValues(data) + if err != nil { + return false } - if a.evalClaude(data) { - return true + recognized := false + for _, v := range values { + if a.evalOpenAI(v) || a.evalClaude(v) || a.evalOpenAIResponse(v) || a.evalGemini(v) { + recognized = true + } else { + a.sawUnknownData = true + } } - if a.evalOpenAIResponse(data) { - return true + return recognized +} + +// decodeJSONValues decodes every top-level JSON value in payload with the +// stdlib decoder until io.EOF, supporting pretty JSON, NDJSON, whitespace +// separated, and directly concatenated values. It requires at least one value +// and a clean EOF; malformed or trailing garbage returns an error. +func decodeJSONValues(payload []byte) ([]json.RawMessage, error) { + dec := json.NewDecoder(bytes.NewReader(payload)) + var values []json.RawMessage + for { + var raw json.RawMessage + if err := dec.Decode(&raw); err != nil { + if err == io.EOF { + break + } + return nil, err + } + values = append(values, raw) + } + if len(values) == 0 { + return nil, io.EOF } - return a.evalGemini(data) + return values, nil } func (a *emptyCompletionAccum) evalOpenAI(data []byte) bool { @@ -209,36 +494,62 @@ func (a *emptyCompletionAccum) evalOpenAI(data []byte) bool { return false } a.recognized = true + a.sawMessageData = true var chunk openAIChunk if err := json.Unmarshal(data, &chunk); err != nil { + // A recognized choices-bearing payload whose shape does not decode + // (for example message.content as an array of content parts) carries + // forward-compatible output we cannot inspect. Treat it as unknown + // data so it passes through instead of being misjudged as an empty + // completion. + a.sawUnknownData = true return true } if chunk.Usage != nil && chunk.Usage.CompletionTokens != nil { a.sawUsage = true - a.completionTokens += *chunk.Usage.CompletionTokens + a.addUsage(*chunk.Usage.CompletionTokens) } for _, ch := range chunk.Choices { if ch.FinishReason != nil { reason := strings.TrimSpace(*ch.FinishReason) - if strings.EqualFold(reason, "stop") { + if strings.EqualFold(reason, "stop") || strings.EqualFold(reason, "tool_calls") || strings.EqualFold(reason, "function_call") { a.terminal = true } else if reason != "" { // content_filter, length, and other non-stop terminal reasons // are not empty completions: the client must see the reason // rather than a silent auth rotation. a.blocked = true + a.terminal = true } } - content := ch.Delta.Content + ch.Message.Content + ch.Delta.ReasoningContent + ch.Message.ReasoningContent + content := ch.Text + ch.Delta.Content + ch.Message.Content + ch.Delta.ReasoningContent + ch.Message.ReasoningContent if strings.TrimSpace(content) != "" { a.hasContent = true } - if ch.Delta.Refusal != nil || ch.Message.Refusal != nil { + if (ch.Delta.Refusal != nil && strings.TrimSpace(*ch.Delta.Refusal) != "") || + (ch.Message.Refusal != nil && strings.TrimSpace(*ch.Message.Refusal) != "") { a.hasContent = true } - if len(ch.Delta.ToolCalls) > 0 || len(ch.Message.ToolCalls) > 0 { + if hasMeaningfulToolCalls(ch.Delta.ToolCalls) || hasMeaningfulToolCalls(ch.Message.ToolCalls) { + a.hasToolCalls = true + } + if nonEmptyFunctionCall(ch.Delta.FunctionCall) || nonEmptyFunctionCall(ch.Message.FunctionCall) { a.hasToolCalls = true } + if nonEmptyAudioPayload(ch.Delta.Audio) || nonEmptyAudioPayload(ch.Message.Audio) { + a.hasContent = true + } + if hasMeaningfulImages(ch.Delta.Images) || hasMeaningfulImages(ch.Message.Images) { + a.hasContent = true + } + } + if len(chunk.Choices) == 0 && chunk.Usage != nil { + // A completed non-streaming payload with zero choices + // ({"choices":[], "usage":...}) never enters the loop above, so + // terminal would never be set and the payload would be accepted as a + // successful response. With usage present the response is complete, so + // the empty judgment can run. + a.terminal = true } return true } @@ -251,7 +562,7 @@ func (a *emptyCompletionAccum) evalClaude(data []byte) bool { isClaude := false switch chunk.Type { - case "message", "message_start", "content_block_start", "content_block_delta", "message_delta", "message_stop": + case "message", "message_start", "content_block_start", "content_block_delta", "message_delta", "message_stop", "ping": isClaude = true default: if chunk.StopReason != nil || (chunk.Message != nil && (chunk.Message.Type == "message" || chunk.Message.StopReason != nil)) { @@ -264,6 +575,15 @@ func (a *emptyCompletionAccum) evalClaude(data []byte) bool { } a.recognized = true + if chunk.Type == "ping" { + a.sawMetadataOnly = true + } else { + a.sawMessageData = true + } + if chunk.Type == "message_stop" { + a.terminal = true + a.claudeTerminal = true + } a.evalClaudeStopReason(chunk.StopReason) if chunk.Message != nil { @@ -275,11 +595,11 @@ func (a *emptyCompletionAccum) evalClaude(data []byte) bool { if chunk.Usage != nil && chunk.Usage.OutputTokens != nil { a.sawUsage = true - a.completionTokens += *chunk.Usage.OutputTokens + a.addUsage(*chunk.Usage.OutputTokens) } if chunk.Message != nil && chunk.Message.Usage != nil && chunk.Message.Usage.OutputTokens != nil { a.sawUsage = true - a.completionTokens += *chunk.Message.Usage.OutputTokens + a.addUsage(*chunk.Message.Usage.OutputTokens) } a.evalClaudeBlocks(chunk.Content) @@ -295,14 +615,24 @@ func (a *emptyCompletionAccum) evalClaude(data []byte) bool { if strings.TrimSpace(chunk.Delta.Text) != "" { a.hasContent = true } - case "thinking_delta", "signature_delta": + case "thinking_delta": if strings.TrimSpace(chunk.Delta.Thinking) != "" { a.hasContent = true } + case "signature_delta": + if strings.TrimSpace(chunk.Delta.Signature) != "" { + a.hasContent = true + } + case "citations_delta": + if nonEmptyJSONPayload(chunk.Delta.Citation) { + a.hasContent = true + } case "input_json_delta": - a.hasToolCalls = true + if hasMeaningfulClaudePartialJSON(chunk.Delta.PartialJSON) { + a.hasToolCalls = true + } default: - if strings.TrimSpace(chunk.Delta.Text) != "" || strings.TrimSpace(chunk.Delta.Thinking) != "" { + if strings.TrimSpace(chunk.Delta.Text) != "" || strings.TrimSpace(chunk.Delta.Thinking) != "" || strings.TrimSpace(chunk.Delta.Signature) != "" || nonEmptyJSONPayload(chunk.Delta.Citation) { a.hasContent = true } } @@ -316,7 +646,7 @@ func (a *emptyCompletionAccum) evalClaudeStopReason(stopReason *string) { return } reason := strings.TrimSpace(*stopReason) - if strings.EqualFold(reason, "end_turn") { + if strings.EqualFold(reason, "end_turn") || strings.EqualFold(reason, "tool_use") { a.terminal = true } else if reason != "" { // Request/output limits, refusals, and control stop reasons must reach the @@ -344,6 +674,7 @@ func (a *emptyCompletionAccum) evalOpenAIResponse(data []byte) bool { return false } a.recognized = true + a.sawMessageData = true var chunk openAIResponseChunk if err := json.Unmarshal(data, &chunk); err != nil { @@ -352,8 +683,11 @@ func (a *emptyCompletionAccum) evalOpenAIResponse(data []byte) bool { switch evType { case "response.completed": + // Terminal Responses-API frames are valid completions even with empty + // output (see codex responses tests); never judge them empty. a.terminal = true - case "response.incomplete", "response.failed": + a.blocked = true + case "response.incomplete", "response.failed", "error": a.terminal = true a.blocked = true } @@ -364,11 +698,11 @@ func (a *emptyCompletionAccum) evalOpenAIResponse(data []byte) bool { if chunk.Usage != nil && chunk.Usage.OutputTokens != nil { a.sawUsage = true - a.completionTokens += *chunk.Usage.OutputTokens + a.addUsage(*chunk.Usage.OutputTokens) } if chunk.Response != nil && chunk.Response.Usage != nil && chunk.Response.Usage.OutputTokens != nil { a.sawUsage = true - a.completionTokens += *chunk.Response.Usage.OutputTokens + a.addUsage(*chunk.Response.Usage.OutputTokens) } switch evType { @@ -380,8 +714,28 @@ func (a *emptyCompletionAccum) evalOpenAIResponse(data []byte) bool { if strings.TrimSpace(chunk.Text) != "" { a.hasContent = true } + case "response.output_item.done": + var item openAIResponseOutputItem + if err := json.Unmarshal(chunk.Item, &item); err == nil { + itemType := strings.ToLower(strings.TrimSpace(item.Type)) + if strings.HasSuffix(itemType, "_call") { + if hasMeaningfulResponsesCallItem(item) { + a.hasToolCalls = true + } + } + } + if err := json.Unmarshal(chunk.Output, &item); err == nil { + itemType := strings.ToLower(strings.TrimSpace(item.Type)) + if strings.HasSuffix(itemType, "_call") { + if hasMeaningfulResponsesCallItem(item) { + a.hasToolCalls = true + } + } + } case "response.function_call_arguments.delta", "response.function_call_arguments.done": - a.hasToolCalls = true + if a.hasToolCalls || hasMeaningfulJSONArguments(chunk.Delta) || hasMeaningfulJSONArguments(chunk.Arguments) { + a.hasToolCalls = true + } } a.evalOpenAIResponseRawOutput(chunk.Output) @@ -397,7 +751,8 @@ func (a *emptyCompletionAccum) evalOpenAIResponseStatus(status string) { switch strings.ToLower(strings.TrimSpace(status)) { case "completed": a.terminal = true - case "incomplete", "failed": + a.blocked = true + case "incomplete", "failed", "error": a.terminal = true a.blocked = true } @@ -418,14 +773,31 @@ func (a *emptyCompletionAccum) evalOpenAIResponseRawOutput(raw json.RawMessage) } } +func hasMeaningfulResponsesCallItem(item openAIResponseOutputItem) bool { + return strings.TrimSpace(item.ID) != "" || + strings.TrimSpace(item.CallID) != "" || + strings.TrimSpace(item.Name) != "" || + hasMeaningfulJSONArguments(item.Arguments) || + strings.TrimSpace(item.Input) != "" || + strings.TrimSpace(item.Result) != "" +} + func (a *emptyCompletionAccum) evalOpenAIResponseOutput(items []openAIResponseOutputItem) { for _, item := range items { itemType := strings.ToLower(strings.TrimSpace(item.Type)) switch { case itemType == "image_generation_call": - a.hasContent = true + if hasMeaningfulResponsesCallItem(item) || strings.TrimSpace(item.Text) != "" { + a.hasContent = true + } case strings.HasSuffix(itemType, "_call"): - a.hasToolCalls = true + if hasMeaningfulResponsesCallItem(item) { + a.hasToolCalls = true + } + case itemType == "reasoning": + if strings.TrimSpace(item.EncryptedContent) != "" || nonEmptyJSONPayload(item.Summary) { + a.hasContent = true + } case itemType != "" && itemType != "message": // Responses may add output item types over time. A complete, typed // non-message item is output unless the protocol proves otherwise. @@ -446,18 +818,30 @@ func (a *emptyCompletionAccum) evalOpenAIResponseOutput(items []openAIResponseOu func (a *emptyCompletionAccum) evalClaudeBlocks(blocks []claudeContentBlock) { for _, b := range blocks { - if b.Type == "tool_use" || len(b.Input) > 0 { + if b.Type == "tool_use" || b.Type == "server_tool_use" || b.Type == "mcp_tool_use" { + if (strings.TrimSpace(b.ID) != "" && strings.TrimSpace(b.Name) != "") || nonEmptyJSONPayload(b.Input) { + a.hasToolCalls = true + } + continue + } + if nonEmptyJSONPayload(b.Input) { a.hasToolCalls = true continue } - if b.Type == "thinking" || b.Type == "redacted_thinking" || strings.TrimSpace(b.Thinking) != "" { - a.hasContent = true + if b.Type == "thinking" || b.Type == "redacted_thinking" || strings.TrimSpace(b.Thinking) != "" || strings.TrimSpace(b.Signature) != "" || strings.TrimSpace(b.Data) != "" { + if strings.TrimSpace(b.Thinking) != "" || strings.TrimSpace(b.Signature) != "" || strings.TrimSpace(b.Data) != "" { + a.hasContent = true + } continue } if strings.TrimSpace(b.Text) != "" { a.hasContent = true continue } + if nonEmptyJSONPayload(b.Citation) { + a.hasContent = true + continue + } if b.Type != "" && b.Type != "text" { a.hasContent = true } @@ -519,11 +903,15 @@ func (a *emptyCompletionAccum) evalGemini(data []byte) bool { // Gemini shape at all. if hasJSONKey(data, "candidates") || hasNestedResponseCandidates(data) { a.recognized = true + a.sawMessageData = true a.terminal = true a.blocked = promptBlocked + if !promptBlocked { + a.geminiTerminal = true + } if usage != nil && usage.CandidatesTokenCount != nil { a.sawUsage = true - a.completionTokens += *usage.CandidatesTokenCount + a.addUsage(*usage.CandidatesTokenCount) } return true } @@ -531,6 +919,7 @@ func (a *emptyCompletionAccum) evalGemini(data []byte) bool { } a.recognized = true + a.sawMessageData = true if promptBlocked { a.blocked = true } @@ -538,7 +927,7 @@ func (a *emptyCompletionAccum) evalGemini(data []byte) bool { if usage != nil { if usage.CandidatesTokenCount != nil { a.sawUsage = true - a.completionTokens += *usage.CandidatesTokenCount + a.addUsage(*usage.CandidatesTokenCount) } } @@ -562,29 +951,32 @@ func (a *emptyCompletionAccum) evalGemini(data []byte) bool { } if cand.Content != nil { for _, part := range cand.Content.Parts { - if len(part.FunctionCall) > 0 { + if isMeaningfulGeminiFunctionCall(part.FunctionCall) { a.hasToolCalls = true } - if len(part.InlineData) > 0 || - len(part.FileData) > 0 || - len(part.FunctionResponse) > 0 { + if nonEmptyJSONPayload(part.InlineData) || + nonEmptyJSONPayload(part.FileData) || + nonEmptyJSONPayload(part.FunctionResponse) { a.hasContent = true } - if len(part.Thought) > 0 { - thoughtStr := strings.TrimSpace(string(part.Thought)) - if thoughtStr != "" && thoughtStr != "false" && thoughtStr != "null" { - a.hasContent = true - } + if nonEmptyJSONPayload(part.ExecutableCode) || nonEmptyJSONPayload(part.CodeExecutionResult) { + a.hasContent = true } if strings.TrimSpace(part.Text) != "" { a.hasContent = true } + if strings.TrimSpace(part.ThoughtSignature) != "" || strings.TrimSpace(part.Thought_Signature) != "" { + a.hasContent = true + } } } } if allTerminal { a.terminal = true + if !blocked { + a.geminiTerminal = true + } } if blocked { a.blocked = true @@ -595,29 +987,34 @@ func (a *emptyCompletionAccum) evalGemini(data []byte) bool { // empty reports whether the accumulated stream is an empty completion. func (a *emptyCompletionAccum) empty() bool { - if !a.recognized || a.sawUnknownData || !a.terminal { + if a.sawUnknownData || a.blocked || a.hasContent || a.hasToolCalls || (a.sawUsage && a.completionTokens > 0) { return false } - if a.blocked { - return false + if a.recognized && a.terminal { + return true } - if a.hasContent || a.hasToolCalls { - return false + if a.recognized { + return true } - if a.sawUsage && a.completionTokens > 0 { - return false + if a.sawMetadataOnly && !a.sawMessageData { + return true } - return true + return false } // isEmptyCompletion reports whether the buffered SSE stream chunks aggregate to // an empty completion. func isEmptyCompletion(chunks []cliproxyexecutor.StreamChunk) bool { - var buf bytes.Buffer + if len(chunks) == 0 { + return false + } + var detector StreamBootstrapDetector for _, c := range chunks { - buf.Write(c.Payload) + if detector.Observe(c.Payload) { + return false + } } - return isEmptyCompletionPayload(buf.Bytes()) + return detector.Finish() } func isEmptyCompletionError(err error) bool { @@ -628,11 +1025,108 @@ func isEmptyCompletionError(err error) bool { // streamBootstrapState incrementally evaluates chunks so a metadata-heavy // prefix is processed once instead of reparsing the entire prefix per chunk. type streamBootstrapState struct { - acc emptyCompletionAccum - bytes int - pending []byte - forward bool - sawSSE bool + acc emptyCompletionAccum + bytes int + pending []byte + dataLines [][]byte + forward bool + sawSSE bool + sawDone bool +} + +func (s *streamBootstrapState) flushData() { + if len(s.dataLines) == 0 { + return + } + data := bytes.Join(s.dataLines, []byte("\n")) + s.dataLines = s.dataLines[:0] + if bytes.Equal(data, []byte("[DONE]")) { + s.acc.recognized = true + s.acc.terminal = true + s.acc.sawMessageData = true + s.sawDone = true + return + } + if len(data) == 0 { + s.acc.sawMetadataOnly = true + return + } + if !s.acc.evalJSON(data) { + s.acc.sawUnknownData = true + } +} + +func isSSEMetadataLine(b []byte) bool { + return bytes.HasPrefix(b, []byte("event:")) || + bytes.HasPrefix(b, []byte("id:")) || + bytes.HasPrefix(b, []byte("retry:")) || + bytes.HasPrefix(b, []byte(":")) || + bytes.Equal(b, []byte("event")) || + bytes.Equal(b, []byte("id")) || + bytes.Equal(b, []byte("retry")) +} + +func isSSEPrefix(b []byte) bool { + return bytes.HasPrefix(b, []byte("data:")) || + bytes.HasPrefix(b, []byte("event:")) || + bytes.HasPrefix(b, []byte("id:")) || + bytes.HasPrefix(b, []byte("retry:")) || + bytes.HasPrefix(b, []byte(":")) || + bytes.Equal(b, []byte("data")) || + bytes.Equal(b, []byte("event")) || + bytes.Equal(b, []byte("id")) || + bytes.Equal(b, []byte("retry")) +} + +func (s *streamBootstrapState) processLine(line []byte) { + line = bytes.TrimSpace(line) + if len(line) == 0 { + s.flushData() + return + } + s.processSingleLine(line) +} + +func (s *streamBootstrapState) processSingleLine(line []byte) { + switch { + case bytes.HasPrefix(line, []byte("event:")): + s.sawSSE = true + event := bytes.TrimSpace(bytes.TrimPrefix(line, []byte("event:"))) + if bytes.Equal(event, []byte("message_stop")) { + s.acc.recognized = true + s.acc.terminal = true + s.acc.sawMessageData = true + s.sawDone = true + } else { + s.acc.sawMetadataOnly = true + } + case bytes.Equal(line, []byte("event")): + s.sawSSE = true + s.acc.sawMetadataOnly = true + case bytes.HasPrefix(line, []byte("id:")), bytes.HasPrefix(line, []byte("retry:")), bytes.HasPrefix(line, []byte(":")): + s.sawSSE = true + s.acc.sawMetadataOnly = true + case bytes.Equal(line, []byte("id")), bytes.Equal(line, []byte("retry")): + s.sawSSE = true + s.acc.sawMetadataOnly = true + case bytes.HasPrefix(line, []byte("data:")): + s.sawSSE = true + s.dataLines = append(s.dataLines, parseSSEDataLine(line)) + case bytes.Equal(line, []byte("data")): + s.sawSSE = true + s.dataLines = append(s.dataLines, []byte("")) + case bytes.HasPrefix(line, []byte("{")), bytes.HasPrefix(line, []byte("[")): + s.sawSSE = true + s.dataLines = append(s.dataLines, line) + default: + if classify := classifyJSONBuffer(line); classify == jsonBufComplete || classify == jsonBufIncomplete { + if !s.acc.evalJSON(line) { + s.acc.sawUnknownData = true + } + } else { + s.acc.sawUnknownData = true + } + } } func (s *streamBootstrapState) observe(fragment []byte) bool { @@ -649,15 +1143,7 @@ func (s *streamBootstrapState) observe(fragment []byte) bool { if newline := bytes.IndexByte(s.pending, '\n'); newline >= 0 { line := bytes.TrimSpace(s.pending[:newline]) s.pending = s.pending[newline+1:] - if len(line) > 0 { - switch { - case bytes.HasPrefix(line, []byte("event:")), bytes.HasPrefix(line, []byte("data:")), bytes.HasPrefix(line, []byte(":")): - s.sawSSE = true - s.acc.evalSSE(line) - default: - s.acc.sawUnknownData = true - } - } + s.processLine(line) if s.shouldForward() { s.forward = true return true @@ -668,34 +1154,170 @@ func (s *streamBootstrapState) observe(fragment []byte) bool { } trimmed := bytes.TrimSpace(s.pending) - if len(trimmed) == 0 || couldBeSSEPrefix(trimmed) { + if len(trimmed) == 0 { return false } - if json.Valid(trimmed) { + + if bytes.HasPrefix(trimmed, []byte("data:")) { + payload := bytes.TrimSpace(trimmed[len("data:"):]) + if len(s.dataLines) == 0 && (bytes.Equal(payload, []byte("[DONE]")) || classifyJSONBuffer(payload) == jsonBufComplete) { + s.sawSSE = true + s.dataLines = append(s.dataLines, parseSSEDataLine(trimmed)) + s.flushData() + s.pending = s.pending[:0] + s.forward = s.shouldForward() + return s.forward + } + return false + } + + if couldBeSSEPrefix(trimmed) { + return false + } + switch classifyJSONBuffer(trimmed) { + case jsonBufComplete: if !s.acc.evalJSON(trimmed) { s.acc.sawUnknownData = true } s.pending = s.pending[:0] - } else if (trimmed[0] == '{' && trimmed[len(trimmed)-1] != '}') || - (trimmed[0] == '[' && trimmed[len(trimmed)-1] != ']') { + case jsonBufEmpty, jsonBufIncomplete: return false - } else { + case jsonBufInvalid: s.acc.sawUnknownData = true } s.forward = s.shouldForward() return s.forward } +func (s *streamBootstrapState) finish() { + if len(s.pending) > 0 { + trimmed := bytes.TrimSpace(s.pending) + s.pending = s.pending[:0] + if len(trimmed) > 0 { + s.processLine(trimmed) + } + } + s.flushData() +} + +func (s *streamBootstrapState) isEmptyCompletion() bool { + return s.acc.empty() +} + +func (s *streamBootstrapState) isTerminalEmpty() bool { + return (s.sawDone || s.acc.geminiTerminal || s.acc.claudeTerminal) && s.acc.empty() +} + +func (s *streamBootstrapState) hasMeaningfulOutput() bool { + if s.forward { + return true + } + if s.acc.hasContent || s.acc.hasToolCalls || s.acc.blocked || (s.acc.sawUsage && s.acc.completionTokens > 0) || s.acc.sawUnknownData { + return true + } + if !s.acc.recognized && !s.sawSSE && s.bytes > 0 { + return true + } + return false +} + func (s *streamBootstrapState) shouldForward() bool { - return s.acc.hasContent || s.acc.hasToolCalls || s.acc.sawUnknownData || (!s.acc.recognized && !s.sawSSE) + return s.acc.hasContent || s.acc.hasToolCalls || s.acc.blocked || (s.acc.sawUsage && s.acc.completionTokens > 0) || s.acc.sawUnknownData || (!s.acc.recognized && !s.sawSSE) +} + +type jsonBufferStatus int + +const ( + jsonBufEmpty jsonBufferStatus = iota + jsonBufComplete + jsonBufIncomplete + jsonBufInvalid +) + +// classifyJSONBuffer classifies an accumulated raw-JSON stream tail as holding +// one or more complete values (jsonBufComplete), a truncated prefix of a value +// (jsonBufIncomplete), malformed or trailing garbage (jsonBufInvalid), or no +// value (jsonBufEmpty). It inspects only the given buffer, so it can be called +// again on each growing chunk without keeping a persistent decoder. +func classifyJSONBuffer(buf []byte) jsonBufferStatus { + if hasTruncatedUTF8Suffix(buf) { + return jsonBufIncomplete + } + dec := json.NewDecoder(bytes.NewReader(buf)) + count := 0 + for { + var raw json.RawMessage + if err := dec.Decode(&raw); err != nil { + if err == io.EOF { + if count == 0 { + return jsonBufEmpty + } + return jsonBufComplete + } + if isTruncatedJSON(err) { + return jsonBufIncomplete + } + return jsonBufInvalid + } + count++ + } +} + +// isTruncatedJSON reports whether a json decoding error is caused by the input +// ending mid-value (a truncated prefix) rather than by malformed contents. +func isTruncatedJSON(err error) bool { + if errors.Is(err, io.ErrUnexpectedEOF) { + return true + } + var syn *json.SyntaxError + if errors.As(err, &syn) { + return strings.Contains(syn.Error(), "unexpected end of JSON input") + } + return false +} + +// hasTruncatedUTF8Suffix reports whether buf ends in the middle of a multi-byte +// UTF-8 sequence, which happens when a raw JSON value is split at a chunk +// boundary inside a string literal. +func hasTruncatedUTF8Suffix(buf []byte) bool { + n := len(buf) + if n == 0 { + return false + } + i := n - 1 + for i >= 0 && buf[i]&0xC0 == 0x80 { + i-- + } + if i < 0 { + return false + } + lead := buf[i] + var need int + switch { + case lead&0xE0 == 0xC0: + need = 1 + case lead&0xF0 == 0xE0: + need = 2 + case lead&0xF8 == 0xF0: + need = 3 + default: + return false + } + return n-i-1 < need } func couldBeSSEPrefix(payload []byte) bool { const dataPrefix = "data:" const eventPrefix = "event:" + const idPrefix = "id:" + const retryPrefix = "retry:" value := string(payload) - return strings.HasPrefix(value, ":") || strings.HasPrefix(dataPrefix, value) || strings.HasPrefix(eventPrefix, value) || - strings.HasPrefix(value, dataPrefix) || strings.HasPrefix(value, eventPrefix) + return strings.HasPrefix(value, ":") || + strings.HasPrefix(dataPrefix, value) || strings.HasPrefix(eventPrefix, value) || + strings.HasPrefix(idPrefix, value) || strings.HasPrefix(retryPrefix, value) || + strings.HasPrefix(value, dataPrefix) || strings.HasPrefix(value, eventPrefix) || + strings.HasPrefix(value, idPrefix) || strings.HasPrefix(value, retryPrefix) || + value == "data" || value == "event" || value == "id" || value == "retry" } // isEmptyCompletionPayload reports whether a payload (aggregated SSE chunks or @@ -703,45 +1325,137 @@ func couldBeSSEPrefix(payload []byte) bool { func isEmptyCompletionPayload(payload []byte) bool { trimmed := bytes.TrimSpace(payload) if len(trimmed) == 0 { - return false + // A zero-length or whitespace-only body on an HTTP success is the + // canonical empty completion: without this, Execute and plugin + // executors returned it as a successful response and never rotated + // credentials. + return true + } + + var jsonAcc emptyCompletionAccum + if jsonAcc.evalJSON(trimmed) { + var probe struct { + Choices json.RawMessage `json:"choices"` + } + if json.Unmarshal(trimmed, &probe) == nil && probe.Choices != nil { + jsonAcc.terminal = true + } + return jsonAcc.empty() } var acc emptyCompletionAccum - if bytes.Contains(trimmed, []byte("data:")) || bytes.HasPrefix(trimmed, []byte("event:")) { + if isSSEPayload(trimmed) { acc.evalSSE(trimmed) return acc.empty() } acc.evalJSON(trimmed) + var probe struct { + Choices json.RawMessage `json:"choices"` + } + if json.Unmarshal(trimmed, &probe) == nil && probe.Choices != nil { + acc.terminal = true + } return acc.empty() } -func (a *emptyCompletionAccum) evalSSE(payload []byte) { - for _, line := range bytes.Split(payload, []byte("\n")) { +func isSSEPayload(trimmed []byte) bool { + for _, line := range bytes.Split(trimmed, []byte("\n")) { line = bytes.TrimSpace(line) - if bytes.HasPrefix(line, []byte("event:")) { - event := bytes.TrimSpace(bytes.TrimPrefix(line, []byte("event:"))) - if bytes.Equal(event, []byte("message_stop")) { - a.recognized = true - } - } - if !bytes.HasPrefix(line, []byte("data:")) { + if len(line) == 0 { continue } - data := bytes.TrimSpace(bytes.TrimPrefix(line, []byte("data:"))) + if isSSEPrefix(line) { + return true + } + } + return false +} + +func parseSSEDataLine(line []byte) []byte { + data := bytes.TrimPrefix(line, []byte("data:")) + if len(data) > 0 && data[0] == ' ' { + data = data[1:] + } + return data +} + +func (a *emptyCompletionAccum) evalSSE(payload []byte) { + var dataLines [][]byte + flush := func() { + if len(dataLines) == 0 { + return + } + data := bytes.Join(dataLines, []byte("\n")) + dataLines = dataLines[:0] if bytes.Equal(data, []byte("[DONE]")) { a.recognized = true a.terminal = true - continue + a.sawMessageData = true + return } if len(data) == 0 { - continue + a.sawMetadataOnly = true + return } if !a.evalJSON(data) { a.sawUnknownData = true } } + + processSingle := func(line []byte) { + if bytes.HasPrefix(line, []byte("event:")) { + event := bytes.TrimSpace(bytes.TrimPrefix(line, []byte("event:"))) + if bytes.Equal(event, []byte("message_stop")) { + a.recognized = true + a.terminal = true + a.sawMessageData = true + } else { + a.sawMetadataOnly = true + } + return + } + if bytes.Equal(line, []byte("event")) { + a.sawMetadataOnly = true + return + } + if bytes.HasPrefix(line, []byte("id:")) || bytes.HasPrefix(line, []byte("retry:")) || bytes.HasPrefix(line, []byte(":")) { + a.sawMetadataOnly = true + return + } + if bytes.Equal(line, []byte("id")) || bytes.Equal(line, []byte("retry")) { + a.sawMetadataOnly = true + return + } + switch { + case bytes.HasPrefix(line, []byte("data:")): + dataLines = append(dataLines, parseSSEDataLine(line)) + case bytes.Equal(line, []byte("data")): + dataLines = append(dataLines, []byte("")) + case bytes.HasPrefix(line, []byte("{")), bytes.HasPrefix(line, []byte("[")): + // Some executors translate upstream SSE into the client format and + // emit raw JSON payloads without SSE framing (the HTTP handler adds + // the data: prefix later). Treat bare JSON lines as chunk data. + dataLines = append(dataLines, line) + default: + a.sawUnknownData = true + } + } + + processLine := func(line []byte) { + line = bytes.TrimSpace(line) + if len(line) == 0 { + flush() + return + } + processSingle(line) + } + + for _, line := range bytes.Split(payload, []byte("\n")) { + processLine(line) + } + flush() } // markEmptyCompletion records a failed retriable empty-completion result and diff --git a/sdk/cliproxy/auth/empty_completion_export.go b/sdk/cliproxy/auth/empty_completion_export.go index 4a6feb228..02b828b97 100644 --- a/sdk/cliproxy/auth/empty_completion_export.go +++ b/sdk/cliproxy/auth/empty_completion_export.go @@ -1,9 +1,5 @@ package auth -import ( - "bytes" -) - // IsEmptyCompletionPayload reports whether a payload (aggregated SSE chunks or // a single non-stream JSON response) represents a terminal but empty // completion. It is the exported form of the internal predicate used by the @@ -21,24 +17,6 @@ func EmptyCompletionError() error { return errEmptyCompletion } -// IsCompletionFormatRecognized reports whether payload uses a wire format the -// empty-completion detection understands (OpenAI chat, OpenAI Responses, -// Anthropic Claude, or Gemini). It supports representative format-contract -// tests without claiming registry-wide executor coverage. -func IsCompletionFormatRecognized(payload []byte) bool { - trimmed := bytes.TrimSpace(payload) - if len(trimmed) == 0 { - return false - } - var acc emptyCompletionAccum - if bytes.Contains(trimmed, []byte("data:")) || bytes.HasPrefix(trimmed, []byte("event:")) { - acc.evalSSE(trimmed) - } else { - acc.evalJSON(trimmed) - } - return acc.recognized -} - // StreamBootstrapDetector incrementally classifies a stream prefix without // reparsing previously observed chunks. Its zero value is ready for use. type StreamBootstrapDetector struct { @@ -54,3 +32,31 @@ func (d *StreamBootstrapDetector) Observe(payload []byte) bool { } return d.state.observe(payload) } + +// HasMeaningfulOutput reports whether any client-visible meaningful output +// (content, tool calls, blocked state, or non-scaffolding data) has been observed. +func (d *StreamBootstrapDetector) HasMeaningfulOutput() bool { + if d == nil { + return false + } + return d.state.hasMeaningfulOutput() +} + +// Finish flushes any trailing pending fragment at EOF and reports whether the +// accumulated stream chunks represent a terminal empty completion. +func (d *StreamBootstrapDetector) Finish() bool { + if d == nil { + return false + } + d.state.finish() + return d.state.isEmptyCompletion() +} + +// IsTerminalEmpty reports whether the accumulated stream has reached a terminal +// marker without any meaningful output. +func (d *StreamBootstrapDetector) IsTerminalEmpty() bool { + if d == nil { + return false + } + return d.state.isTerminalEmpty() +} diff --git a/sdk/cliproxy/auth/empty_completion_formats_test.go b/sdk/cliproxy/auth/empty_completion_formats_test.go index 2596e56d2..551af9f7b 100644 --- a/sdk/cliproxy/auth/empty_completion_formats_test.go +++ b/sdk/cliproxy/auth/empty_completion_formats_test.go @@ -1,7 +1,9 @@ -package auth +package auth_test import ( "testing" + + "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" ) // TestSupportedCompletionFormatsRecognized covers representative wire formats @@ -20,10 +22,11 @@ import ( // candidates-shaped) — voice/media channel, not a text completion stream. func TestSupportedCompletionFormatsRecognized(t *testing.T) { cases := []struct { - name string - executors []string - nonEmpty []byte - empty []byte + name string + executors []string + nonEmpty []byte + empty []byte + neverEmpty bool }{ { // OpenAI chat-completions wire. Emitted by the OpenAI-compatible @@ -36,10 +39,11 @@ func TestSupportedCompletionFormatsRecognized(t *testing.T) { { // OpenAI Responses-API wire (codex agent format). Emitted by the // codex-family executors (requestToFormat is FormatCodex). - name: "codex-responses", - executors: []string{"codex", "home_codex", "xai"}, - nonEmpty: []byte("data: {\"type\":\"response.output_text.delta\",\"item_id\":\"1\",\"output_index\":0,\"content_index\":0,\"delta\":\"hello\"}\n\ndata: {\"type\":\"response.completed\",\"response\":{\"id\":\"r\",\"status\":\"completed\",\"output\":[{\"type\":\"message\",\"content\":[{\"type\":\"output_text\",\"text\":\"hello\"}]}],\"usage\":{\"output_tokens\":5}}}\n\ndata: [DONE]\n\n"), - empty: []byte("data: {\"type\":\"response.completed\",\"response\":{\"id\":\"r\",\"status\":\"completed\",\"output\":[],\"usage\":{\"output_tokens\":0}}}\n\ndata: [DONE]\n\n"), + name: "codex-responses", + executors: []string{"codex", "home_codex", "xai"}, + nonEmpty: []byte("data: {\"type\":\"response.output_text.delta\",\"item_id\":\"1\",\"output_index\":0,\"content_index\":0,\"delta\":\"hello\"}\n\ndata: {\"type\":\"response.completed\",\"response\":{\"id\":\"r\",\"status\":\"completed\",\"output\":[{\"type\":\"message\",\"content\":[{\"type\":\"output_text\",\"text\":\"hello\"}]}],\"usage\":{\"output_tokens\":5}}}\n\ndata: [DONE]\n\n"), + empty: []byte("data: {\"type\":\"response.completed\",\"response\":{\"id\":\"r\",\"status\":\"completed\",\"output\":[],\"usage\":{\"output_tokens\":0}}}\n\ndata: [DONE]\n\n"), + neverEmpty: true, }, { // Anthropic Claude wire. Emitted by the Claude executor @@ -71,13 +75,19 @@ func TestSupportedCompletionFormatsRecognized(t *testing.T) { } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { - if !IsCompletionFormatRecognized(tc.nonEmpty) { + if !auth.IsCompletionFormatRecognized(tc.nonEmpty) { t.Fatalf("non-empty chunk for executors %v was NOT recognized; a new executor emitting this format would silently bypass empty-completion detection", tc.executors) } - if !IsEmptyCompletionPayload(tc.empty) { + if tc.neverEmpty { + // Responses-API terminal frames pass through by contract (existing + // repo tests define them as valid completions even with no output). + if auth.IsEmptyCompletionPayload(tc.empty) { + t.Fatalf("terminal variant for executors %v must pass through (never empty), but was judged empty", tc.executors) + } + } else if !auth.IsEmptyCompletionPayload(tc.empty) { t.Fatalf("empty-terminal variant for executors %v was not judged empty", tc.executors) } - if IsEmptyCompletionPayload(tc.nonEmpty) { + if auth.IsEmptyCompletionPayload(tc.nonEmpty) { t.Fatalf("non-empty chunk for executors %v was wrongly judged empty", tc.executors) } }) diff --git a/sdk/cliproxy/auth/empty_completion_test.go b/sdk/cliproxy/auth/empty_completion_test.go index 60e3b20fa..895524e21 100644 --- a/sdk/cliproxy/auth/empty_completion_test.go +++ b/sdk/cliproxy/auth/empty_completion_test.go @@ -6,6 +6,7 @@ import ( "net/http" "strings" "testing" + "time" "github.com/google/uuid" "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" @@ -35,6 +36,7 @@ type emptyCompletionTestExecutor struct { // non-OpenAI stream formats). emptyStreamPayload [][]byte contentStreamPayload [][]byte + leaveStreamOpen bool } func (e *emptyCompletionTestExecutor) Identifier() string { return "claude" } @@ -56,7 +58,7 @@ func (e *emptyCompletionTestExecutor) Execute(ctx context.Context, auth *Auth, _ // The first auth picked returns an empty completion; every subsequent auth // returns real content. This guarantees the rotation test exercises the // empty-completion failure path regardless of global selector state. - if e.firstExecute == auth.ID { + if len(e.executePayloads) == 0 && e.firstExecute == auth.ID { return cliproxyexecutor.Response{Payload: []byte(`{"choices":[{"message":{"content":""},"finish_reason":"stop"}],"usage":{"completion_tokens":0}}`)}, nil } if p, ok := e.executePayloads[auth.ID]; ok { @@ -93,7 +95,9 @@ func (e *emptyCompletionTestExecutor) ExecuteStream(ctx context.Context, auth *A for _, p := range empty { chunks <- cliproxyexecutor.StreamChunk{Payload: p} } - close(chunks) + if !e.leaveStreamOpen { + close(chunks) + } return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil } if payloads, ok := e.streamPayloads[auth.ID]; ok && len(payloads) > 0 { @@ -191,6 +195,71 @@ func TestEmptyCompletionPredicate(t *testing.T) { payload: []byte("data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"id\":\"x\"}]},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n"), expected: false, }, + { + name: "openai sse semantically empty tool_calls null", + payload: []byte("data: {\"choices\":[{\"delta\":{\"tool_calls\":[null]},\"finish_reason\":\"tool_calls\"}]}\n\ndata: [DONE]\n\n"), + expected: true, + }, + { + name: "openai sse semantically empty tool_calls empty object", + payload: []byte("data: {\"choices\":[{\"delta\":{\"tool_calls\":[{}]},\"finish_reason\":\"tool_calls\"}]}\n\ndata: [DONE]\n\n"), + expected: true, + }, + { + name: "openai json semantically empty tool_calls empty fields", + payload: []byte(`{"choices":[{"message":{"tool_calls":[{"id":"","type":"function","function":{"name":"","arguments":""}}]},"finish_reason":"tool_calls"}]}`), + expected: true, + }, + { + name: "openai json semantically empty tool_calls empty object args", + payload: []byte(`{"choices":[{"message":{"tool_calls":[{"id":"","type":"function","function":{"name":"","arguments":"{}"}}]},"finish_reason":"tool_calls"}]}`), + expected: true, + }, + { + name: "openai json semantically empty tool_calls empty array args", + payload: []byte(`{"choices":[{"message":{"tool_calls":[{"id":"","type":"function","function":{"name":"","arguments":"[]"}}]},"finish_reason":"tool_calls"}]}`), + expected: true, + }, + { + name: "openai json semantically empty tool_calls null args", + payload: []byte(`{"choices":[{"message":{"tool_calls":[{"id":"","type":"function","function":{"name":"","arguments":"null"}}]},"finish_reason":"tool_calls"}]}`), + expected: true, + }, + { + name: "openai sse semantically empty tool_calls empty object args", + payload: []byte("data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"id\":\"\",\"function\":{\"name\":\"\",\"arguments\":\"{}\"}}]},\"finish_reason\":\"tool_calls\"}]}\n\ndata: [DONE]\n\n"), + expected: true, + }, + { + name: "openai json semantically empty legacy function_call empty object args", + payload: []byte(`{"choices":[{"message":{"function_call":{"name":"","arguments":"{}"}},"finish_reason":"function_call"}]}`), + expected: true, + }, + { + name: "openai json semantically empty legacy function_call null args", + payload: []byte(`{"choices":[{"message":{"function_call":{"name":"","arguments":"null"}},"finish_reason":"function_call"}]}`), + expected: true, + }, + { + name: "openai json meaningful tool_calls with real args", + payload: []byte(`{"choices":[{"message":{"tool_calls":[{"id":"","type":"function","function":{"name":"","arguments":"{\"location\":\"Paris\"}"}}]},"finish_reason":"tool_calls"}]}`), + expected: false, + }, + { + name: "openai json meaningful legacy function_call with real args", + payload: []byte(`{"choices":[{"message":{"function_call":{"name":"","arguments":"{\"query\":\"test\"}"}},"finish_reason":"function_call"}]}`), + expected: false, + }, + { + name: "openai json meaningful tool_calls", + payload: []byte(`{"choices":[{"message":{"tool_calls":[{"id":"call_1","type":"function","function":{"name":"get_weather","arguments":"{}"}}]},"finish_reason":"tool_calls"}]}`), + expected: false, + }, + { + name: "openai sse meaningful tool_calls", + payload: []byte("data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"id\":\"call_1\",\"function\":{\"name\":\"get_weather\",\"arguments\":\"{}\"}}]},\"finish_reason\":\"tool_calls\"}]}\n\ndata: [DONE]\n\n"), + expected: false, + }, { name: "openai sse reasoning only is not empty", payload: []byte("data: {\"choices\":[{\"delta\":{\"reasoning_content\":\"thinking step by step\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n"), @@ -202,14 +271,14 @@ func TestEmptyCompletionPredicate(t *testing.T) { expected: false, }, { - name: "unterminated is not empty", + name: "unterminated is empty", payload: []byte("data: {\"choices\":[{\"delta\":{},\"finish_reason\":null}]}\n\n"), - expected: false, + expected: true, }, { - name: "claude sse message_stop without end_turn is not empty", + name: "claude sse message_stop without end_turn is empty", payload: []byte("event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"), - expected: false, + expected: true, }, { name: "unrecognized format is not empty", @@ -221,6 +290,16 @@ func TestEmptyCompletionPredicate(t *testing.T) { payload: []byte("data: {\"vendor_event\":\"usable-or-unknown\"}\n\ndata: [DONE]\n\n"), expected: false, }, + { + name: "openai sse with id and retry metadata then empty is empty", + payload: []byte("id: 12345\nretry: 3000\ndata: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":{\"completion_tokens\":0}}\n\ndata: [DONE]\n\n"), + expected: true, + }, + { + name: "openai sse with unknown field then empty is not empty", + payload: []byte("x-unknown: 123\ndata: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":{\"completion_tokens\":0}}\n\ndata: [DONE]\n\n"), + expected: false, + }, { name: "done-only stream remains intentionally empty", payload: []byte("data: [DONE]\n\n"), @@ -231,6 +310,21 @@ func TestEmptyCompletionPredicate(t *testing.T) { payload: []byte(`{"choices":[{"message":{"content":""},"finish_reason":"stop"}],"usage":{"completion_tokens":0}}`), expected: true, }, + { + name: "openai json semantically empty tool_calls null", + payload: []byte(`{"choices":[{"message":{"tool_calls":[null]},"finish_reason":"tool_calls"}]}`), + expected: true, + }, + { + name: "openai json semantically empty tool_calls empty object", + payload: []byte(`{"choices":[{"message":{"tool_calls":[{}]},"finish_reason":"tool_calls"}]}`), + expected: true, + }, + { + name: "openai json semantically empty tool_calls empty fields", + payload: []byte(`{"choices":[{"message":{"tool_calls":[{"id":"","type":"function","function":{"name":"","arguments":""}}]},"finish_reason":"tool_calls"}]}`), + expected: true, + }, { name: "non stream content is not empty", payload: []byte(`{"choices":[{"message":{"content":"hi"},"finish_reason":"stop"}]}`), @@ -301,11 +395,61 @@ func TestEmptyCompletionPredicate(t *testing.T) { payload: []byte(`{"candidates":[{"content":{"role":"model","parts":[]},"finishReason":"STOP"}]}`), expected: true, }, + { + name: "gemini non-stream null functionCall part is empty", + payload: []byte(`{"candidates":[{"content":{"role":"model","parts":[{"functionCall":null}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":0}}`), + expected: true, + }, + { + name: "gemini non-stream empty inlineData object part is empty", + payload: []byte(`{"candidates":[{"content":{"role":"model","parts":[{"inlineData":{}}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":0}}`), + expected: true, + }, + { + name: "gemini non-stream whitespace object inlineData part is empty", + payload: []byte(`{"candidates":[{"content":{"role":"model","parts":[{"inlineData":{ }}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":0}}`), + expected: true, + }, + { + name: "gemini non-stream whitespace array functionCall part is empty", + payload: []byte(`{"candidates":[{"content":{"role":"model","parts":[{"functionCall":[ ]}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":0}}`), + expected: true, + }, + { + name: "gemini sse whitespace inlineData stream is empty", + payload: []byte("data: {\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"inlineData\":{ }}]},\"finishReason\":\"STOP\"}],\"usageMetadata\":{\"candidatesTokenCount\":0}}\n\n"), + expected: true, + }, + { + name: "gemini sse whitespace functionCall array stream is empty", + payload: []byte("data: {\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"functionCall\":[ ]}]},\"finishReason\":\"STOP\"}],\"usageMetadata\":{\"candidatesTokenCount\":0}}\n\n"), + expected: true, + }, + { + name: "gemini non-stream null functionResponse part is empty", + payload: []byte(`{"candidates":[{"content":{"role":"model","parts":[{"functionResponse":null}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":0}}`), + expected: true, + }, { name: "gemini non-stream with functionCall part is not empty", payload: []byte(`{"candidates":[{"content":{"role":"model","parts":[{"functionCall":{"name":"search","args":{}}}]},"finishReason":"STOP"}]}`), expected: false, }, + { + name: "gemini non-stream with empty-name empty-args functionCall is empty", + payload: []byte(`{"candidates":[{"content":{"role":"model","parts":[{"functionCall":{"name":"","args":{}}}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":0}}`), + expected: true, + }, + { + name: "gemini sse with empty-name empty-args functionCall is empty", + payload: []byte("data: {\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"functionCall\":{\"name\":\"\",\"args\":{}}}]},\"finishReason\":\"STOP\"}],\"usageMetadata\":{\"candidatesTokenCount\":0}}\n\n"), + expected: true, + }, + { + name: "gemini non-stream with functionCall args and empty name is not empty", + payload: []byte(`{"candidates":[{"content":{"role":"model","parts":[{"functionCall":{"name":"","args":{"query":"hello"}}}]},"finishReason":"STOP"}]}`), + expected: false, + }, { name: "gemini non-stream with text content is not empty", payload: []byte(`{"candidates":[{"content":{"role":"model","parts":[{"text":"hello"}]},"finishReason":"STOP"}]}`), @@ -316,6 +460,26 @@ func TestEmptyCompletionPredicate(t *testing.T) { payload: []byte(`{"candidates":[{"content":{"role":"model","parts":[{"thought":true,"text":"thinking"}]},"finishReason":"STOP"}]}`), expected: false, }, + { + name: "gemini non-stream with empty text and thought flag is empty", + payload: []byte(`{"candidates":[{"content":{"role":"model","parts":[{"thought":true,"text":""}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":0}}`), + expected: true, + }, + { + name: "gemini non-stream with thought flag only is empty", + payload: []byte(`{"candidates":[{"content":{"role":"model","parts":[{"thought":true}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":0}}`), + expected: true, + }, + { + name: "gemini sse stream with empty text and thought flag is empty", + payload: []byte("data: {\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"thought\":true,\"text\":\"\"}]},\"finishReason\":\"STOP\"}],\"usageMetadata\":{\"candidatesTokenCount\":0}}\n\n"), + expected: true, + }, + { + name: "antigravity stream with empty text and thought flag is empty", + payload: []byte("data: {\"response\":{\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"thought\":true,\"text\":\"\"}]},\"finishReason\":\"STOP\"}]}}\n\n"), + expected: true, + }, { name: "gemini sse empty stream is empty", payload: []byte("data: {\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[]},\"finishReason\":\"STOP\"}],\"usageMetadata\":{\"candidatesTokenCount\":0}}\n\n"), @@ -382,20 +546,55 @@ func TestEmptyCompletionPredicate(t *testing.T) { expected: true, }, { - name: "codex responses-api sse completed with empty output is empty", - payload: []byte("data: {\"type\":\"response.completed\",\"response\":{\"id\":\"r\",\"status\":\"completed\",\"output\":[],\"usage\":{\"output_tokens\":0}}}\n\ndata: [DONE]\n\n"), + name: "openai legacy non-stream text content is not empty", + payload: []byte(`{"choices":[{"text":"hello","finish_reason":"stop"}]}`), + expected: false, + }, + { + name: "openai legacy non-stream text content with zero usage is not empty", + payload: []byte(`{"choices":[{"text":"hello","finish_reason":"stop"}],"usage":{"completion_tokens":0}}`), + expected: false, + }, + { + name: "openai legacy non-stream empty text is empty", + payload: []byte(`{"choices":[{"text":"","finish_reason":"stop"}],"usage":{"completion_tokens":0}}`), expected: true, }, { - name: "codex responses-api non-stream completed with empty output is empty", - payload: []byte(`{"object":"response","id":"r","status":"completed","output":[],"usage":{"output_tokens":0}}`), + name: "openai legacy non-stream whitespace text is empty", + payload: []byte(`{"choices":[{"text":" ","finish_reason":"stop"}],"usage":{"completion_tokens":0}}`), expected: true, }, { - name: "codex responses-api sse output_item message empty then completed is empty", - payload: []byte("data: {\"type\":\"response.output_item.done\",\"output\":{\"type\":\"message\",\"content\":[{\"type\":\"output_text\",\"text\":\"\",\"annotations\":[]}],\"status\":\"completed\"}}\n\ndata: {\"type\":\"response.completed\",\"response\":{\"id\":\"r\",\"status\":\"completed\",\"output\":[{\"type\":\"message\",\"content\":[{\"type\":\"output_text\",\"text\":\"\",\"annotations\":[]}]}],\"usage\":{\"output_tokens\":0}}}\n\ndata: [DONE]\n\n"), + name: "openai legacy sse text content stream is not empty", + payload: []byte("data: {\"choices\":[{\"text\":\"hello\",\"finish_reason\":null}]}\n\ndata: {\"choices\":[{\"text\":\"\",\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n"), + expected: false, + }, + { + name: "openai legacy sse empty text stream is empty", + payload: []byte("data: {\"choices\":[{\"text\":\"\",\"finish_reason\":\"stop\"}],\"usage\":{\"completion_tokens\":0}}\n\ndata: [DONE]\n\n"), + expected: true, + }, + { + name: "openai legacy sse whitespace text stream is empty", + payload: []byte("data: {\"choices\":[{\"text\":\" \",\"finish_reason\":\"stop\"}],\"usage\":{\"completion_tokens\":0}}\n\ndata: [DONE]\n\n"), expected: true, }, + { + name: "codex responses-api sse completed with empty output passes through (never empty by contract)", + payload: []byte("data: {\"type\":\"response.completed\",\"response\":{\"id\":\"r\",\"status\":\"completed\",\"output\":[],\"usage\":{\"output_tokens\":0}}}\n\ndata: [DONE]\n\n"), + expected: false, + }, + { + name: "codex responses-api non-stream completed with empty output passes through (never empty by contract)", + payload: []byte(`{"object":"response","id":"r","status":"completed","output":[],"usage":{"output_tokens":0}}`), + expected: false, + }, + { + name: "codex responses-api sse output_item message empty then completed passes through", + payload: []byte("data: {\"type\":\"response.output_item.done\",\"output\":{\"type\":\"message\",\"content\":[{\"type\":\"output_text\",\"text\":\"\",\"annotations\":[]}],\"status\":\"completed\"}}\n\ndata: {\"type\":\"response.completed\",\"response\":{\"id\":\"r\",\"status\":\"completed\",\"output\":[{\"type\":\"message\",\"content\":[{\"type\":\"output_text\",\"text\":\"\",\"annotations\":[]}]}],\"usage\":{\"output_tokens\":0}}}\n\ndata: [DONE]\n\n"), + expected: false, + }, { name: "codex responses-api non-stream with function_call is not empty", payload: []byte(`{"object":"response","id":"r","status":"completed","output":[{"type":"function_call","name":"get_weather","arguments":"{}","call_id":"call_1"}],"usage":{"output_tokens":5}}`), @@ -475,85 +674,617 @@ func TestEmptyCompletionPredicate(t *testing.T) { }) } } +func TestEmptyCompletionTolerantUsage(t *testing.T) { + cases := []struct { + name string + payload []byte + expected bool + }{ + { + name: "openai completion_tokens 1e2 positive is not empty", + payload: []byte(`{"choices":[{"message":{"content":""},"finish_reason":"stop"}],"usage":{"completion_tokens":1e2}}`), + expected: false, + }, + { + name: "openai completion_tokens 1.5 positive is not empty", + payload: []byte(`{"choices":[{"message":{"content":""},"finish_reason":"stop"}],"usage":{"completion_tokens":1.5}}`), + expected: false, + }, + { + name: "openai completion_tokens 100.0 positive is not empty", + payload: []byte(`{"choices":[{"message":{"content":""},"finish_reason":"stop"}],"usage":{"completion_tokens":100.0}}`), + expected: false, + }, + { + name: "openai completion_tokens zero stays empty", + payload: []byte(`{"choices":[{"message":{"content":""},"finish_reason":"stop"}],"usage":{"completion_tokens":0}}`), + expected: true, + }, + { + name: "openai completed payload with empty choices array is empty", + payload: []byte(`{"id":"chatcmpl-x","object":"chat.completion","choices":[],"usage":{"prompt_tokens":10,"completion_tokens":0,"total_tokens":10}}`), + expected: true, + }, + { + name: "openai empty choices without usage is terminal and empty", + payload: []byte(`{"choices":[]}`), + expected: true, + }, + { + name: "openai empty choices with null usage is terminal and empty", + payload: []byte(`{"choices":[],"usage":null}`), + expected: true, + }, + { + name: "openai array content parts pass through as unknown data", + payload: []byte(`{"id":"chatcmpl-x","object":"chat.completion","choices":[{"index":0,"message":{"role":"assistant","content":[{"type":"output_text","text":"hello"}]},"finish_reason":"stop"}]}`), + expected: false, + }, + { + name: "openai empty refusal string is terminal and empty", + payload: []byte(`{"choices":[{"index":0,"message":{"role":"assistant","content":"","refusal":""},"finish_reason":"stop"}],"usage":{"prompt_tokens":5,"completion_tokens":0,"total_tokens":5}}`), + expected: true, + }, + { + name: "openai real refusal string is not empty", + payload: []byte(`{"choices":[{"index":0,"message":{"role":"assistant","content":null,"refusal":"I cannot help with that"},"finish_reason":"stop"}],"usage":{"prompt_tokens":5,"completion_tokens":0,"total_tokens":5}}`), + expected: false, + }, + { + name: "zero-length body is an empty completion", + payload: []byte(``), + expected: true, + }, + { + name: "whitespace-only body is an empty completion", + payload: []byte(" \n\t "), + expected: true, + }, + { + name: "openai completion_tokens negative stays empty", + payload: []byte(`{"choices":[{"message":{"content":""},"finish_reason":"stop"}],"usage":{"completion_tokens":-5}}`), + expected: true, + }, + { + name: "openai completion_tokens overflow stays empty", + payload: []byte(`{"choices":[{"message":{"content":""},"finish_reason":"stop"}],"usage":{"completion_tokens":1e999}}`), + expected: true, + }, + { + name: "openai malformed completion_tokens with content is not empty", + payload: []byte(`{"choices":[{"message":{"content":"hello"},"finish_reason":"stop"}],"usage":{"completion_tokens":"abc"}}`), + expected: false, + }, + { + name: "openai malformed completion_tokens alone stays empty", + payload: []byte(`{"choices":[{"message":{"content":""},"finish_reason":"stop"}],"usage":{"completion_tokens":"abc"}}`), + expected: true, + }, + { + name: "claude message usage exponent positive is not empty", + payload: []byte("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"stop_reason\":null,\"usage\":{\"output_tokens\":1e2}}}\n\nevent: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":0}}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"), + expected: false, + }, + { + name: "openai responses output_tokens decimal positive keeps terminal blocking", + payload: []byte(`{"object":"response","id":"r","status":"completed","output":[],"usage":{"output_tokens":1.5}}`), + expected: false, + }, + { + name: "gemini candidatesTokenCount exponent positive is not empty", + payload: []byte(`{"candidates":[{"content":{"role":"model","parts":[]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":1e2}}`), + expected: false, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := isEmptyCompletionPayload(tc.payload); got != tc.expected { + t.Fatalf("isEmptyCompletionPayload() = %v, want %v", got, tc.expected) + } + }) + } +} -func TestStreamBootstrapStateForwardsAtMetadataLimit(t *testing.T) { - var state streamBootstrapState - metadata := []byte("data: {\"type\":\"response.in_progress\",\"response\":{\"status\":\"in_progress\"}}\n\n") - for state.bytes+len(metadata) <= maxStreamBootstrapBytes { - if state.observe(metadata) { - t.Fatal("bootstrap forwarded recognized metadata before reaching its byte limit") - } +// TestStreamBootstrapDetectorClaudePing is a regression guard for the codex +// P2 finding on PR #4881: Claude streams may emit {"type":"ping"} keep-alive +// events. evalClaude used to treat them as unknown payloads, permanently +// switching the detector to forwarding mode, so a terminally empty completion +// after a ping bypassed failover and surfaced as a successful empty stream. +func TestStreamBootstrapDetectorClaudePing(t *testing.T) { + var detector StreamBootstrapDetector + if detector.Observe([]byte("event: ping\ndata: {\"type\":\"ping\"}\n\n")) { + t.Fatal("Observe() forwarded after Claude ping keep-alive") } - if !state.observe(metadata) { - t.Fatal("bootstrap did not conservatively forward after reaching its byte limit") + if detector.Observe([]byte("event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n")) { + t.Fatal("Observe() forwarded terminal empty Claude stream preceded by ping") } } -func TestStreamBootstrapDetector(t *testing.T) { +func TestStreamBootstrapDetectorSSEMetadataFields(t *testing.T) { var detector StreamBootstrapDetector - if detector.Observe([]byte("data: {\"type\":\"response.created\",\"response\":{\"status\":\"in_progress\"}}\n\n")) { - t.Fatal("StreamBootstrapDetector.Observe() = true for metadata-only prefix") + if detector.Observe([]byte("id: evt_12345\nretry: 5000\n")) { + t.Fatal("Observe() forwarded after standard SSE id/retry metadata") } - if !detector.Observe([]byte("data: {\"type\":\"response.output_item.done\",\"item\":{\"type\":\"custom_tool_call\"}}\n\n")) { - t.Fatal("StreamBootstrapDetector.Observe() = false after complete custom tool output") + if detector.Observe([]byte("data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":{\"completion_tokens\":0}}\n\n")) { + t.Fatal("Observe() forwarded recognized empty terminal chunk") + } + if detector.Observe([]byte("data: [DONE]\n\n")) { + t.Fatal("Observe() forwarded [DONE]") + } + if !detector.Finish() { + t.Fatal("Finish() = false, want empty completion recognized despite id/retry metadata") + } + + var detectorUnknown StreamBootstrapDetector + if !detectorUnknown.Observe([]byte("x-unknown-metadata: foo\n")) { + t.Fatal("Observe() = false, want unknown SSE metadata to force forwarding") } } -func TestStreamBootstrapDetectorHandlesSplitSSEFrames(t *testing.T) { - t.Run("terminal empty remains buffered", func(t *testing.T) { +func TestStreamBootstrapDetectorMetadataOnlyEOF(t *testing.T) { + t.Run("comments and keepalive only then EOF classifies as empty", func(t *testing.T) { var detector StreamBootstrapDetector - fragments := [][]byte{ - []byte("da"), - []byte("ta: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n"), - []byte("\nda"), - []byte("ta: [DO"), - []byte("NE]\n\n"), + if detector.Observe([]byte(": keep-alive\n\n")) { + t.Fatal("Observe() forwarded keep-alive comment") } - for i, fragment := range fragments { - if detector.Observe(fragment) { - t.Fatalf("Observe(fragment %d) forwarded terminal empty stream", i) - } + if !detector.Finish() { + t.Fatal("Finish() = false, want metadata-only stream recognized as empty completion at EOF") } }) - t.Run("meaningful output forwards after complete line", func(t *testing.T) { + t.Run("id and retry metadata only then EOF classifies as empty", func(t *testing.T) { var detector StreamBootstrapDetector - if detector.Observe([]byte("da")) { - t.Fatal("Observe() forwarded incomplete SSE prefix") - } - if detector.Observe([]byte("ta: {\"type\":\"response.output_text.delta\",\"delta\":\"hel")) { - t.Fatal("Observe() forwarded incomplete meaningful SSE line") + if detector.Observe([]byte("id: evt_12345\nretry: 5000\n\n")) { + t.Fatal("Observe() forwarded id/retry metadata") } - if !detector.Observe([]byte("lo\"}\n\n")) { - t.Fatal("Observe() did not forward completed meaningful SSE line") + if !detector.Finish() { + t.Fatal("Finish() = false, want id/retry-only stream recognized as empty completion at EOF") } }) - t.Run("opaque payload forwards promptly", func(t *testing.T) { + t.Run("claude ping only then EOF classifies as empty", func(t *testing.T) { var detector StreamBootstrapDetector - if !detector.Observe([]byte("opaque-provider-payload")) { - t.Fatal("Observe() buffered definitely unrecognized payload") + if detector.Observe([]byte("event: ping\ndata: {\"type\":\"ping\"}\n\n")) { + t.Fatal("Observe() forwarded ping metadata") + } + if !detector.Finish() { + t.Fatal("Finish() = false, want ping-only stream recognized as empty completion at EOF") } }) - t.Run("event line waits for empty claude data", func(t *testing.T) { + t.Run("data-bearing stream still forwards and does not classify as empty", func(t *testing.T) { var detector StreamBootstrapDetector - fragments := [][]byte{ - []byte("event: message_start\n"), - []byte("data: {\"type\":\"message_start\",\"message\":{\"type\":\"message\",\"content\":[],\"stop_reason\":null,\"usage\":{\"output_tokens\":0}}}\n\n"), - []byte("event: message_delta\n"), - []byte("data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":0}}\n\n"), - []byte("event: message_stop\n"), - []byte("data: {\"type\":\"message_stop\"}\n\n"), + if !detector.Observe([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"hello\"}}]}\n\n")) { + t.Fatal("Observe() = false, want data-bearing stream to forward") } - for i, fragment := range fragments { - if detector.Observe(fragment) { - t.Fatalf("Observe(fragment %d) forwarded empty Claude stream", i) - } + if detector.Finish() { + t.Fatal("Finish() = true, want data-bearing stream not recognized as empty completion") } }) - t.Run("split comment waits for terminal empty data", func(t *testing.T) { + t.Run("unknown non-SSE format keeps existing behavior", func(t *testing.T) { + var detector StreamBootstrapDetector + if !detector.Observe([]byte("{\"status\":\"running\"}")) { + t.Fatal("Observe() = false, want unknown format to forward") + } + if detector.Finish() { + t.Fatal("Finish() = true, want unknown format not recognized as empty completion") + } + }) + + t.Run("isEmptyCompletionPayload classifies metadata-only SSE payloads as empty", func(t *testing.T) { + if !IsEmptyCompletionPayload([]byte(": keep-alive\n\n")) { + t.Fatal("IsEmptyCompletionPayload() = false for comment-only SSE") + } + if !IsEmptyCompletionPayload([]byte("id: evt_12345\nretry: 5000\n\n")) { + t.Fatal("IsEmptyCompletionPayload() = false for id/retry-only SSE") + } + if !IsEmptyCompletionPayload([]byte("event: ping\ndata: {\"type\":\"ping\"}\n\n")) { + t.Fatal("IsEmptyCompletionPayload() = false for ping-only SSE") + } + }) +} + +func TestStreamBootstrapDetectorTerminalBlockedForwardsImmediately(t *testing.T) { + tests := []struct { + name string + payload string + }{ + { + name: "openai content_filter", + payload: "data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"content_filter\"}]}\n\n", + }, + { + name: "openai length", + payload: "data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"length\"}]}\n\n", + }, + { + name: "gemini safety candidate", + payload: "data: {\"candidates\":[{\"finishReason\":\"SAFETY\"}]}\n\n", + }, + { + name: "gemini prompt feedback block", + payload: "data: {\"promptFeedback\":{\"blockReason\":\"SAFETY\"}}\n\n", + }, + { + name: "claude refusal", + payload: "event: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"refusal\"}}\n\n", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var detector StreamBootstrapDetector + if !detector.Observe([]byte(tt.payload)) { + t.Fatalf("Observe() = false, want terminal blocked frame to forward immediately without waiting for EOF") + } + if detector.Finish() { + t.Fatalf("Finish() = true, want terminal blocked stream not to be classified as empty completion") + } + }) + } +} + +func TestStreamBootstrapDetectorEmptyDataEventsClassifyAsEmpty(t *testing.T) { + t.Run("single empty data event", func(t *testing.T) { + var detector StreamBootstrapDetector + if detector.Observe([]byte("data:\n\n")) { + t.Fatal("Observe() forwarded empty data event") + } + if !detector.Finish() { + t.Fatal("Finish() = false, want empty data event stream classified as empty completion at EOF") + } + }) + + t.Run("empty data event with whitespace", func(t *testing.T) { + var detector StreamBootstrapDetector + if detector.Observe([]byte("data: \n\n")) { + t.Fatal("Observe() forwarded whitespace empty data event") + } + if !detector.Finish() { + t.Fatal("Finish() = false, want whitespace empty data event stream classified as empty completion at EOF") + } + }) + + t.Run("multiple empty data events", func(t *testing.T) { + var detector StreamBootstrapDetector + if detector.Observe([]byte("data:\n\ndata:\n\n")) { + t.Fatal("Observe() forwarded multiple empty data events") + } + if !detector.Finish() { + t.Fatal("Finish() = false, want multiple empty data events classified as empty completion at EOF") + } + }) + + t.Run("isEmptyCompletionPayload classifies empty data event as empty", func(t *testing.T) { + if !IsEmptyCompletionPayload([]byte("data:\n\n")) { + t.Fatal("IsEmptyCompletionPayload() = false for empty data: event") + } + if !IsEmptyCompletionPayload([]byte("data: \n\n")) { + t.Fatal("IsEmptyCompletionPayload() = false for whitespace empty data: event") + } + }) +} + +func TestStreamBootstrapDetectorOpaqueSSEMetadata(t *testing.T) { + t.Run("event containing data: substring does not parse suffix as data", func(t *testing.T) { + var detector StreamBootstrapDetector + payload := []byte("event: metadata:ping\ndata: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n") + if detector.Observe(payload) { + t.Fatal("Observe() forwarded empty completion stream with event: metadata:ping") + } + if !detector.Finish() { + t.Fatal("Finish() = false, want empty completion recognized despite event: metadata:ping") + } + if detector.state.acc.sawUnknownData { + t.Fatal("sawUnknownData = true, want event field value to remain opaque") + } + }) + + t.Run("comment containing data: substring does not parse suffix as data", func(t *testing.T) { + var detector StreamBootstrapDetector + payload := []byte(": data: keep-alive\ndata: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n") + if detector.Observe(payload) { + t.Fatal("Observe() forwarded empty completion stream with : data: keep-alive comment") + } + if !detector.Finish() { + t.Fatal("Finish() = false, want empty completion recognized despite : data: keep-alive comment") + } + if detector.state.acc.sawUnknownData { + t.Fatal("sawUnknownData = true, want comment field value to remain opaque") + } + }) + + t.Run("isEmptyCompletionPayload classifies payload with metadata containing data: as empty", func(t *testing.T) { + payload := []byte("event: metadata:ping\n: data: keep-alive\ndata: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n") + if !IsEmptyCompletionPayload(payload) { + t.Fatal("IsEmptyCompletionPayload() = false for payload with metadata containing data:") + } + }) + + t.Run("control: real data field with data: inside JSON value parses correctly", func(t *testing.T) { + var detector StreamBootstrapDetector + payload := []byte("data: {\"choices\":[{\"delta\":{\"content\":\"data: hello\"},\"finish_reason\":null}]}\n\n") + if !detector.Observe(payload) { + t.Fatal("Observe() = false, want content payload to forward") + } + if detector.Finish() { + t.Fatal("Finish() = true, want content stream not classified as empty") + } + }) + + t.Run("split metadata line across chunk boundary followed by data-like content", func(t *testing.T) { + var detector StreamBootstrapDetector + // Chunk 1 has partial metadata line: "event:" without newline + // Chunk 2 has continuation of event name "data:ping\n" followed by empty completion data line + chunk1 := []byte("event:") + chunk2 := []byte("data:ping\n") + chunk3 := []byte("data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n") + + if detector.Observe(chunk1) { + t.Fatal("Observe(chunk1) forwarded partial event line") + } + if detector.Observe(chunk2) { + t.Fatal("Observe(chunk2) forwarded event line continuation") + } + if detector.state.acc.sawUnknownData { + t.Fatal("sawUnknownData = true after event: continuation, want metadata field value to remain opaque") + } + if detector.Observe(chunk3) { + t.Fatal("Observe(chunk3) forwarded empty completion stream") + } + if !detector.Finish() { + t.Fatal("Finish() = false, want empty completion recognized when metadata line split across chunk boundary") + } + }) + + t.Run("arbitrary split points of metadata lines do not set sawUnknownData", func(t *testing.T) { + fullPayload := "event: metadata:ping_data:123\n: data: keep-alive\ndata: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n" + for split := 1; split < 40; split++ { + var detector StreamBootstrapDetector + c1 := []byte(fullPayload[:split]) + c2 := []byte(fullPayload[split:]) + if detector.Observe(c1) { + t.Fatalf("split %d: Observe(c1) forwarded unexpectedly", split) + } + if detector.Observe(c2) { + t.Fatalf("split %d: Observe(c2) forwarded unexpectedly", split) + } + if detector.state.acc.sawUnknownData { + t.Fatalf("split %d: sawUnknownData = true, want metadata value to remain opaque across split", split) + } + if !detector.Finish() { + t.Fatalf("split %d: Finish() = false, want empty completion recognized", split) + } + } + }) +} + +func TestStreamBootstrapDetectorMultilineSSE(t *testing.T) { + t.Run("empty completion split across data fields remains buffered and recognized", func(t *testing.T) { + var detector StreamBootstrapDetector + fragments := [][]byte{ + []byte("data: {\n"), + []byte("data: \"id\": \"chatcmpl-test\",\n"), + []byte("data: \"choices\": [\n"), + []byte("data: {\n"), + []byte("data: \"index\": 0,\n"), + []byte("data: \"delta\": {},\n"), + []byte("data: \"finish_reason\": \"stop\"\n"), + []byte("data: }\n"), + []byte("data: ],\n"), + []byte("data: \"usage\": {\n"), + []byte("data: \"prompt_tokens\": 5,\n"), + []byte("data: \"completion_tokens\": 0,\n"), + []byte("data: \"total_tokens\": 5\n"), + []byte("data: }\n"), + []byte("data: }\n\n"), + []byte("data: [DONE]\n\n"), + } + for i, f := range fragments { + if detector.Observe(f) { + t.Fatalf("Observe(fragment %d: %q) forwarded empty completion stream", i, string(f)) + } + } + if !detector.Finish() { + t.Fatal("Finish() = false, want multiline empty completion recognized") + } + }) + + t.Run("multiline SSE with content forwards at event boundary", func(t *testing.T) { + var detector StreamBootstrapDetector + fragments := [][]byte{ + []byte("data: {\n"), + []byte("data: \"choices\": [\n"), + []byte("data: {\n"), + []byte("data: \"delta\": {\n"), + []byte("data: \"content\": \"Hello\"\n"), + []byte("data: }\n"), + []byte("data: }\n"), + []byte("data: ]\n"), + []byte("data: }\n\n"), + } + forwarded := false + for _, f := range fragments { + if detector.Observe(f) { + forwarded = true + break + } + } + if !forwarded { + t.Fatal("Observe() = false, want multiline content event to forward") + } + if detector.Finish() { + t.Fatal("Finish() = true, want non-empty multiline stream not recognized as empty completion") + } + }) + + t.Run("isEmptyCompletionPayload handles multiline SSE payloads", func(t *testing.T) { + openaiEmpty := []byte("data: {\ndata: \"choices\": [{\"delta\":{},\"finish_reason\":\"stop\"}],\ndata: \"usage\": {\"completion_tokens\": 0}\ndata: }\n\ndata: [DONE]\n\n") + if !IsEmptyCompletionPayload(openaiEmpty) { + t.Fatal("IsEmptyCompletionPayload() = false for multiline OpenAI empty completion") + } + + geminiEmpty := []byte("data: {\ndata: \"candidates\": [\ndata: {\"finishReason\": \"STOP\"}\ndata: ],\ndata: \"usageMetadata\": {\"candidatesTokenCount\": 0}\ndata: }\n\n") + if !IsEmptyCompletionPayload(geminiEmpty) { + t.Fatal("IsEmptyCompletionPayload() = false for multiline Gemini empty completion") + } + + malformed := []byte("data: {\ndata: not valid json\ndata: }\n\n") + if IsEmptyCompletionPayload(malformed) { + t.Fatal("IsEmptyCompletionPayload() = true for multiline malformed SSE") + } + }) + + t.Run("event field between data fragments does not flush partial data prematurely", func(t *testing.T) { + var detector StreamBootstrapDetector + fragments := [][]byte{ + []byte("data: {\"choices\":[\n"), + []byte("event: message\n"), + []byte("id: evt_999\n"), + []byte("data: ]}\n\n"), + []byte("data: [DONE]\n\n"), + } + for i, f := range fragments { + if detector.Observe(f) { + t.Fatalf("Observe(fragment %d: %q) forwarded stream with interleaved event/id fields", i, string(f)) + } + } + if !detector.Finish() { + t.Fatal("Finish() = false, want empty completion recognized when event: field is interleaved between data lines") + } + + interleavedPayload := []byte("data: {\"choices\":[\nevent: message\nid: evt_999\ndata: ]}\n\ndata: [DONE]\n\n") + if !IsEmptyCompletionPayload(interleavedPayload) { + t.Fatal("IsEmptyCompletionPayload() = false for payload with event: field between data: lines") + } + }) + + t.Run("split event metadata and data without newline", func(t *testing.T) { + var detector StreamBootstrapDetector + fragments := [][]byte{ + []byte("event: response.completed"), + []byte("data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp-1\",\"output\":[]}}"), + } + for _, f := range fragments { + detector.Observe(f) + } + // Without a newline between chunks, "event: response.completeddata: ..." is an event line whose value happens to contain "data: ...". + // Because SSE metadata is opaque and chunks without newlines are buffered as a single line, Finish() recognizes the stream as metadata-only (empty completion). + if !detector.Finish() { + t.Fatal("Finish() = false, want response.completed without newline recognized as metadata-only empty completion") + } + + // When concatenated without a newline, "event: response.completeddata: ..." is a single event header with value "response.completeddata: ...". + // Because SSE metadata values are treated as opaque, it must NOT split on the internal "data:" substring. + singlePayload := []byte("event: response.completeddata: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp-1\",\"output\":[]}}\n\n") + if !IsEmptyCompletionPayload(singlePayload) { + t.Fatal("IsEmptyCompletionPayload() = false for metadata-only payload without newline between event and data prefix") + } + }) +} + +func TestStreamBootstrapStateForwardsAtMetadataLimit(t *testing.T) { + var state streamBootstrapState + metadata := []byte("data: {\"type\":\"response.in_progress\",\"response\":{\"status\":\"in_progress\"}}\n\n") + for state.bytes+len(metadata) <= maxStreamBootstrapBytes { + if state.observe(metadata) { + t.Fatal("bootstrap forwarded recognized metadata before reaching its byte limit") + } + } + if !state.observe(metadata) { + t.Fatal("bootstrap did not conservatively forward after reaching its byte limit") + } +} + +func TestStreamBootstrapDetector(t *testing.T) { + var detector StreamBootstrapDetector + if detector.Observe([]byte("data: {\"type\":\"response.created\",\"response\":{\"status\":\"in_progress\"}}\n\n")) { + t.Fatal("StreamBootstrapDetector.Observe() = true for metadata-only prefix") + } + if !detector.Observe([]byte("data: {\"type\":\"response.output_item.done\",\"item\":{\"type\":\"custom_tool_call\",\"name\":\"shell\",\"input\":\"pwd\"}}\n\n")) { + t.Fatal("StreamBootstrapDetector.Observe() = false after complete custom tool output") + } +} + +func TestStreamBootstrapDetectorRequiresResponsesDiscriminator(t *testing.T) { + t.Run("status-only custom JSON forwards", func(t *testing.T) { + var detector StreamBootstrapDetector + if !detector.Observe([]byte(`{"status":"running"}`)) { + t.Fatal("StreamBootstrapDetector.Observe() buffered status-only custom JSON") + } + }) + + t.Run("responses object remains buffered", func(t *testing.T) { + var detector StreamBootstrapDetector + if detector.Observe([]byte(`{"object":"response","status":"in_progress"}`)) { + t.Fatal("StreamBootstrapDetector.Observe() forwarded Responses metadata") + } + }) + + t.Run("known responses event remains buffered", func(t *testing.T) { + var detector StreamBootstrapDetector + if detector.Observe([]byte(`{"type":"response.in_progress","status":"in_progress"}`)) { + t.Fatal("StreamBootstrapDetector.Observe() forwarded known Responses event") + } + }) +} + +func TestStreamBootstrapDetectorHandlesSplitSSEFrames(t *testing.T) { + t.Run("terminal empty remains buffered", func(t *testing.T) { + var detector StreamBootstrapDetector + fragments := [][]byte{ + []byte("da"), + []byte("ta: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n"), + []byte("\nda"), + []byte("ta: [DO"), + []byte("NE]\n\n"), + } + for i, fragment := range fragments { + if detector.Observe(fragment) { + t.Fatalf("Observe(fragment %d) forwarded terminal empty stream", i) + } + } + }) + + t.Run("meaningful output forwards after complete line", func(t *testing.T) { + var detector StreamBootstrapDetector + if detector.Observe([]byte("da")) { + t.Fatal("Observe() forwarded incomplete SSE prefix") + } + if detector.Observe([]byte("ta: {\"type\":\"response.output_text.delta\",\"delta\":\"hel")) { + t.Fatal("Observe() forwarded incomplete meaningful SSE line") + } + if !detector.Observe([]byte("lo\"}\n\n")) { + t.Fatal("Observe() did not forward completed meaningful SSE line") + } + }) + + t.Run("opaque payload forwards promptly", func(t *testing.T) { + var detector StreamBootstrapDetector + if !detector.Observe([]byte("opaque-provider-payload")) { + t.Fatal("Observe() buffered definitely unrecognized payload") + } + }) + + t.Run("event line waits for empty claude data", func(t *testing.T) { + var detector StreamBootstrapDetector + fragments := [][]byte{ + []byte("event: message_start\n"), + []byte("data: {\"type\":\"message_start\",\"message\":{\"type\":\"message\",\"content\":[],\"stop_reason\":null,\"usage\":{\"output_tokens\":0}}}\n\n"), + []byte("event: message_delta\n"), + []byte("data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":0}}\n\n"), + []byte("event: message_stop\n"), + []byte("data: {\"type\":\"message_stop\"}\n\n"), + } + for i, fragment := range fragments { + if detector.Observe(fragment) { + t.Fatalf("Observe(fragment %d) forwarded empty Claude stream", i) + } + } + }) + + t.Run("split comment waits for terminal empty data", func(t *testing.T) { var detector StreamBootstrapDetector fragments := [][]byte{ []byte(":"), @@ -737,3 +1468,1647 @@ func assertRotatesToContent(t *testing.T, ids []string, emptyFirst, gotPayload, t.Fatalf("content auth %q was not recorded as a success result; results=%v", other, capture.Results()) } } +func TestEmptyCompletionAudio(t *testing.T) { + cases := []struct { + name string + payload []byte + expected bool + }{ + { + name: "delta audio transcript plus data is not empty", + payload: []byte(`data: {"id":"1","choices":[{"index":0,"delta":{"audio":{"transcript":"hi","data":"AQID"}},"finish_reason":"stop"}]}` + "\n\n" + `data: [DONE]` + "\n\n"), + expected: false, + }, + { + name: "delta audio transcript only is not empty", + payload: []byte(`data: {"id":"1","choices":[{"index":0,"delta":{"audio":{"transcript":"hi"}},"finish_reason":"stop"}]}` + "\n\n" + `data: [DONE]` + "\n\n"), + expected: false, + }, + { + name: "delta audio data only is not empty", + payload: []byte(`data: {"id":"1","choices":[{"index":0,"delta":{"audio":{"data":"AQID"}},"finish_reason":"stop"}]}` + "\n\n" + `data: [DONE]` + "\n\n"), + expected: false, + }, + { + name: "message audio non-stream is not empty", + payload: []byte(`{"id":"1","choices":[{"index":0,"message":{"role":"assistant","content":"","audio":{"transcript":"hi","data":"AQID"}},"finish_reason":"stop"}],"usage":{"completion_tokens":0}}`), + expected: false, + }, + { + name: "delta audio null stays empty", + payload: []byte(`data: {"id":"1","choices":[{"index":0,"delta":{"audio":null},"finish_reason":"stop"}]}` + "\n\n" + `data: [DONE]` + "\n\n"), + expected: true, + }, + { + name: "delta audio empty object stays empty", + payload: []byte(`data: {"id":"1","choices":[{"index":0,"delta":{"audio":{}},"finish_reason":"stop"}]}` + "\n\n" + `data: [DONE]` + "\n\n"), + expected: true, + }, + { + name: "delta audio empty fields stay empty", + payload: []byte(`data: {"id":"1","choices":[{"index":0,"delta":{"audio":{"transcript":"","data":""}},"finish_reason":"stop"}]}` + "\n\n" + `data: [DONE]` + "\n\n"), + expected: true, + }, + { + name: "message audio recursively empty stays empty", + payload: []byte(`{"id":"1","choices":[{"index":0,"message":{"audio":{"transcript":" ","nested":{"items":[null,false,0,"",{},[]]}}},"finish_reason":"stop"}]}`), + expected: true, + }, + { + name: "delta audio id only is not empty", + payload: []byte(`data: {"id":"1","choices":[{"index":0,"delta":{"audio":{"id":"audio-1"}},"finish_reason":"stop"}]}` + "\n\n" + `data: [DONE]` + "\n\n"), + expected: false, + }, + { + name: "delta audio positive expires at is not empty", + payload: []byte(`data: {"id":"1","choices":[{"index":0,"delta":{"audio":{"expires_at":1}},"finish_reason":"stop"}]}` + "\n\n" + `data: [DONE]` + "\n\n"), + expected: false, + }, + { + name: "delta audio malformed frame fails safe as non-empty", + payload: []byte(`data: {"id":"1","choices":[{"index":0,"delta":{"audio":"unterminated},"finish_reason":"stop"}]}` + "\n\n" + `data: [DONE]` + "\n\n"), + expected: false, + }, + { + name: "delta audio malformed with text stays not empty", + payload: []byte(`data: {"id":"1","choices":[{"index":0,"delta":{"content":"text","audio":"unterminated},"finish_reason":"stop"}]}` + "\n\n" + `data: [DONE]` + "\n\n"), + expected: false, + }, + { + name: "audio with malformed usage stays not empty", + payload: []byte(`data: {"id":"1","choices":[{"index":0,"delta":{"audio":{"transcript":"hi"}},"finish_reason":"stop"}],"usage":{"completion_tokens":"abc"}}` + "\n\n" + `data: [DONE]` + "\n\n"), + expected: false, + }, + { + name: "raw json audio frame is not empty", + payload: []byte(`{"id":"1","choices":[{"index":0,"delta":{"audio":{"transcript":"hi","data":"AQID"}},"finish_reason":"stop"}]}`), + expected: false, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := isEmptyCompletionPayload(tc.payload); got != tc.expected { + t.Fatalf("isEmptyCompletionPayload() = %v, want %v", got, tc.expected) + } + }) + } +} + +// TestEmptyCompletionMeaningfulFields covers the targeted meaningful-content +// fields: Gemini executableCode/codeExecutionResult parts and OpenAI legacy +// message.function_call (and its streaming delta.function_call form). A value +// is meaningful only when it carries actual payload; null, empty string, empty +// object, and empty array stay empty. +func TestEmptyCompletionMeaningfulFields(t *testing.T) { + cases := []struct { + name string + payload []byte + expected bool // true = empty completion + }{ + { + name: "gemini executableCode with payload is not empty", + payload: []byte(`data: {"candidates":[{"content":{"role":"model","parts":[{"executableCode":{"language":"python","code":"print(1)"}}]},"finishReason":"STOP"}]}`), + expected: false, + }, + { + name: "gemini executableCode null stays empty", + payload: []byte(`data: {"candidates":[{"content":{"role":"model","parts":[{"executableCode":null}]},"finishReason":"STOP"}]}`), + expected: true, + }, + { + name: "gemini executableCode empty object stays empty", + payload: []byte(`data: {"candidates":[{"content":{"role":"model","parts":[{"executableCode":{}}]},"finishReason":"STOP"}]}`), + expected: true, + }, + { + name: "gemini executableCode whitespace object stays empty", + payload: []byte(`data: {"candidates":[{"content":{"role":"model","parts":[{"executableCode":{ }}]},"finishReason":"STOP"}]}`), + expected: true, + }, + { + name: "gemini codeExecutionResult with payload is not empty", + payload: []byte(`data: {"candidates":[{"content":{"role":"model","parts":[{"codeExecutionResult":{"outcome":"OK","output":"1"}}]},"finishReason":"STOP"}]}`), + expected: false, + }, + { + name: "gemini codeExecutionResult null stays empty", + payload: []byte(`data: {"candidates":[{"content":{"role":"model","parts":[{"codeExecutionResult":null}]},"finishReason":"STOP"}]}`), + expected: true, + }, + { + name: "gemini codeExecutionResult empty object stays empty", + payload: []byte(`data: {"candidates":[{"content":{"role":"model","parts":[{"codeExecutionResult":{}}]},"finishReason":"STOP"}]}`), + expected: true, + }, + { + name: "gemini codeExecutionResult whitespace object stays empty", + payload: []byte(`data: {"candidates":[{"content":{"role":"model","parts":[{"codeExecutionResult":{ }}]},"finishReason":"STOP"}]}`), + expected: true, + }, + { + name: "openai non-stream message function_call name only is not empty", + payload: []byte(`{"choices":[{"message":{"function_call":{"name":"get_weather","arguments":""}},"finish_reason":"stop"}]}`), + expected: false, + }, + { + name: "openai non-stream message function_call arguments only is not empty", + payload: []byte(`{"choices":[{"message":{"function_call":{"name":"","arguments":"{\"city\":\"x\"}"}},"finish_reason":"stop"}]}`), + expected: false, + }, + { + name: "openai non-stream message function_call empty object stays empty", + payload: []byte(`{"choices":[{"message":{"function_call":{}},"finish_reason":"stop"}]}`), + expected: true, + }, + { + name: "openai non-stream message function_call null stays empty", + payload: []byte(`{"choices":[{"message":{"function_call":null},"finish_reason":"stop"}]}`), + expected: true, + }, + { + name: "openai non-stream message function_call whitespace fields stay empty", + payload: []byte(`{"choices":[{"message":{"function_call":{"name":" ","arguments":" "}},"finish_reason":"stop"}]}`), + expected: true, + }, + { + name: "openai sse delta function_call is not empty", + payload: []byte("data: {\"choices\":[{\"delta\":{\"function_call\":{\"name\":\"get_weather\",\"arguments\":\"{}\"}},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n"), + expected: false, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got := IsEmptyCompletionPayload(tc.payload) + if got != tc.expected { + t.Fatalf("IsEmptyCompletionPayload = %v, want %v\npayload: %s", got, tc.expected, tc.payload) + } + }) + } +} + +// TestEmptyCompletionFraming covers aggregated raw JSON payloads that carry one +// or more top-level values (NDJSON, whitespace/concat, pretty) evaluated through +// the protocol evaluators. Malformed or trailing garbage must stay non-empty +// (safe to forward). +func TestEmptyCompletionFraming(t *testing.T) { + cases := []struct { + name string + payload []byte + expected bool // true = empty completion + }{ + { + name: "ndjson second frame meaningful", + payload: []byte("{\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":{\"completion_tokens\":0}}\n{\"choices\":[{\"delta\":{\"content\":\"hi\"},\"finish_reason\":\"stop\"}],\"usage\":{\"completion_tokens\":5}}"), + expected: false, + }, + { + name: "concatenated second frame meaningful", + payload: []byte("{\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":{\"completion_tokens\":0}}{\"choices\":[{\"delta\":{\"content\":\"hi\"},\"finish_reason\":\"stop\"}],\"usage\":{\"completion_tokens\":5}}"), + expected: false, + }, + { + name: "ndjson all empty terminal", + payload: []byte("{\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":{\"completion_tokens\":0}}\n{\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":{\"completion_tokens\":0}}"), + expected: true, + }, + { + name: "pretty multiline meaningful", + payload: []byte("{\n \"choices\": [\n {\"delta\": {\"content\": \"hi\"}, \"finish_reason\": \"stop\"}\n ],\n \"usage\": {\"completion_tokens\": 5}\n}"), + expected: false, + }, + { + name: "ndjson trailing garbage not empty", + payload: []byte("{\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":{\"completion_tokens\":0}}\nnot-json"), + expected: false, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got := IsEmptyCompletionPayload(tc.payload) + if got != tc.expected { + t.Fatalf("IsEmptyCompletionPayload = %v, want %v\npayload: %s", got, tc.expected, tc.payload) + } + }) + } +} + +// TestStreamBootstrapDetectorRawJSON verifies a single raw JSON value split at +// every byte boundary (including inside string escapes and multi-byte UTF-8) +// never forwards prematurely and forwards promptly once complete. +func TestStreamBootstrapDetectorRawJSON(t *testing.T) { + raws := []struct { + name string + raw []byte + }{ + {"ascii", []byte(`{"choices":[{"delta":{"content":"hi"},"finish_reason":"stop"}]}`)}, + {"string-escape", []byte(`{"choices":[{"delta":{"content":"a\nb"},"finish_reason":"stop"}]}`)}, + {"utf8", []byte(`{"choices":[{"delta":{"content":"😀"},"finish_reason":"stop"}]}`)}, + } + for _, r := range raws { + for i := 0; i <= len(r.raw); i++ { + d := &StreamBootstrapDetector{} + first := d.Observe(r.raw[:i]) + second := d.Observe(r.raw[i:]) + if i == len(r.raw) { + if !first { + t.Fatalf("%s split at %d: expected forward after full value, got %v", r.name, i, first) + } + continue + } + if first { + t.Fatalf("%s split at %d: premature forward on prefix %q", r.name, i, r.raw[:i]) + } + if !second { + t.Fatalf("%s split at %d: expected forward after completion, got %v", r.name, i, second) + } + } + } +} + +// TestStreamBootstrapDetectorRawConcatenated verifies two raw JSON frames +// delivered as concatenated values (no newline); the detector must not forward +// on the empty first frame and must forward once the meaningful second lands. +func TestStreamBootstrapDetectorRawConcatenated(t *testing.T) { + d := &StreamBootstrapDetector{} + if got := d.Observe([]byte(`{"choices":[{"delta":{},"finish_reason":"stop"}],"usage":{"completion_tokens":0}}`)); got != false { + t.Fatalf("Observe(first empty frame) = %v, want false", got) + } + if got := d.Observe([]byte(`{"choices":[{"delta":{"content":"hi"},"finish_reason":"stop"}],"usage":{"completion_tokens":5}}`)); got != true { + t.Fatalf("Observe(second meaningful frame) = %v, want true", got) + } +} + +// TestStreamBootstrapDetectorRawSSEPrefixes verifies incomplete SSE command +// prefixes (d/da/data/data:/: and a split [DONE]) keep buffering, preserving +// the current SSE bootstrap contract. +func TestStreamBootstrapDetectorRawSSEPrefixes(t *testing.T) { + for _, p := range [][]byte{[]byte("d"), []byte("da"), []byte("data"), []byte("data:"), []byte(":")} { + d := &StreamBootstrapDetector{} + if got := d.Observe(p); got != false { + t.Fatalf("Observe(%q) = %v, want false (incomplete SSE prefix)", p, got) + } + } + d := &StreamBootstrapDetector{} + if got := d.Observe([]byte("data: [DO")); got != false { + t.Fatalf("Observe(split [DONE]) = %v, want false", got) + } + if got := d.Observe([]byte("NE]\n\n")); got != false { + t.Fatalf("Observe(completed [DONE]) = %v, want false (empty terminal stays buffered)", got) + } +} + +func TestStreamBootstrapDetectorNewlineLessSSE(t *testing.T) { + t.Run("complete newline-less content frame forwards immediately", func(t *testing.T) { + d := &StreamBootstrapDetector{} + payload := []byte(`data: {"choices":[{"delta":{"content":"hello"}}],"finish_reason":null}`) + if got := d.Observe(payload); got != true { + t.Fatalf("Observe(newline-less content) = %v, want true", got) + } + if !d.state.forward { + t.Fatal("state.forward = false, want true") + } + }) + + t.Run("complete newline-less empty terminal frame stays buffered", func(t *testing.T) { + d := &StreamBootstrapDetector{} + payload := []byte(`data: {"choices":[{"delta":{},"finish_reason":"stop"}]}`) + if got := d.Observe(payload); got != false { + t.Fatalf("Observe(newline-less empty terminal) = %v, want false", got) + } + if d.state.forward { + t.Fatal("state.forward = true, want false") + } + if !d.state.acc.empty() { + t.Fatal("acc.empty() = false, want true for empty terminal frame") + } + }) + + t.Run("following newline-less [DONE] remains terminal-empty", func(t *testing.T) { + d := &StreamBootstrapDetector{} + emptyFrame := []byte("data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n") + if got := d.Observe(emptyFrame); got != false { + t.Fatalf("Observe(empty frame) = %v, want false", got) + } + doneFrame := []byte("data: [DONE]") + if got := d.Observe(doneFrame); got != false { + t.Fatalf("Observe(newline-less [DONE]) = %v, want false", got) + } + if d.state.forward { + t.Fatal("state.forward = true, want false") + } + if !d.state.acc.empty() { + t.Fatal("acc.empty() = false, want true after [DONE]") + } + }) + + t.Run("split truncated JSON and split [DONE] do not forward prematurely", func(t *testing.T) { + d := &StreamBootstrapDetector{} + if got := d.Observe([]byte(`data: {"choices":[{"delta":{"content":"hel`)); got != false { + t.Fatalf("Observe(truncated JSON) = %v, want false", got) + } + if got := d.Observe([]byte(`lo"}}],"finish_reason":null}`)); got != true { + t.Fatalf("Observe(completed JSON remainder) = %v, want true", got) + } + + d2 := &StreamBootstrapDetector{} + if got := d2.Observe([]byte("data: [DO")); got != false { + t.Fatalf("Observe(split [DONE] part 1) = %v, want false", got) + } + if got := d2.Observe([]byte("NE]")); got != false { + t.Fatalf("Observe(split [DONE] part 2) = %v, want false", got) + } + if d2.state.forward { + t.Fatal("state.forward after split [DONE] = true, want false") + } + if !d2.state.acc.empty() { + t.Fatal("acc.empty() = false, want true after complete [DONE]") + } + }) +} + +func TestEmptyCompletion_MultiChunkBoundarySafety(t *testing.T) { + t.Run("two complete newline-less data chunks plus terminal DONE classify empty", func(t *testing.T) { + chunks := []cliproxyexecutor.StreamChunk{ + {Payload: []byte("data: {\"id\":\"1\",\"choices\":[{\"delta\":{\"content\":\"\"}}]}")}, + {Payload: []byte("data: [DONE]")}, + } + if !isEmptyCompletion(chunks) { + t.Fatal("isEmptyCompletion = false, want true") + } + }) + + t.Run("split JSON string fragments concatenate and classify non-empty", func(t *testing.T) { + chunks := []cliproxyexecutor.StreamChunk{ + {Payload: []byte("data: {\"id\":\"1\",\"choices\":[{\"delta\":{\"content\":\"hello ")}, + {Payload: []byte("world\"}}]}\n")}, + } + if isEmptyCompletion(chunks) { + t.Fatal("isEmptyCompletion = true, want false") + } + }) + + t.Run("boundary before nested object remains valid", func(t *testing.T) { + chunks := []cliproxyexecutor.StreamChunk{ + {Payload: []byte("data: ")}, + {Payload: []byte("{\"id\":\"1\",\"choices\":[{\"delta\":{\"content\":\"\"}}]}")}, + {Payload: []byte("\ndata: [DONE]\n")}, + } + if !isEmptyCompletion(chunks) { + t.Fatal("isEmptyCompletion = false, want true") + } + }) + + t.Run("unknown custom stream remains unrecognized and non-empty", func(t *testing.T) { + chunks := []cliproxyexecutor.StreamChunk{ + {Payload: []byte("custom_binary_payload_format")}, + } + if isEmptyCompletion(chunks) { + t.Fatal("isEmptyCompletion = true, want false for unrecognized stream") + } + }) + + t.Run("detector finish flushes pending at EOF", func(t *testing.T) { + var detector StreamBootstrapDetector + if detector.Observe([]byte("data: {\"id\":\"1\",\"choices\":[{\"delta\":{\"content\":\"\"}}]}")) { + t.Fatal("Observe() = true, want false") + } + if detector.Observe([]byte("data: [DONE]")) { + t.Fatal("Observe() = true, want false") + } + if !detector.Finish() { + t.Fatal("detector.Finish() = false, want true") + } + }) +} + +func TestClaudeToolBlocksEmptyCompletion(t *testing.T) { + t.Run("empty tool block without id name or input is empty completion", func(t *testing.T) { + payload := []byte(`{"type":"message","role":"assistant","content":[{"type":"tool_use","input":null}],"stop_reason":"end_turn"}`) + if !IsEmptyCompletionPayload(payload) { + t.Fatal("IsEmptyCompletionPayload() = false for tool_use with null input and no name/id, want true") + } + }) + + t.Run("empty tool block with empty input object and no id/name is empty completion", func(t *testing.T) { + payload := []byte(`{"type":"message","role":"assistant","content":[{"type":"tool_use","input":{}}],"stop_reason":"end_turn"}`) + if !IsEmptyCompletionPayload(payload) { + t.Fatal("IsEmptyCompletionPayload() = false for tool_use with empty input and no name/id, want true") + } + }) + + t.Run("tool block with valid name is recognized as tool call", func(t *testing.T) { + payload := []byte(`{"type":"message","role":"assistant","content":[{"type":"tool_use","id":"toolu_1","name":"get_weather","input":{}}],"stop_reason":"tool_use"}`) + if IsEmptyCompletionPayload(payload) { + t.Fatal("IsEmptyCompletionPayload() = true for valid tool_use with name/id, want false") + } + }) + + t.Run("text block with lexical null input does not treat null as tool call", func(t *testing.T) { + payload := []byte(`{"type":"message","role":"assistant","content":[{"type":"text","text":"","input":null}],"stop_reason":"end_turn"}`) + if !IsEmptyCompletionPayload(payload) { + t.Fatal("IsEmptyCompletionPayload() = false for text block with lexical null input, want true") + } + }) +} + +func TestClaudeToolUseStopReasonEmptyCompletion(t *testing.T) { + t.Run("empty tool_use blocks with stop_reason tool_use is empty completion", func(t *testing.T) { + payload := []byte(`{"type":"message","role":"assistant","content":[{"type":"tool_use","input":null}],"stop_reason":"tool_use"}`) + if !IsEmptyCompletionPayload(payload) { + t.Fatal("IsEmptyCompletionPayload() = false for empty tool_use block with stop_reason tool_use, want true") + } + }) + + t.Run("empty tool_use blocks in sse stream with stop_reason tool_use is empty completion", func(t *testing.T) { + payload := []byte("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"model\":\"claude-3\",\"usage\":{\"output_tokens\":0}}}\n\nevent: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"tool_use\",\"input\":null}}\n\nevent: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"tool_use\"},\"usage\":{\"output_tokens\":0}}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n") + if !IsEmptyCompletionPayload(payload) { + t.Fatal("IsEmptyCompletionPayload() = false for stream with empty tool_use and stop_reason tool_use, want true") + } + }) + + t.Run("control real tool_use with stop_reason tool_use is not empty", func(t *testing.T) { + payload := []byte(`{"type":"message","role":"assistant","content":[{"type":"tool_use","id":"toolu_1","name":"get_weather","input":{"location":"San Francisco"}}],"stop_reason":"tool_use"}`) + if IsEmptyCompletionPayload(payload) { + t.Fatal("IsEmptyCompletionPayload() = true for real tool_use with stop_reason tool_use, want false") + } + }) + + t.Run("claude mcp_tool_use with id and name is not empty completion", func(t *testing.T) { + payload := []byte(`{"type":"message","role":"assistant","content":[{"type":"mcp_tool_use","id":"mcp_1","name":"server__tool","input":{}}],"stop_reason":"tool_use"}`) + if IsEmptyCompletionPayload(payload) { + t.Fatal("IsEmptyCompletionPayload() = true for mcp_tool_use with id and name, want false") + } + }) + + t.Run("claude mcp_tool_use in sse stream with id and name is not empty completion", func(t *testing.T) { + payload := []byte("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"model\":\"claude-3\",\"usage\":{\"output_tokens\":0}}}\n\nevent: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"mcp_tool_use\",\"id\":\"mcp_1\",\"name\":\"server__tool\"}}\n\nevent: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"tool_use\"},\"usage\":{\"output_tokens\":0}}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n") + if IsEmptyCompletionPayload(payload) { + t.Fatal("IsEmptyCompletionPayload() = true for mcp_tool_use stream with id and name, want false") + } + }) + + t.Run("claude mcp_tool_use missing id is empty completion", func(t *testing.T) { + payload := []byte(`{"type":"message","role":"assistant","content":[{"type":"mcp_tool_use","id":"","name":"server__tool","input":{}}],"stop_reason":"tool_use"}`) + if !IsEmptyCompletionPayload(payload) { + t.Fatal("IsEmptyCompletionPayload() = false for mcp_tool_use missing id, want true") + } + }) + + t.Run("claude mcp_tool_use missing name is empty completion", func(t *testing.T) { + payload := []byte(`{"type":"message","role":"assistant","content":[{"type":"mcp_tool_use","id":"mcp_1","name":"","input":{}}],"stop_reason":"tool_use"}`) + if !IsEmptyCompletionPayload(payload) { + t.Fatal("IsEmptyCompletionPayload() = false for mcp_tool_use missing name, want true") + } + }) + + t.Run("control stop_reason max_tokens without content is blocked and not empty completion", func(t *testing.T) { + payload := []byte(`{"type":"message","role":"assistant","content":[],"stop_reason":"max_tokens"}`) + if IsEmptyCompletionPayload(payload) { + t.Fatal("IsEmptyCompletionPayload() = true for stop_reason max_tokens, want false (blocked)") + } + }) + + t.Run("control stop_reason refusal without content is blocked and not empty completion", func(t *testing.T) { + payload := []byte(`{"type":"message","role":"assistant","content":[],"stop_reason":"refusal"}`) + if IsEmptyCompletionPayload(payload) { + t.Fatal("IsEmptyCompletionPayload() = true for stop_reason refusal, want false (blocked)") + } + }) +} + +func TestPrettyPrintedJSONWithDataSubstringEmptyCompletion(t *testing.T) { + t.Run("pretty-printed json with data substring is evaluated as empty completion", func(t *testing.T) { + payload := []byte("{\n \"id\": \"msg-data:123\",\n \"choices\": [\n {\n \"message\": {\n \"role\": \"assistant\",\n \"content\": \"\"\n },\n \"finish_reason\": \"stop\"\n }\n ]\n}") + if !IsEmptyCompletionPayload(payload) { + t.Fatal("IsEmptyCompletionPayload() = false for pretty-printed JSON with data: substring, want true") + } + }) +} + +func TestClaudeInputJSONDeltaEmptyCompletion(t *testing.T) { + t.Run("empty input_json_delta with empty partial_json and no preceding tool id/name is empty completion", func(t *testing.T) { + payload := []byte("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"model\":\"claude-3\",\"usage\":{\"output_tokens\":0}}}\n\nevent: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"tool_use\",\"input\":null}}\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"\"}}\n\nevent: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"tool_use\"},\"usage\":{\"output_tokens\":0}}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n") + if !IsEmptyCompletionPayload(payload) { + t.Fatal("IsEmptyCompletionPayload() = false for stream with empty input_json_delta, want true") + } + }) + + t.Run("meaningful input_json_delta sets tool calls", func(t *testing.T) { + payload := []byte("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"model\":\"claude-3\",\"usage\":{\"output_tokens\":0}}}\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"{\\\"location\\\":\\\"SF\\\"}\"}}\n\nevent: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"tool_use\"},\"usage\":{\"output_tokens\":0}}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n") + if IsEmptyCompletionPayload(payload) { + t.Fatal("IsEmptyCompletionPayload() = true for stream with meaningful input_json_delta, want false") + } + }) +} + +func TestClaudeEmptyThinkingBlockStartEmptyCompletion(t *testing.T) { + t.Run("empty thinking content_block_start followed by message_stop is empty completion", func(t *testing.T) { + payload := []byte("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"model\":\"claude-3\",\"usage\":{\"output_tokens\":0}}}\n\nevent: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"thinking\",\"thinking\":\"\"}}\n\nevent: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":0}}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n") + if !IsEmptyCompletionPayload(payload) { + t.Fatal("IsEmptyCompletionPayload() = false for empty thinking block start with message_stop, want true") + } + }) + + t.Run("thinking content_block_start followed by thinking_delta with text is not empty", func(t *testing.T) { + payload := []byte("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"model\":\"claude-3\",\"usage\":{\"output_tokens\":0}}}\n\nevent: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"thinking\",\"thinking\":\"\"}}\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"thinking_delta\",\"thinking\":\"thinking step\"}}\n\nevent: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":0}}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n") + if IsEmptyCompletionPayload(payload) { + t.Fatal("IsEmptyCompletionPayload() = true for thinking block with thinking_delta text, want false") + } + }) + + t.Run("empty redacted_thinking content_block_start followed by message_stop is empty completion", func(t *testing.T) { + payload := []byte("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"model\":\"claude-3\",\"usage\":{\"output_tokens\":0}}}\n\nevent: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"redacted_thinking\",\"data\":\"\"}}\n\nevent: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":0}}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n") + if !IsEmptyCompletionPayload(payload) { + t.Fatal("IsEmptyCompletionPayload() = false for empty redacted_thinking block start with message_stop, want true") + } + }) + + t.Run("non-empty redacted_thinking content_block_start followed by message_stop is not empty", func(t *testing.T) { + payload := []byte("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"model\":\"claude-3\",\"usage\":{\"output_tokens\":0}}}\n\nevent: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"redacted_thinking\",\"data\":\"abc123encryptedpayload\"}}\n\nevent: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":0}}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n") + if IsEmptyCompletionPayload(payload) { + t.Fatal("IsEmptyCompletionPayload() = true for non-empty redacted_thinking block, want false") + } + }) +} + +func TestRecognizedContentlessEOFEmptyStream(t *testing.T) { + t.Run("OpenAI role-only delta stream closed at EOF without [DONE] is empty completion", func(t *testing.T) { + var detector StreamBootstrapDetector + payload := []byte("data: {\"choices\":[{\"delta\":{\"role\":\"assistant\"},\"finish_reason\":null}]}\n\n") + if detector.Observe(payload) { + t.Fatal("Observe() = true for role-only delta, want false") + } + if !detector.Finish() { + t.Fatal("Finish() = false at EOF for recognized role-only stream without content, want true") + } + }) + + t.Run("Claude message_start stream closed at EOF without message_stop is empty completion", func(t *testing.T) { + var detector StreamBootstrapDetector + payload := []byte("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"model\":\"claude-3\",\"usage\":{\"output_tokens\":0}}}\n\n") + if detector.Observe(payload) { + t.Fatal("Observe() = true for message_start, want false") + } + if !detector.Finish() { + t.Fatal("Finish() = false at EOF for recognized message_start stream without content, want true") + } + }) + + t.Run("OpenAI delta stream with content closed at EOF is not empty", func(t *testing.T) { + var detector StreamBootstrapDetector + payload := []byte("data: {\"choices\":[{\"delta\":{\"content\":\"hello\"}}]}\n\n") + if !detector.Observe(payload) { + t.Fatal("Observe() = false for stream with content, want true") + } + if detector.Finish() { + t.Fatal("Finish() = true for stream with content, want false") + } + }) + + t.Run("unknown-format stream closed at EOF remains non-empty and forwards", func(t *testing.T) { + var detector StreamBootstrapDetector + payload := []byte("data: {\"unknown_payload\":true}\n\n") + if !detector.Observe(payload) { + t.Fatal("Observe() = false for unknown-format, want true (force forward)") + } + if detector.Finish() { + t.Fatal("Finish() = true for unknown-format, want false") + } + }) +} + +func TestColonlessSSEFields(t *testing.T) { + t.Run("stream with colonless event and id fields then empty data is recognized as empty completion", func(t *testing.T) { + var detector StreamBootstrapDetector + payload := []byte("event\nid\nretry\ndata: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n") + if detector.Observe(payload) { + t.Fatal("Observe() = true for stream with colonless metadata lines, want false") + } + if !detector.Finish() { + t.Fatal("Finish() = false, want empty completion recognized for colonless metadata") + } + if detector.state.acc.sawUnknownData { + t.Fatal("sawUnknownData = true for colonless metadata lines, want false") + } + if !IsEmptyCompletionPayload(payload) { + t.Fatal("IsEmptyCompletionPayload() = false for colonless metadata lines") + } + }) + + t.Run("colonless data field treated as empty data event", func(t *testing.T) { + var detector StreamBootstrapDetector + payload := []byte("data\n\ndata: [DONE]\n\n") + if detector.Observe(payload) { + t.Fatal("Observe() = true for colonless data event, want false") + } + if !detector.Finish() { + t.Fatal("Finish() = false, want empty completion recognized for colonless data event") + } + if detector.state.acc.sawUnknownData { + t.Fatal("sawUnknownData = true for colonless data, want false") + } + if !IsEmptyCompletionPayload(payload) { + t.Fatal("IsEmptyCompletionPayload() = false for colonless data event payload") + } + }) + + t.Run("couldBeSSEPrefix recognizes colonless prefixes", func(t *testing.T) { + for _, prefix := range []string{"data", "event", "id", "retry"} { + if !couldBeSSEPrefix([]byte(prefix)) { + t.Fatalf("couldBeSSEPrefix(%q) = false, want true", prefix) + } + } + }) +} + +func TestStreamBootstrapDetectorMeaningfulOutput(t *testing.T) { + t.Run("openai role-only delta is not meaningful output", func(t *testing.T) { + var detector StreamBootstrapDetector + if detector.Observe([]byte("data: {\"choices\":[{\"delta\":{\"role\":\"assistant\"}}]}\n\n")) { + t.Fatal("Observe() = true for role-only delta") + } + if detector.HasMeaningfulOutput() { + t.Fatal("HasMeaningfulOutput() = true for role-only delta") + } + }) + + t.Run("claude message_start is not meaningful output", func(t *testing.T) { + var detector StreamBootstrapDetector + if detector.Observe([]byte("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"role\":\"assistant\"}}\n\n")) { + t.Fatal("Observe() = true for message_start") + } + if detector.HasMeaningfulOutput() { + t.Fatal("HasMeaningfulOutput() = true for message_start") + } + }) + + t.Run("responses response.created is not meaningful output", func(t *testing.T) { + var detector StreamBootstrapDetector + if detector.Observe([]byte(`{"type":"response.created","response":{"id":"r1"}}`)) { + t.Fatal("Observe() = true for response.created") + } + if detector.HasMeaningfulOutput() { + t.Fatal("HasMeaningfulOutput() = true for response.created") + } + }) + + t.Run("sse ping comment is not meaningful output", func(t *testing.T) { + var detector StreamBootstrapDetector + if detector.Observe([]byte(": ping\n\n")) { + t.Fatal("Observe() = true for SSE ping comment") + } + if detector.HasMeaningfulOutput() { + t.Fatal("HasMeaningfulOutput() = true for SSE ping comment") + } + }) + + t.Run("openai content delta is meaningful output", func(t *testing.T) { + var detector StreamBootstrapDetector + if !detector.Observe([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"hello\"}}]}\n\n")) { + t.Fatal("Observe() = false for content delta") + } + if !detector.HasMeaningfulOutput() { + t.Fatal("HasMeaningfulOutput() = false for content delta") + } + }) + + t.Run("openai tool call is meaningful output", func(t *testing.T) { + var detector StreamBootstrapDetector + if !detector.Observe([]byte("data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"id\":\"call_1\",\"type\":\"function\",\"function\":{\"name\":\"search\"}}]}}]}\n\n")) { + t.Fatal("Observe() = false for tool_calls") + } + if !detector.HasMeaningfulOutput() { + t.Fatal("HasMeaningfulOutput() = false for tool_calls") + } + }) + + t.Run("content filter block is meaningful output", func(t *testing.T) { + var detector StreamBootstrapDetector + if !detector.Observe([]byte("data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"content_filter\"}]}\n\n")) { + t.Fatal("Observe() = false for content_filter finish_reason") + } + if !detector.HasMeaningfulOutput() { + t.Fatal("HasMeaningfulOutput() = false for content_filter finish_reason") + } + }) +} + +func TestReadStreamBootstrapErrorHandling(t *testing.T) { + errUpstream := errors.New("upstream failed") + + t.Run("error following openai role delta propagates as failover error", func(t *testing.T) { + ch := make(chan cliproxyexecutor.StreamChunk, 2) + ch <- cliproxyexecutor.StreamChunk{Payload: []byte("data: {\"choices\":[{\"delta\":{\"role\":\"assistant\"}}]}\n\n")} + ch <- cliproxyexecutor.StreamChunk{Err: errUpstream} + close(ch) + + buffered, closed, err := readStreamBootstrap(context.Background(), ch) + if err == nil { + t.Fatal("readStreamBootstrap error = nil, want errUpstream propagated for failover") + } + if !errors.Is(err, errUpstream) { + t.Fatalf("readStreamBootstrap error = %v, want %v", err, errUpstream) + } + if len(buffered) != 0 { + t.Fatalf("len(buffered) = %d, want 0 when error propagates", len(buffered)) + } + if closed { + t.Fatal("closed = true, want false") + } + }) + + t.Run("error following claude message_start propagates as failover error", func(t *testing.T) { + ch := make(chan cliproxyexecutor.StreamChunk, 2) + ch <- cliproxyexecutor.StreamChunk{Payload: []byte("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"m1\",\"role\":\"assistant\"}}\n\n")} + ch <- cliproxyexecutor.StreamChunk{Err: errUpstream} + close(ch) + + buffered, _, err := readStreamBootstrap(context.Background(), ch) + if err == nil { + t.Fatal("readStreamBootstrap error = nil, want errUpstream propagated") + } + if !errors.Is(err, errUpstream) { + t.Fatalf("readStreamBootstrap error = %v, want %v", err, errUpstream) + } + if len(buffered) != 0 { + t.Fatalf("len(buffered) = %d, want 0", len(buffered)) + } + }) + + t.Run("error following zero payload chunk propagates as failover error", func(t *testing.T) { + ch := make(chan cliproxyexecutor.StreamChunk, 2) + ch <- cliproxyexecutor.StreamChunk{Payload: nil} + ch <- cliproxyexecutor.StreamChunk{Err: errUpstream} + close(ch) + + buffered, _, err := readStreamBootstrap(context.Background(), ch) + if err == nil { + t.Fatal("readStreamBootstrap error = nil, want errUpstream propagated") + } + if !errors.Is(err, errUpstream) { + t.Fatalf("readStreamBootstrap error = %v, want %v", err, errUpstream) + } + if len(buffered) != 0 { + t.Fatalf("len(buffered) = %d, want 0", len(buffered)) + } + }) + + t.Run("error following responses created event propagates as failover error", func(t *testing.T) { + ch := make(chan cliproxyexecutor.StreamChunk, 2) + ch <- cliproxyexecutor.StreamChunk{Payload: []byte(`{"type":"response.created","response":{"id":"r1"}}`)} + ch <- cliproxyexecutor.StreamChunk{Err: errUpstream} + close(ch) + + buffered, _, err := readStreamBootstrap(context.Background(), ch) + if err == nil { + t.Fatal("readStreamBootstrap error = nil, want errUpstream propagated") + } + if !errors.Is(err, errUpstream) { + t.Fatalf("readStreamBootstrap error = %v, want %v", err, errUpstream) + } + if len(buffered) != 0 { + t.Fatalf("len(buffered) = %d, want 0", len(buffered)) + } + }) + + t.Run("meaningful content starts stream immediately", func(t *testing.T) { + ch := make(chan cliproxyexecutor.StreamChunk, 2) + ch <- cliproxyexecutor.StreamChunk{Payload: []byte("data: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\n")} + ch <- cliproxyexecutor.StreamChunk{Payload: []byte("data: [DONE]\n\n")} + close(ch) + + buffered, closed, err := readStreamBootstrap(context.Background(), ch) + if err != nil { + t.Fatalf("readStreamBootstrap error = %v, want nil", err) + } + if len(buffered) != 1 { + t.Fatalf("len(buffered) = %d, want 1", len(buffered)) + } + if closed { + t.Fatal("closed = true, want false (started stream)") + } + }) +} + +func TestExecuteLegacyOpenAICompletionNotRotated(t *testing.T) { + executor := &emptyCompletionTestExecutor{ + executePayloads: map[string][]byte{}, + executeCalls: map[string]int{}, + } + manager, ids, model, _ := newEmptyCompletionTestManager(t, executor) + + legacyPayload := []byte(`{"choices":[{"text":"hello legacy completion","finish_reason":"stop"}]}`) + executor.executePayloads[ids[0]] = legacyPayload + executor.executePayloads[ids[1]] = legacyPayload + + resp, err := manager.Execute(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}) + if err != nil { + t.Fatalf("Execute() error = %v", err) + } + if !strings.Contains(string(resp.Payload), "hello legacy completion") { + t.Fatalf("resp payload = %q, want legacy completion text", string(resp.Payload)) + } + if auth, ok := manager.GetByID(ids[0]); ok && auth != nil { + if auth.Unavailable || !auth.NextRetryAfter.IsZero() { + t.Fatalf("auth %q was cooled despite returning legacy completion text", ids[0]) + } + } +} + +func TestStreamBootstrapDetectorLegacyOpenAI(t *testing.T) { + d := &StreamBootstrapDetector{} + if got := d.Observe([]byte("data: {\"choices\":[{\"text\":\"hello\",\"finish_reason\":null}]}\n\n")); got != true { + t.Fatalf("Observe(legacy choices.text chunk) = %v, want true (forwarded immediately)", got) + } + if !d.state.forward { + t.Fatal("state.forward = false, want true") + } + if d.state.isEmptyCompletion() { + t.Fatal("isEmptyCompletion = true, want false") + } +} + +func TestClaudeSignatureDeltaEmptyCompletion(t *testing.T) { + t.Run("thinking content_block_start followed by signature_delta with signature is not empty", func(t *testing.T) { + payload := []byte("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"model\":\"claude-3\",\"usage\":{\"output_tokens\":0}}}\n\nevent: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"thinking\",\"thinking\":\"\"}}\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"signature_delta\",\"signature\":\"sig_encrypted_carrier_payload\"}}\n\nevent: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":0}}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n") + if IsEmptyCompletionPayload(payload) { + t.Fatal("IsEmptyCompletionPayload() = true for thinking stream with non-empty signature_delta, want false") + } + }) + + t.Run("thinking content_block_start followed by empty signature_delta is empty completion", func(t *testing.T) { + payload := []byte("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"model\":\"claude-3\",\"usage\":{\"output_tokens\":0}}}\n\nevent: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"thinking\",\"thinking\":\"\"}}\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"signature_delta\",\"signature\":\"\"}}\n\nevent: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":0}}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n") + if !IsEmptyCompletionPayload(payload) { + t.Fatal("IsEmptyCompletionPayload() = false for thinking stream with empty signature_delta, want true") + } + }) +} + +func TestOpenAIResponsesFunctionCallArgumentsEmptyCompletion(t *testing.T) { + t.Run("empty function_call_arguments delta without prior call item does not set tool calls", func(t *testing.T) { + var detector StreamBootstrapDetector + chunk := []byte("event: response.function_call_arguments.delta\ndata: {\"type\":\"response.function_call_arguments.delta\",\"delta\":\"\"}\n\n") + if detector.Observe(chunk) { + t.Fatal("Observe() = true for empty function_call_arguments.delta, want false") + } + if detector.HasMeaningfulOutput() { + t.Fatal("HasMeaningfulOutput() = true for empty function_call_arguments.delta, want false") + } + }) + + t.Run("non-empty function_call_arguments delta sets tool calls", func(t *testing.T) { + var detector StreamBootstrapDetector + chunk := []byte("event: response.function_call_arguments.delta\ndata: {\"type\":\"response.function_call_arguments.delta\",\"delta\":\"{\\\"q\\\":\\\"search\\\"}\"}\n\n") + if !detector.Observe(chunk) { + t.Fatal("Observe() = false for meaningful function_call_arguments.delta, want true") + } + if !detector.HasMeaningfulOutput() { + t.Fatal("HasMeaningfulOutput() = false for meaningful function_call_arguments.delta, want true") + } + }) + + t.Run("empty function_call_arguments delta with prior established call item retains tool calls", func(t *testing.T) { + var detector StreamBootstrapDetector + itemChunk := []byte("event: response.output_item.added\ndata: {\"type\":\"response.output_item.added\",\"item\":{\"type\":\"function_call\",\"name\":\"search\",\"arguments\":\"\"}}\n\n") + if !detector.Observe(itemChunk) { + t.Fatal("Observe() = false for output_item.added function_call, want true") + } + deltaChunk := []byte("event: response.function_call_arguments.delta\ndata: {\"type\":\"response.function_call_arguments.delta\",\"delta\":\"\"}\n\n") + if !detector.Observe(deltaChunk) { + t.Fatal("Observe() = false for stream with prior function_call item, want true") + } + if !detector.HasMeaningfulOutput() { + t.Fatal("HasMeaningfulOutput() = false for stream with prior function_call item, want true") + } + }) + + t.Run("semantically empty function_call_arguments delta ({}) without prior call item does not set tool calls", func(t *testing.T) { + var detector StreamBootstrapDetector + chunk := []byte("event: response.function_call_arguments.delta\ndata: {\"type\":\"response.function_call_arguments.delta\",\"delta\":\"{}\"}\n\n") + if detector.Observe(chunk) { + t.Fatal("Observe() = true for semantically empty function_call_arguments.delta, want false") + } + if detector.HasMeaningfulOutput() { + t.Fatal("HasMeaningfulOutput() = true for semantically empty function_call_arguments.delta, want false") + } + }) + + t.Run("semantically empty function_call_arguments done ({}) without prior call item does not set tool calls", func(t *testing.T) { + var detector StreamBootstrapDetector + chunk := []byte("event: response.function_call_arguments.done\ndata: {\"type\":\"response.function_call_arguments.done\",\"arguments\":\"{}\"}\n\n") + if detector.Observe(chunk) { + t.Fatal("Observe() = true for semantically empty function_call_arguments.done, want false") + } + if detector.HasMeaningfulOutput() { + t.Fatal("HasMeaningfulOutput() = true for semantically empty function_call_arguments.done, want false") + } + }) + + t.Run("semantically empty function_call_arguments done ([]) without prior call item does not set tool calls", func(t *testing.T) { + var detector StreamBootstrapDetector + chunk := []byte("event: response.function_call_arguments.done\ndata: {\"type\":\"response.function_call_arguments.done\",\"arguments\":\"[]\"}\n\n") + if detector.Observe(chunk) { + t.Fatal("Observe() = true for semantically empty function_call_arguments.done, want false") + } + if detector.HasMeaningfulOutput() { + t.Fatal("HasMeaningfulOutput() = true for semantically empty function_call_arguments.done, want false") + } + }) + + t.Run("semantically empty function_call_arguments done (null) without prior call item does not set tool calls", func(t *testing.T) { + var detector StreamBootstrapDetector + chunk := []byte("event: response.function_call_arguments.done\ndata: {\"type\":\"response.function_call_arguments.done\",\"arguments\":\"null\"}\n\n") + if detector.Observe(chunk) { + t.Fatal("Observe() = true for semantically empty function_call_arguments.done, want false") + } + if detector.HasMeaningfulOutput() { + t.Fatal("HasMeaningfulOutput() = true for semantically empty function_call_arguments.done, want false") + } + }) + + t.Run("meaningful function_call_arguments done with real args sets tool calls", func(t *testing.T) { + var detector StreamBootstrapDetector + chunk := []byte("event: response.function_call_arguments.done\ndata: {\"type\":\"response.function_call_arguments.done\",\"arguments\":\"{\\\"location\\\":\\\"Paris\\\"}\"}\n\n") + if !detector.Observe(chunk) { + t.Fatal("Observe() = false for meaningful function_call_arguments.done, want true") + } + if !detector.HasMeaningfulOutput() { + t.Fatal("HasMeaningfulOutput() = false for meaningful function_call_arguments.done, want true") + } + }) +} + +func TestClaudeStreamBootstrapShortCircuitsOnMessageStop(t *testing.T) { + t.Run("empty claude stream message_stop marks terminal empty without waiting for channel close", func(t *testing.T) { + var detector StreamBootstrapDetector + startChunk := []byte("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"model\":\"claude-3-5-sonnet\",\"stop_reason\":null,\"stop_sequence\":null,\"usage\":{\"input_tokens\":10,\"output_tokens\":0}}}\n\n") + stopChunk := []byte("event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n") + + if detector.Observe(startChunk) { + t.Fatal("Observe(message_start) = true, want false") + } + if detector.Observe(stopChunk) { + t.Fatal("Observe(message_stop) = true, want false") + } + if !detector.IsTerminalEmpty() { + t.Fatal("IsTerminalEmpty() = false on message_stop, want true (sawDone equivalent)") + } + }) +} + +func TestMultiValueJSONMixedUnknownEmptyCompletion(t *testing.T) { + t.Run("multi-value json with recognized empty and unknown object is not empty completion", func(t *testing.T) { + payload := []byte(`{"choices":[{"message":{"content":""},"finish_reason":"stop"}]}{"custom_provider_event":{"data":"foo"}}`) + if IsEmptyCompletionPayload(payload) { + t.Fatal("IsEmptyCompletionPayload() = true for mixed recognized-empty and unknown JSON values, want false") + } + }) + + t.Run("multi-value json with only recognized empty completions is empty completion", func(t *testing.T) { + payload := []byte(`{"choices":[{"message":{"content":""},"finish_reason":"stop"}]}{"choices":[{"message":{"content":""},"finish_reason":"stop"}]}`) + if !IsEmptyCompletionPayload(payload) { + t.Fatal("IsEmptyCompletionPayload() = false for multiple recognized empty completions, want true") + } + }) +} + +func TestGeminiThoughtSignatureEmptyCompletion(t *testing.T) { + t.Run("gemini STOP with thoughtSignature and omitted candidatesTokenCount is not empty", func(t *testing.T) { + payload := []byte(`{"candidates":[{"content":{"role":"model","parts":[{"thoughtSignature":"sig_gemini_thought_123"}]},"finishReason":"STOP"}]}`) + if IsEmptyCompletionPayload(payload) { + t.Fatal("IsEmptyCompletionPayload() = true for non-empty thoughtSignature with omitted token count, want false") + } + }) + + t.Run("gemini STOP with thought_signature and omitted candidatesTokenCount is not empty", func(t *testing.T) { + payload := []byte(`{"candidates":[{"content":{"role":"model","parts":[{"thought_signature":"sig_gemini_thought_123"}]},"finishReason":"STOP"}]}`) + if IsEmptyCompletionPayload(payload) { + t.Fatal("IsEmptyCompletionPayload() = true for non-empty thought_signature with omitted token count, want false") + } + }) + + t.Run("gemini STOP with thoughtSignature and zero candidatesTokenCount is not empty", func(t *testing.T) { + payload := []byte(`{"candidates":[{"content":{"role":"model","parts":[{"text":"","thoughtSignature":"sig_gemini_thought_123"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":0}}`) + if IsEmptyCompletionPayload(payload) { + t.Fatal("IsEmptyCompletionPayload() = true for non-empty thoughtSignature with zero token count, want false") + } + }) + + t.Run("gemini STOP with thought_signature and zero candidatesTokenCount is not empty", func(t *testing.T) { + payload := []byte(`{"candidates":[{"content":{"role":"model","parts":[{"text":"","thought_signature":"sig_gemini_thought_123"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":0}}`) + if IsEmptyCompletionPayload(payload) { + t.Fatal("IsEmptyCompletionPayload() = true for non-empty thought_signature with zero token count, want false") + } + }) + + t.Run("gemini STOP with empty thoughtSignature is empty", func(t *testing.T) { + payload := []byte(`{"candidates":[{"content":{"role":"model","parts":[{"text":"","thoughtSignature":""}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":0}}`) + if !IsEmptyCompletionPayload(payload) { + t.Fatal("IsEmptyCompletionPayload() = false for empty thoughtSignature, want true") + } + }) + + t.Run("gemini STOP with empty thought_signature is empty", func(t *testing.T) { + payload := []byte(`{"candidates":[{"content":{"role":"model","parts":[{"text":"","thought_signature":""}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":0}}`) + if !IsEmptyCompletionPayload(payload) { + t.Fatal("IsEmptyCompletionPayload() = false for empty thought_signature, want true") + } + }) +} + +func TestReadStreamBootstrapForwardsPositiveUsageTerminalFrameImmediately(t *testing.T) { + t.Run("positive completion tokens forwards immediately without stream close", func(t *testing.T) { + ch := make(chan cliproxyexecutor.StreamChunk, 2) + ch <- cliproxyexecutor.StreamChunk{ + Payload: []byte("data: {\"choices\":[],\"usage\":{\"completion_tokens\":1}}\n\n"), + } + ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) + defer cancel() + + buffered, closed, err := readStreamBootstrap(ctx, ch) + if err != nil { + t.Fatalf("readStreamBootstrap error = %v, want immediate forward", err) + } + if closed { + t.Fatalf("readStreamBootstrap returned closed = true, want false (channel still open)") + } + if len(buffered) != 1 { + t.Fatalf("buffered chunks count = %d, want 1", len(buffered)) + } + }) + + t.Run("zero completion tokens is withheld and not forwarded while stream open", func(t *testing.T) { + ch := make(chan cliproxyexecutor.StreamChunk, 2) + ch <- cliproxyexecutor.StreamChunk{ + Payload: []byte("data: {\"choices\":[],\"usage\":{\"completion_tokens\":0}}\n\n"), + } + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + + _, _, err := readStreamBootstrap(ctx, ch) + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("readStreamBootstrap error = %v, want context.DeadlineExceeded (withheld)", err) + } + }) +} + +func TestResponsesReasoningOutputItemBootstrap(t *testing.T) { + emptyReasoning := []byte("data: {\"type\":\"response.output_item.added\",\"sequence_number\":0,\"output_index\":0,\"item\":{\"id\":\"rs_1\",\"type\":\"reasoning\",\"status\":\"in_progress\",\"encrypted_content\":\"\",\"summary\":[]}}\n\n") + encryptedReasoning := []byte("data: {\"type\":\"response.output_item.added\",\"sequence_number\":0,\"output_index\":0,\"item\":{\"id\":\"rs_1\",\"type\":\"reasoning\",\"status\":\"in_progress\",\"encrypted_content\":\"gAAAA_signature_123\",\"summary\":[]}}\n\n") + summaryReasoning := []byte("data: {\"type\":\"response.output_item.added\",\"sequence_number\":0,\"output_index\":0,\"item\":{\"id\":\"rs_1\",\"type\":\"reasoning\",\"status\":\"in_progress\",\"encrypted_content\":\"\",\"summary\":[{\"type\":\"summary_text\",\"text\":\"reasoning step\"}]}}\n\n") + + t.Run("empty reasoning item does not mark meaningful and allows bootstrap error failover", func(t *testing.T) { + var detector StreamBootstrapDetector + if detector.Observe(emptyReasoning) { + t.Fatal("detector.Observe() = true for empty reasoning scaffolding, want false") + } + if detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = true for empty reasoning scaffolding, want false") + } + + errUpstream := errors.New("upstream failed immediately after scaffolding") + ch := make(chan cliproxyexecutor.StreamChunk, 2) + ch <- cliproxyexecutor.StreamChunk{Payload: emptyReasoning} + ch <- cliproxyexecutor.StreamChunk{Err: errUpstream} + close(ch) + + buffered, closed, err := readStreamBootstrap(context.Background(), ch) + if !errors.Is(err, errUpstream) { + t.Fatalf("readStreamBootstrap error = %v, want %v for failover", err, errUpstream) + } + if len(buffered) != 0 { + t.Fatalf("readStreamBootstrap buffered = %d, want 0 on failover error", len(buffered)) + } + if closed { + t.Fatal("readStreamBootstrap returned closed = true, want false on error") + } + }) + + t.Run("reasoning item with encrypted_content marks meaningful and forwards", func(t *testing.T) { + var detector StreamBootstrapDetector + if !detector.Observe(encryptedReasoning) { + t.Fatal("detector.Observe() = false for reasoning with encrypted_content, want true") + } + if !detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = false for reasoning with encrypted_content, want true") + } + }) + + t.Run("reasoning item with summary marks meaningful and forwards", func(t *testing.T) { + var detector StreamBootstrapDetector + if !detector.Observe(summaryReasoning) { + t.Fatal("detector.Observe() = false for reasoning with summary, want true") + } + if !detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = false for reasoning with summary, want true") + } + }) +} + +func TestClaudeInputJSONDeltaSemanticallyEmpty(t *testing.T) { + emptyObjectDelta := []byte(`{"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"{ }"}}`) + emptyArrayDelta := []byte(`{"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"[]"}}`) + nullSpaceDelta := []byte(`{"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"null "}}`) + validCompleteDelta := []byte(`{"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"{\"query\":\"test\"}"}}`) + validIncompleteDelta := []byte(`{"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"{\"query\":"}}`) + + t.Run("whitespace empty object does not mark meaningful and allows bootstrap error failover", func(t *testing.T) { + var detector StreamBootstrapDetector + if detector.Observe(emptyObjectDelta) { + t.Fatal("detector.Observe() = true for empty object partial_json, want false") + } + if detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = true for empty object partial_json, want false") + } + errUpstream := errors.New("upstream failed after empty arg delta") + ch := make(chan cliproxyexecutor.StreamChunk, 2) + ch <- cliproxyexecutor.StreamChunk{Payload: emptyObjectDelta} + ch <- cliproxyexecutor.StreamChunk{Err: errUpstream} + close(ch) + buffered, closed, err := readStreamBootstrap(context.Background(), ch) + if !errors.Is(err, errUpstream) { + t.Fatalf("readStreamBootstrap error = %v, want %v for failover", err, errUpstream) + } + if len(buffered) != 0 { + t.Fatalf("readStreamBootstrap buffered = %d, want 0 on failover error", len(buffered)) + } + if closed { + t.Fatal("readStreamBootstrap returned closed = true, want false on error") + } + }) + + t.Run("empty array does not mark meaningful and allows bootstrap error failover", func(t *testing.T) { + var detector StreamBootstrapDetector + if detector.Observe(emptyArrayDelta) { + t.Fatal("detector.Observe() = true for empty array partial_json, want false") + } + if detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = true for empty array partial_json, want false") + } + }) + + t.Run("null with space does not mark meaningful and allows bootstrap error failover", func(t *testing.T) { + var detector StreamBootstrapDetector + if detector.Observe(nullSpaceDelta) { + t.Fatal("detector.Observe() = true for null space partial_json, want false") + } + if detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = true for null space partial_json, want false") + } + }) + + t.Run("valid complete partial_json marks meaningful and forwards", func(t *testing.T) { + var detector StreamBootstrapDetector + if !detector.Observe(validCompleteDelta) { + t.Fatal("detector.Observe() = false for valid complete partial_json, want true") + } + if !detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = false for valid complete partial_json, want true") + } + }) + + t.Run("valid incomplete partial_json marks meaningful and forwards", func(t *testing.T) { + var detector StreamBootstrapDetector + if !detector.Observe(validIncompleteDelta) { + t.Fatal("detector.Observe() = false for valid incomplete partial_json, want true") + } + if !detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = false for valid incomplete partial_json, want true") + } + }) +} + +func TestExecuteStream_TerminalDoneWithoutClosingChannelRotatesAuth(t *testing.T) { + executor := &emptyCompletionTestExecutor{ + streamPayloads: map[string][][]byte{}, + streamCalls: map[string]int{}, + leaveStreamOpen: true, + emptyStreamPayload: [][]byte{ + []byte(": keep-alive\n\n"), + []byte("data: [DONE]\n\n"), + }, + } + manager, ids, model, capture := newEmptyCompletionTestManager(t, executor) + + ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond) + defer cancel() + + stream, err := manager.ExecuteStream(ctx, []string{"claude"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{Stream: true}) + if err != nil { + t.Fatalf("ExecuteStream() error = %v", err) + } + if stream == nil { + t.Fatal("ExecuteStream() returned nil stream") + } + + var got strings.Builder + for chunk := range stream.Chunks { + got.Write(chunk.Payload) + } + assertRotatesToContent(t, ids, executor.firstStream, got.String(), "hello", capture) +} + +func TestExecuteStream_MeaningfulContentWithOpenChannelForwardsImmediately(t *testing.T) { + chunks := make(chan cliproxyexecutor.StreamChunk, 2) + chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("data: {\"choices\":[{\"delta\":{\"content\":\"meaningful_content\"}}]}\n\n")} + chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("data: [DONE]\n\n")} + // Leave channel open + + customExec := &customStreamOpenChannelExecutor{chunks: chunks} + + manager := NewManager(nil, nil, nil) + manager.SetRetryConfig(3, 5*time.Second, 3) + model := "open-channel-meaningful-" + uuid.NewString() + + auth := &Auth{ID: "auth-1", Provider: "claude", Status: StatusActive} + if _, err := manager.Register(context.Background(), auth); err != nil { + t.Fatalf("Register error: %v", err) + } + reg := registry.GetGlobalRegistry() + reg.RegisterClient(auth.ID, "claude", []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { reg.UnregisterClient(auth.ID) }) + manager.RefreshSchedulerEntry(auth.ID) + manager.RegisterExecutor(customExec) + + ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond) + defer cancel() + + stream, err := manager.ExecuteStream(ctx, []string{"claude"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{Stream: true}) + if err != nil { + t.Fatalf("ExecuteStream() error = %v", err) + } + if stream == nil { + t.Fatal("ExecuteStream() returned nil stream") + } + + firstChunk := <-stream.Chunks + if !strings.Contains(string(firstChunk.Payload), "meaningful_content") { + t.Fatalf("first chunk payload = %q, want meaningful_content", string(firstChunk.Payload)) + } +} + +type customStreamOpenChannelExecutor struct { + chunks chan cliproxyexecutor.StreamChunk +} + +func (e *customStreamOpenChannelExecutor) Identifier() string { return "claude" } + +func (e *customStreamOpenChannelExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{Payload: []byte("ok")}, nil +} + +func (e *customStreamOpenChannelExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{Payload: []byte("ok")}, nil +} + +func (e *customStreamOpenChannelExecutor) Refresh(_ context.Context, a *Auth) (*Auth, error) { + return a, nil +} + +func (e *customStreamOpenChannelExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +func (e *customStreamOpenChannelExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + return &cliproxyexecutor.StreamResult{Chunks: e.chunks}, nil +} + +func TestResponsesEmptyToolCallScaffold(t *testing.T) { + emptyFuncScaffold := []byte("event: response.output_item.added\ndata: {\"type\":\"response.output_item.added\",\"sequence_number\":0,\"output_index\":0,\"item\":{\"id\":\"\",\"type\":\"function_call\",\"status\":\"in_progress\",\"arguments\":\"\",\"call_id\":\"\",\"name\":\"\"}}\n\n") + emptyCustomToolScaffold := []byte("event: response.output_item.added\ndata: {\"type\":\"response.output_item.added\",\"sequence_number\":0,\"output_index\":0,\"item\":{\"id\":\"\",\"type\":\"custom_tool_call\",\"status\":\"in_progress\",\"input\":\"\",\"call_id\":\"\",\"name\":\"\"}}\n\n") + funcWithID := []byte("event: response.output_item.added\ndata: {\"type\":\"response.output_item.added\",\"sequence_number\":0,\"output_index\":0,\"item\":{\"id\":\"fc_123\",\"type\":\"function_call\",\"status\":\"in_progress\",\"arguments\":\"\",\"call_id\":\"\",\"name\":\"\"}}\n\n") + funcWithCallID := []byte("event: response.output_item.added\ndata: {\"type\":\"response.output_item.added\",\"sequence_number\":0,\"output_index\":0,\"item\":{\"id\":\"\",\"type\":\"function_call\",\"status\":\"in_progress\",\"arguments\":\"\",\"call_id\":\"call_123\",\"name\":\"\"}}\n\n") + funcWithName := []byte("event: response.output_item.added\ndata: {\"type\":\"response.output_item.added\",\"sequence_number\":0,\"output_index\":0,\"item\":{\"id\":\"\",\"type\":\"function_call\",\"status\":\"in_progress\",\"arguments\":\"\",\"call_id\":\"\",\"name\":\"lookup\"}}\n\n") + funcWithArgs := []byte("event: response.output_item.added\ndata: {\"type\":\"response.output_item.added\",\"sequence_number\":0,\"output_index\":0,\"item\":{\"id\":\"\",\"type\":\"function_call\",\"status\":\"in_progress\",\"arguments\":\"{\\\"q\\\":\\\"search\\\"}\",\"call_id\":\"\",\"name\":\"\"}}\n\n") + customToolWithInput := []byte("event: response.output_item.added\ndata: {\"type\":\"response.output_item.added\",\"sequence_number\":0,\"output_index\":0,\"item\":{\"id\":\"\",\"type\":\"custom_tool_call\",\"status\":\"in_progress\",\"input\":\"{\\\"cmd\\\":\\\"run\\\"}\",\"call_id\":\"\",\"name\":\"\"}}\n\n") + + t.Run("empty function_call scaffold does not mark meaningful and allows bootstrap error failover", func(t *testing.T) { + var detector StreamBootstrapDetector + if detector.Observe(emptyFuncScaffold) { + t.Fatal("detector.Observe() = true for empty function_call scaffold, want false") + } + if detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = true for empty function_call scaffold, want false") + } + + errUpstream := errors.New("upstream failed immediately after function_call scaffold") + ch := make(chan cliproxyexecutor.StreamChunk, 2) + ch <- cliproxyexecutor.StreamChunk{Payload: emptyFuncScaffold} + ch <- cliproxyexecutor.StreamChunk{Err: errUpstream} + close(ch) + buffered, closed, err := readStreamBootstrap(context.Background(), ch) + if !errors.Is(err, errUpstream) { + t.Fatalf("readStreamBootstrap error = %v, want %v for failover", err, errUpstream) + } + if len(buffered) != 0 { + t.Fatalf("readStreamBootstrap buffered = %d, want 0 on failover error", len(buffered)) + } + if closed { + t.Fatal("readStreamBootstrap returned closed = true, want false on error") + } + }) + + t.Run("empty custom_tool_call scaffold does not mark meaningful and allows bootstrap error failover", func(t *testing.T) { + var detector StreamBootstrapDetector + if detector.Observe(emptyCustomToolScaffold) { + t.Fatal("detector.Observe() = true for empty custom_tool_call scaffold, want false") + } + if detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = true for empty custom_tool_call scaffold, want false") + } + }) + + t.Run("scaffold with non-empty id marks meaningful and forwards", func(t *testing.T) { + var detector StreamBootstrapDetector + if !detector.Observe(funcWithID) { + t.Fatal("detector.Observe() = false for function_call with id, want true") + } + if !detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = false for function_call with id, want true") + } + }) + + t.Run("scaffold with non-empty call_id marks meaningful and forwards", func(t *testing.T) { + var detector StreamBootstrapDetector + if !detector.Observe(funcWithCallID) { + t.Fatal("detector.Observe() = false for function_call with call_id, want true") + } + if !detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = false for function_call with call_id, want true") + } + }) + + t.Run("scaffold with non-empty name marks meaningful and forwards", func(t *testing.T) { + var detector StreamBootstrapDetector + if !detector.Observe(funcWithName) { + t.Fatal("detector.Observe() = false for function_call with name, want true") + } + if !detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = false for function_call with name, want true") + } + }) + + t.Run("scaffold with non-empty arguments marks meaningful and forwards", func(t *testing.T) { + var detector StreamBootstrapDetector + if !detector.Observe(funcWithArgs) { + t.Fatal("detector.Observe() = false for function_call with arguments, want true") + } + if !detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = false for function_call with arguments, want true") + } + }) + + t.Run("custom tool with non-empty input marks meaningful and forwards", func(t *testing.T) { + var detector StreamBootstrapDetector + if !detector.Observe(customToolWithInput) { + t.Fatal("detector.Observe() = false for custom_tool_call with input, want true") + } + if !detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = false for custom_tool_call with input, want true") + } + }) + + t.Run("output_item.done with empty function_call does not mark meaningful and allows bootstrap error failover", func(t *testing.T) { + emptyDoneFuncItems := [][]byte{ + []byte("event: response.output_item.done\ndata: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"type\":\"function_call\"}}\n\n"), + []byte("event: response.output_item.done\ndata: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"id\":\"\",\"type\":\"function_call\",\"status\":\"completed\",\"arguments\":\"\",\"call_id\":\"\",\"name\":\"\"}}\n\n"), + []byte("event: response.output_item.done\ndata: {\"type\":\"response.output_item.done\",\"output_index\":0,\"output\":{\"type\":\"function_call\"}}\n\n"), + []byte("event: response.output_item.done\ndata: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"type\":\"custom_tool_call\"}}\n\n"), + } + for i, payload := range emptyDoneFuncItems { + var detector StreamBootstrapDetector + if detector.Observe(payload) { + t.Fatalf("case %d: detector.Observe() = true for empty output_item.done, want false", i) + } + if detector.HasMeaningfulOutput() { + t.Fatalf("case %d: detector.HasMeaningfulOutput() = true for empty output_item.done, want false", i) + } + } + + errUpstream := errors.New("upstream failed immediately after output_item.done empty scaffold") + ch := make(chan cliproxyexecutor.StreamChunk, 2) + ch <- cliproxyexecutor.StreamChunk{Payload: emptyDoneFuncItems[0]} + ch <- cliproxyexecutor.StreamChunk{Err: errUpstream} + close(ch) + buffered, closed, err := readStreamBootstrap(context.Background(), ch) + if !errors.Is(err, errUpstream) { + t.Fatalf("readStreamBootstrap error = %v, want %v for failover", err, errUpstream) + } + if len(buffered) != 0 { + t.Fatalf("readStreamBootstrap buffered = %d, want 0 on failover error", len(buffered)) + } + if closed { + t.Fatal("readStreamBootstrap returned closed = true, want false on error") + } + }) + + t.Run("output_item.done with valid function_call marks meaningful and forwards", func(t *testing.T) { + validDoneFunc := []byte("event: response.output_item.done\ndata: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"id\":\"fc_123\",\"type\":\"function_call\",\"status\":\"completed\",\"arguments\":\"{\\\"query\\\":\\\"go\\\"}\",\"call_id\":\"call_123\",\"name\":\"search\"}}\n\n") + var detector StreamBootstrapDetector + if !detector.Observe(validDoneFunc) { + t.Fatal("detector.Observe() = false for valid output_item.done function_call, want true") + } + if !detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = false for valid output_item.done function_call, want true") + } + }) +} + +func TestGeminiStreamBootstrapTerminalEmptyOnSTOP(t *testing.T) { + t.Run("empty gemini stream STOP marks terminal empty without waiting for channel close", func(t *testing.T) { + var detector StreamBootstrapDetector + stopChunk := []byte("data: {\"candidates\":[{\"finishReason\":\"STOP\"}]}\n\n") + + if detector.Observe(stopChunk) { + t.Fatal("Observe(gemini STOP) = true, want false") + } + if !detector.IsTerminalEmpty() { + t.Fatal("IsTerminalEmpty() = false on gemini STOP, want true") + } + }) + + t.Run("gemini stream with content then STOP is not terminal empty", func(t *testing.T) { + var detector StreamBootstrapDetector + contentChunk := []byte("data: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"hi\"}]}}]}\n\n") + stopChunk := []byte("data: {\"candidates\":[{\"finishReason\":\"STOP\"}]}\n\n") + + if !detector.Observe(contentChunk) { + t.Fatal("Observe(gemini content) = false, want true") + } + if !detector.Observe(stopChunk) { + t.Fatal("Observe(gemini STOP after content) = false, want true") + } + if detector.IsTerminalEmpty() { + t.Fatal("IsTerminalEmpty() = true for stream with content, want false") + } + }) + + t.Run("gemini stream with blocked finishReason is forwarded and not terminal empty", func(t *testing.T) { + var detector StreamBootstrapDetector + safetyChunk := []byte("data: {\"candidates\":[{\"finishReason\":\"SAFETY\"}]}\n\n") + + if !detector.Observe(safetyChunk) { + t.Fatal("Observe(gemini SAFETY) = false, want true (blocked reasons must reach client)") + } + if detector.IsTerminalEmpty() { + t.Fatal("IsTerminalEmpty() = true on gemini SAFETY, want false") + } + }) + + t.Run("conductor readStreamBootstrap with empty gemini STOP over open channel classifies empty", func(t *testing.T) { + ch := make(chan cliproxyexecutor.StreamChunk, 2) + ch <- cliproxyexecutor.StreamChunk{Payload: []byte("data: {\"candidates\":[{\"finishReason\":\"STOP\"}]}\n\n")} + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + buffered, terminalEmpty, err := readStreamBootstrap(ctx, ch) + if err != nil { + t.Fatalf("readStreamBootstrap error = %v, want nil", err) + } + if !terminalEmpty { + t.Fatal("readStreamBootstrap terminalEmpty = false, want true") + } + if len(buffered) != 1 { + t.Fatalf("len(buffered) = %d, want 1", len(buffered)) + } + }) +} + +func TestClaudeDataOnlyMessageStopTerminalEmpty(t *testing.T) { + t.Run("empty claude data-only message_stop marks terminal empty without waiting for channel close", func(t *testing.T) { + var detector StreamBootstrapDetector + stopChunk := []byte("data: {\"type\":\"message_stop\"}\n\n") + + if detector.Observe(stopChunk) { + t.Fatal("Observe(claude data-only message_stop) = true, want false") + } + if !detector.IsTerminalEmpty() { + t.Fatal("IsTerminalEmpty() = false on claude data-only message_stop, want true") + } + }) + + t.Run("claude stream with content then data-only message_stop is not terminal empty", func(t *testing.T) { + var detector StreamBootstrapDetector + contentChunk := []byte("data: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"text_delta\",\"text\":\"hi\"}}\n\n") + stopChunk := []byte("data: {\"type\":\"message_stop\"}\n\n") + + if !detector.Observe(contentChunk) { + t.Fatal("Observe(claude content) = false, want true") + } + if !detector.Observe(stopChunk) { + t.Fatal("Observe(claude message_stop after content) = false, want true") + } + if detector.IsTerminalEmpty() { + t.Fatal("IsTerminalEmpty() = true for stream with content, want false") + } + }) + + t.Run("conductor readStreamBootstrap with data-only claude message_stop over open channel classifies empty", func(t *testing.T) { + ch := make(chan cliproxyexecutor.StreamChunk, 2) + ch <- cliproxyexecutor.StreamChunk{Payload: []byte("data: {\"type\":\"message_stop\"}\n\n")} + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + buffered, terminalEmpty, err := readStreamBootstrap(ctx, ch) + if err != nil { + t.Fatalf("readStreamBootstrap error = %v, want nil", err) + } + if !terminalEmpty { + t.Fatal("readStreamBootstrap terminalEmpty = false, want true") + } + if len(buffered) != 1 { + t.Fatalf("len(buffered) = %d, want 1", len(buffered)) + } + }) +} + +func TestEmptyCompletionImages(t *testing.T) { + cases := []struct { + name string + payload []byte + expected bool + }{ + { + name: "delta images with image_url is not empty", + payload: []byte(`data: {"id":"1","choices":[{"index":0,"delta":{"images":[{"type":"image_url","image_url":{"url":"data:image/png;base64,AQID"}}]},"finish_reason":"stop"}]}` + "\n\n" + `data: [DONE]` + "\n\n"), + expected: false, + }, + { + name: "message images non-stream is not empty", + payload: []byte(`{"id":"1","choices":[{"index":0,"message":{"role":"assistant","content":"","images":[{"type":"image_url","image_url":{"url":"data:image/png;base64,AQID"}}]},"finish_reason":"stop"}],"usage":{"completion_tokens":0}}`), + expected: false, + }, + { + name: "delta images empty array stays empty", + payload: []byte(`data: {"id":"1","choices":[{"index":0,"delta":{"images":[]},"finish_reason":"stop"}]}` + "\n\n" + `data: [DONE]` + "\n\n"), + expected: true, + }, + { + name: "delta images null stays empty", + payload: []byte(`data: {"id":"1","choices":[{"index":0,"delta":{"images":null},"finish_reason":"stop"}]}` + "\n\n" + `data: [DONE]` + "\n\n"), + expected: true, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := isEmptyCompletionPayload(tc.payload); got != tc.expected { + t.Fatalf("isEmptyCompletionPayload() = %v, want %v", got, tc.expected) + } + }) + } +} + +func TestEmptyCompletionClaudeCitations(t *testing.T) { + cases := []struct { + name string + payload []byte + expected bool + }{ + { + name: "citations_delta with non-empty citation object is not empty", + payload: []byte("event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"citations_delta\",\"citation\":{\"type\":\"char_location\",\"cited_text\":\"some cited text\",\"document_index\":0}}}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"), + expected: false, + }, + { + name: "citations_delta with empty citation object is empty", + payload: []byte("event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"citations_delta\",\"citation\":{}}}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"), + expected: true, + }, + { + name: "citations_delta with null citation is empty", + payload: []byte("event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"citations_delta\",\"citation\":null}}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"), + expected: true, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := isEmptyCompletionPayload(tc.payload); got != tc.expected { + t.Fatalf("isEmptyCompletionPayload() = %v, want %v", got, tc.expected) + } + }) + } +} + +func TestEmptyCompletionResponsesImageGenerationCallResult(t *testing.T) { + meaningfulPayload := []byte("data: {\"type\":\"response.output_item.done\",\"item\":{\"type\":\"image_generation_call\",\"status\":\"completed\",\"result\":\"image-data\"}}\n\ndata: {\"type\":\"response.completed\",\"response\":{\"status\":\"completed\",\"output\":[],\"usage\":{\"output_tokens\":0}}}\n\ndata: [DONE]\n\n") + if got := isEmptyCompletionPayload(meaningfulPayload); got != false { + t.Fatalf("isEmptyCompletionPayload(meaningful image_generation_call result) = %v, want false", got) + } + + detector := &StreamBootstrapDetector{} + if got := detector.Observe([]byte("data: {\"type\":\"response.output_item.done\",\"item\":{\"type\":\"image_generation_call\",\"status\":\"completed\",\"result\":\"image-data\"}}\n\n")); got != true { + t.Fatalf("StreamBootstrapDetector.Observe(meaningful result) = %v, want true", got) + } + + emptyDetector := &StreamBootstrapDetector{} + if got := emptyDetector.Observe([]byte("data: {\"type\":\"response.output_item.done\",\"item\":{\"type\":\"image_generation_call\",\"status\":\"completed\",\"result\":\"\"}}\n\n")); got != false { + t.Fatalf("StreamBootstrapDetector.Observe(empty result) = %v, want false", got) + } + + wsDetector := &StreamBootstrapDetector{} + if got := wsDetector.Observe([]byte("data: {\"type\":\"response.output_item.done\",\"item\":{\"type\":\"image_generation_call\",\"status\":\"completed\",\"result\":\" \"}}\n\n")); got != false { + t.Fatalf("StreamBootstrapDetector.Observe(whitespace result) = %v, want false", got) + } +} diff --git a/sdk/cliproxy/auth/export_test.go b/sdk/cliproxy/auth/export_test.go new file mode 100644 index 000000000..d8e6a0225 --- /dev/null +++ b/sdk/cliproxy/auth/export_test.go @@ -0,0 +1,23 @@ +package auth + +import ( + "bytes" +) + +// IsCompletionFormatRecognized reports whether payload uses a wire format the +// empty-completion detection understands (OpenAI chat, OpenAI Responses, +// Anthropic Claude, or Gemini). It supports representative format-contract +// tests without claiming registry-wide executor coverage. +func IsCompletionFormatRecognized(payload []byte) bool { + trimmed := bytes.TrimSpace(payload) + if len(trimmed) == 0 { + return false + } + var acc emptyCompletionAccum + if bytes.Contains(trimmed, []byte("data:")) || bytes.HasPrefix(trimmed, []byte("event:")) { + acc.evalSSE(trimmed) + } else { + acc.evalJSON(trimmed) + } + return acc.recognized +} diff --git a/sdk/cliproxy/auth/home_concurrency.go b/sdk/cliproxy/auth/home_concurrency.go index d3f91774e..3bfdcefd4 100644 --- a/sdk/cliproxy/auth/home_concurrency.go +++ b/sdk/cliproxy/auth/home_concurrency.go @@ -265,6 +265,13 @@ func verifyAccountedHomeConcurrencyIdentity(tuple homeConcurrencyTuple, auth *Au // SafeResponseHeaders returns trusted response headers only for CPA's concrete Home busy error. func SafeResponseHeaders(err error) http.Header { + if err == nil { + return nil + } + var carrier interface{ SafeResponseHeaders() http.Header } + if errors.As(err, &carrier) && carrier != nil { + return carrier.SafeResponseHeaders() + } var busy *HomeConcurrencyBusyError if !errors.As(err, &busy) || busy == nil { return nil diff --git a/sdk/cliproxy/auth/home_execution_paths_test.go b/sdk/cliproxy/auth/home_execution_paths_test.go index 29bb89348..5e2c359fd 100644 --- a/sdk/cliproxy/auth/home_execution_paths_test.go +++ b/sdk/cliproxy/auth/home_execution_paths_test.go @@ -640,7 +640,7 @@ func (*alwaysEmptyHomeStreamExecutor) Execute(context.Context, *Auth, cliproxyex func (e *alwaysEmptyHomeStreamExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { e.calls.Add(1) chunks := make(chan cliproxyexecutor.StreamChunk, 1) - chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("data: {\"type\":\"response.completed\",\"response\":{\"status\":\"completed\",\"output\":[],\"usage\":{\"output_tokens\":0}}}\n\n")} + chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":{\"completion_tokens\":0}}\n\n")} close(chunks) return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil } @@ -714,7 +714,7 @@ func (e *alternatingEmptyHomeExecutor) ExecuteStream(_ context.Context, auth *Au return nil, &Error{Code: "transient", Message: "transient stream failure", Retryable: true} } chunks := make(chan cliproxyexecutor.StreamChunk, 1) - chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("data: {\"type\":\"response.completed\",\"response\":{\"status\":\"completed\",\"output\":[]}}\n\n")} + chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":{\"completion_tokens\":0}}\n\n")} close(chunks) return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil } @@ -734,7 +734,7 @@ func TestHomeStreamDoesNotRevisitEmptyAuthAfterAnotherFailure(t *testing.T) { executor := &alternatingEmptyHomeExecutor{} manager.RegisterExecutor(executor) - _, errExecute := manager.executeStreamMixedOnce(context.Background(), []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{Stream: true}, 0) + _, errExecute := manager.executeStreamMixedOnce(context.Background(), []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{Stream: true}, 0, nil) if errExecute == nil || !strings.Contains(errExecute.Error(), "transient stream failure") { t.Fatalf("executeStreamMixedOnce() error = %v, want last transient failure", errExecute) } diff --git a/sdk/cliproxy/auth/outer_retry_exclusions_test.go b/sdk/cliproxy/auth/outer_retry_exclusions_test.go new file mode 100644 index 000000000..bad5ea3bf --- /dev/null +++ b/sdk/cliproxy/auth/outer_retry_exclusions_test.go @@ -0,0 +1,255 @@ +package auth + +import ( + "context" + "errors" + "net/http" + "sync" + "testing" + "time" + + "github.com/google/uuid" + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +type outerRetryTestExecutor struct { + mu sync.Mutex + executeCalls map[string]int + streamCalls map[string]int + totalCalls int + failFirstN int + executeErrs map[string]error + streamErrs map[string]error + responses map[string]cliproxyexecutor.Response +} + +func newOuterRetryTestExecutor() *outerRetryTestExecutor { + return &outerRetryTestExecutor{ + executeCalls: make(map[string]int), + streamCalls: make(map[string]int), + executeErrs: make(map[string]error), + streamErrs: make(map[string]error), + responses: make(map[string]cliproxyexecutor.Response), + } +} + +func (e *outerRetryTestExecutor) Identifier() string { return "claude" } +func (*outerRetryTestExecutor) ShouldPrepareRequestAuth(*Auth) bool { return false } +func (e *outerRetryTestExecutor) PrepareRequestAuth(_ context.Context, auth *Auth) (*Auth, error) { + return auth, nil +} + +func (e *outerRetryTestExecutor) Execute(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + e.mu.Lock() + defer e.mu.Unlock() + e.executeCalls[auth.ID]++ + e.totalCalls++ + if err, ok := e.executeErrs[auth.ID]; ok && err != nil { + return cliproxyexecutor.Response{}, err + } + if e.totalCalls <= e.failFirstN { + return cliproxyexecutor.Response{}, &Error{Code: "service_unavailable", Message: "503 Service Unavailable"} + } + if resp, ok := e.responses[auth.ID]; ok { + return resp, nil + } + return cliproxyexecutor.Response{Payload: []byte(`{"choices":[{"message":{"content":"ok"}}]}`)}, nil +} + +func (e *outerRetryTestExecutor) ExecuteStream(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + e.mu.Lock() + defer e.mu.Unlock() + e.streamCalls[auth.ID]++ + if err, ok := e.streamErrs[auth.ID]; ok && err != nil { + return nil, err + } + ch := make(chan cliproxyexecutor.StreamChunk, 1) + ch <- cliproxyexecutor.StreamChunk{Payload: []byte(`data: {"choices":[{"delta":{"content":"ok"}}]}\n\n`)} + close(ch) + return &cliproxyexecutor.StreamResult{Chunks: ch}, nil +} + +func (e *outerRetryTestExecutor) Refresh(_ context.Context, auth *Auth) (*Auth, error) { + return auth, nil +} + +func (*outerRetryTestExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, errors.New("not implemented") +} + +func (e *outerRetryTestExecutor) CountTokens(_ context.Context, _ *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} + +func newOuterRetryTestManager(t *testing.T, executor *outerRetryTestExecutor, authCount int, disableCooling bool) (*Manager, []string, string) { + t.Helper() + model := "outer-retry-model-" + uuid.NewString() + manager := NewManager(nil, nil, nil) + manager.SetRetryConfig(3, 0, 3) + manager.RegisterExecutor(executor) + + var ids []string + for i := 0; i < authCount; i++ { + authID := "outer-retry-auth-" + uuid.NewString() + auth := &Auth{ + ID: authID, + Provider: "claude", + Attributes: map[string]string{"auth_kind": "oauth"}, + Metadata: map[string]any{ + "access_token": "access-token", + "refresh_token": "refresh-token", + "disable_cooling": disableCooling, + }, + } + registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: model}}) + if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { + t.Fatalf("Register() error = %v", errRegister) + } + ids = append(ids, auth.ID) + } + return manager, ids, model +} + +func TestOuterRetryExclusions_NonStream_DisableCooling_InvokedOnce(t *testing.T) { + exec := newOuterRetryTestExecutor() + manager, ids, model := newOuterRetryTestManager(t, exec, 1, true) + // A 429 carrying Retry-After plus a positive max-wait interval creates a + // real outer retry opportunity (each outer attempt would re-invoke the + // executor before the exclusion fix). disable_cooling keeps the auth + // eligible across attempts so the loop is not short-circuited by cooldown. + manager.SetRetryConfig(3, 100*time.Millisecond, 3) + exec.executeErrs[ids[0]] = &retryAfterStatusError{ + status: http.StatusTooManyRequests, + message: "quota exhausted", + retryAfter: 5 * time.Millisecond, + } + + req := cliproxyexecutor.Request{Model: model} + opts := cliproxyexecutor.Options{} + + _, err := manager.Execute(context.Background(), []string{"claude"}, req, opts) + if err == nil { + t.Fatal("Execute() expected error, got nil") + } + + exec.mu.Lock() + calls := exec.executeCalls[ids[0]] + exec.mu.Unlock() + + if calls != 1 { + t.Fatalf("auth executed %d times across outer retries, want exactly 1", calls) + } +} + +func TestOuterRetryExclusions_Stream_DisableCooling_InvokedOnce(t *testing.T) { + exec := newOuterRetryTestExecutor() + manager, ids, model := newOuterRetryTestManager(t, exec, 1, true) + // Same non-vacuous setup as the non-stream variant: a 429 + Retry-After + // with a positive max-wait interval guarantees the outer retry loop has a + // real opportunity to re-invoke the executor before the exclusion fix. + manager.SetRetryConfig(3, 100*time.Millisecond, 3) + exec.streamErrs[ids[0]] = &retryAfterStatusError{ + status: http.StatusTooManyRequests, + message: "quota exhausted", + retryAfter: 5 * time.Millisecond, + } + + req := cliproxyexecutor.Request{Model: model} + opts := cliproxyexecutor.Options{} + + _, err := manager.ExecuteStream(context.Background(), []string{"claude"}, req, opts) + if err == nil { + t.Fatal("ExecuteStream() expected error, got nil") + } + + exec.mu.Lock() + calls := exec.streamCalls[ids[0]] + exec.mu.Unlock() + + if calls != 1 { + t.Fatalf("auth streamed %d times across outer retries, want exactly 1", calls) + } +} + +func TestOuterRetryExclusions_FailedAuth_FallbackToHealthyNextAuth(t *testing.T) { + exec := newOuterRetryTestExecutor() + exec.failFirstN = 1 + manager, _, model := newOuterRetryTestManager(t, exec, 2, true) + + req := cliproxyexecutor.Request{Model: model} + opts := cliproxyexecutor.Options{} + + resp, err := manager.Execute(context.Background(), []string{"claude"}, req, opts) + if err != nil { + t.Fatalf("Execute() error = %v, expected success on second auth", err) + } + if string(resp.Payload) == "" { + t.Fatal("Execute() returned empty payload") + } + + exec.mu.Lock() + total := exec.totalCalls + exec.mu.Unlock() + + if total != 2 { + t.Fatalf("total execute calls = %d, want 2 (1 failed + 1 healthy)", total) + } +} + +func TestOuterRetryExclusions_PreservesCallerSuppliedExclusions(t *testing.T) { + exec := newOuterRetryTestExecutor() + exec.failFirstN = 1 + manager, ids, model := newOuterRetryTestManager(t, exec, 3, false) + + req := cliproxyexecutor.Request{Model: model} + opts := cliproxyexecutor.Options{ + Metadata: map[string]any{ + cliproxyexecutor.ExcludedAuthIDsMetadataKey: map[string]struct{}{ + ids[0]: {}, + }, + }, + } + + resp, err := manager.Execute(context.Background(), []string{"claude"}, req, opts) + if err != nil { + t.Fatalf("Execute() error = %v, expected success on third auth", err) + } + if string(resp.Payload) == "" { + t.Fatal("Execute() returned empty payload") + } + + exec.mu.Lock() + c0 := exec.executeCalls[ids[0]] + total := exec.totalCalls + exec.mu.Unlock() + + if c0 != 0 { + t.Fatalf("caller-excluded auth executed %d times, want 0", c0) + } + if total != 2 { + t.Fatalf("total execute calls = %d, want 2", total) + } +} + +func TestOuterRetryExclusions_CoolingEnabled_RemainsGreen(t *testing.T) { + exec := newOuterRetryTestExecutor() + exec.failFirstN = 1 + manager, _, model := newOuterRetryTestManager(t, exec, 2, false) + + req := cliproxyexecutor.Request{Model: model} + opts := cliproxyexecutor.Options{} + + _, err := manager.Execute(context.Background(), []string{"claude"}, req, opts) + if err != nil { + t.Fatalf("Execute() error = %v, expected fallback success with cooling enabled", err) + } + + exec.mu.Lock() + total := exec.totalCalls + exec.mu.Unlock() + + if total != 2 { + t.Fatalf("total execute calls = %d, want 2", total) + } +} diff --git a/sdk/cliproxy/auth/request_termination_test.go b/sdk/cliproxy/auth/request_termination_test.go index 3ebd92e3b..16c1ff87f 100644 --- a/sdk/cliproxy/auth/request_termination_test.go +++ b/sdk/cliproxy/auth/request_termination_test.go @@ -1,6 +1,8 @@ package auth import ( + "context" + "errors" "net/http" "testing" @@ -16,3 +18,61 @@ func TestRequestTerminatedErrorSkipsCreditsFallback(t *testing.T) { t.Fatal("terminated request must not use Antigravity credits fallback") } } + +type afterAuthTestExecutor struct{} + +func (*afterAuthTestExecutor) Identifier() string { return "afterauth" } + +func (*afterAuthTestExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} + +func (*afterAuthTestExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + ch := make(chan cliproxyexecutor.StreamChunk) + close(ch) + return &cliproxyexecutor.StreamResult{Chunks: ch}, nil +} + +func (*afterAuthTestExecutor) Refresh(_ context.Context, auth *Auth) (*Auth, error) { return auth, nil } + +func (*afterAuthTestExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} + +func (*afterAuthTestExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +func TestRequestTerminatedErrorZeroValueIsUntrustedAndApplyAfterAuthSetsTrusted(t *testing.T) { + // Zero-value Trusted must be false (safe default for untrusted upstreams). + zero := &cliproxyexecutor.RequestTerminatedError{HTTPStatus: http.StatusBadGateway} + if zero.Trusted { + t.Fatal("zero-value RequestTerminatedError.Trusted must be false") + } + + // A termination produced by applyRequestAfterAuthInterceptor must be trusted. + executor := &afterAuthTestExecutor{} + appliedTerminate := false + interceptor := func(context.Context, cliproxyexecutor.RequestAfterAuthInterceptRequest) cliproxyexecutor.RequestAfterAuthInterceptResponse { + appliedTerminate = true + return cliproxyexecutor.RequestAfterAuthInterceptResponse{ + Terminate: true, + StatusCode: http.StatusTeapot, + ResponseHeaders: http.Header{"X-AfterAuth": []string{"yes"}}, + ResponseBody: []byte(`{"ok":true}`), + } + } + _, _, errIntercept := applyRequestAfterAuthInterceptor(context.Background(), executor, "test", cliproxyexecutor.Request{}, cliproxyexecutor.Options{ + RequestAfterAuthInterceptor: interceptor, + }, "model") + var terminated *cliproxyexecutor.RequestTerminatedError + if !errors.As(errIntercept, &terminated) || terminated == nil { + t.Fatalf("applyRequestAfterAuthInterceptor error = %v, want RequestTerminatedError", errIntercept) + } + if !terminated.Trusted { + t.Fatal("after-auth interceptor termination must set Trusted=true") + } + if !appliedTerminate { + t.Fatal("interceptor was not applied") + } +} diff --git a/sdk/cliproxy/auth/route_exhaustion_test.go b/sdk/cliproxy/auth/route_exhaustion_test.go new file mode 100644 index 000000000..5905c7480 --- /dev/null +++ b/sdk/cliproxy/auth/route_exhaustion_test.go @@ -0,0 +1,782 @@ +package auth + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "strings" + "testing" + "time" + + "github.com/google/uuid" + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + executionregistry "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +type routeExhaustionTestExecutor struct { + provider string + failErrors map[string]error +} + +func newRouteExhaustionTestExecutor(provider string) *routeExhaustionTestExecutor { + return &routeExhaustionTestExecutor{ + provider: provider, + failErrors: make(map[string]error), + } +} + +func (e *routeExhaustionTestExecutor) Identifier() string { return e.provider } +func (*routeExhaustionTestExecutor) ShouldPrepareRequestAuth(*Auth) bool { return false } +func (e *routeExhaustionTestExecutor) PrepareRequestAuth(_ context.Context, auth *Auth) (*Auth, error) { + return auth, nil +} +func (e *routeExhaustionTestExecutor) Execute(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + if err, ok := e.failErrors[auth.ID]; ok && err != nil { + return cliproxyexecutor.Response{}, err + } + return cliproxyexecutor.Response{Payload: []byte(`{"choices":[{"message":{"content":"ok"}}]}`)}, nil +} +func (e *routeExhaustionTestExecutor) ExecuteStream(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + if err, ok := e.failErrors[auth.ID]; ok && err != nil { + return nil, err + } + ch := make(chan cliproxyexecutor.StreamChunk, 1) + ch <- cliproxyexecutor.StreamChunk{Payload: []byte(`data: {"choices":[{"delta":{"content":"ok"}}]}\n\n`)} + close(ch) + return &cliproxyexecutor.StreamResult{Chunks: ch}, nil +} +func (e *routeExhaustionTestExecutor) Refresh(_ context.Context, auth *Auth) (*Auth, error) { + return auth, nil +} +func (*routeExhaustionTestExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, errors.New("not implemented") +} +func (e *routeExhaustionTestExecutor) CountTokens(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + if err, ok := e.failErrors[auth.ID]; ok && err != nil { + return cliproxyexecutor.Response{}, err + } + return cliproxyexecutor.Response{}, nil +} + +func registerRouteTestAuth(t *testing.T, mgr *Manager, provider string, model string, errToReturn error) string { + t.Helper() + authID := fmt.Sprintf("%s-auth-%s", provider, uuid.NewString()) + auth := &Auth{ + ID: authID, + Provider: provider, + Attributes: map[string]string{"disable_cooling": "true"}, + } + registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: model}}) + if _, err := mgr.Register(context.Background(), auth); err != nil { + t.Fatalf("Register() error = %v", err) + } + return authID +} + +// A. Non-stream public exhaustion with three routes/classes/statuses +func TestRouteExhaustion_ThreeRoutes(t *testing.T) { + mgr := NewManager(nil, nil, nil) + mgr.SetRetryConfig(3, 5*time.Second, 3) + + execClaude := newRouteExhaustionTestExecutor("claude") + execGemini := newRouteExhaustionTestExecutor("gemini") + execCodex := newRouteExhaustionTestExecutor("codex") + + mgr.RegisterExecutor(execClaude) + mgr.RegisterExecutor(execGemini) + mgr.RegisterExecutor(execCodex) + + model := "test-model-" + uuid.NewString() + + id1 := registerRouteTestAuth(t, mgr, "claude", model, nil) + id2 := registerRouteTestAuth(t, mgr, "gemini", model, nil) + id3 := registerRouteTestAuth(t, mgr, "codex", model, nil) + + execClaude.failErrors[id1] = &Error{Code: "rate_limit", Message: "429 Rate Limit", HTTPStatus: 429} + execGemini.failErrors[id2] = &Error{Code: "rate_limit", Message: "429 Rate Limit", HTTPStatus: 429} + execCodex.failErrors[id3] = &Error{Code: "rate_limit", Message: "429 Rate Limit", HTTPStatus: 429} + + _, err := mgr.Execute(context.Background(), []string{"claude", "gemini", "codex"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}) + if err == nil { + t.Fatalf("Execute() expected error, got nil") + } + + errStr := err.Error() + if !strings.Contains(errStr, "attempted routes:") || !strings.Contains(errStr, "claude:429") || !strings.Contains(errStr, "gemini:429") || !strings.Contains(errStr, "codex:429") { + t.Errorf("unexpected error string: %s", errStr) + } + + var authErr *Error + if !errors.As(err, &authErr) || authErr == nil { + t.Fatalf("errors.As(*Error) failed") + } + if authErr.HTTPStatus != 429 { + t.Errorf("expected status 429, got %d", authErr.HTTPStatus) + } +} + +// B. Security redaction +func TestRouteExhaustion_SecurityRedaction(t *testing.T) { + mgr := NewManager(nil, nil, nil) + mgr.SetRetryConfig(3, 5*time.Second, 3) + + exec := newRouteExhaustionTestExecutor("gemini") + mgr.RegisterExecutor(exec) + + model := "redaction-model-" + uuid.NewString() + + secretAuthID := "secret-auth-id-999" + fakeKey := "sk-secret-api-key-12345" + emailFilename := "user@secret.com.json" + baseURL := "https://internal.secret.net/v1" + privateModel := "private-alias-99" + + auth := &Auth{ + ID: secretAuthID, + Provider: "gemini", + Attributes: map[string]string{"api_key": fakeKey, "credential_file": emailFilename, "base_url": baseURL, "disable_cooling": "true"}, + Metadata: map[string]any{"private_model": privateModel}, + } + registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: model}}) + if _, err := mgr.Register(context.Background(), auth); err != nil { + t.Fatalf("Register() error = %v", err) + } + + exec.failErrors[secretAuthID] = &Error{Code: "upstream_error", Message: "502 Bad Gateway to " + baseURL, HTTPStatus: 502} + + _, err := mgr.Execute(context.Background(), []string{"gemini"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}) + if err == nil { + t.Fatalf("Execute() expected error") + } + + summaryIdx := strings.Index(err.Error(), "attempted routes:") + if summaryIdx == -1 { + t.Fatalf("summary missing in error: %s", err.Error()) + } + summaryPart := err.Error()[summaryIdx:] + + sensitiveTokens := []string{secretAuthID, fakeKey, emailFilename, baseURL, privateModel} + for _, tok := range sensitiveTokens { + if strings.Contains(summaryPart, tok) { + t.Errorf("sensitive token %q leaked in summary part: %s", tok, summaryPart) + } + } +} + +// C. Duplicate and Cap behavior +func TestRouteExhaustion_DedupAndCap(t *testing.T) { + tracker := newRouteAttemptTracker() + + // Dedup test + authGemini := &Auth{Provider: "gemini"} + authCodex := &Auth{Provider: "codex"} + err502 := &Error{HTTPStatus: 502} + err429 := &Error{HTTPStatus: 429} + + tracker.Record(authGemini, err502) + tracker.Record(authGemini, err502) // dup + tracker.Record(authCodex, err429) + tracker.Record(authCodex, err429) // dup + + if summary := tracker.Summary(); summary != "attempted routes: [gemini:502, codex:429]" { + t.Errorf("dedup failed: %s", summary) + } + + // Cap test (> 16 unique attempts) + bigTracker := newRouteAttemptTracker() + for i := 0; i < 20; i++ { + provider := fmt.Sprintf("provider-%d", i) + bigTracker.Record(&Auth{Provider: provider}, &Error{HTTPStatus: 400 + i}) + } + bigSummary := bigTracker.Summary() + if !strings.Contains(bigSummary, "... (+4 omitted)") { + t.Errorf("expected omitted count in summary, got: %s", bigSummary) + } +} + +// D. No-candidate path: exact existing auth_not_found behavior unchanged +func TestRouteExhaustion_NoCandidatePath(t *testing.T) { + mgr := NewManager(nil, nil, nil) + _, err := mgr.Execute(context.Background(), []string{"unknown-provider"}, cliproxyexecutor.Request{Model: "nonexistent"}, cliproxyexecutor.Options{}) + if err == nil { + t.Fatalf("expected error") + } + if strings.Contains(err.Error(), "attempted routes:") { + t.Errorf("no candidate path should not contain attempted routes summary, got: %s", err.Error()) + } +} + +// E. ExecuteCount full exhaustion summary +func TestRouteExhaustion_ExecuteCount(t *testing.T) { + mgr := NewManager(nil, nil, nil) + exec := newRouteExhaustionTestExecutor("openai") + mgr.RegisterExecutor(exec) + + model := "count-model-" + uuid.NewString() + authID := registerRouteTestAuth(t, mgr, "openai", model, nil) + exec.failErrors[authID] = &Error{Code: "bad_gateway", Message: "502 Bad Gateway", HTTPStatus: 502} + + _, err := mgr.ExecuteCount(context.Background(), []string{"openai"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}) + if err == nil { + t.Fatalf("expected error") + } + if !strings.Contains(err.Error(), "attempted routes: [openai:502]") { + t.Errorf("ExecuteCount error missing summary: %s", err.Error()) + } +} + +// F. Pre-commit ExecuteStream exhaustion summary +func TestRouteExhaustion_ExecuteStream(t *testing.T) { + mgr := NewManager(nil, nil, nil) + exec := newRouteExhaustionTestExecutor("claude") + mgr.RegisterExecutor(exec) + + model := "stream-model-" + uuid.NewString() + authID := registerRouteTestAuth(t, mgr, "claude", model, nil) + exec.failErrors[authID] = &Error{Code: "rate_limit", Message: "429 Too Many Requests", HTTPStatus: 429} + + _, err := mgr.ExecuteStream(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}) + if err == nil { + t.Fatalf("ExecuteStream() expected error") + } + if !strings.Contains(err.Error(), "attempted routes: [claude:429]") { + t.Errorf("ExecuteStream error missing summary, got err=%v", err) + } +} + +// G. Healthy fallback success: no summary leaks into successful response +func TestRouteExhaustion_HealthyFallbackSuccess(t *testing.T) { + mgr := NewManager(nil, nil, nil) + execClaude := newRouteExhaustionTestExecutor("claude") + execGemini := newRouteExhaustionTestExecutor("gemini") + + mgr.RegisterExecutor(execClaude) + mgr.RegisterExecutor(execGemini) + + model := "fallback-model-" + uuid.NewString() + id1 := registerRouteTestAuth(t, mgr, "claude", model, nil) + _ = registerRouteTestAuth(t, mgr, "gemini", model, nil) + + execClaude.failErrors[id1] = &Error{Code: "bad_gateway", Message: "502 Bad Gateway", HTTPStatus: 502} + + resp, err := mgr.Execute(context.Background(), []string{"claude", "gemini"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}) + if err != nil { + t.Fatalf("Execute() unexpected error = %v", err) + } + if len(resp.Payload) == 0 { + t.Fatalf("empty response payload") + } +} + +// H. Request-invalid / cancellation early abort remains unchanged +func TestRouteExhaustion_RequestInvalidEarlyAbort(t *testing.T) { + mgr := NewManager(nil, nil, nil) + exec := newRouteExhaustionTestExecutor("claude") + mgr.RegisterExecutor(exec) + + model := "invalid-model-" + uuid.NewString() + id1 := registerRouteTestAuth(t, mgr, "claude", model, nil) + id2 := registerRouteTestAuth(t, mgr, "claude", model, nil) + + exec.failErrors[id1] = &Error{Code: "invalid_request", Message: "400 Bad Request", HTTPStatus: 400} + exec.failErrors[id2] = &Error{Code: "invalid_request", Message: "400 Bad Request", HTTPStatus: 400} + + _, err := mgr.Execute(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}) + if err == nil { + t.Fatalf("expected error") + } + var authErr *Error + if !errors.As(err, &authErr) || authErr.HTTPStatus != 400 { + t.Errorf("expected status 400, got: %v", err) + } + if strings.Contains(err.Error(), "attempted routes:") { + t.Errorf("request invalid early abort should not contain attempted routes summary, got: %s", err.Error()) + } +} + +type routeExhaustionHomeDispatcher struct { + responses map[int]string + callCount int +} + +func (d *routeExhaustionHomeDispatcher) HeartbeatOK() bool { return true } +func (d *routeExhaustionHomeDispatcher) AbortAmbiguousDispatch() {} + +func (d *routeExhaustionHomeDispatcher) RPopAuth(_ context.Context, _ string, _ string, _ http.Header, _ int) ([]byte, error) { + resp, ok := d.responses[d.callCount] + d.callCount++ + if !ok || resp == "" { + if d.callCount > 1 && d.responses[d.callCount-2] != "" { + return []byte(d.responses[d.callCount-2]), nil + } + return nil, errors.New("no home auth available") + } + return []byte(resp), nil +} + +func TestRouteExhaustion_HomeMode(t *testing.T) { + mgr := NewManager(nil, nil, nil) + mgr.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + mgr.SetRetryConfig(3, 5*time.Second, 3) + + execClaude := newRouteExhaustionTestExecutor("claude") + execGemini := newRouteExhaustionTestExecutor("gemini") + mgr.RegisterExecutor(execClaude) + mgr.RegisterExecutor(execGemini) + + model := "home-model-" + uuid.NewString() + + authID1 := "home-secret-auth-1" + key1 := "sk-home-secret-key-111" + file1 := "user1@home-secret.com.json" + url1 := "https://home1.secret.net/v1" + alias1 := "secret-home-alias-1" + + authID2 := "home-secret-auth-2" + key2 := "sk-home-secret-key-222" + file2 := "user2@home-secret.com.json" + url2 := "https://home2.secret.net/v1" + alias2 := "secret-home-alias-2" + + payload1 := fmt.Sprintf(`{"provider":"claude","auth":{"id":%q,"provider":"claude","status":"active","attributes":{"api_key":%q,"credential_file":%q,"base_url":%q,"disable_cooling":"true"},"metadata":{"private_model":%q}}}`, authID1, key1, file1, url1, alias1) + payload2 := fmt.Sprintf(`{"provider":"gemini","auth":{"id":%q,"provider":"gemini","status":"active","attributes":{"api_key":%q,"credential_file":%q,"base_url":%q,"disable_cooling":"true"},"metadata":{"private_model":%q}}}`, authID2, key2, file2, url2, alias2) + + execClaude.failErrors[authID1] = &Error{Code: "rate_limit", Message: "429 Rate Limit", HTTPStatus: 429} + execGemini.failErrors[authID2] = &Error{Code: "bad_gateway", Message: "502 Bad Gateway", HTTPStatus: 502} + + t.Run("Exhaustion", func(t *testing.T) { + dispatcher := &routeExhaustionHomeDispatcher{ + responses: map[int]string{ + 0: payload1, + 1: payload2, + }, + } + mgr.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + + _, err := mgr.Execute(context.Background(), []string{"claude", "gemini"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}) + if err == nil { + t.Fatalf("Execute() expected error on home route exhaustion") + } + + errStr := err.Error() + summaryIdx := strings.Index(errStr, "attempted routes:") + if summaryIdx == -1 { + t.Fatalf("summary missing in home route exhaustion error: %s", errStr) + } + summaryPart := errStr[summaryIdx:] + + if !strings.Contains(summaryPart, "claude:429") || !strings.Contains(summaryPart, "gemini:502") { + t.Errorf("expected attempted routes [claude:429, gemini:502] in summary, got: %s", summaryPart) + } + + var authErr *Error + if !errors.As(err, &authErr) || authErr == nil { + t.Fatalf("errors.As(*Error) failed on home route exhaustion error") + } + if authErr.HTTPStatus != 502 { + t.Errorf("expected preserved cause status 502, got %d", authErr.HTTPStatus) + } + if authErr.Code != "bad_gateway" { + t.Errorf("expected preserved cause code bad_gateway, got %s", authErr.Code) + } + + secrets := []string{authID1, authID2, key1, key2, file1, file2, url1, url2, alias1, alias2} + for _, sec := range secrets { + if strings.Contains(summaryPart, sec) { + t.Errorf("sensitive secret %q leaked in home route summary: %s", sec, summaryPart) + } + } + }) + + t.Run("HealthyFallbackSuccess", func(t *testing.T) { + execGemini.failErrors[authID2] = nil + dispatcher := &routeExhaustionHomeDispatcher{ + responses: map[int]string{ + 0: payload1, + 1: payload2, + }, + } + mgr.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + + resp, err := mgr.Execute(context.Background(), []string{"claude", "gemini"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}) + if err != nil { + t.Fatalf("Execute() unexpected error on healthy home fallback: %v", err) + } + if len(resp.Payload) == 0 { + t.Errorf("expected non-empty payload on home fallback success") + } + }) +} + +type routeExhaustionHomeNoModelDispatcher struct { + payloads []string + count int +} + +func (d *routeExhaustionHomeNoModelDispatcher) HeartbeatOK() bool { return true } +func (d *routeExhaustionHomeNoModelDispatcher) AbortAmbiguousDispatch() {} +func (d *routeExhaustionHomeNoModelDispatcher) RPopAuth(_ context.Context, _ string, _ string, _ http.Header, _ int) ([]byte, error) { + if d.count >= len(d.payloads) { + return nil, errors.New("no home auth available") + } + next := d.payloads[d.count] + d.count++ + if next == "" { + return nil, errors.New("no home auth available") + } + return []byte(next), nil +} + +// I. Home no-execution-models path: a dispatched auth blocked for its upstream model yields a +// sanitized attempted-routes diagnostic (records the no_model outcome) instead of an empty error. +func TestRouteExhaustion_HomeNoExecutionModelsDiagnostic(t *testing.T) { + mgr := NewManager(nil, nil, nil) + mgr.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + mgr.SetRetryConfig(3, 5*time.Second, 3) + + exec := newRouteExhaustionTestExecutor("claude") + mgr.RegisterExecutor(exec) + + authID := "home-no-model-auth" + fakeKey := "sk-home-no-model-secret-123" + credFile := "user-no-model@home-secret.com.json" + + payload := fmt.Sprintf( + `{"model":"upstream-blocked-1","provider":"claude","auth":{"id":%q,"provider":"claude","status":"active","attributes":{"api_key":%q,"credential_file":%q,"disable_cooling":"true"},"unavailable":true,"next_retry_after":"2030-01-02T03:04:05Z"}}`, + authID, fakeKey, credFile, + ) + + dispatcher := &routeExhaustionHomeNoModelDispatcher{payloads: []string{payload}} + mgr.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + + _, err := mgr.Execute(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: "claude-3-5-sonnet"}, cliproxyexecutor.Options{}) + if err == nil { + t.Fatalf("Execute() expected error on home no-execution-models route") + } + + summaryIdx := strings.Index(err.Error(), "attempted routes:") + if summaryIdx == -1 { + t.Fatalf("no-model home path missing attempted-routes diagnostic: %s", err.Error()) + } + summaryPart := err.Error()[summaryIdx:] + + if !strings.Contains(summaryPart, "claude:") { + t.Errorf("expected claude provider recorded in summary, got: %s", summaryPart) + } + for _, secret := range []string{authID, fakeKey, credFile} { + if strings.Contains(summaryPart, secret) { + t.Errorf("sensitive token %q leaked in home no-model summary: %s", secret, summaryPart) + } + } +} + +// routeExhaustionHeaderCause is a generic non-*Error cause that exposes +// headers the same way upstream errors (streamBootstrapError, +// modelCooldownError) do, so handlers collecting passthrough headers from the +// final routed error can still surface them through route exhaustion. +type routeExhaustionHeaderCause struct { + msg string + headers http.Header +} + +func (e *routeExhaustionHeaderCause) Error() string { return e.msg } +func (e *routeExhaustionHeaderCause) Headers() http.Header { + return e.headers.Clone() +} + +// routeExhaustionNoHeaderCause is a generic non-*Error cause that exposes no +// headers, mirroring ordinary upstream failures. +type routeExhaustionNoHeaderCause struct{ msg string } + +func (e *routeExhaustionNoHeaderCause) Error() string { return e.msg } + +// J. Wrapper contract: a wrapped cause exposing headers retains them, while +// Error/Unwrap and the sanitized route summary stay intact. +func TestRouteExhaustion_HeadersForwarded(t *testing.T) { + tracker := newRouteAttemptTracker() + tracker.Record(&Auth{Provider: "gemini"}, &Error{HTTPStatus: 429}) + + cause := &routeExhaustionHeaderCause{ + msg: "upstream retry-after", + headers: http.Header{"Retry-After": {"1"}, "X-Request-Id": {"req-123"}}, + } + err := wrapRouteExhaustion(cause, tracker) + + // errors.As / errors.Is must still traverse the wrapper. + var unwrapped *routeExhaustionHeaderCause + if !errors.As(err, &unwrapped) || unwrapped == nil { + t.Fatalf("errors.As(*routeExhaustionHeaderCause) failed, err=%v", err) + } + if !errors.Is(err, cause) { + t.Fatalf("errors.Is(err, cause) = false") + } + + // Sanitized route summary retained. + if !strings.Contains(err.Error(), "attempted routes: [gemini") { + t.Errorf("unexpected error string, summary missing routes: %s", err.Error()) + } + + // Headers readable via the same assertion handlers use; values exact. + he, ok := err.(interface{ Headers() http.Header }) + if !ok || he == nil { + t.Fatalf("routeExhaustionError must implement Headers(), err=%T", err) + } + hdr := he.Headers() + if hdr.Get("Retry-After") != "1" { + t.Errorf("Headers().Get(Retry-After) = %q, want 1", hdr.Get("Retry-After")) + } + if hdr.Get("X-Request-Id") != "req-123" { + t.Errorf("Headers().Get(X-Request-Id) = %q, want req-123", hdr.Get("X-Request-Id")) + } + + // Forwarded map is a copy: mutating it must not touch the caller's map. + hdr.Set("Retry-After", "999") + if cause.headers.Get("Retry-After") != "1" { + t.Errorf("wrapped headers mutated caller map: got %q", cause.headers.Get("Retry-After")) + } +} + +// K. Wrapper contract: a cause without Headers yields nil, matching the +// convention other header-carriers follow for absent headers. +func TestRouteExhaustion_HeadersAbsentNil(t *testing.T) { + tracker := newRouteAttemptTracker() + tracker.Record(&Auth{Provider: "openai"}, &Error{HTTPStatus: 502}) + + err := wrapRouteExhaustion(&routeExhaustionNoHeaderCause{msg: "502 Bad Gateway"}, tracker) + + he, ok := err.(interface{ Headers() http.Header }) + if !ok || he == nil { + t.Fatalf("routeExhaustionError must implement Headers(), err=%T", err) + } + if got := he.Headers(); got != nil { + t.Errorf("Headers() = %v, want nil for cause without headers", got) + } +} + +// L. Stream-level: headers reach the returned route-exhaustion error so the +// stream/error handlers can surface passthrough headers. +func TestRouteExhaustion_ExecuteStreamHeaders(t *testing.T) { + mgr := NewManager(nil, nil, nil) + mgr.SetRetryConfig(3, 5*time.Second, 3) + + exec := newRouteExhaustionTestExecutor("claude") + mgr.RegisterExecutor(exec) + + model := "stream-hdr-model-" + uuid.NewString() + authID := registerRouteTestAuth(t, mgr, "claude", model, nil) + exec.failErrors[authID] = &routeExhaustionHeaderCause{ + msg: "429 Too Many Requests", + headers: http.Header{"Retry-After": {"1"}, "X-Request-Id": {"req-777"}}, + } + + _, err := mgr.ExecuteStream(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}) + if err == nil { + t.Fatalf("ExecuteStream() expected error on route exhaustion") + } + if !strings.Contains(err.Error(), "attempted routes: [claude") { + t.Errorf("ExecuteStream error missing route summary: %v", err) + } + + he, ok := err.(interface{ Headers() http.Header }) + if !ok || he == nil { + t.Fatalf("ExecuteStream() error must implement Headers(), got %T", err) + } + hdr := he.Headers() + if hdr.Get("Retry-After") != "1" { + t.Errorf("Headers().Get(Retry-After) = %q, want 1", hdr.Get("Retry-After")) + } + if hdr.Get("X-Request-Id") != "req-777" { + t.Errorf("Headers().Get(X-Request-Id) = %q, want req-777", hdr.Get("X-Request-Id")) + } +} + +// routeExhaustionNestedCause is a header-carrier that also unwraps to an inner +// header-carrier, so the first/outermost carrier must win per errors.As. +type routeExhaustionNestedCause struct { + inner *routeExhaustionHeaderCause + headers http.Header +} + +func (e *routeExhaustionNestedCause) Error() string { return "outer wrapped cause" } +func (e *routeExhaustionNestedCause) Unwrap() error { return e.inner } +func (e *routeExhaustionNestedCause) Headers() http.Header { return e.headers } + +// M. Wrapper contract: nil receiver returns nil, not a panic. +func TestRouteExhaustion_HeadersNilReceiver(t *testing.T) { + var e *routeExhaustionClonedError + if hdr := e.Headers(); hdr != nil { + t.Errorf("Headers() = %v, want nil for nil receiver", hdr) + } + if hdr := e.SafeResponseHeaders(); hdr != nil { + t.Errorf("SafeResponseHeaders() = %v, want nil for nil receiver", hdr) + } +} + +// N. Wrapper contract: errors.As starts at the cause and returns the +// first/outermost carrier; an inner carrier must not shadow it, and the +// forwarded map is a fresh clone even when the outer carrier returns raw. +func TestRouteExhaustion_HeadersNestedOutermostWins(t *testing.T) { + tracker := newRouteAttemptTracker() + tracker.Record(&Auth{Provider: "gemini"}, &Error{HTTPStatus: 429}) + + inner := &routeExhaustionHeaderCause{ + msg: "inner retry-after", + headers: http.Header{"Retry-After": {"inner"}, "Inner": {"1"}}, + } + outer := &routeExhaustionNestedCause{ + inner: inner, + headers: http.Header{"Retry-After": {"outer"}, "Outer": {"1"}}, + } + err := wrapRouteExhaustion(outer, tracker) + + he, ok := err.(interface{ Headers() http.Header }) + if !ok || he == nil { + t.Fatalf("routeExhaustionError must implement Headers(), err=%T", err) + } + hdr := he.Headers() + if hdr.Get("Retry-After") != "outer" { + t.Errorf("Headers().Get(Retry-After) = %q, want outer", hdr.Get("Retry-After")) + } + if hdr.Get("Outer") != "1" { + t.Errorf("Headers().Get(Outer) = %q, want 1", hdr.Get("Outer")) + } + if hdr.Get("Inner") != "" { + t.Errorf("headers from inner carrier leaked, outermost must win: %q", hdr.Get("Inner")) + } + + // The inner carrier remains reachable through the Unwrap chain. + var unwrapped *routeExhaustionHeaderCause + if !errors.As(err, &unwrapped) || unwrapped == nil { + t.Fatalf("errors.As(*routeExhaustionHeaderCause) failed, err=%v", err) + } + + // Forwarded map is a fresh clone of the outer cause's raw map. + hdr.Set("Retry-After", "999") + if outer.headers.Get("Retry-After") != "outer" { + t.Errorf("wrapped headers mutated caller map: got %q", outer.headers.Get("Retry-After")) + } + if inner.headers.Get("Retry-After") != "inner" { + t.Errorf("inner caller map mutated: got %q", inner.headers.Get("Retry-After")) + } +} + +// O. Wrapper contract: the Error message appends the sanitized summary exactly +// once and never mutates the cause's *Error.Message. +func TestRouteExhaustion_SummaryAppendedOnceMessageUnchanged(t *testing.T) { + tracker := newRouteAttemptTracker() + tracker.Record(&Auth{Provider: "gemini"}, &Error{HTTPStatus: 429}) + tracker.Record(&Auth{Provider: "claude"}, &Error{HTTPStatus: 502}) + + original := &Error{ + Code: "bad_gateway", + Message: "502 Bad Gateway", + Retryable: true, + HTTPStatus: 502, + } + var originalMessage = original.Message + + err := wrapRouteExhaustion(original, tracker) + if err == nil || err == original { + t.Fatalf("wrapRouteExhaustion must return a wrapper (got %T, same=%v)", err, err == original) + } + + // Message contains the summary exactly once, appended after the cause. + got := err.Error() + if strings.Count(got, "attempted routes:") != 1 { + t.Errorf("summary appended more than once, got: %s", got) + } + wantPrefix := original.Error() + if !strings.HasPrefix(got, wantPrefix+"; ") { + t.Errorf("message should start with cause message, got: %s", got) + } + if !strings.Contains(got, "attempted routes: [gemini:429, claude:502]") { + t.Errorf("expected both recorded routes in summary, got: %s", got) + } + + // The cause's stored Message is never mutated. + if original.Message != originalMessage { + t.Errorf("cause *Error.Message mutated: %q -> %q", originalMessage, original.Message) + } + + // The wrapper unwraps the ORIGINAL cause, so identity/typed access is preserved. + var authErr *Error + if !errors.As(err, &authErr) || authErr == nil { + t.Fatalf("errors.As(*Error) failed on wrapped error") + } + if authErr.HTTPStatus != 502 || !authErr.Retryable || authErr.Code != "bad_gateway" { + t.Errorf("wrapped cause lost typed fields, got code=%q status=%d retry=%v", authErr.Code, authErr.HTTPStatus, authErr.Retryable) + } + if !errors.Is(err, original) { + t.Errorf("errors.Is(wrappedError, original) = false, want true (original cause preserved)") + } +} + +// P. Wrapper contract: SafeResponseHeaders forwards the Home busy error's trusted +// Retry-After through route exhaustion via the same access path handlers use. +func TestRouteExhaustion_SafeResponseHeadersForwarded(t *testing.T) { + tracker := newRouteAttemptTracker() + tracker.Record(&Auth{Provider: "gemini"}, &Error{HTTPStatus: 429}) + + cause := NewHomeConcurrencyBusyError("credential concurrency limit exceeded", 7*time.Second) + err := wrapRouteExhaustion(cause, tracker) + + se, ok := err.(interface{ SafeResponseHeaders() http.Header }) + if !ok || se == nil { + t.Fatalf("wrapped error must implement SafeResponseHeaders(), err=%T", err) + } + if got := se.SafeResponseHeaders().Get("Retry-After"); got != "7" { + t.Errorf("SafeResponseHeaders().Get(Retry-After) = %q, want 7", got) + } + + // Each call yields a fresh map: mutating one result must not affect the next + // (the underlying Home busy error keeps its own Retry-After). + first := se.SafeResponseHeaders() + first.Set("Retry-After", "999") + if got := se.SafeResponseHeaders().Get("Retry-After"); got != "7" { + t.Errorf("SafeResponseHeaders returned a shared map: got %q after mutating a prior call", got) + } +} + +// Q. Wrapper contract: a cause without SafeResponseHeaders yields nil. +func TestRouteExhaustion_SafeResponseHeadersAbsentNil(t *testing.T) { + tracker := newRouteAttemptTracker() + tracker.Record(&Auth{Provider: "openai"}, &Error{HTTPStatus: 502}) + + err := wrapRouteExhaustion(&routeExhaustionNoHeaderCause{msg: "502 Bad Gateway"}, tracker) + + se, ok := err.(interface{ SafeResponseHeaders() http.Header }) + if !ok || se == nil { + t.Fatalf("wrapped error must implement SafeResponseHeaders(), err=%T", err) + } + if got := se.SafeResponseHeaders(); got != nil { + t.Errorf("SafeResponseHeaders() = %v, want nil for cause without it", got) + } + + // Header cause (generic Headers only) also yields nil for SafeResponseHeaders. + headerErr := wrapRouteExhaustion(&routeExhaustionHeaderCause{msg: "429", headers: http.Header{"Retry-After": {"1"}}}, tracker) + if hse, ok := headerErr.(interface{ SafeResponseHeaders() http.Header }); ok && hse != nil { + if got := hse.SafeResponseHeaders(); got != nil { + t.Errorf("SafeResponseHeaders() = %v, want nil for Headers-only cause", got) + } + } +} + +func TestRouteExhaustion_PreservesStructuredJSONRequestFault(t *testing.T) { + tracker := newRouteAttemptTracker() + tracker.Record(&Auth{Provider: "openai"}, &Error{HTTPStatus: 502}) + + rawJSON := `{"error":{"type":"invalid_request_error","code":"cyber_policy","message":"blocked"}}` + cause := errors.New(rawJSON) + wrapped := wrapRouteExhaustion(cause, tracker) + + if !json.Valid([]byte(wrapped.Error())) { + t.Fatalf("wrapped.Error() corrupted structured JSON: %s", wrapped.Error()) + } + if wrapped.Error() != rawJSON { + t.Fatalf("wrapped.Error() = %q, want original %q", wrapped.Error(), rawJSON) + } +} diff --git a/sdk/cliproxy/auth/route_tracker.go b/sdk/cliproxy/auth/route_tracker.go new file mode 100644 index 000000000..871a582f5 --- /dev/null +++ b/sdk/cliproxy/auth/route_tracker.go @@ -0,0 +1,199 @@ +package auth + +import ( + "encoding/json" + "errors" + "fmt" + "net/http" + "strconv" + "strings" +) + +const maxRouteAttemptsRecorded = 16 + +type routeAttempt struct { + providerClass string + status string +} + +type routeAttemptTracker struct { + attempts []routeAttempt + seen map[routeAttempt]bool + omitted int +} + +func newRouteAttemptTracker() *routeAttemptTracker { + return &routeAttemptTracker{ + attempts: make([]routeAttempt, 0, 8), + seen: make(map[routeAttempt]bool), + } +} + +func (t *routeAttemptTracker) Record(auth *Auth, err error) { + if t == nil { + return + } + pClass := sanitizeProviderClass(authProviderName(auth)) + statusStr := sanitizeStatus(err) + entry := routeAttempt{ + providerClass: pClass, + status: statusStr, + } + if t.seen[entry] { + return + } + t.seen[entry] = true + if len(t.attempts) >= maxRouteAttemptsRecorded { + t.omitted++ + return + } + t.attempts = append(t.attempts, entry) +} + +func (t *routeAttemptTracker) Summary() string { + if t == nil || len(t.attempts) == 0 { + return "" + } + var sb strings.Builder + sb.WriteString("attempted routes: [") + for i, a := range t.attempts { + if i > 0 { + sb.WriteString(", ") + } + sb.WriteString(a.providerClass) + sb.WriteString(":") + sb.WriteString(a.status) + } + if t.omitted > 0 { + sb.WriteString(fmt.Sprintf(", ... (+%d omitted)", t.omitted)) + } + sb.WriteString("]") + return sb.String() +} + +func authProviderName(auth *Auth) string { + if auth == nil { + return "" + } + return auth.Provider +} + +func sanitizeProviderClass(p string) string { + switch strings.ToLower(strings.TrimSpace(p)) { + case "gemini": + return "gemini" + case "claude", "anthropic": + return "claude" + case "openai": + return "openai" + case "codex": + return "codex" + case "antigravity": + return "antigravity" + case "aistudio": + return "aistudio" + case "vertex", "vertexai", "vertex_ai": + return "vertex" + default: + return "other" + } +} + +func sanitizeStatus(err error) string { + if err == nil { + return "error" + } + var authErr *Error + if errors.As(err, &authErr) && authErr != nil { + if authErr.HTTPStatus > 0 && authErr.HTTPStatus < 1000 { + return strconv.Itoa(authErr.HTTPStatus) + } + } + if sc := statusCodeFromError(err); sc > 0 && sc < 1000 { + return strconv.Itoa(sc) + } + return "error" +} + +type routeExhaustionClonedError struct { + cause error + summary string +} + +func wrapRouteExhaustion(cause error, tracker *routeAttemptTracker) error { + if cause == nil { + return nil + } + if tracker == nil { + return cause + } + summary := tracker.Summary() + if summary == "" { + return cause + } + return &routeExhaustionClonedError{ + cause: cause, + summary: summary, + } +} + +func (e *routeExhaustionClonedError) Error() string { + if e == nil { + return "" + } + if e.cause == nil { + return e.summary + } + if e.summary == "" { + return e.cause.Error() + } + causeStr := e.cause.Error() + if isStructuredJSON(causeStr) { + return causeStr + } + return causeStr + "; " + e.summary +} + +func isStructuredJSON(s string) bool { + trimmed := strings.TrimSpace(s) + if len(trimmed) < 2 { + return false + } + if (trimmed[0] == '{' && trimmed[len(trimmed)-1] == '}') || (trimmed[0] == '[' && trimmed[len(trimmed)-1] == ']') { + return json.Valid([]byte(trimmed)) + } + return false +} + +func (e *routeExhaustionClonedError) Unwrap() error { + if e == nil { + return nil + } + return e.cause +} + +// Headers forwards the wrapped cause's error headers if it exposes them, so +// handlers that collect passthrough headers from the final routed error do not +// lose them when the cause is wrapped by route exhaustion. It returns a fresh +// copy of the cause's map and never mutates the caller's headers. +func (e *routeExhaustionClonedError) Headers() http.Header { + if e == nil { + return nil + } + var carrier interface{ Headers() http.Header } + if errors.As(e.cause, &carrier) && carrier != nil { + return cloneHTTPHeader(carrier.Headers()) + } + return nil +} + +// SafeResponseHeaders forwards trusted response headers from the wrapped cause +// if it exposes them, so handlers reading SafeResponseHeaders from the final +// routed error keep e.g. the Home busy error's Retry-After through route +// exhaustion. It returns a fresh copy and never mutates the caller's headers. +func (e *routeExhaustionClonedError) SafeResponseHeaders() http.Header { + if e == nil { + return nil + } + return SafeResponseHeaders(e.cause) +} diff --git a/sdk/cliproxy/auth/selected_auth_metadata_test.go b/sdk/cliproxy/auth/selected_auth_metadata_test.go index 2a7433e44..9c9087243 100644 --- a/sdk/cliproxy/auth/selected_auth_metadata_test.go +++ b/sdk/cliproxy/auth/selected_auth_metadata_test.go @@ -1,8 +1,11 @@ package auth import ( + "context" + "net/http" "testing" + registry "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" ) @@ -38,3 +41,153 @@ func TestPublishSelectedAuthMetadataIncludesStableIndex(t *testing.T) { t.Fatalf("selected auth index metadata = %#v, want %q", got, auth.Index) } } + +type dummySelExecutor struct { + provider string +} + +func (e *dummySelExecutor) Identifier() string { return e.provider } +func (e *dummySelExecutor) Execute(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} +func (e *dummySelExecutor) ExecuteStream(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + return nil, nil +} +func (e *dummySelExecutor) Refresh(ctx context.Context, auth *Auth) (*Auth, error) { + return auth, nil +} +func (e *dummySelExecutor) CountTokens(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} +func (e *dummySelExecutor) HttpRequest(ctx context.Context, auth *Auth, req *http.Request) (*http.Response, error) { + return nil, nil +} + +func TestManagerSelection_NilMetadataPreservesAffinityNamespace(t *testing.T) { + ctx := context.Background() + manager := NewManager(nil, nil, nil) + manager.RegisterExecutor(&dummySelExecutor{provider: "claude"}) + manager.RegisterExecutor(&dummySelExecutor{provider: "openai"}) + reg := registry.GetGlobalRegistry() + reg.RegisterClient("auth-1", "claude", []*registry.ModelInfo{{ID: "claude-3-5-sonnet"}}) + t.Cleanup(func() { + reg.UnregisterClient("auth-1") + }) + auth1 := &Auth{ + ID: "auth-1", + Provider: "claude", + FileName: "auth-1.json", + Status: StatusActive, + } + if _, err := manager.Register(ctx, auth1); err != nil { + t.Fatalf("manager.Register() error = %v", err) + } + affinity := NewSessionAffinitySelector(&WeightedRoundRobinSelector{}) + manager.SetSelector(affinity) + + t.Run("single provider pickNext with nil Metadata populates and preserves affinity metadata", func(t *testing.T) { + opts := cliproxyexecutor.Options{Metadata: make(map[string]any)} + auth, _, err := manager.pickNextLegacy(ctx, "claude", "claude-3-5-sonnet", opts, nil) + if err != nil { + t.Fatalf("pickNextLegacy() error = %v", err) + } + if auth == nil { + t.Fatal("pickNextLegacy() returned nil auth") + } + if opts.Metadata == nil { + t.Fatal("opts.Metadata is nil after pickNextLegacy, expected initialized map") + } + providerMeta, ok := opts.Metadata[cliproxyexecutor.SessionAffinityProviderMetadataKey].(string) + if !ok || providerMeta != "claude" { + t.Fatalf("SessionAffinityProviderMetadataKey = %q, %v; want \"claude\", true", providerMeta, ok) + } + modelMeta, ok := opts.Metadata[cliproxyexecutor.SessionAffinityModelMetadataKey].(string) + if !ok || modelMeta != "claude-3-5-sonnet" { + t.Fatalf("SessionAffinityModelMetadataKey = %q, %v; want \"claude-3-5-sonnet\", true", modelMeta, ok) + } + }) + + t.Run("mixed provider pickNextMixed with nil Metadata populates mixed namespace", func(t *testing.T) { + opts := cliproxyexecutor.Options{Metadata: make(map[string]any)} + auth, _, _, err := manager.pickNextMixedLegacy(ctx, []string{"claude", "openai"}, "claude-3-5-sonnet", opts, nil) + if err != nil { + t.Fatalf("pickNextMixedLegacy() error = %v", err) + } + if auth == nil { + t.Fatal("pickNextMixedLegacy() returned nil auth") + } + if opts.Metadata == nil { + t.Fatal("opts.Metadata is nil after pickNextMixedLegacy, expected initialized map") + } + providerMeta, ok := opts.Metadata[cliproxyexecutor.SessionAffinityProviderMetadataKey].(string) + if !ok || providerMeta != "mixed" { + t.Fatalf("SessionAffinityProviderMetadataKey = %q, %v; want \"mixed\", true", providerMeta, ok) + } + + res := Result{ + AuthID: auth.ID, + Provider: "rewritten-provider", + Model: "rewritten-model", + Success: true, + Options: opts, + } + affinity.OnResult(res) + }) +} + +// Manager-level regression: with nonempty caller-supplied exclusions, the mixed-pool affinity +// namespace ("mixed" provider + requested model) stamped during selection must survive into the +// success Result.Options, and a subsequent same-session request must bind to the same auth. +func TestManagerExecute_MixedAffinityNamespaceRetainedThroughExclusions(t *testing.T) { + model := "claude-3-5-sonnet" + reg := registry.GetGlobalRegistry() + reg.RegisterClient("affinity-auth-1", "claude", []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { + reg.UnregisterClient("affinity-auth-1") + }) + auth1 := &Auth{ + ID: "affinity-auth-1", + Provider: "claude", + FileName: "affinity-auth-1.json", + Status: StatusActive, + } + executor := newOuterRetryTestExecutor() + manager := NewManager(nil, nil, nil) + manager.SetRetryConfig(0, 0, 1) + manager.RegisterExecutor(executor) + if _, err := manager.Register(context.Background(), auth1); err != nil { + t.Fatalf("manager.Register() error = %v", err) + } + affinity := NewSessionAffinitySelector(&WeightedRoundRobinSelector{}) + manager.SetSelector(affinity) + + opts := cliproxyexecutor.Options{ + Headers: http.Header{"X-Session-Id": []string{"mixed-affinity-session-123"}}, + Metadata: map[string]any{}, + } + // Caller supplies an exclusion so withExcludedAuthIDs clones the metadata map. + opts.Metadata[cliproxyexecutor.ExcludedAuthIDsMetadataKey] = map[string]struct{}{"some-other-auth": {}} + + // First request: success must bind the mixed + requested-model affinity namespace so the + // second same-session request resolves to the same auth (the "no failover" requirement). + // withExcludedAuthIDs clones the metadata map, so the affinity keys live on that clone; the + // caller-visible opts.Metadata does not necessarily carry them. Assert the semantic effect. + if _, err := manager.Execute(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: model}, opts); err != nil { + t.Fatalf("Execute() error = %v, expected success", err) + } + if cached, ok := affinity.cache.Get("mixed::header:mixed-affinity-session-123::" + model); !ok { + t.Logf("mixed affinity binding not found; opts.Metadata=%#v", opts.Metadata) + t.Fatalf("expected mixed affinity binding under mixed::header:mixed-affinity-session-123::%s", model) + } else if cached != auth1.ID { + t.Fatalf("mixed affinity binding = %q, want %q", cached, auth1.ID) + } + + // Second same-session request: must hit the same mixed binding (i.e. not rotate empty/fail). + resp, err := manager.Execute(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: model}, opts) + if err != nil { + t.Fatalf("Execute() second request error = %v, want same binding success", err) + } + if len(resp.Payload) == 0 { + t.Fatal("Execute() second request returned empty payload") + } +} diff --git a/sdk/cliproxy/auth/selector.go b/sdk/cliproxy/auth/selector.go index c77513727..9f29440fd 100644 --- a/sdk/cliproxy/auth/selector.go +++ b/sdk/cliproxy/auth/selector.go @@ -248,10 +248,19 @@ func preferCodexWebsocketAuths(ctx context.Context, provider string, available [ return available } -func collectAvailableByPriority(auths []*Auth, model string, now time.Time) (available map[int][]*Auth, cooldownCount int, earliest time.Time) { +func collectAvailableByPriority(auths []*Auth, model string, now time.Time, excluded map[string]struct{}) (available map[int][]*Auth, cooldownCount int, eligibleCount int, earliest time.Time) { available = make(map[int][]*Auth) for i := 0; i < len(auths); i++ { candidate := auths[i] + // Skip nil candidates before consulting the exclusion map: + // candidate.ID on a nil entry would panic. + if candidate == nil { + continue + } + if _, skip := excluded[candidate.ID]; skip { + continue + } + eligibleCount++ blocked, reason, next := isAuthBlockedForModel(candidate, model, now) if !blocked { priority := authPriority(candidate) @@ -265,25 +274,33 @@ func collectAvailableByPriority(auths []*Auth, model string, now time.Time) (ava } } } - return available, cooldownCount, earliest + return available, cooldownCount, eligibleCount, earliest } -func getAvailableAuths(auths []*Auth, provider, model string, now time.Time) ([]*Auth, error) { - return getAvailableAuthsWithPriorityMode(auths, provider, model, now, false) +func getAvailableAuths(auths []*Auth, provider, model string, now time.Time, excluded ...map[string]struct{}) ([]*Auth, error) { + return getAvailableAuthsWithPriorityMode(auths, provider, model, now, false, excluded...) } -func getAvailableAuthsAcrossPriorities(auths []*Auth, provider, model string, now time.Time) ([]*Auth, error) { - return getAvailableAuthsWithPriorityMode(auths, provider, model, now, true) +func getAvailableAuthsAcrossPriorities(auths []*Auth, provider, model string, now time.Time, excluded ...map[string]struct{}) ([]*Auth, error) { + return getAvailableAuthsWithPriorityMode(auths, provider, model, now, true, excluded...) } -func getAvailableAuthsWithPriorityMode(auths []*Auth, provider, model string, now time.Time, allPriorities bool) ([]*Auth, error) { +func getAvailableAuthsWithPriorityMode(auths []*Auth, provider, model string, now time.Time, allPriorities bool, excluded ...map[string]struct{}) ([]*Auth, error) { if len(auths) == 0 { return nil, &Error{Code: "auth_not_found", Message: "no auth candidates"} } + var ex map[string]struct{} + if len(excluded) > 0 { + ex = excluded[0] + } - availableByPriority, cooldownCount, earliest := collectAvailableByPriority(auths, model, now) + availableByPriority, cooldownCount, eligibleCount, earliest := collectAvailableByPriority(auths, model, now, ex) if len(availableByPriority) == 0 { - if cooldownCount == len(auths) && !earliest.IsZero() { + // Count only eligible (non-excluded) auths: excluded entries are not + // part of the cooldown decision, otherwise the caller would get a + // non-retryable auth_unavailable instead of the cooldown error with + // Retry-After when every pickable auth is in fact cooling. + if eligibleCount > 0 && cooldownCount == eligibleCount && !earliest.IsZero() { providerForError := provider if providerForError == "mixed" { providerForError = "" @@ -370,7 +387,7 @@ func highestPriorityAuths(auths []*Auth) []*Auth { func (s *RoundRobinSelector) Pick(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth) (*Auth, error) { _ = opts now := time.Now() - available, err := getAvailableAuths(auths, provider, model, now) + available, err := getAvailableAuths(auths, provider, model, now, extractExcludedAuthIDs(opts.Metadata)) if err != nil { return nil, err } @@ -416,7 +433,7 @@ func positiveWeightAuths(auths []*Auth) []*Auth { // Pick selects the next available auth using smooth weighted round-robin. func (s *WeightedRoundRobinSelector) Pick(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth) (*Auth, error) { _ = opts - available, errAvailable := getAvailableAuths(positiveWeightAuths(auths), provider, model, time.Now()) + available, errAvailable := getAvailableAuths(positiveWeightAuths(auths), provider, model, time.Now(), extractExcludedAuthIDs(opts.Metadata)) if errAvailable != nil { return nil, errAvailable } @@ -526,7 +543,7 @@ func saturatingAddInt64(value, delta int64) int64 { func (s *FillFirstSelector) Pick(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth) (*Auth, error) { _ = opts now := time.Now() - available, err := getAvailableAuths(auths, provider, model, now) + available, err := getAvailableAuths(auths, provider, model, now, extractExcludedAuthIDs(opts.Metadata)) if err != nil { return nil, err } @@ -541,6 +558,9 @@ func isAuthBlockedForModel(auth *Auth, model string, now time.Time) (bool, block if auth.Disabled || auth.Status == StatusDisabled { return true, blockReasonDisabled, time.Time{} } + if auth.Quota.Exceeded && auth.Quota.Reason == "credential_quota" && auth.Quota.NextRecoverAt.After(now) { + return true, blockReasonCooldown, auth.Quota.NextRecoverAt + } if model != "" { if len(auth.ModelStates) > 0 { modelKey := canonicalModelKey(model) @@ -608,8 +628,10 @@ func availabilityBlock(unavailable, quotaExceeded bool, nextRetryAfter, nextReco // It extracts session ID from multiple sources and maintains session-to-auth // mappings with automatic failover when the bound auth becomes unavailable. type SessionAffinitySelector struct { - fallback Selector - cache *SessionCache + fallback Selector + cache *SessionCache + quarantine *SessionCache + bindMu sync.Mutex } // SessionAffinityConfig configures the session affinity selector. @@ -635,8 +657,9 @@ func NewSessionAffinitySelectorWithConfig(cfg SessionAffinityConfig) *SessionAff cfg.TTL = time.Hour } return &SessionAffinitySelector{ - fallback: cfg.Fallback, - cache: NewSessionCache(cfg.TTL), + fallback: cfg.Fallback, + cache: NewSessionCache(cfg.TTL), + quarantine: NewSessionCache(cfg.TTL), } } @@ -654,14 +677,21 @@ func NewSessionAffinitySelectorWithConfig(cfg SessionAffinityConfig) *SessionAff // that may be supported by different auth credentials, and to avoid cross-provider conflicts. func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth) (*Auth, error) { entry := selectorLogEntry(ctx) + if opts.Metadata == nil { + opts.Metadata = make(map[string]any) + } + opts.Metadata[cliproxyexecutor.SessionAffinityProviderMetadataKey] = provider + opts.Metadata[cliproxyexecutor.SessionAffinityModelMetadataKey] = model + primaryID, fallbackID := extractSessionIDs(opts.Headers, opts.OriginalRequest, opts.Metadata) now := time.Now() + excluded := extractExcludedAuthIDs(opts.Metadata) availabilityCandidates := auths if _, weighted := s.fallback.(*WeightedRoundRobinSelector); weighted { availabilityCandidates = positiveWeightAuths(auths) } if primaryID == "" { - fallbackAuths, errAvailable := getAvailableAuths(availabilityCandidates, provider, model, now) + fallbackAuths, errAvailable := getAvailableAuths(availabilityCandidates, provider, model, now, excluded) if errAvailable != nil { return nil, errAvailable } @@ -671,17 +701,18 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri // A single availability pass serves both lookups: the bound credential is validated against // every priority tier, while the fallback selector keeps seeing only the highest tier. - available, err := getAvailableAuthsAcrossPriorities(availabilityCandidates, provider, model, now) + available, err := getAvailableAuthsAcrossPriorities(availabilityCandidates, provider, model, now, excluded) if err != nil { return nil, err } - fallbackAuths := highestPriorityAuths(available) cacheKey := provider + "::" + primaryID + "::" + model fallbackKey := "" if fallbackID != "" && fallbackID != primaryID { fallbackKey = provider + "::" + fallbackID + "::" + model } + available = s.excludeSessionQuarantine(cacheKey, fallbackKey, available) + fallbackAuths := highestPriorityAuths(available) bind := func(authID string) { if fallbackKey != "" { s.cache.SetAliases(authID, cacheKey, fallbackKey) @@ -690,6 +721,7 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri s.cache.Set(cacheKey, authID) } + // Fast path outside bindMu: reuse valid cached binding without holding bindMu. if cachedAuthID, ok := s.cache.GetAndRefresh(cacheKey); ok { for _, auth := range available { if auth.ID == cachedAuthID { @@ -698,37 +730,264 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri return auth, nil } } - // Cached auth not available, reselect via fallback selector for even distribution - auth, err := s.fallback.Pick(ctx, provider, model, opts, fallbackAuths) - if err != nil { - return nil, err + } else if fallbackKey != "" { + if cachedAuthID, ok := s.cache.Get(fallbackKey); ok { + for _, auth := range available { + if auth.ID == cachedAuthID { + bind(auth.ID) + entry.Infof("session-affinity: cache hit | session=%s auth=%s provider=%s model=%s", truncateSessionID(primaryID), auth.ID, provider, model) + return auth, nil + } + } } - bind(auth.ID) - entry.Infof("session-affinity: cache hit but auth unavailable, reselected | session=%s auth=%s provider=%s model=%s", truncateSessionID(primaryID), auth.ID, provider, model) - return auth, nil } - if fallbackKey != "" { + s.bindMu.Lock() + defer s.bindMu.Unlock() + + // Under bindMu, re-check if a concurrent request refreshed or rebound the session. + if cachedAuthID, ok := s.cache.GetAndRefresh(cacheKey); ok { + for _, auth := range available { + if auth.ID == cachedAuthID { + entry.Infof("session-affinity: concurrent cache hit | session=%s auth=%s provider=%s model=%s", truncateSessionID(primaryID), auth.ID, provider, model) + bind(auth.ID) + return auth, nil + } + } + } else if fallbackKey != "" { if cachedAuthID, ok := s.cache.Get(fallbackKey); ok { for _, auth := range available { if auth.ID == cachedAuthID { + entry.Infof("session-affinity: concurrent cache hit | session=%s auth=%s provider=%s model=%s", truncateSessionID(primaryID), auth.ID, provider, model) bind(auth.ID) - entry.Infof("session-affinity: fallback cache hit | session=%s fallback=%s auth=%s provider=%s model=%s", truncateSessionID(primaryID), truncateSessionID(fallbackID), auth.ID, provider, model) return auth, nil } } } } + // Authoritative stale observation conducted under bindMu using non-refreshing token read. + // Observe both alias groups: they may be split across different auths, in + // which case failover must reconcile both, not just the first one found. + staleKey := cacheKey + staleAuthID, staleGen, staleAliases, hasStale := s.cache.GetWithGeneration(cacheKey) + splitAuthID := "" + var splitGen uint64 + var splitAliases []string + hasSplit := false + if fallbackKey != "" { + splitAuthID, splitGen, splitAliases, hasSplit = s.cache.GetWithGeneration(fallbackKey) + } + splitGroups := hasStale && hasSplit && staleAuthID != splitAuthID + if !hasStale && hasSplit { + staleKey = fallbackKey + staleAuthID, staleGen, staleAliases, hasStale = splitAuthID, splitGen, splitAliases, true + } + auth, err := s.fallback.Pick(ctx, provider, model, opts, fallbackAuths) if err != nil { return nil, err } + + if hasStale { + if splitGroups { + // Split alias groups (prompt-cache and conversation aliases bound to + // different auths): merge BOTH alias sets into a single group bound + // to the selected auth. Rebinding the groups separately would leave + // two groups on the same auth, and later housekeeping (OnResult) + // processes only the group holding the request's primary key — the + // surviving split group would keep selecting a failed auth. + if !s.mergeSplitAliasGroupsCAS(cacheKey, fallbackKey, auth.ID) { + entry.Infof("session-affinity: split-group merge lost to concurrent writer after retries | session=%s auth=%s provider=%s model=%s", truncateSessionID(primaryID), auth.ID, provider, model) + } + } else { + additional := []string{cacheKey} + if fallbackKey != "" { + additional = append(additional, fallbackKey) + } + if s.rebindAliasGroupCAS(staleKey, staleAuthID, staleGen, staleAliases, auth.ID, additional) { + entry.Infof("session-affinity: rebound stale alias group | session=%s oldAuth=%s newAuth=%s gen=%d", truncateSessionID(primaryID), staleAuthID, auth.ID, staleGen) + } else { + entry.Infof("session-affinity: CAS rebind aborted due to concurrent mutation, serving selected auth statelessly | session=%s auth=%s", truncateSessionID(primaryID), auth.ID) + } + } + return auth, nil + } + bind(auth.ID) - entry.Infof("session-affinity: cache miss, new binding | session=%s auth=%s provider=%s model=%s", truncateSessionID(primaryID), auth.ID, provider, model) + entry.Infof("session-affinity: cache miss, bound candidate | session=%s auth=%s provider=%s model=%s", truncateSessionID(primaryID), auth.ID, provider, model) return auth, nil } +// rebindAliasGroupCAS atomically rebinds a session alias group to newAuthID. +// When the compare-and-swap loses to a concurrent writer, the group is +// re-observed and the binding retried (bounded), so the auth selected for +// this request is not silently dropped by a stale generation. +func (s *SessionAffinitySelector) rebindAliasGroupCAS(sessionKey string, expectedAuthID string, expectedGen uint64, expectedAliases []string, newAuthID string, additionalAliases []string) bool { + for attempt := 0; attempt < 3; attempt++ { + if s.cache.CompareAndReplaceAliases(expectedAuthID, expectedGen, expectedAliases, newAuthID, additionalAliases...) { + return true + } + authID, gen, aliases, ok := s.cache.GetWithGeneration(sessionKey) + if !ok { + return false + } + expectedAuthID, expectedGen, expectedAliases = authID, gen, aliases + } + return false +} + +// mergeSplitAliasGroupsCAS reconciles two split session alias groups (a +// prompt-cache alias and a conversation alias previously bound to different +// auths) into a single group bound to authID. Merging matters because later +// housekeeping (OnResult) processes only the group holding the request's +// primary key: two surviving groups would let the conversation-only alias +// keep selecting a failed auth. The merge is retried with fresh observation +// when a concurrent writer invalidates the expectations (bounded). +func (s *SessionAffinitySelector) mergeSplitAliasGroupsCAS(cacheKey, fallbackKey string, authID string) bool { + // retainedF holds the fallback group's aliases once its delete has + // committed. A retry after a lost primary CAS would otherwise re-observe + // the (now deleted) fallback entry and rebuild merged from cacheKey and + // fallbackKey alone, permanently dropping the fallback group's + // historical aliases from the rebound group. + // + // Mirror of CLIProxyAPI e768fba9. + var retainedF []string + var deletedAuthF string + for attempt := 0; attempt < 3; attempt++ { + authP, genP, aliasesP, okP := s.cache.GetWithGeneration(cacheKey) + authF, genF, aliasesF, okF := s.cache.GetWithGeneration(fallbackKey) + if !okF { + aliasesF = retainedF + } + merged := mergeSessionAliases(aliasesP, aliasesF...) + merged = mergeSessionAliases(merged, cacheKey, fallbackKey) + if okF && authF != authID { + removed := s.cache.CompareAndDeleteGroup(fallbackKey, authF, genF, aliasesF) + if removed == nil { + continue + } + deletedAuthF = authF + retainedF = mergeSessionAliases(retainedF, removed...) + } + if okP { + if s.cache.CompareAndReplaceAliases(authP, genP, aliasesP, authID, merged...) { + return true + } + continue + } + s.cache.SetAliases(authID, merged...) + return true + } + if len(retainedF) > 0 && deletedAuthF != "" { + s.cache.RestoreAliasesIfAbsent(deletedAuthF, retainedF...) + } + return false +} + +// OnResult handles session affinity binding or release based on execution outcome. +func (s *SessionAffinitySelector) OnResult(res Result) { + if s == nil || s.cache == nil || res.AuthID == "" { + return + } + primaryID, fallbackID := extractSessionIDs(res.Options.Headers, res.Options.OriginalRequest, res.Options.Metadata) + if primaryID == "" && fallbackID == "" { + return + } + + // Use the affinity selection namespace when present so mixed pools bind under + // the same key selection read (the literal "mixed" pool key); otherwise fall + // back to the auth's actual provider for single-provider callers. + ns := res.Provider + if raw, ok := res.Options.Metadata[cliproxyexecutor.SessionAffinityProviderMetadataKey].(string); ok && raw != "" { + ns = raw + } + // Use the affinity model namespace (the normalized Pick-time model) when present; + // fall back to the rewritten result model for metadata-absent callers. + nsModel := res.Model + if raw, ok := res.Options.Metadata[cliproxyexecutor.SessionAffinityModelMetadataKey].(string); ok && raw != "" { + nsModel = raw + } + + cacheKey := ns + "::" + primaryID + "::" + nsModel + var fallbackKey string + if fallbackID != "" && fallbackID != primaryID { + fallbackKey = ns + "::" + fallbackID + "::" + nsModel + } + + if res.Success { + if fallbackKey != "" { + s.cache.SetAliases(res.AuthID, cacheKey, fallbackKey) + } else { + s.cache.Set(cacheKey, res.AuthID) + } + return + } + + if res.Error != nil && shouldSkipCredentialCooldown(res.Error) { + return + } + + var aliases []string + if authID, _, groupAliases, ok := s.cache.GetWithGeneration(cacheKey); ok && authID == res.AuthID { + aliases = groupAliases + s.cache.Invalidate(cacheKey) + } else if fallbackKey != "" { + if authID, _, groupAliases, ok := s.cache.GetWithGeneration(fallbackKey); ok && authID == res.AuthID { + aliases = groupAliases + s.cache.Invalidate(fallbackKey) + } + } + if len(aliases) == 0 { + aliases = []string{cacheKey, fallbackKey} + } + s.quarantineSessionAuth(aliases, res.AuthID, res.RetryAfter) +} + +func (s *SessionAffinitySelector) excludeSessionQuarantine(cacheKey, fallbackKey string, auths []*Auth) []*Auth { + if s == nil || s.quarantine == nil || len(auths) == 0 { + return auths + } + filtered := make([]*Auth, 0, len(auths)) + for _, auth := range auths { + if auth == nil { + continue + } + blocked := false + for _, key := range []string{cacheKey, fallbackKey} { + if key == "" { + continue + } + if _, ok := s.quarantine.Get(key + "::failed::" + auth.ID); ok { + blocked = true + break + } + } + if !blocked { + filtered = append(filtered, auth) + } + } + return filtered +} + +func (s *SessionAffinitySelector) quarantineSessionAuth(cacheKeys []string, authID string, retryAfter *time.Duration) { + if s == nil || s.quarantine == nil || authID == "" { + return + } + delay := 5 * time.Second + if retryAfter != nil && *retryAfter > 0 { + delay = *retryAfter + } + expiresAt := time.Now().Add(delay) + for _, key := range cacheKeys { + if key == "" { + continue + } + quarantineKey := key + "::failed::" + authID + s.quarantine.setAliasesUntil(authID, expiresAt, quarantineKey) + } +} + func selectorLogEntry(ctx context.Context) *log.Entry { if ctx == nil { return log.NewEntry(log.StandardLogger()) @@ -752,6 +1011,9 @@ func (s *SessionAffinitySelector) Stop() { if s.cache != nil { s.cache.Stop() } + if s.quarantine != nil { + s.quarantine.Stop() + } } // InvalidateAuth removes all session bindings for a specific auth. @@ -760,6 +1022,9 @@ func (s *SessionAffinitySelector) InvalidateAuth(authID string) { if s.cache != nil { s.cache.InvalidateAuth(authID) } + if s.quarantine != nil { + s.quarantine.InvalidateAuth(authID) + } } // normalizedSessionCandidate validates an explicit client-provided session signal. diff --git a/sdk/cliproxy/auth/selector_review_p2_test.go b/sdk/cliproxy/auth/selector_review_p2_test.go new file mode 100644 index 000000000..2dedce380 --- /dev/null +++ b/sdk/cliproxy/auth/selector_review_p2_test.go @@ -0,0 +1,393 @@ +package auth + +import ( + "context" + "errors" + "fmt" + "net/http" + "slices" + "testing" + "time" + + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +// Regression tests mirrored from CLIProxyAPI PR #4881 follow-up +// (codex pullrequestreview-4943660625, findings 1 and 2). Finding 3 (ignored +// CompareAndReplaceGroup result) is CPA-specific: the CPAPlus affinity path +// uses CompareAndReplaceAliases and already checks its result, serving the +// selected auth statelessly when the CAS loses. + +// TestCooldownErrorCountsOnlyEligibleAuths is a regression guard for the +// first finding: cooldownCount used to be compared against len(auths), +// including request-excluded entries, so a pool where every pickable auth was +// cooling reported the non-retryable auth_unavailable instead of +// model_cooldown with Retry-After. +func TestCooldownErrorCountsOnlyEligibleAuths(t *testing.T) { + t.Parallel() + + model := "test-model" + now := time.Now() + next := now.Add(60 * time.Second) + cooled := &Auth{ + ID: "auth-cooled", + ModelStates: map[string]*ModelState{ + model: { + Status: StatusActive, + Unavailable: true, + NextRetryAfter: next, + Quota: QuotaState{ + Exceeded: true, + NextRecoverAt: next, + }, + }, + }, + } + excluded := &Auth{ + ID: "auth-excluded", + ModelStates: map[string]*ModelState{ + model: {Status: StatusActive}, + }, + } + + _, err := getAvailableAuths([]*Auth{cooled, excluded}, "gemini", model, now, map[string]struct{}{"auth-excluded": {}}) + if err == nil { + t.Fatal("getAvailableAuths() error = nil") + } + var mce *modelCooldownError + if !errors.As(err, &mce) { + t.Fatalf("getAvailableAuths() error = %T (%v), want *modelCooldownError: excluded auths must not count toward the cooldown decision", err, err) + } + if mce.StatusCode() != http.StatusTooManyRequests { + t.Fatalf("StatusCode() = %d, want %d", mce.StatusCode(), http.StatusTooManyRequests) + } + if got := mce.Headers().Get("Retry-After"); got == "" { + t.Fatal("Headers().Get(Retry-After) = empty, want a value") + } +} + +// TestGetAvailableAuthsSkipsNilCandidates is a regression guard for the +// second finding: a nil entry in the auth list used to panic on candidate.ID +// when consulting the exclusion map. +func TestGetAvailableAuthsSkipsNilCandidates(t *testing.T) { + t.Parallel() + + model := "test-model" + active := &Auth{ + ID: "auth-active", + ModelStates: map[string]*ModelState{ + model: {Status: StatusActive}, + }, + } + + got, err := getAvailableAuths([]*Auth{nil, active}, "gemini", model, time.Now()) + if err != nil { + t.Fatalf("getAvailableAuths() error = %v, want nil", err) + } + if len(got) != 1 || got[0] != active { + t.Fatalf("getAvailableAuths() = %v, want [auth-active]", got) + } + + _, err = getAvailableAuths([]*Auth{nil}, "gemini", model, time.Now(), map[string]struct{}{"anything": {}}) + if err == nil { + t.Fatal("getAvailableAuths() with only a nil candidate: error = nil, want auth_unavailable") + } + var mce *modelCooldownError + if errors.As(err, &mce) { + t.Fatalf("getAvailableAuths() with only a nil candidate: error = %v, must not be modelCooldownError", err) + } +} + +// TestPickRebindsSplitAffinityGroupsOnFailover mirrors the CPA regression +// guard for the codex P2 finding on PR #4881. The CPAPlus binding design has +// no splitConflict skip: on a miss it rebinds the observed stale group via +// CompareAndReplaceAliases and absorbs both session keys into it, which +// converges the split groups onto the selected auth. This test locks that +// convergence in. +func TestPickRebindsSplitAffinityGroupsOnFailover(t *testing.T) { + t.Parallel() + + model := "test-model" + provider := "gemini" + primaryKey := provider + "::pck:pk1::" + model + fallbackKey := provider + "::conv:c1::" + model + + cooled := func(id string) *Auth { + return &Auth{ + ID: id, + ModelStates: map[string]*ModelState{ + model: { + Status: StatusActive, + Unavailable: true, + NextRetryAfter: time.Now().Add(60 * time.Second), + Quota: QuotaState{ + Exceeded: true, + NextRecoverAt: time.Now().Add(60 * time.Second), + }, + }, + }, + } + } + authA := cooled("auth-a") + authB := cooled("auth-b") + authC := &Auth{ + ID: "auth-c", + ModelStates: map[string]*ModelState{ + model: {Status: StatusActive}, + }, + } + + selector := NewSessionAffinitySelector(&FillFirstSelector{}) + selector.cache.SetAliases("auth-a", primaryKey) + selector.cache.SetAliases("auth-b", fallbackKey) + + payload := []byte(`{"prompt_cache_key":"pk1","conversation":{"id":"c1"}}`) + opts := cliproxyexecutor.Options{OriginalRequest: payload, Metadata: map[string]any{}} + auth, err := selector.Pick(context.Background(), provider, model, opts, []*Auth{authA, authB, authC}) + if err != nil { + t.Fatalf("Pick() error = %v, want nil", err) + } + if auth != authC { + t.Fatalf("Pick() = %v, want auth-c (only available auth)", auth.ID) + } + + gotPrimary, genP, aliasesPrimary, okPrimary := selector.cache.GetWithGeneration(primaryKey) + if !okPrimary || gotPrimary != "auth-c" { + t.Fatalf("primary group after failover = %q (ok=%v), want auth-c", gotPrimary, okPrimary) + } + gotFallback, genF, _, okFallback := selector.cache.GetWithGeneration(fallbackKey) + if !okFallback || gotFallback != "auth-c" { + t.Fatalf("fallback group after failover = %q (ok=%v), want auth-c", gotFallback, okFallback) + } + if genP == 0 || genP != genF { + t.Fatalf("split groups not merged into one: primary gen=%d, fallback gen=%d", genP, genF) + } + if !slices.Contains(aliasesPrimary, fallbackKey) { + t.Fatalf("primary group aliases %v missing fallback key %q", aliasesPrimary, fallbackKey) + } +} + +// TestCompareAndDeleteGroupRejectsStaleObservation covers the codex P2 +// follow-up on PR #4881: when a concurrent request refreshes or extends the +// fallback group between the observation and the delete, a stale merge +// observation must not remove the newer group, otherwise the newly attached +// aliases lose their affinity binding. +// +// Mirror of CLIProxyAPI dd8c72a3. +func TestCompareAndDeleteGroupRejectsStaleObservation(t *testing.T) { + t.Parallel() + + cache := NewSessionCache(time.Minute) + key := "gemini::conv:c1::test-model" + extended := "gemini::conv:c2::test-model" + cache.SetAliases("auth-a", key) + + authID, gen, aliases, ok := cache.GetWithGeneration(key) + if !ok || authID != "auth-a" { + t.Fatalf("GetWithGeneration() = %q, ok=%v; want auth-a bound", authID, ok) + } + + // A concurrent request extends the same group on the same auth. + cache.SetAliases("auth-a", key, extended) + + if removed := cache.CompareAndDeleteGroup(key, "auth-a", gen, aliases); removed != nil { + t.Fatalf("CompareAndDeleteGroup with stale observation removed %v; want nil (group must survive)", removed) + } + if got, _, gotAliases, ok := cache.GetWithGeneration(key); !ok || got != "auth-a" || !slices.Contains(gotAliases, extended) { + t.Fatalf("group after stale delete = %q, aliases=%v, ok=%v; want intact auth-a group", got, gotAliases, ok) + } + + // A fresh observation still deletes successfully. + authID, gen, aliases, ok = cache.GetWithGeneration(key) + if !ok { + t.Fatal("GetWithGeneration() after extension lost the group") + } + if removed := cache.CompareAndDeleteGroup(key, authID, gen, aliases); removed == nil { + t.Fatal("CompareAndDeleteGroup with fresh observation returned nil; want removed aliases") + } + if _, _, _, ok := cache.GetWithGeneration(key); ok { + t.Fatal("group still present after fresh CompareAndDeleteGroup") + } +} + +// TestMergeSplitGroupsRetainsFallbackAliasesUnderCASContention covers the +// codex P2 follow-up on PR #4881: the fallback delete commits before the +// primary CAS, so a primary CAS that loses to a concurrent writer must not +// rebuild merged from cacheKey and fallbackKey alone — the fallback group's +// historical aliases have to survive via the retained observation. +// +// The contending writer uses CompareAndReplaceAliases itself, so its bump +// only lands while the primary group is still in its pre-merge state; once +// the merge commits, the contender's CAS refuses and cannot corrupt the +// result. Both interleavings are therefore valid and the assertions hold +// either way. +// +// Mirror of CLIProxyAPI e768fba9. +func TestMergeSplitGroupsRetainsFallbackAliasesUnderCASContention(t *testing.T) { + t.Parallel() + + for i := 0; i < 2000; i++ { + selector := NewSessionAffinitySelector(&FillFirstSelector{}) + cacheKey := fmt.Sprintf("gemini::pck:pk%d::test-model", i) + fallbackKey := fmt.Sprintf("gemini::conv:c1-%d::test-model", i) + historical := fmt.Sprintf("gemini::conv:c0-%d::test-model", i) + scratch := fmt.Sprintf("gemini::scratch:%d::test-model", i) + selector.cache.SetAliases("auth-a", cacheKey) + selector.cache.SetAliases("auth-b", fallbackKey, historical) + + // A concurrent writer attaches a scratch alias to the primary group + // right after the fallback group disappears — the interleaving that + // makes the merge's first primary CAS lose. + done := make(chan struct{}) + finished := make(chan struct{}) + go func() { + defer close(done) + bumps := 0 + for { + select { + case <-finished: + return + default: + } + if _, _, _, ok := selector.cache.GetWithGeneration(fallbackKey); ok { + continue + } + if bumps >= 2 { + return + } + authP, genP, aliasesP, okP := selector.cache.GetWithGeneration(cacheKey) + if okP { + bumped := append(append([]string(nil), aliasesP...), scratch) + if selector.cache.CompareAndReplaceAliases(authP, genP, aliasesP, authP, bumped...) { + bumps++ + } + } + } + }() + + merged := selector.mergeSplitAliasGroupsCAS(cacheKey, fallbackKey, "auth-c") + close(finished) + <-done + if !merged { + t.Fatalf("iteration %d: mergeSplitAliasGroupsCAS() = false under single-bump contention, want true", i) + } + got, _, aliases, ok := selector.cache.GetWithGeneration(cacheKey) + if !ok || got != "auth-c" { + t.Fatalf("iteration %d: merged group = %q, ok=%v; want auth-c", i, got, ok) + } + if !slices.Contains(aliases, historical) { + t.Fatalf("iteration %d: merged aliases %v missing historical fallback alias %q", i, aliases, historical) + } + } +} + +func TestSessionAffinitySelector_SplitGroupMergeExhaustionRestoresFallbackAliases(t *testing.T) { + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: &RoundRobinSelector{}, + TTL: time.Minute, + }) + defer selector.Stop() + provider := "responses-split-exhaustion" + model := "gpt-test" + + cacheKey := provider + "::pck:shared-prompt::" + model + fallbackKey := provider + "::conv:conversation-session::" + model + extraFallbackAlias := provider + "::extra:alias::" + model + missingPrimaryAlias := provider + "::conv:missing::" + model + + // 1. Group 1 on cacheKey + missingPrimaryAlias bound to auth-a + selector.cache.SetAliases("auth-a", cacheKey, missingPrimaryAlias) + + // Invalidate missingPrimaryAlias from entries table while leaving cacheKey's entry + // expecting it, so CompareAndReplaceAliases on cacheKey fails on all CAS attempts. + selector.cache.mu.Lock() + delete(selector.cache.entries, missingPrimaryAlias) + selector.cache.mu.Unlock() + + // 2. Group 2 on fallbackKey + extraFallbackAlias bound to auth-b + selector.cache.SetAliases("auth-b", fallbackKey, extraFallbackAlias) + + // 3. mergeSplitAliasGroupsCAS attempts to merge cacheKey and fallbackKey into auth-c. + // Since cacheKey expects missingPrimaryAlias which is missing from entries, + // CompareAndReplaceAliases fails on all 3 attempts. + merged := selector.mergeSplitAliasGroupsCAS(cacheKey, fallbackKey, "auth-c") + if merged { + t.Fatalf("mergeSplitAliasGroupsCAS must fail due to CAS exhaustion") + } + + // Since mergeSplitAliasGroupsCAS exhausted retries after deleting fallbackKey, + // the fallback group (fallbackKey and extraFallbackAlias) must be restored to auth-b. + if got, ok := selector.cache.Get(fallbackKey); !ok || got != "auth-b" { + t.Fatalf("fallbackKey must be restored to auth-b, got %q, %v", got, ok) + } + if got, ok := selector.cache.Get(extraFallbackAlias); !ok || got != "auth-b" { + t.Fatalf("extraFallbackAlias must be restored to auth-b, got %q, %v", got, ok) + } +} + +func TestSessionAffinitySelector_SplitGroupMergeExhaustionDoesNotClobberConcurrentRebind(t *testing.T) { + for i := 0; i < 500; i++ { + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: &RoundRobinSelector{}, + TTL: time.Minute, + }) + provider := fmt.Sprintf("responses-split-exhaustion-concurrent-%d", i) + model := "gpt-test" + + cacheKey := provider + "::pck:shared-prompt::" + model + fallbackKey := provider + "::conv:conversation-session::" + model + extraFallbackAlias := provider + "::extra:alias::" + model + missingPrimaryAlias := provider + "::conv:missing::" + model + + // 1. Group 1 on cacheKey + missingPrimaryAlias bound to auth-a + selector.cache.SetAliases("auth-a", cacheKey, missingPrimaryAlias) + + // Invalidate missingPrimaryAlias from entries table while leaving cacheKey's entry + // expecting it, so CompareAndReplaceAliases on cacheKey fails on all CAS attempts. + selector.cache.mu.Lock() + delete(selector.cache.entries, missingPrimaryAlias) + selector.cache.mu.Unlock() + + // 2. Group 2 on fallbackKey + extraFallbackAlias bound to auth-b + selector.cache.SetAliases("auth-b", fallbackKey, extraFallbackAlias) + + // Start a concurrent goroutine that observes when fallbackKey is deleted by attempt 0, + // and immediately rebinds extraFallbackAlias to auth-x. + done := make(chan struct{}) + finished := make(chan struct{}) + rebound := false + go func() { + defer close(done) + for { + select { + case <-finished: + return + default: + } + if _, ok := selector.cache.Get(fallbackKey); ok { + continue + } + selector.cache.SetAliases("auth-x", extraFallbackAlias) + rebound = true + return + } + }() + + // 3. mergeSplitAliasGroupsCAS attempts to merge cacheKey and fallbackKey into auth-c. + // Since cacheKey expects missingPrimaryAlias which is missing from entries, + // CompareAndReplaceAliases fails on all 3 attempts. + merged := selector.mergeSplitAliasGroupsCAS(cacheKey, fallbackKey, "auth-c") + close(finished) + <-done + if merged { + t.Fatalf("iteration %d: mergeSplitAliasGroupsCAS must fail due to CAS exhaustion", i) + } + + if rebound { + if got, ok := selector.cache.Get(extraFallbackAlias); !ok || got != "auth-x" { + t.Fatalf("iteration %d: concurrent rebind extraFallbackAlias clobbered; got %q, %v, want auth-x", i, got, ok) + } + } + selector.Stop() + } +} diff --git a/sdk/cliproxy/auth/selector_test.go b/sdk/cliproxy/auth/selector_test.go index 6024b0cca..205892ce2 100644 --- a/sdk/cliproxy/auth/selector_test.go +++ b/sdk/cliproxy/auth/selector_test.go @@ -766,6 +766,7 @@ func TestSessionAffinitySelector_SameSessionSameAuth(t *testing.T) { if first == nil { t.Fatalf("Pick() returned nil") } + selector.OnResult(Result{AuthID: first.ID, Provider: "claude", Model: "claude-3", Options: opts, Success: true}) // Verify consistency: same session, same auths -> same result for i := 0; i < 10; i++ { @@ -794,6 +795,7 @@ func TestSessionAffinitySelector_WeightedBindingRebindsAfterWeightBecomesZero(t if errFirst != nil { t.Fatalf("first Pick() error = %v", errFirst) } + selector.OnResult(Result{AuthID: first.ID, Provider: "claude", Model: "claude-3", Options: opts, Success: true}) if first.ID != authA.ID { t.Fatalf("first Pick() auth.ID = %q, want %q", first.ID, authA.ID) } @@ -803,6 +805,7 @@ func TestSessionAffinitySelector_WeightedBindingRebindsAfterWeightBecomesZero(t if errSecond != nil { t.Fatalf("Pick() after weight update error = %v", errSecond) } + selector.OnResult(Result{AuthID: second.ID, Provider: "claude", Model: "claude-3", Options: opts, Success: true}) if second.ID != authB.ID { t.Fatalf("Pick() after weight update auth.ID = %q, want %q", second.ID, authB.ID) } @@ -832,6 +835,7 @@ func TestSessionAffinitySelector_WeightedNewSessionsResetAfterWeightChange(t *te if errPick != nil { t.Fatalf("Pick(session-%d) error = %v", index, errPick) } + selector.OnResult(Result{AuthID: picked.ID, Provider: "claude", Model: "claude-3", Options: opts, Success: true}) return picked } for index := 0; index < 1000; index++ { @@ -893,7 +897,9 @@ func TestSessionAffinitySelector_DifferentSessionsDifferentAuths(t *testing.T) { opts2 := cliproxyexecutor.Options{OriginalRequest: session2} auth1, _ := selector.Pick(context.Background(), "claude", "claude-3", opts1, auths) + selector.OnResult(Result{AuthID: auth1.ID, Provider: "claude", Model: "claude-3", Options: opts1, Success: true}) auth2, _ := selector.Pick(context.Background(), "claude", "claude-3", opts2, auths) + selector.OnResult(Result{AuthID: auth2.ID, Provider: "claude", Model: "claude-3", Options: opts2, Success: true}) // Different sessions may or may not pick different auths (depends on hash collision) // But each session should be consistent @@ -933,6 +939,7 @@ func TestSessionAffinitySelector_FailoverWhenAuthUnavailable(t *testing.T) { if err != nil { t.Fatalf("Pick() error = %v", err) } + selector.OnResult(Result{AuthID: first.ID, Provider: "claude", Model: "claude-3", Options: opts, Success: true}) // Remove the bound auth from available list (simulating rate limit) availableWithoutFirst := make([]*Auth, 0, len(auths)-1) @@ -950,6 +957,7 @@ func TestSessionAffinitySelector_FailoverWhenAuthUnavailable(t *testing.T) { if second.ID == first.ID { t.Fatalf("Pick() after failover returned same auth %q, expected different", first.ID) } + selector.OnResult(Result{AuthID: second.ID, Provider: "claude", Model: "claude-3", Options: opts, Success: true}) // Subsequent picks should consistently return the new binding for i := 0; i < 5; i++ { @@ -1269,6 +1277,42 @@ func TestExtractSessionID_OpenAIResponsesAPI(t *testing.T) { } } +func TestSessionCache_RestoreAliasesIfAbsent_IndependentRestoration(t *testing.T) { + cache := NewSessionCache(time.Minute) + defer cache.Stop() + + // Initial fallback group with 3 aliases bound to auth-b + cache.SetAliases("auth-b", "fallbackKey", "reboundAlias", "stillAbsentAlias") + + // Fallback group gets deleted during merge attempt + removed := cache.CompareAndDeleteGroup("fallbackKey", "auth-b", 1, []string{"fallbackKey", "reboundAlias", "stillAbsentAlias"}) + if len(removed) == 0 { + t.Fatalf("CompareAndDeleteGroup failed") + } + + // Concurrent writer rebinds reboundAlias to auth-x + cache.SetAliases("auth-x", "reboundAlias") + + // Merge exhaustion attempts to restore remaining aliases + restored := cache.RestoreAliasesIfAbsent("auth-b", removed...) + if !restored { + t.Fatalf("RestoreAliasesIfAbsent should return true when some aliases are still absent") + } + + // reboundAlias must remain bound to auth-x + if got, ok := cache.Get("reboundAlias"); !ok || got != "auth-x" { + t.Fatalf("reboundAlias must remain auth-x, got %q, %v", got, ok) + } + + // fallbackKey and stillAbsentAlias must be restored to auth-b + if got, ok := cache.Get("fallbackKey"); !ok || got != "auth-b" { + t.Fatalf("fallbackKey must be restored to auth-b, got %q, %v", got, ok) + } + if got, ok := cache.Get("stillAbsentAlias"); !ok || got != "auth-b" { + t.Fatalf("stillAbsentAlias must be restored to auth-b, got %q, %v", got, ok) + } +} + func TestSessionAffinitySelector_ThreeScenarios(t *testing.T) { t.Parallel() @@ -1356,6 +1400,7 @@ func TestSessionAffinitySelector_ThreeScenarios(t *testing.T) { opts3 := cliproxyexecutor.Options{OriginalRequest: openaiS3} picked2, _ := selector.Pick(context.Background(), "test", "model", opts2, auths) + selector.OnResult(Result{AuthID: picked2.ID, Provider: "test", Model: "model", Options: opts2, Success: true}) picked3, _ := selector.Pick(context.Background(), "test", "model", opts3, auths) if picked2.ID != picked3.ID { @@ -1371,6 +1416,7 @@ func TestSessionAffinitySelector_ThreeScenarios(t *testing.T) { opts2 := cliproxyexecutor.Options{OriginalRequest: s2} picked1, _ := selector.Pick(context.Background(), "inherit", "model", opts1, auths) + selector.OnResult(Result{AuthID: picked1.ID, Provider: "inherit", Model: "model", Options: opts1, Success: true}) picked2, _ := selector.Pick(context.Background(), "inherit", "model", opts2, auths) if picked1.ID != picked2.ID { @@ -1408,6 +1454,7 @@ func TestSessionAffinitySelectorBodyIdentifierTransitionsPreserveBinding(t *test if err != nil { t.Fatalf("first Pick() error = %v", err) } + selector.OnResult(Result{AuthID: first.ID, Provider: provider, Model: "gpt-test", Options: cliproxyexecutor.Options{OriginalRequest: tt.firstPayload}, Success: true}) second, err := selector.Pick(context.Background(), provider, "gpt-test", cliproxyexecutor.Options{OriginalRequest: bothPayload}, auths) if err != nil { t.Fatalf("combined-identifier Pick() error = %v", err) @@ -1436,6 +1483,7 @@ func TestSessionAffinitySelectorCombinedIdentifiersBindConversationFallback(t *t if err != nil { t.Fatalf("combined-identifier Pick() error = %v", err) } + selector.OnResult(Result{AuthID: first.ID, Provider: provider, Model: "gpt-test", Options: cliproxyexecutor.Options{OriginalRequest: combined}, Success: true}) second, err := selector.Pick(context.Background(), provider, "gpt-test", cliproxyexecutor.Options{OriginalRequest: conversationOnly}, auths) if err != nil { t.Fatalf("conversation-only Pick() error = %v", err) @@ -1445,6 +1493,459 @@ func TestSessionAffinitySelectorCombinedIdentifiersBindConversationFallback(t *t } } +func TestSessionAffinitySelectorFailureQuarantinesAllAliases(t *testing.T) { + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: &RoundRobinSelector{}, + TTL: time.Minute, + }) + defer selector.Stop() + auths := []*Auth{{ID: "auth-a"}, {ID: "auth-b"}} + provider := "responses-alias-group-failure" + model := "gpt-test" + + combined := cliproxyexecutor.Options{OriginalRequest: []byte(`{"conversation":{"id":"conversation-session"},"prompt_cache_key":"shared-cache-bucket"}`), Metadata: map[string]any{}} + first, err := selector.Pick(context.Background(), provider, model, combined, auths) + if err != nil { + t.Fatalf("combined-identifier Pick() error = %v", err) + } + selector.OnResult(Result{AuthID: first.ID, Provider: provider, Model: model, Options: combined, Success: true}) + + promptOnly := cliproxyexecutor.Options{OriginalRequest: []byte(`{"prompt_cache_key":"shared-cache-bucket"}`), Metadata: map[string]any{}} + failed, err := selector.Pick(context.Background(), provider, model, promptOnly, auths) + if err != nil { + t.Fatalf("prompt-only Pick() error = %v", err) + } + if failed.ID != first.ID { + t.Fatalf("prompt-only alias selected %q, want %q", failed.ID, first.ID) + } + selector.OnResult(Result{AuthID: failed.ID, Provider: provider, Model: model, Options: promptOnly, Error: &Error{Code: "upstream_failed", Message: "upstream failed", Retryable: true}}) + + conversationOnly := cliproxyexecutor.Options{OriginalRequest: []byte(`{"conversation":{"id":"conversation-session"}}`), Metadata: map[string]any{}} + next, err := selector.Pick(context.Background(), provider, model, conversationOnly, auths) + if err != nil { + t.Fatalf("conversation-only Pick() error = %v", err) + } + if next.ID == failed.ID { + t.Fatalf("conversation alias reused failed auth %q", failed.ID) + } +} + +func TestSessionCacheCompareAndReplaceAliasesPreservesNewerBinding(t *testing.T) { + cache := NewSessionCache(time.Minute) + defer cache.Stop() + cache.SetAliases("auth-a", "prompt", "conversation") + _, gen1, aliases1, _ := cache.GetWithGeneration("prompt") + cache.SetAliases("auth-b", "prompt", "conversation") + + if replaced := cache.CompareAndReplaceAliases("auth-a", gen1, aliases1, "auth-c"); replaced { + t.Fatal("CompareAndReplaceAliases() succeeded for stale auth, want false") + } + for _, key := range []string{"prompt", "conversation"} { + if got, ok := cache.Get(key); !ok || got != "auth-b" { + t.Fatalf("cache.Get(%q) = %q, %v; want auth-b, true", key, got, ok) + } + } +} + +func TestSessionAffinitySelectorUnavailableCachedAuthRebindsFullAliasGroup(t *testing.T) { + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: &RoundRobinSelector{}, + TTL: time.Minute, + }) + defer selector.Stop() + // Three auths so the fallback replacement is stable and never collides with the stale cached auth. + auths := []*Auth{{ID: "auth-a"}, {ID: "auth-b"}, {ID: "auth-c"}} + provider := "responses-alias-group-recovery" + model := "gpt-test" + + combined := []byte(`{"conversation":{"id":"conversation-session"},"prompt_cache_key":"shared-cache-bucket"}`) + combinedOpts := cliproxyexecutor.Options{OriginalRequest: combined} + first, err := selector.Pick(context.Background(), provider, model, combinedOpts, auths) + if err != nil { + t.Fatalf("combined Pick() error = %v", err) + } + bound := first.ID + selector.OnResult(Result{AuthID: bound, Provider: provider, Model: model, Options: combinedOpts, Success: true}) + + // Make the bound auth unavailable so it leaves the fallback reselect candidate set. + for _, a := range auths { + if a.ID == bound { + a.Unavailable = true + a.NextRetryAfter = time.Now().Add(time.Hour) + } + } + + // Prompt-cache-only request triggers the stale-primary recovery path. + promptOpts := cliproxyexecutor.Options{OriginalRequest: []byte(`{"prompt_cache_key":"shared-cache-bucket"}`)} + replacement, err := selector.Pick(context.Background(), provider, model, promptOpts, auths) + if err != nil { + t.Fatalf("prompt-only recovery Pick() error = %v", err) + } + if replacement.ID == bound { + t.Fatalf("recovery reused unavailable cached auth %q", bound) + } + + // The full alias group must now point at the replacement auth: a conversation-only request + // (touching the sibling alias the cache key likely chose as primary) must also recover to it. + conversationOpts := cliproxyexecutor.Options{OriginalRequest: []byte(`{"conversation":{"id":"conversation-session"}}`)} + for i := 0; i < 3; i++ { + picked, errPick := selector.Pick(context.Background(), provider, model, conversationOpts, auths) + if errPick != nil { + t.Fatalf("conversation-only Pick() #%d error = %v", i, errPick) + } + if picked.ID != replacement.ID { + t.Fatalf("conversation alias did not inherit replacement auth %q: got %q", replacement.ID, picked.ID) + } + } +} + +func TestSessionAffinitySelectorCachedAuthUnavailableRebindsWholeAliasGroup(t *testing.T) { + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: &RoundRobinSelector{}, + TTL: time.Minute, + }) + defer selector.Stop() + auths := []*Auth{{ID: "auth-a"}, {ID: "auth-b"}} + provider := "responses-cache-rebound" + model := "gpt-test" + + combined := cliproxyexecutor.Options{OriginalRequest: []byte(`{"conversation":{"id":"conversation-session"},"prompt_cache_key":"shared-cache-bucket"}`)} + first, err := selector.Pick(context.Background(), provider, model, combined, auths) + if err != nil { + t.Fatalf("Pick() combined error = %v", err) + } + if first.ID != "auth-a" { + t.Fatalf("first Pick() = %q, want auth-a", first.ID) + } + selector.OnResult(Result{AuthID: first.ID, Provider: provider, Model: model, Options: combined, Success: true}) + + // The bound auth leaves the candidate set entirely, mirroring an unavailable cached auth. + availableWithoutFirst := []*Auth{{ID: "auth-b"}} + for _, payload := range []struct { + name string + request []byte + }{ + {name: "primary prompt", request: []byte(`{"prompt_cache_key":"shared-cache-bucket"}`)}, + {name: "fallback conversation", request: []byte(`{"conversation":{"id":"conversation-session"}}`)}, + } { + opts := cliproxyexecutor.Options{OriginalRequest: payload.request} + picked, errPick := selector.Pick(context.Background(), provider, model, opts, availableWithoutFirst) + if errPick != nil { + t.Fatalf("%s Pick() error = %v", payload.name, errPick) + } + if picked.ID != "auth-b" { + t.Fatalf("%s alias selected %q, want rebound %q", payload.name, picked.ID, "auth-b") + } + selector.OnResult(Result{AuthID: picked.ID, Provider: provider, Model: model, Options: opts, Success: true}) + } +} + +func TestSessionAffinitySelectorCachedAuthUnavailableRebindsSharedPromptGroup(t *testing.T) { + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: &RoundRobinSelector{}, + TTL: time.Minute, + }) + defer selector.Stop() + auths := []*Auth{{ID: "auth-a"}, {ID: "auth-b"}} + provider := "responses-cache-rebound-shared" + model := "gpt-test" + + combinedA := cliproxyexecutor.Options{OriginalRequest: []byte(`{"conversation":{"id":"conversation-a"},"prompt_cache_key":"shared-cache-bucket"}`)} + combinedB := cliproxyexecutor.Options{OriginalRequest: []byte(`{"conversation":{"id":"conversation-b"},"prompt_cache_key":"shared-cache-bucket"}`)} + first, err := selector.Pick(context.Background(), provider, model, combinedA, auths) + if err != nil { + t.Fatalf("Pick() A error = %v", err) + } + selector.OnResult(Result{AuthID: first.ID, Provider: provider, Model: model, Options: combinedA, Success: true}) + if first.ID != "auth-a" { + t.Fatalf("first Pick() = %q, want auth-a", first.ID) + } + second, err := selector.Pick(context.Background(), provider, model, combinedB, auths) + if err != nil { + t.Fatalf("Pick() B error = %v", err) + } + selector.OnResult(Result{AuthID: second.ID, Provider: provider, Model: model, Options: combinedB, Success: true}) + if second.ID != first.ID { + t.Fatalf("shared prompt key changed auth from %q to %q", first.ID, second.ID) + } + + availableWithoutFirst := []*Auth{{ID: "auth-b"}} + for _, payload := range []struct { + name string + request []byte + }{ + {name: "conversation A", request: []byte(`{"conversation":{"id":"conversation-a"}}`)}, + {name: "conversation B", request: []byte(`{"conversation":{"id":"conversation-b"}}`)}, + } { + opts := cliproxyexecutor.Options{OriginalRequest: payload.request} + picked, errPick := selector.Pick(context.Background(), provider, model, opts, availableWithoutFirst) + if errPick != nil { + t.Fatalf("%s Pick() error = %v", payload.name, errPick) + } + if picked.ID != "auth-b" { + t.Fatalf("%s alias selected %q, want rebound %q", payload.name, picked.ID, "auth-b") + } + selector.OnResult(Result{AuthID: picked.ID, Provider: provider, Model: model, Options: opts, Success: true}) + } +} + +func TestSessionAffinitySelectorCachedAuthUnavailableConcurrencyNewerBinding(t *testing.T) { + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: &RoundRobinSelector{}, + TTL: time.Minute, + }) + defer selector.Stop() + auths := []*Auth{{ID: "auth-a"}, {ID: "auth-b"}} + provider := "responses-cache-rebound-concurrent" + model := "gpt-test" + + combined := cliproxyexecutor.Options{OriginalRequest: []byte(`{"conversation":{"id":"conversation-session"},"prompt_cache_key":"shared-cache-bucket"}`)} + first, err := selector.Pick(context.Background(), provider, model, combined, auths) + if err != nil { + t.Fatalf("Pick() error = %v", err) + } + selector.OnResult(Result{AuthID: first.ID, Provider: provider, Model: model, Options: combined, Success: true}) + + // A concurrent request sees the bound auth unavailable and rebounds the group to auth-b. + availableWithoutFirst := []*Auth{{ID: "auth-b"}} + picked, err := selector.Pick(context.Background(), provider, model, combined, availableWithoutFirst) + if err != nil { + t.Fatalf("unavailable Pick() error = %v", err) + } + selector.OnResult(Result{AuthID: picked.ID, Provider: provider, Model: model, Options: combined, Success: true}) + if picked.ID != "auth-b" { + t.Fatalf("unavailable Pick() = %q, want auth-b", picked.ID) + } + + // The newer binding must survive subsequent requests without reverting to auth-a. + for i := 0; i < 5; i++ { + got, errPick := selector.Pick(context.Background(), provider, model, combined, availableWithoutFirst) + if errPick != nil { + t.Fatalf("concurrent Pick() #%d error = %v", i, errPick) + } + if got.ID != "auth-b" { + t.Fatalf("concurrent Pick() #%d = %q, want auth-b (newer binding preserved)", i, got.ID) + } + } +} + +type interceptingFallbackSelector struct { + inner Selector + onPick func() +} + +func (s *interceptingFallbackSelector) Pick(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth) (*Auth, error) { + if s.onPick != nil { + s.onPick() + } + return s.inner.Pick(ctx, provider, model, opts, auths) +} + +func TestSessionAffinitySelector_CachedAuthUnavailableConcurrencyFallbackGroupRebind(t *testing.T) { + fallback := &interceptingFallbackSelector{inner: &RoundRobinSelector{}} + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: fallback, + TTL: time.Minute, + }) + defer selector.Stop() + auths := []*Auth{{ID: "auth-a"}, {ID: "auth-b"}} + provider := "responses-fallback-rebound-concurrent" + model := "gpt-test" + + // 1. Initial request with conversation ID only (bound to auth-a) + convOpts := cliproxyexecutor.Options{OriginalRequest: []byte(`{"conversation":{"id":"conversation-session"}}`)} + first, err := selector.Pick(context.Background(), provider, model, convOpts, auths) + if err != nil { + t.Fatalf("first Pick() error = %v", err) + } + if first.ID != "auth-a" { + t.Fatalf("first Pick() = %q, want auth-a", first.ID) + } + selector.OnResult(Result{AuthID: first.ID, Provider: provider, Model: model, Options: convOpts, Success: true}) + + // 2. Prepare rebind with prompt_cache_key (primary) + conversation.id (fallback). + // Primary key is absent in cache; fallback key is present. + // When fallback.Pick is invoked (after selector observes the fallback key), + // simulate concurrent modification of the fallback group in cache to invalidate expectedGen. + combinedOpts := cliproxyexecutor.Options{OriginalRequest: []byte(`{"conversation":{"id":"conversation-session"},"prompt_cache_key":"shared-cache-bucket"}`)} + fallbackKey := provider + "::conv:conversation-session::" + model + + fallback.onPick = func() { + authID, gen, aliases, ok := selector.cache.GetWithGeneration(fallbackKey) + if !ok { + t.Fatalf("expected fallbackKey %q in cache", fallbackKey) + } + // Bump generation concurrently + if !selector.cache.CompareAndReplaceAliases(authID, gen, aliases, authID, fallbackKey) { + t.Fatalf("CompareAndReplaceAliases failed in onPick") + } + } + + availableWithoutFirst := []*Auth{{ID: "auth-b"}} + picked, err := selector.Pick(context.Background(), provider, model, combinedOpts, availableWithoutFirst) + if err != nil { + t.Fatalf("rebind Pick() error = %v", err) + } + if picked.ID != "auth-b" { + t.Fatalf("rebind Pick() = %q, want auth-b", picked.ID) + } + selector.OnResult(Result{AuthID: picked.ID, Provider: provider, Model: model, Options: combinedOpts, Success: true}) + + // 3. Subsequent request with conversation only must route to rebound auth-b + fallback.onPick = nil + convQueryOpts := cliproxyexecutor.Options{OriginalRequest: []byte(`{"conversation":{"id":"conversation-session"}}`)} + queryPicked, err := selector.Pick(context.Background(), provider, model, convQueryOpts, auths) + if err != nil { + t.Fatalf("subsequent conversation Pick() error = %v", err) + } + if queryPicked.ID != "auth-b" { + t.Fatalf("subsequent conversation Pick() = %q, want rebound auth-b", queryPicked.ID) + } +} + +func TestSessionAffinitySelector_ABAMutationCycleRejected(t *testing.T) { + cache := NewSessionCache(time.Minute) + defer cache.Stop() + + sessionA := "openai::conv:session-aba::gpt-test" + promptA := "openai::pck:prompt-aba::gpt-test" + + // 1. Initial binding: auth-a (Generation 1) + cache.SetAliases("auth-a", sessionA, promptA) + authID, gen1, aliases1, ok := cache.GetWithGeneration(sessionA) + if !ok || authID != "auth-a" || gen1 == 0 { + t.Fatalf("initial binding failed: authID=%q, gen=%d, ok=%v", authID, gen1, ok) + } + + // 2. Mutation to auth-b (Generation 2) + cache.SetAliases("auth-b", sessionA, promptA) + authID2, gen2, _, ok2 := cache.GetWithGeneration(sessionA) + if !ok2 || authID2 != "auth-b" || gen2 <= gen1 { + t.Fatalf("second binding failed: authID=%q, gen=%d", authID2, gen2) + } + + // 3. Mutation back to auth-a (Generation 3 - ABA cycle) + cache.SetAliases("auth-a", sessionA, promptA) + authID3, gen3, _, ok3 := cache.GetWithGeneration(sessionA) + if !ok3 || authID3 != "auth-a" || gen3 <= gen2 { + t.Fatalf("third binding failed: authID=%q, gen=%d", authID3, gen3) + } + + // 4. Stale CAS attempting to replace gen1 auth-a with auth-c MUST fail + casSuccess := cache.CompareAndReplaceAliases("auth-a", gen1, aliases1, "auth-c") + if casSuccess { + t.Fatal("ABA CAS replacement succeeded unexpectedly with stale generation token") + } + + // Verify cache still retains gen3 auth-a intact + currentAuth, currentGen, _, okCurrent := cache.GetWithGeneration(sessionA) + if !okCurrent || currentAuth != "auth-a" || currentGen != gen3 { + t.Fatalf("cache was corrupted by failed ABA CAS: auth=%q gen=%d", currentAuth, currentGen) + } +} + +func TestSessionAffinitySelector_PartialGroupSplitAbortsWithoutMutation(t *testing.T) { + cache := NewSessionCache(time.Minute) + defer cache.Stop() + + k1 := "openai::conv:session-split-1::gpt-test" + k2 := "openai::conv:session-split-2::gpt-test" + + // 1. Bind group {k1, k2} to auth-a + cache.SetAliases("auth-a", k1, k2) + _, gen1, aliases1, ok := cache.GetWithGeneration(k1) + if !ok { + t.Fatal("initial SetAliases failed") + } + + // 2. Invalidate k1 (splits group, k2 gets updated aliases and gen2) + cache.Invalidate(k1) + + // 3. Stale CAS expecting {k1, k2} at gen1 must fail because k1 is gone and k2 generation/aliases changed + casSuccess := cache.CompareAndReplaceAliases("auth-a", gen1, aliases1, "auth-b") + if casSuccess { + t.Fatal("partial group split CAS succeeded unexpectedly") + } + + // k2 must remain bound to auth-a with its updated group + authK2, _, aliasesK2, okK2 := cache.GetWithGeneration(k2) + if !okK2 || authK2 != "auth-a" || len(aliasesK2) != 1 || aliasesK2[0] != k2 { + t.Fatalf("k2 corrupted after aborted CAS: ok=%v auth=%q aliases=%v", okK2, authK2, aliasesK2) + } +} + +func TestSessionAffinitySelector_AdditionalAliasBelongsToOtherActiveGroupAbortsCAS(t *testing.T) { + cache := NewSessionCache(time.Minute) + defer cache.Stop() + + g1Session := "openai::conv:group1::gpt-test" + g2Session := "openai::conv:group2::gpt-test" + + // Group 1 -> auth-a + cache.SetAliases("auth-a", g1Session) + _, gen1, aliases1, _ := cache.GetWithGeneration(g1Session) + + // Group 2 -> auth-b + cache.SetAliases("auth-b", g2Session) + + // CAS attempting to rebind group 1 to auth-c while adding g2Session (which belongs to live group 2) must abort + casSuccess := cache.CompareAndReplaceAliases("auth-a", gen1, aliases1, "auth-c", g2Session) + if casSuccess { + t.Fatal("CAS with conflicting foreign active alias succeeded unexpectedly") + } + + // Verify group 2 remains intact on auth-b + authG2, _ := cache.Get(g2Session) + if authG2 != "auth-b" { + t.Fatalf("group 2 was corrupted: auth=%q, want auth-b", authG2) + } +} + +type failingTestSelector struct{} + +func (f *failingTestSelector) Pick(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth) (*Auth, error) { + return nil, fmt.Errorf("fallback selection failed") +} + +func TestSessionAffinitySelector_FallbackFailureKeepsOriginalGroupIntact(t *testing.T) { + failingFallback := &failingTestSelector{} + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: failingFallback, + TTL: time.Minute, + }) + defer selector.Stop() + + provider := "responses-fallback-failure" + model := "gpt-test" + sessionID := "conv:fallback-fail-test" + cacheKey := provider + "::" + sessionID + "::" + model + + // Pre-seed cache with auth-a + selector.cache.Set(cacheKey, "auth-a") + authBefore, okBefore := selector.cache.Get(cacheKey) + if !okBefore || authBefore != "auth-a" { + t.Fatalf("failed to seed cache: got %q, %v", authBefore, okBefore) + } + + opts := cliproxyexecutor.Options{ + OriginalRequest: []byte(`{"conversation":{"id":"fallback-fail-test"}}`), + } + + // Request with only auth-b available (auth-a unavailable). + // Fallback selector will fail with error. + availableAuths := []*Auth{{ID: "auth-b"}} + picked, err := selector.Pick(context.Background(), provider, model, opts, availableAuths) + if err == nil { + t.Fatalf("expected error from failing fallback selector, got auth=%v", picked) + } + + // Cache MUST still contain auth-a (no eager delete). + authAfter, okAfter := selector.cache.Get(cacheKey) + if !okAfter || authAfter != "auth-a" { + t.Fatalf("cache was eagerly deleted or corrupted on fallback failure: got %q, %v", authAfter, okAfter) + } +} + func TestSessionAffinitySelectorPrimaryTrafficKeepsConversationAliasAlive(t *testing.T) { selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ Fallback: &RoundRobinSelector{}, @@ -1462,6 +1963,7 @@ func TestSessionAffinitySelectorPrimaryTrafficKeepsConversationAliasAlive(t *tes if err != nil { t.Fatalf("combined Pick() error = %v", err) } + selector.OnResult(Result{AuthID: first.ID, Provider: provider, Model: model, Options: cliproxyexecutor.Options{OriginalRequest: combined}, Success: true}) conversationKey := provider + "::conv:conversation-session::" + model selector.cache.mu.Lock() conversationEntry := selector.cache.entries[conversationKey] @@ -1473,6 +1975,7 @@ func TestSessionAffinitySelectorPrimaryTrafficKeepsConversationAliasAlive(t *tes if err != nil { t.Fatalf("prompt-only Pick() error = %v", err) } + selector.OnResult(Result{AuthID: primary.ID, Provider: provider, Model: model, Options: cliproxyexecutor.Options{OriginalRequest: promptOnly}, Success: true}) if primary.ID != first.ID { t.Fatalf("prompt-only auth = %q, want %q", primary.ID, first.ID) } @@ -1504,10 +2007,12 @@ func TestSessionAffinitySelectorSharedPromptKeyPreservesConversationAliases(t *t if err != nil { t.Fatalf("conversation A combined Pick() error = %v", err) } + selector.OnResult(Result{AuthID: first.ID, Provider: provider, Model: model, Options: cliproxyexecutor.Options{OriginalRequest: combinedA}, Success: true}) second, err := selector.Pick(context.Background(), provider, model, cliproxyexecutor.Options{OriginalRequest: combinedB}, auths) if err != nil { t.Fatalf("conversation B combined Pick() error = %v", err) } + selector.OnResult(Result{AuthID: second.ID, Provider: provider, Model: model, Options: cliproxyexecutor.Options{OriginalRequest: combinedB}, Success: true}) if second.ID != first.ID { t.Fatalf("shared prompt key changed auth from %q to %q", first.ID, second.ID) } @@ -1538,6 +2043,7 @@ func TestSessionAffinitySelectorConversationIDContainingPromptMarkerRemainsStabl if err != nil { t.Fatalf("combined Pick() error = %v", err) } + selector.OnResult(Result{AuthID: first.ID, Provider: provider, Model: model, Options: cliproxyexecutor.Options{OriginalRequest: combined}, Success: true}) second, err := selector.Pick(context.Background(), provider, model, cliproxyexecutor.Options{OriginalRequest: conversationOnly}, auths) if err != nil { t.Fatalf("conversation-only Pick() error = %v", err) @@ -1627,6 +2133,7 @@ func TestSessionAffinitySelector_MultiModelSession(t *testing.T) { if pickedA.ID != "auth-a" { t.Fatalf("Pick() for model-a = %q, want auth-a", pickedA.ID) } + selector.OnResult(Result{AuthID: pickedA.ID, Provider: "provider", Model: "model-a", Options: opts, Success: true}) // Request model-b with only auth-b available for that model authsForModelB := []*Auth{authB} @@ -1637,6 +2144,7 @@ func TestSessionAffinitySelector_MultiModelSession(t *testing.T) { if pickedB.ID != "auth-b" { t.Fatalf("Pick() for model-b = %q, want auth-b", pickedB.ID) } + selector.OnResult(Result{AuthID: pickedB.ID, Provider: "provider", Model: "model-b", Options: opts, Success: true}) // Switch back to model-a - should still get auth-a (separate binding per model) pickedA2, err := selector.Pick(context.Background(), "provider", "model-a", opts, authsForModelA) @@ -1723,6 +2231,7 @@ func TestSessionAffinitySelector_CrossProviderIsolation(t *testing.T) { if pickedClaude.ID != "auth-claude" { t.Fatalf("Pick() for claude = %q, want auth-claude", pickedClaude.ID) } + selector.OnResult(Result{AuthID: pickedClaude.ID, Provider: "claude", Model: "claude-3", Options: opts, Success: true}) // Same session but via gemini provider should get different auth pickedGemini, err := selector.Pick(context.Background(), "gemini", "gemini-2.5-pro", opts, []*Auth{authGemini}) @@ -1732,6 +2241,7 @@ func TestSessionAffinitySelector_CrossProviderIsolation(t *testing.T) { if pickedGemini.ID != "auth-gemini" { t.Fatalf("Pick() for gemini = %q, want auth-gemini", pickedGemini.ID) } + selector.OnResult(Result{AuthID: pickedGemini.ID, Provider: "gemini", Model: "gemini-2.5-pro", Options: opts, Success: true}) // Verify both bindings remain stable for i := 0; i < 5; i++ { @@ -1845,6 +2355,7 @@ func TestSessionAffinitySelector_Concurrent(t *testing.T) { t.Fatalf("Initial Pick() error = %v", err) } expectedID := first.ID + selector.OnResult(Result{AuthID: first.ID, Provider: "claude", Model: "claude-3", Options: opts, Success: true}) start := make(chan struct{}) var wg sync.WaitGroup @@ -1887,6 +2398,58 @@ func TestSessionAffinitySelector_Concurrent(t *testing.T) { } } +func TestSessionAffinitySelector_ConcurrentCacheMissBindsOneAuth(t *testing.T) { + t.Parallel() + + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: &RoundRobinSelector{}, + TTL: time.Minute, + }) + defer selector.Stop() + + auths := []*Auth{{ID: "auth-a"}, {ID: "auth-b"}, {ID: "auth-c"}} + opts := cliproxyexecutor.Options{OriginalRequest: []byte(`{"metadata":{"user_id":"user_xxx_account__session_concurrent-cache-miss"}}`)} + + const goroutines = 64 + start := make(chan struct{}) + results := make(chan string, goroutines) + errs := make(chan error, goroutines) + var wg sync.WaitGroup + for i := 0; i < goroutines; i++ { + wg.Add(1) + go func() { + defer wg.Done() + <-start + auth, err := selector.Pick(context.Background(), "claude", "claude-3", opts, auths) + if err != nil { + errs <- err + return + } + results <- auth.ID + }() + } + close(start) + wg.Wait() + close(results) + close(errs) + + for err := range errs { + t.Fatalf("concurrent cache-miss Pick() error = %v", err) + } + var expected string + for authID := range results { + if expected == "" { + expected = authID + } + if authID != expected { + t.Fatalf("concurrent cache-miss Pick() returned %q after %q was bound", authID, expected) + } + } + if expected == "" { + t.Fatal("concurrent cache-miss Pick() returned no auth") + } +} + func TestExtractSessionIDNativeSignals(t *testing.T) { t.Parallel() tests := []struct { @@ -2123,6 +2686,7 @@ func TestSessionAffinitySelectorUsesRequestPayloadWhenOriginalRequestMissing(t * if errFirst != nil { t.Fatalf("first Pick() error = %v", errFirst) } + selector.OnResult(Result{AuthID: first.ID, Provider: "openai", Model: request.Model, Options: opts, Success: true}) second, errSecond := selector.Pick(context.Background(), "openai", request.Model, opts, auths) if errSecond != nil { t.Fatalf("second Pick() error = %v", errSecond) @@ -2186,6 +2750,7 @@ func TestSessionAffinitySelector_FallbackReselectReceivesOnlyAvailable(t *testin if err != nil { t.Fatalf("Pick() error = %v", err) } + selector.OnResult(Result{AuthID: first.ID, Provider: "claude", Model: "claude-3", Options: opts, Success: true}) // Make the bound auth unavailable so it is filtered out of `available`. bound := first.ID @@ -2213,3 +2778,53 @@ func TestSessionAffinitySelector_FallbackReselectReceivesOnlyAvailable(t *testin } } } + +// TestSessionAffinitySelector_RequestScopedExclusionBreaksCarousel is a +// regression test for the live bug where an auth that just failed (429/5xx/empty) +// was re-picked for the SAME request repeatedly. In home/mixed mode per-attempt +// failures never update local availability (reportHomeResult does not set a local +// cooldown), so `getAvailableAuths` kept reporting the freshly-failed auth as +// available and the session-affinity cache returned it on every retry. The fix +// threads the request-scoped set of failed auth IDs through to the selector so a +// failed auth is never re-picked for the remainder of that request, even though +// it is still locally "available". +func TestSessionAffinitySelector_RequestScopedExclusionBreaksCarousel(t *testing.T) { + t.Parallel() + + rec := &recordingFallbackSelector{inner: &RoundRobinSelector{}} + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: rec, + TTL: time.Minute, + }) + defer selector.Stop() + + auths := []*Auth{ + {ID: "auth-a"}, + {ID: "auth-b"}, + {ID: "auth-c"}, + } + + payload := []byte(`{"metadata":{"user_id":"user_xxx_account__session_carousel-test"}}`) + opts := cliproxyexecutor.Options{OriginalRequest: payload} + + first, err := selector.Pick(context.Background(), "claude", "claude-3", opts, auths) + if err != nil { + t.Fatalf("Pick() error = %v", err) + } + selector.OnResult(Result{AuthID: first.ID, Provider: "claude", Model: "claude-3", Options: opts, Success: true}) + + failed := first.ID + opts.Metadata = map[string]any{ + cliproxyexecutor.ExcludedAuthIDsMetadataKey: map[string]struct{}{failed: {}}, + } + + for attempt := 0; attempt < 20; attempt++ { + got, errPick := selector.Pick(context.Background(), "claude", "claude-3", opts, auths) + if errPick != nil { + t.Fatalf("Pick() attempt %d error = %v", attempt, errPick) + } + if got.ID == failed { + t.Fatalf("attempt %d re-picked auth %q that already failed in this request; want a different auth", attempt, failed) + } + } +} diff --git a/sdk/cliproxy/auth/session_affinity_fix_test.go b/sdk/cliproxy/auth/session_affinity_fix_test.go new file mode 100644 index 000000000..d8be2cd7b --- /dev/null +++ b/sdk/cliproxy/auth/session_affinity_fix_test.go @@ -0,0 +1,938 @@ +package auth + +import ( + "context" + "net/http" + "testing" + "time" + + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +func newTestRoundRobinSelector() *RoundRobinSelector { + return &RoundRobinSelector{ + cursors: make(map[string]int), + } +} + +func TestSessionAffinity_InitialPickBindsBeforeSuccess(t *testing.T) { + authA := &Auth{ID: "auth-a"} + authB := &Auth{ID: "auth-b"} + auths := []*Auth{authA, authB} + + fallback := newTestRoundRobinSelector() + selector := NewSessionAffinitySelector(fallback) + + opts := cliproxyexecutor.Options{ + Headers: http.Header{"X-Session-Id": []string{"sess-12345678"}}, + } + + picked, err := selector.Pick(context.Background(), "provider", "model", opts, auths) + if err != nil || picked == nil { + t.Fatalf("Pick failed: err=%v, picked=%v", err, picked) + } + + cacheKey := "provider::header:sess-12345678::model" + bound, ok := selector.cache.Get(cacheKey) + if !ok || bound != picked.ID { + t.Fatalf("expected cacheKey %q to be pre-bound to %q, got bound=%q ok=%v", cacheKey, picked.ID, bound, ok) + } +} + +func TestSessionAffinity_SuccessfulResultBindsChosenAuth(t *testing.T) { + authA := &Auth{ID: "auth-a"} + authB := &Auth{ID: "auth-b"} + auths := []*Auth{authA, authB} + + fallback := newTestRoundRobinSelector() + selector := NewSessionAffinitySelector(fallback) + + opts := cliproxyexecutor.Options{ + Headers: http.Header{"X-Session-Id": []string{"sess-12345678"}}, + } + + picked, err := selector.Pick(context.Background(), "provider", "model", opts, auths) + if err != nil || picked == nil { + t.Fatalf("Pick failed: err=%v, picked=%v", err, picked) + } + + selector.OnResult(Result{ + AuthID: picked.ID, + Provider: "provider", + Model: "model", + Success: true, + Options: opts, + }) + + cacheKey := "provider::header:sess-12345678::model" + bound, ok := selector.cache.Get(cacheKey) + if !ok || bound != picked.ID { + t.Fatalf("expected cacheKey %q to be bound to %q, got bound=%q ok=%v", cacheKey, picked.ID, bound, ok) + } +} + +func TestSessionAffinity_RetryableFailureInvalidatesMatchingBinding(t *testing.T) { + authA := &Auth{ID: "auth-a"} + + fallback := newTestRoundRobinSelector() + selector := NewSessionAffinitySelector(fallback) + + opts := cliproxyexecutor.Options{ + Headers: http.Header{"X-Session-Id": []string{"sess-12345678"}}, + } + + // Bind authA first + selector.OnResult(Result{ + AuthID: authA.ID, + Provider: "provider", + Model: "model", + Success: true, + Options: opts, + }) + + cacheKey := "provider::header:sess-12345678::model" + if _, ok := selector.cache.Get(cacheKey); !ok { + t.Fatalf("precondition failed: cache key should be bound") + } + + // Retryable failure (429 Rate Limit) + selector.OnResult(Result{ + AuthID: authA.ID, + Provider: "provider", + Model: "model", + Success: false, + Error: &Error{HTTPStatus: http.StatusTooManyRequests, Message: "rate limit"}, + Options: opts, + }) + + if bound, ok := selector.cache.Get(cacheKey); ok { + t.Fatalf("expected cacheKey %q to be invalidated on 429 failure, but still bound to %q", cacheKey, bound) + } +} + +func TestSessionAffinity_StaleFailureCannotDeleteNewerSuccess(t *testing.T) { + authA := &Auth{ID: "auth-a"} + authB := &Auth{ID: "auth-b"} + + fallback := newTestRoundRobinSelector() + selector := NewSessionAffinitySelector(fallback) + + opts := cliproxyexecutor.Options{ + Headers: http.Header{"X-Session-Id": []string{"sess-shared-12345"}}, + } + + // Request 2 succeeds on Auth B and binds it + selector.OnResult(Result{ + AuthID: authB.ID, + Provider: "provider", + Model: "model", + Success: true, + Options: opts, + }) + + // Request 1 (older execution) fails on Auth A with retryable error + selector.OnResult(Result{ + AuthID: authA.ID, + Provider: "provider", + Model: "model", + Success: false, + Error: &Error{HTTPStatus: http.StatusTooManyRequests, Message: "rate limit"}, + Options: opts, + }) + + // Cache MUST still point to Auth B + cacheKey := "provider::header:sess-shared-12345::model" + bound, ok := selector.cache.Get(cacheKey) + if !ok || bound != authB.ID { + t.Fatalf("stale failure for %q deleted newer binding %q; cache state bound=%q ok=%v", authA.ID, authB.ID, bound, ok) + } +} + +func TestSessionAffinity_ExhaustedRequestDoesNotLeaveLastFailedAuthBound(t *testing.T) { + authA := &Auth{ID: "auth-a"} + authB := &Auth{ID: "auth-b"} + auths := []*Auth{authA, authB} + + fallback := newTestRoundRobinSelector() + selector := NewSessionAffinitySelector(fallback) + + opts := cliproxyexecutor.Options{ + Headers: http.Header{"X-Session-Id": []string{"sess-exhausted-123"}}, + } + + // Attempt 1: picks Auth A, fails 429 + picked1, _ := selector.Pick(context.Background(), "provider", "model", opts, auths) + selector.OnResult(Result{ + AuthID: picked1.ID, + Provider: "provider", + Model: "model", + Success: false, + Error: &Error{HTTPStatus: http.StatusTooManyRequests}, + Options: opts, + }) + + // Attempt 2 within request: excludes Auth A, picks Auth B, fails 429 + opts2 := cliproxyexecutor.Options{ + Headers: http.Header{"X-Session-Id": []string{"sess-exhausted-123"}}, + Metadata: map[string]any{ + cliproxyexecutor.ExcludedAuthIDsMetadataKey: map[string]struct{}{picked1.ID: {}}, + }, + } + picked2, _ := selector.Pick(context.Background(), "provider", "model", opts2, auths) + selector.OnResult(Result{ + AuthID: picked2.ID, + Provider: "provider", + Model: "model", + Success: false, + Error: &Error{HTTPStatus: http.StatusTooManyRequests}, + Options: opts2, + }) + + // Verify session cache is left clean (unbound) + cacheKey := "provider::header:sess-exhausted-123::model" + if bound, ok := selector.cache.Get(cacheKey); ok { + t.Fatalf("exhausted request left last failed auth bound=%q in cache", bound) + } +} + +func TestSessionAffinity_ExistingWithinRequestExclusionPreventsRepeat(t *testing.T) { + authA := &Auth{ID: "auth-a"} + authB := &Auth{ID: "auth-b"} + auths := []*Auth{authA, authB} + + fallback := newTestRoundRobinSelector() + selector := NewSessionAffinitySelector(fallback) + + opts1 := cliproxyexecutor.Options{ + Headers: http.Header{"X-Session-Id": []string{"sess-repeat-12345"}}, + } + + picked1, _ := selector.Pick(context.Background(), "provider", "model", opts1, auths) + + // Next attempt excludes Auth A + opts2 := cliproxyexecutor.Options{ + Headers: http.Header{"X-Session-Id": []string{"sess-repeat-12345"}}, + Metadata: map[string]any{ + cliproxyexecutor.ExcludedAuthIDsMetadataKey: map[string]struct{}{picked1.ID: {}}, + }, + } + + picked2, _ := selector.Pick(context.Background(), "provider", "model", opts2, auths) + if picked2.ID == picked1.ID { + t.Fatalf("within-request exclusion failed: pick returned excluded auth %q", picked1.ID) + } +} + +func TestSessionAffinity_ClientValidationErrorDoesNotRotateAffinity(t *testing.T) { + authA := &Auth{ID: "auth-a"} + + fallback := newTestRoundRobinSelector() + selector := NewSessionAffinitySelector(fallback) + + opts := cliproxyexecutor.Options{ + Headers: http.Header{"X-Session-Id": []string{"sess-client-err-1"}}, + } + + // Bind authA + selector.OnResult(Result{ + AuthID: authA.ID, + Provider: "provider", + Model: "model", + Success: true, + Options: opts, + }) + + cacheKey := "provider::header:sess-client-err-1::model" + + // Client request validation error (400 Bad Request, request-scoped) + selector.OnResult(Result{ + AuthID: authA.ID, + Provider: "provider", + Model: "model", + Success: false, + Error: &Error{Code: requestScopedErrorCode, HTTPStatus: http.StatusBadRequest, Message: "invalid param"}, + Options: opts, + }) + + // Binding MUST be preserved for client request validation errors + bound, ok := selector.cache.Get(cacheKey) + if !ok || bound != authA.ID { + t.Fatalf("client validation error removed session affinity; expected bound=%q, got bound=%q ok=%v", authA.ID, bound, ok) + } +} + +type pickFuncSelector func(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, available []*Auth) (*Auth, error) + +func (f pickFuncSelector) Pick(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, available []*Auth) (*Auth, error) { + return f(ctx, provider, model, opts, available) +} + +func closedStreamChunks() <-chan cliproxyexecutor.StreamChunk { + ch := make(chan cliproxyexecutor.StreamChunk) + close(ch) + return ch +} + +func TestSessionAffinity_CachedAuthUnavailableRebindsFallback(t *testing.T) { + authA := &Auth{ID: "auth-a"} + authB := &Auth{ID: "auth-b"} + + fallback := pickFuncSelector(func(_ context.Context, _, _ string, _ cliproxyexecutor.Options, available []*Auth) (*Auth, error) { + return available[0], nil + }) + selector := NewSessionAffinitySelector(fallback) + + opts := cliproxyexecutor.Options{Headers: http.Header{"X-Session-Id": []string{"sess-unavail-1234"}}} + cacheKey := "provider::header:sess-unavail-1234::model" + + selector.OnResult(Result{AuthID: authA.ID, Provider: "provider", Model: "model", Success: true, Options: opts}) + picked, err := selector.Pick(context.Background(), "provider", "model", opts, []*Auth{authB}) + if err != nil || picked == nil || picked.ID != authB.ID { + t.Fatalf("Pick = %v/%v, want B", picked, err) + } + + bound, ok := selector.cache.Get(cacheKey) + if !ok || bound != authB.ID { + t.Fatalf("expected stale A to be replaced by pre-bound B; cache=%q ok=%v", bound, ok) + } +} + +func TestSessionAffinity_FallbackBFailsLeavesCacheEmpty(t *testing.T) { + authA := &Auth{ID: "auth-a"} + authB := &Auth{ID: "auth-b"} + + fallback := pickFuncSelector(func(_ context.Context, _, _ string, _ cliproxyexecutor.Options, available []*Auth) (*Auth, error) { + return available[0], nil + }) + selector := NewSessionAffinitySelector(fallback) + + opts := cliproxyexecutor.Options{Headers: http.Header{"X-Session-Id": []string{"sess-bfail-1234"}}} + cacheKey := "provider::header:sess-bfail-1234::model" + + selector.OnResult(Result{AuthID: authA.ID, Provider: "provider", Model: "model", Success: true, Options: opts}) + + // A unavailable -> fallback picks B. + picked, _ := selector.Pick(context.Background(), "provider", "model", opts, []*Auth{authB}) + if picked.ID != authB.ID { + t.Fatalf("Pick = %q, want B", picked.ID) + } + // B fails with retryable error. + selector.OnResult(Result{AuthID: picked.ID, Provider: "provider", Model: "model", Success: false, Error: &Error{HTTPStatus: http.StatusTooManyRequests}, Options: opts}) + + // Cache must be empty (no stale A, no B). + if bound, ok := selector.cache.Get(cacheKey); ok { + t.Fatalf("cache should be empty after B failure, got %q", bound) + } + + // Immediate second request starts from normal fallback (A), not stale B affinity. + second, _ := selector.Pick(context.Background(), "provider", "model", opts, []*Auth{authA, authB}) + if second.ID != authA.ID { + t.Fatalf("second request should reselect from fallback, got %q", second.ID) + } +} + +func TestSessionAffinity_FallbackBSucceedsBindsB(t *testing.T) { + authA := &Auth{ID: "auth-a"} + authB := &Auth{ID: "auth-b"} + + fallback := pickFuncSelector(func(_ context.Context, _, _ string, _ cliproxyexecutor.Options, available []*Auth) (*Auth, error) { + return available[0], nil + }) + selector := NewSessionAffinitySelector(fallback) + + opts := cliproxyexecutor.Options{Headers: http.Header{"X-Session-Id": []string{"sess-bok-123456"}}} + cacheKey := "provider::header:sess-bok-123456::model" + + selector.OnResult(Result{AuthID: authA.ID, Provider: "provider", Model: "model", Success: true, Options: opts}) + + picked, _ := selector.Pick(context.Background(), "provider", "model", opts, []*Auth{authB}) + if picked.ID != authB.ID { + t.Fatalf("Pick = %q, want B", picked.ID) + } + selector.OnResult(Result{AuthID: picked.ID, Provider: "provider", Model: "model", Success: true, Options: opts}) + + bound, ok := selector.cache.Get(cacheKey) + if !ok || bound != authB.ID { + t.Fatalf("B should be bound after success; bound=%q ok=%v", bound, ok) + } +} + +func TestSessionAffinity_StreamSuccessThroughWrapperBinds(t *testing.T) { + ctx := context.Background() + manager := NewManager(nil, nil, nil) + affinity := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{Fallback: &RoundRobinSelector{}, TTL: time.Hour}) + defer affinity.Stop() + manager.SetSelector(affinity) + + auth := &Auth{ID: "stream-auth", Provider: "stream-provider", Status: StatusActive} + if _, err := manager.Register(WithSkipPersist(ctx), auth); err != nil { + t.Fatalf("Register: %v", err) + } + + opts := cliproxyexecutor.Options{Headers: http.Header{"X-Session-Id": []string{"stream-sess-12345"}}} + cacheKey := "stream-provider::header:stream-sess-12345::stream-model" + + chunk := cliproxyexecutor.StreamChunk{Payload: []byte("data: {\"id\":\"x\"}\n\n")} + res := manager.wrapStreamResult(ctx, auth, "stream-provider", "stream-model", opts, nil, []cliproxyexecutor.StreamChunk{chunk}, closedStreamChunks(), OAuthModelAliasResult{}, false) + for range res.Chunks { + } + + bound, ok := affinity.cache.Get(cacheKey) + if !ok || bound != auth.ID { + t.Fatalf("stream success should bind auth; bound=%q ok=%v", bound, ok) + } +} + +func TestSessionAffinity_StreamFailureThroughWrapperInvalidates(t *testing.T) { + ctx := context.Background() + manager := NewManager(nil, nil, nil) + affinity := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{Fallback: &RoundRobinSelector{}, TTL: time.Hour}) + defer affinity.Stop() + manager.SetSelector(affinity) + + auth := &Auth{ID: "stream-auth", Provider: "stream-provider", Status: StatusActive} + if _, err := manager.Register(WithSkipPersist(ctx), auth); err != nil { + t.Fatalf("Register: %v", err) + } + + opts := cliproxyexecutor.Options{Headers: http.Header{"X-Session-Id": []string{"stream-sess-12345"}}} + cacheKey := "stream-provider::header:stream-sess-12345::stream-model" + + // Bind first. + affinity.OnResult(Result{AuthID: auth.ID, Provider: "stream-provider", Model: "stream-model", Success: true, Options: opts}) + if _, ok := affinity.cache.Get(cacheKey); !ok { + t.Fatalf("precondition: bound") + } + + // Stream fails with retryable upstream error (503). + errChunk := cliproxyexecutor.StreamChunk{Err: &Error{HTTPStatus: http.StatusServiceUnavailable}} + res := manager.wrapStreamResult(ctx, auth, "stream-provider", "stream-model", opts, nil, []cliproxyexecutor.StreamChunk{errChunk}, closedStreamChunks(), OAuthModelAliasResult{}, false) + for range res.Chunks { + } + + if bound, ok := affinity.cache.Get(cacheKey); ok { + t.Fatalf("stream failure should invalidate affinity; still bound=%q", bound) + } +} +func optsWithMixedNamespace(opts cliproxyexecutor.Options) cliproxyexecutor.Options { + if opts.Metadata == nil { + opts.Metadata = make(map[string]any) + } + opts.Metadata[cliproxyexecutor.SessionAffinityProviderMetadataKey] = "mixed" + return opts +} + +func TestSessionAffinity_MixedNamespace_PickRecordsAndOnResultBindsCanonicalKey(t *testing.T) { + gemini := &Auth{ID: "gemini-auth", Provider: "gemini"} + fallback := pickFuncSelector(func(_ context.Context, _, _ string, _ cliproxyexecutor.Options, available []*Auth) (*Auth, error) { + return available[0], nil + }) + selector := NewSessionAffinitySelector(fallback) + + // Mixed pool: selection provider is literally "mixed", actual auth provider is "gemini". + opts := cliproxyexecutor.Options{Headers: http.Header{"X-Session-Id": []string{"mixed-sess-12345"}}, Metadata: map[string]any{}} + picked, err := selector.Pick(context.Background(), "mixed", "model", opts, []*Auth{gemini}) + if err != nil || picked == nil || picked.ID != gemini.ID { + t.Fatalf("Pick = %v/%v, want gemini", picked, err) + } + // Pick must have recorded the "mixed" namespace into the shared request-local metadata map. + ns, _ := opts.Metadata[cliproxyexecutor.SessionAffinityProviderMetadataKey].(string) + if ns != "mixed" { + t.Fatalf("namespace metadata = %q, want %q", ns, "mixed") + } + + // Successful OnResult binds under the canonical "mixed" key, not the actual provider. + selector.OnResult(Result{AuthID: picked.ID, Provider: "gemini", Model: "model", Success: true, Options: opts}) + + mixedKey := "mixed::header:mixed-sess-12345::model" + geminiKey := "gemini::header:mixed-sess-12345::model" + if _, ok := selector.cache.Get(mixedKey); !ok { + t.Fatalf("expected binding under canonical mixed key") + } + if _, ok := selector.cache.Get(geminiKey); ok { + t.Fatalf("must NOT bind under actual provider key") + } + + // Next Pick under "mixed" hits the bound auth. + next, _ := selector.Pick(context.Background(), "mixed", "model", opts, []*Auth{gemini}) + if next.ID != gemini.ID { + t.Fatalf("second Pick = %q, want gemini", next.ID) + } +} + +func TestSessionAffinity_MixedNamespace_FailureLeavesCacheEmpty(t *testing.T) { + gemini := &Auth{ID: "gemini-auth", Provider: "gemini"} + fallback := pickFuncSelector(func(_ context.Context, _, _ string, _ cliproxyexecutor.Options, available []*Auth) (*Auth, error) { + return available[0], nil + }) + selector := NewSessionAffinitySelector(fallback) + + opts := optsWithMixedNamespace(cliproxyexecutor.Options{Headers: http.Header{"X-Session-Id": []string{"mixed-fail-12345"}}}) + cacheKey := "mixed::header:mixed-fail-12345::model" + + picked, _ := selector.Pick(context.Background(), "mixed", "model", opts, []*Auth{gemini}) + selector.OnResult(Result{AuthID: picked.ID, Provider: "gemini", Model: "model", Success: false, Error: &Error{HTTPStatus: http.StatusTooManyRequests}, Options: opts}) + + if bound, ok := selector.cache.Get(cacheKey); ok { + t.Fatalf("mixed cache should be empty after failure; still bound=%q", bound) + } +} + +func TestSessionAffinity_MixedNamespace_StaleAuthRebindsFallback(t *testing.T) { + authA := &Auth{ID: "auth-a", Provider: "gemini"} + authB := &Auth{ID: "auth-b", Provider: "antigravity"} + fallback := pickFuncSelector(func(_ context.Context, _, _ string, _ cliproxyexecutor.Options, available []*Auth) (*Auth, error) { + return available[0], nil + }) + selector := NewSessionAffinitySelector(fallback) + + opts := optsWithMixedNamespace(cliproxyexecutor.Options{Headers: http.Header{"X-Session-Id": []string{"mixed-stale-12345"}}}) + cacheKey := "mixed::header:mixed-stale-12345::model" + + selector.OnResult(Result{AuthID: authA.ID, Provider: "gemini", Model: "model", Success: true, Options: opts}) + picked, _ := selector.Pick(context.Background(), "mixed", "model", opts, []*Auth{authB}) + if picked.ID != authB.ID { + t.Fatalf("Pick = %q, want B", picked.ID) + } + bound, ok := selector.cache.Get(cacheKey) + if !ok || bound != authB.ID { + t.Fatalf("expected stale A to be replaced by pre-bound B; cache=%q ok=%v", bound, ok) + } +} + +func TestSessionAffinity_MixedNamespace_StaleFailureCannotDeleteNewerSuccess(t *testing.T) { + authA := &Auth{ID: "auth-a", Provider: "gemini"} + authB := &Auth{ID: "auth-b", Provider: "gemini"} + fallback := pickFuncSelector(func(_ context.Context, _, _ string, _ cliproxyexecutor.Options, available []*Auth) (*Auth, error) { + return available[0], nil + }) + selector := NewSessionAffinitySelector(fallback) + + opts := optsWithMixedNamespace(cliproxyexecutor.Options{Headers: http.Header{"X-Session-Id": []string{"mixed-stale-newer-1"}}}) + cacheKey := "mixed::header:mixed-stale-newer-1::model" + + // Newer request succeeds on B under the canonical mixed key. + selector.OnResult(Result{AuthID: authB.ID, Provider: "gemini", Model: "model", Success: true, Options: opts}) + // Older request fails on A; must not delete B's newer binding. + selector.OnResult(Result{AuthID: authA.ID, Provider: "gemini", Model: "model", Success: false, Error: &Error{HTTPStatus: http.StatusTooManyRequests}, Options: opts}) + + bound, ok := selector.cache.Get(cacheKey) + if !ok || bound != authB.ID { + t.Fatalf("stale failure for %q deleted newer binding %q; bound=%q ok=%v", authA.ID, authB.ID, bound, ok) + } +} + +func TestSessionAffinity_MixedNamespace_StreamBindsAndInvalidatesCanonicalKey(t *testing.T) { + ctx := context.Background() + manager := NewManager(nil, nil, nil) + affinity := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{Fallback: &RoundRobinSelector{}, TTL: time.Hour}) + defer affinity.Stop() + manager.SetSelector(affinity) + + auth := &Auth{ID: "stream-mixed-auth", Provider: "gemini", Status: StatusActive} + if _, err := manager.Register(WithSkipPersist(ctx), auth); err != nil { + t.Fatalf("Register: %v", err) + } + + opts := optsWithMixedNamespace(cliproxyexecutor.Options{Headers: http.Header{"X-Session-Id": []string{"stream-mixed-12345"}}}) + cacheKey := "mixed::header:stream-mixed-12345::stream-model" + + // Success binds under the canonical key. + chunk := cliproxyexecutor.StreamChunk{Payload: []byte("data: {\"id\":\"x\"}\n\n")} + res := manager.wrapStreamResult(ctx, auth, "gemini", "stream-model", opts, nil, []cliproxyexecutor.StreamChunk{chunk}, closedStreamChunks(), OAuthModelAliasResult{}, false) + for range res.Chunks { + } + bound, ok := affinity.cache.Get(cacheKey) + if !ok || bound != auth.ID { + t.Fatalf("stream mixed success should bind canonical key; bound=%q ok=%v", bound, ok) + } + + // Failure invalidates the same canonical key. + errChunk := cliproxyexecutor.StreamChunk{Err: &Error{HTTPStatus: http.StatusServiceUnavailable}} + res2 := manager.wrapStreamResult(ctx, auth, "gemini", "stream-model", opts, nil, []cliproxyexecutor.StreamChunk{errChunk}, closedStreamChunks(), OAuthModelAliasResult{}, false) + for range res2.Chunks { + } + if bound, ok := affinity.cache.Get(cacheKey); ok { + t.Fatalf("stream mixed failure should invalidate canonical key; still bound=%q", bound) + } +} + +func TestSessionAffinity_SingleProviderStillUsesActualProviderKey(t *testing.T) { + auth := &Auth{ID: "single-auth", Provider: "claude"} + fallback := pickFuncSelector(func(_ context.Context, _, _ string, _ cliproxyexecutor.Options, available []*Auth) (*Auth, error) { + return available[0], nil + }) + selector := NewSessionAffinitySelector(fallback) + + opts := cliproxyexecutor.Options{Headers: http.Header{"X-Session-Id": []string{"single-sess-12345"}}, Metadata: map[string]any{}} + + // Single provider: Pick records "claude" as namespace; OnResult binds under "claude". + picked, _ := selector.Pick(context.Background(), "claude", "model", opts, []*Auth{auth}) + selector.OnResult(Result{AuthID: picked.ID, Provider: "claude", Model: "model", Success: true, Options: opts}) + + cacheKey := "claude::header:single-sess-12345::model" + bound, ok := selector.cache.Get(cacheKey) + if !ok || bound != auth.ID { + t.Fatalf("single-provider bind failed; bound=%q ok=%v", bound, ok) + } +} + +func TestSessionAffinity_MixedNamespace_SecondRequestSkipsUnavailable(t *testing.T) { + authA := &Auth{ID: "auth-a", Provider: "gemini"} + authB := &Auth{ID: "auth-b", Provider: "gemini"} + fallback := pickFuncSelector(func(_ context.Context, _, _ string, _ cliproxyexecutor.Options, available []*Auth) (*Auth, error) { + return available[0], nil + }) + selector := NewSessionAffinitySelector(fallback) + + opts := optsWithMixedNamespace(cliproxyexecutor.Options{Headers: http.Header{"X-Session-Id": []string{"mixed-second-12345"}}}) + cacheKey := "mixed::header:mixed-second-12345::model" + + // Bind A under the canonical key. + selector.OnResult(Result{AuthID: authA.ID, Provider: "gemini", Model: "model", Success: true, Options: opts}) + + // A deterministic-unavailable (only B in list) -> Pick skips A, gets B; B fails -> cache empty. + picked, _ := selector.Pick(context.Background(), "mixed", "model", opts, []*Auth{authB}) + if picked.ID != authB.ID { + t.Fatalf("Pick = %q, want B", picked.ID) + } + selector.OnResult(Result{AuthID: picked.ID, Provider: "gemini", Model: "model", Success: false, Error: &Error{HTTPStatus: http.StatusServiceUnavailable}, Options: opts}) + if _, ok := selector.cache.Get(cacheKey); ok { + t.Fatalf("cache should be empty after B failure") + } + + // Immediate second request with both available reselects from fallback (A), not stale B. + second, _ := selector.Pick(context.Background(), "mixed", "model", opts, []*Auth{authA, authB}) + if second.ID != authA.ID { + t.Fatalf("second request should reselect from fallback, got %q", second.ID) + } +} +func optsWithAffinityNamespaces(opts cliproxyexecutor.Options, provider, model string) cliproxyexecutor.Options { + if opts.Metadata == nil { + opts.Metadata = make(map[string]any) + } + opts.Metadata[cliproxyexecutor.SessionAffinityProviderMetadataKey] = provider + opts.Metadata[cliproxyexecutor.SessionAffinityModelMetadataKey] = model + return opts +} + +func TestSessionAffinity_ModelNamespace_MixedRouteModelBindsCanonicalKey(t *testing.T) { + gemini := &Auth{ID: "gemini-auth", Provider: "gemini"} + fallback := pickFuncSelector(func(_ context.Context, _, _ string, _ cliproxyexecutor.Options, available []*Auth) (*Auth, error) { + return available[0], nil + }) + selector := NewSessionAffinitySelector(fallback) + + opts := cliproxyexecutor.Options{Headers: http.Header{"X-Session-Id": []string{"model-mixed-12345"}}, Metadata: map[string]any{}} + + // Pick under the route model records the model namespace before any upstream rewrite. + picked, err := selector.Pick(context.Background(), "mixed", ".gemini-flash", opts, []*Auth{gemini}) + if err != nil || picked == nil || picked.ID != gemini.ID { + t.Fatalf("Pick = %v/%v, want gemini", picked, err) + } + if ns, _ := opts.Metadata[cliproxyexecutor.SessionAffinityModelMetadataKey].(string); ns != ".gemini-flash" { + t.Fatalf("model namespace = %q, want .gemini-flash", ns) + } + + // Success result carries a rewritten upstream model; OnResult must key by route model. + selector.OnResult(Result{AuthID: picked.ID, Provider: "gemini", Model: "gemini-3.5-flash-lite", Success: true, Options: opts}) + + routeKey := "mixed::header:model-mixed-12345::.gemini-flash" + rewrittenKey := "mixed::header:model-mixed-12345::gemini-3.5-flash-lite" + if _, ok := selector.cache.Get(routeKey); !ok { + t.Fatalf("expected binding under route-model canonical key") + } + if _, ok := selector.cache.Get(rewrittenKey); ok { + t.Fatalf("must NOT bind under rewritten upstream model key") + } + + // Next Pick for the same route model hits the bound auth. + next, _ := selector.Pick(context.Background(), "mixed", ".gemini-flash", opts, []*Auth{gemini}) + if next.ID != gemini.ID { + t.Fatalf("second Pick = %q, want gemini", next.ID) + } +} + +func TestSessionAffinity_ModelNamespace_SingleProviderAliasRewrite(t *testing.T) { + auth := &Auth{ID: "single-auth", Provider: "claude"} + fallback := pickFuncSelector(func(_ context.Context, _, _ string, _ cliproxyexecutor.Options, available []*Auth) (*Auth, error) { + return available[0], nil + }) + selector := NewSessionAffinitySelector(fallback) + + opts := cliproxyexecutor.Options{Headers: http.Header{"X-Session-Id": []string{"model-single-12345"}}, Metadata: map[string]any{}} + + picked, _ := selector.Pick(context.Background(), "claude", "op-4-mini", opts, []*Auth{auth}) + selector.OnResult(Result{AuthID: picked.ID, Provider: "claude", Model: "op-4-mini-rewritten", Success: true, Options: opts}) + + routeKey := "claude::header:model-single-12345::op-4-mini" + if _, ok := selector.cache.Get(routeKey); !ok { + t.Fatalf("expected binding under route-model alias key") + } + next, _ := selector.Pick(context.Background(), "claude", "op-4-mini", opts, []*Auth{auth}) + if next.ID != auth.ID { + t.Fatalf("second Pick = %q, want auth", next.ID) + } +} + +func TestSessionAffinity_ModelNamespace_FailureClearsRouteBinding(t *testing.T) { + auth := &Auth{ID: "auth-a", Provider: "gemini"} + fallback := pickFuncSelector(func(_ context.Context, _, _ string, _ cliproxyexecutor.Options, available []*Auth) (*Auth, error) { + return available[0], nil + }) + selector := NewSessionAffinitySelector(fallback) + + opts := optsWithAffinityNamespaces(cliproxyexecutor.Options{Headers: http.Header{"X-Session-Id": []string{"model-fail-12345"}}}, "mixed", ".gemini-flash") + routeKey := "mixed::header:model-fail-12345::.gemini-flash" + + // Bind under the route-model canonical key. + selector.OnResult(Result{AuthID: auth.ID, Provider: "gemini", Model: "gemini-3.5-flash-lite", Success: true, Options: opts}) + if _, ok := selector.cache.Get(routeKey); !ok { + t.Fatalf("precondition: route binding should exist") + } + + // Failure with a rewritten Result model clears the canonical route binding. + selector.OnResult(Result{AuthID: auth.ID, Provider: "gemini", Model: "gemini-3.5-flash-lite", Success: false, Error: &Error{HTTPStatus: http.StatusServiceUnavailable}, Options: opts}) + if bound, ok := selector.cache.Get(routeKey); ok { + t.Fatalf("route binding not cleared after failure; bound=%q", bound) + } +} + +func TestSessionAffinity_ModelNamespace_StaleFailureCannotDeleteNewerSuccess(t *testing.T) { + authA := &Auth{ID: "auth-a", Provider: "gemini"} + authB := &Auth{ID: "auth-b", Provider: "gemini"} + fallback := pickFuncSelector(func(_ context.Context, _, _ string, _ cliproxyexecutor.Options, available []*Auth) (*Auth, error) { + return available[0], nil + }) + selector := NewSessionAffinitySelector(fallback) + + opts := optsWithAffinityNamespaces(cliproxyexecutor.Options{Headers: http.Header{"X-Session-Id": []string{"model-stale-newer"}}}, "mixed", ".gemini-flash") + routeKey := "mixed::header:model-stale-newer::.gemini-flash" + + // Newer request succeeds on B under the route-model key. + selector.OnResult(Result{AuthID: authB.ID, Provider: "gemini", Model: "gemini-3.5-flash-lite", Success: true, Options: opts}) + // Older request fails on A; must not delete B's newer binding. + selector.OnResult(Result{AuthID: authA.ID, Provider: "gemini", Model: "gemini-3.5-flash-lite", Success: false, Error: &Error{HTTPStatus: http.StatusTooManyRequests}, Options: opts}) + + bound, ok := selector.cache.Get(routeKey) + if !ok || bound != authB.ID { + t.Fatalf("stale failure for %q deleted newer binding %q; bound=%q ok=%v", authA.ID, authB.ID, bound, ok) + } +} + +func TestSessionAffinity_ModelNamespace_StreamRewriteBindsRouteKey(t *testing.T) { + ctx := context.Background() + manager := NewManager(nil, nil, nil) + affinity := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{Fallback: &RoundRobinSelector{}, TTL: time.Hour}) + defer affinity.Stop() + manager.SetSelector(affinity) + + auth := &Auth{ID: "stream-model-auth", Provider: "gemini", Status: StatusActive} + if _, err := manager.Register(WithSkipPersist(ctx), auth); err != nil { + t.Fatalf("Register: %v", err) + } + + opts := optsWithAffinityNamespaces(cliproxyexecutor.Options{Headers: http.Header{"X-Session-Id": []string{"stream-model-12345"}}}, "mixed", ".gemini-flash") + routeKey := "mixed::header:stream-model-12345::.gemini-flash" + + // Stream succeeds with a rewritten upstream model; must bind the route-model key. + chunk := cliproxyexecutor.StreamChunk{Payload: []byte("data: {\"id\":\"x\"}\n\n")} + res := manager.wrapStreamResult(ctx, auth, "gemini", "gemini-3.5-flash-lite", opts, nil, []cliproxyexecutor.StreamChunk{chunk}, closedStreamChunks(), OAuthModelAliasResult{}, false) + for range res.Chunks { + } + bound, ok := affinity.cache.Get(routeKey) + if !ok || bound != auth.ID { + t.Fatalf("stream rewrite should bind route key; bound=%q ok=%v", bound, ok) + } + + // Stream failure with rewritten model clears the route key. + errChunk := cliproxyexecutor.StreamChunk{Err: &Error{HTTPStatus: http.StatusServiceUnavailable}} + res2 := manager.wrapStreamResult(ctx, auth, "gemini", "gemini-3.5-flash-lite", opts, nil, []cliproxyexecutor.StreamChunk{errChunk}, closedStreamChunks(), OAuthModelAliasResult{}, false) + for range res2.Chunks { + } + if bound, ok := affinity.cache.Get(routeKey); ok { + t.Fatalf("stream failure should clear route key; still bound=%q", bound) + } +} + +func TestSessionAffinity_ModelNamespace_MetadataAbsentUsesResultModel(t *testing.T) { + auth := &Auth{ID: "auth-a", Provider: "gemini"} + fallback := pickFuncSelector(func(_ context.Context, _, _ string, _ cliproxyexecutor.Options, available []*Auth) (*Auth, error) { + return available[0], nil + }) + selector := NewSessionAffinitySelector(fallback) + + // No namespace metadata: OnResult must fall back to Result.Provider/Result.Model. + opts := cliproxyexecutor.Options{Headers: http.Header{"X-Session-Id": []string{"model-compat-12345"}}} + selector.OnResult(Result{AuthID: auth.ID, Provider: "gemini", Model: "gemini-model", Success: true, Options: opts}) + + cacheKey := "gemini::header:model-compat-12345::gemini-model" + bound, ok := selector.cache.Get(cacheKey) + if !ok || bound != auth.ID { + t.Fatalf("metadata-absent should key by Result.Provider/Model; bound=%q ok=%v", bound, ok) + } +} + +func TestSessionAffinity_QuarantinesRetryAfterForSameSessionOnly(t *testing.T) { + authA := &Auth{ID: "auth-a", Provider: "gemini"} + authB := &Auth{ID: "auth-b", Provider: "gemini"} + fallback := pickFuncSelector(func(_ context.Context, _, _ string, _ cliproxyexecutor.Options, available []*Auth) (*Auth, error) { + return available[0], nil + }) + selector := NewSessionAffinitySelector(fallback) + defer selector.Stop() + + opts := optsWithAffinityNamespaces(cliproxyexecutor.Options{Headers: http.Header{"X-Session-Id": []string{"cooldown-session-one"}}}, "mixed", ".gemini-flash") + first, err := selector.Pick(context.Background(), "mixed", ".gemini-flash", opts, []*Auth{authA, authB}) + if err != nil || first.ID != authA.ID { + t.Fatalf("first Pick = %v/%v, want auth-a", first, err) + } + + retryAfter := 53 * time.Second + selector.OnResult(Result{ + AuthID: authA.ID, + Provider: "gemini", + Model: "gemini-3.6-flash", + Success: false, + Error: &Error{HTTPStatus: http.StatusTooManyRequests}, + RetryAfter: &retryAfter, + Options: opts, + }) + + second, err := selector.Pick(context.Background(), "mixed", ".gemini-flash", opts, []*Auth{authA, authB}) + if err != nil || second.ID != authB.ID { + t.Fatalf("same-session retry Pick = %v/%v, want auth-b", second, err) + } + + otherOpts := optsWithAffinityNamespaces(cliproxyexecutor.Options{Headers: http.Header{"X-Session-Id": []string{"cooldown-session-two"}}}, "mixed", ".gemini-flash") + other, err := selector.Pick(context.Background(), "mixed", ".gemini-flash", otherOpts, []*Auth{authA, authB}) + if err != nil || other.ID != authA.ID { + t.Fatalf("other-session Pick = %v/%v, want auth-a", other, err) + } +} + +func TestSessionAffinity_QuarantinesMultipleFailedAuths(t *testing.T) { + authA := &Auth{ID: "auth-a", Provider: "gemini"} + authB := &Auth{ID: "auth-b", Provider: "gemini"} + authC := &Auth{ID: "auth-c", Provider: "gemini"} + fallback := pickFuncSelector(func(_ context.Context, _, _ string, _ cliproxyexecutor.Options, available []*Auth) (*Auth, error) { + return available[0], nil + }) + selector := NewSessionAffinitySelector(fallback) + defer selector.Stop() + + opts := optsWithAffinityNamespaces(cliproxyexecutor.Options{Headers: http.Header{"X-Session-Id": []string{"cooldown-multiple"}}}, "mixed", ".gemini-flash") + retryAfter := 53 * time.Second + for _, auth := range []*Auth{authA, authB} { + picked, err := selector.Pick(context.Background(), "mixed", ".gemini-flash", opts, []*Auth{authA, authB, authC}) + if err != nil || picked.ID != auth.ID { + t.Fatalf("Pick before failing %s = %v/%v", auth.ID, picked, err) + } + selector.OnResult(Result{ + AuthID: auth.ID, + Provider: "gemini", + Model: "gemini-3.6-flash", + Success: false, + Error: &Error{HTTPStatus: http.StatusTooManyRequests}, + RetryAfter: &retryAfter, + Options: opts, + }) + } + + third, err := selector.Pick(context.Background(), "mixed", ".gemini-flash", opts, []*Auth{authA, authB, authC}) + if err != nil || third.ID != authC.ID { + t.Fatalf("third Pick = %v/%v, want auth-c", third, err) + } +} + +func TestSessionAffinity_QuarantineExpires(t *testing.T) { + authA := &Auth{ID: "auth-a", Provider: "gemini"} + authB := &Auth{ID: "auth-b", Provider: "gemini"} + fallback := pickFuncSelector(func(_ context.Context, _, _ string, _ cliproxyexecutor.Options, available []*Auth) (*Auth, error) { + return available[0], nil + }) + selector := NewSessionAffinitySelector(fallback) + defer selector.Stop() + + opts := optsWithAffinityNamespaces(cliproxyexecutor.Options{Headers: http.Header{"X-Session-Id": []string{"cooldown-expiry"}}}, "mixed", ".gemini-flash") + retryAfter := 20 * time.Millisecond + selector.OnResult(Result{ + AuthID: authA.ID, + Provider: "gemini", + Model: "gemini-3.6-flash", + Success: false, + Error: &Error{HTTPStatus: http.StatusTooManyRequests}, + RetryAfter: &retryAfter, + Options: opts, + }) + + before, err := selector.Pick(context.Background(), "mixed", ".gemini-flash", opts, []*Auth{authA, authB}) + if err != nil || before.ID != authB.ID { + t.Fatalf("Pick before expiry = %v/%v, want auth-b", before, err) + } + selector.OnResult(Result{AuthID: authB.ID, Provider: "gemini", Model: "gemini-3.6-flash", Success: false, Error: &Error{HTTPStatus: http.StatusBadGateway}, Options: opts}) + time.Sleep(30 * time.Millisecond) + after, err := selector.Pick(context.Background(), "mixed", ".gemini-flash", opts, []*Auth{authA, authB}) + if err != nil || after.ID != authA.ID { + t.Fatalf("Pick after expiry = %v/%v, want auth-a", after, err) + } +} + +func TestSessionAffinity_StaleSuccessDoesNotClearNewerQuarantine(t *testing.T) { + authA := &Auth{ID: "auth-a", Provider: "gemini"} + authB := &Auth{ID: "auth-b", Provider: "gemini"} + fallback := pickFuncSelector(func(_ context.Context, _, _ string, _ cliproxyexecutor.Options, available []*Auth) (*Auth, error) { + return available[0], nil + }) + selector := NewSessionAffinitySelector(fallback) + defer selector.Stop() + + opts := optsWithAffinityNamespaces(cliproxyexecutor.Options{Headers: http.Header{"X-Session-Id": []string{"cooldown-stale-success"}}}, "mixed", ".gemini-flash") + retryAfter := 53 * time.Second + selector.OnResult(Result{ + AuthID: authA.ID, + Provider: "gemini", + Model: "gemini-3.6-flash", + Success: false, + Error: &Error{HTTPStatus: http.StatusTooManyRequests}, + RetryAfter: &retryAfter, + Options: opts, + }) + selector.OnResult(Result{ + AuthID: authA.ID, + Provider: "gemini", + Model: "gemini-3.6-flash", + Success: true, + Options: opts, + }) + + got, err := selector.Pick(context.Background(), "mixed", ".gemini-flash", opts, []*Auth{authA, authB}) + if err != nil || got.ID != authB.ID { + t.Fatalf("Pick after stale success = %v/%v, want auth-b while auth-a remains quarantined", got, err) + } +} + +func TestSessionAffinity_RequestScoped400DoesNotQuarantine(t *testing.T) { + authA := &Auth{ID: "auth-a", Provider: "gemini"} + authB := &Auth{ID: "auth-b", Provider: "gemini"} + fallback := pickFuncSelector(func(_ context.Context, _, _ string, _ cliproxyexecutor.Options, available []*Auth) (*Auth, error) { + return available[0], nil + }) + selector := NewSessionAffinitySelector(fallback) + defer selector.Stop() + + opts := optsWithAffinityNamespaces(cliproxyexecutor.Options{Headers: http.Header{"X-Session-Id": []string{"cooldown-client-error"}}}, "mixed", ".gemini-flash") + selector.OnResult(Result{ + AuthID: authA.ID, + Provider: "gemini", + Model: "gemini-3.6-flash", + Success: false, + Error: &Error{Code: requestScopedErrorCode, HTTPStatus: http.StatusBadRequest}, + Options: opts, + }) + + got, err := selector.Pick(context.Background(), "mixed", ".gemini-flash", opts, []*Auth{authA, authB}) + if err != nil || got.ID != authA.ID { + t.Fatalf("Pick after request-scoped 400 = %v/%v, want auth-a", got, err) + } +} diff --git a/sdk/cliproxy/auth/session_affinity_priority_test.go b/sdk/cliproxy/auth/session_affinity_priority_test.go index adb1c67bf..7426cf270 100644 --- a/sdk/cliproxy/auth/session_affinity_priority_test.go +++ b/sdk/cliproxy/auth/session_affinity_priority_test.go @@ -75,22 +75,42 @@ func TestManagerSessionAffinityPreservesBindingAcrossHigherPriorityRecovery(t *t return auth } - if got := pick(opts); got.ID != highID { - t.Fatalf("cold binding = %q, want high priority %q", got.ID, highID) + highAuth := pick(opts) + if highAuth.ID != highID { + t.Fatalf("cold binding = %q, want high priority %q", highAuth.ID, highID) } + manager.MarkResult(ctx, Result{ + AuthID: highAuth.ID, + Provider: highAuth.Provider, + Model: model, + Success: true, + Options: opts, + }) manager.MarkResult(ctx, Result{ - AuthID: highID, - Provider: provider, + AuthID: highAuth.ID, + Provider: highAuth.Provider, Model: model, Success: false, Error: &Error{HTTPStatus: http.StatusTooManyRequests, Message: "quota"}, + Options: opts, }) - if got := pick(opts); got.ID != lowID { - t.Fatalf("failover binding = %q, want %q", got.ID, lowID) + lowAuth := pick(opts) + if lowAuth.ID != lowID { + t.Fatalf("failover binding = %q, want %q", lowAuth.ID, lowID) } + manager.MarkResult(ctx, Result{ + AuthID: lowAuth.ID, + Provider: lowAuth.Provider, + Model: model, + Success: true, + Options: opts, + }) expireSessionAffinityPriorityModelCooldown(t, manager, highID, model) + // The affinity namespace fix makes the mixed selection path bind and read + // under the canonical pool key, so the lowID binding is retained across + // higher-priority recovery in both the single- and mixed-provider subtests. if got := pick(opts); got.ID != lowID { t.Fatalf("binding after higher-priority recovery = %q, want sticky %q", got.ID, lowID) } @@ -103,14 +123,16 @@ func TestManagerSessionAffinityPreservesBindingAcrossHigherPriorityRecovery(t *t } manager.MarkResult(ctx, Result{ - AuthID: lowID, - Provider: provider, + AuthID: lowAuth.ID, + Provider: lowAuth.Provider, Model: model, Success: false, Error: &Error{HTTPStatus: http.StatusTooManyRequests, Message: "quota"}, + Options: opts, }) - if got := pick(opts); got.ID != highID { - t.Fatalf("binding after bound auth became unavailable = %q, want %q", got.ID, highID) + got, errPick := testCase.pick(manager, ctx, provider, model, opts) + if errPick == nil || got != nil { + t.Fatalf("binding after all session auths failed = %v/%v, want no candidate until quarantine expires", got, errPick) } }) } diff --git a/sdk/cliproxy/auth/session_cache.go b/sdk/cliproxy/auth/session_cache.go index 54bf3867b..0dae249ab 100644 --- a/sdk/cliproxy/auth/session_cache.go +++ b/sdk/cliproxy/auth/session_cache.go @@ -10,17 +10,19 @@ const maxStableSessionAliases = 64 // sessionEntry stores an auth binding, its identifier aliases, and expiration. type sessionEntry struct { - authID string - expiresAt time.Time - aliases []string + authID string + expiresAt time.Time + aliases []string + generation uint64 } // SessionCache provides TTL-based session to auth mapping with automatic cleanup. type SessionCache struct { - mu sync.RWMutex - entries map[string]sessionEntry - ttl time.Duration - stopCh chan struct{} + mu sync.RWMutex + entries map[string]sessionEntry + ttl time.Duration + stopCh chan struct{} + generation uint64 } // NewSessionCache creates a cache with the specified TTL. @@ -92,6 +94,22 @@ func (c *SessionCache) GetAndRefresh(sessionID string) (string, bool) { return entry.authID, true } +// GetWithGeneration retrieves the auth ID, monotonic generation token, and alias list +// bound to a session without refreshing the TTL. +func (c *SessionCache) GetWithGeneration(sessionID string) (string, uint64, []string, bool) { + if sessionID == "" { + return "", 0, nil, false + } + now := time.Now() + c.mu.RLock() + defer c.mu.RUnlock() + entry, ok := c.entries[sessionID] + if !ok || !now.Before(entry.expiresAt) { + return "", 0, nil, false + } + return entry.authID, entry.generation, append([]string(nil), entry.aliases...), true +} + // Set binds a session to an auth ID with TTL refresh. Existing aliases for the // same logical session remain attached when the binding is refreshed or moved. func (c *SessionCache) Set(sessionID, authID string) { @@ -100,10 +118,54 @@ func (c *SessionCache) Set(sessionID, authID string) { // SetAliases binds multiple identifiers for one logical session to an auth ID. func (c *SessionCache) SetAliases(authID string, sessionIDs ...string) { - if authID == "" { + c.setAliasesUntil(authID, time.Now().Add(c.ttl), sessionIDs...) +} + +// RestoreAliasesIfAbsent atomically sets the still-absent aliases to authID. +// Any alias that is currently live (bound to another active group) is left untouched. +// Returns true if at least one alias was restored, false otherwise. +func (c *SessionCache) RestoreAliasesIfAbsent(authID string, sessionIDs ...string) bool { + if c == nil || authID == "" || len(sessionIDs) == 0 { + return false + } + now := time.Now() + c.mu.Lock() + defer c.mu.Unlock() + + var absent []string + for _, sid := range sessionIDs { + if sid == "" { + continue + } + if entry, ok := c.entries[sid]; !ok || !now.Before(entry.expiresAt) { + absent = append(absent, sid) + } + } + aliases := compactSessionAliases(absent) + if len(aliases) == 0 { + return false + } + c.generation++ + entry := sessionEntry{ + authID: authID, + expiresAt: now.Add(c.ttl), + aliases: aliases, + generation: c.generation, + } + for _, alias := range aliases { + c.entries[alias] = entry + } + return true +} + +func (c *SessionCache) setAliasesUntil(authID string, expiresAt time.Time, sessionIDs ...string) { + if authID == "" || expiresAt.IsZero() { return } now := time.Now() + if !now.Before(expiresAt) { + return + } c.mu.Lock() defer c.mu.Unlock() @@ -125,14 +187,16 @@ func (c *SessionCache) SetAliases(authID string, sessionIDs ...string) { if len(aliases) == 0 { return } - c.replaceAliasGroupsLocked(authID, now.Add(c.ttl), aliases, previousGroups...) + c.replaceAliasGroupsLocked(authID, expiresAt, aliases, previousGroups...) } func (c *SessionCache) replaceAliasGroupsLocked(authID string, expiresAt time.Time, aliases []string, previousGroups ...sessionEntry) { + c.generation++ + gen := c.generation for _, previous := range previousGroups { c.removeAliasGroupLocked(previous) } - entry := sessionEntry{authID: authID, expiresAt: expiresAt, aliases: aliases} + entry := sessionEntry{authID: authID, expiresAt: expiresAt, aliases: aliases, generation: gen} for _, alias := range aliases { c.entries[alias] = entry } @@ -229,28 +293,152 @@ func (c *SessionCache) Invalidate(sessionID string) { return } c.mu.Lock() + defer c.mu.Unlock() entry, ok := c.entries[sessionID] + if !ok { + return + } delete(c.entries, sessionID) - if ok { - for _, alias := range entry.aliases { - if alias == sessionID { - continue - } - current, exists := c.entries[alias] - if !exists || current.authID != entry.authID { - continue - } - filtered := make([]string, 0, len(current.aliases)) - for _, candidate := range current.aliases { - if candidate != sessionID { - filtered = append(filtered, candidate) - } + c.generation++ + gen := c.generation + for _, alias := range entry.aliases { + if alias == sessionID { + continue + } + current, exists := c.entries[alias] + if !exists || current.authID != entry.authID { + continue + } + filtered := make([]string, 0, len(current.aliases)) + for _, candidate := range current.aliases { + if candidate != sessionID { + filtered = append(filtered, candidate) } - current.aliases = filtered - c.entries[alias] = current + } + current.aliases = filtered + current.generation = gen + c.entries[alias] = current + } +} + +// CompareAndDeleteAliases removes the alias group holding sessionID when it is +// still bound to expectedAuthID, and returns the group's aliases. A stale +// expectation cannot remove a newer group. +func (c *SessionCache) CompareAndDeleteAliases(sessionID, expectedAuthID string) []string { + if c == nil || sessionID == "" || expectedAuthID == "" { + return nil + } + c.mu.Lock() + defer c.mu.Unlock() + + entry, ok := c.entries[sessionID] + if !ok || entry.authID != expectedAuthID { + return nil + } + aliases := append([]string(nil), entry.aliases...) + c.generation++ + for _, alias := range aliases { + if current, exists := c.entries[alias]; exists && current.authID == expectedAuthID && equalSessionAliases(current.aliases, entry.aliases) { + delete(c.entries, alias) + } + } + return aliases +} + +// CompareAndDeleteGroup removes a binding only when its auth ID, generation, +// and alias set all still match the observed values, and returns the removed +// aliases. A concurrent refresh or extension of the group bumps the +// generation or changes the aliases, so a stale observation cannot delete +// newer state; callers retry their merge on a nil result. +// +// Mirror of CLIProxyAPI dd8c72a3. +func (c *SessionCache) CompareAndDeleteGroup(sessionID, expectedAuthID string, expectedGen uint64, expectedAliases []string) []string { + if c == nil || sessionID == "" || expectedAuthID == "" { + return nil + } + c.mu.Lock() + defer c.mu.Unlock() + + entry, ok := c.entries[sessionID] + if !ok || entry.authID != expectedAuthID || entry.generation != expectedGen { + return nil + } + if !equalSessionAliases(compactSessionAliases(entry.aliases), compactSessionAliases(expectedAliases)) { + return nil + } + removed := append([]string(nil), entry.aliases...) + c.generation++ + for _, alias := range removed { + if current, exists := c.entries[alias]; exists && current.authID == expectedAuthID && equalSessionAliases(current.aliases, entry.aliases) { + delete(c.entries, alias) } } - c.mu.Unlock() + return removed +} + +// CompareAndReplaceAliases atomically validates that every observed alias still +// maps to expectedAuthID with expectedGen, that all observed aliases belong to the +// exact same group, and that no additional alias is currently bound to +// another active group. Upon validation, it replaces the entire alias group with +// newAuthID, a refreshed TTL, and an incremented monotonic generation token. +// It returns true if replaced, or false if the CAS precondition failed. +func (c *SessionCache) CompareAndReplaceAliases( + expectedAuthID string, + expectedGen uint64, + observedAliases []string, + newAuthID string, + additionalAliases ...string, +) bool { + if c == nil || expectedAuthID == "" || expectedGen == 0 || len(observedAliases) == 0 || newAuthID == "" { + return false + } + now := time.Now() + c.mu.Lock() + defer c.mu.Unlock() + + var matchedEntry *sessionEntry + for _, alias := range observedAliases { + entry, exists := c.entries[alias] + if !exists || !now.Before(entry.expiresAt) { + return false + } + if entry.authID != expectedAuthID || entry.generation != expectedGen { + return false + } + if !equalSessionAliases(entry.aliases, observedAliases) { + return false + } + if matchedEntry == nil { + entryCopy := entry + matchedEntry = &entryCopy + } + } + if matchedEntry == nil || len(matchedEntry.aliases) != len(observedAliases) { + return false + } + + allAliases := mergeSessionAliases(observedAliases, additionalAliases...) + allAliases = compactSessionAliases(allAliases) + if len(allAliases) == 0 { + return false + } + + observedSet := make(map[string]struct{}, len(observedAliases)) + for _, a := range observedAliases { + observedSet[a] = struct{}{} + } + + for _, alias := range allAliases { + if _, isObserved := observedSet[alias]; isObserved { + continue + } + if existing, exists := c.entries[alias]; exists && now.Before(existing.expiresAt) { + return false + } + } + + c.replaceAliasGroupsLocked(newAuthID, now.Add(c.ttl), allAliases, *matchedEntry) + return true } // InvalidateAuth removes all sessions bound to a specific auth ID. @@ -260,12 +448,12 @@ func (c *SessionCache) InvalidateAuth(authID string) { return } c.mu.Lock() + defer c.mu.Unlock() for sid, entry := range c.entries { if entry.authID == authID { delete(c.entries, sid) } } - c.mu.Unlock() } // Stop terminates the background cleanup goroutine. @@ -293,10 +481,10 @@ func (c *SessionCache) cleanupLoop() { func (c *SessionCache) cleanup() { now := time.Now() c.mu.Lock() + defer c.mu.Unlock() for sid, entry := range c.entries { if !now.Before(entry.expiresAt) { delete(c.entries, sid) } } - c.mu.Unlock() } diff --git a/sdk/cliproxy/auth/stream_ttft_test.go b/sdk/cliproxy/auth/stream_ttft_test.go new file mode 100644 index 000000000..f82cd8705 --- /dev/null +++ b/sdk/cliproxy/auth/stream_ttft_test.go @@ -0,0 +1,728 @@ +package auth + +import ( + "context" + "net/http" + "strings" + "sync" + "testing" + "time" + + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +type ttftTestExecutor struct { + id string + + mu sync.Mutex + streamCalls []string + delayAuthA time.Duration + delayAuthB time.Duration +} + +func (e *ttftTestExecutor) Identifier() string { return e.id } + +func (e *ttftTestExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusNotImplemented, Message: "not implemented"} +} + +func (e *ttftTestExecutor) ExecuteStream(ctx context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + e.mu.Lock() + e.streamCalls = append(e.streamCalls, auth.ID) + delayA := e.delayAuthA + delayB := e.delayAuthB + e.mu.Unlock() + + if auth.ID == "auth-a" && delayA > 0 { + select { + case <-time.After(delayA): + case <-ctx.Done(): + return nil, ctx.Err() + } + } + + if auth.ID == "auth-b" && delayB > 0 { + select { + case <-time.After(delayB): + case <-ctx.Done(): + return nil, ctx.Err() + } + } + + ch := make(chan cliproxyexecutor.StreamChunk, 2) + ch <- cliproxyexecutor.StreamChunk{Payload: []byte(`data: {"type":"content_block_delta","delta":{"text":"chunk-from-` + auth.ID + `"}}` + "\n\n")} + close(ch) + return &cliproxyexecutor.StreamResult{Headers: http.Header{"X-Auth": {auth.ID}}, Chunks: ch}, nil +} + +func (e *ttftTestExecutor) Refresh(context.Context, *Auth) (*Auth, error) { + return nil, nil +} + +func (e *ttftTestExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} + +func (e *ttftTestExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +func (e *ttftTestExecutor) StreamCalls() []string { + e.mu.Lock() + defer e.mu.Unlock() + out := make([]string, len(e.streamCalls)) + copy(out, e.streamCalls) + return out +} + +func TestManagerExecuteStream_TTFTTimeoutFailsOverToNextAuth(t *testing.T) { + model := "gpt-5.5" + authA := &Auth{ID: "auth-a", Provider: "codex"} + authB := &Auth{ID: "auth-b", Provider: "codex"} + + executor := &ttftTestExecutor{ + id: "codex", + delayAuthA: 500 * time.Millisecond, + } + + m := NewManager(nil, nil, nil) + m.RegisterExecutor(executor) + + reg := registry.GetGlobalRegistry() + reg.RegisterClient(authA.ID, "codex", []*registry.ModelInfo{{ID: model}}) + reg.RegisterClient(authB.ID, "codex", []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { + reg.UnregisterClient(authA.ID) + reg.UnregisterClient(authB.ID) + }) + + if _, err := m.Register(context.Background(), authA); err != nil { + t.Fatalf("register authA: %v", err) + } + if _, err := m.Register(context.Background(), authB); err != nil { + t.Fatalf("register authB: %v", err) + } + + start := time.Now() + opts := cliproxyexecutor.Options{ + Metadata: map[string]any{ + "stream_first_chunk_timeout_ms": 50, + }, + } + + stream, errStream := m.ExecuteStream(context.Background(), []string{"codex"}, cliproxyexecutor.Request{Model: model}, opts) + elapsed := time.Since(start) + + if errStream != nil { + t.Fatalf("ExecuteStream error = %v, expected failover to authB", errStream) + } + if stream == nil || stream.Chunks == nil { + t.Fatalf("expected stream result from authB") + } + + var chunks []cliproxyexecutor.StreamChunk + for chunk := range stream.Chunks { + chunks = append(chunks, chunk) + } + + if len(chunks) == 0 { + t.Fatalf("expected chunks from authB") + } + if got := string(chunks[0].Payload); !strings.Contains(got, "chunk-from-auth-b") { + t.Fatalf("chunk payload = %q, expected bytes ONLY from authB", got) + } + + calls := executor.StreamCalls() + if len(calls) != 2 || calls[0] != "auth-a" || calls[1] != "auth-b" { + t.Fatalf("executor stream calls = %v, expected [auth-a, auth-b] called once each", calls) + } + + if elapsed > 400*time.Millisecond { + t.Fatalf("elapsed time = %v, expected under 400ms", elapsed) + } + + updatedA, ok := m.GetByID("auth-a") + if !ok || updatedA == nil { + t.Fatalf("auth-a missing from manager") + } + if !updatedA.Unavailable { + t.Fatalf("expected auth-a to be marked unavailable after TTFT timeout") + } +} + +func TestManagerExecuteStream_PostFirstChunkDelayNotCutOffByTTFT(t *testing.T) { + model := "gpt-5.5" + authA := &Auth{ID: "auth-a", Provider: "codex"} + + delayedStreamExec := &postFirstChunkExecutor{id: "codex", authID: "auth-a"} + + m := NewManager(nil, nil, nil) + m.RegisterExecutor(delayedStreamExec) + + reg := registry.GetGlobalRegistry() + reg.RegisterClient(authA.ID, "codex", []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { + reg.UnregisterClient(authA.ID) + }) + + if _, err := m.Register(context.Background(), authA); err != nil { + t.Fatalf("register authA: %v", err) + } + + opts := cliproxyexecutor.Options{ + Metadata: map[string]any{ + "stream_first_chunk_timeout_ms": 40, + }, + } + + stream, errStream := m.ExecuteStream(context.Background(), []string{"codex"}, cliproxyexecutor.Request{Model: model}, opts) + if errStream != nil { + t.Fatalf("ExecuteStream error = %v, want success", errStream) + } + + var chunks []cliproxyexecutor.StreamChunk + for chunk := range stream.Chunks { + if chunk.Err != nil { + t.Fatalf("unexpected chunk error: %v", chunk.Err) + } + chunks = append(chunks, chunk) + } + + if len(chunks) != 2 { + t.Fatalf("got %d chunks, want 2 chunks (both first and delayed second chunk)", len(chunks)) + } +} + +type postFirstChunkExecutor struct { + id string + authID string +} + +func (e *postFirstChunkExecutor) Identifier() string { return e.id } +func (e *postFirstChunkExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} +func (e *postFirstChunkExecutor) Refresh(context.Context, *Auth) (*Auth, error) { return nil, nil } +func (e *postFirstChunkExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} +func (e *postFirstChunkExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +func (e *postFirstChunkExecutor) ExecuteStream(ctx context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + ch := make(chan cliproxyexecutor.StreamChunk) + go func() { + defer close(ch) + ch <- cliproxyexecutor.StreamChunk{Payload: []byte(`data: {"type":"content_block_delta","delta":{"text":"chunk1"}}` + "\n\n")} + select { + case <-time.After(120 * time.Millisecond): + case <-ctx.Done(): + return + } + ch <- cliproxyexecutor.StreamChunk{Payload: []byte(`data: {"type":"content_block_delta","delta":{"text":"chunk2"}}` + "\n\n")} + }() + return &cliproxyexecutor.StreamResult{Chunks: ch}, nil +} + +type metadataFirstPostCommitExecutor struct { + id string + authID string + + mu sync.Mutex + streamCalls []string + metadataWait time.Duration + postCommitErr bool + postCommitDelay time.Duration + // provideContent controls whether a final erroring stream also emits a + // content chunk first, exercising the true "post-commit" path. + emitContentBeforeErr bool +} + +func (e *metadataFirstPostCommitExecutor) Identifier() string { return e.id } + +func (e *metadataFirstPostCommitExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} + +func (e *metadataFirstPostCommitExecutor) Refresh(context.Context, *Auth) (*Auth, error) { + return nil, nil +} + +func (e *metadataFirstPostCommitExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} + +func (e *metadataFirstPostCommitExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +func (e *metadataFirstPostCommitExecutor) ExecuteStream(ctx context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + e.mu.Lock() + e.streamCalls = append(e.streamCalls, auth.ID) + wait := e.metadataWait + postErr := e.postCommitErr + postDelay := e.postCommitDelay + emitContent := e.emitContentBeforeErr + e.mu.Unlock() + + ch := make(chan cliproxyexecutor.StreamChunk, 4) + go func() { + defer close(ch) + // Recognized metadata/comment first: liveness but no semantic content. + ch <- cliproxyexecutor.StreamChunk{Payload: []byte(": keepalive\n\n")} + if wait > 0 { + select { + case <-time.After(wait): + case <-ctx.Done(): + return + } + } + if emitContent { + ch <- cliproxyexecutor.StreamChunk{Payload: []byte(`data: {"type":"content_block_delta","delta":{"text":"chunk-from-` + auth.ID + `"}}` + "\n\n")} + if postDelay > 0 { + select { + case <-time.After(postDelay): + case <-ctx.Done(): + return + } + } + } + if postErr { + ch <- cliproxyexecutor.StreamChunk{Err: &Error{HTTPStatus: http.StatusBadGateway, Message: "upstream failed after first chunk"}} + } + }() + return &cliproxyexecutor.StreamResult{Chunks: ch}, nil +} + +func (e *metadataFirstPostCommitExecutor) StreamCalls() []string { + e.mu.Lock() + defer e.mu.Unlock() + out := make([]string, len(e.streamCalls)) + copy(out, e.streamCalls) + return out +} + +func TestManagerExecuteStream_MetadataFirstPrefixStopsTTFTWithoutFailover(t *testing.T) { + model := "gpt-5.5" + authA := &Auth{ID: "auth-a", Provider: "codex"} + authB := &Auth{ID: "auth-b", Provider: "codex"} + + executor := &metadataFirstPostCommitExecutor{id: "codex", authID: "auth-a", metadataWait: 150 * time.Millisecond, emitContentBeforeErr: true} + + m := NewManager(nil, nil, nil) + m.RegisterExecutor(executor) + + reg := registry.GetGlobalRegistry() + reg.RegisterClient(authA.ID, "codex", []*registry.ModelInfo{{ID: model}}) + reg.RegisterClient(authB.ID, "codex", []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { + reg.UnregisterClient(authA.ID) + reg.UnregisterClient(authB.ID) + }) + + if _, err := m.Register(context.Background(), authA); err != nil { + t.Fatalf("register authA: %v", err) + } + if _, err := m.Register(context.Background(), authB); err != nil { + t.Fatalf("register authB: %v", err) + } + + opts := cliproxyexecutor.Options{ + Metadata: map[string]any{ + "stream_first_chunk_timeout_ms": 40, + }, + } + + stream, errStream := m.ExecuteStream(context.Background(), []string{"codex"}, cliproxyexecutor.Request{Model: model}, opts) + if errStream != nil { + t.Fatalf("ExecuteStream error = %v, want success from auth-a (metadata kept TTFT alive)", errStream) + } + if stream == nil || stream.Chunks == nil { + t.Fatal("expected stream result") + } + + var chunks []cliproxyexecutor.StreamChunk + for chunk := range stream.Chunks { + if chunk.Err != nil { + t.Fatalf("unexpected chunk error: %v", chunk.Err) + } + chunks = append(chunks, chunk) + } + + if len(chunks) != 2 { + t.Fatalf("got %d chunks, want 2 (buffered metadata + content)", len(chunks)) + } + if got := string(chunks[0].Payload); got != ": keepalive\n\n" { + t.Fatalf("first chunk payload = %q, want buffered metadata kept first", got) + } + if got := string(chunks[1].Payload); !strings.Contains(got, "chunk-from-auth-a") { + t.Fatalf("second chunk payload = %q, want content from auth-a", got) + } + + calls := executor.StreamCalls() + if len(calls) != 1 || calls[0] != "auth-a" { + t.Fatalf("executor stream calls = %v, want [auth-a] only (no failover to auth-b)", calls) + } + + updatedA, ok := m.GetByID("auth-a") + if !ok || updatedA == nil { + t.Fatal("auth-a missing from manager") + } + if updatedA.Unavailable { + t.Fatal("expected auth-a to remain available after metadata-first stream") + } +} + +func TestManagerExecuteStream_PostCommitErrorNotRetried(t *testing.T) { + model := "gpt-5.5" + authA := &Auth{ID: "auth-a", Provider: "codex"} + authB := &Auth{ID: "auth-b", Provider: "codex"} + + executor := &metadataFirstPostCommitExecutor{ + id: "codex", + authID: "auth-a", + metadataWait: 10 * time.Millisecond, + postCommitErr: true, + postCommitDelay: 10 * time.Millisecond, + emitContentBeforeErr: true, + } + + m := NewManager(nil, nil, nil) + m.RegisterExecutor(executor) + + reg := registry.GetGlobalRegistry() + reg.RegisterClient(authA.ID, "codex", []*registry.ModelInfo{{ID: model}}) + reg.RegisterClient(authB.ID, "codex", []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { + reg.UnregisterClient(authA.ID) + reg.UnregisterClient(authB.ID) + }) + + if _, err := m.Register(context.Background(), authA); err != nil { + t.Fatalf("register authA: %v", err) + } + if _, err := m.Register(context.Background(), authB); err != nil { + t.Fatalf("register authB: %v", err) + } + + opts := cliproxyexecutor.Options{ + Metadata: map[string]any{ + "stream_first_chunk_timeout_ms": 40, + }, + } + + stream, errStream := m.ExecuteStream(context.Background(), []string{"codex"}, cliproxyexecutor.Request{Model: model}, opts) + if errStream != nil { + t.Fatalf("ExecuteStream error = %v, want success (stream handed to consumer)", errStream) + } + + var chunks []cliproxyexecutor.StreamChunk + var terminalErr error + seenContent := false + for chunk := range stream.Chunks { + if chunk.Err != nil { + if terminalErr != nil { + t.Fatalf("duplicate terminal chunk error: %v after %v", chunk.Err, terminalErr) + } + terminalErr = chunk.Err + continue + } + if len(chunk.Payload) > 0 && strings.Contains(string(chunk.Payload), "chunk-from-auth-a") { + seenContent = true + chunks = append(chunks, chunk) + } + } + + if terminalErr == nil { + t.Fatal("expected post-commit error to reach the result stream") + } + if !seenContent { + t.Fatal("expected content chunk from auth-a before the terminal error") + } + // No duplicate content replay from a backend retry. + if len(chunks) != 1 { + t.Fatalf("got %d content chunks, want exactly 1 (no backend replay)", len(chunks)) + } + + calls := executor.StreamCalls() + if len(calls) != 1 || calls[0] != "auth-a" { + t.Fatalf("executor stream calls = %v, want [auth-a] only (post-commit error must not retry)", calls) + } +} + +func TestStreamFirstChunkTimeout_ConfigAndMetadata(t *testing.T) { + m := NewManager(nil, nil, nil) + + if got := m.streamFirstChunkTimeout(cliproxyexecutor.Options{}); got != 0 { + t.Fatalf("default streamFirstChunkTimeout = %v, want disabled", got) + } + + cfg := &internalconfig.Config{ + SDKConfig: internalconfig.SDKConfig{ + Streaming: internalconfig.StreamingConfig{ + StreamFirstChunkTimeoutSeconds: 10, + }, + }, + } + m.runtimeConfig.Store(cfg) + if got := m.streamFirstChunkTimeout(cliproxyexecutor.Options{}); got != 10*time.Second { + t.Fatalf("custom streamFirstChunkTimeout = %v, want 10s", got) + } + + cfgDisabled := &internalconfig.Config{ + SDKConfig: internalconfig.SDKConfig{ + Streaming: internalconfig.StreamingConfig{ + StreamFirstChunkTimeoutSeconds: -1, + }, + }, + } + m.runtimeConfig.Store(cfgDisabled) + if got := m.streamFirstChunkTimeout(cliproxyexecutor.Options{}); got != 0 { + t.Fatalf("disabled streamFirstChunkTimeout = %v, want 0", got) + } + + opts := cliproxyexecutor.Options{ + Metadata: map[string]any{ + "stream_first_chunk_timeout_ms": 75, + }, + } + if got := m.streamFirstChunkTimeout(opts); got != 75*time.Millisecond { + t.Fatalf("metadata streamFirstChunkTimeout = %v, want 75ms", got) + } +} + +type ttftRefreshExecutor struct { + id string + + mu sync.Mutex + streamCalls int + refreshCalls int + tokenInvalid map[string]struct{} + refreshTokens map[string]string + delayedStreams []string // auth IDs whose retry stream must delay its first chunk + staleDelay time.Duration + retryDelay time.Duration +} + +func (e *ttftRefreshExecutor) Identifier() string { return e.id } + +func (e *ttftRefreshExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusNotImplemented, Message: "not implemented"} +} + +func (e *ttftRefreshExecutor) ExecuteStream(ctx context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + e.mu.Lock() + e.streamCalls++ + call := e.streamCalls + token := authAccessToken(auth) + _, invalid := e.tokenInvalid[token] + delayed := false + for _, id := range e.delayedStreams { + if id == auth.ID { + delayed = true + break + } + } + staleDelay := e.staleDelay + retryDelay := e.retryDelay + e.mu.Unlock() + + if invalid { + // The stale-token attempt consumes TTFT budget before failing, so a + // shared (non-fresh) timer would be near-firing by the time the retry + // starts its own (possibly long) first-chunk wait. + if call == 1 && staleDelay > 0 { + select { + case <-time.After(staleDelay): + case <-ctx.Done(): + return nil, ctx.Err() + } + } + return nil, &Error{ + HTTPStatus: http.StatusUnauthorized, + Message: "Your authentication token has been invalidated. Please try signing in again.", + } + } + + if delayed && call > 1 && retryDelay > 0 { + select { + case <-time.After(retryDelay): + case <-ctx.Done(): + return nil, ctx.Err() + } + } + ch := make(chan cliproxyexecutor.StreamChunk, 1) + ch <- cliproxyexecutor.StreamChunk{Payload: []byte(auth.ID + ":" + token)} + close(ch) + return &cliproxyexecutor.StreamResult{Chunks: ch}, nil +} + +func (e *ttftRefreshExecutor) Refresh(ctx context.Context, auth *Auth) (*Auth, error) { + e.mu.Lock() + defer e.mu.Unlock() + e.refreshCalls++ + if auth.Metadata == nil { + auth.Metadata = make(map[string]any) + } + next := e.refreshTokens[auth.ID] + if next == "" { + next = "refreshed-access-token" + } + auth.Metadata["access_token"] = next + return auth, nil +} + +func (e *ttftRefreshExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} + +func (e *ttftRefreshExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +func (e *ttftRefreshExecutor) StreamCalls() int { + e.mu.Lock() + defer e.mu.Unlock() + return e.streamCalls +} + +func (e *ttftRefreshExecutor) RefreshCalls() int { + e.mu.Lock() + defer e.mu.Unlock() + return e.refreshCalls +} + +func newTTFTRefreshFixture(t *testing.T, staleDelay, retryDelay time.Duration) (*Manager, *ttftRefreshExecutor, *Auth, string) { + t.Helper() + + model := "gpt-5.5" + primary := &Auth{ + ID: "ttft-primary", + Provider: "codex", + Metadata: map[string]any{ + "access_token": "ttft-stale-token", + "refresh_token": "ttft-refresh-token", + }, + } + + executor := &ttftRefreshExecutor{ + id: "codex", + tokenInvalid: map[string]struct{}{ + "ttft-stale-token": {}, + }, + staleDelay: staleDelay, + retryDelay: retryDelay, + delayedStreams: []string{primary.ID}, + refreshTokens: map[string]string{ + primary.ID: "ttft-fresh-token", + }, + } + + m := NewManager(nil, nil, nil) + m.RegisterExecutor(executor) + + reg := registry.GetGlobalRegistry() + reg.RegisterClient(primary.ID, "codex", []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { + reg.UnregisterClient(primary.ID) + }) + + if _, errRegister := m.Register(context.Background(), primary); errRegister != nil { + t.Fatalf("register primary: %v", errRegister) + } + + return m, executor, primary, model +} + +func TestManagerExecuteStream_RefreshRetryGetsFreshTTFTTimer(t *testing.T) { + // First attempt consumes 90ms of a 120ms budget before failing with the + // stale-token 401, then triggers a refresh. The retry must receive a fresh + // TTFT budget: its first chunk arrives 70ms into the retry, well inside a + // fresh 120ms window. A shared (non-fresh) timer would have already elapsed + // 90ms and fired at ~50ms into the retry, cutting the 70ms chunk off. + m, executor, primary, model := newTTFTRefreshFixture(t, 90*time.Millisecond, 70*time.Millisecond) + + opts := cliproxyexecutor.Options{ + Metadata: map[string]any{ + "stream_first_chunk_timeout_ms": 120, + }, + } + + stream, errStream := m.ExecuteStream(context.Background(), []string{"codex"}, cliproxyexecutor.Request{Model: model}, opts) + if errStream != nil { + t.Fatalf("ExecuteStream error = %v, want success on refreshed retry", errStream) + } + if stream == nil || stream.Chunks == nil { + t.Fatal("expected stream result") + } + for chunk := range stream.Chunks { + if chunk.Err != nil { + t.Fatalf("stream chunk error = %v", chunk.Err) + } + if got := string(chunk.Payload); got != primary.ID+":ttft-fresh-token" { + t.Fatalf("payload = %q, want refreshed primary response (refreshCalls=%d streamCalls=%d)", got, executor.RefreshCalls(), executor.StreamCalls()) + } + } + + if got := executor.RefreshCalls(); got != 1 { + t.Fatalf("Refresh calls = %d, want 1", got) + } + if got := executor.StreamCalls(); got != 2 { + t.Fatalf("Stream calls = %d, want 2 (initial + refreshed retry)", got) + } +} + +func TestStreamConnectTimeout_ConfigAndMetadata(t *testing.T) { + m := NewManager(nil, nil, nil) + + // Canonical config key + cfg := &internalconfig.Config{ + SDKConfig: internalconfig.SDKConfig{ + Streaming: internalconfig.StreamingConfig{ + StreamConnectTimeoutSeconds: 15, + }, + }, + } + m.runtimeConfig.Store(cfg) + if got := m.streamFirstChunkTimeout(cliproxyexecutor.Options{}); got != 15*time.Second { + t.Fatalf("canonical StreamConnectTimeoutSeconds = %v, want 15s", got) + } + + // Precedence: canonical config overrides legacy alias + cfgBoth := &internalconfig.Config{ + SDKConfig: internalconfig.SDKConfig{ + Streaming: internalconfig.StreamingConfig{ + StreamConnectTimeoutSeconds: 20, + StreamFirstChunkTimeoutSeconds: 5, + }, + }, + } + m.runtimeConfig.Store(cfgBoth) + if got := m.streamFirstChunkTimeout(cliproxyexecutor.Options{}); got != 20*time.Second { + t.Fatalf("StreamConnectTimeoutSeconds precedence = %v, want 20s", got) + } + + // Canonical metadata key + opts := cliproxyexecutor.Options{ + Metadata: map[string]any{ + "stream_connect_timeout_ms": 60, + }, + } + if got := m.streamFirstChunkTimeout(opts); got != 60*time.Millisecond { + t.Fatalf("canonical stream_connect_timeout_ms = %v, want 60ms", got) + } + + // Precedence: canonical metadata overrides legacy metadata alias + optsBoth := cliproxyexecutor.Options{ + Metadata: map[string]any{ + "stream_connect_timeout_ms": 90, + "stream_first_chunk_timeout_ms": 30, + }, + } + if got := m.streamFirstChunkTimeout(optsBoth); got != 90*time.Millisecond { + t.Fatalf("stream_connect_timeout_ms metadata precedence = %v, want 90ms", got) + } +} diff --git a/sdk/cliproxy/executor/types.go b/sdk/cliproxy/executor/types.go index 85e5fdc23..b59573f4a 100644 --- a/sdk/cliproxy/executor/types.go +++ b/sdk/cliproxy/executor/types.go @@ -48,6 +48,19 @@ const ( DerivedSessionIDMetadataKey = "derived_session_id" // CallerScopeMetadataKey isolates inferred session identities between downstream callers. CallerScopeMetadataKey = "caller_scope" + // ExcludedAuthIDsMetadataKey carries the set of auth IDs that already failed + // (429/5xx/empty) within the current request and must never be re-selected + // for the remainder of that request. Value is map[string]struct{} or []string. + ExcludedAuthIDsMetadataKey = "request_excluded_auth_ids" + // SessionAffinityProviderMetadataKey carries the affinity selection namespace + // (provider string, e.g. the literal "mixed" pool key) used by SessionAffinitySelector.Pick, + // so OnResult keys the session cache identically to how selection read it. + SessionAffinityProviderMetadataKey = "session_affinity_provider" + // SessionAffinityModelMetadataKey carries the normalized model argument used by + // SessionAffinitySelector.Pick to build the session cache key, before any + // executor/model-pool/home upstream rewrite, so OnResult keys the session cache + // identically to how selection read it. + SessionAffinityModelMetadataKey = "session_affinity_model" ) // Request encapsulates the translated payload that will be sent to a provider executor. @@ -108,6 +121,9 @@ type RequestTerminatedError struct { HTTPStatus int Header http.Header Body []byte + // Trusted reports that the termination originated locally (plugin/interceptor) + // rather than from an untrusted upstream. Zero value false is the safe default. + Trusted bool } func (e *RequestTerminatedError) Error() string {