From 2166ca48aa60fd4c419ddd0b05297e78a7b7e72c Mon Sep 17 00:00:00 2001 From: warelik Date: Sat, 22 Aug 2026 21:56:58 +0300 Subject: [PATCH 1/9] fix(auth): treat empty upstream completions as retriable failures --- internal/clienterror/client_error.go | 9 +- internal/clienterror/client_error_test.go | 17 + internal/config/sdk_config.go | 7 + internal/interfaces/error_message.go | 4 + internal/pluginhost/executor_route.go | 164 +- .../executor_route_close_order_test.go | 44 + .../executor_route_stream_codex_test.go | 167 + .../pluginhost/executor_route_stream_test.go | 337 ++ internal/pluginhost/model_router_test.go | 82 + internal/registry/model_registry.go | 87 +- .../model_registry_resume_reason_test.go | 85 + .../claude_executor_ratelimit_test.go | 2 +- .../executor/home_codex_terminal_test.go | 17 +- sdk/api/handlers/gemini/gemini_handlers.go | 12 + .../gemini/gemini_handlers_error_test.go | 36 + .../handlers/handlers_error_response_test.go | 21 + sdk/api/handlers/handlers_errors.go | 2 +- sdk/api/handlers/handlers_execution.go | 11 +- sdk/api/handlers/handlers_interceptors.go | 9 +- .../handlers/handlers_interceptors_test.go | 6 +- .../handlers_stream_bootstrap_test.go | 8 +- sdk/api/handlers/model_execution_test.go | 13 +- .../auth/conductor_availability_test.go | 384 ++ sdk/cliproxy/auth/conductor_cooldown.go | 167 +- .../conductor_cooldown_retry_reset_test.go | 316 + sdk/cliproxy/auth/conductor_execution.go | 213 +- .../auth/conductor_fast_error_test.go | 4 +- .../auth/conductor_force_mapping_test.go | 7 + sdk/cliproxy/auth/conductor_home.go | 6 + sdk/cliproxy/auth/conductor_home_execution.go | 44 +- .../conductor_invalid_key_sibling_test.go | 305 + sdk/cliproxy/auth/conductor_overrides_test.go | 369 +- sdk/cliproxy/auth/conductor_selection.go | 21 +- sdk/cliproxy/auth/conductor_stream.go | 244 +- .../auth/conductor_stream_drain_test.go | 159 + .../auth/conductor_stream_eof_test.go | 36 + .../auth/conductor_stream_ttft_test.go | 474 ++ .../conductor_unauthorized_refresh_test.go | 68 +- sdk/cliproxy/auth/empty_completion.go | 2584 ++++++++ sdk/cliproxy/auth/empty_completion_export.go | 158 + .../auth/empty_completion_formats_test.go | 96 + sdk/cliproxy/auth/empty_completion_test.go | 5178 +++++++++++++++++ sdk/cliproxy/auth/export_test.go | 23 + sdk/cliproxy/auth/home_concurrency.go | 7 + .../auth/home_execution_paths_test.go | 128 +- sdk/cliproxy/auth/home_retry_contract_test.go | 24 +- .../auth/home_selected_auth_callback_test.go | 5 +- .../auth/outer_retry_exclusions_test.go | 237 + sdk/cliproxy/auth/route_exhaustion_test.go | 779 +++ sdk/cliproxy/auth/route_tracker.go | 210 + .../selected_auth_failover_metadata_test.go | 275 + .../auth/selected_auth_metadata_test.go | 211 + sdk/cliproxy/auth/selector.go | 351 +- sdk/cliproxy/auth/selector_review_p2_test.go | 304 + sdk/cliproxy/auth/selector_test.go | 750 +++ .../auth/session_affinity_priority_test.go | 42 +- .../auth/session_affinity_quarantine_test.go | 136 + sdk/cliproxy/auth/session_cache.go | 205 +- sdk/cliproxy/auth/stream_ttft_test.go | 349 ++ sdk/cliproxy/auth/types_test.go | 9 +- sdk/cliproxy/executor/types.go | 12 +- 61 files changed, 15819 insertions(+), 211 deletions(-) create mode 100644 internal/pluginhost/executor_route_close_order_test.go create mode 100644 internal/pluginhost/executor_route_stream_codex_test.go create mode 100644 internal/pluginhost/executor_route_stream_test.go create mode 100644 internal/registry/model_registry_resume_reason_test.go create mode 100644 sdk/api/handlers/gemini/gemini_handlers_error_test.go create mode 100644 sdk/cliproxy/auth/conductor_cooldown_retry_reset_test.go create mode 100644 sdk/cliproxy/auth/conductor_invalid_key_sibling_test.go create mode 100644 sdk/cliproxy/auth/conductor_stream_drain_test.go create mode 100644 sdk/cliproxy/auth/conductor_stream_eof_test.go create mode 100644 sdk/cliproxy/auth/conductor_stream_ttft_test.go create mode 100644 sdk/cliproxy/auth/empty_completion.go create mode 100644 sdk/cliproxy/auth/empty_completion_export.go create mode 100644 sdk/cliproxy/auth/empty_completion_formats_test.go create mode 100644 sdk/cliproxy/auth/empty_completion_test.go create mode 100644 sdk/cliproxy/auth/export_test.go create mode 100644 sdk/cliproxy/auth/outer_retry_exclusions_test.go create mode 100644 sdk/cliproxy/auth/route_exhaustion_test.go create mode 100644 sdk/cliproxy/auth/route_tracker.go create mode 100644 sdk/cliproxy/auth/selected_auth_failover_metadata_test.go create mode 100644 sdk/cliproxy/auth/selector_review_p2_test.go create mode 100644 sdk/cliproxy/auth/session_affinity_quarantine_test.go create mode 100644 sdk/cliproxy/auth/stream_ttft_test.go diff --git a/internal/clienterror/client_error.go b/internal/clienterror/client_error.go index 51db164a32d..58d3451bf63 100644 --- a/internal/clienterror/client_error.go +++ b/internal/clienterror/client_error.go @@ -87,9 +87,12 @@ func IsRequestFault(status int, err error) bool { return false } // DeepSeek reports an invalid API key as 401 with the authentication_error - // type alongside the same generic code. Preserve that credential failure - // classification without weakening generic request-fault handling. - if status == http.StatusUnauthorized && hasAuthenticationErrorBody(err) { + // type alongside the same generic code. Other providers also surface an + // invalid or expired credential as the authentication_error body on 403. + // Preserve that credential failure classification without weakening generic + // request-fault handling: a request-fault-looking code on the same body does + // not turn a credential rejection into a request fault. + if (status == http.StatusUnauthorized || status == http.StatusForbidden) && hasAuthenticationErrorBody(err) { return false } if hasRequestFaultBody(err) { diff --git a/internal/clienterror/client_error_test.go b/internal/clienterror/client_error_test.go index 758efda55a4..489971535ea 100644 --- a/internal/clienterror/client_error_test.go +++ b/internal/clienterror/client_error_test.go @@ -197,6 +197,23 @@ func TestIsRequestFault(t *testing.T) { err: errors.New(`{"error":{"code":"invalid_request_error","message":"Authentication Fails, Your api key: ****heck is invalid","param":null,"type":"authentication_error"}}`), want: false, }, + { + // Codex surfaces an invalid or expired API key as 403 with the + // authentication_error body. The credential must stay eligible for + // rotation, so this is not a request fault. + name: "codex invalid or expired key on 403 is credential failure", + status: http.StatusForbidden, + err: errors.New(`{"error":{"message":"invalid or expired token","type":"authentication_error","code":"invalid_api_key"}}`), + want: false, + }, + { + // A credential rejection must not be reclassified as a request fault + // merely because the same body carries a request-fault-looking code. + name: "authentication body with generic code on 403 is credential failure", + status: http.StatusForbidden, + err: errors.New(`{"error":{"code":"invalid_request_error","message":"Authentication Fails, Your api key: ****heck is invalid","param":null,"type":"authentication_error"}}`), + want: false, + }, { name: "deepseek insufficient balance is payment failure", status: http.StatusPaymentRequired, diff --git a/internal/config/sdk_config.go b/internal/config/sdk_config.go index c7a53ffb307..24aafa751ea 100644 --- a/internal/config/sdk_config.go +++ b/internal/config/sdk_config.go @@ -79,4 +79,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. + // Zero or a negative value disables the timeout. + 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/interfaces/error_message.go b/internal/interfaces/error_message.go index 93fa3acbee2..f1720041240 100644 --- a/internal/interfaces/error_message.go +++ b/internal/interfaces/error_message.go @@ -21,6 +21,10 @@ 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 from a + // trusted local interceptor rather than an untrusted upstream error. + 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 be6138db82b..e7fabc0973b 100644 --- a/internal/pluginhost/executor_route.go +++ b/internal/pluginhost/executor_route.go @@ -88,7 +88,14 @@ func (h *Host) ExecutePluginExecutor(ctx context.Context, pluginID string, req c if errAdapter != nil { return coreexecutor.Response{}, errAdapter } - return adapter.Execute(ctx, (*coreauth.Auth)(nil), req, opts) + resp, err := adapter.Execute(ctx, (*coreauth.Auth)(nil), req, opts) + if err != nil { + return coreexecutor.Response{}, err + } + if coreauth.IsEmptyCompletionPayload(resp.Payload) { + return coreexecutor.Response{}, coreauth.EmptyCompletionError() + } + return resp, nil } // ExecutePluginExecutorStream executes a streaming request with the named plugin executor without changing the requested model. @@ -97,7 +104,160 @@ func (h *Host) ExecutePluginExecutorStream(ctx context.Context, pluginID string, if errAdapter != nil { return nil, errAdapter } - return adapter.ExecuteStream(ctx, (*coreauth.Auth)(nil), req, opts) + streamResult, err := adapter.ExecuteStream(ctx, (*coreauth.Auth)(nil), req, opts) + if err != nil { + return nil, err + } + return wrapStreamEmptyCompletion(ctx, streamResult, req.Payload, opts.OriginalRequest), nil +} + +// wrapStreamEmptyCompletion wraps a plugin stream so that a terminal but empty +// completion (no content, no tool calls) surfaces as an empty-completion error +// instead of a clean stream end, mirroring the conductor's aggregate-at-close +// judgment. Recognized protocol framing is buffered only until meaningful output +// appears or the stream closes; unrecognized streams remain pass-through. +func wrapStreamEmptyCompletion(ctx context.Context, streamResult *coreexecutor.StreamResult, requestPayloads ...[]byte) *coreexecutor.StreamResult { + if streamResult == nil || streamResult.Chunks == nil { + return streamResult + } + if ctx == nil { + ctx = context.Background() + } + src := streamResult.Chunks + wrapped := make(chan coreexecutor.StreamChunk) + go func() { + defer close(wrapped) + buffered := make([]coreexecutor.StreamChunk, 0, 1) + var detector coreauth.StreamBootstrapDetector + for _, p := range requestPayloads { + if n := coreauth.ExtractExpectedChoices(p); n > 1 { + detector.SetExpectedChoices(n) + break + } + } + forwarding := false + forward := func(chunk coreexecutor.StreamChunk) bool { + select { + case <-ctx.Done(): + return false + case wrapped <- chunk: + return true + } + } + flush := func() bool { + for _, chunk := range buffered { + if !forward(chunk) { + return false + } + } + buffered = nil + return true + } + + for { + var ( + chunk coreexecutor.StreamChunk + ok bool + ) + select { + case <-ctx.Done(): + return + case chunk, ok = <-src: + } + if !ok { + 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. + // Finish() parses the trailing fragment, so a provider error that only + // lands in that final unterminated frame is not known until after it + // runs: consult StreamError() before reporting terminal emptiness. + terminalEmpty := detector.Finish() + if streamErr := detector.StreamError(); streamErr != nil { + _ = forward(coreexecutor.StreamChunk{Err: streamErr}) + return + } + if terminalEmpty { + _ = forward(coreexecutor.StreamChunk{Err: coreauth.EmptyCompletionError()}) + return + } + } + _ = flush() + return + } + if forwarding { + if !forward(chunk) { + return + } + if chunk.Err != nil { + return + } + continue + } + + buffered = append(buffered, chunk) + if chunk.Err != nil { + // Before any semantic output, protocol framing is not client-visible. + // Surface the upstream failure first so the HTTP layer can still + // choose an error response instead of committing a successful stream. + buffered = buffered[:0] + forwarding = true + if !forward(chunk) { + return + } + return + } + if detector.Observe(chunk.Payload) { + forwarding = true + if !flush() { + return + } + } + if streamErr := detector.StreamError(); streamErr != nil { + discardStreamChunks(src) + _ = forward(coreexecutor.StreamChunk{Err: streamErr}) + 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 { + payload = append(payload, chunk.Payload...) + } + return payload } // CountPluginExecutor executes a count-tokens request with the named plugin executor without changing the requested model. diff --git a/internal/pluginhost/executor_route_close_order_test.go b/internal/pluginhost/executor_route_close_order_test.go new file mode 100644 index 00000000000..10b5757ab0e --- /dev/null +++ b/internal/pluginhost/executor_route_close_order_test.go @@ -0,0 +1,44 @@ +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" +) + +// TestWrapStreamEmptyCompletion_PrefersDetectedErrorOverTerminalEmptiness covers a +// stream whose recognized frames carry no content and whose real provider error +// arrives as an SSE error event that is newline-terminated but never followed by the +// blank separator line. flushData() only runs on that blank line or from Finish(), +// so the detected provider error does not exist yet while the stream is being +// observed; judging emptiness before consulting it would replace a routable +// invalid_api_key with a generic empty_completion. +func TestWrapStreamEmptyCompletion_PrefersDetectedErrorOverTerminalEmptiness(t *testing.T) { + chunks := make(chan coreexecutor.StreamChunk, 3) + chunks <- coreexecutor.StreamChunk{Payload: []byte("data: {\"choices\":[{\"delta\":{\"role\":\"assistant\"}}]}\n\n")} + chunks <- coreexecutor.StreamChunk{Payload: []byte("event: error\ndata: {\"error\":{\"code\":\"invalid_api_key\",\"message\":\"invalid api key\"}}\n")} + close(chunks) + + res := wrapStreamEmptyCompletion(context.Background(), &coreexecutor.StreamResult{Chunks: chunks}) + var received []coreexecutor.StreamChunk + for c := range res.Chunks { + received = append(received, c) + } + + if len(received) != 1 { + t.Fatalf("expected 1 chunk carrying the detected provider error, got %d", len(received)) + } + if received[0].Err == nil { + t.Fatalf("expected an error chunk, got payload: %s", string(received[0].Payload)) + } + var authErr *coreauth.Error + if !errors.As(received[0].Err, &authErr) { + t.Fatalf("expected *coreauth.Error, got %v", received[0].Err) + } + if authErr.Code != "invalid_api_key" { + t.Fatalf("expected the provider error to survive terminal emptiness, got code %q (%v)", authErr.Code, received[0].Err) + } +} 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 00000000000..d6844161a5b --- /dev/null +++ b/internal/pluginhost/executor_route_stream_codex_test.go @@ -0,0 +1,167 @@ +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 for codex pullrequestreview-4943660625 on PR #4881 +// (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) + } +} + +// TestWrapStreamEmptyCompletionMultiChoiceWithholdsTerminalUntilAllChoicesFinish +// verifies that when n=2 is requested, an early terminal chunk for choice 0 +// does not cause an immediate empty_completion error if choice 1 subsequently emits content. +func TestWrapStreamEmptyCompletionMultiChoiceForwardsContentWhenChoice1HasContent(t *testing.T) { + src := make(chan coreexecutor.StreamChunk, 3) + src <- coreexecutor.StreamChunk{Payload: []byte("data: {\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n")} + src <- coreexecutor.StreamChunk{Payload: []byte("data: {\"choices\":[{\"index\":1,\"delta\":{\"content\":\"hello\"}}]}\n\n")} + src <- coreexecutor.StreamChunk{Payload: []byte("data: [DONE]\n\n")} + close(src) + + reqPayload := []byte(`{"model":"gpt-4o","n":2,"messages":[{"role":"user","content":"hi"}]}`) + wrapped := wrapStreamEmptyCompletion(context.Background(), &coreexecutor.StreamResult{Chunks: src}, reqPayload) + + first, ok := <-wrapped.Chunks + if !ok { + t.Fatal("wrapped stream closed prematurely") + } + if first.Err != nil { + t.Fatalf("unexpected error on first chunk: %v", first.Err) + } + second, ok := <-wrapped.Chunks + if !ok { + t.Fatal("wrapped stream missing second chunk") + } + if string(second.Payload) != "data: {\"choices\":[{\"index\":1,\"delta\":{\"content\":\"hello\"}}]}\n\n" { + t.Fatalf("second payload = %q, want choice 1 content", second.Payload) + } +} + +// TestWrapStreamEmptyCompletionMultiChoiceDetectsEmptyWhenAllChoicesFinishEmpty +// verifies that when n=2 is requested and both choices finish empty, +// wrapStreamEmptyCompletion correctly detects empty completion. +func TestWrapStreamEmptyCompletionMultiChoiceDetectsEmptyWhenAllChoicesFinishEmpty(t *testing.T) { + src := make(chan coreexecutor.StreamChunk, 2) + src <- coreexecutor.StreamChunk{Payload: []byte("data: {\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n")} + src <- coreexecutor.StreamChunk{Payload: []byte("data: {\"choices\":[{\"index\":1,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n")} + close(src) + + reqPayload := []byte(`{"model":"gpt-4o","n":2,"messages":[{"role":"user","content":"hi"}]}`) + wrapped := wrapStreamEmptyCompletion(context.Background(), &coreexecutor.StreamResult{Chunks: src}, reqPayload) + + first, ok := <-wrapped.Chunks + if !ok { + t.Fatal("wrapped multi-choice empty 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) + } +} + +// TestWrapStreamEmptyCompletionGeminiMultiCandidateForwardsContentWhenCandidate1HasContent +// verifies that when generationConfig.candidateCount=2 is requested, an early STOP chunk for candidate 0 +// does not cause an immediate empty_completion error if candidate 1 subsequently emits content. +func TestWrapStreamEmptyCompletionGeminiMultiCandidateForwardsContentWhenCandidate1HasContent(t *testing.T) { + src := make(chan coreexecutor.StreamChunk, 2) + src <- coreexecutor.StreamChunk{Payload: []byte("data: {\"candidates\":[{\"index\":0,\"content\":{\"parts\":[]},\"finishReason\":\"STOP\"}]}\n\n")} + src <- coreexecutor.StreamChunk{Payload: []byte("data: {\"candidates\":[{\"index\":1,\"content\":{\"parts\":[{\"text\":\"gemini answer\"}]}}]}\n\n")} + close(src) + + reqPayload := []byte(`{"contents":[{"parts":[{"text":"hi"}]}],"generationConfig":{"candidateCount":2}}`) + wrapped := wrapStreamEmptyCompletion(context.Background(), &coreexecutor.StreamResult{Chunks: src}, reqPayload) + + first, ok := <-wrapped.Chunks + if !ok { + t.Fatal("wrapped stream closed prematurely") + } + if first.Err != nil { + t.Fatalf("unexpected error on first chunk: %v", first.Err) + } + second, ok := <-wrapped.Chunks + if !ok { + t.Fatal("wrapped stream missing second chunk") + } + if string(second.Payload) != "data: {\"candidates\":[{\"index\":1,\"content\":{\"parts\":[{\"text\":\"gemini answer\"}]}}]}\n\n" { + t.Fatalf("second payload = %q, want candidate 1 content", second.Payload) + } +} + +// TestWrapStreamEmptyCompletionGeminiMultiCandidateDetectsEmptyWhenAllCandidatesFinishEmpty +// verifies that when candidateCount=2 is requested and both candidates finish empty, +// wrapStreamEmptyCompletion correctly detects empty completion. +func TestWrapStreamEmptyCompletionGeminiMultiCandidateDetectsEmptyWhenAllCandidatesFinishEmpty(t *testing.T) { + src := make(chan coreexecutor.StreamChunk, 2) + src <- coreexecutor.StreamChunk{Payload: []byte("data: {\"candidates\":[{\"index\":0,\"content\":{\"parts\":[]},\"finishReason\":\"STOP\"}]}\n\n")} + src <- coreexecutor.StreamChunk{Payload: []byte("data: {\"candidates\":[{\"index\":1,\"content\":{\"parts\":[]},\"finishReason\":\"STOP\"}]}\n\n")} + close(src) + + reqPayload := []byte(`{"contents":[{"parts":[{"text":"hi"}]}],"generationConfig":{"candidateCount":2}}`) + wrapped := wrapStreamEmptyCompletion(context.Background(), &coreexecutor.StreamResult{Chunks: src}, reqPayload) + + first, ok := <-wrapped.Chunks + if !ok { + t.Fatal("wrapped multi-candidate empty 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) + } +} diff --git a/internal/pluginhost/executor_route_stream_test.go b/internal/pluginhost/executor_route_stream_test.go new file mode 100644 index 00000000000..20d5bf0ba56 --- /dev/null +++ b/internal/pluginhost/executor_route_stream_test.go @@ -0,0 +1,337 @@ +package pluginhost + +import ( + "context" + "errors" + "net/http" + "testing" + "time" + + coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +func TestWrapStreamEmptyCompletionWithholdsTerminalFrames(t *testing.T) { + src := make(chan coreexecutor.StreamChunk, 2) + 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) + + wrapped := wrapStreamEmptyCompletion(context.Background(), &coreexecutor.StreamResult{Chunks: src}) + first, ok := <-wrapped.Chunks + if !ok { + t.Fatal("wrapped stream closed without empty_completion error") + } + if len(first.Payload) != 0 { + t.Fatalf("first payload = %q, want no terminal bytes before error", first.Payload) + } + 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 _, ok = <-wrapped.Chunks; ok { + t.Fatal("wrapped stream emitted chunks after empty_completion error") + } +} + +func TestWrapStreamEmptyCompletionRejectsZeroChunkStream(t *testing.T) { + src := make(chan coreexecutor.StreamChunk) + 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") + } +} + +func TestWrapStreamEmptyCompletionWithholdsSplitTerminalFrames(t *testing.T) { + fragments := [][]byte{ + []byte("da"), + []byte("ta: {\"choices\":[{\"delta\":{},\"finish_rea"), + []byte("son\":\"stop\"}],\"usage\":{\"completion_tokens\":0}}\n\n"), + []byte("data: [DO"), + []byte("NE]\n"), + []byte("\n"), + } + src := make(chan coreexecutor.StreamChunk, len(fragments)) + for _, fragment := range fragments { + src <- coreexecutor.StreamChunk{Payload: fragment} + } + close(src) + + wrapped := wrapStreamEmptyCompletion(context.Background(), &coreexecutor.StreamResult{Chunks: src}) + first, ok := <-wrapped.Chunks + if !ok { + t.Fatal("wrapped split stream closed without empty_completion error") + } + if len(first.Payload) != 0 { + t.Fatalf("first payload = %q, want no split terminal bytes before error", first.Payload) + } + 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 _, ok = <-wrapped.Chunks; ok { + t.Fatal("wrapped stream emitted chunks after empty_completion error") + } +} + +func TestWrapStreamEmptyCompletionForwardsSplitMeaningfulOutputInOrder(t *testing.T) { + fragments := [][]byte{ + []byte("da"), + []byte("ta: {\"choices\":[{\"delta\":{\"content\":"), + []byte("\"hello\"},\"finish_reason\":null}]}\n"), + []byte("\n"), + } + src := make(chan coreexecutor.StreamChunk, len(fragments)) + for _, fragment := range fragments { + src <- coreexecutor.StreamChunk{Payload: fragment} + } + + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + wrapped := wrapStreamEmptyCompletion(ctx, &coreexecutor.StreamResult{Chunks: src}) + for _, fragment := range fragments { + assertStreamPayload(t, wrapped.Chunks, fragment) + } + close(src) + if _, ok := <-wrapped.Chunks; ok { + t.Fatal("wrapped stream emitted an unexpected trailing chunk") + } +} + +func TestWrapStreamEmptyCompletionForwardsMeaningfulOutputInOrder(t *testing.T) { + emptyFrame := []byte("data: {\"choices\":[{\"delta\":{},\"finish_reason\":null}]}\n\n") + contentFrame := []byte("data: {\"choices\":[{\"delta\":{\"content\":\"hello\"},\"finish_reason\":null}]}\n\n") + src := make(chan coreexecutor.StreamChunk, 2) + src <- coreexecutor.StreamChunk{Payload: emptyFrame} + src <- coreexecutor.StreamChunk{Payload: contentFrame} + + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + wrapped := wrapStreamEmptyCompletion(ctx, &coreexecutor.StreamResult{Chunks: src}) + assertStreamPayload(t, wrapped.Chunks, emptyFrame) + assertStreamPayload(t, wrapped.Chunks, contentFrame) + close(src) + if _, ok := <-wrapped.Chunks; ok { + t.Fatal("wrapped stream emitted an unexpected trailing chunk") + } +} + +func TestWrapStreamEmptyCompletionForwardsUnrecognizedStreamPromptly(t *testing.T) { + firstPayload := []byte("opaque: first\n") + secondPayload := []byte("opaque: second\n") + src := make(chan coreexecutor.StreamChunk, 1) + src <- coreexecutor.StreamChunk{Payload: firstPayload} + + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + wrapped := wrapStreamEmptyCompletion(ctx, &coreexecutor.StreamResult{Chunks: src}) + assertStreamPayload(t, wrapped.Chunks, firstPayload) + src <- coreexecutor.StreamChunk{Payload: secondPayload} + assertStreamPayload(t, wrapped.Chunks, secondPayload) + close(src) + if _, ok := <-wrapped.Chunks; ok { + t.Fatal("wrapped stream emitted an unexpected trailing chunk") + } +} + +func TestWrapStreamEmptyCompletionWithholdsMetadataBeforeUpstreamError(t *testing.T) { + upstreamErr := errors.New("upstream failed") + src := make(chan coreexecutor.StreamChunk, 2) + src <- coreexecutor.StreamChunk{Payload: []byte("data: {\"choices\":[{\"delta\":{},\"finish_reason\":null}]}\n\n")} + src <- coreexecutor.StreamChunk{Err: upstreamErr} + close(src) + + wrapped := wrapStreamEmptyCompletion(context.Background(), &coreexecutor.StreamResult{Chunks: src}) + first, ok := <-wrapped.Chunks + if !ok { + t.Fatal("wrapped stream closed without upstream error") + } + if len(first.Payload) != 0 { + t.Fatalf("first payload = %q, want no metadata before upstream error", first.Payload) + } + if !errors.Is(first.Err, upstreamErr) { + t.Fatalf("first error = %v, want %v", first.Err, upstreamErr) + } + if _, ok = <-wrapped.Chunks; ok { + t.Fatal("wrapped stream emitted chunks after upstream error") + } +} + +func TestWrapStreamEmptyCompletionPreservesContentBeforeUpstreamError(t *testing.T) { + metadata := []byte("data: {\"choices\":[{\"delta\":{},\"finish_reason\":null}]}\n\n") + content := []byte("data: {\"choices\":[{\"delta\":{\"content\":\"hello\"},\"finish_reason\":null}]}\n\n") + upstreamErr := errors.New("upstream failed") + src := make(chan coreexecutor.StreamChunk, 3) + src <- coreexecutor.StreamChunk{Payload: metadata} + src <- coreexecutor.StreamChunk{Payload: content} + src <- coreexecutor.StreamChunk{Err: upstreamErr} + close(src) + + wrapped := wrapStreamEmptyCompletion(context.Background(), &coreexecutor.StreamResult{Chunks: src}) + assertStreamPayload(t, wrapped.Chunks, metadata) + assertStreamPayload(t, wrapped.Chunks, content) + errorChunk, ok := <-wrapped.Chunks + if !ok { + t.Fatal("wrapped stream closed before upstream error") + } + if !errors.Is(errorChunk.Err, upstreamErr) { + t.Fatalf("error = %v, want %v", errorChunk.Err, upstreamErr) + } + if _, ok = <-wrapped.Chunks; ok { + t.Fatal("wrapped stream emitted chunks after upstream error") + } +} + +func TestWrapStreamEmptyCompletionPreservesNilResults(t *testing.T) { + if got := wrapStreamEmptyCompletion(context.Background(), nil); got != nil { + t.Fatalf("wrapStreamEmptyCompletion(nil) = %#v, want nil", got) + } + + result := &coreexecutor.StreamResult{Headers: http.Header{"X-Test": []string{"value"}}} + if got := wrapStreamEmptyCompletion(context.Background(), result); got != result { + t.Fatalf("wrapStreamEmptyCompletion(nil chunks) = %#v, want original result", got) + } +} + +func TestWrapStreamEmptyCompletionStopsWhenContextCanceled(t *testing.T) { + src := make(chan coreexecutor.StreamChunk, 1) + src <- coreexecutor.StreamChunk{Payload: []byte("data: {\"choices\":[{\"delta\":{},\"finish_reason\":null}]}\n\n")} + ctx, cancel := context.WithCancel(context.Background()) + wrapped := wrapStreamEmptyCompletion(ctx, &coreexecutor.StreamResult{Chunks: src}) + cancel() + + select { + case _, ok := <-wrapped.Chunks: + if ok { + t.Fatal("wrapped stream emitted a chunk after cancellation") + } + case <-time.After(time.Second): + t.Fatal("wrapped stream did not close after cancellation") + } +} + +func assertStreamPayload(t *testing.T, chunks <-chan coreexecutor.StreamChunk, want []byte) { + t.Helper() + select { + case chunk, ok := <-chunks: + if !ok { + t.Fatalf("stream closed before payload %q", want) + } + if chunk.Err != nil { + t.Fatalf("chunk error = %v, want payload %q", chunk.Err, want) + } + if string(chunk.Payload) != string(want) { + t.Fatalf("payload = %q, want %q", chunk.Payload, want) + } + case <-time.After(time.Second): + 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/pluginhost/model_router_test.go b/internal/pluginhost/model_router_test.go index eacb4cc3132..c523a4a45ca 100644 --- a/internal/pluginhost/model_router_test.go +++ b/internal/pluginhost/model_router_test.go @@ -4,6 +4,7 @@ import ( "context" "errors" "fmt" + "net/http" "testing" coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" @@ -162,6 +163,87 @@ func TestHostExecutePluginExecutorByPluginIDPreservesModel(t *testing.T) { } } +func TestHostExecutePluginExecutorRejectsEmptyCompletion(t *testing.T) { + executor := &fakeExecutor{ + identifier: "plugin-provider", + execute: func(ctx context.Context, req pluginapi.ExecutorRequest) (pluginapi.ExecutorResponse, error) { + return pluginapi.ExecutorResponse{Payload: []byte(`{"choices":[{"message":{"content":""},"finish_reason":"stop"}],"usage":{"completion_tokens":0}}`)}, nil + }, + } + host := newRouteModelHostWithRecords(capabilityRecord{ + id: "executor", + plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{ + Executor: executor, + ExecutorInputFormats: []string{"openai"}, + ExecutorOutputFormats: []string{"openai"}, + }}, + }) + + _, errExecute := host.ExecutePluginExecutor(context.Background(), "executor", coreexecutor.Request{Model: "client-model", Payload: []byte(`{"model":"client-model"}`)}, coreexecutor.Options{}) + if errExecute == nil { + t.Fatal("ExecutePluginExecutor() with empty completion = nil, want retriable error") + } + var authErr *coreauth.Error + if !errors.As(errExecute, &authErr) { + t.Fatalf("error = %v (%T), want *coreauth.Error", errExecute, errExecute) + } + if !authErr.Retryable || authErr.Code != "empty_completion" { + t.Fatalf("error = %+v, want retriable empty_completion", authErr) + } + if authErr.StatusCode() != http.StatusServiceUnavailable { + t.Fatalf("error status = %d, want %d", authErr.StatusCode(), http.StatusServiceUnavailable) + } +} + +func TestHostExecutePluginExecutorStreamRejectsEmptyCompletion(t *testing.T) { + streamChunks := make(chan pluginapi.ExecutorStreamChunk) + executor := &fakeExecutor{ + identifier: "plugin-provider", + executeStream: func(ctx context.Context, req pluginapi.ExecutorRequest) (pluginapi.ExecutorStreamResponse, error) { + return pluginapi.ExecutorStreamResponse{Chunks: streamChunks}, nil + }, + } + host := newRouteModelHostWithRecords(capabilityRecord{ + id: "executor", + plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{ + Executor: executor, + ExecutorInputFormats: []string{"openai"}, + ExecutorOutputFormats: []string{"openai"}, + }}, + }) + + streamResult, errStream := host.ExecutePluginExecutorStream(context.Background(), "executor", coreexecutor.Request{Model: "client-model", Payload: []byte(`{"model":"client-model"}`)}, coreexecutor.Options{}) + if errStream != nil { + t.Fatalf("ExecutePluginExecutorStream() unexpected error = %v", errStream) + } + + go func() { + streamChunks <- pluginapi.ExecutorStreamChunk{Payload: []byte("data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n")} + streamChunks <- pluginapi.ExecutorStreamChunk{Payload: []byte("data: [DONE]\n\n")} + close(streamChunks) + }() + + var aggregated []byte + var emptyErr error + for chunk := range streamResult.Chunks { + if chunk.Err != nil { + emptyErr = chunk.Err + break + } + aggregated = append(aggregated, chunk.Payload...) + } + if emptyErr == nil { + t.Fatalf("stream chunks (%q) closed clean, want empty_completion error", aggregated) + } + var authErr *coreauth.Error + if !errors.As(emptyErr, &authErr) { + t.Fatalf("error = %v (%T), want *coreauth.Error", emptyErr, emptyErr) + } + if !authErr.Retryable || authErr.Code != "empty_completion" { + t.Fatalf("error = %+v, want retriable empty_completion", authErr) + } +} + func TestHostRouteModelDefaultsHandledRouterToOwnExecutor(t *testing.T) { host := newRouteModelHostWithRecords(capabilityRecord{ id: "router", diff --git a/internal/registry/model_registry.go b/internal/registry/model_registry.go index ed904bc2716..0cd5b4feed8 100644 --- a/internal/registry/model_registry.go +++ b/internal/registry/model_registry.go @@ -736,11 +736,25 @@ func (r *ModelRegistry) ClearModelQuotaExceeded(clientID, modelID string) { } // SuspendClientModel marks a client's model as temporarily unavailable until explicitly resumed. +// When the client is already suspended for the model, the existing reason is never replaced. // Parameters: // - clientID: The client to suspend // - modelID: The model affected by the suspension // - reason: Optional description for observability func (r *ModelRegistry) SuspendClientModel(clientID, modelID, reason string) { + r.SuspendClientModelReplacingReasons(clientID, modelID, reason) +} + +// SuspendClientModelReplacingReasons marks a client's model as temporarily unavailable. When the +// client is already suspended for the model, the stored reason is replaced only if the existing +// reason matches one of replaceable; otherwise the existing (more specific) reason is preserved. +// Calling it without replaceable reasons is equivalent to SuspendClientModel. +// Parameters: +// - clientID: The client to suspend +// - modelID: The model affected by the suspension +// - reason: Optional description for observability +// - replaceable: Suspension reasons that may be overwritten by reason +func (r *ModelRegistry) SuspendClientModelReplacingReasons(clientID, modelID, reason string, replaceable ...string) { if clientID == "" || modelID == "" { return } @@ -755,8 +769,20 @@ func (r *ModelRegistry) SuspendClientModel(clientID, modelID, reason string) { if registration.SuspendedClients == nil { registration.SuspendedClients = make(map[string]string) } - if _, already := registration.SuspendedClients[clientID]; already { - return + if existingReason, already := registration.SuspendedClients[clientID]; already { + if existingReason == reason { + return + } + canReplace := false + for _, rep := range replaceable { + if existingReason == rep { + canReplace = true + break + } + } + if !canReplace { + return + } } registration.SuspendedClients[clientID] = reason registration.LastUpdated = time.Now() @@ -793,6 +819,63 @@ func (r *ModelRegistry) ResumeClientModel(clientID, modelID string) { log.Debugf("Resumed client %s for model %s", clientID, modelID) } +// ResumeClientModelIfReason atomically verifies that clientID is suspended for modelID with one +// of the given reason(s) and, only if so, resumes it (removing the suspension) under a single +// lock. It reports whether a resume happened. This avoids the TOCTOU of a separate +// GetClientModelSuspensionReason check followed by ResumeClientModel racing with a newer +// suspension recorded between the two. +func (r *ModelRegistry) ResumeClientModelIfReason(clientID, modelID string, resumableReasons ...string) bool { + clientID = strings.TrimSpace(clientID) + modelID = strings.TrimSpace(modelID) + if clientID == "" || modelID == "" { + return false + } + r.mutex.Lock() + defer r.mutex.Unlock() + r.ensureAvailableModelsCacheLocked() + + registration, exists := r.models[modelID] + if !exists || registration == nil || registration.SuspendedClients == nil { + return false + } + reason, suspended := registration.SuspendedClients[clientID] + if !suspended { + return false + } + resumable := false + for _, rr := range resumableReasons { + if reason == rr { + resumable = true + break + } + } + if !resumable { + return false + } + delete(registration.SuspendedClients, clientID) + registration.LastUpdated = time.Now() + r.invalidateAvailableModelsCacheLocked() + log.Debugf("Resumed client %s for model %s (reason %s)", clientID, modelID, reason) + return true +} + +// 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/registry/model_registry_resume_reason_test.go b/internal/registry/model_registry_resume_reason_test.go new file mode 100644 index 00000000000..05ee21dd9d6 --- /dev/null +++ b/internal/registry/model_registry_resume_reason_test.go @@ -0,0 +1,85 @@ +package registry + +import "testing" + +// TestResumeClientModelIfReason_RacePreservesNewerSuspension is a red-proof regression test for +// the TOCTOU in the cooldown resume path. Under concurrent requests for the same credential, an +// explicit reason check (GetClientModelSuspensionReason == "invalid_api_key") followed by a +// separate ResumeClientModel transaction would delete a newer suspension recorded between the +// two. ResumeClientModelIfReason verifies and removes the suspension under one registry lock, so +// a suspension whose reason changed in the interim is left intact. +func TestResumeClientModelIfReason_RacePreservesNewerSuspension(t *testing.T) { + r := newTestModelRegistry() + const ( + clientID = "auth-1" + modelID = "provider/model-1" + ) + r.RegisterClient(clientID, "provider", []*ModelInfo{{ID: modelID}}) + r.SuspendClientModel(clientID, modelID, "invalid_api_key") + + // A concurrent request records a model-specific failure after the initial reason read but + // before the resume transaction, changing the suspension reason. SuspendClientModel refuses to + // overwrite an existing suspension, so this mirrors the interleaved-failure window directly. + r.mutex.Lock() + r.models[modelID].SuspendedClients[clientID] = "budget_exceeded" + r.mutex.Unlock() + + // The conditional resume must refuse: the current reason is no longer invalid_api_key. + if resumed := r.ResumeClientModelIfReason(clientID, modelID, "invalid_api_key"); resumed { + t.Fatalf("ResumeClientModelIfReason resumed a model whose current suspension should remain, want no-op") + } + if reason := r.GetClientModelSuspensionReason(clientID, modelID); reason != "budget_exceeded" { + t.Fatalf("newer suspension reason = %q, want budget_exceeded (must survive)", reason) + } + + // A matching reason resumes normally (after the earlier budget_exceeded suspension clears). + r.SuspendClientModel(clientID, modelID, "budget_exceeded") + r.mutex.Lock() + r.models[modelID].SuspendedClients[clientID] = "invalid_api_key" + r.mutex.Unlock() + if resumed := r.ResumeClientModelIfReason(clientID, modelID, "invalid_api_key"); !resumed { + t.Fatalf("ResumeClientModelIfReason with matching reason did not resume") + } + if reason := r.GetClientModelSuspensionReason(clientID, modelID); reason != "" { + t.Fatalf("suspension reason after resume = %q, want empty", reason) + } +} + +func TestSuspendClientModelReplacingReasons(t *testing.T) { + r := newTestModelRegistry() + const ( + clientID = "auth-1" + modelID = "provider/model-1" + ) + r.RegisterClient(clientID, "provider", []*ModelInfo{{ID: modelID}}) + + // 1. Initial suspension inserts reason. + r.SuspendClientModelReplacingReasons(clientID, modelID, "invalid_api_key") + if reason := r.GetClientModelSuspensionReason(clientID, modelID); reason != "invalid_api_key" { + t.Fatalf("initial suspension reason = %q, want invalid_api_key", reason) + } + + // 2. Already suspended with a replaceable reason -> reason replaced. + r.SuspendClientModelReplacingReasons(clientID, modelID, "not_found", "invalid_api_key") + if reason := r.GetClientModelSuspensionReason(clientID, modelID); reason != "not_found" { + t.Fatalf("suspension reason after replacement = %q, want not_found", reason) + } + + // 3. Already suspended with a non-replaceable reason -> reason preserved. + r.SuspendClientModelReplacingReasons(clientID, modelID, "invalid_api_key", "quota") + if reason := r.GetClientModelSuspensionReason(clientID, modelID); reason != "not_found" { + t.Fatalf("suspension reason after non-matching replacement = %q, want not_found (preserved)", reason) + } + + // 4. Already suspended with no replaceable args -> reason preserved (behaves like SuspendClientModel). + r.SuspendClientModelReplacingReasons(clientID, modelID, "invalid_api_key") + if reason := r.GetClientModelSuspensionReason(clientID, modelID); reason != "not_found" { + t.Fatalf("suspension reason without replaceable args = %q, want not_found (preserved)", reason) + } + + // 5. SuspendClientModel delegation preserves existing reason. + r.SuspendClientModel(clientID, modelID, "invalid_api_key") + if reason := r.GetClientModelSuspensionReason(clientID, modelID); reason != "not_found" { + t.Fatalf("suspension reason via SuspendClientModel = %q, want not_found (preserved)", reason) + } +} diff --git a/internal/runtime/executor/claude_executor_ratelimit_test.go b/internal/runtime/executor/claude_executor_ratelimit_test.go index 0c1a852fb35..50bb6ef4fb1 100644 --- a/internal/runtime/executor/claude_executor_ratelimit_test.go +++ b/internal/runtime/executor/claude_executor_ratelimit_test.go @@ -721,7 +721,7 @@ func TestClaudeExecutor_AuthManager_MultiModelPoolStreamStopsProbingOn429(t *tes attemptsCred2.Add(1) w.Header().Set("Content-Type", "text/event-stream") w.WriteHeader(http.StatusOK) - _, _ = w.Write([]byte("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg-1\",\"model\":\"claude-3-5-sonnet-20241022\"}}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n")) + _, _ = w.Write([]byte("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg-1\",\"model\":\"claude-3-5-sonnet-20241022\"}}\n\nevent: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"ok\"}}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n")) })) defer server2.Close() diff --git a/internal/runtime/executor/home_codex_terminal_test.go b/internal/runtime/executor/home_codex_terminal_test.go index c6f7342563c..cd901929084 100644 --- a/internal/runtime/executor/home_codex_terminal_test.go +++ b/internal/runtime/executor/home_codex_terminal_test.go @@ -43,9 +43,24 @@ 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 { - _ = conn.WriteJSON(map[string]any{"type": "response.completed", "response": map[string]any{"id": "response-2", "output": []any{}}}) + writeCompletion := func() { + _ = conn.WriteJSON(map[string]any{"type": "response.completed", "response": map[string]any{ + "id": "response-2", "object": "response", "status": "completed", + "output": []any{map[string]any{ + "type": "message", "status": "completed", "role": "assistant", + "content": []any{map[string]any{"type": "output_text", "text": "ok"}}, + }}, + "usage": map[string]any{"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + }}) + } + writeCompletion() + if _, _, errRead := conn.ReadMessage(); errRead != nil { + return + } + writeCompletion() } for { if _, _, errRead := conn.ReadMessage(); errRead != nil { diff --git a/sdk/api/handlers/gemini/gemini_handlers.go b/sdk/api/handlers/gemini/gemini_handlers.go index f01dd9b067b..16c8f655f11 100644 --- a/sdk/api/handlers/gemini/gemini_handlers.go +++ b/sdk/api/handlers/gemini/gemini_handlers.go @@ -163,6 +163,18 @@ 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. +func pendingGeminiStreamError(errChan <-chan *interfaces.ErrorMessage) *interfaces.ErrorMessage { + errMsg, hasPendingError := handlers.PendingStreamError(errChan) + if !hasPendingError { + return nil + } + return errMsg +} + // 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 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 00000000000..f08b999ea01 --- /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 c525390166c..6ca1f645d67 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" ) @@ -213,6 +214,26 @@ func TestEnrichAuthSelectionError_IgnoresOtherErrors(t *testing.T) { } } +func TestExecutionErrorMessageMapsTrustedProvenance(t *testing.T) { + trusted := &coreexecutor.RequestTerminatedError{HTTPStatus: http.StatusTooManyRequests, Body: []byte("trusted"), Trusted: true} + got := executionErrorMessage(trusted) + if got == nil || !got.DirectResponse || !got.TrustedDirectResponse { + t.Fatalf("trusted terminated error = %#v", got) + } + if got.StatusCode != http.StatusTooManyRequests || string(got.Body) != "trusted" { + t.Fatalf("trusted terminated body/status = %#v", got) + } + + untrusted := &coreexecutor.RequestTerminatedError{HTTPStatus: http.StatusBadGateway, Body: []byte("upstream")} + got = executionErrorMessage(untrusted) + if got == nil || !got.DirectResponse || got.TrustedDirectResponse { + t.Fatalf("untrusted terminated error = %#v", got) + } + if got.StatusCode != http.StatusBadGateway || string(got.Body) != "upstream" { + t.Fatalf("untrusted terminated body/status = %#v", got) + } +} + func TestExecutionErrorMessageMapsContextStatuses(t *testing.T) { tests := []struct { name string diff --git a/sdk/api/handlers/handlers_errors.go b/sdk/api/handlers/handlers_errors.go index 57df50e0437..47f7a1568bb 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 e4a4c254858..5429f56203f 100644 --- a/sdk/api/handlers/handlers_execution.go +++ b/sdk/api/handlers/handlers_execution.go @@ -315,11 +315,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 dfed4310e8e..b179884bff1 100644 --- a/sdk/api/handlers/handlers_interceptors.go +++ b/sdk/api/handlers/handlers_interceptors.go @@ -154,10 +154,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_interceptors_test.go b/sdk/api/handlers/handlers_interceptors_test.go index c328a620979..426fb250327 100644 --- a/sdk/api/handlers/handlers_interceptors_test.go +++ b/sdk/api/handlers/handlers_interceptors_test.go @@ -265,7 +265,7 @@ func TestHandlerRequestInterceptorTerminatesBeforeAuth(t *testing.T) { if body != nil || headers != nil { t.Fatalf("terminated response body = %q, headers = %#v", body, headers) } - if errMsg == nil || !errMsg.DirectResponse || errMsg.StatusCode != http.StatusForbidden { + if errMsg == nil || !errMsg.DirectResponse || !errMsg.TrustedDirectResponse || errMsg.StatusCode != http.StatusForbidden { t.Fatalf("termination error = %#v", errMsg) } if string(errMsg.Body) != `{"error":"blocked"}` || errMsg.Headers.Get("X-Policy") != "blocked" { @@ -312,7 +312,7 @@ func TestHandlerRequestInterceptorTerminatesAfterAuth(t *testing.T) { }) _, _, errMsg := handler.ExecuteWithAuthManager(context.Background(), "openai", model, []byte(`{"model":"`+model+`"}`), "") - if errMsg == nil || !errMsg.DirectResponse || errMsg.StatusCode != http.StatusTooManyRequests { + if errMsg == nil || !errMsg.DirectResponse || !errMsg.TrustedDirectResponse || errMsg.StatusCode != http.StatusTooManyRequests { t.Fatalf("termination error = %#v", errMsg) } if beforeRequestID == "" || afterRequestID != beforeRequestID || completion.RequestID != beforeRequestID { @@ -358,7 +358,7 @@ func TestHandlerAfterAuthTerminationSkipsCountAndStreamExecutors(t *testing.T) { } errMsg = <-errChan } - if errMsg == nil || !errMsg.DirectResponse || errMsg.StatusCode != http.StatusForbidden { + if errMsg == nil || !errMsg.DirectResponse || !errMsg.TrustedDirectResponse || errMsg.StatusCode != http.StatusForbidden { t.Fatalf("termination error = %#v", errMsg) } if afterCalls != 1 { diff --git a/sdk/api/handlers/handlers_stream_bootstrap_test.go b/sdk/api/handlers/handlers_stream_bootstrap_test.go index f10d74743ba..d8679f314cd 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 } @@ -1178,10 +1178,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 e83337a2153..64d96eed8d1 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) + startupErrText := "" + if errMsg.Error != nil { + startupErrText = errMsg.Error.Error() + } + if !strings.HasPrefix(startupErrText, "startup failed") { + t.Fatalf("error = %q, want startup failed prefix", startupErrText) + } + if !strings.Contains(startupErrText, "attempted routes: [codex:error]") { + t.Fatalf("error = %q, want sanitized route summary attempted routes: [codex:error]", startupErrText) } 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/cliproxy/auth/conductor_availability_test.go b/sdk/cliproxy/auth/conductor_availability_test.go index 7e07cc07148..02ea63947f5 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,386 @@ 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) + } +} + +func TestManager_ModelNotSupportedSuspensionResumesOnOwnSuccess(t *testing.T) { + manager := NewManager(nil, nil, nil) + ctx := context.Background() + authID := "model-not-supported-resume-auth" + model := "model-a" + + reg := registry.GetGlobalRegistry() + reg.RegisterClient(authID, "openai", []*registry.ModelInfo{ + {ID: model}, + }) + t.Cleanup(func() { + reg.UnregisterClient(authID) + }) + + if _, errRegister := manager.Register(ctx, &Auth{ + ID: authID, + Provider: "openai", + Status: StatusActive, + ModelStates: map[string]*ModelState{ + model: {Status: StatusActive}, + }, + }); errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + + // Fail model with model_not_supported -> model should be suspended in registry + manager.MarkResult(ctx, Result{ + AuthID: authID, + Provider: "openai", + Model: model, + Success: false, + Error: &Error{HTTPStatus: http.StatusBadRequest, Code: "model_not_supported", Message: "model not supported"}, + }) + + if count := reg.GetModelCount(model); count != 0 { + t.Fatalf("registry model count for model after suspension = %d, want 0", count) + } + if reason := reg.GetClientModelSuspensionReason(authID, model); reason != "model_not_supported" { + t.Fatalf("suspension reason = %q, want model_not_supported", reason) + } + + // Success on model -> should resume model in registry + manager.MarkResult(ctx, Result{ + AuthID: authID, + Provider: "openai", + Model: model, + Success: true, + }) + + if count := reg.GetModelCount(model); count != 1 { + t.Fatalf("registry model count for model after success = %d, want 1", count) + } + if reason := reg.GetClientModelSuspensionReason(authID, model); reason != "" { + t.Fatalf("suspension reason after success = %q, want empty", reason) + } +} + +// TestManager_ModelSpecificResumableSiblingSuspensionSurvivesSiblingSuccess verifies that a +// sibling model suspended for a resumable model-specific reason (not_found, quota, +// payment_required) keeps its registry suspension when a different model of the same credential +// succeeds. Only credential-wide reasons like invalid_api_key justify cross-model resumption. +func TestManager_ModelSpecificResumableSiblingSuspensionSurvivesSiblingSuccess(t *testing.T) { + manager := NewManager(nil, nil, nil) + ctx := context.Background() + authID := "sibling-resumable-suspension-auth" + modelA := "model-a2" + modelB := "model-b2" + + 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 a resumable, model-specific reason (not_found). + manager.MarkResult(ctx, Result{ + AuthID: authID, + Provider: "openai", + Model: modelB, + Success: false, + Error: &Error{HTTPStatus: http.StatusNotFound, Code: "not_found", Message: "model b not found"}, + }) + if count := reg.GetModelCount(modelB); count != 0 { + t.Fatalf("registry model count for modelB after suspension = %d, want 0", count) + } + + // Success on modelA -> must NOT resume modelB (model-specific reason). + manager.MarkResult(ctx, Result{ + AuthID: authID, + Provider: "openai", + Model: modelA, + Success: true, + }) + if count := reg.GetModelCount(modelB); count != 0 { + t.Fatalf("registry model count for modelB after modelA success = %d, want 0 (model-specific suspension should survive)", count) + } + if count := reg.GetModelCount(modelA); count != 1 { + t.Fatalf("registry model count for modelA after success = %d, want 1", count) + } +} + +// TestManager_CredentialWideSiblingSuspensionResumesOnSiblingSuccess verifies that a sibling +// suspended for a credential-wide reason (invalid_api_key) is resumed by a successful request on +// another model of the same credential. +func TestManager_CredentialWideSiblingSuspensionResumesOnSiblingSuccess(t *testing.T) { + manager := NewManager(nil, nil, nil) + ctx := context.Background() + authID := "sibling-credentialwide-resume-auth" + modelA := "model-a3" + modelB := "model-b3" + + 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) + } + + // Suspend sibling modelB with a genuine credential-wide reason (invalid_api_key) directly in + // the registry. A live invalid_api_key failure would also mark the credential with an active + // credential_quota cooldown that legitimately suppresses the resume path; suspending the sibling + // directly isolates the sibling-resume loop's reason scoping. + reg.SuspendClientModel(authID, modelB, "invalid_api_key") + if count := reg.GetModelCount(modelB); count != 0 { + t.Fatalf("registry model count for modelB after suspension = %d, want 0", count) + } + + // Success on modelA -> resumes modelB (credential-wide reason). + manager.MarkResult(ctx, Result{ + AuthID: authID, + Provider: "openai", + Model: modelA, + Success: true, + }) + if count := reg.GetModelCount(modelB); count != 1 { + t.Fatalf("registry model count for modelB after modelA success = %d, want 1 (credential-wide suspension should resume)", count) + } +} + +// TestManager_ModelSpecificFailureOverwritesCredentialWideSuspension reproduces the scenario +// where a credential-wide invalid_api_key fanout suspends all models, and model B later encounters +// a model-specific failure (e.g. not_found). Model B's suspension reason must be updated to the +// model-specific reason so that a subsequent success on model A does not erroneously resume model B. +func TestManager_ModelSpecificFailureOverwritesCredentialWideSuspension(t *testing.T) { + manager := NewManager(nil, nil, nil) + ctx := context.Background() + authID := "sibling-overwrite-suspension-auth" + modelA := "model-a4" + modelB := "model-b4" + + 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) + } + + // 1. Initial credential-wide suspension on all models (invalid_api_key). + // To isolate the suspension reason overwrite and sibling-resume behavior without active + // credential_quota cooldown gating, suspend modelB directly with invalid_api_key as the + // fanout produces. + reg.SuspendClientModel(authID, modelB, "invalid_api_key") + if reason := reg.GetClientModelSuspensionReason(authID, modelB); reason != "invalid_api_key" { + t.Fatalf("modelB initial suspension reason = %q, want invalid_api_key", reason) + } + + // 2. Model B records a model-specific failure (404 not_found). + manager.MarkResult(ctx, Result{ + AuthID: authID, + Provider: "openai", + Model: modelB, + Success: false, + Error: &Error{HTTPStatus: http.StatusNotFound, Code: "not_found", Message: "model b not found"}, + }) + if reason := reg.GetClientModelSuspensionReason(authID, modelB); reason != "not_found" { + t.Fatalf("modelB suspension reason after 404 = %q, want not_found", reason) + } + if count := reg.GetModelCount(modelB); count != 0 { + t.Fatalf("registry model count for modelB after 404 = %d, want 0", count) + } + + // 3. Model A succeeds -> sibling-resume loop runs. Model B must STAY suspended. + manager.MarkResult(ctx, Result{ + AuthID: authID, + Provider: "openai", + Model: modelA, + Success: true, + }) + if reason := reg.GetClientModelSuspensionReason(authID, modelB); reason != "not_found" { + t.Fatalf("modelB suspension reason after modelA success = %q, want not_found (must survive)", reason) + } + if count := reg.GetModelCount(modelB); count != 0 { + t.Fatalf("registry model count for modelB after modelA success = %d, want 0 (must stay suspended)", count) + } + if count := reg.GetModelCount(modelA); count != 1 { + t.Fatalf("registry model count for modelA after success = %d, want 1", count) + } +} diff --git a/sdk/cliproxy/auth/conductor_cooldown.go b/sdk/cliproxy/auth/conductor_cooldown.go index 815ce70e32e..87bbadca35f 100644 --- a/sdk/cliproxy/auth/conductor_cooldown.go +++ b/sdk/cliproxy/auth/conductor_cooldown.go @@ -24,6 +24,26 @@ var quotaCooldownDisabled atomic.Bool var transientErrorCooldownSeconds atomic.Int64 +// resumableCooldownReasons are the registry suspension reasons a successful result can clear for +// the model that just succeeded (they include model-specific reasons like not_found and quota). +var resumableCooldownReasons = []string{ + "invalid_api_key", + "invalid_grant", + "unauthorized", + "payment_required", + "not_found", + "model_not_supported", + "quota", +} + +// credentialWideCooldownReasons are the suspension reasons that span every model of a credential +// and may therefore be cleared on sibling models when a different model of the same credential +// succeeds. Only invalid_api_key is propagated credential-wide by SuspendClientModel; invalid_grant, +// unauthorized, and model-specific reasons are recorded per-model and must not resume siblings. +var credentialWideCooldownReasons = []string{ + "invalid_api_key", +} + // SetQuotaCooldownDisabled toggles auth/model cooldown scheduling globally. func SetQuotaCooldownDisabled(disable bool) { quotaCooldownDisabled.Store(disable) @@ -103,6 +123,19 @@ func recoverableFailureRetryAfter(now time.Time, disableCooling bool) time.Time return nextTransientErrorRetryAfter(now) } +func laterTime(a, b time.Time) time.Time { + if a.IsZero() { + return b + } + if b.IsZero() { + return a + } + if a.After(b) { + return a + } + return b +} + // SetConfig updates the runtime config snapshot used by request-time helpers. // Callers should provide the latest config on reload so per-credential alias mapping stays in sync. func (m *Manager) SetConfig(cfg *internalconfig.Config) { @@ -793,6 +826,48 @@ 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 = laterTime(state.NextRetryAfter, next) + if !(state.Quota.Exceeded && state.Quota.NextRecoverAt.After(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 = laterTime(otherState.NextRetryAfter, next) + if !(otherState.Quota.Exceeded && otherState.Quota.NextRecoverAt.After(next)) { + otherState.Quota = QuotaState{ + Exceeded: true, + Reason: "credential_quota", + NextRecoverAt: next, + } + } + } + } + auth.Unavailable = true + auth.Status = StatusError + auth.StatusMessage = "invalid_api_key" + if !(auth.Quota.Exceeded && auth.Quota.NextRecoverAt.After(next)) { + auth.Quota = QuotaState{ + Exceeded: true, + Reason: "credential_quota", + NextRecoverAt: next, + } + } + auth.NextRetryAfter = laterTime(auth.NextRetryAfter, next) + suspendReason = "invalid_api_key" + shouldSuspendModel = true + } } else { switch statusCode { case 401: @@ -927,9 +1002,28 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { registry.GetGlobalRegistry().SetModelQuotaExceeded(result.AuthID, modelKey) } if shouldResumeModel { - registry.GetGlobalRegistry().ResumeClientModel(result.AuthID, modelKey) + // Sibling models resume only for credential-wide reasons (invalid_api_key); model-specific + // suspensions (not_found, quota, payment_required, ...) must survive until that sibling + // succeeds on its own. + for _, m := range modelsForRegisteredAuth(result.AuthID) { + registry.GetGlobalRegistry().ResumeClientModelIfReason(result.AuthID, m, credentialWideCooldownReasons...) + } + if modelKey != "" { + registry.GetGlobalRegistry().ResumeClientModelIfReason(result.AuthID, modelKey, resumableCooldownReasons...) + } } 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 { + // A model-specific reason must overwrite a stale credential-wide one, otherwise the + // sibling-resume loop above would clear this suspension when another model succeeds. + registry.GetGlobalRegistry().SuspendClientModelReplacingReasons(result.AuthID, modelKey, suspendReason, credentialWideCooldownReasons...) + } } m.hook.OnResult(ctx, result) @@ -1452,7 +1546,7 @@ func isCredentialScopedError(err error) bool { IsCredentialScoped() bool } var csp credentialScopedProvider - return errors.As(err, &csp) && csp != nil && csp.IsCredentialScoped() + return (errors.As(err, &csp) && csp != nil && csp.IsCredentialScoped()) || isInvalidAPIKeyError(err) } func statusCodeFromResult(err *Error) int { @@ -1524,6 +1618,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 @@ -1850,6 +1976,9 @@ func isRequestInvalidError(err error) bool { if isInvalidGrantError(err) { return false } + if isInvalidAPIKeyError(err) { + return false + } if isModelSupportError(err) { return false } @@ -1912,6 +2041,38 @@ 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 = laterTime(auth.NextRetryAfter, next) + if !(auth.Quota.Exceeded && auth.Quota.NextRecoverAt.After(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 = laterTime(state.NextRetryAfter, next) + if !(state.Quota.Exceeded && state.Quota.NextRecoverAt.After(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 00000000000..c130aa35ae2 --- /dev/null +++ b/sdk/cliproxy/auth/conductor_cooldown_retry_reset_test.go @@ -0,0 +1,316 @@ +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 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) { + disabled := true + manager := NewManager(nil, nil, nil) + manager.SetConfigSnapshot(&internalconfig.Config{ + OpenAICompatibility: []internalconfig.OpenAICompatibility{ + { + Name: "custom-openai", + DisableCooling: &disabled, + }, + }, + }) + 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 c98508ab2ee..d36e1398b8e 100644 --- a/sdk/cliproxy/auth/conductor_execution.go +++ b/sdk/cliproxy/auth/conductor_execution.go @@ -44,22 +44,27 @@ 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() { - resp, errHome := m.executeHome(ctx, normalized, req, opts, false) - return resp, unwrapRequestStopError(errHome) + resp, errHome := m.executeHome(ctx, normalized, req, opts, false, tracker) + return resp, wrapRouteExhaustion(unwrapRequestStopError(errHome), tracker) } defaultRequestRetry, maxRetryCredentials, maxWait := m.retrySettings() var lastErr error retryModel := authSelectionModelFromOptions(opts, req.Model) + tried := make(map[string]struct{}) for attempt := 0; ; attempt++ { - resp, errExec := m.executeMixedOnce(ctx, normalized, req, opts, maxRetryCredentials, attempt, defaultRequestRetry) + resp, errExec := m.executeMixedOnce(ctx, normalized, req, opts, maxRetryCredentials, attempt, defaultRequestRetry, tried, tracker) if errExec == nil { return resp, nil } + if lastErr != nil && isAuthUnavailableOrNotFound(errExec) { + errExec = lastErr + } if isRequestTerminatedError(errExec) || isRequestStopError(errExec) { - return cliproxyexecutor.Response{}, unwrapRequestStopError(errExec) + return cliproxyexecutor.Response{}, wrapRouteExhaustion(unwrapRequestStopError(errExec), tracker) } lastErr = errExec wait, shouldRetry := m.shouldRetryAfterErrorWithHomeRetryLimit(ctx, opts, errExec, attempt, normalized, retryModel, maxWait, -1, defaultRequestRetry) @@ -74,14 +79,14 @@ func (m *Manager) Execute(ctx context.Context, providers []string, req cliproxye lastErr = unwrapRequestStopError(lastErr) 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. @@ -91,22 +96,27 @@ 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() { - resp, errHome := m.executeHome(ctx, normalized, req, opts, true) - return resp, unwrapRequestStopError(errHome) + resp, errHome := m.executeHome(ctx, normalized, req, opts, true, tracker) + return resp, wrapRouteExhaustion(unwrapRequestStopError(errHome), tracker) } defaultRequestRetry, maxRetryCredentials, maxWait := m.retrySettings() var lastErr error retryModel := authSelectionModelFromOptions(opts, req.Model) + tried := make(map[string]struct{}) for attempt := 0; ; attempt++ { - resp, errExec := m.executeCountMixedOnce(ctx, normalized, req, opts, maxRetryCredentials, attempt, defaultRequestRetry) + resp, errExec := m.executeCountMixedOnce(ctx, normalized, req, opts, maxRetryCredentials, attempt, defaultRequestRetry, tried, tracker) if errExec == nil { return resp, nil } + if lastErr != nil && isAuthUnavailableOrNotFound(errExec) { + errExec = lastErr + } if isRequestTerminatedError(errExec) || isRequestStopError(errExec) { - return cliproxyexecutor.Response{}, unwrapRequestStopError(errExec) + return cliproxyexecutor.Response{}, wrapRouteExhaustion(unwrapRequestStopError(errExec), tracker) } lastErr = errExec wait, shouldRetry := m.shouldRetryAfterErrorWithHomeRetryLimit(ctx, opts, errExec, attempt, normalized, retryModel, maxWait, -1, defaultRequestRetry) @@ -118,9 +128,9 @@ func (m *Manager) ExecuteCount(ctx context.Context, providers []string, req clip } } if lastErr != nil { - return cliproxyexecutor.Response{}, unwrapRequestStopError(lastErr) + return cliproxyexecutor.Response{}, wrapRouteExhaustion(unwrapRequestStopError(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. @@ -145,11 +155,16 @@ func (m *Manager) ExecuteStream(ctx context.Context, providers []string, req cli attempt := 0 retryRoundPending := false retryRoundWaited := false + tracker := newRouteAttemptTracker() + tried := make(map[string]struct{}) for { - result, errStream := m.executeStreamMixedOnce(ctx, normalized, req, opts, maxRetryCredentials, &homeRetryLimit, attempt, defaultRequestRetry) + result, errStream := m.executeStreamMixedOnce(ctx, normalized, req, opts, maxRetryCredentials, &homeRetryLimit, attempt, defaultRequestRetry, tried, tracker) if errStream == nil { return result, nil } + if lastErr != nil && isAuthUnavailableOrNotFound(errStream) { + errStream = lastErr + } if m.HomeEnabled() && retryRoundPending { if wait, okWait := pendingHomeRetryRoundDelay(errStream, maxWait, &homeRetryLimit, pinnedAuthIDFromMetadata(opts.Metadata) == ""); okWait && m.homeRetryAllowed(attempt-1, homeRetryLimit) { if retryRoundWaited { @@ -165,7 +180,7 @@ func (m *Manager) ExecuteStream(ctx context.Context, providers []string, req cli retryRoundPending = false retryRoundWaited = false if isRequestTerminatedError(errStream) || isRequestStopError(errStream) { - return nil, unwrapRequestStopError(errStream) + return nil, wrapRouteExhaustion(unwrapRequestStopError(errStream), tracker) } lastErr = errStream wait, shouldRetry := m.shouldRetryAfterErrorWithHomeRetryLimit(ctx, opts, errStream, attempt, normalized, retryModel, maxWait, homeRetryLimit, defaultRequestRetry) @@ -183,18 +198,18 @@ func (m *Manager) ExecuteStream(ctx context.Context, providers []string, req cli lastErr = unwrapRequestStopError(lastErr) 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(), lastErr), nil + return streamErrorResult(bootstrapErr.Headers(), wrapRouteExhaustion(lastErr, 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 { @@ -231,6 +246,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 @@ -300,7 +316,32 @@ 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, retryRound int, defaultRequestRetry int) (cliproxyexecutor.Response, error) { +func isAuthUnavailableOrNotFound(err error) bool { + var authErr *Error + if !errors.As(err, &authErr) || authErr == nil { + return false + } + return authErr.Code == "auth_not_found" || authErr.Code == "auth_unavailable" || authErr.Code == "home_unavailable" +} + +func persistExcludedAuthForRetry(m *Manager, auth *Auth, err error, retryRound int, defaultRequestRetry int, excluded map[string]struct{}) { + if auth == nil || excluded == nil { + return + } + if effectiveRequestRetryLimit(auth, defaultRequestRetry) <= retryRound { + excluded[auth.ID] = struct{}{} + return + } + if m.cooldownDisabledForAuth(auth) && retryAfterFromError(err) != nil { + excluded[auth.ID] = struct{}{} + } +} + +func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, maxRetryCredentials int, retryRound int, defaultRequestRetry int, excluded map[string]struct{}, optionalTracker ...*routeAttemptTracker) (cliproxyexecutor.Response, error) { + var tracker *routeAttemptTracker + if len(optionalTracker) > 0 { + tracker = optionalTracker[0] + } if len(providers) == 0 { return cliproxyexecutor.Response{}, &Error{Code: "provider_not_found", Message: "no provider supplied"} } @@ -309,11 +350,20 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req opts = ensureRequestedModelMetadata(opts, routeModel) homeMode := m.HomeEnabled() homeAuthCount := 1 + if excluded == nil { + excluded = make(map[string]struct{}) + } tried := make(map[string]struct{}) + for authID := range excluded { + tried[authID] = struct{}{} + } if !homeMode { for authID := range m.requestRetryRoundExclusions(retryRound, defaultRequestRetry) { tried[authID] = struct{}{} } + for authID := range extractExcludedAuthIDs(opts.Metadata) { + tried[authID] = struct{}{} + } } attempted := make(map[string]struct{}) var lastErr error @@ -364,6 +414,8 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req result := Result{AuthID: auth.ID, Provider: provider, Model: routeModel, Success: false, Error: resultErrorFromError(errPrepare), Options: pickOpts} m.MarkResult(execCtx, result) lastErr = errPrepare + tracker.Record(auth, errPrepare) + persistExcludedAuthForRetry(m, auth, errPrepare, retryRound, defaultRequestRetry, excluded) continue } var authErr error @@ -413,6 +465,7 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: errExec == nil, Options: execOpts} if errExec != nil { result.Error = resultErrorFromError(errExec) + tracker.Record(auth, errExec) if ra := retryAfterFromError(errExec); ra != nil { result.RetryAfter = ra } @@ -431,6 +484,7 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req return cliproxyexecutor.Response{}, wrapRequestStopError(errExec) } authErr = errExec + tracker.Record(auth, errExec) if result.CredentialScope { break } @@ -440,12 +494,26 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req return cliproxyexecutor.Response{}, errExec } authErr = errExec + tracker.Record(auth, errExec) + persistExcludedAuthForRetry(m, auth, errExec, retryRound, defaultRequestRetry, excluded) if result.CredentialScope { break } continue } m.MarkResult(execCtx, result) + if isEmptyCompletionPayload(resp.Payload) { + result.Success = false + result.Error = errEmptyCompletion + m.MarkResult(execCtx, result) + lastErr = errEmptyCompletion + tracker.Record(auth, errEmptyCompletion) + persistExcludedAuthForRetry(m, auth, errEmptyCompletion, retryRound, defaultRequestRetry, excluded) + if homeMode { + homeAuthCount++ + } + continue + } attemptAliasResult := resolveAttemptAliasResult(routing, auth, routeModel, upstreamModel, aliasResult) rewriteForceMappedResponse(&resp, attemptAliasResult) return resp, nil @@ -466,6 +534,7 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req return cliproxyexecutor.Response{}, authErr } lastErr = authErr + persistExcludedAuthForRetry(m, auth, authErr, retryRound, defaultRequestRetry, excluded) if homeMode { homeAuthCount++ } @@ -474,7 +543,11 @@ 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, retryRound int, defaultRequestRetry int) (cliproxyexecutor.Response, error) { +func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, maxRetryCredentials int, retryRound int, defaultRequestRetry int, excluded map[string]struct{}, optionalTracker ...*routeAttemptTracker) (cliproxyexecutor.Response, error) { + var tracker *routeAttemptTracker + if len(optionalTracker) > 0 { + tracker = optionalTracker[0] + } if len(providers) == 0 { return cliproxyexecutor.Response{}, &Error{Code: "provider_not_found", Message: "no provider supplied"} } @@ -483,11 +556,20 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, opts = ensureRequestedModelMetadata(opts, routeModel) homeMode := m.HomeEnabled() homeAuthCount := 1 + if excluded == nil { + excluded = make(map[string]struct{}) + } tried := make(map[string]struct{}) + for authID := range excluded { + tried[authID] = struct{}{} + } if !homeMode { for authID := range m.requestRetryRoundExclusions(retryRound, defaultRequestRetry) { tried[authID] = struct{}{} } + for authID := range extractExcludedAuthIDs(opts.Metadata) { + tried[authID] = struct{}{} + } } attempted := make(map[string]struct{}) var lastErr error @@ -538,6 +620,8 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, result := Result{AuthID: auth.ID, Provider: provider, Model: routeModel, Success: false, Error: resultErrorFromError(errPrepare), Options: pickOpts} m.MarkResult(execCtx, result) lastErr = errPrepare + tracker.Record(auth, errPrepare) + persistExcludedAuthForRetry(m, auth, errPrepare, retryRound, defaultRequestRetry, excluded) continue } var authErr error @@ -587,6 +671,7 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: errExec == nil, Options: execOpts} if errExec != nil { result.Error = resultErrorFromError(errExec) + tracker.Record(auth, errExec) if ra := retryAfterFromError(errExec); ra != nil { result.RetryAfter = ra } @@ -609,6 +694,7 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, return cliproxyexecutor.Response{}, wrapRequestStopError(errExec) } authErr = errExec + tracker.Record(auth, errExec) if result.CredentialScope { break } @@ -618,12 +704,26 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, return cliproxyexecutor.Response{}, errExec } authErr = errExec + tracker.Record(auth, errExec) + persistExcludedAuthForRetry(m, auth, errExec, retryRound, defaultRequestRetry, excluded) if result.CredentialScope { break } continue } m.MarkResult(execCtx, result) + if isEmptyCompletionPayload(resp.Payload) { + result.Success = false + result.Error = errEmptyCompletion + m.MarkResult(execCtx, result) + lastErr = errEmptyCompletion + tracker.Record(auth, errEmptyCompletion) + persistExcludedAuthForRetry(m, auth, errEmptyCompletion, retryRound, defaultRequestRetry, excluded) + if homeMode { + homeAuthCount++ + } + continue + } attemptAliasResult := resolveAttemptAliasResult(routing, auth, routeModel, upstreamModel, aliasResult) rewriteForceMappedResponse(&resp, attemptAliasResult) return resp, nil @@ -644,6 +744,7 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, return cliproxyexecutor.Response{}, authErr } lastErr = authErr + persistExcludedAuthForRetry(m, auth, authErr, retryRound, defaultRequestRetry, excluded) if homeMode { homeAuthCount++ } @@ -652,7 +753,11 @@ 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, homeRetryLimit *int, retryRound int, defaultRequestRetry int) (*cliproxyexecutor.StreamResult, error) { +func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, maxRetryCredentials int, homeRetryLimit *int, retryRound int, defaultRequestRetry int, excluded map[string]struct{}, optionalTracker ...*routeAttemptTracker) (*cliproxyexecutor.StreamResult, error) { + var tracker *routeAttemptTracker + if len(optionalTracker) > 0 { + tracker = optionalTracker[0] + } if len(providers) == 0 { return nil, &Error{Code: "provider_not_found", Message: "no provider supplied"} } @@ -662,11 +767,20 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string opts = ensureRequestedModelMetadata(opts, routeModel) homeMode := m.HomeEnabled() homeAuthCount := 1 + if excluded == nil { + excluded = make(map[string]struct{}) + } tried := make(map[string]struct{}) + for authID := range excluded { + tried[authID] = struct{}{} + } if !homeMode { for authID := range m.requestRetryRoundExclusions(retryRound, defaultRequestRetry) { tried[authID] = struct{}{} } + for authID := range extractExcludedAuthIDs(opts.Metadata) { + tried[authID] = struct{}{} + } } homeExcludedAuthIDs := make(map[string]struct{}) homeSameAuthRetries := make(map[string]int) @@ -674,6 +788,7 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string homeSameAuthRetryPending := false attempted := make(map[string]struct{}) unauthorizedRefreshTried := make(map[string]struct{}) + emptyCompletionTried := make(map[string]struct{}) var lastErr error var roundTiming homeRetryRoundTiming for { @@ -745,6 +860,24 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string homeSameAuthRetryPending = false } if selection != nil { + // #4881: an auth that already returned an empty completion must rotate. + if _, emptyAlready := emptyCompletionTried[auth.ID]; emptyAlready { + selection.End("repeated_empty_completion_auth") + if lastErr != nil { + return nil, lastErr + } + return nil, errEmptyCompletion + } + + // #4881: an auth that already had an unauthorized refresh must rotate. + if _, refreshedAlready := unauthorizedRefreshTried[auth.ID]; refreshedAlready { + selection.End("repeated_refresh_auth") + if lastErr != nil { + return nil, lastErr + } + return nil, errEmptyCompletion + } + // A legacy Home may ignore excluded_auth_ids and return the same // credential again. Reject credentials explicitly excluded from this // round while retaining the explicit same-auth retry path, which @@ -857,6 +990,8 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string m.MarkResult(execCtx, result) } lastErr = errPrepare + tracker.Record(auth, errPrepare) + persistExcludedAuthForRetry(m, auth, errPrepare, retryRound, defaultRequestRetry, excluded) if homeMode { roundTiming.Observe(lastErr) } @@ -908,10 +1043,13 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string return nil, wrapRequestStopError(errStream) } lastErr = errStream + tracker.Record(auth, errStream) + persistExcludedAuthForRetry(m, auth, errStream, retryRound, defaultRequestRetry, excluded) if homeMode { roundTiming.Observe(lastErr) - } - if homeMode { + if isEmptyCompletionError(errStream) { + emptyCompletionTried[auth.ID] = struct{}{} + } homeAuthCount++ } continue @@ -920,10 +1058,13 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string return nil, errStream } lastErr = errStream + tracker.Record(auth, errStream) + persistExcludedAuthForRetry(m, auth, errStream, retryRound, defaultRequestRetry, excluded) if homeMode { roundTiming.Observe(lastErr) - } - if homeMode { + if isEmptyCompletionError(errStream) { + emptyCompletionTried[auth.ID] = struct{}{} + } homeAuthCount++ } continue @@ -951,6 +1092,26 @@ func shouldExcludeHomeAuthAfterStreamError(ctx context.Context, auth *Auth, err return !isUnauthorizedError(err) || auth == nil || auth.AuthKind() != AuthKindOAuth } +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 cloneRequestMetadata(src map[string]any) map[string]any { if len(src) == 0 { return make(map[string]any, 4) diff --git a/sdk/cliproxy/auth/conductor_fast_error_test.go b/sdk/cliproxy/auth/conductor_fast_error_test.go index 7956bdb39ba..366e101d81b 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 ce6cf915f32..8512b6ed0c0 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 ce956a1706a..417eec58e0a 100644 --- a/sdk/cliproxy/auth/conductor_home.go +++ b/sdk/cliproxy/auth/conductor_home.go @@ -1362,6 +1362,12 @@ func (m *Manager) tryAntigravityCreditsExecute(ctx context.Context, req cliproxy } continue } + if isEmptyCompletionPayload(resp.Payload) { + result.Success = false + result.Error = errEmptyCompletion + m.recordExecutionResult(creditsCtx, result, c.auth, false) + continue + } m.MarkResult(creditsCtx, result) attemptAliasResult := resolveAttemptAliasResult(routing, c.auth, routeModel, upstreamModel, aliasResult) rewriteForceMappedResponse(&resp, attemptAliasResult) diff --git a/sdk/cliproxy/auth/conductor_home_execution.go b/sdk/cliproxy/auth/conductor_home_execution.go index d965712a0c4..17068b38f5e 100644 --- a/sdk/cliproxy/auth/conductor_home_execution.go +++ b/sdk/cliproxy/auth/conductor_home_execution.go @@ -11,7 +11,31 @@ 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 unwrapHomeRetryRoundExhausted(err error) error { + var exhausted *homeRetryRoundExhaustedError + if errors.As(err, &exhausted) && exhausted != nil { + return exhausted.Unwrap() + } + return err +} + +func isHomeRetryHaltError(err error) bool { + var authErr *Error + if !errors.As(err, &authErr) || authErr == nil { + return false + } + switch authErr.Code { + case "home_unavailable", "auth_not_found", "auth_unavailable": + return true + } + return false +} + +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() } @@ -21,11 +45,15 @@ func (m *Manager) executeHome(ctx context.Context, providers []string, req clipr attempt := 0 retryRoundPending := false retryRoundWaited := false + var lastErr error for { - response, errExecute := m.executeHomeOnce(ctx, providers, req, opts, countTokens, maxRetryCredentials, &homeRetryLimit, attempt) + response, errExecute := m.executeHomeOnce(ctx, providers, req, opts, countTokens, maxRetryCredentials, &homeRetryLimit, attempt, tracker) if errExecute == nil { return response, nil } + if isHomeRetryHaltError(errExecute) && lastErr != nil { + errExecute = unwrapHomeRetryRoundExhausted(lastErr) + } if retryRoundPending { if wait, okWait := pendingHomeRetryRoundDelay(errExecute, maxWait, &homeRetryLimit, pinnedAuthIDFromMetadata(opts.Metadata) == ""); okWait && m.homeRetryAllowed(attempt-1, homeRetryLimit) { if retryRoundWaited { @@ -50,16 +78,17 @@ func (m *Manager) executeHome(ctx context.Context, providers []string, req clipr if errWait := waitForCooldown(ctx, wait, maxWait); errWait != nil { return cliproxyexecutor.Response{}, errWait } + lastErr = errExecute attempt++ retryRoundPending = true retryRoundWaited = false } } -func (m *Manager) executeHomeOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, countTokens bool, maxRetryCredentials int, homeRetryLimit *int, retryRounds ...int) (cliproxyexecutor.Response, error) { - retryRound := 0 - if len(retryRounds) > 0 { - retryRound = retryRounds[0] +func (m *Manager) executeHomeOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, countTokens bool, maxRetryCredentials int, homeRetryLimit *int, retryRound int, optionalTracker ...*routeAttemptTracker) (cliproxyexecutor.Response, error) { + var tracker *routeAttemptTracker + if len(optionalTracker) > 0 { + tracker = optionalTracker[0] } routeModel := authSelectionModelFromOptions(opts, req.Model) responseAlias := requestedModelAliasFromOptions(opts, routeModel) @@ -141,6 +170,7 @@ func (m *Manager) executeHomeOnce(ctx context.Context, providers []string, req c } lastErr = &Error{Code: "auth_not_found", Message: "no execution models available"} roundTiming.Observe(lastErr) + tracker.Record(auth, lastErr) continue } preparedAuth, errPrepare := m.prepareHomeRequestAuth(execCtx, selection.Executor, selection) @@ -151,6 +181,7 @@ func (m *Manager) executeHomeOnce(ctx context.Context, providers []string, req c return cliproxyexecutor.Response{}, errEnd } lastErr = errPrepare + tracker.Record(preparedAuth, errPrepare) roundTiming.Observe(lastErr) continue } @@ -264,6 +295,7 @@ func (m *Manager) executeHomeOnce(ctx context.Context, providers []string, req c action, okAction := matchRequestScopedErrorAction(preparedAuth, errExecute, m.runtimeConfigSnapshot()) applyRequestScopedActionToResult(action, okAction, &result) m.reportHomeResult(execCtx, result, preparedAuth) + tracker.Record(preparedAuth, errExecute) lastErr = errExecute if okAction { if isRequestScopedStop(action, okAction) { diff --git a/sdk/cliproxy/auth/conductor_invalid_key_sibling_test.go b/sdk/cliproxy/auth/conductor_invalid_key_sibling_test.go new file mode 100644 index 00000000000..e6492137288 --- /dev/null +++ b/sdk/cliproxy/auth/conductor_invalid_key_sibling_test.go @@ -0,0 +1,305 @@ +package auth + +import ( + "context" + "net/http" + "testing" + "time" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" +) + +func TestManager_InvalidAPIKey_PreservesLongerSiblingCooldown(t *testing.T) { + manager := NewManager(nil, nil, nil) + ctx := context.Background() + authID := "sibling-longer-cooldown-auth" + modelInvalid := "model-invalid" + modelLong := "model-long" + + reg := registry.GetGlobalRegistry() + reg.RegisterClient(authID, "openai", []*registry.ModelInfo{ + {ID: modelInvalid}, + {ID: modelLong}, + }) + t.Cleanup(func() { + reg.UnregisterClient(authID) + }) + + now := time.Now() + longerRecovery := now.Add(12 * time.Hour) + + auth := &Auth{ + ID: authID, + Provider: "openai", + Status: StatusActive, + ModelStates: map[string]*ModelState{ + modelInvalid: {Status: StatusActive}, + modelLong: { + Unavailable: true, + Status: StatusError, + StatusMessage: "model_not_supported", + NextRetryAfter: longerRecovery, + Quota: QuotaState{ + Exceeded: true, + Reason: "model_not_supported", + NextRecoverAt: longerRecovery, + }, + }, + }, + } + + if _, err := manager.Register(ctx, auth); err != nil { + t.Fatalf("register auth: %v", err) + } + + // Fail modelInvalid with invalid_api_key error + manager.MarkResult(ctx, Result{ + AuthID: authID, + Provider: "openai", + Model: modelInvalid, + Success: false, + Error: &Error{HTTPStatus: http.StatusUnauthorized, Code: "api_key_invalid", Message: "api key not valid"}, + }) + + updated, ok := manager.GetByID(authID) + if !ok || updated == nil { + t.Fatalf("auth not found after MarkResult") + } + + stateLong := updated.ModelStates[modelLong] + if stateLong == nil { + t.Fatalf("modelState for %s missing", modelLong) + } + + if !stateLong.Unavailable { + t.Errorf("modelLong Unavailable = false, want true") + } + if stateLong.Status != StatusError { + t.Errorf("modelLong Status = %v, want %v", stateLong.Status, StatusError) + } + if stateLong.StatusMessage != "invalid_api_key" { + t.Errorf("modelLong StatusMessage = %q, want invalid_api_key", stateLong.StatusMessage) + } + + // Sibling longer cooldown must NOT be shortened to 30 minutes. + if !stateLong.NextRetryAfter.Equal(longerRecovery) { + t.Errorf("modelLong NextRetryAfter = %v, want %v", stateLong.NextRetryAfter, longerRecovery) + } + if stateLong.Quota.Reason != "model_not_supported" { + t.Errorf("modelLong Quota.Reason = %q, want model_not_supported", stateLong.Quota.Reason) + } + if !stateLong.Quota.NextRecoverAt.Equal(longerRecovery) { + t.Errorf("modelLong Quota.NextRecoverAt = %v, want %v", stateLong.Quota.NextRecoverAt, longerRecovery) + } + + // After the 30-minute window elapses, the sibling must STILL be blocked by its 12-hour cooldown. + after30m := now.Add(35 * time.Minute) + if !stateLong.NextRetryAfter.After(after30m) { + t.Errorf("modelLong NextRetryAfter %v not after 35m mark %v (shortened cooldown)", stateLong.NextRetryAfter, after30m) + } +} + +func TestManager_InvalidAPIKey_AppliesCredentialBlockToSiblingWithNoPriorCooldown(t *testing.T) { + manager := NewManager(nil, nil, nil) + ctx := context.Background() + authID := "sibling-no-cooldown-auth" + modelInvalid := "model-invalid" + modelClean := "model-clean" + + reg := registry.GetGlobalRegistry() + reg.RegisterClient(authID, "openai", []*registry.ModelInfo{ + {ID: modelInvalid}, + {ID: modelClean}, + }) + t.Cleanup(func() { + reg.UnregisterClient(authID) + }) + + auth := &Auth{ + ID: authID, + Provider: "openai", + Status: StatusActive, + ModelStates: map[string]*ModelState{ + modelInvalid: {Status: StatusActive}, + modelClean: {Status: StatusActive}, + }, + } + + if _, err := manager.Register(ctx, auth); err != nil { + t.Fatalf("register auth: %v", err) + } + + now := time.Now() + manager.MarkResult(ctx, Result{ + AuthID: authID, + Provider: "openai", + Model: modelInvalid, + Success: false, + Error: &Error{HTTPStatus: http.StatusUnauthorized, Code: "api_key_invalid", Message: "api key not valid"}, + }) + + updated, ok := manager.GetByID(authID) + if !ok || updated == nil { + t.Fatalf("auth not found after MarkResult") + } + + stateClean := updated.ModelStates[modelClean] + if stateClean == nil { + t.Fatalf("modelState for %s missing", modelClean) + } + + if !stateClean.Unavailable { + t.Errorf("modelClean Unavailable = false, want true") + } + if stateClean.Status != StatusError { + t.Errorf("modelClean Status = %v, want %v", stateClean.Status, StatusError) + } + if stateClean.StatusMessage != "invalid_api_key" { + t.Errorf("modelClean StatusMessage = %q, want invalid_api_key", stateClean.StatusMessage) + } + if stateClean.NextRetryAfter.IsZero() || !stateClean.NextRetryAfter.After(now.Add(29*time.Minute)) { + t.Errorf("modelClean NextRetryAfter = %v, want ~30m from now", stateClean.NextRetryAfter) + } + if !stateClean.Quota.Exceeded || stateClean.Quota.Reason != "credential_quota" || stateClean.Quota.NextRecoverAt.IsZero() { + t.Errorf("modelClean Quota = %+v, want Exceeded: true, Reason: credential_quota", stateClean.Quota) + } +} + +func TestManager_InvalidAPIKey_ExtendsShorterSiblingCooldown(t *testing.T) { + manager := NewManager(nil, nil, nil) + ctx := context.Background() + authID := "sibling-short-cooldown-auth" + modelInvalid := "model-invalid" + modelShort := "model-short" + + reg := registry.GetGlobalRegistry() + reg.RegisterClient(authID, "openai", []*registry.ModelInfo{ + {ID: modelInvalid}, + {ID: modelShort}, + }) + t.Cleanup(func() { + reg.UnregisterClient(authID) + }) + + now := time.Now() + shortRecovery := now.Add(5 * time.Minute) + + auth := &Auth{ + ID: authID, + Provider: "openai", + Status: StatusActive, + ModelStates: map[string]*ModelState{ + modelInvalid: {Status: StatusActive}, + modelShort: { + Unavailable: true, + Status: StatusError, + StatusMessage: "short_error", + NextRetryAfter: shortRecovery, + Quota: QuotaState{ + Exceeded: true, + Reason: "quota", + NextRecoverAt: shortRecovery, + }, + }, + }, + } + + if _, err := manager.Register(ctx, auth); err != nil { + t.Fatalf("register auth: %v", err) + } + + manager.MarkResult(ctx, Result{ + AuthID: authID, + Provider: "openai", + Model: modelInvalid, + Success: false, + Error: &Error{HTTPStatus: http.StatusUnauthorized, Code: "api_key_invalid", Message: "api key not valid"}, + }) + + updated, ok := manager.GetByID(authID) + if !ok || updated == nil { + t.Fatalf("auth not found after MarkResult") + } + + stateShort := updated.ModelStates[modelShort] + if stateShort == nil { + t.Fatalf("modelState for %s missing", modelShort) + } + + if !stateShort.Unavailable { + t.Errorf("modelShort Unavailable = false, want true") + } + if stateShort.Status != StatusError { + t.Errorf("modelShort Status = %v, want %v", stateShort.Status, StatusError) + } + if stateShort.StatusMessage != "invalid_api_key" { + t.Errorf("modelShort StatusMessage = %q, want invalid_api_key", stateShort.StatusMessage) + } + if !stateShort.NextRetryAfter.After(shortRecovery) || !stateShort.NextRetryAfter.After(now.Add(29*time.Minute)) { + t.Errorf("modelShort NextRetryAfter = %v, want extended to ~30m (> %v)", stateShort.NextRetryAfter, shortRecovery) + } + if !stateShort.Quota.Exceeded || stateShort.Quota.Reason != "credential_quota" || !stateShort.Quota.NextRecoverAt.After(shortRecovery) { + t.Errorf("modelShort Quota = %+v, want extended credential_quota (> %v)", stateShort.Quota, shortRecovery) + } +} + +func TestManager_InvalidAPIKey_PreservesLongerAuthLevelCooldown(t *testing.T) { + manager := NewManager(nil, nil, nil) + ctx := context.Background() + authID := "auth-longer-cooldown" + modelInvalid := "model-invalid" + + reg := registry.GetGlobalRegistry() + reg.RegisterClient(authID, "openai", []*registry.ModelInfo{ + {ID: modelInvalid}, + }) + t.Cleanup(func() { + reg.UnregisterClient(authID) + }) + + now := time.Now() + longerRecovery := now.Add(12 * time.Hour) + + auth := &Auth{ + ID: authID, + Provider: "openai", + Status: StatusError, + Unavailable: true, + NextRetryAfter: longerRecovery, + Quota: QuotaState{ + Exceeded: true, + Reason: "credential_quota", + NextRecoverAt: longerRecovery, + }, + ModelStates: map[string]*ModelState{ + modelInvalid: {Status: StatusActive}, + }, + } + + if _, err := manager.Register(ctx, auth); err != nil { + t.Fatalf("register auth: %v", err) + } + + manager.MarkResult(ctx, Result{ + AuthID: authID, + Provider: "openai", + Model: modelInvalid, + Success: false, + Error: &Error{HTTPStatus: http.StatusUnauthorized, Code: "api_key_invalid", Message: "api key not valid"}, + }) + + updated, ok := manager.GetByID(authID) + if !ok || updated == nil { + t.Fatalf("auth not found after MarkResult") + } + + if !updated.NextRetryAfter.Equal(longerRecovery) { + t.Errorf("auth NextRetryAfter = %v, want %v", updated.NextRetryAfter, longerRecovery) + } + if updated.Quota.Reason != "credential_quota" { + t.Errorf("auth Quota.Reason = %q, want credential_quota", updated.Quota.Reason) + } + if !updated.Quota.NextRecoverAt.Equal(longerRecovery) { + t.Errorf("auth Quota.NextRecoverAt = %v, want %v", updated.Quota.NextRecoverAt, longerRecovery) + } +} diff --git a/sdk/cliproxy/auth/conductor_overrides_test.go b/sdk/cliproxy/auth/conductor_overrides_test.go index e8e635afb56..dcb2501f653 100644 --- a/sdk/cliproxy/auth/conductor_overrides_test.go +++ b/sdk/cliproxy/auth/conductor_overrides_test.go @@ -885,6 +885,163 @@ func TestManagerExecuteStream_AntigravityInvalidGrantFallsBackAndSuspendsAuth(t } } +// 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_ModelSupportBadRequestFallsBackAndSuspendsAuth(t *testing.T) { m := NewManager(nil, nil, nil) executor := &authFallbackExecutor{ @@ -1381,7 +1538,13 @@ func TestManager_Execute_DisableCooling_DoesNotBlackoutAfter429RetryAfter(t *tes } } -func TestManager_Execute_DisableCooling_RetriesAfter429RetryAfter(t *testing.T) { +// Retry config (3 attempts, 100ms cap) plus a 5ms Retry-After gives a genuine +// nonzero retry opportunity here: shouldRetryAfterError returns (5ms, true) for +// attempts 0-2. Pre-fix the same auth was re-picked on every outer retry +// (4 executor calls); post-fix the request-scoped exclusion of failed auths +// stops re-invocation, so a single auth with a retryable 429 is executed +// exactly once per request. +func TestManager_Execute_DisableCooling_NoReinvokeOnRetryable429(t *testing.T) { prev := quotaCooldownDisabled.Load() quotaCooldownDisabled.Store(false) t.Cleanup(func() { quotaCooldownDisabled.Store(prev) }) @@ -1427,8 +1590,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)) } } @@ -2486,3 +2649,203 @@ func TestManager_MarkResult_RequestFaultBodyDoesNotCooldownModelOrAuth(t *testin t.Fatalf("expected real 401 authentication error to set model cooldown state, got %#v", state) } } + +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() + t.Logf("calls: %v", calls) + 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 cc9dc751fc1..e0ed812d49b 100644 --- a/sdk/cliproxy/auth/conductor_selection.go +++ b/sdk/cliproxy/auth/conductor_selection.go @@ -788,7 +788,7 @@ func credentialRetryRoundStateEligible(lastErr *Error, quotaExceeded bool) bool return isCredentialRetryRoundStatus(statusCodeFromResult(lastErr)) } -func (m *Manager) closestCooldownWait(providers []string, model string, attempt int, eligibility authSelectionEligibility, pinnedAuthID string, defaultRequestRetry int) (time.Duration, bool) { +func (m *Manager) closestCooldownWait(providers []string, model string, attempt int, eligibility authSelectionEligibility, pinnedAuthID string, defaultRequestRetry int, excluded map[string]struct{}) (time.Duration, bool) { if m == nil || len(providers) == 0 { return 0, false } @@ -828,6 +828,9 @@ func (m *Manager) closestCooldownWait(providers []string, model string, attempt if model != "" && !m.authSupportsRouteModel(registryRef, auth, model) { continue } + if _, ok := excluded[auth.ID]; ok { + continue + } effectiveRetry := effectiveRequestRetryLimit(auth, defaultRequestRetry) if attempt >= effectiveRetry { continue @@ -855,7 +858,7 @@ func (m *Manager) closestCooldownWait(providers []string, model string, attempt return minWait, found } -func (m *Manager) retryAllowed(attempt int, providers []string, model string, eligibility authSelectionEligibility, pinnedAuthID string, defaultRequestRetry int) bool { +func (m *Manager) retryAllowed(attempt int, providers []string, model string, eligibility authSelectionEligibility, pinnedAuthID string, defaultRequestRetry int, excluded map[string]struct{}) bool { if m == nil || attempt < 0 || len(providers) == 0 { return false } @@ -895,6 +898,9 @@ func (m *Manager) retryAllowed(attempt int, providers []string, model string, el if model != "" && !m.authSupportsRouteModel(registryRef, auth, model) { continue } + if _, ok := excluded[auth.ID]; ok { + continue + } effectiveRetry := effectiveRequestRetryLimit(auth, defaultRequestRetry) if attempt >= effectiveRetry { continue @@ -971,10 +977,11 @@ func (m *Manager) shouldRetryAfterErrorWithHomeRetryLimit(ctx context.Context, o } eligibility := authSelectionEligibilityForRequest(ctx, opts) pinnedAuthID := pinnedAuthIDFromMetadata(opts.Metadata) - if !isCredentialRetryRoundStatus(status) || !m.retryAllowed(attempt, providers, model, eligibility, pinnedAuthID, defaultRequestRetry) { + excluded := extractExcludedAuthIDs(opts.Metadata) + if !isCredentialRetryRoundStatus(status) || !m.retryAllowed(attempt, providers, model, eligibility, pinnedAuthID, defaultRequestRetry, excluded) { return 0, false } - wait, found := m.closestCooldownWait(providers, model, attempt, eligibility, pinnedAuthID, defaultRequestRetry) + wait, found := m.closestCooldownWait(providers, model, attempt, eligibility, pinnedAuthID, defaultRequestRetry, excluded) if found { if wait > 0 && (maxWait <= 0 || wait > maxWait) { return 0, false @@ -1222,6 +1229,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 @@ -1539,6 +1549,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 6efb2bd6781..a1f14ad2ecc 100644 --- a/sdk/cliproxy/auth/conductor_stream.go +++ b/sdk/cliproxy/auth/conductor_stream.go @@ -2,13 +2,57 @@ package auth import ( "context" + "fmt" "net/http" "strings" + "sync" + "sync/atomic" "time" + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" ) +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 @@ -82,11 +126,18 @@ 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, requestPayloads ...[]byte) ([]cliproxyexecutor.StreamChunk, bool, error) { if ch == nil { return nil, true, nil } buffered := make([]cliproxyexecutor.StreamChunk, 0, 1) + var bootstrap streamBootstrapState + for _, p := range requestPayloads { + if n := ExtractExpectedChoices(p); n > 1 { + bootstrap.setExpectedChoices(n) + break + } + } for { var ( chunk cliproxyexecutor.StreamChunk @@ -102,24 +153,50 @@ func readStreamBootstrap(ctx context.Context, ch <-chan cliproxyexecutor.StreamC chunk, ok = <-ch } if !ok { + // A final frame without its blank-line delimiter is only parsed by + // finish(): without it a provider error carried by that last frame stays + // pending and the stream looks like a clean close. + bootstrap.finish() + if err := bootstrap.streamError(); err != nil && !bootstrap.hasMeaningfulOutput() { + return nil, false, err + } return buffered, true, nil } if chunk.Err != nil { + if bootstrap.hasMeaningfulOutput() { + buffered = append(buffered, chunk) + return buffered, false, nil + } return nil, false, chunk.Err } buffered = append(buffered, chunk) - if len(chunk.Payload) > 0 { + if bootstrap.observe(chunk.Payload) { return buffered, false, nil } + if err := bootstrap.streamError(); err != nil { + if bootstrap.hasMeaningfulOutput() { + return buffered, false, nil + } + return nil, false, err + } + 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, opts cliproxyexecutor.Options) *cliproxyexecutor.StreamResult { +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, opts cliproxyexecutor.Options, cleanups ...func()) *cliproxyexecutor.StreamResult { out := make(chan cliproxyexecutor.StreamChunk) streamStart := time.Now() go func() { defer close(out) + for _, cleanup := range cleanups { + if cleanup != nil { + defer cleanup() + } + } var failed bool + var errorDetector streamPayloadErrorDetector forward := true var rewriter *StreamRewriter if aliasResult.ForceMapping && strings.TrimSpace(aliasResult.OriginalAlias) != "" { @@ -136,6 +213,18 @@ func (m *Manager) wrapStreamResult(ctx context.Context, auth *Auth, provider, re applyRequestScopedActionToResult(action, okAction, &result) m.recordExecutionResult(ctx, result, auth, ephemeralResult) } + if !failed && len(chunk.Payload) > 0 { + if streamErr := errorDetector.Observe(chunk.Payload); streamErr != nil { + failed = true + entry := logEntryWithRequestID(ctx) + warnLogUpstreamFailure(ctx, entry, provider, resultModel, auth, time.Since(streamStart), streamErr) + rerr := resultErrorFromError(streamErr) + action, okAction := matchRequestScopedErrorAction(auth, streamErr, m.runtimeConfigSnapshot()) + result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: rerr, Options: opts} + applyRequestScopedActionToResult(action, okAction, &result) + m.recordExecutionResult(ctx, result, auth, ephemeralResult) + } + } if !forward { return false } @@ -190,6 +279,18 @@ func (m *Manager) wrapStreamResult(ctx context.Context, auth *Auth, provider, re return } } + if !failed { + if streamErr := errorDetector.Finish(); streamErr != nil { + failed = true + entry := logEntryWithRequestID(ctx) + warnLogUpstreamFailure(ctx, entry, provider, resultModel, auth, time.Since(streamStart), streamErr) + rerr := resultErrorFromError(streamErr) + action, okAction := matchRequestScopedErrorAction(auth, streamErr, m.runtimeConfigSnapshot()) + result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: rerr, Options: opts} + applyRequestScopedActionToResult(action, okAction, &result) + m.recordExecutionResult(ctx, result, auth, ephemeralResult) + } + } if !failed && (ephemeralResult || claudeOAuthRequestCancellation(ctx, auth, nil) == nil) { m.recordExecutionResult(ctx, Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: true, Options: opts}, auth, ephemeralResult) } @@ -216,6 +317,29 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi _, didRefreshOnUnauthorized = unauthorizedRefreshTried[auth.ID] } for idx, execModel := range execModels { + ttftTimeout := m.streamFirstChunkTimeout(opts) + attemptCtx, cancelAttempt := context.WithCancel(ctx) + var timer *time.Timer + var timedOut atomic.Bool + var attemptMu sync.Mutex + var attemptSeq uint64 + + stopTTFT := func() { + if timer != nil { + timer.Stop() + } + attemptMu.Lock() + attemptSeq++ + attemptMu.Unlock() + } + + checkTTFTErr := func(err error) error { + if timedOut.Load() { + return newTTFTTimeoutError(ttftTimeout) + } + return err + } + resultModel := m.stateModelForExecution(auth, routeModel, execModel, pooled) execReq := req execReq.Model = execModel @@ -226,23 +350,64 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi var errIntercept error execReq, execOpts, errIntercept = applyRequestAfterAuthInterceptor(ctx, executor, provider, execReq, execOpts, requestedModelAliasFromOptions(execOpts, routeModel)) if errIntercept != nil { + stopTTFT() + cancelAttempt() return nil, errIntercept } if executionModel == "" { execReq = attachResolvedAPIKeyModelInfo(routing, execReq, auth, routeModel, execModel) } if errCtx := ctx.Err(); errCtx != nil { + stopTTFT() + cancelAttempt() return nil, errCtx } + // Arm the TTFT timer 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. + armTTFT := func() { + if ttftTimeout > 0 { + currentSeq := attemptSeq + currentCancel := cancelAttempt + timer = time.AfterFunc(ttftTimeout, func() { + attemptMu.Lock() + defer attemptMu.Unlock() + if currentSeq != attemptSeq { + return + } + timedOut.Store(true) + currentCancel() + }) + } + } + armTTFT() + // The unauthorized-refresh retries below re-execute behind a credential + // refresh, which may consume the whole TTFT budget (or fire the timer + // and cancel attemptCtx). Restart the first-chunk window on a fresh + // attempt context so the refreshed upstream request gets a full budget. + restartAttempt := func() { + stopTTFT() + attemptMu.Lock() + cancelAttempt() + attemptCtx, cancelAttempt = context.WithCancel(ctx) + timedOut.Store(false) + attemptMu.Unlock() + armTTFT() + } entry := logEntryWithRequestID(ctx) startStream := time.Now() - streamResult, errStream := executor.ExecuteStream(ctx, auth, execReq, execOpts) + streamResult, errStream := executor.ExecuteStream(attemptCtx, auth, execReq, execOpts) durationStream := time.Since(startStream) if errStream != nil { if errCtx := ctx.Err(); errCtx != nil { + stopTTFT() + cancelAttempt() return nil, errCtx } + errStream = checkTTFTErr(errStream) if allowRetry { + stopTTFT() alreadyTried := didRefreshOnUnauthorized willAttemptHomeRefresh := ephemeralResult && !alreadyTried && auth != nil && auth.AuthKind() == AuthKindOAuth && isUnauthorizedError(errStream) refreshed, okRefresh, errRefresh := m.tryRefreshExecutionAuthAfterUnauthorized(ctx, executor, auth, errStream, alreadyTried, ephemeralResult) @@ -260,12 +425,22 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi m.replaceHomeExecutionLifecycleAuth(execOpts.ExecutionLifecycle, auth) publishSelectedAuthMetadata(execOpts.Metadata, auth) didRefreshOnUnauthorized = true + restartAttempt() + if streamResult != nil { + discardStreamChunks(streamResult.Chunks) + } startRetry := time.Now() - streamResult, errStream = executor.ExecuteStream(ctx, auth, execReq, execOpts) + streamResult, errStream = executor.ExecuteStream(attemptCtx, auth, execReq, execOpts) durationRetry := time.Since(startRetry) + errStream = checkTTFTErr(errStream) if errStream != nil { warnLogUpstreamFailure(ctx, entry, provider, execModel, auth, durationRetry, errStream) if errCtx := ctx.Err(); errCtx != nil { + stopTTFT() + cancelAttempt() + if streamResult != nil { + discardStreamChunks(streamResult.Chunks) + } return nil, errCtx } } @@ -278,11 +453,22 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi } if !ephemeralResult { if errCancel := claudeOAuthRequestCancellation(ctx, auth, errStream); errCancel != nil { + stopTTFT() + cancelAttempt() + if streamResult != nil { + discardStreamChunks(streamResult.Chunks) + } return nil, errCancel } } streamResult, errStream = validateStreamResult(streamResult, errStream) if errStream != nil { + stopTTFT() + cancelAttempt() + if streamResult != nil { + discardStreamChunks(streamResult.Chunks) + } + errStream = checkTTFTErr(errStream) rerr := resultErrorFromError(errStream) action, okAction := matchRequestScopedErrorAction(auth, errStream, m.runtimeConfigSnapshot()) result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: rerr, Options: execOpts} @@ -311,14 +497,19 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi } continue } + stopTTFT() - buffered, closed, bootstrapErr := readStreamBootstrap(ctx, streamResult.Chunks) + buffered, closed, bootstrapErr := readStreamBootstrap(attemptCtx, streamResult.Chunks, execReq.Payload, execOpts.OriginalRequest) if bootstrapErr != nil { if errCtx := ctx.Err(); errCtx != nil { + stopTTFT() + cancelAttempt() discardStreamChunks(streamResult.Chunks) return nil, errCtx } + bootstrapErr = checkTTFTErr(bootstrapErr) if allowRetry { + stopTTFT() alreadyTried := didRefreshOnUnauthorized willAttemptHomeRefresh := ephemeralResult && !alreadyTried && auth != nil && auth.AuthKind() == AuthKindOAuth && isUnauthorizedError(bootstrapErr) refreshed, okRefresh, errRefresh := m.tryRefreshExecutionAuthAfterUnauthorized(ctx, executor, auth, bootstrapErr, alreadyTried, ephemeralResult) @@ -339,11 +530,19 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi m.replaceHomeExecutionLifecycleAuth(execOpts.ExecutionLifecycle, auth) publishSelectedAuthMetadata(execOpts.Metadata, auth) didRefreshOnUnauthorized = true + restartAttempt() startRetry := time.Now() - retryStream, retryErr := executor.ExecuteStream(ctx, auth, execReq, execOpts) + retryStream, retryErr := executor.ExecuteStream(attemptCtx, auth, execReq, execOpts) retryStream, retryErr = validateStreamResult(retryStream, retryErr) + stopTTFT() + retryErr = checkTTFTErr(retryErr) if retryErr != nil { + if retryStream != nil { + discardStreamChunks(retryStream.Chunks) + } if errCtx := ctx.Err(); errCtx != nil { + stopTTFT() + cancelAttempt() return nil, errCtx } bootstrapErr = retryErr @@ -351,7 +550,8 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi streamResult = &cliproxyexecutor.StreamResult{} } else { streamResult = retryStream - buffered, closed, bootstrapErr = readStreamBootstrap(ctx, streamResult.Chunks) + buffered, closed, bootstrapErr = readStreamBootstrap(attemptCtx, streamResult.Chunks, execReq.Payload, execOpts.OriginalRequest) + bootstrapErr = checkTTFTErr(bootstrapErr) if bootstrapErr != nil { warnLogUpstreamFailure(ctx, entry, provider, execModel, auth, time.Since(startRetry), bootstrapErr) } @@ -365,11 +565,16 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi } if !ephemeralResult { if errCancel := claudeOAuthRequestCancellation(ctx, auth, bootstrapErr); errCancel != nil { + stopTTFT() + cancelAttempt() discardStreamChunks(streamResult.Chunks) return nil, errCancel } } if bootstrapErr != nil { + stopTTFT() + cancelAttempt() + bootstrapErr = checkTTFTErr(bootstrapErr) action, okAction := matchRequestScopedErrorAction(auth, bootstrapErr, m.runtimeConfigSnapshot()) if okAction { rerr := resultErrorFromError(bootstrapErr) @@ -427,11 +632,25 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi return nil, newStreamBootstrapError(bootstrapErr, streamResult.Headers) } - if closed && len(buffered) == 0 { - emptyErr := &Error{Code: "empty_stream", Message: "upstream stream closed before first payload", Retryable: true} + 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)) { + stopTTFT() + cancelAttempt() + emptyErr := errEmptyCompletion + if payloadBytes == 0 { + emptyErr = &Error{Code: "empty_stream", Message: "upstream stream closed before first payload", Retryable: true} + } warnLogUpstreamFailure(ctx, entry, provider, execModel, auth, time.Since(startStream), 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 @@ -439,14 +658,17 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi return nil, newStreamBootstrapError(emptyErr, streamResult.Headers) } + stopTTFT() + 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, execOpts), nil + return m.wrapStreamResult(ctx, auth.Clone(), provider, resultModel, streamResult.Headers, buffered, remaining, attemptAliasResult, ephemeralResult, execOpts, cancelAttempt), 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 00000000000..d76bcae8acd --- /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_eof_test.go b/sdk/cliproxy/auth/conductor_stream_eof_test.go new file mode 100644 index 00000000000..54cf0afc1ca --- /dev/null +++ b/sdk/cliproxy/auth/conductor_stream_eof_test.go @@ -0,0 +1,36 @@ +package auth + +import ( + "context" + "strings" + "testing" + + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +// TestReadStreamBootstrapFinalizesDetectorAtEOF covers an upstream that closes the +// channel right after an SSE error event whose data line is newline-terminated but +// never followed by the blank separator line. flushData() only runs on that blank +// line or from finish(), so without finalizing the bootstrap state the provider +// error stays buffered, the bootstrap reports a clean close, and the caller gets an +// empty stream instead of a routable failure it can fail over on. +func TestReadStreamBootstrapFinalizesDetectorAtEOF(t *testing.T) { + ch := make(chan cliproxyexecutor.StreamChunk, 2) + ch <- cliproxyexecutor.StreamChunk{Payload: []byte("data: {\"choices\":[{\"delta\":{\"role\":\"assistant\"}}]}\n\n")} + ch <- cliproxyexecutor.StreamChunk{Payload: []byte("event: error\ndata: {\"error\":{\"code\":\"invalid_api_key\",\"message\":\"invalid api key\"}}\n")} + close(ch) + + buffered, closed, err := readStreamBootstrap(context.Background(), ch) + if err == nil { + t.Fatalf("readStreamBootstrap() error = nil, want the in-band provider error (closed=%v, buffered=%d)", closed, len(buffered)) + } + if !strings.Contains(err.Error(), "invalid api key") { + t.Fatalf("readStreamBootstrap() error = %v, want the invalid api key provider error", err) + } + if closed { + t.Fatal("closed = true, want false so the caller can fail over") + } + if len(buffered) != 0 { + t.Fatalf("len(buffered) = %d, want 0 when the provider error propagates", len(buffered)) + } +} 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 00000000000..5ba49223119 --- /dev/null +++ b/sdk/cliproxy/auth/conductor_stream_ttft_test.go @@ -0,0 +1,474 @@ +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 is a regression guard +// for the codex P2 finding on PR #4881: the TTFT timer used to stay armed +// across the unauthorized-refresh retry, so a refresh slower than the +// first-chunk budget left the retried ExecuteStream with an already-canceled +// context, surfacing a spurious 504 although the refreshed upstream request +// never ran. The retry must restart the window on a fresh attempt 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) + } +} + +// ttftRaceProbeExecutor is designed to trigger the race where timer 1 fires +// right around restartAttempt during unauthorized refresh. +type ttftRaceProbeExecutor struct { + calls atomic.Int32 + refreshCalls atomic.Int32 + retryCtxErr atomic.Value // error +} + +func (e *ttftRaceProbeExecutor) Identifier() string { return "gemini" } + +func (e *ttftRaceProbeExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} + +func (e *ttftRaceProbeExecutor) ExecuteStream(ctx context.Context, _ *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + call := e.calls.Add(1) + if call == 1 { + return nil, errors.New("upstream returned status 401") + } + if err := ctx.Err(); err != nil { + e.retryCtxErr.Store(err) + return nil, err + } + chunks := make(chan cliproxyexecutor.StreamChunk, 1) + chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("data: success\n\n")} + close(chunks) + return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil +} + +func (e *ttftRaceProbeExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} + +func (e *ttftRaceProbeExecutor) Refresh(_ context.Context, a *Auth) (*Auth, error) { + e.refreshCalls.Add(1) + time.Sleep(5 * time.Millisecond) + return a, nil +} + +func (e *ttftRaceProbeExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +// TestStreamTTFTCallbackBoundToAttemptAcrossRefreshRace tests that an in-flight +// TTFT timeout callback from the first attempt does not cancel the refreshed +// attempt or mark it as timed out. +func TestStreamTTFTCallbackBoundToAttemptAcrossRefreshRace(t *testing.T) { + for i := 0; i < 50; i++ { + manager := NewManager(nil, nil, nil) + auth := &Auth{ + ID: "auth-oauth-race", + 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"}}) + manager.RefreshSchedulerEntry(auth.ID) + exec := &ttftRaceProbeExecutor{} + manager.RegisterExecutor(exec) + + opts := cliproxyexecutor.Options{ + Metadata: map[string]any{"stream_first_chunk_timeout_ms": 5}, + } + _, err := manager.ExecuteStream(context.Background(), []string{"gemini"}, cliproxyexecutor.Request{Model: "test-model"}, opts) + reg.UnregisterClient(auth.ID) + + if err != nil { + t.Fatalf("iteration %d: ExecuteStream() error = %v, retryCtxErr = %v", i, err, exec.retryCtxErr.Load()) + } + } +} + +// ttftNonTimeoutErrExecutor returns 401 on first call, refreshes, and then +// returns a 500 error on the second call. +type ttftNonTimeoutErrExecutor struct { + calls atomic.Int32 + refreshCalls atomic.Int32 +} + +func (e *ttftNonTimeoutErrExecutor) Identifier() string { return "gemini" } + +func (e *ttftNonTimeoutErrExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} + +func (e *ttftNonTimeoutErrExecutor) ExecuteStream(ctx context.Context, _ *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + call := e.calls.Add(1) + if call == 1 { + return nil, errors.New("upstream returned status 401") + } + return nil, &Error{Code: "internal_error", Message: "upstream 500 internal server error", HTTPStatus: 500} +} + +func (e *ttftNonTimeoutErrExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} + +func (e *ttftNonTimeoutErrExecutor) Refresh(_ context.Context, a *Auth) (*Auth, error) { + e.refreshCalls.Add(1) + time.Sleep(5 * time.Millisecond) + return a, nil +} + +func (e *ttftNonTimeoutErrExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +// TestStreamTTFTCallbackDoesNotMarkSubsequentNonTimeoutErrorAs504 tests that a +// stale callback from attempt 1 does not reset timedOut to true and turn a +// non-timeout error on attempt 2 into a 504 stream_first_chunk_timeout. +func TestStreamTTFTCallbackDoesNotMarkSubsequentNonTimeoutErrorAs504(t *testing.T) { + for i := 0; i < 20; i++ { + manager := NewManager(nil, nil, nil) + auth := &Auth{ + ID: "auth-oauth-500", + 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"}}) + manager.RefreshSchedulerEntry(auth.ID) + exec := &ttftNonTimeoutErrExecutor{} + manager.RegisterExecutor(exec) + + opts := cliproxyexecutor.Options{ + Metadata: map[string]any{"stream_first_chunk_timeout_ms": 5}, + } + _, err := manager.ExecuteStream(context.Background(), []string{"gemini"}, cliproxyexecutor.Request{Model: "test-model"}, opts) + reg.UnregisterClient(auth.ID) + + if err == nil { + t.Fatalf("iteration %d: expected error, got nil", i) + } + if strings.Contains(err.Error(), "stream_first_chunk_timeout") || statusCodeFromError(err) == http.StatusGatewayTimeout { + t.Fatalf("iteration %d: stale TTFT callback marked attempt 2 as 504: %v", i, err) + } + if !strings.Contains(err.Error(), "internal_error") && !strings.Contains(err.Error(), "500") { + t.Fatalf("iteration %d: expected internal_error 500, got: %v", i, err) + } + } +} + +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) + } +} diff --git a/sdk/cliproxy/auth/conductor_unauthorized_refresh_test.go b/sdk/cliproxy/auth/conductor_unauthorized_refresh_test.go index 37451925076..b9bad3297c4 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 new file mode 100644 index 00000000000..e48ef4cee97 --- /dev/null +++ b/sdk/cliproxy/auth/empty_completion.go @@ -0,0 +1,2584 @@ +package auth + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "io" + "math" + "net/http" + "strconv" + "strings" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/clienterror" + 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 +// rotates to the next auth/model. +var errEmptyCompletion = &Error{ + Code: "empty_completion", + Message: "upstream returned an empty completion", + Retryable: true, + HTTPStatus: http.StatusServiceUnavailable, +} + +// maxStreamBootstrapBytes bounds how much metadata a stream can accumulate +// before the conductor conservatively forwards it. Empty-completion detection +// must never create an unbounded pre-output buffer. +const maxStreamBootstrapBytes = 1 << 20 + +// openAIChunk is the minimal OpenAI-style SSE/JSON shape used to detect empty +// completions. +type openAIChunk struct { + Choices []struct { + Index *int `json:"index"` + Text string `json:"text"` + Delta struct { + Content string `json:"content"` + ReasoningContent string `json:"reasoning_content"` + Reasoning string `json:"reasoning"` + 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"` + Reasoning string `json:"reasoning"` + 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 *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 +} + +// hasMeaningfulGeminiMediaPayload reports whether a Gemini media part contains usable content. +// Matching the translator semantics, inlineData counts only when data is non-blank and fileData +// only when fileUri is non-blank; a scaffold object with only a mimeType carries no media. +func hasMeaningfulGeminiMediaPayload(inlineData, fileData json.RawMessage) bool { + if len(bytes.TrimSpace(inlineData)) > 0 { + var v struct { + Data string `json:"data"` + } + if json.Unmarshal(inlineData, &v) == nil && strings.TrimSpace(v.Data) != "" { + return true + } + } + if len(bytes.TrimSpace(fileData)) > 0 { + var v struct { + FileURI string `json:"fileUri"` + } + if json.Unmarshal(fileData, &v) == nil && strings.TrimSpace(v.FileURI) != "" { + 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) +} + +type geminiGroundingChunk struct { + Web *struct { + URI string `json:"uri"` + Title string `json:"title"` + } `json:"web"` + RetrievedContext *struct { + URI string `json:"uri"` + Title string `json:"title"` + Text string `json:"text"` + } `json:"retrievedContext"` +} + +type geminiGroundingMetadata struct { + WebSearchQueries []string `json:"webSearchQueries"` + GroundingChunks []geminiGroundingChunk `json:"groundingChunks"` + SearchEntryPoint *struct { + RenderedContent string `json:"renderedContent"` + } `json:"searchEntryPoint"` + RetrievalQueries []string `json:"retrievalQueries"` +} + +func hasMeaningfulGroundingMetadata(raw json.RawMessage) bool { + trimmed := bytes.TrimSpace(raw) + if len(trimmed) == 0 || bytes.Equal(trimmed, []byte("null")) || bytes.Equal(trimmed, []byte("{}")) || bytes.Equal(trimmed, []byte("[]")) { + return false + } + var meta geminiGroundingMetadata + if err := json.Unmarshal(trimmed, &meta); err == nil { + for _, q := range meta.WebSearchQueries { + if strings.TrimSpace(q) != "" { + return true + } + } + for _, q := range meta.RetrievalQueries { + if strings.TrimSpace(q) != "" { + return true + } + } + for _, chunk := range meta.GroundingChunks { + if chunk.Web != nil { + if strings.TrimSpace(chunk.Web.URI) != "" || strings.TrimSpace(chunk.Web.Title) != "" { + return true + } + } + if chunk.RetrievedContext != nil { + if strings.TrimSpace(chunk.RetrievedContext.URI) != "" || + strings.TrimSpace(chunk.RetrievedContext.Title) != "" || + strings.TrimSpace(chunk.RetrievedContext.Text) != "" { + return true + } + } + } + if meta.SearchEntryPoint != nil && strings.TrimSpace(meta.SearchEntryPoint.RenderedContent) != "" { + return true + } + } + var generic map[string]any + if err := json.Unmarshal(trimmed, &generic); err == nil { + for k, v := range generic { + if k == "groundingChunks" || k == "webSearchQueries" || k == "retrievalQueries" || k == "searchEntryPoint" { + continue + } + switch val := v.(type) { + case nil: + case string: + if strings.TrimSpace(val) != "" { + return true + } + case map[string]any: + if len(val) > 0 { + for _, mv := range val { + if s, ok := mv.(string); ok && strings.TrimSpace(s) != "" { + return true + } + } + } + case []any: + if len(val) > 0 { + for _, ev := range val { + if s, ok := ev.(string); ok && strings.TrimSpace(s) != "" { + return true + } + } + } + default: + return true + } + } + } + return false +} + +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.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 { + 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 { + Type string `json:"type"` + StopReason *string `json:"stop_reason"` + Content []claudeContentBlock `json:"content"` + Usage *struct { + 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 *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"` + 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"` + ExecutableCode json.RawMessage `json:"executableCode"` + CodeExecutionResult json.RawMessage `json:"codeExecutionResult"` + ThoughtSignature string `json:"thoughtSignature"` + Thought_Signature string `json:"thought_signature"` +} + +type geminiCandidate struct { + Index *int `json:"index"` + Content *struct { + Parts []geminiPart `json:"parts"` + } `json:"content"` + FinishReason *string `json:"finishReason"` + GroundingMetadata json.RawMessage `json:"groundingMetadata"` +} + +type geminiUsageMetadata struct { + CandidatesTokenCount *tokenCount `json:"candidatesTokenCount"` +} + +type geminiPromptFeedback struct { + BlockReason string `json:"blockReason"` +} + +type geminiChunk struct { + Candidates []geminiCandidate `json:"candidates"` + UsageMetadata *geminiUsageMetadata `json:"usageMetadata"` + PromptFeedback *geminiPromptFeedback `json:"promptFeedback"` + GroundingMetadata json.RawMessage `json:"groundingMetadata"` + Response *struct { + Candidates []geminiCandidate `json:"candidates"` + UsageMetadata *geminiUsageMetadata `json:"usageMetadata"` + PromptFeedback *geminiPromptFeedback `json:"promptFeedback"` + GroundingMetadata json.RawMessage `json:"groundingMetadata"` + } `json:"response"` +} + +// openAIResponseUsage is the usage block of the OpenAI Responses-API shape +// (used by codex/xai executors). +type openAIResponseUsage struct { + OutputTokens *tokenCount `json:"output_tokens"` +} + +type openAIResponseContentPart struct { + Type string `json:"type"` + Text string `json:"text"` + Refusal string `json:"refusal"` + Annotations []json.RawMessage `json:"annotations"` +} + +type openAIResponseOutputItem struct { + 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"` + Action json.RawMessage `json:"action"` + Results json.RawMessage `json:"results"` + Content []openAIResponseContentPart `json:"content"` + EncryptedContent string `json:"encrypted_content"` + Summary json.RawMessage `json:"summary"` +} + +type openAIResponseObject struct { + Status string `json:"status"` + Output json.RawMessage `json:"output"` + Usage *openAIResponseUsage `json:"usage"` +} + +type openAIResponsePart struct { + Type string `json:"type"` + Text string `json:"text"` + EncryptedContent string `json:"encrypted_content"` + Annotations []json.RawMessage `json:"annotations"` +} + +type openAIResponseChunk struct { + Type string `json:"type"` + Object string `json:"object"` + Status string `json:"status"` + Output json.RawMessage `json:"output"` + Item json.RawMessage `json:"item"` + Part json.RawMessage `json:"part"` + Usage *openAIResponseUsage `json:"usage"` + Response *openAIResponseObject `json:"response"` + Delta string `json:"delta"` + Text string `json:"text"` + Arguments string `json:"arguments"` +} + +// openAIResponseEventTypes is the conservative set of Responses-API streamed +// event types we recognize. Unknown/partial sub-shapes are left unrecognized so +// the stream is forwarded rather than judged empty. +var openAIResponseEventTypes = map[string]bool{ + "response.created": true, + "response.in_progress": true, + "response.completed": true, + "response.incomplete": true, + "response.failed": true, + "response.output_item.added": true, + "response.output_item.done": true, + "response.content_part.added": true, + "response.content_part.done": true, + "response.output_text.delta": true, + "response.output_text.done": true, + "response.reasoning_summary_part.added": true, + "response.reasoning_summary_part.done": true, + "response.reasoning_summary_text.delta": true, + "response.reasoning_summary_text.done": true, + "response.reasoning_text.delta": true, + "response.reasoning_text.done": true, + "response.function_call_arguments.delta": true, + "response.function_call_arguments.done": true, + "response.web_search_call.in_progress": true, + "response.web_search_call.searching": true, + "response.web_search_call.completed": true, + "error": true, + "codex.rate_limits": true, + "codex.response.metadata": true, +} + +var interactionsEventTypes = map[string]bool{ + "interaction.created": true, + "interaction.status_update": true, + "interaction.completed": true, + "interaction.failed": true, + "finish": true, + "step.start": true, + "step.delta": true, + "step.stop": true, +} + +type interactionsChunk struct { + Object string `json:"object"` + EventType string `json:"event_type"` + Type string `json:"type"` + Status string `json:"status"` + InteractionID string `json:"interaction_id"` + Steps []interactionsStep `json:"steps"` + Step *interactionsStep `json:"step"` + Delta *interactionsDelta `json:"delta"` + Usage *interactionsUsage `json:"usage"` + Metadata *interactionsMeta `json:"metadata"` + Interaction *struct { + ID string `json:"id"` + Status string `json:"status"` + Object string `json:"object"` + Steps []interactionsStep `json:"steps"` + Usage *interactionsUsage `json:"usage"` + } `json:"interaction"` +} + +type interactionsMeta struct { + TotalUsage *interactionsUsage `json:"total_usage"` + Usage *interactionsUsage `json:"usage"` +} + +type interactionsUsage struct { + OutputTokens *tokenCount `json:"output_tokens"` + TotalOutputTokens *tokenCount `json:"total_output_tokens"` + CompletionTokens *tokenCount `json:"completion_tokens"` +} + +type interactionsStep struct { + ID string `json:"id"` + CallID string `json:"call_id"` + Type string `json:"type"` + Name string `json:"name"` + Arguments json.RawMessage `json:"arguments"` + Content []interactionsContent `json:"content"` + Result json.RawMessage `json:"result"` + Signature string `json:"signature"` + ThoughtSignature string `json:"thought_signature"` + ThoughtSignatureCamel string `json:"thoughtSignature"` + EncryptedContent string `json:"encrypted_content"` + ExtraContent *interactionsExtraContent `json:"extra_content"` +} + +// interactionsExtraContent carries the vendor-specific envelope Gemini uses to +// ship a thought signature alongside a step. +type interactionsExtraContent struct { + Google *struct { + ThoughtSignature string `json:"thought_signature"` + } `json:"google"` +} + +// hasSignature reports whether the step carries a reasoning signature. A step +// that only carries a signature is still a meaningful upstream answer: dropping +// it makes the turn look empty and costs the signature on the next request. +func (s *interactionsStep) hasSignature() bool { + if s == nil { + return false + } + if strings.TrimSpace(s.Signature) != "" || + strings.TrimSpace(s.ThoughtSignature) != "" || + strings.TrimSpace(s.ThoughtSignatureCamel) != "" || + strings.TrimSpace(s.EncryptedContent) != "" { + return true + } + if s.ExtraContent != nil && s.ExtraContent.Google != nil { + return strings.TrimSpace(s.ExtraContent.Google.ThoughtSignature) != "" + } + return false +} + +type interactionsContent struct { + Type string `json:"type"` + Text string `json:"text"` + Data string `json:"data"` + FileURI string `json:"file_uri"` + FileUri string `json:"fileUri"` + URL string `json:"url"` + MimeType string `json:"mime_type"` + Mime_Type string `json:"mimeType"` + Signature string `json:"signature"` + ThoughtSignature string `json:"thought_signature"` + ThoughtSignatureCamel string `json:"thoughtSignature"` +} + +func (c *interactionsContent) hasMeaningfulContent() bool { + if c == nil { + return false + } + if strings.TrimSpace(c.Text) != "" || + strings.TrimSpace(c.Data) != "" || + strings.TrimSpace(c.FileURI) != "" || + strings.TrimSpace(c.FileUri) != "" || + strings.TrimSpace(c.URL) != "" { + return true + } + if strings.TrimSpace(c.Signature) != "" || + strings.TrimSpace(c.ThoughtSignature) != "" || + strings.TrimSpace(c.ThoughtSignatureCamel) != "" { + return true + } + return false +} + +type interactionsDelta struct { + Type string `json:"type"` + Text string `json:"text"` + Data string `json:"data"` + FileURI string `json:"file_uri"` + FileUri string `json:"fileUri"` + URL string `json:"url"` + Arguments json.RawMessage `json:"arguments"` + Signature string `json:"signature"` + ThoughtSignature string `json:"thought_signature"` + ThoughtSignatureCamel string `json:"thoughtSignature"` + Name string `json:"name"` + Content *interactionsContent `json:"content"` + Result json.RawMessage `json:"result"` +} + +func hasMeaningfulInteractionsArguments(raw json.RawMessage) bool { + trimmed := bytes.TrimSpace(raw) + if len(trimmed) == 0 || bytes.Equal(trimmed, []byte("null")) { + return false + } + var str string + if err := json.Unmarshal(trimmed, &str); err == nil { + return hasMeaningfulJSONArguments(str) + } + return nonEmptyJSONPayload(raw) +} + +// emptyCompletionAccum accumulates the properties relevant to deciding whether +// an OpenAI-, Claude-, or Gemini-style completion is empty. +type emptyCompletionAccum struct { + expectedChoices int + recognized bool + sawUnknownData bool + terminal bool + hasContent bool + hasToolCalls bool + completionTokens int + sawUsage bool + blocked bool + sawMetadataOnly bool + sawMessageData bool + geminiTerminal bool + claudeTerminal bool + openAITerminal bool + interactionsTerminal bool + openAIChoicesSeen map[int]bool + openAIChoicesFinished map[int]bool + geminiCandidatesSeen map[int]bool + geminiCandidatesFinished map[int]bool +} + +func (a *emptyCompletionAccum) evalJSON(data []byte) bool { + values, err := decodeJSONValues(data) + if err != nil { + return false + } + recognized := false + for _, v := range values { + if evalProviderError(v, "") != nil { + recognized = true + a.blocked = true + a.terminal = true + } else if a.evalOpenAI(v) || a.evalClaude(v) || a.evalOpenAIResponse(v) || a.evalGemini(v) || a.evalInteractions(v) { + recognized = true + } else { + a.sawUnknownData = 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 values, nil +} + +func (a *emptyCompletionAccum) evalOpenAI(data []byte) bool { + // Recognize the OpenAI shape by the presence of a "choices" key, even when + // the array is empty (e.g. {"choices":[]}). Such prefixes must still be + // judged at stream close instead of being forwarded immediately. + if !hasJSONKey(data, "choices") { + 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.addUsage(*chunk.Usage.CompletionTokens) + } + if a.openAIChoicesSeen == nil { + a.openAIChoicesSeen = make(map[int]bool) + a.openAIChoicesFinished = make(map[int]bool) + } + for i, ch := range chunk.Choices { + idx := i + if ch.Index != nil { + idx = *ch.Index + } + a.openAIChoicesSeen[idx] = true + if ch.FinishReason != nil { + reason := strings.TrimSpace(*ch.FinishReason) + if strings.EqualFold(reason, "stop") || strings.EqualFold(reason, "tool_calls") || strings.EqualFold(reason, "function_call") { + a.openAIChoicesFinished[idx] = true + 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.Text + ch.Delta.Content + ch.Message.Content + ch.Delta.ReasoningContent + ch.Message.ReasoningContent + ch.Delta.Reasoning + ch.Message.Reasoning + if strings.TrimSpace(content) != "" { + a.hasContent = true + } + if (ch.Delta.Refusal != nil && strings.TrimSpace(*ch.Delta.Refusal) != "") || + (ch.Message.Refusal != nil && strings.TrimSpace(*ch.Message.Refusal) != "") { + a.hasContent = true + } + 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 + } + } + expected := a.expectedChoices + if expected <= 0 { + expected = 1 + } + targetChoices := expected + if len(a.openAIChoicesSeen) > targetChoices { + targetChoices = len(a.openAIChoicesSeen) + } + if len(a.openAIChoicesFinished) >= targetChoices && len(a.openAIChoicesFinished) >= len(a.openAIChoicesSeen) && !a.blocked { + a.openAITerminal = true + } else { + a.openAITerminal = false + } + 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. (Streamed zero-choices chunks without + // usage are mid-stream signals and must not mark terminal here.) + a.terminal = true + } + return true +} + +func (a *emptyCompletionAccum) evalClaude(data []byte) bool { + var chunk claudeChunk + if err := json.Unmarshal(data, &chunk); err != nil { + return false + } + + isClaude := false + switch chunk.Type { + 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)) { + isClaude = true + } + } + + if !isClaude { + return false + } + + 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 { + a.evalClaudeStopReason(chunk.Message.StopReason) + } + if chunk.Delta != nil { + a.evalClaudeStopReason(chunk.Delta.StopReason) + } + + if chunk.Usage != nil && chunk.Usage.OutputTokens != nil { + a.sawUsage = true + a.addUsage(*chunk.Usage.OutputTokens) + } + if chunk.Message != nil && chunk.Message.Usage != nil && chunk.Message.Usage.OutputTokens != nil { + a.sawUsage = true + a.addUsage(*chunk.Message.Usage.OutputTokens) + } + + a.evalClaudeBlocks(chunk.Content) + if chunk.Message != nil { + a.evalClaudeBlocks(chunk.Message.Content) + } + if chunk.ContentBlock != nil { + a.evalClaudeBlocks([]claudeContentBlock{*chunk.ContentBlock}) + } + if chunk.Delta != nil { + switch chunk.Delta.Type { + case "text_delta": + if strings.TrimSpace(chunk.Delta.Text) != "" { + a.hasContent = true + } + 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": + if hasMeaningfulClaudePartialJSON(chunk.Delta.PartialJSON) { + a.hasToolCalls = true + } + default: + if strings.TrimSpace(chunk.Delta.Text) != "" || strings.TrimSpace(chunk.Delta.Thinking) != "" || strings.TrimSpace(chunk.Delta.Signature) != "" || nonEmptyJSONPayload(chunk.Delta.Citation) { + a.hasContent = true + } + } + } + + return true +} + +func (a *emptyCompletionAccum) evalClaudeStopReason(stopReason *string) { + if stopReason == nil { + return + } + reason := strings.TrimSpace(*stopReason) + if strings.EqualFold(reason, "end_turn") || strings.EqualFold(reason, "tool_use") { + a.terminal = true + a.claudeTerminal = true + } else if reason != "" { + // Request/output limits, refusals, and control stop reasons must reach the + // client instead of being converted into a credential failure. + a.blocked = true + } +} + +func (a *emptyCompletionAccum) evalOpenAIResponse(data []byte) bool { + var probe map[string]json.RawMessage + if err := json.Unmarshal(data, &probe); err != nil { + return false + } + + var evType string + if raw := probe["type"]; raw != nil { + _ = json.Unmarshal(raw, &evType) + } + var objName string + if raw := probe["object"]; raw != nil { + _ = json.Unmarshal(raw, &objName) + } + + if objName != "response" && !openAIResponseEventTypes[evType] { + return false + } + a.recognized = true + a.sawMessageData = true + + var chunk openAIResponseChunk + if err := json.Unmarshal(data, &chunk); err != nil { + return true + } + + 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 + a.blocked = true + case "response.incomplete", "response.failed", "error": + a.terminal = true + a.blocked = true + } + a.evalOpenAIResponseStatus(chunk.Status) + if chunk.Response != nil { + a.evalOpenAIResponseStatus(chunk.Response.Status) + } + + if chunk.Usage != nil && chunk.Usage.OutputTokens != nil { + a.sawUsage = true + a.addUsage(*chunk.Usage.OutputTokens) + } + if chunk.Response != nil && chunk.Response.Usage != nil && chunk.Response.Usage.OutputTokens != nil { + a.sawUsage = true + a.addUsage(*chunk.Response.Usage.OutputTokens) + } + + switch evType { + case "response.output_text.delta", "response.reasoning_summary_text.delta", "response.reasoning_text.delta": + if strings.TrimSpace(chunk.Delta) != "" { + a.hasContent = true + } + case "response.output_text.done", "response.reasoning_summary_text.done", "response.reasoning_text.done": + if strings.TrimSpace(chunk.Text) != "" { + a.hasContent = true + } + case "response.reasoning_summary_part.added", "response.reasoning_summary_part.done", "response.content_part.added", "response.content_part.done": + var part openAIResponsePart + if err := json.Unmarshal(chunk.Part, &part); err == nil { + if strings.TrimSpace(part.Text) != "" || strings.TrimSpace(part.EncryptedContent) != "" || hasMeaningfulAnnotations(part.Annotations) { + 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") && itemType != "image_generation_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") && itemType != "image_generation_call" { + if hasMeaningfulResponsesCallItem(item) { + a.hasToolCalls = true + } + } + } + case "response.function_call_arguments.delta", "response.function_call_arguments.done": + if a.hasToolCalls || hasMeaningfulJSONArguments(chunk.Delta) || hasMeaningfulJSONArguments(chunk.Arguments) { + a.hasToolCalls = true + } + case "response.web_search_call.in_progress", "response.web_search_call.searching", "response.web_search_call.completed", + "codex.rate_limits", "codex.response.metadata": + } + + a.evalOpenAIResponseRawOutput(chunk.Output) + a.evalOpenAIResponseRawOutput(chunk.Item) + if chunk.Response != nil { + a.evalOpenAIResponseRawOutput(chunk.Response.Output) + } + + return true +} + +func (a *emptyCompletionAccum) evalOpenAIResponseStatus(status string) { + switch strings.ToLower(strings.TrimSpace(status)) { + case "completed": + a.terminal = true + a.blocked = true + case "incomplete", "failed", "error": + a.terminal = true + a.blocked = true + } +} + +func (a *emptyCompletionAccum) evalOpenAIResponseRawOutput(raw json.RawMessage) { + if len(raw) == 0 { + return + } + var items []openAIResponseOutputItem + if err := json.Unmarshal(raw, &items); err == nil { + a.evalOpenAIResponseOutput(items) + return + } + var item openAIResponseOutputItem + if err := json.Unmarshal(raw, &item); err == nil { + a.evalOpenAIResponseOutput([]openAIResponseOutputItem{item}) + } +} + +func hasMeaningfulResponsesCallItem(item openAIResponseOutputItem) bool { + return strings.TrimSpace(item.Name) != "" || + hasMeaningfulJSONArguments(item.Arguments) || + strings.TrimSpace(item.Input) != "" || + strings.TrimSpace(item.Result) != "" || + nonEmptyJSONPayload(item.Action) || + nonEmptyJSONPayload(item.Results) +} + +func hasMeaningfulResponsesImageGenerationCallItem(item openAIResponseOutputItem) bool { + return strings.TrimSpace(item.Result) != "" || strings.TrimSpace(item.Text) != "" +} + +func hasMeaningfulResponsesReasoningSummary(raw json.RawMessage) bool { + trimmed := bytes.TrimSpace(raw) + if len(trimmed) == 0 || bytes.Equal(trimmed, []byte("null")) { + return false + } + var parts []openAIResponsePart + if err := json.Unmarshal(trimmed, &parts); err == nil { + for _, part := range parts { + if strings.TrimSpace(part.Text) != "" || strings.TrimSpace(part.EncryptedContent) != "" { + return true + } + } + return false + } + var single openAIResponsePart + if err := json.Unmarshal(trimmed, &single); err == nil { + return strings.TrimSpace(single.Text) != "" || strings.TrimSpace(single.EncryptedContent) != "" + } + var strSlice []string + if err := json.Unmarshal(trimmed, &strSlice); err == nil { + for _, s := range strSlice { + if strings.TrimSpace(s) != "" { + return true + } + } + return false + } + var str string + if err := json.Unmarshal(trimmed, &str); err == nil { + return strings.TrimSpace(str) != "" + } + return false +} + +func hasMeaningfulAnnotations(annotations []json.RawMessage) bool { + if len(annotations) == 0 { + return false + } + for _, raw := range annotations { + trimmed := bytes.TrimSpace(raw) + if len(trimmed) == 0 || bytes.Equal(trimmed, []byte("null")) || bytes.Equal(trimmed, []byte("{}")) || bytes.Equal(trimmed, []byte("[]")) { + continue + } + var obj map[string]any + if err := json.Unmarshal(trimmed, &obj); err == nil { + hasField := false + for _, v := range obj { + switch val := v.(type) { + case nil: + case string: + if strings.TrimSpace(val) != "" { + hasField = true + } + default: + hasField = true + } + if hasField { + break + } + } + if hasField { + return true + } + continue + } + var str string + if err := json.Unmarshal(trimmed, &str); err == nil { + if strings.TrimSpace(str) != "" { + return true + } + continue + } + return true + } + return false +} + +func (a *emptyCompletionAccum) evalOpenAIResponseOutput(items []openAIResponseOutputItem) { + for _, item := range items { + itemType := strings.ToLower(strings.TrimSpace(item.Type)) + switch { + case itemType == "image_generation_call": + if hasMeaningfulResponsesImageGenerationCallItem(item) { + a.hasContent = true + } + case strings.HasSuffix(itemType, "_call"): + if hasMeaningfulResponsesCallItem(item) { + a.hasToolCalls = true + } + case itemType == "reasoning": + if strings.TrimSpace(item.EncryptedContent) != "" || hasMeaningfulResponsesReasoningSummary(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. + a.hasContent = true + } + if strings.TrimSpace(item.Text) != "" { + a.hasContent = true + } + for _, part := range item.Content { + partType := strings.ToLower(strings.TrimSpace(part.Type)) + if strings.TrimSpace(part.Text) != "" || strings.TrimSpace(part.Refusal) != "" || + hasMeaningfulAnnotations(part.Annotations) || + partType == "refusal" || (partType != "" && partType != "output_text" && partType != "text") { + a.hasContent = true + } + } + } +} + +func (a *emptyCompletionAccum) evalClaudeBlocks(blocks []claudeContentBlock) { + for _, b := range blocks { + 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" || b.Type == "reasoning" || 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 b.Type == "text" || strings.TrimSpace(b.Text) != "" { + if strings.TrimSpace(b.Text) != "" { + a.hasContent = true + } + continue + } + if nonEmptyJSONPayload(b.Citation) { + a.hasContent = true + continue + } + } +} + +// hasJSONKey reports whether the given JSON object contains name as a top-level +// key. It returns false for non-object or malformed input. +func hasJSONKey(data []byte, name string) bool { + var probe map[string]json.RawMessage + if err := json.Unmarshal(data, &probe); err != nil { + return false + } + _, ok := probe[name] + return ok +} + +// hasNestedResponseCandidates reports whether the payload's response object +// contains a candidates key (the Gemini streaming wrapper shape). +func hasNestedResponseCandidates(data []byte) bool { + var probe struct { + Response map[string]json.RawMessage `json:"response"` + } + if err := json.Unmarshal(data, &probe); err != nil { + return false + } + _, ok := probe.Response["candidates"] + return ok +} + +func (a *emptyCompletionAccum) evalGemini(data []byte) bool { + var chunk geminiChunk + if err := json.Unmarshal(data, &chunk); err != nil { + return false + } + + candidates := chunk.Candidates + usage := chunk.UsageMetadata + promptFeedback := chunk.PromptFeedback + + if chunk.Response != nil { + if len(candidates) == 0 { + candidates = chunk.Response.Candidates + } + if usage == nil { + usage = chunk.Response.UsageMetadata + } + if promptFeedback == nil { + promptFeedback = chunk.Response.PromptFeedback + } + } + promptBlocked := promptFeedback != nil && strings.TrimSpace(promptFeedback.BlockReason) != "" + + if len(candidates) == 0 { + // Only treat an empty candidates array as a recognized empty completion + // when the candidates key is actually present (e.g. a Gemini + // safety/empty response with zero candidate tokens, or a stream + // aggregate with nothing else). An absent candidates key is not a + // 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.addUsage(*usage.CandidatesTokenCount) + } + return true + } + return false + } + + a.recognized = true + a.sawMessageData = true + if promptBlocked { + a.blocked = true + } + + if usage != nil { + if usage.CandidatesTokenCount != nil { + a.sawUsage = true + a.addUsage(*usage.CandidatesTokenCount) + } + } + if hasMeaningfulGroundingMetadata(chunk.GroundingMetadata) { + a.hasContent = true + } + if chunk.Response != nil && hasMeaningfulGroundingMetadata(chunk.Response.GroundingMetadata) { + a.hasContent = true + } + + if a.geminiCandidatesSeen == nil { + a.geminiCandidatesSeen = make(map[int]bool) + a.geminiCandidatesFinished = make(map[int]bool) + } + + blocked := false + for i, cand := range candidates { + idx := i + if cand.Index != nil { + idx = *cand.Index + } + a.geminiCandidatesSeen[idx] = true + if cand.FinishReason != nil { + reason := strings.TrimSpace(*cand.FinishReason) + if reason != "" { + if strings.EqualFold(reason, "STOP") { + a.geminiCandidatesFinished[idx] = true + a.terminal = true + } else { + // A blocking or other terminal reason (SAFETY, RECITATION, + // MAX_TOKENS, BLOCKLIST, PROHIBITED_CONTENT, OTHER) is not an + // empty completion: the client must see the stop/block reason + // rather than a silent auth rotation. + blocked = true + a.terminal = true + } + } + } + if hasMeaningfulGroundingMetadata(cand.GroundingMetadata) { + a.hasContent = true + } + if cand.Content != nil { + for _, part := range cand.Content.Parts { + if isMeaningfulGeminiFunctionCall(part.FunctionCall) { + a.hasToolCalls = true + } + if hasMeaningfulGeminiMediaPayload(part.InlineData, part.FileData) || + nonEmptyJSONPayload(part.FunctionResponse) { + 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 blocked { + a.blocked = true + } + + expected := a.expectedChoices + if expected <= 0 { + expected = 1 + } + targetCandidates := expected + if len(a.geminiCandidatesSeen) > targetCandidates { + targetCandidates = len(a.geminiCandidatesSeen) + } + if len(a.geminiCandidatesFinished) >= targetCandidates && len(a.geminiCandidatesFinished) >= len(a.geminiCandidatesSeen) && !a.blocked { + a.geminiTerminal = true + } else { + a.geminiTerminal = false + } + + return true +} + +func (a *emptyCompletionAccum) evalInteractions(data []byte) bool { + var probe map[string]json.RawMessage + if err := json.Unmarshal(data, &probe); err != nil { + return false + } + + var evType string + if raw := probe["event_type"]; raw != nil { + _ = json.Unmarshal(raw, &evType) + } + if evType == "" { + if raw := probe["type"]; raw != nil { + _ = json.Unmarshal(raw, &evType) + } + } + + var objName string + if raw := probe["object"]; raw != nil { + _ = json.Unmarshal(raw, &objName) + } + + isInteractions := objName == "interaction" || + interactionsEventTypes[evType] || + hasJSONKey(data, "interaction") || + (hasJSONKey(data, "steps") && (hasJSONKey(data, "status") || hasJSONKey(data, "interaction_id"))) + + if !isInteractions { + return false + } + + a.recognized = true + a.sawMessageData = true + + var chunk interactionsChunk + if err := json.Unmarshal(data, &chunk); err != nil { + return true + } + + status := chunk.Status + if chunk.Interaction != nil && chunk.Interaction.Status != "" { + status = chunk.Interaction.Status + } + + switch strings.ToLower(strings.TrimSpace(status)) { + case "completed": + a.terminal = true + a.interactionsTerminal = true + case "failed", "cancelled", "error", "blocked", "incomplete": + a.terminal = true + a.blocked = true + case "requires_action": + a.terminal = true + a.blocked = true + } + + if evType == "interaction.completed" { + a.terminal = true + if !a.blocked { + a.interactionsTerminal = true + } + } else if evType == "finish" { + // The Interactions protocol ends a turn with a bare "finish" event whose + // usage lives under metadata.total_usage instead of the top-level usage + // field the other terminal events carry. + a.terminal = true + if !a.blocked { + a.interactionsTerminal = true + } + if chunk.Metadata != nil { + if chunk.Metadata.TotalUsage != nil { + a.evalInteractionsUsage(chunk.Metadata.TotalUsage) + } else { + a.evalInteractionsUsage(chunk.Metadata.Usage) + } + } + } else if evType == "interaction.failed" { + a.terminal = true + a.blocked = true + } + + a.evalInteractionsUsage(chunk.Usage) + if chunk.Interaction != nil { + a.evalInteractionsUsage(chunk.Interaction.Usage) + } + + if len(chunk.Steps) == 0 && chunk.Interaction != nil { + a.evalInteractionsSteps(chunk.Interaction.Steps) + } else { + a.evalInteractionsSteps(chunk.Steps) + } + if chunk.Step != nil { + a.evalInteractionsSteps([]interactionsStep{*chunk.Step}) + } + + if chunk.Delta != nil { + if strings.TrimSpace(chunk.Delta.Text) != "" || + strings.TrimSpace(chunk.Delta.Data) != "" || + strings.TrimSpace(chunk.Delta.FileURI) != "" || + strings.TrimSpace(chunk.Delta.FileUri) != "" || + strings.TrimSpace(chunk.Delta.URL) != "" { + a.hasContent = true + } + if chunk.Delta.Content != nil && chunk.Delta.Content.hasMeaningfulContent() { + a.hasContent = true + } + if strings.TrimSpace(chunk.Delta.Signature) != "" || + strings.TrimSpace(chunk.Delta.ThoughtSignature) != "" || + strings.TrimSpace(chunk.Delta.ThoughtSignatureCamel) != "" { + a.hasContent = true + } + if strings.TrimSpace(chunk.Delta.Name) != "" || hasMeaningfulInteractionsArguments(chunk.Delta.Arguments) { + a.hasToolCalls = true + } + if nonEmptyJSONPayload(chunk.Delta.Result) { + a.hasContent = true + } + } + + return true +} + +func (a *emptyCompletionAccum) evalInteractionsUsage(usage *interactionsUsage) { + if usage == nil { + return + } + if usage.OutputTokens != nil { + a.sawUsage = true + a.addUsage(*usage.OutputTokens) + } + if usage.TotalOutputTokens != nil { + a.sawUsage = true + a.addUsage(*usage.TotalOutputTokens) + } + if usage.CompletionTokens != nil { + a.sawUsage = true + a.addUsage(*usage.CompletionTokens) + } +} + +func (a *emptyCompletionAccum) evalInteractionsSteps(steps []interactionsStep) { + for _, step := range steps { + stepType := strings.ToLower(strings.TrimSpace(step.Type)) + switch stepType { + case "function_call": + if strings.TrimSpace(step.Name) != "" || hasMeaningfulInteractionsArguments(step.Arguments) { + a.hasToolCalls = true + } + case "function_result": + if strings.TrimSpace(step.Name) != "" || nonEmptyJSONPayload(step.Result) { + a.hasContent = true + } + default: + if strings.TrimSpace(step.Name) != "" || hasMeaningfulInteractionsArguments(step.Arguments) { + a.hasToolCalls = true + } + if nonEmptyJSONPayload(step.Result) { + a.hasContent = true + } + } + if step.hasSignature() { + a.hasContent = true + } + for _, content := range step.Content { + if content.hasMeaningfulContent() { + a.hasContent = true + } + } + } +} + +// empty reports whether the accumulated stream is an empty completion. +func (a *emptyCompletionAccum) empty() bool { + if a.sawUnknownData || a.blocked || a.hasContent || a.hasToolCalls || (a.sawUsage && a.completionTokens > 0) { + return false + } + if a.recognized && a.terminal { + return true + } + if a.recognized { + return true + } + if a.sawMetadataOnly && !a.sawMessageData { + return true + } + return false +} + +// isEmptyCompletion reports whether the buffered SSE stream chunks aggregate to +// an empty completion. +func isEmptyCompletion(chunks []cliproxyexecutor.StreamChunk) bool { + if len(chunks) == 0 { + return false + } + var detector StreamBootstrapDetector + for _, c := range chunks { + if detector.Observe(c.Payload) { + return false + } + } + return detector.Finish() +} + +func isEmptyCompletionError(err error) bool { + var authErr *Error + return errors.As(err, &authErr) && authErr != nil && authErr.Code == errEmptyCompletion.Code +} + +// 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 + dataLines [][]byte + forward bool + sawSSE bool + sawDone bool + currentEvent string + streamErr *Error +} + +func (s *streamBootstrapState) streamError() error { + if s == nil || s.streamErr == nil { + return nil + } + return s.streamErr +} + +func (s *streamBootstrapState) flushData() { + if len(s.dataLines) == 0 { + s.currentEvent = "" + return + } + data := bytes.Join(s.dataLines, []byte("\n")) + s.dataLines = s.dataLines[:0] + currentEvent := s.currentEvent + s.currentEvent = "" + 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 { + if currentEvent != "error" { + s.acc.sawMetadataOnly = true + } + return + } + if err := evalProviderError(data, currentEvent); err != nil { + s.streamErr = err + 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 { + if len(s.dataLines) > 0 && classifyJSONBuffer(bytes.Join(s.dataLines, []byte("\n"))) == jsonBufIncomplete { + return + } + s.flushData() + return + } + s.processSingleLine(line) +} + +func (s *streamBootstrapState) processSingleLine(line []byte) { + switch { + case bytes.HasPrefix(line, []byte("event:")): + s.sawSSE = true + event := strings.TrimSpace(string(bytes.TrimPrefix(line, []byte("event:")))) + s.currentEvent = event + if event == "message_stop" { + s.acc.recognized = true + s.acc.terminal = true + s.acc.sawMessageData = true + s.sawDone = true + } else if event == "error" { + // Do not mark metadata only as success signal on error event + } 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 + // Raw JSONL/NDJSON frames are newline-terminated and never followed by a + // blank separator line, so buffering them would defer evaluation until the + // bootstrap byte cap is hit. Evaluate a self-contained frame immediately; + // keep buffering only when a multi-line JSON value is already in progress. + if len(s.dataLines) == 0 && classifyJSONBuffer(line) == jsonBufComplete { + if err := evalProviderError(line, ""); err != nil { + s.streamErr = err + } else if !s.acc.evalJSON(line) { + s.acc.sawUnknownData = true + } + return + } + s.dataLines = append(s.dataLines, line) + default: + // A pretty-printed raw JSON frame arrives one line at a time: the opening + // brace lands in dataLines and every continuation line looks like an + // isolated, invalid JSON value on its own. Append the closing line first, + // then classify the joined buffer; evaluate immediately when it becomes + // complete instead of buffering until a blank line or EOF that may not come. + if len(s.dataLines) > 0 { + s.dataLines = append(s.dataLines, line) + joined := bytes.Join(s.dataLines, []byte("\n")) + s.dataLines = s.dataLines[:0] + switch classifyJSONBuffer(joined) { + case jsonBufComplete: + if err := evalProviderError(joined, ""); err != nil { + s.streamErr = err + } else if !s.acc.evalJSON(joined) { + s.acc.sawUnknownData = true + } + case jsonBufIncomplete: + // The buffered value is still incomplete; restore the accumulated + // lines and wait for the next continuation. + s.dataLines = append([][]byte(nil), joined) + default: + s.acc.sawUnknownData = true + } + return + } + if classify := classifyJSONBuffer(line); classify == jsonBufComplete || classify == jsonBufIncomplete { + if err := evalProviderError(line, ""); err != nil { + s.streamErr = err + } else if !s.acc.evalJSON(line) { + s.acc.sawUnknownData = true + } + } else { + s.acc.sawUnknownData = true + } + } +} + +func (s *streamBootstrapState) observe(fragment []byte) bool { + if s.forward { + return true + } + s.bytes += len(fragment) + if s.bytes > maxStreamBootstrapBytes { + s.forward = true + return true + } + s.pending = append(s.pending, fragment...) + for { + if newline := bytes.IndexByte(s.pending, '\n'); newline >= 0 { + line := bytes.TrimSpace(s.pending[:newline]) + s.pending = s.pending[newline+1:] + s.processLine(line) + if s.shouldForward() { + s.forward = true + return true + } + continue + } + break + } + + trimmed := bytes.TrimSpace(s.pending) + if len(trimmed) == 0 { + return false + } + + 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 err := evalProviderError(trimmed, ""); err != nil { + s.streamErr = err + } else if !s.acc.evalJSON(trimmed) { + s.acc.sawUnknownData = true + } + s.pending = s.pending[:0] + case jsonBufEmpty, jsonBufIncomplete: + return false + 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.openAITerminal || s.acc.interactionsTerminal) && s.acc.empty() +} + +func (s *streamBootstrapState) setExpectedChoices(n int) { + if n <= 0 { + n = 1 + } + s.acc.expectedChoices = n +} + +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.streamErr != nil { + return false + } + if !s.acc.recognized && !s.sawSSE && s.bytes > 0 { + return true + } + return false +} + +func (s *streamBootstrapState) shouldForward() bool { + if s.streamErr != nil { + return false + } + 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 streamErrorEnvelope struct { + Type string `json:"type"` + EventType string `json:"event_type"` + Error json.RawMessage `json:"error"` + Message string `json:"message"` + Code json.RawMessage `json:"code"` + Status string `json:"status"` + Response *streamErrorEnvelope `json:"response,omitempty"` + Interaction *streamErrorEnvelope `json:"interaction,omitempty"` +} + +func inferHTTPStatus(typeStr, codeStr, statusStr string) int { + if statusStr != "" { + switch strings.ToUpper(strings.TrimSpace(statusStr)) { + case "RESOURCE_EXHAUSTED": + return http.StatusTooManyRequests + case "UNAUTHENTICATED": + return http.StatusUnauthorized + case "PERMISSION_DENIED": + return http.StatusForbidden + case "UNAVAILABLE": + return http.StatusServiceUnavailable + case "DEADLINE_EXCEEDED": + return http.StatusGatewayTimeout + case "INTERNAL": + return http.StatusInternalServerError + case "INVALID_ARGUMENT", "FAILED_PRECONDITION": + return http.StatusBadRequest + case "NOT_FOUND": + return http.StatusNotFound + case "ALREADY_EXISTS": + return http.StatusConflict + } + } + for _, s := range []string{typeStr, codeStr} { + switch strings.ToLower(strings.TrimSpace(s)) { + case "overloaded_error", "overloaded": + return http.StatusServiceUnavailable + case "rate_limit_error", "rate_limit_exceeded", "insufficient_quota", "quota_exceeded", "requests": + return http.StatusTooManyRequests + case "authentication_error", "invalid_api_key", "unauthorized": + return http.StatusUnauthorized + case "permission_error", "forbidden": + return http.StatusForbidden + case "not_found_error": + return http.StatusNotFound + case "invalid_request_error", "bad_request_error", "invalid_prompt", "cyber_policy", "context_length_exceeded": + return http.StatusBadRequest + case "api_error", "internal_server_error": + return http.StatusInternalServerError + } + } + return 0 +} + +// isInteractionsFailureEnvelope reports whether an Interactions event carries a +// provider failure. The failure detail is nested under "interaction", so without +// this the stream is only marked blocked and the request never fails over. +func isInteractionsFailureEnvelope(envelope streamErrorEnvelope) bool { + if strings.EqualFold(envelope.EventType, "interaction.failed") || strings.EqualFold(envelope.Type, "interaction.failed") { + return true + } + if envelope.Interaction == nil { + return false + } + if len(envelope.Interaction.Error) > 0 && !bytes.Equal(envelope.Interaction.Error, []byte("null")) { + return true + } + return strings.EqualFold(envelope.Interaction.Status, "failed") +} + +func parseStreamErrorFromEnvelope(data []byte, envelope streamErrorEnvelope) *Error { + if envelope.Interaction != nil { + if (len(envelope.Error) == 0 || bytes.Equal(envelope.Error, []byte("null"))) && len(envelope.Interaction.Error) > 0 { + envelope.Error = envelope.Interaction.Error + } + if envelope.Message == "" { + envelope.Message = envelope.Interaction.Message + } + if len(envelope.Code) == 0 { + envelope.Code = envelope.Interaction.Code + } + } + + if envelope.Response != nil { + if (len(envelope.Error) == 0 || bytes.Equal(envelope.Error, []byte("null"))) && len(envelope.Response.Error) > 0 { + envelope.Error = envelope.Response.Error + } + if envelope.Message == "" { + envelope.Message = envelope.Response.Message + } + if len(envelope.Code) == 0 { + envelope.Code = envelope.Response.Code + } + if envelope.Status == "" { + envelope.Status = envelope.Response.Status + } + if (envelope.Type == "" || strings.EqualFold(envelope.Type, "response.failed")) && envelope.Response.Type != "" { + envelope.Type = envelope.Response.Type + } + } + + var detail struct { + Message string `json:"message"` + Type string `json:"type"` + Code json.RawMessage `json:"code"` + Status string `json:"status"` + } + + var rawErrorString string + if len(envelope.Error) > 0 { + trimmedErr := bytes.TrimSpace(envelope.Error) + if bytes.HasPrefix(trimmedErr, []byte("{")) { + _ = json.Unmarshal(trimmedErr, &detail) + } else if bytes.HasPrefix(trimmedErr, []byte("\"")) { + _ = json.Unmarshal(trimmedErr, &rawErrorString) + } + } + + message := detail.Message + if message == "" { + message = envelope.Message + } + if message == "" { + message = rawErrorString + } + if message == "" && detail.Type != "" { + message = detail.Type + } + if message == "" && detail.Status != "" { + message = detail.Status + } + if message == "" && len(data) > 0 && !bytes.HasPrefix(data, []byte("{")) { + message = string(data) + } + if message == "" { + message = "upstream stream error" + } + + code := "" + var rawCodeInt int + + extractCode := func(raw json.RawMessage) { + if len(raw) == 0 { + return + } + var strCode string + if json.Unmarshal(raw, &strCode) == nil && strings.TrimSpace(strCode) != "" { + code = strings.TrimSpace(strCode) + if num, err := strconv.Atoi(code); err == nil && num > 0 { + rawCodeInt = num + } + return + } + var num json.Number + if json.Unmarshal(raw, &num) == nil { + if n, err := num.Int64(); err == nil && n > 0 { + rawCodeInt = int(n) + code = strconv.Itoa(rawCodeInt) + } + } + } + + extractCode(detail.Code) + if code == "" { + extractCode(envelope.Code) + } + if code == "" && detail.Type != "" { + code = detail.Type + } + if code == "" && detail.Status != "" { + code = detail.Status + } + if code == "" && envelope.Type != "" && !strings.EqualFold(envelope.Type, "error") { + code = envelope.Type + } + + status := 0 + if rawCodeInt >= 100 && rawCodeInt <= 599 { + status = rawCodeInt + } + statusStr := strings.TrimSpace(detail.Status) + if statusStr == "" { + statusStr = strings.TrimSpace(envelope.Status) + } + typeStr := strings.TrimSpace(detail.Type) + if typeStr == "" { + typeStr = strings.TrimSpace(envelope.Type) + } + + if status == 0 { + status = inferHTTPStatus(typeStr, code, statusStr) + } + + if status == 0 { + lowerMsg := strings.ToLower(message) + switch { + case strings.Contains(lowerMsg, "rate limit") || strings.Contains(lowerMsg, "resource exhausted") || strings.Contains(lowerMsg, "too many requests") || strings.Contains(lowerMsg, "quota"): + status = http.StatusTooManyRequests + case strings.Contains(lowerMsg, "overloaded"): + status = http.StatusServiceUnavailable + case strings.Contains(lowerMsg, "unauthorized") || strings.Contains(lowerMsg, "invalid api key") || strings.Contains(lowerMsg, "invalid x-api-key") || strings.Contains(lowerMsg, "unauthenticated"): + status = http.StatusUnauthorized + case strings.Contains(lowerMsg, "permission denied") || strings.Contains(lowerMsg, "forbidden"): + status = http.StatusForbidden + default: + status = http.StatusBadGateway + } + } + + err := &Error{ + Code: code, + Message: message, + HTTPStatus: status, + } + + if isRequestInvalidError(err) || clienterror.IsRequestFault(status, errors.New(string(data))) { + err.Retryable = false + } else { + err.Retryable = true + } + + return err +} + +func evalProviderError(data []byte, sseEvent string) *Error { + trimmed := bytes.TrimSpace(data) + if len(trimmed) == 0 { + return nil + } + + var envelope streamErrorEnvelope + isError := strings.EqualFold(sseEvent, "error") + + if bytes.HasPrefix(trimmed, []byte("{")) { + if err := json.Unmarshal(trimmed, &envelope); err == nil { + if len(envelope.Error) > 0 && !bytes.Equal(envelope.Error, []byte("null")) { + isError = true + } else if strings.EqualFold(envelope.Type, "error") || strings.EqualFold(envelope.Type, "response.failed") { + isError = true + } else if isInteractionsFailureEnvelope(envelope) { + isError = true + } else if envelope.Response != nil { + if len(envelope.Response.Error) > 0 && !bytes.Equal(envelope.Response.Error, []byte("null")) { + isError = true + } else if strings.EqualFold(envelope.Response.Type, "error") || strings.EqualFold(envelope.Response.Status, "failed") { + isError = true + } + } + } + } + + if !isError { + return nil + } + + return parseStreamErrorFromEnvelope(trimmed, envelope) +} + +type streamPayloadErrorDetector struct { + pending []byte + dataLines [][]byte + currentEvent string + err *Error +} + +func (d *streamPayloadErrorDetector) Observe(chunk []byte) *Error { + if d.err != nil { + return d.err + } + if len(chunk) == 0 { + return nil + } + d.pending = append(d.pending, chunk...) + for { + newline := bytes.IndexByte(d.pending, '\n') + if newline < 0 { + break + } + line := bytes.TrimSpace(d.pending[:newline]) + d.pending = d.pending[newline+1:] + if len(line) == 0 { + d.flushData() + if d.err != nil { + return d.err + } + continue + } + d.processLine(line) + if d.err != nil { + return d.err + } + } + trimmed := bytes.TrimSpace(d.pending) + if len(trimmed) > 0 && !isSSEPrefix(trimmed) && !couldBeSSEPrefix(trimmed) { + if classifyJSONBuffer(trimmed) == jsonBufComplete { + if values, err := decodeJSONValues(trimmed); err == nil { + for _, v := range values { + if streamErr := evalProviderError(v, ""); streamErr != nil { + d.err = streamErr + d.pending = d.pending[:0] + return d.err + } + } + d.pending = d.pending[:0] + } + } + } + return d.err +} + +func (d *streamPayloadErrorDetector) processLine(line []byte) { + switch { + case bytes.HasPrefix(line, []byte("event:")): + d.currentEvent = strings.TrimSpace(string(bytes.TrimPrefix(line, []byte("event:")))) + case bytes.Equal(line, []byte("event")): + case bytes.HasPrefix(line, []byte("id:")), bytes.HasPrefix(line, []byte("retry:")), bytes.HasPrefix(line, []byte(":")): + case bytes.Equal(line, []byte("id")), bytes.Equal(line, []byte("retry")): + case bytes.HasPrefix(line, []byte("data:")): + d.dataLines = append(d.dataLines, parseSSEDataLine(line)) + case bytes.Equal(line, []byte("data")): + d.dataLines = append(d.dataLines, []byte("")) + case bytes.HasPrefix(line, []byte("{")), bytes.HasPrefix(line, []byte("[")): + // Same JSONL/NDJSON framing as the bootstrap detector: a self-contained + // frame is never followed by a blank line, so evaluate it now instead of + // buffering it until a separator that will not arrive. + if len(d.dataLines) == 0 && classifyJSONBuffer(line) == jsonBufComplete { + d.evalCompleteJSONLine(line) + return + } + d.dataLines = append(d.dataLines, line) + default: + // Mirror of the bootstrap detector: a pretty-printed raw JSON frame must + // append the closing line first, then classify the joined buffer, and + // evaluate immediately when it becomes complete. + if len(d.dataLines) > 0 { + d.dataLines = append(d.dataLines, line) + joined := bytes.Join(d.dataLines, []byte("\n")) + d.dataLines = nil + switch classifyJSONBuffer(joined) { + case jsonBufComplete: + d.evalCompleteJSONLine(joined) + case jsonBufIncomplete: + // Still incomplete; restore and wait for the next line. + d.dataLines = [][]byte{joined} + } + return + } + if classifyJSONBuffer(line) == jsonBufComplete { + d.evalCompleteJSONLine(line) + } + } +} + +func (d *streamPayloadErrorDetector) evalCompleteJSONLine(line []byte) { + values, err := decodeJSONValues(line) + if err != nil { + if streamErr := evalProviderError(line, ""); streamErr != nil { + d.err = streamErr + } + return + } + for _, v := range values { + if streamErr := evalProviderError(v, ""); streamErr != nil { + d.err = streamErr + return + } + } +} + +func (d *streamPayloadErrorDetector) flushData() { + if len(d.dataLines) == 0 { + return + } + data := bytes.Join(d.dataLines, []byte("\n")) + currentEvent := d.currentEvent + d.dataLines = nil + d.currentEvent = "" + if bytes.Equal(data, []byte("[DONE]")) { + return + } + if len(data) == 0 { + return + } + if err := evalProviderError(data, currentEvent); err != nil { + d.err = err + } +} + +func (d *streamPayloadErrorDetector) Finish() *Error { + if d.err != nil { + return d.err + } + if len(d.pending) > 0 { + trimmed := bytes.TrimSpace(d.pending) + d.pending = d.pending[:0] + if len(trimmed) > 0 { + if !isSSEPrefix(trimmed) && !couldBeSSEPrefix(trimmed) && classifyJSONBuffer(trimmed) == jsonBufComplete { + if values, err := decodeJSONValues(trimmed); err == nil { + for _, v := range values { + if err := evalProviderError(v, ""); err != nil { + d.err = err + return d.err + } + } + } + } else { + d.processLine(trimmed) + } + } + } + d.flushData() + return d.err +} + +func detectStreamPayloadError(payload []byte) *Error { + var d streamPayloadErrorDetector + if err := d.Observe(payload); err != nil { + return err + } + return d.Finish() +} + +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(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 +// a single non-stream JSON response) represents an empty completion. +func isEmptyCompletionPayload(payload []byte) bool { + trimmed := bytes.TrimSpace(payload) + if len(trimmed) == 0 { + // 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 isSSEPayload(trimmed) { + acc.evalSSE(trimmed) + return acc.empty() + } + + acc.evalJSON(trimmed) + // A complete non-SSE OpenAI chat completion body is terminal by + // construction: zero-choice payloads such as {"choices":[]} or + // {"choices":[],"usage":null} never enter the per-choice terminal paths, + // so without this they would be accepted as successful responses instead + // of being judged as empty completions. Other recognized shapes (for + // example Claude messages) keep their per-shape terminal rules. + var probe struct { + Choices json.RawMessage `json:"choices"` + } + if json.Unmarshal(trimmed, &probe) == nil && probe.Choices != nil { + acc.terminal = true + } + return acc.empty() +} + +func isSSEPayload(trimmed []byte) bool { + for _, line := range bytes.Split(trimmed, []byte("\n")) { + line = bytes.TrimSpace(line) + if len(line) == 0 { + continue + } + 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 + a.sawMessageData = true + return + } + if len(data) == 0 { + 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 +// returns the error to propagate. The mixed duty execution path rotates on an +// empty completion; the home (credits) path reports it via reportHomeResult. +func (m *Manager) markEmptyCompletion(ctx context.Context, result *Result) error { + result.Success = false + result.Error = errEmptyCompletion + m.MarkResult(ctx, *result) + return errEmptyCompletion +} diff --git a/sdk/cliproxy/auth/empty_completion_export.go b/sdk/cliproxy/auth/empty_completion_export.go new file mode 100644 index 00000000000..60a67f8f249 --- /dev/null +++ b/sdk/cliproxy/auth/empty_completion_export.go @@ -0,0 +1,158 @@ +package auth + +import "encoding/json" + +// 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 +// conductor, exposed so the plugin-executor path can reject empty completions +// before they reach the client. +func IsEmptyCompletionPayload(payload []byte) bool { + return isEmptyCompletionPayload(payload) +} + +// EmptyCompletionError returns the retriable error used when upstream returns +// a terminal but empty completion. The plugin-executor path returns it so the +// client receives an error instead of a silent empty response, matching how +// the conductor surfaces empty completions. +func EmptyCompletionError() error { + return errEmptyCompletion +} + +type choiceExtractionPayload struct { + N *int `json:"n"` + CandidateCount *int `json:"candidateCount"` + Candidate_Count *int `json:"candidate_count"` + GenerationConfig *struct { + CandidateCount *int `json:"candidateCount"` + Candidate_Count *int `json:"candidate_count"` + } `json:"generationConfig"` + Generation_Config *struct { + CandidateCount *int `json:"candidateCount"` + Candidate_Count *int `json:"candidate_count"` + } `json:"generation_config"` + Request *struct { + N *int `json:"n"` + CandidateCount *int `json:"candidateCount"` + Candidate_Count *int `json:"candidate_count"` + GenerationConfig *struct { + CandidateCount *int `json:"candidateCount"` + Candidate_Count *int `json:"candidate_count"` + } `json:"generationConfig"` + Generation_Config *struct { + CandidateCount *int `json:"candidateCount"` + Candidate_Count *int `json:"candidate_count"` + } `json:"generation_config"` + } `json:"request"` +} + +// ExtractExpectedChoices parses the request payload to extract the choice count parameter +// (top-level "n" for OpenAI or "generationConfig.candidateCount" for Gemini). +// Returns 1 if payload is empty, invalid, or choice count is omitted/<=0. +func ExtractExpectedChoices(payload []byte) int { + if len(payload) == 0 { + return 1 + } + var req choiceExtractionPayload + if err := json.Unmarshal(payload, &req); err != nil { + return 1 + } + maxChoices := 1 + updateMax := func(ptr *int) { + if ptr != nil && *ptr > maxChoices { + maxChoices = *ptr + } + } + updateMax(req.N) + updateMax(req.CandidateCount) + updateMax(req.Candidate_Count) + if req.GenerationConfig != nil { + updateMax(req.GenerationConfig.CandidateCount) + updateMax(req.GenerationConfig.Candidate_Count) + } + if req.Generation_Config != nil { + updateMax(req.Generation_Config.CandidateCount) + updateMax(req.Generation_Config.Candidate_Count) + } + if req.Request != nil { + updateMax(req.Request.N) + updateMax(req.Request.CandidateCount) + updateMax(req.Request.Candidate_Count) + if req.Request.GenerationConfig != nil { + updateMax(req.Request.GenerationConfig.CandidateCount) + updateMax(req.Request.GenerationConfig.Candidate_Count) + } + if req.Request.Generation_Config != nil { + updateMax(req.Request.Generation_Config.CandidateCount) + updateMax(req.Request.Generation_Config.Candidate_Count) + } + } + return maxChoices +} + +// StreamBootstrapDetector incrementally classifies a stream prefix without +// reparsing previously observed chunks. Its zero value is ready for use. +type StreamBootstrapDetector struct { + state streamBootstrapState +} + +// SetExpectedChoices sets the number of expected choices for multi-choice streams. +// When n <= 0, it defaults to 1. +func (d *StreamBootstrapDetector) SetExpectedChoices(n int) { + if d != nil { + d.state.setExpectedChoices(n) + } +} + +// SetRequestPayload parses the request payload to configure expected choice count. +func (d *StreamBootstrapDetector) SetRequestPayload(payload []byte) { + if d != nil { + d.state.setExpectedChoices(ExtractExpectedChoices(payload)) + } +} + +// Observe records an arbitrary stream byte fragment and reports whether +// buffered data should now be forwarded. It retains incomplete SSE lines across +// calls and forwards conservatively at the bootstrap byte limit. +func (d *StreamBootstrapDetector) Observe(payload []byte) bool { + if d == nil { + return true + } + 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() +} + +// StreamError returns the parsed in-band provider error detected so far, if any. +func (d *StreamBootstrapDetector) StreamError() error { + if d == nil { + return nil + } + return d.state.streamError() +} + +// 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 new file mode 100644 index 00000000000..126d934628a --- /dev/null +++ b/sdk/cliproxy/auth/empty_completion_formats_test.go @@ -0,0 +1,96 @@ +package auth + +import ( + "testing" +) + +// TestSupportedCompletionFormatsRecognized covers representative wire formats +// handled by empty-completion detection. Executor names document current users +// of each format; this manually maintained table is not an executor-registry +// completeness check. +// +// Each case lists executors that emit a given wire format plus a +// representative NON-empty chunk in that format (asserted recognized=true) and +// the corresponding empty-terminal variant (asserted empty). +// +// Documented exclusions (NOT in this table, by design): +// - codex-live: realtime bidirectional voice/media relay, not a text +// completion stream — empty-completion handling does not apply. +// - gemini-interactions: Gemini Live realtime relay (parts/role frames, not +// 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 + // neverEmpty documents formats whose terminal events are valid even with + // no output (existing callers rely on pass-through); the empty variant + // is then asserted to NOT be judged empty. + neverEmpty bool + }{ + { + // OpenAI chat-completions wire. Emitted by the OpenAI-compatible + // proxy executors (their requestToFormat is FormatOpenAI). + name: "openai-chat", + executors: []string{"kimi", "kiro", "kilo", "cursor", "github-copilot", "codebuddy", "gitlab", "qoder", "openai-compatibility"}, + nonEmpty: []byte("data: {\"choices\":[{\"delta\":{\"content\":\"hello\"},\"finish_reason\":null}]}\n\ndata: [DONE]\n\n"), + empty: []byte("data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":{\"completion_tokens\":0}}\n\ndata: [DONE]\n\n"), + }, + { + // 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"), + neverEmpty: true, + }, + { + // Anthropic Claude wire. Emitted by the Claude executor + // (requestToFormat is FormatClaude). + name: "claude", + executors: []string{"claude"}, + nonEmpty: []byte("data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"hello\"}}\n\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"}}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"), + empty: []byte("data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"}}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"), + }, + { + // Gemini wire (top-level candidates). Emitted by the Gemini-family + // executors (requestToFormat is FormatGemini). aistudio is listed + // here as its representative case: it emits the client-requested + // SDK format (body.toFormat), and Gemini is one of its valid + // outputs — all of which are recognized formats. + name: "gemini", + executors: []string{"gemini", "gemini-cli", "vertex", "aistudio"}, + nonEmpty: []byte("data: {\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"text\":\"hello\"}]},\"finishReason\":\"STOP\"}],\"usageMetadata\":{\"candidatesTokenCount\":5}}\n\n"), + empty: []byte("data: {\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[]},\"finishReason\":\"STOP\"}],\"usageMetadata\":{\"candidatesTokenCount\":0}}\n\n"), + }, + { + // Antigravity emits Gemini-shaped wire (nested response.candidates + // wrapper), recognized through the same Gemini predicate. + name: "antigravity", + executors: []string{"antigravity"}, + nonEmpty: []byte("data: {\"response\":{\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"text\":\"hello\"}]},\"finishReason\":\"STOP\"}]}}\n\n"), + empty: []byte("data: {\"response\":{\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[]},\"finishReason\":\"STOP\"}]}}\n\n"), + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if !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 tc.neverEmpty { + // Responses-API terminal frames pass through by contract (existing + // repo tests define them as valid completions even with no output). + if IsEmptyCompletionPayload(tc.empty) { + t.Fatalf("terminal variant for executors %v must pass through (never empty), but was judged empty", tc.executors) + } + } else if !IsEmptyCompletionPayload(tc.empty) { + t.Fatalf("empty-terminal variant for executors %v was not judged empty", tc.executors) + } + if 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 new file mode 100644 index 00000000000..84a619d6b68 --- /dev/null +++ b/sdk/cliproxy/auth/empty_completion_test.go @@ -0,0 +1,5178 @@ +package auth + +import ( + "context" + "errors" + "net/http" + "strings" + "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" +) + +// emptyCompletionTestExecutor returns a configurable payload per auth, allowing +// tests to make one auth produce an empty completion and another a real one. +type emptyCompletionTestExecutor struct { + executePayloads map[string][]byte // auth ID -> non-stream payload + streamPayloads map[string][][]byte // auth ID -> SSE chunk payloads + executeErr map[string]error // auth ID -> forced execute error + streamErr map[string]error // auth ID -> forced stream error + executeCalls map[string]int // auth ID -> call count (non-stream) + streamCalls map[string]int // auth ID -> call count (stream) + hook func(authID, kind string) + + // firstExecute records the first auth that was picked for a non-stream + // execution, so tests can deterministically wire the empty payload to it + // regardless of global selector state. + firstExecute string + // firstStream records the first auth picked for a stream execution. + firstStream string + + // emptyPayload/contentPayload override the default first-auth empty payload + // and subsequent-auth content payload (used to exercise non-OpenAI formats). + emptyPayload []byte + contentPayload []byte + + // emptyStreamPayload/contentStreamPayload override the default first-auth + // empty stream and subsequent-auth content stream (used to exercise + // non-OpenAI stream formats). + emptyStreamPayload [][]byte + contentStreamPayload [][]byte + leaveStreamOpen bool +} + +func (e *emptyCompletionTestExecutor) Identifier() string { return "claude" } + +func (*emptyCompletionTestExecutor) ShouldPrepareRequestAuth(*Auth) bool { return false } + +func (e *emptyCompletionTestExecutor) PrepareRequestAuth(ctx context.Context, auth *Auth) (*Auth, error) { + return auth, nil +} + +func (e *emptyCompletionTestExecutor) Execute(ctx context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + e.executeCalls[auth.ID]++ + if e.firstExecute == "" { + e.firstExecute = auth.ID + } + if err := e.executeErr[auth.ID]; err != nil { + return cliproxyexecutor.Response{}, err + } + // 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 len(e.executePayloads) == 0 && e.firstExecute == auth.ID { + empty := e.emptyPayload + if len(empty) == 0 { + empty = []byte(`{"choices":[{"message":{"content":""},"finish_reason":"stop"}],"usage":{"completion_tokens":0}}`) + } + return cliproxyexecutor.Response{Payload: empty}, nil + } + if p, ok := e.executePayloads[auth.ID]; ok { + return cliproxyexecutor.Response{Payload: p}, nil + } + content := e.contentPayload + if len(content) == 0 { + content = []byte(`{"choices":[{"message":{"content":"real"},"finish_reason":"stop"}]}`) + } + return cliproxyexecutor.Response{Payload: content}, nil +} + +func (e *emptyCompletionTestExecutor) CountTokens(ctx context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{Payload: []byte("ok")}, nil +} + +func (e *emptyCompletionTestExecutor) ExecuteStream(ctx context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + e.streamCalls[auth.ID]++ + if e.firstStream == "" { + e.firstStream = auth.ID + } + if err := e.streamErr[auth.ID]; err != nil { + return nil, err + } + // When the test pre-wires explicit payloads (e.g. the thinking-then-content + // positive control), honor them. Otherwise force the first auth to stream an + // empty completion and subsequent auths to stream real content, so rotation + // tests are deterministic regardless of global selector state. + if len(e.streamPayloads) == 0 && e.firstStream == auth.ID { + empty := e.emptyStreamPayload + if len(empty) == 0 { + empty = [][]byte{ + []byte("data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n"), + []byte("data: [DONE]\n\n"), + } + } + chunks := make(chan cliproxyexecutor.StreamChunk, len(empty)) + for _, p := range empty { + chunks <- cliproxyexecutor.StreamChunk{Payload: p} + } + if !e.leaveStreamOpen { + close(chunks) + } + return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil + } + if payloads, ok := e.streamPayloads[auth.ID]; ok && len(payloads) > 0 { + chunks := make(chan cliproxyexecutor.StreamChunk, len(payloads)) + for _, p := range payloads { + chunks <- cliproxyexecutor.StreamChunk{Payload: p} + } + close(chunks) + return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil + } + content := e.contentStreamPayload + if len(content) == 0 { + content = [][]byte{ + []byte("data: {\"choices\":[{\"delta\":{\"content\":\"hello\"},\"finish_reason\":null}]}\n\n"), + []byte("data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n"), + []byte("data: [DONE]\n\n"), + } + } + chunks := make(chan cliproxyexecutor.StreamChunk, len(content)) + for _, p := range content { + chunks <- cliproxyexecutor.StreamChunk{Payload: p} + } + close(chunks) + return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil +} + +func (e *emptyCompletionTestExecutor) Refresh(ctx context.Context, auth *Auth) (*Auth, error) { + return auth, nil +} + +func (*emptyCompletionTestExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, errors.New("not implemented") +} + +// newEmptyCompletionTestManager registers two auths for the same model and +// returns the manager, the auth IDs, the model name, and a result-capture hook. +func newEmptyCompletionTestManager(t *testing.T, executor *emptyCompletionTestExecutor) (*Manager, []string, string, *resultCaptureHook) { + t.Helper() + model := "empty-completion-model-" + uuid.NewString() + capture := &resultCaptureHook{} + manager := NewManager(nil, nil, capture) + manager.SetRetryConfig(0, 0, 0) + manager.RegisterExecutor(executor) + + var ids []string + for i := 0; i < 2; i++ { + auth := &Auth{ + ID: "empty-completion-auth-" + uuid.NewString(), + Provider: "claude", + Attributes: map[string]string{"auth_kind": "oauth"}, + Metadata: map[string]any{ + "access_token": "access-token", + "refresh_token": "refresh-token", + "request_retry": float64(0), + }, + } + registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth.ID) }) + if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { + t.Fatalf("Register(%s) error = %v", auth.ID, errRegister) + } + ids = append(ids, auth.ID) + } + return manager, ids, model, capture +} + +func TestEmptyCompletionPredicate(t *testing.T) { + cases := []struct { + name string + payload []byte + expected bool + }{ + { + name: "openai sse empty", + payload: []byte("data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n"), + expected: true, + }, + { + name: "openai sse thinking then content is not empty", + payload: []byte("data: {\"choices\":[{\"delta\":{\"content\":\"\"},\"finish_reason\":null}]}\n\ndata: {\"choices\":[{\"delta\":{\"content\":\"hello\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n"), + expected: false, + }, + { + name: "whitespace only zero tokens is empty", + payload: []byte("data: {\"choices\":[{\"delta\":{\"content\":\" \"},\"finish_reason\":\"stop\"}],\"usage\":{\"completion_tokens\":0}}\n\ndata: [DONE]\n\n"), + expected: true, + }, + { + name: "non zero tokens is not empty", + payload: []byte("data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":{\"completion_tokens\":5}}\n\ndata: [DONE]\n\n"), + expected: false, + }, + { + name: "tool calls are not empty", + payload: []byte("data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"id\":\"x\",\"function\":{\"name\":\"lookup\"}}]},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n"), + expected: false, + }, + { + name: "openai sse semantically empty tool_calls id only", + payload: []byte("data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"id\":\"x\"}]},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n"), + 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: "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 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 sse semantically empty tool_calls empty fields", + 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 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"), + expected: false, + }, + { + name: "openai sse refusal is not credential empty", + payload: []byte("data: {\"choices\":[{\"delta\":{\"refusal\":\"I cannot help with that\"},\"finish_reason\":null}]}\n\ndata: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n"), + expected: false, + }, + { + name: "unterminated is empty", + payload: []byte("data: {\"choices\":[{\"delta\":{},\"finish_reason\":null}]}\n\n"), + expected: true, + }, + { + name: "claude sse message_stop without end_turn is empty", + payload: []byte("event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"), + expected: true, + }, + { + name: "unrecognized format is not empty", + payload: []byte("data: {\"unknown_payload\":true}\n\n"), + expected: false, + }, + { + name: "unknown sse data followed by done is not empty", + 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"), + expected: true, + }, + { + name: "non stream empty json", + payload: []byte(`{"choices":[{"message":{"content":""},"finish_reason":"stop"}],"usage":{"completion_tokens":0}}`), + expected: true, + }, + { + name: "non stream content is not empty", + payload: []byte(`{"choices":[{"message":{"content":"hi"},"finish_reason":"stop"}]}`), + expected: false, + }, + { + name: "openai non-stream content_filter is not empty", + payload: []byte(`{"choices":[{"message":{"content":""},"finish_reason":"content_filter"}]}`), + expected: false, + }, + { + name: "openai non-stream length is not empty", + payload: []byte(`{"choices":[{"message":{"content":""},"finish_reason":"length"}]}`), + expected: false, + }, + { + name: "openai sse content_filter is not empty", + payload: []byte("data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"content_filter\"}]}\n\ndata: [DONE]\n\n"), + expected: false, + }, + { + name: "openai sse length is not empty", + payload: []byte("data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"length\"}]}\n\ndata: [DONE]\n\n"), + expected: false, + }, + { + name: "non stream reasoning only is not empty", + payload: []byte(`{"choices":[{"message":{"reasoning_content":"thinking"},"finish_reason":"stop"}]}`), + expected: false, + }, + { + name: "openai non-stream refusal is not credential empty", + payload: []byte(`{"choices":[{"message":{"content":"","refusal":"I cannot help with that"},"finish_reason":"stop"}],"usage":{"completion_tokens":0}}`), + expected: false, + }, + { + name: "claude non-stream empty message is empty", + payload: []byte(`{"type":"message","id":"msg_1","role":"assistant","content":[],"stop_reason":"end_turn","usage":{"input_tokens":10,"output_tokens":0}}`), + expected: true, + }, + { + name: "claude non-stream empty message without usage is empty", + payload: []byte(`{"type":"message","id":"msg_1","role":"assistant","content":[],"stop_reason":"end_turn"}`), + expected: true, + }, + { + name: "claude non-stream max_tokens with empty content is not empty", + payload: []byte(`{"type":"message","id":"msg_1","role":"assistant","content":[],"stop_reason":"max_tokens"}`), + expected: false, + }, + { + name: "claude non-stream refusal with empty content is not empty", + payload: []byte(`{"type":"message","id":"msg_1","role":"assistant","content":[],"stop_reason":"refusal"}`), + expected: false, + }, + { + name: "claude sse max_tokens stream is not empty", + payload: []byte("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"stop_reason\":null,\"usage\":{\"output_tokens\":0}}}\n\nevent: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"max_tokens\"},\"usage\":{\"output_tokens\":0}}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"), + expected: false, + }, + { + name: "claude sse refusal stream is not empty", + payload: []byte("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"stop_reason\":null,\"usage\":{\"output_tokens\":0}}}\n\nevent: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"refusal\"},\"usage\":{\"output_tokens\":0}}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"), + expected: false, + }, + { + name: "claude unknown content block type with empty text does not flip hasContent", + payload: []byte(`{"type":"message","id":"msg_1","role":"assistant","content":[{"type":"unknown_custom_block","text":""}],"stop_reason":"end_turn"}`), + expected: true, + }, + { + name: "claude non-stream with tool_use is not empty", + payload: []byte(`{"type":"message","id":"msg_1","role":"assistant","content":[{"type":"tool_use","id":"call_1","name":"get_weather","input":{}}],"stop_reason":"tool_use","usage":{"input_tokens":10,"output_tokens":5}}`), + expected: false, + }, + { + name: "claude non-stream thinking-block-only is not empty", + payload: []byte(`{"type":"message","id":"msg_1","role":"assistant","content":[{"type":"thinking","thinking":"let me think","signature":"sig"}],"stop_reason":"end_turn","usage":{"input_tokens":10,"output_tokens":5}}`), + expected: false, + }, + { + name: "claude non-stream with text content is not empty", + payload: []byte(`{"type":"message","id":"msg_1","role":"assistant","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","usage":{"input_tokens":10,"output_tokens":1}}`), + expected: false, + }, + { + name: "claude non-stream max_tokens is not credential empty", + payload: []byte(`{"type":"message","content":[],"stop_reason":"max_tokens","usage":{"output_tokens":0}}`), + expected: false, + }, + { + name: "claude sse refusal is not credential empty", + payload: []byte("event: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"refusal\"},\"usage\":{\"output_tokens\":0}}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"), + expected: false, + }, + { + name: "claude sse empty stream is empty", + payload: []byte("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"stop_reason\":null,\"usage\":{\"output_tokens\":0}}}\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: true, + }, + { + name: "gemini non-stream empty candidates is empty", + payload: []byte(`{"candidates":[{"content":{"role":"model","parts":[]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":0}}`), + expected: true, + }, + { + name: "gemini non-stream whitespace parts is empty", + payload: []byte(`{"candidates":[{"content":{"role":"model","parts":[{"text":" "}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":0}}`), + expected: true, + }, + { + name: "gemini non-stream without usage is empty", + 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 inlineData empty data scaffold is empty", + payload: []byte(`{"candidates":[{"content":{"role":"model","parts":[{"inlineData":{"mimeType":"image/png","data":""}}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":0}}`), + expected: true, + }, + { + name: "gemini non-stream inlineData real data is not empty", + payload: []byte(`{"candidates":[{"content":{"role":"model","parts":[{"inlineData":{"mimeType":"image/png","data":"aW1n"}}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":0}}`), + expected: false, + }, + { + name: "gemini non-stream fileData empty fileUri scaffold is empty", + payload: []byte(`{"candidates":[{"content":{"role":"model","parts":[{"fileData":{"mimeType":"application/pdf","fileUri":""}}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":0}}`), + expected: true, + }, + { + name: "gemini non-stream fileData real fileUri is not empty", + payload: []byte(`{"candidates":[{"content":{"role":"model","parts":[{"fileData":{"mimeType":"application/pdf","fileUri":"gs://bucket/doc.pdf"}}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":0}}`), + expected: false, + }, + { + 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"}]}`), + expected: false, + }, + { + name: "gemini non-stream with thought part is not empty", + 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"), + expected: true, + }, + { + name: "gemini blocked safety is not empty", + payload: []byte(`{"candidates":[{"content":{"role":"model","parts":[]},"finishReason":"SAFETY"}]}`), + expected: false, + }, + { + name: "gemini blocked recitation is not empty", + payload: []byte(`{"candidates":[{"content":{"role":"model","parts":[]},"finishReason":"RECITATION"}]}`), + expected: false, + }, + { + name: "gemini max tokens with empty content is not empty", + payload: []byte(`{"candidates":[{"content":{"role":"model","parts":[]},"finishReason":"MAX_TOKENS"}]}`), + expected: false, + }, + { + name: "gemini sse blocked safety closed by done is not empty", + payload: []byte("data: {\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[]},\"finishReason\":\"SAFETY\"}]}\n\ndata: [DONE]\n\n"), + expected: false, + }, + { + name: "gemini prompt feedback safety is not credential empty", + payload: []byte(`{"promptFeedback":{"blockReason":"SAFETY"},"candidates":[]}`), + expected: false, + }, + { + name: "gemini sse prompt feedback safety is not credential empty", + payload: []byte("data: {\"promptFeedback\":{\"blockReason\":\"SAFETY\"},\"candidates\":[]}\n\ndata: [DONE]\n\n"), + expected: false, + }, + { + name: "openai sse empty choices array then done is empty", + payload: []byte("data: {\"choices\":[]}\n\ndata: [DONE]\n\n"), + expected: true, + }, + { + name: "openai sse content_filter with empty content is not empty", + payload: []byte("data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"content_filter\"}]}\n\ndata: [DONE]\n\n"), + expected: false, + }, + { + name: "openai sse length with empty content is not empty", + payload: []byte("data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"length\"}]}\n\ndata: [DONE]\n\n"), + expected: false, + }, + { + name: "openai non-stream content_filter with empty content is not empty", + payload: []byte(`{"choices":[{"message":{"content":""},"finish_reason":"content_filter"}],"usage":{"completion_tokens":0}}`), + expected: false, + }, + { + name: "openai non-stream length with empty content is not empty", + payload: []byte(`{"choices":[{"message":{"content":""},"finish_reason":"length"}],"usage":{"completion_tokens":0}}`), + expected: false, + }, + { + name: "openai non-stream stop empty is empty", + payload: []byte(`{"choices":[{"message":{"content":""},"finish_reason":"stop"}],"usage":{"completion_tokens":0}}`), + expected: true, + }, + { + 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: "openai legacy non-stream whitespace text is empty", + payload: []byte(`{"choices":[{"text":" ","finish_reason":"stop"}],"usage":{"completion_tokens":0}}`), + expected: true, + }, + { + 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}}`), + expected: false, + }, + { + name: "codex responses-api sse with function_call is not empty", + payload: []byte("data: {\"type\":\"response.output_item.done\",\"output\":{\"type\":\"function_call\",\"name\":\"get_weather\",\"arguments\":\"{}\",\"call_id\":\"call_1\"}}\n\ndata: {\"type\":\"response.completed\",\"response\":{\"id\":\"r\",\"status\":\"completed\",\"output\":[{\"type\":\"function_call\",\"name\":\"get_weather\",\"arguments\":\"{}\",\"call_id\":\"call_1\"}],\"usage\":{\"output_tokens\":5}}}\n\ndata: [DONE]\n\n"), + expected: false, + }, + { + name: "codex responses-api non-stream with custom_tool_call is not empty", + payload: []byte(`{"object":"response","status":"completed","output":[{"type":"custom_tool_call","name":"shell","input":"pwd"}],"usage":{"output_tokens":0}}`), + expected: false, + }, + { + name: "codex responses-api sse with custom_tool_call is not empty", + payload: []byte("data: {\"type\":\"response.output_item.done\",\"item\":{\"type\":\"custom_tool_call\",\"name\":\"shell\",\"input\":\"pwd\"}}\n\ndata: {\"type\":\"response.completed\",\"response\":{\"status\":\"completed\",\"output\":[],\"usage\":{\"output_tokens\":0}}}\n\ndata: [DONE]\n\n"), + expected: false, + }, + { + name: "codex responses-api non-stream with image_generation_call is not empty", + payload: []byte(`{"object":"response","status":"completed","output":[{"type":"image_generation_call","status":"completed","result":"image-data"}],"usage":{"output_tokens":0}}`), + expected: false, + }, + { + name: "codex responses-api sse with image_generation_call is not empty", + payload: []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"), + expected: false, + }, + { + name: "codex responses-api sse with reasoning item is not empty", + payload: []byte("data: {\"type\":\"response.output_item.done\",\"item\":{\"type\":\"reasoning\",\"status\":\"completed\",\"summary\":[]}}\n\ndata: {\"type\":\"response.completed\",\"response\":{\"status\":\"completed\",\"output\":[],\"usage\":{\"output_tokens\":0}}}\n\ndata: [DONE]\n\n"), + expected: false, + }, + { + name: "codex responses-api non-stream refusal is not credential empty", + payload: []byte(`{"object":"response","status":"completed","output":[{"type":"message","role":"assistant","content":[{"type":"refusal","refusal":"I cannot help with that"}]}],"usage":{"output_tokens":0}}`), + expected: false, + }, + { + name: "codex responses-api sse refusal item is not credential empty", + payload: []byte("data: {\"type\":\"response.output_item.done\",\"item\":{\"type\":\"message\",\"role\":\"assistant\",\"content\":[{\"type\":\"refusal\",\"refusal\":\"I cannot help with that\"}]}}\n\ndata: {\"type\":\"response.completed\",\"response\":{\"status\":\"completed\",\"output\":[],\"usage\":{\"output_tokens\":0}}}\n\ndata: [DONE]\n\n"), + expected: false, + }, + { + name: "codex responses-api non-stream incomplete is not credential empty", + payload: []byte(`{"object":"response","status":"incomplete","output":[],"usage":{"output_tokens":0}}`), + expected: false, + }, + { + name: "codex responses-api sse incomplete is not credential empty", + payload: []byte("data: {\"type\":\"response.incomplete\",\"response\":{\"status\":\"incomplete\",\"output\":[],\"usage\":{\"output_tokens\":0}}}\n\ndata: [DONE]\n\n"), + expected: false, + }, + { + name: "codex responses-api sse failed is not credential empty", + payload: []byte("data: {\"type\":\"response.failed\",\"response\":{\"status\":\"failed\",\"output\":[],\"usage\":{\"output_tokens\":0}}}\n\ndata: [DONE]\n\n"), + expected: false, + }, + { + name: "gemini non-stream empty candidates array is empty", + payload: []byte(`{"candidates":[],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":0}}`), + expected: true, + }, + { + name: "image_generation_call stream empty result is empty", + payload: []byte("data: {\"type\":\"response.output_item.done\",\"item\":{\"type\":\"image_generation_call\",\"id\":\"call_1\",\"call_id\":\"call_1\",\"result\":\"\"}}\n\ndata: [DONE]\n\n"), + expected: true, + }, + { + name: "image_generation_call stream result is not empty", + payload: []byte("data: {\"type\":\"response.output_item.done\",\"item\":{\"type\":\"image_generation_call\",\"id\":\"call_1\",\"call_id\":\"call_1\",\"result\":\"{\\\"image_ref\\\":\\\"file-abc\\\"}\"}}\n\ndata: [DONE]\n\n"), + expected: false, + }, + { + name: "image_generation_call stream omitted result is empty", + payload: []byte("data: {\"type\":\"response.output_item.done\",\"item\":{\"type\":\"image_generation_call\",\"id\":\"call_1\",\"call_id\":\"call_1\"}}\n\ndata: [DONE]\n\n"), + expected: true, + }, + { + name: "web_search_call stream populated action is not empty", + payload: []byte("data: {\"type\":\"response.output_item.done\",\"item\":{\"type\":\"web_search_call\",\"status\":\"completed\",\"action\":{\"type\":\"search\",\"query\":\"weather\"}}}\n\ndata: [DONE]\n\n"), + expected: false, + }, + { + name: "computer_call stream populated action is not empty", + payload: []byte("data: {\"type\":\"response.output_item.done\",\"item\":{\"type\":\"computer_call\",\"status\":\"completed\",\"action\":{\"type\":\"click\",\"x\":100,\"y\":200}}}\n\ndata: [DONE]\n\n"), + expected: false, + }, + { + name: "web_search_call stream id only and empty action object is empty", + payload: []byte("data: {\"type\":\"response.output_item.done\",\"item\":{\"id\":\"call_1\",\"call_id\":\"call_1\",\"type\":\"web_search_call\",\"status\":\"in_progress\",\"action\":{}}}\n\ndata: [DONE]\n\n"), + expected: true, + }, + { + name: "web_search_call stream id only and no action is empty", + payload: []byte("data: {\"type\":\"response.output_item.done\",\"item\":{\"id\":\"call_1\",\"call_id\":\"call_1\",\"type\":\"web_search_call\",\"status\":\"in_progress\"}}\n\ndata: [DONE]\n\n"), + expected: true, + }, + { + name: "web_search_call non-stream populated results is not empty", + payload: []byte(`{"object":"response","status":"in_progress","output":[{"type":"web_search_call","results":[{"url":"https://example.com"}]}],"usage":{"output_tokens":0}}`), + expected: false, + }, + { + name: "web_search_call non-stream empty results array is empty", + payload: []byte(`{"object":"response","status":"in_progress","output":[{"type":"web_search_call","id":"call_1","call_id":"call_1","results":[]}],"usage":{"output_tokens":0}}`), + expected: true, + }, + { + name: "web_search_call non-stream null results is empty", + payload: []byte(`{"object":"response","status":"in_progress","output":[{"type":"web_search_call","id":"call_1","call_id":"call_1","results":null}],"usage":{"output_tokens":0}}`), + expected: true, + }, + { + name: "web_search_call stream populated results is not empty", + payload: []byte("data: {\"type\":\"response.output_item.done\",\"item\":{\"type\":\"web_search_call\",\"status\":\"completed\",\"results\":[{\"url\":\"https://example.com\"}]}}\n\ndata: [DONE]\n\n"), + expected: false, + }, + { + name: "web_search_call stream id only and empty results array is empty", + payload: []byte("data: {\"type\":\"response.output_item.done\",\"item\":{\"id\":\"call_1\",\"call_id\":\"call_1\",\"type\":\"web_search_call\",\"status\":\"in_progress\",\"results\":[]}}\n\ndata: [DONE]\n\n"), + expected: true, + }, + { + name: "web_search_call stream id only and empty results object is empty", + payload: []byte("data: {\"type\":\"response.output_item.done\",\"item\":{\"id\":\"call_1\",\"call_id\":\"call_1\",\"type\":\"web_search_call\",\"status\":\"in_progress\",\"results\":{}}}\n\ndata: [DONE]\n\n"), + expected: true, + }, + { + name: "web_search_call stream id only and null results is empty", + payload: []byte("data: {\"type\":\"response.output_item.done\",\"item\":{\"id\":\"call_1\",\"call_id\":\"call_1\",\"type\":\"web_search_call\",\"status\":\"in_progress\",\"results\":null}}\n\ndata: [DONE]\n\n"), + expected: true, + }, + { + name: "gemini sse empty candidates array is empty", + payload: []byte("data: {\"candidates\":[],\"usageMetadata\":{\"candidatesTokenCount\":0}}\n\n"), + expected: true, + }, + { + name: "interactions json empty steps is empty", + payload: []byte(`{"id":"interaction_1","object":"interaction","status":"completed","steps":[]}`), + expected: true, + }, + { + name: "interactions json empty steps zero usage is empty", + payload: []byte(`{"id":"interaction_1","object":"interaction","status":"completed","steps":[],"usage":{"output_tokens":0,"total_output_tokens":0}}`), + expected: true, + }, + { + name: "interactions json with model_output text is not empty", + payload: []byte(`{"id":"interaction_1","object":"interaction","status":"completed","steps":[{"type":"model_output","content":[{"type":"text","text":"hello"}]}]}`), + expected: false, + }, + { + name: "interactions json with function_call is not empty", + payload: []byte(`{"id":"interaction_1","object":"interaction","status":"completed","steps":[{"type":"function_call","name":"get_weather","arguments":{"location":"Paris"}}]}`), + expected: false, + }, + { + name: "interactions json with output tokens is not empty", + payload: []byte(`{"id":"interaction_1","object":"interaction","status":"completed","steps":[],"usage":{"output_tokens":5}}`), + expected: false, + }, + { + name: "interactions sse stream scaffold and empty completion is empty", + payload: []byte("event: interaction.created\ndata: {\"event_type\":\"interaction.created\",\"interaction\":{\"id\":\"int_1\",\"object\":\"interaction\",\"status\":\"in_progress\"}}\n\nevent: step.start\ndata: {\"event_type\":\"step.start\",\"index\":0,\"step\":{\"type\":\"model_output\"}}\n\nevent: interaction.completed\ndata: {\"event_type\":\"interaction.completed\",\"interaction\":{\"id\":\"int_1\",\"status\":\"completed\",\"steps\":[],\"usage\":{\"output_tokens\":0}}}\n\n"), + expected: true, + }, + { + name: "interactions sse stream with step delta is not empty", + payload: []byte("event: interaction.created\ndata: {\"event_type\":\"interaction.created\",\"interaction\":{\"id\":\"int_1\",\"object\":\"interaction\",\"status\":\"in_progress\"}}\n\nevent: step.delta\ndata: {\"event_type\":\"step.delta\",\"index\":0,\"delta\":{\"type\":\"text\",\"text\":\"hello\"}}\n\nevent: interaction.completed\ndata: {\"event_type\":\"interaction.completed\",\"interaction\":{\"id\":\"int_1\",\"status\":\"completed\",\"steps\":[],\"usage\":{\"output_tokens\":1}}}\n\n"), + expected: false, + }, + { + name: "interactions json with media data is not empty", + payload: []byte(`{"id":"interaction_1","object":"interaction","status":"completed","steps":[{"type":"model_output","content":[{"type":"image","mime_type":"image/png","data":"iVBORw0KGgo="}]}]}`), + expected: false, + }, + { + name: "interactions json with file_uri is not empty", + payload: []byte(`{"id":"interaction_1","object":"interaction","status":"completed","steps":[{"type":"model_output","content":[{"type":"document","file_uri":"files/123"}]}]}`), + expected: false, + }, + { + name: "interactions sse stream with media delta is not empty", + payload: []byte("event: interaction.created\ndata: {\"event_type\":\"interaction.created\",\"interaction\":{\"id\":\"int_1\",\"object\":\"interaction\",\"status\":\"in_progress\"}}\n\nevent: step.delta\ndata: {\"event_type\":\"step.delta\",\"index\":0,\"delta\":{\"type\":\"content\",\"content\":{\"type\":\"image\",\"data\":\"iVBORw0KGgo=\"}}}\n\nevent: interaction.completed\ndata: {\"event_type\":\"interaction.completed\",\"interaction\":{\"id\":\"int_1\",\"status\":\"completed\",\"steps\":[],\"usage\":{\"output_tokens\":0}}}\n\n"), + expected: false, + }, + { + name: "interactions sse stream ending in bare finish with zero output is empty", + payload: []byte("event: interaction.created\ndata: {\"event_type\":\"interaction.created\",\"interaction\":{\"id\":\"int_1\",\"object\":\"interaction\",\"status\":\"in_progress\"}}\n\nevent: step.start\ndata: {\"event_type\":\"step.start\",\"index\":0,\"step\":{\"type\":\"model_output\"}}\n\nevent: finish\ndata: {\"event_type\":\"finish\",\"metadata\":{\"total_usage\":{\"total_output_tokens\":0}}}\n\n"), + expected: true, + }, + { + name: "interactions sse stream ending in bare finish with output is not empty", + payload: []byte("event: interaction.created\ndata: {\"event_type\":\"interaction.created\",\"interaction\":{\"id\":\"int_1\",\"object\":\"interaction\",\"status\":\"in_progress\"}}\n\nevent: step.start\ndata: {\"event_type\":\"step.start\",\"index\":0,\"step\":{\"type\":\"model_output\"}}\n\nevent: finish\ndata: {\"event_type\":\"finish\",\"metadata\":{\"total_usage\":{\"total_output_tokens\":7}}}\n\n"), + 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 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 TestStreamBootstrapDetectorClaudePing(t *testing.T) { + var detector StreamBootstrapDetector + // Standard Claude keep-alive prefix; ping is non-output metadata and must + // not poison the bootstrap detector with sawUnknownData. + if detector.Observe([]byte("event: ping\ndata: {\"type\":\"ping\"}\n\n")) { + t.Fatal("Observe() forwarded after Claude ping keep-alive") + } + // A terminal-but-empty Claude message after the ping must still be + // withheld as an empty completion instead of bypassing failover. + 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 TestStreamBootstrapDetectorSSEMetadataFields(t *testing.T) { + var detector StreamBootstrapDetector + 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: {\"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 TestStreamBootstrapDetectorMetadataOnlyEOF(t *testing.T) { + t.Run("comments and keepalive only then EOF classifies as empty", func(t *testing.T) { + var detector StreamBootstrapDetector + if detector.Observe([]byte(": keep-alive\n\n")) { + t.Fatal("Observe() forwarded keep-alive comment") + } + if !detector.Finish() { + t.Fatal("Finish() = false, want metadata-only stream recognized as empty completion at EOF") + } + }) + + t.Run("id and retry metadata only then EOF classifies as empty", func(t *testing.T) { + var detector StreamBootstrapDetector + if detector.Observe([]byte("id: evt_12345\nretry: 5000\n\n")) { + t.Fatal("Observe() forwarded id/retry metadata") + } + if !detector.Finish() { + t.Fatal("Finish() = false, want id/retry-only stream recognized as empty completion at EOF") + } + }) + + t.Run("claude ping only then EOF classifies as empty", func(t *testing.T) { + var detector StreamBootstrapDetector + 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("data-bearing stream still forwards and does not classify as empty", func(t *testing.T) { + var detector StreamBootstrapDetector + if !detector.Observe([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"hello\"}}]}\n\n")) { + t.Fatal("Observe() = false, want data-bearing stream to forward") + } + if detector.Finish() { + t.Fatal("Finish() = true, want data-bearing stream not recognized as empty completion") + } + }) + + 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 TestStreamBootstrapStateEvaluatesNewlineTerminatedJSONFrameImmediately(t *testing.T) { + var state streamBootstrapState + state.observe([]byte("{\"error\":{\"message\":\"quota exceeded\",\"code\":429}}\n")) + streamErr := state.streamError() + if streamErr == nil { + t.Fatal("bootstrap did not surface a provider error carried by a newline-terminated raw JSON frame") + } + if !strings.Contains(streamErr.Error(), "quota exceeded") { + t.Fatalf("bootstrap stream error = %q, want it to carry the provider message", streamErr.Error()) + } +} + +func TestStreamPayloadErrorDetectorEvaluatesNewlineTerminatedJSONFrameImmediately(t *testing.T) { + var d streamPayloadErrorDetector + streamErr := d.Observe([]byte("{\"error\":{\"message\":\"quota exceeded\",\"code\":429}}\n")) + if streamErr == nil { + t.Fatal("payload detector did not surface a provider error carried by a newline-terminated raw JSON frame") + } + if !strings.Contains(streamErr.Message, "quota exceeded") { + t.Fatalf("payload detector stream error message = %q, want it to carry the provider message", streamErr.Message) + } +} + +func TestStreamBootstrapStateBuffersPrettyPrintedJSONFrame(t *testing.T) { + var state streamBootstrapState + frame := "{\n \"error\": {\n \"message\": \"quota exceeded\",\n \"code\": 429\n }\n}\n" + if state.observe([]byte(frame)) { + t.Fatal("bootstrap forwarded a pretty-printed raw JSON frame instead of buffering it to completion") + } + state.finish() + streamErr := state.streamError() + if streamErr == nil { + t.Fatal("bootstrap did not surface a provider error carried by a pretty-printed raw JSON frame") + } + if !strings.Contains(streamErr.Error(), "quota exceeded") { + t.Fatalf("bootstrap stream error = %q, want it to carry the provider message", streamErr.Error()) + } +} + +func TestStreamBootstrapStateBuffersPrettyPrintedJSONFrameWithBlankLine(t *testing.T) { + var state streamBootstrapState + frame := "{\n\n \"error\": {\n\n \"message\": \"quota exceeded\",\n \"code\": 429\n }\n}\n" + if state.observe([]byte(frame)) { + t.Fatal("bootstrap forwarded a pretty-printed raw JSON frame with a blank line instead of buffering it to completion") + } + state.finish() + streamErr := state.streamError() + if streamErr == nil { + t.Fatal("bootstrap did not surface a provider error carried by a pretty-printed raw JSON frame with a blank line") + } + if !strings.Contains(streamErr.Error(), "quota exceeded") { + t.Fatalf("bootstrap stream error = %q, want it to carry the provider message", streamErr.Error()) + } +} + +func TestStreamPayloadErrorDetectorBuffersPrettyPrintedJSONFrame(t *testing.T) { + var d streamPayloadErrorDetector + frame := "{\n \"error\": {\n \"message\": \"quota exceeded\",\n \"code\": 429\n }\n}\n" + d.Observe([]byte(frame)) + streamErr := d.Finish() + if streamErr == nil { + t.Fatal("payload detector did not surface a provider error carried by a pretty-printed raw JSON frame") + } + if !strings.Contains(streamErr.Message, "quota exceeded") { + t.Fatalf("payload detector stream error message = %q, want it to carry the provider message", streamErr.Message) + } +} + +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(":"), + []byte(" ping\n"), + []byte("data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n"), + []byte("data: [DONE]\n\n"), + } + for i, fragment := range fragments { + if detector.Observe(fragment) { + t.Fatalf("Observe(fragment %d) forwarded comment-prefixed empty stream", i) + } + } + }) + + t.Run("comment then opaque line forwards", func(t *testing.T) { + var detector StreamBootstrapDetector + if detector.Observe([]byte(":")) || detector.Observe([]byte(" heartbeat\n")) { + t.Fatal("Observe() forwarded a valid split SSE comment") + } + if !detector.Observe([]byte("opaque-provider-payload\n")) { + t.Fatal("Observe() buffered a definitely non-SSE line after a comment") + } + }) +} + +func TestReadStreamBootstrapWithholdsSplitClaudeEmptyCompletion(t *testing.T) { + 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"), + } + chunks := make(chan cliproxyexecutor.StreamChunk, len(fragments)) + for _, fragment := range fragments { + chunks <- cliproxyexecutor.StreamChunk{Payload: fragment} + } + close(chunks) + + buffered, closed, err := readStreamBootstrap(context.Background(), chunks) + if err != nil { + t.Fatalf("readStreamBootstrap() error = %v", err) + } + if !closed { + t.Fatal("readStreamBootstrap() forwarded empty Claude stream") + } + if !isEmptyCompletion(buffered) { + t.Fatal("split Claude stream was not classified as empty at close") + } +} + +func TestExecuteEmptyCompletionRotatesAuth(t *testing.T) { + executor := &emptyCompletionTestExecutor{ + executePayloads: map[string][]byte{}, + executeCalls: map[string]int{}, + } + manager, ids, model, capture := newEmptyCompletionTestManager(t, executor) + + resp, err := manager.Execute(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}) + if err != nil { + t.Fatalf("Execute() error = %v", err) + } + assertRotatesToContent(t, ids, executor.firstExecute, string(resp.Payload), "real", capture) +} + +func TestExecuteStreamEmptyCompletionRotatesAuth(t *testing.T) { + executor := &emptyCompletionTestExecutor{ + streamPayloads: map[string][][]byte{}, + streamCalls: map[string]int{}, + } + manager, ids, model, capture := newEmptyCompletionTestManager(t, executor) + + stream, err := manager.ExecuteStream(context.Background(), []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 TestExecuteStreamEmptyGeminiStreamRotatesAuth(t *testing.T) { + executor := &emptyCompletionTestExecutor{ + streamPayloads: map[string][][]byte{}, + streamCalls: map[string]int{}, + emptyStreamPayload: [][]byte{ + []byte("data: {\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"text\":\"\"}]},\"finishReason\":\"STOP\"}],\"usageMetadata\":{\"candidatesTokenCount\":0}}\n\n"), + }, + contentStreamPayload: [][]byte{ + []byte("data: {\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"functionCall\":{\"name\":\"search\",\"args\":{}}}]},\"finishReason\":\"STOP\"}],\"usageMetadata\":{\"candidatesTokenCount\":5}}\n\n"), + }, + } + manager, ids, model, capture := newEmptyCompletionTestManager(t, executor) + + stream, err := manager.ExecuteStream(context.Background(), []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(), "functionCall", capture) +} + +func TestExecuteStreamEmptyRawJSONChunksRotateAuth(t *testing.T) { + executor := &emptyCompletionTestExecutor{ + streamPayloads: map[string][][]byte{}, + streamCalls: map[string]int{}, + emptyStreamPayload: [][]byte{ + // Executors emit raw JSON payloads without SSE framing. + []byte("{\"choices\":[{\"delta\":{\"role\":\"assistant\",\"content\":\"\"},\"finish_reason\":\"stop\"}],\"usage\":{\"completion_tokens\":0}}"), + }, + contentStreamPayload: [][]byte{ + []byte("{\"choices\":[{\"delta\":{\"content\":\"hello\"},\"finish_reason\":null}]}"), + []byte("{\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":{\"completion_tokens\":3}}"), + }, + } + manager, ids, model, capture := newEmptyCompletionTestManager(t, executor) + + stream, err := manager.ExecuteStream(context.Background(), []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) + } + if !strings.Contains(got.String(), "hello") { + t.Fatalf("stream payload = %q, want content from the non-empty auth", got.String()) + } + emptyFirst := executor.firstStream + if emptyFirst == "" { + t.Fatal("executor never streamed any auth") + } + other := ids[0] + if emptyFirst == ids[0] { + other = ids[1] + } + var emptyRecorded bool + var otherSucceeded bool + for _, r := range capture.Results() { + if r.AuthID == emptyFirst && !r.Success { + emptyRecorded = true + } + if r.AuthID == other && r.Success { + otherSucceeded = true + } + } + if !emptyRecorded { + t.Fatalf("expected failure recorded for empty auth %s, results: %+v", emptyFirst, capture.Results()) + } + if !otherSucceeded { + t.Fatalf("expected success recorded for non-empty auth %s, results: %+v", other, capture.Results()) + } +} + +func TestExecuteStreamThinkingThenContentNotRotated(t *testing.T) { + executor := &emptyCompletionTestExecutor{ + streamPayloads: map[string][][]byte{}, + streamCalls: map[string]int{}, + } + manager, ids, model, _ := newEmptyCompletionTestManager(t, executor) + + // Positive control: a thinking-first-then-content stream must NOT be + // treated as empty, so it should not rotate to the second auth. + content := [][]byte{ + []byte("data: {\"choices\":[{\"delta\":{\"reasoning_content\":\"thinking\"},\"finish_reason\":null}]}\n\n"), + []byte("data: {\"choices\":[{\"delta\":{\"content\":\"answer\"},\"finish_reason\":\"stop\"}]}\n\n"), + []byte("data: [DONE]\n\n"), + } + executor.streamPayloads[ids[0]] = content + executor.streamPayloads[ids[1]] = content + + var streamCalls int + stream, err := manager.ExecuteStream(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{Stream: true}) + if err != nil { + t.Fatalf("ExecuteStream() error = %v", err) + } + var got strings.Builder + for chunk := range stream.Chunks { + got.Write(chunk.Payload) + } + _ = streamCalls + if !strings.Contains(got.String(), "answer") { + t.Fatalf("stream payload = %q, want thinking-then-content stream to pass through", got.String()) + } + // The first auth must NOT have been cooled (it produced a real completion). + if auth, ok := manager.GetByID(ids[0]); ok && auth != nil { + if auth.Unavailable || !auth.NextRetryAfter.IsZero() { + t.Fatalf("auth %q was cooled despite producing a real completion", ids[0]) + } + } +} + +func TestExecuteStreamInStreamGemini429ErrorRotatesAuth(t *testing.T) { + executor := &emptyCompletionTestExecutor{ + streamPayloads: map[string][][]byte{}, + streamCalls: map[string]int{}, + emptyStreamPayload: [][]byte{ + []byte("data: {\"error\":{\"code\":429,\"message\":\"Resource exhausted\",\"status\":\"RESOURCE_EXHAUSTED\"}}\n\n"), + }, + contentStreamPayload: [][]byte{ + []byte("data: {\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"text\":\"gemini response\"}]},\"finishReason\":\"STOP\"}],\"usageMetadata\":{\"candidatesTokenCount\":5}}\n\n"), + }, + } + manager, ids, model, capture := newEmptyCompletionTestManager(t, executor) + + stream, err := manager.ExecuteStream(context.Background(), []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(), "gemini response", capture) + + if auth, ok := manager.GetByID(executor.firstStream); ok && auth != nil { + if !auth.Unavailable && auth.NextRetryAfter.IsZero() && !auth.Quota.Exceeded { + t.Fatalf("auth %q was not marked unavailable or quota exceeded after in-stream 429 error", executor.firstStream) + } + } +} + +func TestExecuteStreamInStreamClaudeOverloadedErrorRotatesAuth(t *testing.T) { + executor := &emptyCompletionTestExecutor{ + streamPayloads: map[string][][]byte{}, + streamCalls: map[string]int{}, + emptyStreamPayload: [][]byte{ + []byte("event: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\",\"message\":\"Overloaded\"}}\n\n"), + }, + contentStreamPayload: [][]byte{ + []byte("data: {\"choices\":[{\"delta\":{\"content\":\"claude response\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n"), + }, + } + manager, ids, model, capture := newEmptyCompletionTestManager(t, executor) + + stream, err := manager.ExecuteStream(context.Background(), []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(), "claude response", capture) +} + +func TestExecuteStreamInStream400InvalidRequestNotRotated(t *testing.T) { + executor := &emptyCompletionTestExecutor{ + streamPayloads: map[string][][]byte{}, + streamCalls: map[string]int{}, + emptyStreamPayload: [][]byte{ + []byte("data: {\"error\":{\"code\":400,\"message\":\"Invalid request prompt\",\"type\":\"invalid_request_error\"}}\n\n"), + }, + contentStreamPayload: [][]byte{ + []byte("data: {\"choices\":[{\"delta\":{\"content\":\"should not reach\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n"), + }, + } + manager, ids, model, capture := newEmptyCompletionTestManager(t, executor) + + _, err := manager.ExecuteStream(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{Stream: true}) + if err == nil { + t.Fatal("ExecuteStream() want error for 400 invalid request, got nil") + } + + other := ids[0] + if executor.firstStream == ids[0] { + other = ids[1] + } + if executor.streamCalls[other] > 0 { + t.Fatalf("second auth %q was called (%d times), want 0 calls (400 must not rotate)", other, executor.streamCalls[other]) + } + _ = capture +} + +func TestExecuteStreamInStreamUnknownJSONForwardedNotRotated(t *testing.T) { + executor := &emptyCompletionTestExecutor{ + streamPayloads: map[string][][]byte{}, + streamCalls: map[string]int{}, + emptyStreamPayload: [][]byte{ + []byte("data: {\"custom_future_protocol_field\":\"forward_me\"}\n\n"), + []byte("data: [DONE]\n\n"), + }, + contentStreamPayload: [][]byte{ + []byte("data: {\"choices\":[{\"delta\":{\"content\":\"should not reach\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n"), + }, + } + manager, ids, model, _ := newEmptyCompletionTestManager(t, executor) + + stream, err := manager.ExecuteStream(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{Stream: true}) + if err != nil { + t.Fatalf("ExecuteStream() error = %v", err) + } + var got strings.Builder + for chunk := range stream.Chunks { + got.Write(chunk.Payload) + } + if !strings.Contains(got.String(), "custom_future_protocol_field") { + t.Fatalf("payload = %q, want unknown JSON forwarded directly", got.String()) + } + other := ids[0] + if executor.firstStream == ids[0] { + other = ids[1] + } + if executor.streamCalls[other] > 0 { + t.Fatalf("second auth %q was called (%d times), want 0 calls for unknown valid JSON", other, executor.streamCalls[other]) + } +} + +func TestExecuteStreamInStreamResponseFailedNestedErrorRotatesAuth(t *testing.T) { + executor := &emptyCompletionTestExecutor{ + streamPayloads: map[string][][]byte{}, + streamCalls: map[string]int{}, + emptyStreamPayload: [][]byte{ + []byte("data: {\"type\":\"response.failed\",\"response\":{\"id\":\"resp_123\",\"status\":\"failed\",\"error\":{\"type\":\"server_error\",\"code\":\"rate_limit_exceeded\",\"message\":\"Rate limit reached\"}}}\n\n"), + }, + contentStreamPayload: [][]byte{ + []byte("data: {\"choices\":[{\"delta\":{\"content\":\"rotated response\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n"), + }, + } + manager, ids, model, capture := newEmptyCompletionTestManager(t, executor) + + stream, err := manager.ExecuteStream(context.Background(), []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(), "rotated response", capture) + + if auth, ok := manager.GetByID(executor.firstStream); ok && auth != nil { + if !auth.Unavailable && auth.NextRetryAfter.IsZero() && !auth.Quota.Exceeded { + t.Fatalf("auth %q was not marked unavailable or quota exceeded after response.failed error", executor.firstStream) + } + } +} + +func TestStreamBootstrapDetectorWebSearchProgressEvents(t *testing.T) { + progressInProgress := []byte("data: {\"type\":\"response.web_search_call.in_progress\",\"item_id\":\"ws_123\"}\n\n") + progressSearching := []byte("data: {\"type\":\"response.web_search_call.searching\",\"item_id\":\"ws_123\"}\n\n") + progressCompleted := []byte("data: {\"type\":\"response.web_search_call.completed\",\"item_id\":\"ws_123\"}\n\n") + + t.Run("web_search_call progress events do not flush bootstrap and allow failover on subsequent error", func(t *testing.T) { + var detector StreamBootstrapDetector + if detector.Observe(progressInProgress) { + t.Fatal("Observe(in_progress) = true, want false (scaffold only)") + } + if detector.Observe(progressSearching) { + t.Fatal("Observe(searching) = true, want false (scaffold only)") + } + if detector.Observe(progressCompleted) { + t.Fatal("Observe(completed) = true, want false (scaffold only)") + } + if detector.HasMeaningfulOutput() { + t.Fatal("HasMeaningfulOutput() = true for progress frames, want false") + } + + errUpstream := errors.New("upstream failed after web_search_call progress") + ch := make(chan cliproxyexecutor.StreamChunk, 4) + ch <- cliproxyexecutor.StreamChunk{Payload: progressInProgress} + ch <- cliproxyexecutor.StreamChunk{Payload: progressSearching} + ch <- cliproxyexecutor.StreamChunk{Payload: progressCompleted} + 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", err, errUpstream) + } + if len(buffered) != 0 { + t.Fatalf("readStreamBootstrap buffered = %d, want 0 on failover", len(buffered)) + } + if closed { + t.Fatal("readStreamBootstrap closed = true, want false") + } + }) +} + +func TestExecuteStreamMidStreamInStreamErrorMarksAuthFailed(t *testing.T) { + midStreamPayload := [][]byte{ + []byte("data: {\"choices\":[{\"delta\":{\"content\":\"valid prefix content\"}}]}\n\n"), + []byte("data: {\"error\":{\"code\":429,\"message\":\"Resource exhausted mid-stream\",\"status\":\"RESOURCE_EXHAUSTED\"}}\n\n"), + } + executor := &emptyCompletionTestExecutor{ + streamPayloads: map[string][][]byte{}, + streamCalls: map[string]int{}, + contentStreamPayload: midStreamPayload, + } + manager, ids, model, capture := newEmptyCompletionTestManager(t, executor) + + stream, err := manager.ExecuteStream(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{Stream: true}) + if err != nil { + t.Fatalf("ExecuteStream() unexpected 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) + } + + payloadStr := got.String() + if !strings.Contains(payloadStr, "valid prefix content") { + t.Fatalf("stream payload missing prefix content, got: %q", payloadStr) + } + if !strings.Contains(payloadStr, "Resource exhausted mid-stream") { + t.Fatalf("stream payload missing mid-stream error, got: %q", payloadStr) + } + + results := capture.Results() + if len(results) == 0 { + t.Fatal("expected at least 1 execution result recorded, got 0") + } + + hasFailed := false + for _, res := range results { + if res.Success { + t.Fatalf("recorded execution result with Success = true for mid-stream error: %+v", res) + } + if !res.Success { + hasFailed = true + } + } + if !hasFailed { + t.Fatal("expected execution result with Success = false, none found") + } + + _ = ids +} + +func TestParseStreamErrorGRPCStatusCodeFallback(t *testing.T) { + payload := []byte(`{"error":{"code":8,"status":"RESOURCE_EXHAUSTED"}}`) + err := evalProviderError(payload, "") + if err == nil { + t.Fatal("evalProviderError() returned nil, want error") + } + if err.HTTPStatus != http.StatusTooManyRequests { + t.Fatalf("err.HTTPStatus = %d, want %d (RESOURCE_EXHAUSTED fallback)", err.HTTPStatus, http.StatusTooManyRequests) + } +} + +func TestStreamSplitSSEEventAndDataAcrossChunksDetectsError(t *testing.T) { + midStreamPayload := [][]byte{ + []byte("data: {\"choices\":[{\"delta\":{\"content\":\"valid prefix content\"}}]}\n\n"), + []byte("event: error\n\n"), + []byte("data: {\"message\":\"overloaded\"}\n\n"), + } + executor := &emptyCompletionTestExecutor{ + streamPayloads: map[string][][]byte{}, + streamCalls: map[string]int{}, + contentStreamPayload: midStreamPayload, + } + manager, ids, model, capture := newEmptyCompletionTestManager(t, executor) + + stream, err := manager.ExecuteStream(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{Stream: true}) + if err != nil { + t.Fatalf("ExecuteStream() unexpected error = %v", err) + } + if stream == nil { + t.Fatal("ExecuteStream() returned nil stream") + } + + for range stream.Chunks { + } + + results := capture.Results() + if len(results) == 0 { + t.Fatal("expected execution results, got 0") + } + for _, res := range results { + if res.Success { + t.Fatalf("recorded execution result with Success = true for split SSE error: %+v", res) + } + } + + _ = ids +} + +func TestStreamRawJSONWithoutNewlinesThenErrorDetectsFailure(t *testing.T) { + midStreamPayload := [][]byte{ + []byte(`{"choices":[{"delta":{"content":"valid prefix content"}}]}`), + []byte(`{"error":{"code":"rate_limit_exceeded","message":"Rate limit reached"}}`), + } + executor := &emptyCompletionTestExecutor{ + streamPayloads: map[string][][]byte{}, + streamCalls: map[string]int{}, + contentStreamPayload: midStreamPayload, + } + manager, ids, model, capture := newEmptyCompletionTestManager(t, executor) + + stream, err := manager.ExecuteStream(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{Stream: true}) + if err != nil { + t.Fatalf("ExecuteStream() unexpected error = %v", err) + } + if stream == nil { + t.Fatal("ExecuteStream() returned nil stream") + } + + for range stream.Chunks { + } + + results := capture.Results() + if len(results) == 0 { + t.Fatal("expected execution results, got 0") + } + for _, res := range results { + if res.Success { + t.Fatalf("recorded execution result with Success = true for raw JSON error: %+v", res) + } + } + + auth, ok := manager.GetByID(executor.firstStream) + if !ok || auth == nil { + t.Fatalf("auth %q not found", executor.firstStream) + } + if !auth.Unavailable && auth.NextRetryAfter.IsZero() && !auth.Quota.Exceeded { + t.Fatalf("auth %q was not marked unavailable or in cooldown after raw JSON error", executor.firstStream) + } + + _ = ids +} + +func TestStreamBootstrapDetectorCodexHandshakeMetadataThenError(t *testing.T) { + rateLimitsFrame := []byte("data: {\"type\":\"codex.rate_limits\",\"rate_limits\":{\"requests_remaining\":100}}\n\n") + metadataFrame := []byte("data: {\"type\":\"codex.response.metadata\",\"session_id\":\"sess_123\"}\n\n") + + t.Run("handshake metadata frames do not commit stream and allow failover on subsequent error", func(t *testing.T) { + var detector StreamBootstrapDetector + if detector.Observe(rateLimitsFrame) { + t.Fatal("Observe(codex.rate_limits) = true, want false (handshake only)") + } + if detector.Observe(metadataFrame) { + t.Fatal("Observe(codex.response.metadata) = true, want false (handshake only)") + } + if detector.HasMeaningfulOutput() { + t.Fatal("HasMeaningfulOutput() = true for handshake frames, want false") + } + + executor := &emptyCompletionTestExecutor{ + streamPayloads: map[string][][]byte{}, + streamCalls: map[string]int{}, + emptyStreamPayload: [][]byte{ + rateLimitsFrame, + metadataFrame, + []byte("data: {\"type\":\"error\",\"error\":{\"code\":\"overloaded_error\",\"message\":\"Server overloaded\"}}\n\n"), + }, + contentStreamPayload: [][]byte{ + []byte("data: {\"choices\":[{\"delta\":{\"content\":\"rotated response\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n"), + }, + } + manager, ids, model, capture := newEmptyCompletionTestManager(t, executor) + + stream, err := manager.ExecuteStream(context.Background(), []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(), "rotated response", capture) + }) +} + +// assertRotatesToContent verifies that an empty-completion from the first-picked +// auth rotates to the other auth, which then succeeds with the given content. +func assertRotatesToContent(t *testing.T, ids []string, emptyFirst, gotPayload, wantSubstr string, capture *resultCaptureHook) { + t.Helper() + if emptyFirst == "" { + t.Fatal("executor never executed/streamed any auth") + } + if !strings.Contains(gotPayload, wantSubstr) { + t.Fatalf("payload = %q, want %q from the non-empty auth", gotPayload, wantSubstr) + } + other := ids[0] + if emptyFirst == ids[0] { + other = ids[1] + } + var emptyRecorded bool + var otherSucceeded bool + for _, r := range capture.Results() { + if r.AuthID == emptyFirst && !r.Success { + emptyRecorded = true + } + if r.AuthID == other && r.Success { + otherSucceeded = true + } + } + if !emptyRecorded { + t.Fatalf("empty auth %q was not recorded as a failure result; results=%v", emptyFirst, capture.Results()) + } + if !otherSucceeded { + 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 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 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 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 positive 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":1}}`) + if IsEmptyCompletionPayload(payload) { + t.Fatal("IsEmptyCompletionPayload() = true for non-empty thoughtSignature with positive token count, want false") + } + }) + + t.Run("gemini STOP with thought_signature and positive 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":1}}`) + if IsEmptyCompletionPayload(payload) { + t.Fatal("IsEmptyCompletionPayload() = true for non-empty thought_signature with positive 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 does not mark meaningful and allows bootstrap error failover", func(t *testing.T) { + var detector StreamBootstrapDetector + if detector.Observe(funcWithID) { + t.Fatal("detector.Observe() = true for function_call with id, want false") + } + if detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = true for function_call with id, want false") + } + + errUpstream := errors.New("upstream failed immediately after function_call id scaffold") + ch := make(chan cliproxyexecutor.StreamChunk, 2) + ch <- cliproxyexecutor.StreamChunk{Payload: funcWithID} + 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("scaffold with non-empty call_id does not mark meaningful and allows bootstrap error failover", func(t *testing.T) { + var detector StreamBootstrapDetector + if detector.Observe(funcWithCallID) { + t.Fatal("detector.Observe() = true for function_call with call_id, want false") + } + if detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = true for function_call with call_id, want false") + } + + errUpstream := errors.New("upstream failed immediately after function_call call_id scaffold") + ch := make(chan cliproxyexecutor.StreamChunk, 2) + ch <- cliproxyexecutor.StreamChunk{Payload: funcWithCallID} + 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("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 TestOpenAIToolCallBootstrapLateNameAndEmptyArgs(t *testing.T) { + idOnlyDelta := []byte("data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"id\":\"call_123\",\"index\":0,\"type\":\"function\"}]}}]}\n\n") + nameDelta := []byte("data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"name\":\"get_weather\"}}]}}]}\n\n") + nameEmptyArgsDelta := []byte("data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"id\":\"call_456\",\"index\":0,\"type\":\"function\",\"function\":{\"name\":\"get_time\",\"arguments\":\"{}\"}}]}}]}\n\n") + + t.Run("id-only delta does not forward and allows failover", func(t *testing.T) { + var detector StreamBootstrapDetector + if detector.Observe(idOnlyDelta) { + t.Fatal("detector.Observe() = true for id-only tool_calls delta, want false") + } + if detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = true for id-only tool_calls delta, want false") + } + + errUpstream := errors.New("upstream failed after id-only delta") + ch := make(chan cliproxyexecutor.StreamChunk, 2) + ch <- cliproxyexecutor.StreamChunk{Payload: idOnlyDelta} + 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", 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") + } + }) + + t.Run("delta that later supplies the name marks meaningful and forwards", func(t *testing.T) { + var detector StreamBootstrapDetector + if detector.Observe(idOnlyDelta) { + t.Fatal("detector.Observe(idOnlyDelta) = true, want false") + } + if detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = true before name, want false") + } + if !detector.Observe(nameDelta) { + t.Fatal("detector.Observe(nameDelta) = false, want true") + } + if !detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = false after name arrives, want true") + } + }) + + t.Run("name with empty-object arguments is meaningful and forwards", func(t *testing.T) { + var detector StreamBootstrapDetector + if !detector.Observe(nameEmptyArgsDelta) { + t.Fatal("detector.Observe(nameEmptyArgsDelta) = false, want true") + } + if !detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = false, want true") + } + }) +} + +func TestResponsesCallItemLateNameAndEmptyArgs(t *testing.T) { + idOnlyItem := []byte("event: response.output_item.added\ndata: {\"type\":\"response.output_item.added\",\"output_index\":0,\"item\":{\"id\":\"fc_123\",\"call_id\":\"call_123\",\"type\":\"function_call\",\"status\":\"in_progress\",\"arguments\":\"\",\"name\":\"\"}}\n\n") + nameLaterItem := []byte("event: response.output_item.added\ndata: {\"type\":\"response.output_item.added\",\"output_index\":0,\"item\":{\"id\":\"fc_123\",\"call_id\":\"call_123\",\"type\":\"function_call\",\"status\":\"in_progress\",\"name\":\"get_weather\"}}\n\n") + nameEmptyArgsItem := []byte("event: response.output_item.added\ndata: {\"type\":\"response.output_item.added\",\"output_index\":0,\"item\":{\"id\":\"fc_456\",\"call_id\":\"call_456\",\"type\":\"function_call\",\"status\":\"in_progress\",\"name\":\"get_time\",\"arguments\":\"{}\"}}\n\n") + + t.Run("responses id and call_id only item does not forward and allows failover", func(t *testing.T) { + var detector StreamBootstrapDetector + if detector.Observe(idOnlyItem) { + t.Fatal("detector.Observe() = true for id-only responses item, want false") + } + if detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = true for id-only responses item, want false") + } + + errUpstream := errors.New("upstream failed after id-only responses item") + ch := make(chan cliproxyexecutor.StreamChunk, 2) + ch <- cliproxyexecutor.StreamChunk{Payload: idOnlyItem} + 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", 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") + } + }) + + t.Run("responses item that later supplies name marks meaningful and forwards", func(t *testing.T) { + var detector StreamBootstrapDetector + if detector.Observe(idOnlyItem) { + t.Fatal("detector.Observe(idOnlyItem) = true, want false") + } + if detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = true before name, want false") + } + if !detector.Observe(nameLaterItem) { + t.Fatal("detector.Observe(nameLaterItem) = false, want true") + } + if !detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = false after name arrives, want true") + } + }) + + t.Run("responses item with name and empty-object arguments is meaningful and forwards", func(t *testing.T) { + var detector StreamBootstrapDetector + if !detector.Observe(nameEmptyArgsItem) { + t.Fatal("detector.Observe(nameEmptyArgsItem) = false, want true") + } + if !detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = false, want true") + } + }) +} + +func TestResponsesCallItemActionPayload(t *testing.T) { + webSearchPopulatedAction := []byte("event: response.output_item.done\ndata: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"id\":\"ws_123\",\"call_id\":\"call_123\",\"type\":\"web_search_call\",\"status\":\"completed\",\"action\":{\"type\":\"search\",\"query\":\"weather\"}}}\n\n") + computerCallPopulatedAction := []byte("event: response.output_item.done\ndata: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"id\":\"cc_123\",\"call_id\":\"call_123\",\"type\":\"computer_call\",\"status\":\"completed\",\"action\":{\"type\":\"click\",\"x\":100,\"y\":200}}}\n\n") + emptyActionItem := []byte("event: response.output_item.done\ndata: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"id\":\"ws_123\",\"call_id\":\"call_123\",\"type\":\"web_search_call\",\"status\":\"in_progress\",\"action\":{}}}\n\n") + noActionItem := []byte("event: response.output_item.done\ndata: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"id\":\"ws_123\",\"call_id\":\"call_123\",\"type\":\"web_search_call\",\"status\":\"in_progress\"}}\n\n") + + t.Run("web_search_call item with populated action marks meaningful and forwards", func(t *testing.T) { + var detector StreamBootstrapDetector + if !detector.Observe(webSearchPopulatedAction) { + t.Fatal("detector.Observe() = false for web_search_call with populated action, want true") + } + if !detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = false for web_search_call with populated action, want true") + } + }) + + t.Run("computer_call item with populated action marks meaningful and forwards", func(t *testing.T) { + var detector StreamBootstrapDetector + if !detector.Observe(computerCallPopulatedAction) { + t.Fatal("detector.Observe() = false for computer_call with populated action, want true") + } + if !detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = false for computer_call with populated action, want true") + } + }) + + t.Run("responses call item with id and empty action object does not mark meaningful and allows failover", func(t *testing.T) { + var detector StreamBootstrapDetector + if detector.Observe(emptyActionItem) { + t.Fatal("detector.Observe() = true for item with empty action object, want false") + } + if detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = true for item with empty action object, want false") + } + + errUpstream := errors.New("upstream failed after empty action item") + ch := make(chan cliproxyexecutor.StreamChunk, 2) + ch <- cliproxyexecutor.StreamChunk{Payload: emptyActionItem} + 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", 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") + } + }) + + t.Run("responses call item with id and no action does not mark meaningful and allows failover", func(t *testing.T) { + var detector StreamBootstrapDetector + if detector.Observe(noActionItem) { + t.Fatal("detector.Observe() = true for item with no action, want false") + } + if detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = true for item with no action, want false") + } + + errUpstream := errors.New("upstream failed after no action item") + ch := make(chan cliproxyexecutor.StreamChunk, 2) + ch <- cliproxyexecutor.StreamChunk{Payload: noActionItem} + 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", 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") + } + }) +} + +func TestResponsesCallItemResultsPayload(t *testing.T) { + webSearchPopulatedResults := []byte("event: response.output_item.done\ndata: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"id\":\"ws_123\",\"call_id\":\"call_123\",\"type\":\"web_search_call\",\"status\":\"completed\",\"results\":[{\"url\":\"https://example.com/search\"}]}}\n\n") + emptyResultsItem := []byte("event: response.output_item.done\ndata: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"id\":\"ws_123\",\"call_id\":\"call_123\",\"type\":\"web_search_call\",\"status\":\"in_progress\",\"results\":[]}}\n\n") + emptyResultsObjectItem := []byte("event: response.output_item.done\ndata: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"id\":\"ws_123\",\"call_id\":\"call_123\",\"type\":\"web_search_call\",\"status\":\"in_progress\",\"results\":{}}}\n\n") + nullResultsItem := []byte("event: response.output_item.done\ndata: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"id\":\"ws_123\",\"call_id\":\"call_123\",\"type\":\"web_search_call\",\"status\":\"in_progress\",\"results\":null}}\n\n") + + t.Run("web_search_call item with populated results marks meaningful and forwards", func(t *testing.T) { + var detector StreamBootstrapDetector + if !detector.Observe(webSearchPopulatedResults) { + t.Fatal("detector.Observe() = false for web_search_call with populated results, want true") + } + if !detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = false for web_search_call with populated results, want true") + } + }) + + t.Run("responses call item with id and empty results array does not mark meaningful and allows failover", func(t *testing.T) { + var detector StreamBootstrapDetector + if detector.Observe(emptyResultsItem) { + t.Fatal("detector.Observe() = true for item with empty results array, want false") + } + if detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = true for item with empty results array, want false") + } + + errUpstream := errors.New("upstream failed after empty results item") + ch := make(chan cliproxyexecutor.StreamChunk, 2) + ch <- cliproxyexecutor.StreamChunk{Payload: emptyResultsItem} + 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", 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") + } + }) + + t.Run("responses call item with id and empty results object does not mark meaningful and allows failover", func(t *testing.T) { + var detector StreamBootstrapDetector + if detector.Observe(emptyResultsObjectItem) { + t.Fatal("detector.Observe() = true for item with empty results object, want false") + } + if detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = true for item with empty results object, want false") + } + }) + + t.Run("responses call item with id and null results does not mark meaningful and allows failover", func(t *testing.T) { + var detector StreamBootstrapDetector + if detector.Observe(nullResultsItem) { + t.Fatal("detector.Observe() = true for item with null results, want false") + } + if detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = true for item with null results, want false") + } + }) +} + +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 TestEmptyCompletion_OpenAIImageGenerationResult(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) + } +} + +func TestEmptyCompletion_OpenAIFinishReasonStopWithoutDoneIsTerminalEmpty(t *testing.T) { + // Case 1: Single choice finish_reason="stop" without [DONE] is terminal empty. + var detector StreamBootstrapDetector + stopChunk := []byte("data: {\"id\":\"1\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n") + got := detector.Observe(stopChunk) + if got { + t.Fatalf("Observe(stopChunk) = %v, want false", got) + } + if !detector.IsTerminalEmpty() { + t.Fatal("IsTerminalEmpty() = false, want true for OpenAI finish_reason:stop without [DONE]") + } + + // Case 2: Multi-choice with partial finish_reason (choice 0 "stop", choice 1 nil) is NOT terminal yet. + var detectorPartial StreamBootstrapDetector + partialChunk := []byte("data: {\"id\":\"1\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"},{\"index\":1,\"delta\":{}}]}\n\n") + gotPartial := detectorPartial.Observe(partialChunk) + if gotPartial { + t.Fatalf("Observe(partialChunk) = %v, want false", gotPartial) + } + if detectorPartial.IsTerminalEmpty() { + t.Fatal("IsTerminalEmpty() = true, want false when not all choices have finish_reason") + } + + // Case 3: Non-stop reason (content_filter) is NOT empty completion and forwards. + var detectorFilter StreamBootstrapDetector + filterChunk := []byte("data: {\"id\":\"1\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"content_filter\"}]}\n\n") + gotFilter := detectorFilter.Observe(filterChunk) + if !gotFilter { + t.Fatalf("Observe(filterChunk) = %v, want true for blocked reason", gotFilter) + } + if detectorFilter.IsTerminalEmpty() { + t.Fatal("IsTerminalEmpty() = true, want false for content_filter") + } + + // Case 4: Multi-choice with separate frames where frame1 has choice 0 finished empty + // and frame2 has choice 1 with content. Frame 1 must NOT trigger IsTerminalEmpty(). + var detectorMulti StreamBootstrapDetector + frame0 := []byte("data: {\"id\":\"1\",\"choices\":[{\"index\":0,\"delta\":{}},{\"index\":1,\"delta\":{}}]}\n\n") + if detectorMulti.Observe(frame0) { + t.Fatal("Observe(frame0) = true, want false for empty metadata") + } + frame1 := []byte("data: {\"id\":\"1\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n") + if detectorMulti.Observe(frame1) { + t.Fatal("Observe(frame1) = true, want false for choice 0 finish") + } + if detectorMulti.IsTerminalEmpty() { + t.Fatal("IsTerminalEmpty() = true on frame1, want false when choice 1 has not finished yet") + } + frame2 := []byte("data: {\"id\":\"1\",\"choices\":[{\"index\":1,\"delta\":{\"content\":\"hello\"}}]}\n\n") + if !detectorMulti.Observe(frame2) { + t.Fatal("Observe(frame2) = false, want true for choice 1 content") + } + if detectorMulti.IsTerminalEmpty() { + t.Fatal("IsTerminalEmpty() = true after content received, want false") + } + + // Case 5 (Round 40): Multi-choice (n=2) stream where the first received frame + // is choice 0 finish:"stop" empty and choice 1 has not appeared yet. + // Must NOT trigger IsTerminalEmpty() early. + var detectorMultiEarly StreamBootstrapDetector + detectorMultiEarly.SetExpectedChoices(2) + frameChoice0Finish := []byte("data: {\"id\":\"1\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n") + if detectorMultiEarly.Observe(frameChoice0Finish) { + t.Fatal("Observe(frameChoice0Finish) = true, want false") + } + if detectorMultiEarly.IsTerminalEmpty() { + t.Fatal("IsTerminalEmpty() = true on first frame when expected n=2, want false until all n choices finish") + } + + // Case 6 (Round 40): Multi-choice (n=2) stream where all n choices finish empty. + frameChoice1Finish := []byte("data: {\"id\":\"1\",\"choices\":[{\"index\":1,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n") + if detectorMultiEarly.Observe(frameChoice1Finish) { + t.Fatal("Observe(frameChoice1Finish) = true, want false") + } + if !detectorMultiEarly.IsTerminalEmpty() { + t.Fatal("IsTerminalEmpty() = false when all n=2 choices finished empty, want true") + } +} + +func TestExtractExpectedChoices(t *testing.T) { + cases := []struct { + name string + payload string + want int + }{ + {name: "nil payload", payload: "", want: 1}, + {name: "empty json", payload: `{}`, want: 1}, + {name: "explicit n=1", payload: `{"n":1}`, want: 1}, + {name: "explicit n=2", payload: `{"model":"gpt-4o","n":2}`, want: 2}, + {name: "explicit n=4", payload: `{"n":4,"prompt":"hello"}`, want: 4}, + {name: "nested request.n=3", payload: `{"request":{"n":3}}`, want: 3}, + {name: "invalid n=0", payload: `{"n":0}`, want: 1}, + {name: "negative n=-1", payload: `{"n":-1}`, want: 1}, + {name: "invalid json", payload: `not json`, want: 1}, + {name: "gemini generationConfig.candidateCount=2", payload: `{"generationConfig":{"candidateCount":2}}`, want: 2}, + {name: "gemini generationConfig.candidate_count=3", payload: `{"generationConfig":{"candidate_count":3}}`, want: 3}, + {name: "gemini generation_config.candidateCount=4", payload: `{"generation_config":{"candidateCount":4}}`, want: 4}, + {name: "gemini nested request.generationConfig.candidateCount=2", payload: `{"request":{"generationConfig":{"candidateCount":2}}}`, want: 2}, + {name: "gemini top-level candidateCount=3", payload: `{"candidateCount":3}`, want: 3}, + {name: "gemini top-level candidate_count=5", payload: `{"candidate_count":5}`, want: 5}, + {name: "both n=2 and candidateCount=3", payload: `{"n":2,"generationConfig":{"candidateCount":3}}`, want: 3}, + {name: "gemini invalid candidateCount=0", payload: `{"generationConfig":{"candidateCount":0}}`, want: 1}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got := ExtractExpectedChoices([]byte(tc.payload)) + if got != tc.want { + t.Fatalf("ExtractExpectedChoices(%q) = %d, want %d", tc.payload, got, tc.want) + } + }) + } +} + +func TestGeminiStreamBootstrapMultiCandidate(t *testing.T) { + t.Run("candidateCount=2 candidate0 STOP empty candidate1 content does not terminate at frame1", func(t *testing.T) { + var detector StreamBootstrapDetector + detector.SetRequestPayload([]byte(`{"generationConfig":{"candidateCount":2}}`)) + + frameCand0 := []byte("data: {\"candidates\":[{\"index\":0,\"content\":{\"parts\":[]},\"finishReason\":\"STOP\"}]}\n\n") + frameCand1 := []byte("data: {\"candidates\":[{\"index\":1,\"content\":{\"parts\":[{\"text\":\"hello from candidate 1\"}]}}]}\n\n") + + if detector.Observe(frameCand0) { + t.Fatal("Observe(frameCand0) = true, want false") + } + if detector.IsTerminalEmpty() { + t.Fatal("IsTerminalEmpty() = true on candidate 0 STOP when expected candidateCount=2, want false until all candidates finish") + } + + if !detector.Observe(frameCand1) { + t.Fatal("Observe(frameCand1 with content) = false, want true") + } + if detector.IsTerminalEmpty() { + t.Fatal("IsTerminalEmpty() = true after candidate 1 emitted content, want false") + } + }) + + t.Run("candidateCount=2 both candidates finish empty is terminal empty", func(t *testing.T) { + var detector StreamBootstrapDetector + detector.SetRequestPayload([]byte(`{"generationConfig":{"candidateCount":2}}`)) + + frameCand0 := []byte("data: {\"candidates\":[{\"index\":0,\"content\":{\"parts\":[]},\"finishReason\":\"STOP\"}]}\n\n") + frameCand1 := []byte("data: {\"candidates\":[{\"index\":1,\"content\":{\"parts\":[]},\"finishReason\":\"STOP\"}]}\n\n") + + if detector.Observe(frameCand0) { + t.Fatal("Observe(frameCand0) = true, want false") + } + if detector.IsTerminalEmpty() { + t.Fatal("IsTerminalEmpty() = true on candidate 0 STOP, want false") + } + + if detector.Observe(frameCand1) { + t.Fatal("Observe(frameCand1) = true, want false") + } + if !detector.IsTerminalEmpty() { + t.Fatal("IsTerminalEmpty() = false when both candidates finished empty, want true") + } + }) + + t.Run("default candidateCount single candidate STOP empty is terminal empty", func(t *testing.T) { + var detector StreamBootstrapDetector + // Default candidateCount is 1 + frameCand0 := []byte("data: {\"candidates\":[{\"finishReason\":\"STOP\"}]}\n\n") + + if detector.Observe(frameCand0) { + t.Fatal("Observe(frameCand0) = true, want false") + } + if !detector.IsTerminalEmpty() { + t.Fatal("IsTerminalEmpty() = false on default single candidate STOP, want true") + } + }) + + t.Run("candidateCount=2 single candidate arrives and stream closes reaches terminal empty on EOF", func(t *testing.T) { + var detector StreamBootstrapDetector + detector.SetRequestPayload([]byte(`{"generationConfig":{"candidateCount":2}}`)) + + frameCand0 := []byte("data: {\"candidates\":[{\"index\":0,\"content\":{\"parts\":[]},\"finishReason\":\"STOP\"}]}\n\n") + + if detector.Observe(frameCand0) { + t.Fatal("Observe(frameCand0) = true, want false") + } + if detector.IsTerminalEmpty() { + t.Fatal("IsTerminalEmpty() = true mid-stream when candidateCount=2 but only candidate 0 finished, want false") + } + + // Channel closes / EOF reached without candidate 1 ever arriving + if !detector.Finish() { + t.Fatal("Finish() at EOF = false for stream with only empty candidate 0, want true") + } + }) +} + +func TestResponsesReasoningSummaryPartScaffoldDoesNotPrematurelyForward(t *testing.T) { + t.Run("empty reasoning_summary_part.added does not forward and preserves buffer for upstream error", func(t *testing.T) { + var detector StreamBootstrapDetector + scaffoldEvent := []byte("data: {\"type\":\"response.reasoning_summary_part.added\",\"sequence_number\":1,\"item_id\":\"rs_1\",\"output_index\":0,\"summary_index\":0,\"part\":{\"type\":\"summary_text\",\"text\":\"\"}}\n\n") + + if detector.Observe(scaffoldEvent) { + t.Fatal("Observe(empty reasoning_summary_part.added) = true, want false (must not prematurely forward)") + } + if detector.HasMeaningfulOutput() { + t.Fatal("HasMeaningfulOutput() = true on empty scaffold, want false") + } + if detector.IsTerminalEmpty() { + t.Fatal("IsTerminalEmpty() = true on scaffold, want false") + } + + // Conductor stream test: when upstream error arrives after empty scaffold, + // readStreamBootstrap surfaces the upstream error instead of committing success. + ch := make(chan cliproxyexecutor.StreamChunk, 2) + ch <- cliproxyexecutor.StreamChunk{Payload: scaffoldEvent} + ch <- cliproxyexecutor.StreamChunk{Err: errors.New("upstream connection reset")} + + buffered, closed, err := readStreamBootstrap(context.Background(), ch) + if err == nil || !strings.Contains(err.Error(), "upstream connection reset") { + t.Fatalf("readStreamBootstrap error = %v, want upstream connection reset", err) + } + if closed { + t.Fatal("readStreamBootstrap closed = true, want false") + } + if len(buffered) != 0 { + t.Fatalf("len(buffered) = %d, want 0 on pre-output error", len(buffered)) + } + }) + + t.Run("reasoning_summary_text.delta with meaningful text forwards stream", func(t *testing.T) { + var detector StreamBootstrapDetector + scaffoldEvent := []byte("data: {\"type\":\"response.reasoning_summary_part.added\",\"sequence_number\":1,\"item_id\":\"rs_1\",\"output_index\":0,\"summary_index\":0,\"part\":{\"type\":\"summary_text\",\"text\":\"\"}}\n\n") + deltaEvent := []byte("data: {\"type\":\"response.reasoning_summary_text.delta\",\"sequence_number\":2,\"item_id\":\"rs_1\",\"output_index\":0,\"summary_index\":0,\"delta\":\"thinking step 1\"}\n\n") + + if detector.Observe(scaffoldEvent) { + t.Fatal("Observe(scaffoldEvent) = true, want false") + } + if !detector.Observe(deltaEvent) { + t.Fatal("Observe(deltaEvent with text) = false, want true") + } + if !detector.HasMeaningfulOutput() { + t.Fatal("HasMeaningfulOutput() = false after delta with text, want true") + } + }) +} + +func TestNonStreamMessageReasoningField(t *testing.T) { + t.Run("non-stream response with message.reasoning is not empty completion", func(t *testing.T) { + payload := []byte(`{"id":"chatcmpl-1","object":"chat.completion","created":12345,"model":"claude-3-5-sonnet","choices":[{"index":0,"message":{"role":"assistant","content":"","reasoning":"let me think about this"},"finish_reason":"stop"}]}`) + if IsEmptyCompletionPayload(payload) { + t.Fatal("IsEmptyCompletionPayload(reasoning payload) = true, want false") + } + }) + + t.Run("non-stream response with empty message.reasoning is empty completion", func(t *testing.T) { + payload := []byte(`{"id":"chatcmpl-1","object":"chat.completion","created":12345,"model":"claude-3-5-sonnet","choices":[{"index":0,"message":{"role":"assistant","content":"","reasoning":""},"finish_reason":"stop"}]}`) + if !IsEmptyCompletionPayload(payload) { + t.Fatal("IsEmptyCompletionPayload(empty reasoning payload) = false, want true") + } + }) +} + +func TestResponsesReasoningOutputItemDoneEmptySummary(t *testing.T) { + emptySummaryDone := []byte("data: {\"type\":\"response.output_item.done\",\"sequence_number\":1,\"output_index\":0,\"item\":{\"id\":\"rs_1\",\"type\":\"reasoning\",\"encrypted_content\":\"\",\"summary\":[{\"type\":\"summary_text\",\"text\":\"\"}]}}\n\n") + meaningfulSummaryDone := []byte("data: {\"type\":\"response.output_item.done\",\"sequence_number\":1,\"output_index\":0,\"item\":{\"id\":\"rs_1\",\"type\":\"reasoning\",\"encrypted_content\":\"\",\"summary\":[{\"type\":\"summary_text\",\"text\":\"real reasoning\"}]}}\n\n") + whitespaceSummaryDone := []byte("data: {\"type\":\"response.output_item.done\",\"sequence_number\":1,\"output_index\":0,\"item\":{\"id\":\"rs_1\",\"type\":\"reasoning\",\"encrypted_content\":\"\",\"summary\":[{\"type\":\"summary_text\",\"text\":\"\\n \\t\"}]}}\n\n") + encryptedSummaryDone := []byte("data: {\"type\":\"response.output_item.done\",\"sequence_number\":1,\"output_index\":0,\"item\":{\"id\":\"rs_1\",\"type\":\"reasoning\",\"encrypted_content\":\"sig_123\",\"summary\":[{\"type\":\"summary_text\",\"text\":\"\"}]}}\n\n") + + t.Run("empty reasoning summary array in output_item.done does not mark meaningful and allows bootstrap error failover", func(t *testing.T) { + var detector StreamBootstrapDetector + if detector.Observe(emptySummaryDone) { + t.Fatal("detector.Observe() = true for empty reasoning summary in output_item.done, want false") + } + if detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = true for empty reasoning summary in output_item.done, want false") + } + + errUpstream := errors.New("upstream failed after empty reasoning summary") + ch := make(chan cliproxyexecutor.StreamChunk, 2) + ch <- cliproxyexecutor.StreamChunk{Payload: emptySummaryDone} + 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("meaningful reasoning summary array in output_item.done marks meaningful and forwards", func(t *testing.T) { + var detector StreamBootstrapDetector + if !detector.Observe(meaningfulSummaryDone) { + t.Fatal("detector.Observe() = false for meaningful reasoning summary in output_item.done, want true") + } + if !detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = false for meaningful reasoning summary in output_item.done, want true") + } + }) + + t.Run("whitespace reasoning summary array in output_item.done does not mark meaningful", func(t *testing.T) { + var detector StreamBootstrapDetector + if detector.Observe(whitespaceSummaryDone) { + t.Fatal("detector.Observe() = true for whitespace reasoning summary in output_item.done, want false") + } + if detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = true for whitespace reasoning summary in output_item.done, want false") + } + }) + + t.Run("reasoning output_item.done with encrypted_content and empty summary marks meaningful", func(t *testing.T) { + var detector StreamBootstrapDetector + if !detector.Observe(encryptedSummaryDone) { + t.Fatal("detector.Observe() = false for reasoning with encrypted_content and empty summary, want true") + } + if !detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = false for reasoning with encrypted_content and empty summary, want true") + } + }) +} + +func TestResponsesAnnotationsOutputItemAndContentPart(t *testing.T) { + citationContentPartDone := []byte("data: {\"type\":\"response.content_part.done\",\"sequence_number\":1,\"item_id\":\"msg_1\",\"output_index\":0,\"content_index\":0,\"part\":{\"type\":\"output_text\",\"annotations\":[{\"type\":\"url_citation\",\"url\":\"https://example.com\",\"title\":\"Example\"}],\"logprobs\":[],\"text\":\"\"}}\n\n") + citationOutputItemDone := []byte("data: {\"type\":\"response.output_item.done\",\"sequence_number\":1,\"output_index\":0,\"item\":{\"id\":\"msg_1\",\"type\":\"message\",\"status\":\"completed\",\"content\":[{\"type\":\"output_text\",\"annotations\":[{\"type\":\"url_citation\",\"url\":\"https://example.com\",\"title\":\"Example\"}],\"logprobs\":[],\"text\":\"\"}],\"role\":\"assistant\"}}\n\n") + emptyAnnotationsPartDone := []byte("data: {\"type\":\"response.content_part.done\",\"sequence_number\":1,\"item_id\":\"msg_1\",\"output_index\":0,\"content_index\":0,\"part\":{\"type\":\"output_text\",\"annotations\":[],\"logprobs\":[],\"text\":\"\"}}\n\n") + emptyAnnotationsItemDone := []byte("data: {\"type\":\"response.output_item.done\",\"sequence_number\":1,\"output_index\":0,\"item\":{\"id\":\"msg_1\",\"type\":\"message\",\"status\":\"completed\",\"content\":[{\"type\":\"output_text\",\"annotations\":[],\"logprobs\":[],\"text\":\"\"}],\"role\":\"assistant\"}}\n\n") + emptyObjectsAnnotationsPartDone := []byte("data: {\"type\":\"response.content_part.done\",\"sequence_number\":1,\"item_id\":\"msg_1\",\"output_index\":0,\"content_index\":0,\"part\":{\"type\":\"output_text\",\"annotations\":[{},{}],\"logprobs\":[],\"text\":\"\"}}\n\n") + emptyObjectsAnnotationsItemDone := []byte("data: {\"type\":\"response.output_item.done\",\"sequence_number\":1,\"output_index\":0,\"item\":{\"id\":\"msg_1\",\"type\":\"message\",\"status\":\"completed\",\"content\":[{\"type\":\"output_text\",\"annotations\":[{}]}],\"role\":\"assistant\"}}\n\n") + normalTextPartDone := []byte("data: {\"type\":\"response.content_part.done\",\"sequence_number\":1,\"item_id\":\"msg_1\",\"output_index\":0,\"content_index\":0,\"part\":{\"type\":\"output_text\",\"annotations\":[],\"logprobs\":[],\"text\":\"Hello world\"}}\n\n") + normalTextItemDone := []byte("data: {\"type\":\"response.output_item.done\",\"sequence_number\":1,\"output_index\":0,\"item\":{\"id\":\"msg_1\",\"type\":\"message\",\"status\":\"completed\",\"content\":[{\"type\":\"output_text\",\"annotations\":[],\"logprobs\":[],\"text\":\"Hello world\"}],\"role\":\"assistant\"}}\n\n") + + t.Run("content_part.done with text empty but annotations non-empty marks meaningful and forwards", func(t *testing.T) { + var detector StreamBootstrapDetector + if !detector.Observe(citationContentPartDone) { + t.Fatal("detector.Observe() = false for content_part.done with annotations, want true") + } + if !detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = false for content_part.done with annotations, want true") + } + }) + + t.Run("output_item.done with text empty but annotations non-empty marks meaningful and forwards", func(t *testing.T) { + var detector StreamBootstrapDetector + if !detector.Observe(citationOutputItemDone) { + t.Fatal("detector.Observe() = false for output_item.done with annotations, want true") + } + if !detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = false for output_item.done with annotations, want true") + } + }) + + t.Run("empty annotations array remains empty-eligible and allows failover", func(t *testing.T) { + var detector StreamBootstrapDetector + if detector.Observe(emptyAnnotationsPartDone) { + t.Fatal("detector.Observe() = true for empty annotations in content_part.done, want false") + } + if detector.Observe(emptyAnnotationsItemDone) { + t.Fatal("detector.Observe() = true for empty annotations in output_item.done, want false") + } + if detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = true for empty annotations, want false") + } + + errUpstream := errors.New("upstream failed after empty annotations") + ch := make(chan cliproxyexecutor.StreamChunk, 2) + ch <- cliproxyexecutor.StreamChunk{Payload: emptyAnnotationsPartDone} + 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("annotations with empty objects remain empty-eligible", func(t *testing.T) { + var detector StreamBootstrapDetector + if detector.Observe(emptyObjectsAnnotationsPartDone) { + t.Fatal("detector.Observe() = true for empty objects in annotations (content_part.done), want false") + } + if detector.Observe(emptyObjectsAnnotationsItemDone) { + t.Fatal("detector.Observe() = true for empty objects in annotations (output_item.done), want false") + } + if detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = true for empty objects in annotations, want false") + } + }) + + t.Run("normal text content path unaffected", func(t *testing.T) { + var d1 StreamBootstrapDetector + if !d1.Observe(normalTextPartDone) { + t.Fatal("d1.Observe() = false for normal text content_part.done, want true") + } + if !d1.HasMeaningfulOutput() { + t.Fatal("d1.HasMeaningfulOutput() = false for normal text content_part.done, want true") + } + + var d2 StreamBootstrapDetector + if !d2.Observe(normalTextItemDone) { + t.Fatal("d2.Observe() = false for normal text output_item.done, want true") + } + if !d2.HasMeaningfulOutput() { + t.Fatal("d2.HasMeaningfulOutput() = false for normal text output_item.done, want true") + } + }) +} + +func TestGeminiGroundingMetadata(t *testing.T) { + groundingChunkStopPayload := []byte(`{"candidates":[{"index":0,"finishReason":"STOP","groundingMetadata":{"groundingChunks":[{"web":{"uri":"https://example.com/weather","title":"Beijing Weather"}}]}}]}`) + emptyGroundingMetaPayload := []byte(`{"candidates":[{"index":0,"finishReason":"STOP","groundingMetadata":{}}]}`) + emptyGroundingChunksPayload := []byte(`{"candidates":[{"index":0,"finishReason":"STOP","groundingMetadata":{"groundingChunks":[{}]}}]}`) + normalContentPayload := []byte(`{"candidates":[{"index":0,"finishReason":"STOP","content":{"parts":[{"text":"hello"}]}}]}`) + + t.Run("candidate STOP with no content but meaningful groundingChunks is not terminal-empty and is forwarded", func(t *testing.T) { + if IsEmptyCompletionPayload(groundingChunkStopPayload) { + t.Fatal("IsEmptyCompletionPayload(groundingChunkStopPayload) = true, want false") + } + + var detector StreamBootstrapDetector + ssePayload := []byte("data: " + string(groundingChunkStopPayload) + "\n\n") + if !detector.Observe(ssePayload) { + t.Fatal("detector.Observe(ssePayload) = false for meaningful grounding chunks, want true") + } + if !detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = false for meaningful grounding chunks, want true") + } + if detector.IsTerminalEmpty() { + t.Fatal("detector.IsTerminalEmpty() = true for meaningful grounding chunks, want false") + } + }) + + t.Run("groundingMetadata empty object remains terminal-empty-eligible", func(t *testing.T) { + if !IsEmptyCompletionPayload(emptyGroundingMetaPayload) { + t.Fatal("IsEmptyCompletionPayload(emptyGroundingMetaPayload) = false, want true") + } + + var detector StreamBootstrapDetector + ssePayload := []byte("data: " + string(emptyGroundingMetaPayload) + "\n\n") + if detector.Observe(ssePayload) { + t.Fatal("detector.Observe(ssePayload) = true for empty grounding metadata, want false") + } + if detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = true for empty grounding metadata, want false") + } + if !detector.IsTerminalEmpty() { + t.Fatal("detector.IsTerminalEmpty() = false for empty grounding metadata at STOP, want true") + } + }) + + t.Run("groundingMetadata with empty chunk objects remains terminal-empty-eligible", func(t *testing.T) { + if !IsEmptyCompletionPayload(emptyGroundingChunksPayload) { + t.Fatal("IsEmptyCompletionPayload(emptyGroundingChunksPayload) = false, want true") + } + + var detector StreamBootstrapDetector + ssePayload := []byte("data: " + string(emptyGroundingChunksPayload) + "\n\n") + if detector.Observe(ssePayload) { + t.Fatal("detector.Observe(ssePayload) = true for empty chunk objects, want false") + } + if detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = true for empty chunk objects, want false") + } + if !detector.IsTerminalEmpty() { + t.Fatal("detector.IsTerminalEmpty() = false for empty chunk objects at STOP, want true") + } + }) + + t.Run("normal content path unaffected", func(t *testing.T) { + if IsEmptyCompletionPayload(normalContentPayload) { + t.Fatal("IsEmptyCompletionPayload(normalContentPayload) = true, want false") + } + + var detector StreamBootstrapDetector + ssePayload := []byte("data: " + string(normalContentPayload) + "\n\n") + if !detector.Observe(ssePayload) { + t.Fatal("detector.Observe(ssePayload) = false for normal content, want true") + } + if !detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = false for normal content, want true") + } + }) + + t.Run("grounding-only with multi-candidate gating marks finished properly", func(t *testing.T) { + multiChunk1 := []byte("data: {\"candidates\":[{\"index\":0,\"finishReason\":\"STOP\",\"groundingMetadata\":{\"groundingChunks\":[{\"web\":{\"uri\":\"https://example.com\",\"title\":\"Example\"}}]}}]}\n\n") + multiChunk2 := []byte("data: {\"candidates\":[{\"index\":1,\"finishReason\":\"STOP\",\"content\":{\"parts\":[{\"text\":\"second candidate text\"}]}}]}\n\n") + + var detector StreamBootstrapDetector + detector.SetExpectedChoices(2) + + if !detector.Observe(multiChunk1) { + t.Fatal("detector.Observe(multiChunk1) = false for candidate 0 with grounding, want true") + } + if !detector.HasMeaningfulOutput() { + t.Fatal("detector.HasMeaningfulOutput() = false after candidate 0 with grounding, want true") + } + if detector.IsTerminalEmpty() { + t.Fatal("detector.IsTerminalEmpty() = true while awaiting candidate 1, want false") + } + + if !detector.Observe(multiChunk2) { + t.Fatal("detector.Observe(multiChunk2) = false for candidate 1, want true") + } + if detector.IsTerminalEmpty() { + t.Fatal("detector.IsTerminalEmpty() = true after both candidates finished, want false") + } + }) +} + +func TestClaudeStreamBootstrapStopReasonWithoutMessageStop(t *testing.T) { + t.Run("claude stream message_delta with stop_reason end_turn without message_stop is terminal empty", 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") + emptyDeltaChunk := []byte("event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"\"}}\n\n") + stopDeltaChunk := []byte("event: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":0}}\n\n") + + if detector.Observe(startChunk) { + t.Fatal("Observe(message_start) = true, want false") + } + if detector.Observe(emptyDeltaChunk) { + t.Fatal("Observe(empty content_block_delta) = true, want false") + } + if detector.Observe(stopDeltaChunk) { + t.Fatal("Observe(message_delta stop_reason:end_turn) = true, want false") + } + if !detector.IsTerminalEmpty() { + t.Fatal("IsTerminalEmpty() = false on message_delta end_turn without message_stop, want true") + } + }) + + t.Run("conductor readStreamBootstrap with claude stop_reason end_turn over open channel classifies empty", func(t *testing.T) { + ch := make(chan cliproxyexecutor.StreamChunk, 3) + ch <- cliproxyexecutor.StreamChunk{Payload: []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")} + ch <- cliproxyexecutor.StreamChunk{Payload: []byte("event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"\"}}\n\n")} + ch <- cliproxyexecutor.StreamChunk{Payload: []byte("event: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":0}}\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) != 3 { + t.Fatalf("len(buffered) = %d, want 3", len(buffered)) + } + }) + + t.Run("claude stream message_delta with stop_reason tool_use without message_stop is terminal empty", func(t *testing.T) { + var detector StreamBootstrapDetector + startChunk := []byte("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_2\",\"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") + emptyToolChunk := []byte("event: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"tool_use\",\"input\":null}}\n\n") + stopDeltaChunk := []byte("event: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"tool_use\"},\"usage\":{\"output_tokens\":0}}\n\n") + + if detector.Observe(startChunk) { + t.Fatal("Observe(message_start) = true, want false") + } + if detector.Observe(emptyToolChunk) { + t.Fatal("Observe(empty tool_use block) = true, want false") + } + if detector.Observe(stopDeltaChunk) { + t.Fatal("Observe(message_delta stop_reason:tool_use) = true, want false") + } + if !detector.IsTerminalEmpty() { + t.Fatal("IsTerminalEmpty() = false on message_delta tool_use without message_stop, want true") + } + }) + + t.Run("conductor readStreamBootstrap with claude stop_reason tool_use over open channel classifies empty", func(t *testing.T) { + ch := make(chan cliproxyexecutor.StreamChunk, 3) + ch <- cliproxyexecutor.StreamChunk{Payload: []byte("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_2\",\"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")} + ch <- cliproxyexecutor.StreamChunk{Payload: []byte("event: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"tool_use\",\"input\":null}}\n\n")} + ch <- cliproxyexecutor.StreamChunk{Payload: []byte("event: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"tool_use\"},\"usage\":{\"output_tokens\":0}}\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) != 3 { + t.Fatalf("len(buffered) = %d, want 3", len(buffered)) + } + }) +} + +func TestExecuteInteractionsEmptyCompletionRotatesAuth(t *testing.T) { + executor := &emptyCompletionTestExecutor{ + executePayloads: map[string][]byte{}, + executeCalls: map[string]int{}, + emptyPayload: []byte(`{"id":"interaction_1","object":"interaction","status":"completed","steps":[],"usage":{"output_tokens":0,"total_output_tokens":0}}`), + contentPayload: []byte(`{"id":"interaction_2","object":"interaction","status":"completed","steps":[{"type":"model_output","content":[{"type":"text","text":"real"}]}],"usage":{"output_tokens":5,"total_output_tokens":5}}`), + } + manager, ids, model, capture := newEmptyCompletionTestManager(t, executor) + + resp, err := manager.Execute(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}) + if err != nil { + t.Fatalf("Execute() error = %v", err) + } + assertRotatesToContent(t, ids, executor.firstExecute, string(resp.Payload), "real", capture) +} + +func TestExecuteInteractionsMeaningfulOutputNotRotated(t *testing.T) { + t.Run("step with text content", func(t *testing.T) { + executor := &emptyCompletionTestExecutor{ + executePayloads: map[string][]byte{}, + executeCalls: map[string]int{}, + emptyPayload: []byte(`{"id":"interaction_1","object":"interaction","status":"completed","steps":[{"type":"model_output","content":[{"type":"text","text":"meaningful text"}]}]}`), + } + manager, ids, model, capture := newEmptyCompletionTestManager(t, executor) + + 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), "meaningful text") { + t.Fatalf("resp.Payload = %s, want meaningful text", string(resp.Payload)) + } + if len(capture.Results()) != 1 || !capture.Results()[0].Success || capture.Results()[0].AuthID != executor.firstExecute { + t.Fatalf("first auth should succeed without rotation, results: %+v, first: %s, ids: %v", capture.Results(), executor.firstExecute, ids) + } + }) + + t.Run("step with function_call", func(t *testing.T) { + executor := &emptyCompletionTestExecutor{ + executePayloads: map[string][]byte{}, + executeCalls: map[string]int{}, + emptyPayload: []byte(`{"id":"interaction_1","object":"interaction","status":"completed","steps":[{"type":"function_call","name":"get_weather","arguments":{"location":"Tokyo"}}]}`), + } + manager, _, model, capture := newEmptyCompletionTestManager(t, executor) + + 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), "get_weather") { + t.Fatalf("resp.Payload = %s, want get_weather", string(resp.Payload)) + } + if len(capture.Results()) != 1 || !capture.Results()[0].Success || capture.Results()[0].AuthID != executor.firstExecute { + t.Fatalf("first auth should succeed without rotation, results: %+v", capture.Results()) + } + }) + + t.Run("non-zero usage with empty steps", func(t *testing.T) { + executor := &emptyCompletionTestExecutor{ + executePayloads: map[string][]byte{}, + executeCalls: map[string]int{}, + emptyPayload: []byte(`{"id":"interaction_1","object":"interaction","status":"completed","steps":[],"usage":{"output_tokens":3,"total_output_tokens":3}}`), + } + manager, _, model, capture := newEmptyCompletionTestManager(t, executor) + + resp, err := manager.Execute(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}) + if err != nil { + t.Fatalf("Execute() error = %v", err) + } + _ = resp + if err != nil { + t.Fatalf("Execute() error = %v", err) + } + if len(capture.Results()) != 1 || !capture.Results()[0].Success || capture.Results()[0].AuthID != executor.firstExecute { + t.Fatalf("first auth should succeed without rotation, results: %+v", capture.Results()) + } + }) +} + +func TestStreamBootstrapDetectorInteractionsScaffoldAllowsFailover(t *testing.T) { + t.Run("scaffold events do not commit stream and classify as empty on completion", func(t *testing.T) { + var detector StreamBootstrapDetector + chunk1 := []byte("event: interaction.created\ndata: {\"event_type\":\"interaction.created\",\"interaction\":{\"id\":\"int_1\",\"object\":\"interaction\",\"status\":\"in_progress\"}}\n\n") + chunk2 := []byte("event: interaction.status_update\ndata: {\"event_type\":\"interaction.status_update\",\"interaction_id\":\"int_1\",\"status\":\"in_progress\"}\n\n") + chunk3 := []byte("event: step.start\ndata: {\"event_type\":\"step.start\",\"index\":0,\"step\":{\"type\":\"model_output\"}}\n\n") + chunk4 := []byte("event: step.stop\ndata: {\"event_type\":\"step.stop\",\"index\":0}\n\n") + chunk5 := []byte("event: interaction.completed\ndata: {\"event_type\":\"interaction.completed\",\"interaction\":{\"id\":\"int_1\",\"status\":\"completed\",\"steps\":[],\"usage\":{\"output_tokens\":0}}}\n\n") + + if detector.Observe(chunk1) { + t.Fatal("Observe(interaction.created) = true, want false") + } + if detector.Observe(chunk2) { + t.Fatal("Observe(interaction.status_update) = true, want false") + } + if detector.Observe(chunk3) { + t.Fatal("Observe(step.start) = true, want false") + } + if detector.Observe(chunk4) { + t.Fatal("Observe(step.stop) = true, want false") + } + if detector.Observe(chunk5) { + t.Fatal("Observe(interaction.completed empty) = true, want false") + } + if !detector.Finish() { + t.Fatal("detector.Finish() = false, want true (empty completion)") + } + }) + + t.Run("scaffold stream rotates auth when upstream fails over", func(t *testing.T) { + executor := &emptyCompletionTestExecutor{ + streamPayloads: map[string][][]byte{}, + streamCalls: map[string]int{}, + emptyStreamPayload: [][]byte{ + []byte("event: interaction.created\ndata: {\"event_type\":\"interaction.created\",\"interaction\":{\"id\":\"int_1\",\"object\":\"interaction\",\"status\":\"in_progress\"}}\n\n"), + []byte("event: step.start\ndata: {\"event_type\":\"step.start\",\"index\":0,\"step\":{\"type\":\"model_output\"}}\n\n"), + []byte("event: interaction.completed\ndata: {\"event_type\":\"interaction.completed\",\"interaction\":{\"id\":\"int_1\",\"status\":\"completed\",\"steps\":[],\"usage\":{\"output_tokens\":0}}}\n\n"), + }, + contentStreamPayload: [][]byte{ + []byte("event: interaction.created\ndata: {\"event_type\":\"interaction.created\",\"interaction\":{\"id\":\"int_2\",\"object\":\"interaction\",\"status\":\"in_progress\"}}\n\n"), + []byte("event: step.start\ndata: {\"event_type\":\"step.start\",\"index\":0,\"step\":{\"type\":\"model_output\"}}\n\n"), + []byte("event: step.delta\ndata: {\"event_type\":\"step.delta\",\"index\":0,\"delta\":{\"type\":\"text\",\"text\":\"hello from stream\"}}\n\n"), + []byte("event: step.stop\ndata: {\"event_type\":\"step.stop\",\"index\":0}\n\n"), + []byte("event: interaction.completed\ndata: {\"event_type\":\"interaction.completed\",\"interaction\":{\"id\":\"int_2\",\"status\":\"completed\",\"steps\":[],\"usage\":{\"output_tokens\":5}}}\n\n"), + }, + } + manager, ids, model, capture := newEmptyCompletionTestManager(t, executor) + + stream, err := manager.ExecuteStream(context.Background(), []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 from stream", capture) + }) +} + +func TestExecuteInteractionsMediaOutputNotRotated(t *testing.T) { + t.Run("step with inline media data", func(t *testing.T) { + executor := &emptyCompletionTestExecutor{ + executePayloads: map[string][]byte{}, + executeCalls: map[string]int{}, + emptyPayload: []byte(`{"id":"interaction_1","object":"interaction","status":"completed","steps":[{"type":"model_output","content":[{"type":"image","mime_type":"image/png","data":"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="}]}],"usage":{"output_tokens":0}}`), + } + manager, _, model, capture := newEmptyCompletionTestManager(t, executor) + + 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), "iVBORw0KGgo") { + t.Fatalf("resp.Payload = %s, want inline image data", string(resp.Payload)) + } + if len(capture.Results()) != 1 || !capture.Results()[0].Success || capture.Results()[0].AuthID != executor.firstExecute { + t.Fatalf("first auth should succeed without rotation for media output, results: %+v", capture.Results()) + } + }) + + t.Run("step with file_uri", func(t *testing.T) { + executor := &emptyCompletionTestExecutor{ + executePayloads: map[string][]byte{}, + executeCalls: map[string]int{}, + emptyPayload: []byte(`{"id":"interaction_1","object":"interaction","status":"completed","steps":[{"type":"model_output","content":[{"type":"document","file_uri":"files/abc123xyz"}]}]}`), + } + manager, _, model, capture := newEmptyCompletionTestManager(t, executor) + + 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), "files/abc123xyz") { + t.Fatalf("resp.Payload = %s, want file_uri", string(resp.Payload)) + } + if len(capture.Results()) != 1 || !capture.Results()[0].Success || capture.Results()[0].AuthID != executor.firstExecute { + t.Fatalf("first auth should succeed without rotation for file_uri media output, results: %+v", capture.Results()) + } + }) + + t.Run("stream delta with media content", func(t *testing.T) { + executor := &emptyCompletionTestExecutor{ + streamPayloads: map[string][][]byte{}, + streamCalls: map[string]int{}, + emptyStreamPayload: [][]byte{ + []byte("event: interaction.created\ndata: {\"event_type\":\"interaction.created\",\"interaction\":{\"id\":\"int_1\",\"object\":\"interaction\",\"status\":\"in_progress\"}}\n\n"), + []byte("event: step.start\ndata: {\"event_type\":\"step.start\",\"index\":0,\"step\":{\"type\":\"model_output\"}}\n\n"), + []byte("event: step.delta\ndata: {\"event_type\":\"step.delta\",\"index\":0,\"delta\":{\"type\":\"content\",\"content\":{\"type\":\"image\",\"mime_type\":\"image/png\",\"data\":\"iVBORw0KGgo=\"}}}\n\n"), + []byte("event: step.stop\ndata: {\"event_type\":\"step.stop\",\"index\":0}\n\n"), + []byte("event: interaction.completed\ndata: {\"event_type\":\"interaction.completed\",\"interaction\":{\"id\":\"int_1\",\"status\":\"completed\",\"steps\":[],\"usage\":{\"output_tokens\":0}}}\n\n"), + }, + } + manager, _, model, capture := newEmptyCompletionTestManager(t, executor) + + stream, err := manager.ExecuteStream(context.Background(), []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) + } + if !strings.Contains(got.String(), "iVBORw0KGgo=") { + t.Fatalf("stream payload = %q, want media data from first auth", got.String()) + } + if len(capture.Results()) != 1 || !capture.Results()[0].Success || capture.Results()[0].AuthID != executor.firstStream { + t.Fatalf("first auth should succeed without rotation for media delta, results: %+v", capture.Results()) + } + }) +} + +func TestExecuteStream_InteractionsTerminalCompletedWithoutClosingChannelRotatesAuth(t *testing.T) { + executor := &emptyCompletionTestExecutor{ + streamPayloads: map[string][][]byte{}, + streamCalls: map[string]int{}, + leaveStreamOpen: true, + emptyStreamPayload: [][]byte{ + []byte("event: interaction.created\ndata: {\"event_type\":\"interaction.created\",\"interaction\":{\"id\":\"int_1\",\"object\":\"interaction\",\"status\":\"in_progress\"}}\n\n"), + []byte("event: step.start\ndata: {\"event_type\":\"step.start\",\"index\":0,\"step\":{\"type\":\"model_output\"}}\n\n"), + []byte("event: interaction.completed\ndata: {\"event_type\":\"interaction.completed\",\"interaction\":{\"id\":\"int_1\",\"status\":\"completed\",\"steps\":[],\"usage\":{\"output_tokens\":0}}}\n\n"), + }, + contentStreamPayload: [][]byte{ + []byte("event: interaction.created\ndata: {\"event_type\":\"interaction.created\",\"interaction\":{\"id\":\"int_2\",\"object\":\"interaction\",\"status\":\"in_progress\"}}\n\n"), + []byte("event: step.start\ndata: {\"event_type\":\"step.start\",\"index\":0,\"step\":{\"type\":\"model_output\"}}\n\n"), + []byte("event: step.delta\ndata: {\"event_type\":\"step.delta\",\"index\":0,\"delta\":{\"type\":\"text\",\"text\":\"content_after_rotation\"}}\n\n"), + []byte("event: step.stop\ndata: {\"event_type\":\"step.stop\",\"index\":0}\n\n"), + []byte("event: interaction.completed\ndata: {\"event_type\":\"interaction.completed\",\"interaction\":{\"id\":\"int_2\",\"status\":\"completed\",\"steps\":[],\"usage\":{\"output_tokens\":5}}}\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(), "content_after_rotation", capture) +} + +func TestStreamBootstrapDetectorInteractionsTerminalCompletedWithoutClose(t *testing.T) { + var detector StreamBootstrapDetector + createdChunk := []byte("event: interaction.created\ndata: {\"event_type\":\"interaction.created\",\"interaction\":{\"id\":\"int_1\",\"object\":\"interaction\",\"status\":\"in_progress\"}}\n\n") + stepStartChunk := []byte("event: step.start\ndata: {\"event_type\":\"step.start\",\"index\":0,\"step\":{\"type\":\"model_output\"}}\n\n") + completedChunk := []byte("event: interaction.completed\ndata: {\"event_type\":\"interaction.completed\",\"interaction\":{\"id\":\"int_1\",\"status\":\"completed\",\"steps\":[],\"usage\":{\"output_tokens\":0}}}\n\n") + + if detector.Observe(createdChunk) { + t.Fatal("Observe(interaction.created) = true, want false") + } + if detector.Observe(stepStartChunk) { + t.Fatal("Observe(step.start) = true, want false") + } + if detector.Observe(completedChunk) { + t.Fatal("Observe(interaction.completed) = true, want false") + } + if !detector.IsTerminalEmpty() { + t.Fatal("IsTerminalEmpty() = false on interaction.completed without closing transport, want true") + } + + ch := make(chan cliproxyexecutor.StreamChunk, 3) + ch <- cliproxyexecutor.StreamChunk{Payload: createdChunk} + ch <- cliproxyexecutor.StreamChunk{Payload: stepStartChunk} + ch <- cliproxyexecutor.StreamChunk{Payload: completedChunk} + // Channel left open + + ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond) + 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) != 3 { + t.Fatalf("len(buffered) = %d, want 3", len(buffered)) + } +} + +func TestEvalProviderErrorInteractionsNestedFailure(t *testing.T) { + cases := []struct { + name string + payload string + wantErr bool + wantStatus int + wantRetryable bool + }{ + { + name: "nested rate limit error", + payload: `{"event_type":"interaction.failed","interaction":{"id":"int_1","status":"failed","error":{"code":429,"message":"Resource has been exhausted","status":"RESOURCE_EXHAUSTED"}}}`, + wantErr: true, + wantStatus: 429, + wantRetryable: true, + }, + { + name: "nested invalid request stays non retryable", + payload: `{"event_type":"interaction.failed","interaction":{"id":"int_1","status":"failed","error":{"code":400,"message":"invalid request","type":"invalid_request_error"}}}`, + wantErr: true, + wantStatus: 400, + wantRetryable: false, + }, + { + name: "failed event without nested detail", + payload: `{"event_type":"interaction.failed","interaction":{"id":"int_1","status":"failed"}}`, + wantErr: true, + wantStatus: 502, + wantRetryable: true, + }, + { + name: "failed status without failed event type", + payload: `{"event_type":"interaction.status_update","interaction":{"id":"int_1","status":"failed"}}`, + wantErr: true, + wantStatus: 502, + wantRetryable: true, + }, + { + name: "completed interaction is not an error", + payload: `{"event_type":"interaction.completed","interaction":{"id":"int_1","status":"completed","steps":[]}}`, + wantErr: false, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got := evalProviderError([]byte(tc.payload), "") + if !tc.wantErr { + if got != nil { + t.Fatalf("evalProviderError() = %+v, want nil", got) + } + return + } + if got == nil { + t.Fatal("evalProviderError() = nil, want a classified provider error so the request can fail over") + } + if got.HTTPStatus != tc.wantStatus { + t.Fatalf("HTTPStatus = %d, want %d (err=%+v)", got.HTTPStatus, tc.wantStatus, got) + } + if got.Retryable != tc.wantRetryable { + t.Fatalf("Retryable = %v, want %v (err=%+v)", got.Retryable, tc.wantRetryable, got) + } + }) + } +} + +func TestReadStreamBootstrapInteractionsNestedFailureSurfacesProviderError(t *testing.T) { + createdChunk := []byte("event: interaction.created\ndata: {\"event_type\":\"interaction.created\",\"interaction\":{\"id\":\"int_1\",\"object\":\"interaction\",\"status\":\"in_progress\"}}\n\n") + failedChunk := []byte("event: interaction.failed\ndata: {\"event_type\":\"interaction.failed\",\"interaction\":{\"id\":\"int_1\",\"status\":\"failed\",\"error\":{\"code\":429,\"message\":\"Resource has been exhausted\",\"status\":\"RESOURCE_EXHAUSTED\"}}}\n\n") + + ch := make(chan cliproxyexecutor.StreamChunk, 2) + ch <- cliproxyexecutor.StreamChunk{Payload: createdChunk} + ch <- cliproxyexecutor.StreamChunk{Payload: failedChunk} + close(ch) + + ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond) + defer cancel() + + _, _, err := readStreamBootstrap(ctx, ch) + if err == nil { + t.Fatal("readStreamBootstrap error = nil on nested interaction failure, want a provider error so the auth is rotated") + } + provErr, ok := err.(*Error) + if !ok { + t.Fatalf("readStreamBootstrap error type = %T, want *Error", err) + } + if provErr.HTTPStatus != 429 { + t.Fatalf("HTTPStatus = %d, want 429 (err=%+v)", provErr.HTTPStatus, provErr) + } + if !provErr.Retryable { + t.Fatalf("Retryable = false, want true (err=%+v)", provErr) + } +} + +func TestStreamBootstrapDetectorInteractionsSignatureOnlyIsNotEmpty(t *testing.T) { + createdChunk := []byte("event: interaction.created\ndata: {\"event_type\":\"interaction.created\",\"interaction\":{\"id\":\"int_1\",\"object\":\"interaction\",\"status\":\"in_progress\"}}\n\n") + completedChunk := []byte("event: interaction.completed\ndata: {\"event_type\":\"interaction.completed\",\"interaction\":{\"id\":\"int_1\",\"status\":\"completed\",\"usage\":{\"output_tokens\":0}}}\n\n") + + cases := []struct { + name string + payload string + }{ + { + name: "step level thought_signature", + payload: "event: step.stop\ndata: {\"event_type\":\"step.stop\",\"index\":0,\"step\":{\"type\":\"model_output\",\"thought_signature\":\"sig-step\"}}\n\n", + }, + { + name: "step level camelCase thoughtSignature", + payload: "event: step.stop\ndata: {\"event_type\":\"step.stop\",\"index\":0,\"step\":{\"type\":\"model_output\",\"thoughtSignature\":\"sig-camel\"}}\n\n", + }, + { + name: "step level signature", + payload: "event: step.stop\ndata: {\"event_type\":\"step.stop\",\"index\":0,\"step\":{\"type\":\"model_output\",\"signature\":\"sig-plain\"}}\n\n", + }, + { + name: "step level encrypted_content", + payload: "event: step.stop\ndata: {\"event_type\":\"step.stop\",\"index\":0,\"step\":{\"type\":\"model_output\",\"encrypted_content\":\"enc-blob\"}}\n\n", + }, + { + name: "step level extra_content google thought_signature", + payload: "event: step.stop\ndata: {\"event_type\":\"step.stop\",\"index\":0,\"step\":{\"type\":\"model_output\",\"extra_content\":{\"google\":{\"thought_signature\":\"sig-extra\"}}}}\n\n", + }, + { + name: "content level thought_signature", + payload: "event: step.stop\ndata: {\"event_type\":\"step.stop\",\"index\":0,\"step\":{\"type\":\"model_output\",\"content\":[{\"type\":\"thinking\",\"thought_signature\":\"sig-content\"}]}}\n\n", + }, + { + name: "delta level thought_signature", + payload: "event: step.delta\ndata: {\"event_type\":\"step.delta\",\"index\":0,\"delta\":{\"type\":\"thinking\",\"thought_signature\":\"sig-delta\"}}\n\n", + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + var detector StreamBootstrapDetector + detector.Observe(createdChunk) + detector.Observe([]byte(tc.payload)) + detector.Observe(completedChunk) + + if detector.IsTerminalEmpty() { + t.Fatalf("IsTerminalEmpty() = true for a signature-carrying interaction, want false; signature would be dropped and the turn retried: %s", tc.payload) + } + if !detector.HasMeaningfulOutput() { + t.Fatalf("HasMeaningfulOutput() = false for a signature-carrying interaction, want true: %s", tc.payload) + } + }) + } +} + +func TestStreamBootstrapDetectorInteractionsWithoutSignatureStaysEmpty(t *testing.T) { + var detector StreamBootstrapDetector + detector.Observe([]byte("event: interaction.created\ndata: {\"event_type\":\"interaction.created\",\"interaction\":{\"id\":\"int_1\",\"object\":\"interaction\",\"status\":\"in_progress\"}}\n\n")) + detector.Observe([]byte("event: step.stop\ndata: {\"event_type\":\"step.stop\",\"index\":0,\"step\":{\"type\":\"model_output\",\"content\":[{\"type\":\"text\",\"text\":\"\"}]}}\n\n")) + detector.Observe([]byte("event: interaction.completed\ndata: {\"event_type\":\"interaction.completed\",\"interaction\":{\"id\":\"int_1\",\"status\":\"completed\",\"usage\":{\"output_tokens\":0}}}\n\n")) + + if !detector.IsTerminalEmpty() { + t.Fatal("IsTerminalEmpty() = false for a genuinely empty interaction, want true") + } +} diff --git a/sdk/cliproxy/auth/export_test.go b/sdk/cliproxy/auth/export_test.go new file mode 100644 index 00000000000..d8e6a0225a9 --- /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 e5e06821e13..8c1d166a327 100644 --- a/sdk/cliproxy/auth/home_concurrency.go +++ b/sdk/cliproxy/auth/home_concurrency.go @@ -277,6 +277,13 @@ func verifyAccountedHomeConcurrencyIdentity(tuple homeConcurrencyTuple, auth *Au // SafeResponseHeaders returns trusted response headers only for concrete // Home-generated retry errors. 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 busy.SafeResponseHeaders() diff --git a/sdk/cliproxy/auth/home_execution_paths_test.go b/sdk/cliproxy/auth/home_execution_paths_test.go index adaa5560c72..151e95ccb72 100644 --- a/sdk/cliproxy/auth/home_execution_paths_test.go +++ b/sdk/cliproxy/auth/home_execution_paths_test.go @@ -5,6 +5,7 @@ import ( "encoding/json" "net/http" "strconv" + "strings" "sync" "sync/atomic" "testing" @@ -524,8 +525,9 @@ func (e *lifecycleRetryExecutor) ExecuteStream(ctx context.Context, _ *Auth, _ c e.firstCtx = ctx return nil, &Error{HTTPStatus: http.StatusUpgradeRequired, Message: "websocket upgrade required"} } - chunks := make(chan cliproxyexecutor.StreamChunk, 1) - chunks <- cliproxyexecutor.StreamChunk{Payload: []byte(`{"type":"response.completed"}`)} + chunks := make(chan cliproxyexecutor.StreamChunk, 2) + chunks <- cliproxyexecutor.StreamChunk{Payload: []byte(`{"type":"response.output_text.delta","delta":"ok"}`)} + chunks <- cliproxyexecutor.StreamChunk{Payload: []byte(`{"type":"response.completed","response":{"status":"completed"}}`)} close(chunks) return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil } @@ -640,8 +642,9 @@ func (e *retryingHomeStreamExecutor) ExecuteStream(_ context.Context, auth *Auth if e.calls.Add(1) == 1 { return nil, &Error{HTTPStatus: http.StatusUnauthorized, Message: "expired"} } - chunks := make(chan cliproxyexecutor.StreamChunk, 1) - chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("data: {\"type\":\"response.completed\"}\n\n")} + chunks := make(chan cliproxyexecutor.StreamChunk, 2) + chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("data: {\"type\":\"response.output_text.delta\",\"delta\":\"ok\"}\n\n")} + chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("data: {\"type\":\"response.completed\",\"response\":{\"status\":\"completed\"}}\n\n")} close(chunks) return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil } @@ -684,6 +687,123 @@ func TestHomeStreamRetryUsesFreshSelection(t *testing.T) { } } +type alwaysEmptyHomeStreamExecutor struct { + calls atomic.Int32 +} + +func (*alwaysEmptyHomeStreamExecutor) Identifier() string { return "home-execution" } +func (*alwaysEmptyHomeStreamExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} +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: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":{\"completion_tokens\":0}}\n\n")} + close(chunks) + return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil +} +func (*alwaysEmptyHomeStreamExecutor) Refresh(context.Context, *Auth) (*Auth, error) { + return nil, nil +} +func (*alwaysEmptyHomeStreamExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} +func (*alwaysEmptyHomeStreamExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +func TestHomeStreamRepeatedEmptyCompletionStopsAtRepeatedAuth(t *testing.T) { + dispatcher := &retainingHomeExecutionDispatcher{} + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + executor := &alwaysEmptyHomeStreamExecutor{} + manager.RegisterExecutor(executor) + + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + result, errExecute := manager.ExecuteStream(ctx, []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{Stream: true}) + if errExecute != nil { + t.Fatalf("ExecuteStream() error = %v", errExecute) + } + var terminalErr error + for chunk := range result.Chunks { + if chunk.Err != nil { + terminalErr = chunk.Err + } + } + if !isEmptyCompletionError(terminalErr) { + t.Fatalf("terminal error = %v, want empty_completion", terminalErr) + } + if got := executor.calls.Load(); got != 1 { + t.Fatalf("executor calls = %d, want 1 before repeated auth is rejected", got) + } + if got := dispatcher.calls.Load(); got != 2 { + t.Fatalf("Home RPOP calls = %d, want 2 to detect the repeated auth", got) + } +} + +type alternatingEmptyHomeDispatcher struct { + calls atomic.Int32 +} + +func (*alternatingEmptyHomeDispatcher) HeartbeatOK() bool { return true } +func (d *alternatingEmptyHomeDispatcher) RPopAuth(context.Context, string, string, http.Header, int) ([]byte, error) { + call := d.calls.Add(1) + authID := "home-auth-a" + if call == 2 { + authID = "home-auth-b" + } + return json.Marshal(homeAuthDispatchResponse{Auth: Auth{ID: authID, Provider: "home-execution", Status: StatusActive}}) +} +func (*alternatingEmptyHomeDispatcher) AbortAmbiguousDispatch() {} + +type alternatingEmptyHomeExecutor struct { + calls atomic.Int32 +} + +func (*alternatingEmptyHomeExecutor) Identifier() string { return "home-execution" } +func (*alternatingEmptyHomeExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} +func (e *alternatingEmptyHomeExecutor) ExecuteStream(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + e.calls.Add(1) + if auth.ID == "home-auth-b" { + return nil, &Error{Code: "transient", Message: "transient stream failure", Retryable: true} + } + chunks := make(chan cliproxyexecutor.StreamChunk, 1) + 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 +} +func (*alternatingEmptyHomeExecutor) Refresh(context.Context, *Auth) (*Auth, error) { return nil, nil } +func (*alternatingEmptyHomeExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} +func (*alternatingEmptyHomeExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +func TestHomeStreamDoesNotRevisitEmptyAuthAfterAnotherFailure(t *testing.T) { + dispatcher := &alternatingEmptyHomeDispatcher{} + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + executor := &alternatingEmptyHomeExecutor{} + manager.RegisterExecutor(executor) + + _, errExecute := manager.executeStreamMixedOnce(context.Background(), []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{Stream: true}, 0, nil, 0, -1, nil) + if errExecute == nil || !strings.Contains(errExecute.Error(), "transient stream failure") { + t.Fatalf("executeStreamMixedOnce() error = %v, want last transient failure", errExecute) + } + if got := executor.calls.Load(); got != 2 { + t.Fatalf("executor calls = %d, want one call per distinct auth", got) + } + if got := dispatcher.calls.Load(); got != 3 { + t.Fatalf("Home RPOP calls = %d, want third dispatch rejected before execution", got) + } +} + type cancellationBarrierExecutor struct { executeCalls atomic.Int32 countCalls atomic.Int32 diff --git a/sdk/cliproxy/auth/home_retry_contract_test.go b/sdk/cliproxy/auth/home_retry_contract_test.go index 5f62c652cf3..627f6362acd 100644 --- a/sdk/cliproxy/auth/home_retry_contract_test.go +++ b/sdk/cliproxy/auth/home_retry_contract_test.go @@ -718,7 +718,7 @@ func TestHomeRetryRoundUsesEarliestCredentialRetryAfter(t *testing.T) { manager.RegisterExecutor(executor) retryLimit := -1 - _, errExecute := manager.executeHomeOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}, false, 2, &retryLimit) + _, errExecute := manager.executeHomeOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}, false, 2, &retryLimit, 0) if !isHomeRetryRoundExhausted(errExecute) { t.Fatalf("executeHomeOnce() error = %v, want exhausted retry round", errExecute) } @@ -772,14 +772,14 @@ func TestHomeRetryRoundUsesAuthoritativeRemoteCooldown(t *testing.T) { { name: "nonstream", execute: func(manager *Manager, retryLimit *int) error { - _, errExecute := manager.executeHomeOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}, false, 2, retryLimit) + _, errExecute := manager.executeHomeOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}, false, 2, retryLimit, 0) return errExecute }, }, { name: "stream", execute: func(manager *Manager, retryLimit *int) error { - _, errExecute := manager.executeStreamMixedOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}, 2, retryLimit, 0, 0) + _, errExecute := manager.executeStreamMixedOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}, 2, retryLimit, 0, 0, nil) return errExecute }, }, @@ -824,14 +824,14 @@ func TestHomeCooldownClassificationPreservesNonRetryableRoundStatus(t *testing.T { name: "nonstream", execute: func(manager *Manager, retryLimit *int) error { - _, errExecute := manager.executeHomeOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}, false, 2, retryLimit) + _, errExecute := manager.executeHomeOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}, false, 2, retryLimit, 0) return errExecute }, }, { name: "stream", execute: func(manager *Manager, retryLimit *int) error { - _, errExecute := manager.executeStreamMixedOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}, 2, retryLimit, 0, 0) + _, errExecute := manager.executeStreamMixedOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}, 2, retryLimit, 0, 0, nil) return errExecute }, }, @@ -875,14 +875,14 @@ func TestHomeRetryRoundStartsImmediatelyWhenHomeReportsAvailableNextRound(t *tes { name: "nonstream", execute: func(manager *Manager, retryLimit *int) error { - _, errExecute := manager.executeHomeOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}, false, 0, retryLimit) + _, errExecute := manager.executeHomeOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}, false, 0, retryLimit, 0) return errExecute }, }, { name: "stream", execute: func(manager *Manager, retryLimit *int) error { - _, errExecute := manager.executeStreamMixedOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}, 0, retryLimit, 0, 0) + _, errExecute := manager.executeStreamMixedOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}, 0, retryLimit, 0, 0, nil) return errExecute }, }, @@ -927,14 +927,14 @@ func TestHomeRetryRoundUsesRemoteCooldownWhenAttemptedErrorHasNoTiming(t *testin { name: "nonstream", execute: func(manager *Manager, retryLimit *int) error { - _, errExecute := manager.executeHomeOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}, false, 2, retryLimit) + _, errExecute := manager.executeHomeOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}, false, 2, retryLimit, 0) return errExecute }, }, { name: "stream", execute: func(manager *Manager, retryLimit *int) error { - _, errExecute := manager.executeStreamMixedOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}, 2, retryLimit, 0, 0) + _, errExecute := manager.executeStreamMixedOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}, 2, retryLimit, 0, 0, nil) return errExecute }, }, @@ -1259,7 +1259,7 @@ func TestHomeLocalSelectionRejectionWaitsForReleaseAcknowledgement(t *testing.T) blockedGroup: executionregistry.ReleaseGroup{CredentialID: "home-retry-a", Model: "gpt"}, blockedSequence: 2, execute: func(manager *Manager, maxRetryCredentials int, retryLimit *int) error { - _, errExecute := manager.executeHomeOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}, false, maxRetryCredentials, retryLimit) + _, errExecute := manager.executeHomeOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}, false, maxRetryCredentials, retryLimit, 0) return errExecute }, }, @@ -1274,7 +1274,7 @@ func TestHomeLocalSelectionRejectionWaitsForReleaseAcknowledgement(t *testing.T) blockedGroup: executionregistry.ReleaseGroup{CredentialID: "home-retry-a", Model: "gpt"}, blockedSequence: 2, execute: func(manager *Manager, maxRetryCredentials int, retryLimit *int) error { - _, errExecute := manager.executeStreamMixedOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}, maxRetryCredentials, retryLimit, 0, 0) + _, errExecute := manager.executeStreamMixedOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}, maxRetryCredentials, retryLimit, 0, 0, nil) return errExecute }, }, @@ -1289,7 +1289,7 @@ func TestHomeLocalSelectionRejectionWaitsForReleaseAcknowledgement(t *testing.T) blockedGroup: executionregistry.ReleaseGroup{CredentialID: "home-retry-b", Model: "gpt"}, blockedSequence: 1, execute: func(manager *Manager, maxRetryCredentials int, retryLimit *int) error { - _, errExecute := manager.executeStreamMixedOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}, maxRetryCredentials, retryLimit, 0, 0) + _, errExecute := manager.executeStreamMixedOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}, maxRetryCredentials, retryLimit, 0, 0, nil) return errExecute }, }, diff --git a/sdk/cliproxy/auth/home_selected_auth_callback_test.go b/sdk/cliproxy/auth/home_selected_auth_callback_test.go index c596071c8d5..deb5b5c5b9e 100644 --- a/sdk/cliproxy/auth/home_selected_auth_callback_test.go +++ b/sdk/cliproxy/auth/home_selected_auth_callback_test.go @@ -42,8 +42,9 @@ func (e *callbackPinHomeExecutor) ExecuteStream(_ context.Context, _ *Auth, _ cl if lifecycle, ok := opts.ExecutionLifecycle.(interface{ Retain() }); ok { lifecycle.Retain() } - chunks := make(chan cliproxyexecutor.StreamChunk, 1) - chunks <- cliproxyexecutor.StreamChunk{Payload: []byte(`{"type":"response.completed"}`)} + chunks := make(chan cliproxyexecutor.StreamChunk, 2) + chunks <- cliproxyexecutor.StreamChunk{Payload: []byte(`{"type":"response.output_text.delta","delta":"ok"}`)} + chunks <- cliproxyexecutor.StreamChunk{Payload: []byte(`{"type":"response.completed","response":{"status":"completed"}}`)} close(chunks) return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil } 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 00000000000..c57b46eff55 --- /dev/null +++ b/sdk/cliproxy/auth/outer_retry_exclusions_test.go @@ -0,0 +1,237 @@ +package auth + +import ( + "context" + "errors" + "net/http" + "sync" + "testing" + + "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) + exec.executeErrs[ids[0]] = &Error{Code: "service_unavailable", Message: "502 Bad Gateway"} + + 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) + exec.streamErrs[ids[0]] = &Error{Code: "service_unavailable", Message: "503 Service Unavailable"} + + 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/route_exhaustion_test.go b/sdk/cliproxy/auth/route_exhaustion_test.go new file mode 100644 index 00000000000..038dab71419 --- /dev/null +++ b/sdk/cliproxy/auth/route_exhaustion_test.go @@ -0,0 +1,779 @@ +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") + } + }) +} + +// TestRouteExhaustion_HomeNoModelDiagnostic proves a home auth with no executable +// models is recorded in the route summary only as a sanitized provider/status: +// no auth ID, credentials, or model leaks. +func TestRouteExhaustion_HomeNoModelDiagnostic(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") + mgr.RegisterExecutor(execClaude) + + authID := "home-no-model-secret-auth" + key := "sk-home-no-model-secret-key" + file := "user@no-model-secret.com.json" + url := "https://no-model.secret.net/v1" + alias := "secret-no-model-alias" + model := "no-model-route" + + payload := fmt.Sprintf(`{"provider":"claude","auth":{"id":%q,"provider":"claude","status":"active","unavailable":true,"attributes":{"api_key":%q,"credential_file":%q,"base_url":%q,"disable_cooling":"true"},"metadata":{"private_model":%q}}}`, authID, key, file, url, alias) + + dispatcher := &routeExhaustionHomeDispatcher{ + responses: map[int]string{0: payload}, + } + mgr.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + + // Empty route model leaves no executable upstream model for this auth, so + // executeHome hits the no_execution_models path and records the attempt. + _, err := mgr.Execute(context.Background(), []string{"claude"}, cliproxyexecutor.Request{}, cliproxyexecutor.Options{}) + if err == nil { + t.Fatalf("Execute() expected error on home no-model route, got nil") + } + + errStr := err.Error() + summaryIdx := strings.Index(errStr, "attempted routes:") + if summaryIdx == -1 { + t.Fatalf("summary missing in home no-model route exhaustion error: %s", errStr) + } + summaryPart := errStr[summaryIdx:] + + if !strings.Contains(summaryPart, "claude:error") { + t.Errorf("expected sanitized claude:error in home no-model summary, got: %s", summaryPart) + } + + leaks := []string{authID, key, file, url, alias, model} + for _, token := range leaks { + if strings.Contains(summaryPart, token) { + t.Errorf("sensitive token %q leaked in home no-model summary: %s", token, 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 } + +// I. 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("routeExhaustionClonedError 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")) + } +} + +// J. 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("routeExhaustionClonedError must implement Headers(), err=%T", err) + } + if got := he.Headers(); got != nil { + t.Errorf("Headers() = %v, want nil for cause without headers", got) + } +} + +// K. 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 } + +// L. 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) + } +} + +// M. 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("routeExhaustionClonedError 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")) + } +} + +// N. Approach B: a wrapped Home busy cause keeps its typed identity, Retry-After +// stays discoverable through the wrapper, SafeResponseHeaders surfaces the +// trusted header, and the route summary is still appended. +func TestRouteExhaustion_HomeBusyRetryAfterAndTypedCause(t *testing.T) { + tracker := newRouteAttemptTracker() + tracker.Record(&Auth{Provider: "gemini"}, &Error{HTTPStatus: 429}) + + cause := NewHomeConcurrencyBusyError("credential busy", 750*time.Millisecond) + err := wrapRouteExhaustion(cause, tracker) + + var busy *HomeConcurrencyBusyError + if !errors.As(err, &busy) || busy == nil { + t.Fatalf("errors.As(*HomeConcurrencyBusyError) failed, err=%v", err) + } + if got := retryAfterFromError(err); got == nil || *got != 750*time.Millisecond { + t.Fatalf("retryAfterFromError(err) = %v, want 750ms", got) + } + if hdr := SafeResponseHeaders(err); hdr == nil || hdr.Get("Retry-After") != "1" { + t.Fatalf("SafeResponseHeaders(err) = %v, want Retry-After 1", hdr) + } + if !strings.Contains(err.Error(), "attempted routes: [gemini") { + t.Errorf("route summary missing from wrapped home busy error: %s", err.Error()) + } +} + +// O. Approach B: wrapping a concrete *Error cause must not clone or mutate it; +// the cause's Message and Error() text stay byte-identical. +func TestRouteExhaustion_MessageNotMutated(t *testing.T) { + tracker := newRouteAttemptTracker() + tracker.Record(&Auth{Provider: "openai"}, &Error{HTTPStatus: 502}) + + cause := &Error{Code: "bad_gateway", Message: "502 Bad Gateway", HTTPStatus: 502} + err := wrapRouteExhaustion(cause, tracker) + + var authErr *Error + if !errors.As(err, &authErr) || authErr == nil { + t.Fatalf("errors.As(*Error) failed, err=%v", err) + } + if authErr.Message != "502 Bad Gateway" { + t.Fatalf("cause Message mutated: got %q, want %q", authErr.Message, "502 Bad Gateway") + } + if authErr.Error() != "bad_gateway: 502 Bad Gateway" { + t.Fatalf("cause Error() altered: got %q", authErr.Error()) + } + if !strings.HasSuffix(err.Error(), "; attempted routes: [openai:502]") { + t.Fatalf("wrapper should append summary to unmutated cause text: %s", err.Error()) + } +} + +// P. Approach B: the sanitized route summary appears exactly once in the +// wrapper error, even when the cause itself is a *Error. +func TestRouteExhaustion_SummaryAppendedExactlyOnce(t *testing.T) { + tracker := newRouteAttemptTracker() + tracker.Record(&Auth{Provider: "openai"}, &Error{HTTPStatus: 502}) + + cause := &Error{Code: "bad_gateway", Message: "502 Bad Gateway", HTTPStatus: 502} + err := wrapRouteExhaustion(cause, tracker) + + errStr := err.Error() + if count := strings.Count(errStr, "attempted routes:"); count != 1 { + t.Fatalf("summary appeared %d times in error, want exactly once: %s", count, errStr) + } + if count := strings.Count(errStr, "; attempted routes:"); count != 1 { + t.Fatalf("summary separator `; attempted routes:` appeared %d times, want exactly once: %s", count, errStr) + } +} + +// Q. Approach B: SafeResponseHeaders is nil-safe on a nil wrapper receiver. +func TestRouteExhaustion_SafeResponseHeadersNilReceiver(t *testing.T) { + var e *routeExhaustionClonedError + if hdr := e.SafeResponseHeaders(); hdr != nil { + t.Errorf("SafeResponseHeaders() = %v, want nil for nil receiver", hdr) + } + if hdr := SafeResponseHeaders(nil); hdr != nil { + t.Errorf("SafeResponseHeaders(nil) = %v, want nil", hdr) + } +} + +// R. Approach B: for a generic (non-*Error) cause chain, Headers forwards the +// first/outermost carrier per errors.As and returns a fresh clone, never the +// cause's own map reference. +func TestRouteExhaustion_GenericOutermostHeadersClone(t *testing.T) { + tracker := newRouteAttemptTracker() + tracker.Record(&Auth{Provider: "gemini"}, &Error{HTTPStatus: 429}) + + inner := &routeExhaustionHeaderCause{ + msg: "inner", + headers: http.Header{"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("routeExhaustionClonedError must implement Headers(), err=%T", err) + } + hdr := he.Headers() + if hdr.Get("Retry-After") != "outer" || hdr.Get("Inner") != "" { + t.Fatalf("outermost carrier did not win over inner, got %v", hdr) + } + hdr.Set("Retry-After", "999") + if outer.headers.Get("Retry-After") != "outer" { + t.Fatalf("wrapped headers mutated caller map: got %q", outer.headers.Get("Retry-After")) + } + if inner.headers.Get("Inner") != "1" { + t.Fatalf("inner caller map mutated: got %q", inner.headers.Get("Inner")) + } +} + +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 00000000000..1f37efd3e7c --- /dev/null +++ b/sdk/cliproxy/auth/route_tracker.go @@ -0,0 +1,210 @@ +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" +} + +// routeExhaustionClonedError wraps the original route-exhaustion cause without +// cloning or mutating it, so typed access (errors.As(*Error), status, retry, +// HomeConcurrencyBusyError) and the sanitized route summary both survive the +// wrapper. The summary is appended in Error() without mutating the underlying +// cause's fields. +type routeExhaustionClonedError struct { + cause error + summary string +} + +func wrapRouteExhaustion(cause error, tracker *routeAttemptTracker) error { + if cause == nil { + return nil + } + if tracker == nil || isRequestInvalidError(cause) || isRequestTerminatedError(cause) { + 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 +} + +func (e *routeExhaustionClonedError) StatusCode() int { + if e == nil || e.cause == nil { + return 0 + } + return statusCodeFromError(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 the +// first/outermost carrier per errors.As and a fresh copy of its map, never +// mutating the caller's or cause'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 the trusted Home busy response headers from the +// underlying cause, nil-safe for the wrapper receiver. +func (e *routeExhaustionClonedError) SafeResponseHeaders() http.Header { + if e == nil { + return nil + } + return SafeResponseHeaders(e.cause) +} diff --git a/sdk/cliproxy/auth/selected_auth_failover_metadata_test.go b/sdk/cliproxy/auth/selected_auth_failover_metadata_test.go new file mode 100644 index 00000000000..e8b4d6fa602 --- /dev/null +++ b/sdk/cliproxy/auth/selected_auth_failover_metadata_test.go @@ -0,0 +1,275 @@ +package auth + +import ( + "context" + "net/http" + "sync" + "testing" + + registry "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +const failoverMetadataModel = "failover-metadata-model" + +type failoverMetadataCall struct { + authID string + authIndex string + metaID string + metaIndex string +} + +// failoverMetadataExecutor fails the first attempt so the conductor rotates to +// the next credential, and records for every attempt which auth it received +// together with the selected-auth metadata carried by the execution options. +type failoverMetadataExecutor struct { + id string + + mu sync.Mutex + calls []failoverMetadataCall +} + +func (e *failoverMetadataExecutor) Identifier() string { return e.id } + +func (e *failoverMetadataExecutor) record(auth *Auth, opts cliproxyexecutor.Options) error { + e.mu.Lock() + defer e.mu.Unlock() + metaID, _ := opts.Metadata[cliproxyexecutor.SelectedAuthMetadataKey].(string) + metaIndex, _ := opts.Metadata[cliproxyexecutor.SelectedAuthIndexMetadataKey].(string) + e.calls = append(e.calls, failoverMetadataCall{ + authID: auth.ID, + authIndex: auth.EnsureIndex(), + metaID: metaID, + metaIndex: metaIndex, + }) + if len(e.calls) == 1 { + return &Error{HTTPStatus: http.StatusInternalServerError, Message: "boom"} + } + return nil +} + +func (e *failoverMetadataExecutor) Calls() []failoverMetadataCall { + e.mu.Lock() + defer e.mu.Unlock() + out := make([]failoverMetadataCall, len(e.calls)) + copy(out, e.calls) + return out +} + +func (e *failoverMetadataExecutor) Execute(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + if err := e.record(auth, opts); err != nil { + return cliproxyexecutor.Response{}, err + } + return cliproxyexecutor.Response{Payload: []byte(auth.ID)}, nil +} + +func (e *failoverMetadataExecutor) ExecuteStream(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + if err := e.record(auth, opts); err != nil { + return nil, err + } + ch := make(chan cliproxyexecutor.StreamChunk, 1) + ch <- cliproxyexecutor.StreamChunk{Payload: []byte(auth.ID)} + close(ch) + return &cliproxyexecutor.StreamResult{Chunks: ch}, nil +} + +func (e *failoverMetadataExecutor) Refresh(_ context.Context, auth *Auth) (*Auth, error) { + return auth, nil +} + +func (e *failoverMetadataExecutor) CountTokens(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + if err := e.record(auth, opts); err != nil { + return cliproxyexecutor.Response{}, err + } + return cliproxyexecutor.Response{Payload: []byte(auth.ID)}, nil +} + +func (e *failoverMetadataExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +// failoverMetadataProbe captures what the after-auth interceptor observed and +// how often the selected-auth callbacks fired. +type failoverMetadataProbe struct { + mu sync.Mutex + interceptedIDs []string + interceptedIndex []string + callbackIDs []string + callbackIndexes []string +} + +func (p *failoverMetadataProbe) observeIntercept(meta map[string]any) { + id, _ := meta[cliproxyexecutor.SelectedAuthMetadataKey].(string) + index, _ := meta[cliproxyexecutor.SelectedAuthIndexMetadataKey].(string) + p.mu.Lock() + defer p.mu.Unlock() + p.interceptedIDs = append(p.interceptedIDs, id) + p.interceptedIndex = append(p.interceptedIndex, index) +} + +func (p *failoverMetadataProbe) snapshot() ([]string, []string, []string, []string) { + p.mu.Lock() + defer p.mu.Unlock() + return append([]string(nil), p.interceptedIDs...), + append([]string(nil), p.interceptedIndex...), + append([]string(nil), p.callbackIDs...), + append([]string(nil), p.callbackIndexes...) +} + +func newFailoverMetadataOptions() (cliproxyexecutor.Options, *failoverMetadataProbe) { + probe := &failoverMetadataProbe{} + opts := cliproxyexecutor.Options{ + Metadata: map[string]any{ + cliproxyexecutor.SelectedAuthCallbackMetadataKey: func(authID string) { + probe.mu.Lock() + probe.callbackIDs = append(probe.callbackIDs, authID) + probe.mu.Unlock() + }, + cliproxyexecutor.SelectedAuthIndexCallbackMetadataKey: func(authIndex string) { + probe.mu.Lock() + probe.callbackIndexes = append(probe.callbackIndexes, authIndex) + probe.mu.Unlock() + }, + }, + RequestAfterAuthInterceptor: func(_ context.Context, req cliproxyexecutor.RequestAfterAuthInterceptRequest) cliproxyexecutor.RequestAfterAuthInterceptResponse { + probe.observeIntercept(req.Metadata) + return cliproxyexecutor.RequestAfterAuthInterceptResponse{} + }, + } + return opts, probe +} + +func newFailoverMetadataManager(t *testing.T, prefix string) (*Manager, *failoverMetadataExecutor) { + t.Helper() + + m := NewManager(nil, nil, nil) + m.SetRetryConfig(0, 0, 0) + executor := &failoverMetadataExecutor{id: "claude"} + m.RegisterExecutor(executor) + + ids := []string{prefix + "-auth-1", prefix + "-auth-2"} + reg := registry.GetGlobalRegistry() + for _, id := range ids { + reg.RegisterClient(id, "claude", []*registry.ModelInfo{{ID: failoverMetadataModel}}) + } + t.Cleanup(func() { + for _, id := range ids { + reg.UnregisterClient(id) + } + }) + for _, id := range ids { + auth := &Auth{ID: id, Provider: "claude", FileName: id + ".json", Status: StatusActive} + if _, errRegister := m.Register(context.Background(), auth); errRegister != nil { + t.Fatalf("register %s: %v", id, errRegister) + } + } + return m, executor +} + +func assertFailoverMetadata(t *testing.T, executor *failoverMetadataExecutor, probe *failoverMetadataProbe) { + t.Helper() + + calls := executor.Calls() + if len(calls) != 2 { + t.Fatalf("executor attempts = %d, want 2 (failover to the second credential)", len(calls)) + } + if calls[0].authID == calls[1].authID { + t.Fatalf("failover reused auth %q for both attempts", calls[0].authID) + } + for i, call := range calls { + if call.metaID != call.authID { + t.Fatalf("attempt %d: execution metadata %s = %q, want %q (stale selected auth reaches after-auth plugins)", + i, cliproxyexecutor.SelectedAuthMetadataKey, call.metaID, call.authID) + } + if call.authIndex == "" { + t.Fatalf("attempt %d: auth index is empty, cannot verify %s", i, cliproxyexecutor.SelectedAuthIndexMetadataKey) + } + if call.metaIndex != call.authIndex { + t.Fatalf("attempt %d: execution metadata %s = %q, want %q", + i, cliproxyexecutor.SelectedAuthIndexMetadataKey, call.metaIndex, call.authIndex) + } + } + + interceptedIDs, interceptedIndexes, callbackIDs, callbackIndexes := probe.snapshot() + if len(interceptedIDs) != len(calls) { + t.Fatalf("after-auth interceptor invocations = %d, want %d", len(interceptedIDs), len(calls)) + } + for i, call := range calls { + if interceptedIDs[i] != call.authID { + t.Fatalf("attempt %d: after-auth interceptor saw %s = %q, want %q", + i, cliproxyexecutor.SelectedAuthMetadataKey, interceptedIDs[i], call.authID) + } + if interceptedIndexes[i] != call.authIndex { + t.Fatalf("attempt %d: after-auth interceptor saw %s = %q, want %q", + i, cliproxyexecutor.SelectedAuthIndexMetadataKey, interceptedIndexes[i], call.authIndex) + } + } + + if len(callbackIDs) != len(calls) { + t.Fatalf("selected auth callback fired %d times, want %d (exactly once per attempt)", len(callbackIDs), len(calls)) + } + if len(callbackIndexes) != len(calls) { + t.Fatalf("selected auth index callback fired %d times, want %d (exactly once per attempt)", len(callbackIndexes), len(calls)) + } + for i, call := range calls { + if callbackIDs[i] != call.authID { + t.Fatalf("attempt %d: selected auth callback got %q, want %q", i, callbackIDs[i], call.authID) + } + if callbackIndexes[i] != call.authIndex { + t.Fatalf("attempt %d: selected auth index callback got %q, want %q", i, callbackIndexes[i], call.authIndex) + } + } +} + +// TestSelectedAuthMetadataFollowsFailover proves that after the first credential +// fails, the cloned per-attempt options handed to the after-auth interceptor and +// to the executor carry the selected auth of the current attempt, not the one of +// the previous attempt, and that the selected-auth callbacks fire exactly once +// per attempt. +func TestSelectedAuthMetadataFollowsFailover(t *testing.T) { + req := cliproxyexecutor.Request{Model: failoverMetadataModel} + testCases := []struct { + name string + invoke func(*testing.T, *Manager, cliproxyexecutor.Options) + }{ + { + name: "execute", + invoke: func(t *testing.T, m *Manager, opts cliproxyexecutor.Options) { + if _, errExecute := m.Execute(context.Background(), []string{"claude"}, req, opts); errExecute != nil { + t.Fatalf("Execute() error = %v", errExecute) + } + }, + }, + { + name: "execute_count", + invoke: func(t *testing.T, m *Manager, opts cliproxyexecutor.Options) { + if _, errExecute := m.ExecuteCount(context.Background(), []string{"claude"}, req, opts); errExecute != nil { + t.Fatalf("ExecuteCount() error = %v", errExecute) + } + }, + }, + { + name: "execute_stream", + invoke: func(t *testing.T, m *Manager, opts cliproxyexecutor.Options) { + result, errExecute := m.ExecuteStream(context.Background(), []string{"claude"}, req, opts) + if errExecute != nil { + t.Fatalf("ExecuteStream() error = %v", errExecute) + } + if result != nil { + for range result.Chunks { + } + } + }, + }, + } + + for _, tc := range testCases { + tc := tc + t.Run(tc.name, func(t *testing.T) { + manager, executor := newFailoverMetadataManager(t, "failover-metadata-"+tc.name) + opts, probe := newFailoverMetadataOptions() + tc.invoke(t, manager, opts) + assertFailoverMetadata(t, executor, probe) + }) + } +} diff --git a/sdk/cliproxy/auth/selected_auth_metadata_test.go b/sdk/cliproxy/auth/selected_auth_metadata_test.go index 2a7433e447d..b296a0c3c4c 100644 --- a/sdk/cliproxy/auth/selected_auth_metadata_test.go +++ b/sdk/cliproxy/auth/selected_auth_metadata_test.go @@ -1,8 +1,12 @@ package auth import ( + "context" + "net/http" + "sync" "testing" + registry "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" ) @@ -38,3 +42,210 @@ 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) + }) +} + +// affinityCaptureExecutor records the execution options and call count per auth so +// tests can prove that session-affinity metadata survives the whole mixed pipeline. +type affinityCaptureExecutor struct { + mu sync.Mutex + provider string + lastOpts cliproxyexecutor.Options + calls map[string]int +} + +func (e *affinityCaptureExecutor) Identifier() string { return e.provider } +func (e *affinityCaptureExecutor) Execute(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + e.mu.Lock() + defer e.mu.Unlock() + e.calls[auth.ID]++ + e.lastOpts = opts + return cliproxyexecutor.Response{Payload: []byte(`{"choices":[{"message":{"content":"ok"}}]}`)}, nil +} +func (e *affinityCaptureExecutor) ExecuteStream(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + return nil, nil +} +func (e *affinityCaptureExecutor) Refresh(ctx context.Context, auth *Auth) (*Auth, error) { + return auth, nil +} +func (e *affinityCaptureExecutor) CountTokens(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} +func (e *affinityCaptureExecutor) HttpRequest(ctx context.Context, auth *Auth, req *http.Request) (*http.Response, error) { + return nil, nil +} + +func (e *affinityCaptureExecutor) captured() (cliproxyexecutor.Options, map[string]int) { + e.mu.Lock() + defer e.mu.Unlock() + calls := make(map[string]int, len(e.calls)) + for id, n := range e.calls { + calls[id] = n + } + return e.lastOpts, calls +} + +// TestMixedAffinityNamespaceSurvivesExecution proves that after a successful mixed +// execution with a caller-supplied exclusion (nonempty tried set), the session +// affinity selector records Result.Options under the "mixed" namespace so a +// subsequent same-session request reuses the same auth instead of re-picking. +// Regression: execOpts must derive from pickOpts (stamped namespace), not opts. +func TestMixedAffinityNamespaceSurvivesExecution(t *testing.T) { + ctx := context.Background() + manager := NewManager(nil, nil, nil) + execClaude := &affinityCaptureExecutor{provider: "claude", calls: make(map[string]int)} + execOpenAI := &affinityCaptureExecutor{provider: "openai", calls: make(map[string]int)} + manager.RegisterExecutor(execClaude) + manager.RegisterExecutor(execOpenAI) + + model := "affinity-mixed-model" + reg := registry.GetGlobalRegistry() + reg.RegisterClient("affinity-claude-1", "claude", []*registry.ModelInfo{{ID: model}}) + reg.RegisterClient("affinity-openai-1", "openai", []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { + reg.UnregisterClient("affinity-claude-1") + reg.UnregisterClient("affinity-openai-1") + }) + if _, err := manager.Register(ctx, &Auth{ID: "affinity-claude-1", Provider: "claude", FileName: "affinity-claude-1.json", Status: StatusActive}); err != nil { + t.Fatalf("Register(claude) error = %v", err) + } + if _, err := manager.Register(ctx, &Auth{ID: "affinity-openai-1", Provider: "openai", FileName: "affinity-openai-1.json", Status: StatusActive}); err != nil { + t.Fatalf("Register(openai) error = %v", err) + } + + affinity := NewSessionAffinitySelector(&WeightedRoundRobinSelector{}) + manager.SetSelector(affinity) + + req := cliproxyexecutor.Request{Model: model} + headers := http.Header{} + headers.Set("X-Claude-Code-Session-Id", "affinity-session-1") + opts := cliproxyexecutor.Options{ + Headers: headers, + Metadata: map[string]any{ + cliproxyexecutor.ExcludedAuthIDsMetadataKey: map[string]struct{}{ + "affinity-claude-1": {}, + }, + }, + } + + // First request: claude is caller-excluded, openai must win. + if _, err := manager.Execute(ctx, []string{"claude", "openai"}, req, opts); err != nil { + t.Fatalf("first Execute() error = %v", err) + } + + lastOpts, calls := execOpenAI.captured() + if got := calls["affinity-openai-1"]; got != 1 { + t.Fatalf("openai executed %d times on first request, want 1", got) + } + if got := calls["affinity-claude-1"]; got != 0 { + t.Fatalf("caller-excluded claude executed %d times, want 0", got) + } + providerMeta, ok := lastOpts.Metadata[cliproxyexecutor.SessionAffinityProviderMetadataKey].(string) + if !ok || providerMeta != "mixed" { + t.Fatalf("SessionAffinityProviderMetadataKey = %q, %v; want \"mixed\", true (Result.Options must keep pickOpts namespace)", providerMeta, ok) + } + modelMeta, ok := lastOpts.Metadata[cliproxyexecutor.SessionAffinityModelMetadataKey].(string) + if !ok || modelMeta != model { + t.Fatalf("SessionAffinityModelMetadataKey = %q, %v; want %q, true", modelMeta, ok, model) + } + + // Second same-session request without exclusions must hit the mixed binding. + if _, err := manager.Execute(ctx, []string{"claude", "openai"}, req, cliproxyexecutor.Options{Headers: headers}); err != nil { + t.Fatalf("second Execute() error = %v", err) + } + _, callsAfter := execOpenAI.captured() + if got := callsAfter["affinity-openai-1"]; got != 2 { + t.Fatalf("openai executed %d times across both same-session requests, want 2 (mixed cache reuse)", got) + } +} diff --git a/sdk/cliproxy/auth/selector.go b/sdk/cliproxy/auth/selector.go index 7a053b78a9a..cd41a754ad7 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 } @@ -610,8 +627,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. @@ -636,9 +655,14 @@ func NewSessionAffinitySelectorWithConfig(cfg SessionAffinityConfig) *SessionAff if cfg.TTL <= 0 { cfg.TTL = time.Hour } + quarantineTTL := 5 * time.Second + if cfg.TTL > 0 && cfg.TTL < quarantineTTL { + quarantineTTL = cfg.TTL + } return &SessionAffinitySelector{ - fallback: cfg.Fallback, - cache: NewSessionCache(cfg.TTL), + fallback: cfg.Fallback, + cache: NewSessionCache(cfg.TTL), + quarantine: NewSessionCache(quarantineTTL), } } @@ -663,12 +687,13 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri 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 } @@ -678,11 +703,10 @@ 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) modelKey := canonicalModelKey(model) cacheKey := provider + "::" + primaryID + "::" + modelKey @@ -690,6 +714,8 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri if fallbackID != "" && fallbackID != primaryID { fallbackKey = provider + "::" + fallbackID + "::" + modelKey } + available = s.excludeSessionQuarantine(cacheKey, fallbackKey, available) + fallbackAuths := highestPriorityAuths(available) bind := func(authID string) { if fallbackKey != "" { s.cache.SetAliases(authID, cacheKey, fallbackKey) @@ -697,34 +723,54 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri } s.cache.Set(cacheKey, authID) } - - if cachedAuthID, ok := s.cache.GetAndRefresh(cacheKey); 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 + pickCached := func() *Auth { + if cachedAuthID, ok := s.cache.GetAndRefresh(cacheKey); ok { + for _, auth := range available { + if auth.ID == cachedAuthID { + bind(auth.ID) + return auth + } } } - // 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 + if fallbackKey != "" { + if cachedAuthID, ok := s.cache.Get(fallbackKey); ok { + for _, auth := range available { + if auth.ID == cachedAuthID { + bind(auth.ID) + return auth + } + } + } } - 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 nil + } + + if auth := pickCached(); auth != nil { + entry.Infof("session-affinity: cache hit | session=%s auth=%s provider=%s model=%s", truncateSessionID(primaryID), auth.ID, provider, model) + return auth, nil + } + + s.bindMu.Lock() + defer s.bindMu.Unlock() + if auth := pickCached(); auth != nil { + entry.Infof("session-affinity: concurrent cache hit | session=%s auth=%s provider=%s model=%s", truncateSessionID(primaryID), auth.ID, provider, model) return auth, nil } + // Under bindMu, authoritatively observe the stale group state using non-refreshing token read. + genPrimary, authPrimary, aliasesPrimary, okPrimary := s.cache.Observe(cacheKey) + var genFallback uint64 + var authFallback string + var aliasesFallback []string + var okFallback bool 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: fallback cache hit | session=%s fallback=%s auth=%s provider=%s model=%s", truncateSessionID(primaryID), truncateSessionID(fallbackID), auth.ID, provider, model) - return auth, nil - } - } + genFallback, authFallback, aliasesFallback, okFallback = s.cache.Observe(fallbackKey) + } + + splitConflict := false + if okPrimary && okFallback { + if genPrimary != genFallback || authPrimary != authFallback || !equalSessionAliases(aliasesPrimary, aliasesFallback) { + splitConflict = true } } @@ -732,42 +778,112 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri if err != nil { return nil, err } - 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) - return auth, nil -} -func selectorLogEntry(ctx context.Context) *log.Entry { - if ctx == nil { - return log.NewEntry(log.StandardLogger()) - } - if reqID := logging.GetRequestID(ctx); reqID != "" { - return log.WithField("request_id", reqID) + if splitConflict { + // The prompt-cache and conversation aliases were previously bound to + // different auths and both cached credentials are unavailable, so + // pickCached missed. Reconcile BOTH alias groups into a single group + // bound to the selected auth. Rebinding them 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.mergeSplitGroupsCAS(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 { + groupKey := cacheKey + var expectedGen uint64 + var expectedAuth string + var expectedAliases []string + if okPrimary { + expectedGen = genPrimary + expectedAuth = authPrimary + expectedAliases = aliasesPrimary + } else if okFallback { + groupKey = fallbackKey + expectedGen = genFallback + expectedAuth = authFallback + expectedAliases = aliasesFallback + } + var newKeys []string + if fallbackKey != "" { + newKeys = []string{cacheKey, fallbackKey} + } else { + newKeys = []string{cacheKey} + } + if !s.rebindGroupCAS(groupKey, expectedGen, expectedAuth, expectedAliases, auth.ID, newKeys) { + entry.Infof("session-affinity: rebind lost to concurrent writer after retries | session=%s auth=%s provider=%s model=%s", truncateSessionID(primaryID), auth.ID, provider, model) + } } - return log.NewEntry(log.StandardLogger()) -} -// truncateSessionID shortens session ID for logging (first 8 chars + "...") -func truncateSessionID(id string) string { - if len(id) <= 20 { - return id - } - return id[:8] + "..." + entry.Infof("session-affinity: cache miss, rebound candidate | session=%s auth=%s provider=%s model=%s", truncateSessionID(primaryID), auth.ID, provider, model) + return auth, nil } -// Stop releases resources held by the selector. -func (s *SessionAffinitySelector) Stop() { - if s.cache != nil { - s.cache.Stop() +// rebindGroupCAS atomically rebinds a session group to authID. When the +// compare-and-replace loses to a concurrent writer, the group is re-observed +// and the binding retried (bounded), so the auth picked for this request is +// not silently dropped by a stale generation. +func (s *SessionAffinitySelector) rebindGroupCAS(sessionKey string, expectedGen uint64, expectedAuth string, expectedAliases []string, authID string, newKeys []string) bool { + for attempt := 0; attempt < 3; attempt++ { + if s.cache.CompareAndReplaceGroup(expectedGen, expectedAuth, expectedAliases, authID, newKeys...) { + return true + } + gen, currentAuth, aliases, ok := s.cache.Observe(sessionKey) + if !ok { + expectedGen, expectedAuth, expectedAliases = 0, "", nil + continue + } + expectedGen, expectedAuth, expectedAliases = gen, currentAuth, aliases } + return false } -// InvalidateAuth removes all session bindings for a specific auth. -// Called when an auth becomes rate-limited or unavailable. -func (s *SessionAffinitySelector) InvalidateAuth(authID string) { - if s.cache != nil { - s.cache.InvalidateAuth(authID) +// mergeSplitGroupsCAS 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) mergeSplitGroupsCAS(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. + var retainedF []string + var deletedAuthF string + for attempt := 0; attempt < 3; attempt++ { + genP, authP, aliasesP, okP := s.cache.Observe(cacheKey) + genF, authF, aliasesF, okF := s.cache.Observe(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.CompareAndReplaceGroup(genP, authP, aliasesP, authID, merged...) { + return true + } + continue + } + if s.cache.CompareAndReplaceGroup(0, "", nil, 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. @@ -780,10 +896,15 @@ func (s *SessionAffinitySelector) OnResult(res Result) { 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 := canonicalModelKey(res.Model) if raw, ok := res.Options.Metadata[cliproxyexecutor.SessionAffinityModelMetadataKey].(string); ok && raw != "" { nsModel = canonicalModelKey(raw) @@ -795,6 +916,9 @@ func (s *SessionAffinitySelector) OnResult(res Result) { fallbackKey = ns + "::" + fallbackID + "::" + nsModel } if res.Success { + // Refresh TTL only while the session is still bound to this auth: + // a late success from a stale attempt must not steal a session that + // was rebound to another auth meanwhile (upstream rebound-safety fix). s.cache.Touch(cacheKey, res.AuthID) if fallbackKey != "" { s.cache.Touch(fallbackKey, res.AuthID) @@ -806,9 +930,96 @@ func (s *SessionAffinitySelector) OnResult(res Result) { return } - s.cache.CompareAndDelete(cacheKey, res.AuthID) - if fallbackKey != "" { - s.cache.CompareAndDelete(fallbackKey, res.AuthID) + aliases := s.cache.CompareAndDeleteAliases(cacheKey, res.AuthID) + if len(aliases) == 0 && fallbackKey != "" { + aliases = s.cache.CompareAndDeleteAliases(fallbackKey, res.AuthID) + } + 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()) + } + if reqID := logging.GetRequestID(ctx); reqID != "" { + return log.WithField("request_id", reqID) + } + return log.NewEntry(log.StandardLogger()) +} + +// truncateSessionID shortens session ID for logging (first 8 chars + "...") +func truncateSessionID(id string) string { + if len(id) <= 20 { + return id + } + return id[:8] + "..." +} + +// Stop releases resources held by the selector. +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. +// Called when an auth becomes rate-limited or unavailable. +func (s *SessionAffinitySelector) InvalidateAuth(authID string) { + if s.cache != nil { + s.cache.InvalidateAuth(authID) + } + if s.quarantine != nil { + s.quarantine.InvalidateAuth(authID) } } 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 00000000000..aba10df8e22 --- /dev/null +++ b/sdk/cliproxy/auth/selector_review_p2_test.go @@ -0,0 +1,304 @@ +package auth + +import ( + "context" + "errors" + "fmt" + "net/http" + "slices" + "testing" + "time" + + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +// Regression tests for codex pullrequestreview-4943660625 on PR #4881 +// (three P2 findings in selector.go). + +// 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) + } +} + +// TestRebindGroupCASRetriesAfterConcurrentWrite is a regression guard for the +// third finding: when CompareAndReplaceGroup loses to a concurrent writer +// between Observe and the rebind, the binding is retried with a fresh +// observation instead of being silently dropped. +func TestRebindGroupCASRetriesAfterConcurrentWrite(t *testing.T) { + t.Parallel() + + selector := NewSessionAffinitySelector(&FillFirstSelector{}) + cacheKey := "pck:test" + keys := []string{cacheKey} + + selector.cache.Set(cacheKey, "auth-old") + gen, authID, aliases, ok := selector.cache.Observe(cacheKey) + if !ok { + t.Fatal("Observe() ok = false after Set") + } + + // Simulate a concurrent writer changing the group between Observe and the + // compare-and-replace: the first CAS attempt now loses the race. + selector.cache.Set(cacheKey, "auth-other") + + if !selector.rebindGroupCAS(cacheKey, gen, authID, aliases, "auth-new", keys) { + t.Fatal("rebindGroupCAS() = false, want success after re-observe and retry") + } + _, finalAuth, _, ok := selector.cache.Observe(cacheKey) + if !ok || finalAuth != "auth-new" { + t.Fatalf("Observe() authID = %q, ok = %v; want auth-new bound", finalAuth, ok) + } +} + +// TestPickRebindsSplitAffinityGroupsOnFailover is a regression guard for the +// codex P2 finding on PR #4881: when the prompt-cache and conversation +// aliases were previously bound to different auths (split groups) and both +// cached credentials are unavailable, pickCached missed and the splitConflict +// branch skipped the rebind, leaving the session pinned to the dead split +// bindings. The fallback auth must now be reconciled onto both groups. +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) + } + + genP, gotPrimary, aliasesPrimary, okPrimary := selector.cache.Observe(primaryKey) + if !okPrimary || gotPrimary != "auth-c" { + t.Fatalf("primary group after failover = %q (ok=%v), want auth-c", gotPrimary, okPrimary) + } + genF, gotFallback, _, okFallback := selector.cache.Observe(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 Observe and the delete, a stale merge observation +// must not remove the newer group, otherwise the newly attached aliases lose +// their affinity binding. +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) + + gen, authID, aliases, ok := cache.Observe(key) + if !ok || authID != "auth-a" { + t.Fatalf("Observe() = %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.Observe(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. + gen, authID, aliases, ok = cache.Observe(key) + if !ok { + t.Fatal("Observe() 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.Observe(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 CompareAndReplaceGroup 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. +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.Observe(fallbackKey); ok { + continue + } + if bumps >= 2 { + return + } + genP, authP, aliasesP, okP := selector.cache.Observe(cacheKey) + if okP { + bumped := append(append([]string(nil), aliasesP...), scratch) + if selector.cache.CompareAndReplaceGroup(genP, authP, aliasesP, authP, bumped...) { + bumps++ + } + } + } + }() + + merged := selector.mergeSplitGroupsCAS(cacheKey, fallbackKey, "auth-c") + close(finished) + <-done + if !merged { + t.Fatalf("iteration %d: mergeSplitGroupsCAS() = false under single-bump contention, want true", i) + } + _, got, aliases, ok := selector.cache.Observe(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) + } + } +} diff --git a/sdk/cliproxy/auth/selector_test.go b/sdk/cliproxy/auth/selector_test.go index 8df0520a795..aaa0a6135b9 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++ { @@ -863,6 +864,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) } @@ -872,6 +874,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) } @@ -901,6 +904,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++ { @@ -962,7 +966,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 @@ -1002,6 +1008,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) @@ -1019,6 +1026,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++ { @@ -1028,6 +1036,353 @@ func TestSessionAffinitySelector_FailoverWhenAuthUnavailable(t *testing.T) { } } } +func TestSessionAffinitySelector_CachedAuthUnavailableRebindsWholeAliasGroupRebindAliases(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}) + + 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, err := selector.Pick(context.Background(), provider, model, opts, availableWithoutFirst) + if err != nil { + t.Fatalf("%s Pick() error = %v", payload.name, err) + } + 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 TestSessionAffinitySelector_CachedAuthUnavailableRebindsSharedPromptGroup(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, err := selector.Pick(context.Background(), provider, model, opts, availableWithoutFirst) + if err != nil { + t.Fatalf("%s Pick() error = %v", payload.name, err) + } + 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 TestSessionAffinitySelector_CachedAuthUnavailableConcurrencyNewerBinding(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}) + + 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) + } + + for i := 0; i < 5; i++ { + got, err := selector.Pick(context.Background(), provider, model, combined, availableWithoutFirst) + if err != nil { + t.Fatalf("concurrent Pick() #%d error = %v", i, err) + } + 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() { + gen, authID, aliases, ok := selector.cache.Observe(fallbackKey) + if !ok { + t.Fatalf("expected fallbackKey %q in cache", fallbackKey) + } + // Bump generation concurrently + if !selector.cache.CompareAndReplaceGroup(gen, authID, aliases, authID, fallbackKey) { + t.Fatalf("CompareAndReplaceGroup 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_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 CompareAndReplaceGroup 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. mergeSplitGroupsCAS attempts to merge cacheKey and fallbackKey into auth-c. + // Since cacheKey expects missingPrimaryAlias which is missing from entries, + // CompareAndReplaceGroup fails on all 3 attempts. + merged := selector.mergeSplitGroupsCAS(cacheKey, fallbackKey, "auth-c") + if merged { + t.Fatalf("mergeSplitGroupsCAS must fail due to CAS exhaustion") + } + + // Since mergeSplitGroupsCAS 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 CompareAndReplaceGroup 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. mergeSplitGroupsCAS attempts to merge cacheKey and fallbackKey into auth-c. + // Since cacheKey expects missingPrimaryAlias which is missing from entries, + // CompareAndReplaceGroup fails on all 3 attempts. + merged := selector.mergeSplitGroupsCAS(cacheKey, fallbackKey, "auth-c") + close(finished) + <-done + if merged { + t.Fatalf("iteration %d: mergeSplitGroupsCAS 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() + } +} + +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 TestExtractSessionID_ClaudeCodePriorityOverHeader(t *testing.T) { t.Parallel() @@ -1425,6 +1780,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 { @@ -1440,6 +1796,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 { @@ -1477,6 +1834,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) @@ -1505,6 +1863,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) @@ -1514,6 +1873,176 @@ 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 TestSessionCacheCompareAndDeleteAliasesPreservesNewerBinding(t *testing.T) { + cache := NewSessionCache(time.Minute) + defer cache.Stop() + cache.SetAliases("auth-a", "prompt", "conversation") + cache.SetAliases("auth-b", "prompt", "conversation") + + if aliases := cache.CompareAndDeleteAliases("prompt", "auth-a"); len(aliases) != 0 { + t.Fatalf("CompareAndDeleteAliases() = %v for stale auth, want none", aliases) + } + 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 TestSessionCacheGenerationToken_ABARejected(t *testing.T) { + cache := NewSessionCache(time.Minute) + defer cache.Stop() + + cache.SetAliases("auth-a", "pck:shared-prompt", "conv:sess-1") + gen1, auth1, aliases1, ok1 := cache.Observe("pck:shared-prompt") + if !ok1 || auth1 != "auth-a" || gen1 == 0 { + t.Fatalf("initial observe failed: gen=%d auth=%q aliases=%v ok=%v", gen1, auth1, aliases1, ok1) + } + + // Mutate A -> B -> A + cache.SetAliases("auth-b", "pck:shared-prompt", "conv:sess-1") + cache.SetAliases("auth-a", "pck:shared-prompt", "conv:sess-1") + + gen3, _, _, _ := cache.Observe("pck:shared-prompt") + if gen3 == gen1 { + t.Fatalf("expected generation to increment after mutations: gen1=%d gen3=%d", gen1, gen3) + } + + // Stale CAS referencing gen1 must fail + replaced := cache.CompareAndReplaceGroup(gen1, "auth-a", aliases1, "auth-c", "pck:shared-prompt", "conv:sess-1") + if replaced { + t.Fatalf("CAS with stale generation must fail on ABA cycle") + } + + for _, key := range []string{"pck:shared-prompt", "conv:sess-1"} { + if got, ok := cache.Get(key); !ok || got != "auth-a" { + t.Fatalf("cache.Get(%q) = %q, %v; want auth-a, true", key, got, ok) + } + } +} + +func TestSessionCacheGenerationToken_PartialSplitAbort(t *testing.T) { + cache := NewSessionCache(time.Minute) + defer cache.Stop() + + // Group 1: primary prompt cache key bound to auth-a + cache.Set("pck:prompt-1", "auth-a") + genA, authA, aliasesA, okA := cache.Observe("pck:prompt-1") + if !okA || authA != "auth-a" { + t.Fatalf("observe prompt-1 failed: %v", okA) + } + + // Group 2: conversation fallback key bound to auth-b + cache.Set("conv:sess-1", "auth-b") + + // Attempting CAS to rebind prompt-1 and attach conv:sess-1 when conv:sess-1 belongs to auth-b + replaced := cache.CompareAndReplaceGroup(genA, "auth-a", aliasesA, "auth-c", "pck:prompt-1", "conv:sess-1") + if replaced { + t.Fatalf("CAS must fail when a new alias belongs to another active live group") + } + + if got, ok := cache.Get("pck:prompt-1"); !ok || got != "auth-a" { + t.Fatalf("prompt-1 must remain auth-a, got %q, %v", got, ok) + } + if got, ok := cache.Get("conv:sess-1"); !ok || got != "auth-b" { + t.Fatalf("conv:sess-1 must remain auth-b, got %q, %v", got, ok) + } +} + +func TestSessionCacheGenerationToken_SharedPromptAndTTLPreservedOnReplace(t *testing.T) { + cache := NewSessionCache(time.Minute) + defer cache.Stop() + + cache.SetAliases("auth-a", "pck:shared-prompt", "conv:sess-1") + gen, authID, aliases, ok := cache.Observe("pck:shared-prompt") + if !ok { + t.Fatalf("observe failed") + } + + replaced := cache.CompareAndReplaceGroup(gen, authID, aliases, "auth-b", "pck:shared-prompt", "conv:sess-1", "conv:sess-2") + if !replaced { + t.Fatalf("CompareAndReplaceGroup failed") + } + + for _, key := range []string{"pck:shared-prompt", "conv:sess-1", "conv:sess-2"} { + 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 TestSessionAffinitySelector_FallbackFailureKeepsOldGroupIntact(t *testing.T) { + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: &RoundRobinSelector{}, + TTL: time.Minute, + }) + defer selector.Stop() + + authA := &Auth{ID: "auth-a", Status: StatusActive} + provider := "test-prov" + model := "gpt-4o" + req := []byte(`{"conversation":{"id":"sess-keep-intact"},"prompt_cache_key":"prompt-keep-intact"}`) + + // Bind initially to auth-a + picked, err := selector.Pick(context.Background(), provider, model, cliproxyexecutor.Options{OriginalRequest: req}, []*Auth{authA}) + if err != nil || picked.ID != "auth-a" { + t.Fatalf("initial pick = %v, err = %v; want auth-a", picked, err) + } + + // Now auth-a becomes unavailable (empty candidate slice), fallback fails + picked2, err2 := selector.Pick(context.Background(), provider, model, cliproxyexecutor.Options{OriginalRequest: req}, []*Auth{}) + if err2 == nil { + t.Fatalf("expected error when no auths available, got %v", picked2) + } + + // Verify that cache still holds auth-a and was NOT eagerly deleted + cacheKey := provider + "::pck:prompt-keep-intact::" + model + convKey := provider + "::conv:sess-keep-intact::" + model + if got, ok := selector.cache.Get(cacheKey); !ok || got != "auth-a" { + t.Fatalf("cacheKey %q = %q, %v; want auth-a intact", cacheKey, got, ok) + } + if got, ok := selector.cache.Get(convKey); !ok || got != "auth-a" { + t.Fatalf("convKey %q = %q, %v; want auth-a intact", convKey, got, ok) + } +} + func TestSessionAffinitySelectorPrimaryTrafficKeepsConversationAliasAlive(t *testing.T) { selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ Fallback: &RoundRobinSelector{}, @@ -1531,6 +2060,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] @@ -1542,6 +2072,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) } @@ -1573,10 +2104,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) } @@ -1607,6 +2140,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) @@ -1696,6 +2230,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} @@ -1706,6 +2241,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) @@ -1792,6 +2328,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}) @@ -1801,6 +2338,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++ { @@ -1914,6 +2452,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 @@ -1956,6 +2495,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 { @@ -2192,6 +2783,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) @@ -2201,6 +2793,164 @@ func TestSessionAffinitySelectorUsesRequestPayloadWhenOriginalRequestMissing(t * } } +// recordingFallbackSelector wraps a Selector and records the auth IDs passed to +// each Pick call, so tests can assert the fallback only receives available auths. +type recordingFallbackSelector struct { + inner Selector + mu sync.Mutex + last []string +} + +func (r *recordingFallbackSelector) Pick(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth) (*Auth, error) { + r.mu.Lock() + r.last = make([]string, 0, len(auths)) + for _, a := range auths { + r.last = append(r.last, a.ID) + } + r.mu.Unlock() + return r.inner.Pick(ctx, provider, model, opts, auths) +} + +func (r *recordingFallbackSelector) lastPick() []string { + r.mu.Lock() + defer r.mu.Unlock() + out := make([]string, len(r.last)) + copy(out, r.last) + return out +} + +// TestSessionAffinitySelector_FallbackReselectReceivesOnlyAvailable is a +// regression test for the fallback reselect path: when the cached auth is no +// longer available, the fallback must only receive the available auths, not the +// full list (which could include the now-unavailable cached auth). +func TestSessionAffinitySelector_FallbackReselectReceivesOnlyAvailable(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_fallback-reselect-test"}}`) + opts := cliproxyexecutor.Options{OriginalRequest: payload} + + // First pick establishes the binding (fallback receives the full list here). + 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}) + + // Make the bound auth unavailable so it is filtered out of `available`. + bound := first.ID + for _, a := range auths { + if a.ID == bound { + a.Unavailable = true + a.NextRetryAfter = time.Now().Add(time.Hour) + } + } + + // Second pick with the SAME full auths list: the cached auth is now filtered + // out of available, so the fallback reselect path fires. It must receive + // only the two still-available auths, not the full list. + if _, err := selector.Pick(context.Background(), "claude", "claude-3", opts, auths); err != nil { + t.Fatalf("Pick() after failover error = %v", err) + } + + got := rec.lastPick() + if len(got) != 2 { + t.Fatalf("fallback reselect received %d auths, want 2 (only available); got %v", len(got), got) + } + for _, id := range got { + if id == bound { + t.Fatalf("fallback reselect received unavailable cached auth %q; received %v", bound, got) + } + } +} + +// 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 pick establishes the session-affinity binding to the first auth. + 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}) + + // The first auth just failed (e.g. 429). Crucially we do NOT set any local + // cooldown/Unavailable flag: that mirrors the home/mixed-mode bug where + // reportHomeResult never updates local availability, so the auth stays + // locally available and would otherwise be re-picked by the cache hit. + failed := first.ID + opts.Metadata = map[string]any{ + cliproxyexecutor.ExcludedAuthIDsMetadataKey: map[string]struct{}{failed: {}}, + } + + // Simulate the retry loop for the SAME request/session: pick again. The + // failed auth must never be returned, even though it is still locally + // available. Before the fix the cache hit returned `failed` on every retry, + // producing the 12x "cache hit but auth unavailable, reselected" carousel. + for attempt := 0; attempt < 20; attempt++ { + got, err := selector.Pick(context.Background(), "claude", "claude-3", opts, auths) + if err != nil { + t.Fatalf("Pick() attempt %d error = %v", attempt, err) + } + if got.ID == failed { + t.Fatalf("attempt %d re-picked auth %q that already failed in this request; want a different auth", attempt, failed) + } + } +} + +func TestSessionAffinitySelector_QuarantineCleanupCadence(t *testing.T) { + selector := NewSessionAffinitySelector(nil) + if selector.quarantine == nil { + t.Fatal("selector.quarantine is nil") + } + if selector.cache == nil { + t.Fatal("selector.cache is nil") + } + if selector.cache.ttl != time.Hour { + t.Fatalf("selector.cache.ttl = %v, want 1h", selector.cache.ttl) + } + if selector.quarantine.ttl > 5*time.Second { + t.Fatalf("selector.quarantine.ttl = %v, want <= 5s to decouple cleanup cadence from 1h affinity TTL", selector.quarantine.ttl) + } +} + func TestSessionCache_StopConcurrent(t *testing.T) { t.Parallel() for iter := 0; iter < 100; iter++ { diff --git a/sdk/cliproxy/auth/session_affinity_priority_test.go b/sdk/cliproxy/auth/session_affinity_priority_test.go index adb1c67bfd2..7426cf27050 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_affinity_quarantine_test.go b/sdk/cliproxy/auth/session_affinity_quarantine_test.go new file mode 100644 index 00000000000..672d4e5e53b --- /dev/null +++ b/sdk/cliproxy/auth/session_affinity_quarantine_test.go @@ -0,0 +1,136 @@ +package auth + +import ( + "context" + "net/http" + "testing" + "time" + + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +type quarantinePickSelector func(context.Context, string, string, cliproxyexecutor.Options, []*Auth) (*Auth, error) + +func (pick quarantinePickSelector) Pick(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth) (*Auth, error) { + return pick(ctx, provider, model, opts, auths) +} + +func quarantineOptions(sessionID string) cliproxyexecutor.Options { + return cliproxyexecutor.Options{ + Headers: http.Header{"X-Session-Id": []string{sessionID}}, + Metadata: map[string]any{ + cliproxyexecutor.SessionAffinityProviderMetadataKey: "mixed", + cliproxyexecutor.SessionAffinityModelMetadataKey: ".gemini-flash", + }, + } +} + +func newQuarantineSelector() *SessionAffinitySelector { + return NewSessionAffinitySelector(quarantinePickSelector(func(_ context.Context, _, _ string, _ cliproxyexecutor.Options, auths []*Auth) (*Auth, error) { + return auths[0], nil + })) +} + +func quarantine429(selector *SessionAffinitySelector, opts cliproxyexecutor.Options, authID string, delay time.Duration) { + selector.OnResult(Result{ + AuthID: authID, + Provider: "gemini", + Model: "gemini-3.6-flash", + Success: false, + Error: &Error{HTTPStatus: http.StatusTooManyRequests}, + RetryAfter: &delay, + Options: opts, + }) +} + +func TestSessionAffinityQuarantineRetryAfterIsSessionLocal(t *testing.T) { + selector := newQuarantineSelector() + defer selector.Stop() + authA := &Auth{ID: "auth-a", Provider: "gemini"} + authB := &Auth{ID: "auth-b", Provider: "gemini"} + opts := quarantineOptions("session-one") + + quarantine429(selector, opts, authA.ID, 53*time.Second) + got, err := selector.Pick(context.Background(), "mixed", ".gemini-flash", opts, []*Auth{authA, authB}) + if err != nil || got.ID != authB.ID { + t.Fatalf("same-session Pick = %v/%v, want auth-b", got, err) + } + + other, err := selector.Pick(context.Background(), "mixed", ".gemini-flash", quarantineOptions("session-two"), []*Auth{authA, authB}) + if err != nil || other.ID != authA.ID { + t.Fatalf("other-session Pick = %v/%v, want auth-a", other, err) + } +} + +func TestSessionAffinityQuarantineTracksMultipleFailures(t *testing.T) { + selector := newQuarantineSelector() + defer selector.Stop() + authA := &Auth{ID: "auth-a", Provider: "gemini"} + authB := &Auth{ID: "auth-b", Provider: "gemini"} + authC := &Auth{ID: "auth-c", Provider: "gemini"} + opts := quarantineOptions("session-multiple") + delay := 53 * time.Second + + quarantine429(selector, opts, authA.ID, delay) + quarantine429(selector, opts, authB.ID, delay) + got, err := selector.Pick(context.Background(), "mixed", ".gemini-flash", opts, []*Auth{authA, authB, authC}) + if err != nil || got.ID != authC.ID { + t.Fatalf("Pick = %v/%v, want auth-c", got, err) + } +} + +func TestSessionAffinityQuarantineExpires(t *testing.T) { + selector := newQuarantineSelector() + defer selector.Stop() + authA := &Auth{ID: "auth-a", Provider: "gemini"} + authB := &Auth{ID: "auth-b", Provider: "gemini"} + opts := quarantineOptions("session-expiry") + + quarantine429(selector, opts, authA.ID, 20*time.Millisecond) + before, _ := selector.Pick(context.Background(), "mixed", ".gemini-flash", opts, []*Auth{authA, authB}) + if before.ID != authB.ID { + t.Fatalf("Pick before expiry = %q, want auth-b", before.ID) + } + selector.OnResult(Result{AuthID: authB.ID, Provider: "gemini", Model: ".gemini-flash", Success: false, Error: &Error{HTTPStatus: http.StatusBadGateway}, Options: opts}) + time.Sleep(30 * time.Millisecond) + after, _ := selector.Pick(context.Background(), "mixed", ".gemini-flash", opts, []*Auth{authA, authB}) + if after.ID != authA.ID { + t.Fatalf("Pick after expiry = %q, want auth-a", after.ID) + } +} + +func TestSessionAffinityQuarantineSurvivesStaleSuccess(t *testing.T) { + selector := newQuarantineSelector() + defer selector.Stop() + authA := &Auth{ID: "auth-a", Provider: "gemini"} + authB := &Auth{ID: "auth-b", Provider: "gemini"} + opts := quarantineOptions("session-stale-success") + + quarantine429(selector, opts, authA.ID, 53*time.Second) + 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", got, err) + } +} + +func TestSessionAffinityQuarantineSkipsRequestScopedErrors(t *testing.T) { + selector := newQuarantineSelector() + defer selector.Stop() + authA := &Auth{ID: "auth-a", Provider: "gemini"} + authB := &Auth{ID: "auth-b", Provider: "gemini"} + opts := quarantineOptions("session-client-error") + + 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 error = %v/%v, want auth-a", got, err) + } +} diff --git a/sdk/cliproxy/auth/session_cache.go b/sdk/cliproxy/auth/session_cache.go index 5dfd9594fb7..a2910029af1 100644 --- a/sdk/cliproxy/auth/session_cache.go +++ b/sdk/cliproxy/auth/session_cache.go @@ -10,18 +10,20 @@ 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{} - stopOnce sync.Once + mu sync.RWMutex + entries map[string]sessionEntry + ttl time.Duration + stopCh chan struct{} + stopOnce sync.Once + generation uint64 } // NewSessionCache creates a cache with the specified TTL. @@ -93,6 +95,22 @@ func (c *SessionCache) GetAndRefresh(sessionID string) (string, bool) { return entry.authID, true } +// Observe returns the current generation token, auth ID, and aliases for a session ID +// without refreshing its TTL or acquiring a write lock. +func (c *SessionCache) Observe(sessionID string) (gen uint64, authID string, aliases []string, ok bool) { + if c == nil || sessionID == "" { + return 0, "", nil, false + } + now := time.Now() + c.mu.RLock() + defer c.mu.RUnlock() + entry, exists := c.entries[sessionID] + if !exists || !now.Before(entry.expiresAt) { + return 0, "", nil, false + } + return entry.generation, entry.authID, 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) { @@ -101,10 +119,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() @@ -126,20 +188,22 @@ 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) { for _, previous := range previousGroups { c.removeAliasGroupLocked(previous) } - entry := sessionEntry{authID: authID, expiresAt: expiresAt, aliases: aliases} + c.generation++ + entry := sessionEntry{authID: authID, expiresAt: expiresAt, aliases: aliases, generation: c.generation} for _, alias := range aliases { c.entries[alias] = entry } } func (c *SessionCache) removeAliasGroupLocked(entry sessionEntry) { + c.generation++ for _, alias := range entry.aliases { current, ok := c.entries[alias] if !ok || current.authID != entry.authID || !current.expiresAt.Equal(entry.expiresAt) || @@ -150,6 +214,76 @@ func (c *SessionCache) removeAliasGroupLocked(entry sessionEntry) { } } +// CompareAndReplaceGroup atomically validates that an observed group has not mutated +// (matching expectedGen, expectedAuthID, and expectedAliases), confirms no requested +// new alias belongs to another active live group, and replaces the whole group with newAuthID. +func (c *SessionCache) CompareAndReplaceGroup(expectedGen uint64, expectedAuthID string, expectedAliases []string, newAuthID string, newSessionIDs ...string) bool { + if c == nil || newAuthID == "" { + return false + } + now := time.Now() + c.mu.Lock() + defer c.mu.Unlock() + + if expectedGen != 0 { + if len(expectedAliases) == 0 { + return false + } + for _, alias := range expectedAliases { + current, ok := c.entries[alias] + if !ok || !now.Before(current.expiresAt) { + return false + } + if current.generation != expectedGen || current.authID != expectedAuthID || !equalSessionAliases(current.aliases, expectedAliases) { + return false + } + } + } else { + for _, sid := range newSessionIDs { + if sid == "" { + continue + } + if current, ok := c.entries[sid]; ok && now.Before(current.expiresAt) { + return false + } + } + } + + candidateAliases := mergeSessionAliases(expectedAliases, newSessionIDs...) + for _, alias := range candidateAliases { + current, ok := c.entries[alias] + if !ok || !now.Before(current.expiresAt) { + continue + } + if expectedGen == 0 || current.generation != expectedGen || current.authID != expectedAuthID { + return false + } + } + + newAliases := compactSessionAliases(candidateAliases) + if len(newAliases) == 0 { + return false + } + + if expectedGen != 0 { + for _, alias := range expectedAliases { + delete(c.entries, alias) + } + } + + c.generation++ + entry := sessionEntry{ + authID: newAuthID, + expiresAt: now.Add(c.ttl), + aliases: newAliases, + generation: c.generation, + } + for _, alias := range newAliases { + c.entries[alias] = entry + } + return true +} + func compactSessionAliases(aliases []string) []string { return compactSessionAliasesWith(aliases, isLocalPromptCacheSessionAlias) } @@ -279,6 +413,8 @@ func (c *SessionCache) Invalidate(sessionID string) { return } c.mu.Lock() + defer c.mu.Unlock() + c.generation++ entry, ok := c.entries[sessionID] delete(c.entries, sessionID) if ok { @@ -297,10 +433,52 @@ func (c *SessionCache) Invalidate(sessionID string) { } } current.aliases = filtered + current.generation = c.generation c.entries[alias] = current } } - c.mu.Unlock() +} + +// CompareAndDeleteAliases removes a binding and returns every alias that still +// belongs to the same expected auth. A stale result 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.removeAliasGroupLocked(entry) + 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. +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.removeAliasGroupLocked(entry) + return removed } // InvalidateAuth removes all sessions bound to a specific auth ID. @@ -310,12 +488,13 @@ func (c *SessionCache) InvalidateAuth(authID string) { return } c.mu.Lock() + defer c.mu.Unlock() + c.generation++ for sid, entry := range c.entries { if entry.authID == authID { delete(c.entries, sid) } } - c.mu.Unlock() } // Stop terminates the background cleanup goroutine. diff --git a/sdk/cliproxy/auth/stream_ttft_test.go b/sdk/cliproxy/auth/stream_ttft_test.go new file mode 100644 index 00000000000..08285c6c4d2 --- /dev/null +++ b/sdk/cliproxy/auth/stream_ttft_test.go @@ -0,0 +1,349 @@ +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 +} + +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) + } +} + +// TestStreamFirstChunkTimeout_UnsupportedKeysIgnored proves the undocumented +// duration form and _seconds key are no longer accepted; only _ms is honored +// with the config fallback verified separately in ConfigAndMetadata. +func TestStreamFirstChunkTimeout_UnsupportedKeysIgnored(t *testing.T) { + m := NewManager(nil, nil, nil) + + opts := cliproxyexecutor.Options{ + Metadata: map[string]any{ + "stream_first_chunk_timeout": time.Second, + "stream_first_chunk_timeout_seconds": 2, + "stream_first_chunk_timeout_ms": 40, + }, + } + if got := m.streamFirstChunkTimeout(opts); got != 40*time.Millisecond { + t.Fatalf("metadata with unsupported keys = %v, want 40ms from _ms", got) + } + + onlyUnsupported := cliproxyexecutor.Options{ + Metadata: map[string]any{ + "stream_first_chunk_timeout": time.Second, + "stream_first_chunk_timeout_seconds": 7, + }, + } + if got := m.streamFirstChunkTimeout(onlyUnsupported); got != 0 { + t.Fatalf("only unsupported metadata keys = %v, want 0 (disabled)", 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/auth/types_test.go b/sdk/cliproxy/auth/types_test.go index 6f8fa28f566..83c7384e4bc 100644 --- a/sdk/cliproxy/auth/types_test.go +++ b/sdk/cliproxy/auth/types_test.go @@ -36,17 +36,12 @@ func TestRequestRetryOverride(t *testing.T) { auth = &Auth{Metadata: map[string]any{"request-retry": 2}} if got, ok := auth.RequestRetryOverride(); !ok || got != 2 { - t.Fatalf("legacy request-retry=2 override = (%d, %t), want (2, true)", got, ok) + t.Fatalf("request-retry=2 override = (%d, %t), want (2, true)", got, ok) } auth = &Auth{Metadata: map[string]any{"request-retry": -2}} if got, ok := auth.RequestRetryOverride(); ok || got != 0 { - t.Fatalf("legacy request-retry=-2 override = (%d, %t), want (0, false)", got, ok) - } - - auth = &Auth{Metadata: map[string]any{"request_retry": 0, "request-retry": 2}} - if got, ok := auth.RequestRetryOverride(); !ok || got != 0 { - t.Fatalf("canonical request_retry precedence = (%d, %t), want (0, true)", got, ok) + t.Fatalf("request-retry=-2 override = (%d, %t), want (0, false)", got, ok) } auth = &Auth{Metadata: map[string]any{"request_retry": "0"}} diff --git a/sdk/cliproxy/executor/types.go b/sdk/cliproxy/executor/types.go index 9839f4ac94d..d00701d633b 100644 --- a/sdk/cliproxy/executor/types.go +++ b/sdk/cliproxy/executor/types.go @@ -48,11 +48,18 @@ 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 model used during session affinity selection. + // 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" ) @@ -114,6 +121,9 @@ type RequestTerminatedError struct { HTTPStatus int Header http.Header Body []byte + // Trusted reports that the response originated from a trusted in-process + // interceptor rather than an untrusted upstream HTTP error. + Trusted bool } func (e *RequestTerminatedError) Error() string { From 239f5fb900afa723a287f57916e46e2574cd66c2 Mon Sep 17 00:00:00 2001 From: warelik Date: Sun, 23 Aug 2026 06:55:17 -0400 Subject: [PATCH 2/9] fix(auth): treat lone signatures as empty Signature without text, tool-call, or positive tokens is not liveness. Record empty completions before success, apply the check on Home Execute, keep blank lines in raw JSON frames, match invalid_api_key, and redact secrets in upstream logs. --- sdk/cliproxy/auth/conductor_cooldown.go | 4 +- sdk/cliproxy/auth/conductor_execution.go | 30 +++++---- sdk/cliproxy/auth/conductor_home_execution.go | 8 +++ sdk/cliproxy/auth/conductor_overrides_test.go | 33 ++++++++++ sdk/cliproxy/auth/empty_completion.go | 48 +++++--------- sdk/cliproxy/auth/empty_completion_test.go | 66 +++++++++++++++---- .../auth/home_execution_paths_test.go | 56 ++++++++++++++++ 7 files changed, 189 insertions(+), 56 deletions(-) diff --git a/sdk/cliproxy/auth/conductor_cooldown.go b/sdk/cliproxy/auth/conductor_cooldown.go index 87bbadca35f..23696b7596b 100644 --- a/sdk/cliproxy/auth/conductor_cooldown.go +++ b/sdk/cliproxy/auth/conductor_cooldown.go @@ -1625,7 +1625,9 @@ func isInvalidGrantResultError(err *Error) bool { // 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") + return strings.Contains(lowered, "api key not valid") || + strings.Contains(lowered, "api_key_invalid") || + strings.Contains(lowered, "invalid_api_key") } func isInvalidAPIKeyError(err error) bool { diff --git a/sdk/cliproxy/auth/conductor_execution.go b/sdk/cliproxy/auth/conductor_execution.go index d36e1398b8e..773e5aba2f1 100644 --- a/sdk/cliproxy/auth/conductor_execution.go +++ b/sdk/cliproxy/auth/conductor_execution.go @@ -8,6 +8,7 @@ import ( "io" "net/http" "path/filepath" + "regexp" "sort" "strconv" "strings" @@ -501,12 +502,8 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req } continue } - m.MarkResult(execCtx, result) if isEmptyCompletionPayload(resp.Payload) { - result.Success = false - result.Error = errEmptyCompletion - m.MarkResult(execCtx, result) - lastErr = errEmptyCompletion + lastErr = m.markEmptyCompletion(execCtx, &result) tracker.Record(auth, errEmptyCompletion) persistExcludedAuthForRetry(m, auth, errEmptyCompletion, retryRound, defaultRequestRetry, excluded) if homeMode { @@ -514,6 +511,7 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req } continue } + m.MarkResult(execCtx, result) attemptAliasResult := resolveAttemptAliasResult(routing, auth, routeModel, upstreamModel, aliasResult) rewriteForceMappedResponse(&resp, attemptAliasResult) return resp, nil @@ -711,12 +709,8 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, } continue } - m.MarkResult(execCtx, result) if isEmptyCompletionPayload(resp.Payload) { - result.Success = false - result.Error = errEmptyCompletion - m.MarkResult(execCtx, result) - lastErr = errEmptyCompletion + lastErr = m.markEmptyCompletion(execCtx, &result) tracker.Record(auth, errEmptyCompletion) persistExcludedAuthForRetry(m, auth, errEmptyCompletion, retryRound, defaultRequestRetry, excluded) if homeMode { @@ -724,6 +718,7 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, } continue } + m.MarkResult(execCtx, result) attemptAliasResult := resolveAttemptAliasResult(routing, auth, routeModel, upstreamModel, aliasResult) rewriteForceMappedResponse(&resp, attemptAliasResult) return resp, nil @@ -1725,11 +1720,24 @@ func formatAuthIdentity(auth *Auth, provider string) string { } } +var ( + logSecretAPIKeyPattern = regexp.MustCompile(`(?i)\bsk-[A-Za-z0-9_-]{8,}`) + logSecretBearerPattern = regexp.MustCompile(`(?i)\b(Bearer|Basic)\s+[A-Za-z0-9._~+/=-]{4,}`) + logSecretLabeledPattern = regexp.MustCompile(`(?i)((?:api[_-]?key|access[_-]?token|secret)\s*[=:]\s*)([^\s"&,;}]+)`) +) + +func redactSecretsForLog(msg string) string { + msg = logSecretAPIKeyPattern.ReplaceAllString(msg, "[REDACTED]") + msg = logSecretBearerPattern.ReplaceAllString(msg, "$1 [REDACTED]") + msg = logSecretLabeledPattern.ReplaceAllString(msg, "${1}[REDACTED]") + return msg +} + func summarizeErrorForLog(err error) string { if err == nil { return "" } - msg := strings.TrimSpace(err.Error()) + msg := redactSecretsForLog(strings.TrimSpace(err.Error())) const maxRunes = 300 runes := []rune(msg) if len(runes) > maxRunes { diff --git a/sdk/cliproxy/auth/conductor_home_execution.go b/sdk/cliproxy/auth/conductor_home_execution.go index 17068b38f5e..37035f05f78 100644 --- a/sdk/cliproxy/auth/conductor_home_execution.go +++ b/sdk/cliproxy/auth/conductor_home_execution.go @@ -278,6 +278,14 @@ func (m *Manager) executeHomeOnce(ctx context.Context, providers []string, req c } result := Result{AuthID: preparedAuth.ID, Provider: selection.Provider, Model: resultModel, Success: errExecute == nil, Options: execOpts} if errExecute == nil { + if !countTokens && isEmptyCompletionPayload(response.Payload) { + result.Success = false + result.Error = errEmptyCompletion + m.reportHomeResult(execCtx, result, preparedAuth) + tracker.Record(preparedAuth, errEmptyCompletion) + lastErr = errEmptyCompletion + break + } m.reportHomeResult(execCtx, result, preparedAuth) releaseAttempt() attemptAliasResult := resolveAttemptAliasResult(routing, preparedAuth, routeModel, upstreamModel, aliasResult) diff --git a/sdk/cliproxy/auth/conductor_overrides_test.go b/sdk/cliproxy/auth/conductor_overrides_test.go index dcb2501f653..e9583700085 100644 --- a/sdk/cliproxy/auth/conductor_overrides_test.go +++ b/sdk/cliproxy/auth/conductor_overrides_test.go @@ -6,6 +6,7 @@ import ( "fmt" "net/http" "slices" + "strings" "sync" "testing" "time" @@ -2667,6 +2668,23 @@ func TestIsCredentialScopedError_InvalidAPIKey(t *testing.T) { t.Fatalf("expected isCredentialScopedError(invalidKey401) = true, got false") } + invalidAPIKeyCode := &Error{ + HTTPStatus: http.StatusUnauthorized, + Code: "invalid_api_key", + Message: "Invalid token", + } + if !isCredentialScopedError(invalidAPIKeyCode) { + t.Fatalf("expected isCredentialScopedError(invalid_api_key code) = true, got false") + } + + invalidAPIKeyBody := &Error{ + HTTPStatus: http.StatusForbidden, + Message: `{"error":{"code":"invalid_api_key","message":"Invalid token"}}`, + } + if !isCredentialScopedError(invalidAPIKeyBody) { + t.Fatalf("expected isCredentialScopedError(invalid_api_key body) = true, got false") + } + normal400 := &Error{ HTTPStatus: http.StatusBadRequest, Message: `invalid argument: field "prompt" cannot be empty`, @@ -2676,6 +2694,21 @@ func TestIsCredentialScopedError_InvalidAPIKey(t *testing.T) { } } +func TestSummarizeErrorForLogRedactsSecrets(t *testing.T) { + got := summarizeErrorForLog(errors.New("Incorrect API key provided: sk-live-secret")) + if strings.Contains(got, "sk-live-secret") { + t.Fatalf("summarizeErrorForLog leaked API key: %q", got) + } + if !strings.Contains(got, "[REDACTED]") { + t.Fatalf("summarizeErrorForLog = %q, want [REDACTED]", got) + } + + got = summarizeErrorForLog(errors.New("authorization failed: Bearer abcdefghijklmnop")) + if strings.Contains(got, "abcdefghijklmnop") { + t.Fatalf("summarizeErrorForLog leaked bearer token: %q", got) + } +} + func TestManagerExecute_InvalidAPIKeyStopsModelPoolRetries(t *testing.T) { m := NewManager(nil, nil, nil) cfg := &internalconfig.Config{ diff --git a/sdk/cliproxy/auth/empty_completion.go b/sdk/cliproxy/auth/empty_completion.go index e48ef4cee97..56bd71e202c 100644 --- a/sdk/cliproxy/auth/empty_completion.go +++ b/sdk/cliproxy/auth/empty_completion.go @@ -636,9 +636,10 @@ type interactionsExtraContent struct { } `json:"google"` } -// hasSignature reports whether the step carries a reasoning signature. A step -// that only carries a signature is still a meaningful upstream answer: dropping -// it makes the turn look empty and costs the signature on the next request. +// hasSignature reports whether the step carries a reasoning signature. +// A signature without text, a tool-call, or positive completion tokens is not +// content: treating it as live would keep a dead Gemini/Claude credential in +// the pool. func (s *interactionsStep) hasSignature() bool { if s == nil { return false @@ -673,19 +674,11 @@ func (c *interactionsContent) hasMeaningfulContent() bool { if c == nil { return false } - if strings.TrimSpace(c.Text) != "" || + return strings.TrimSpace(c.Text) != "" || strings.TrimSpace(c.Data) != "" || strings.TrimSpace(c.FileURI) != "" || strings.TrimSpace(c.FileUri) != "" || - strings.TrimSpace(c.URL) != "" { - return true - } - if strings.TrimSpace(c.Signature) != "" || - strings.TrimSpace(c.ThoughtSignature) != "" || - strings.TrimSpace(c.ThoughtSignatureCamel) != "" { - return true - } - return false + strings.TrimSpace(c.URL) != "" } type interactionsDelta struct { @@ -941,9 +934,8 @@ func (a *emptyCompletionAccum) evalClaude(data []byte) bool { a.hasContent = true } case "signature_delta": - if strings.TrimSpace(chunk.Delta.Signature) != "" { - a.hasContent = true - } + // A lonely signature is not visible content and does not spend + // completion tokens; do not treat it as a live completion. case "citations_delta": if nonEmptyJSONPayload(chunk.Delta.Citation) { a.hasContent = true @@ -953,7 +945,7 @@ func (a *emptyCompletionAccum) evalClaude(data []byte) bool { a.hasToolCalls = true } default: - if strings.TrimSpace(chunk.Delta.Text) != "" || strings.TrimSpace(chunk.Delta.Thinking) != "" || strings.TrimSpace(chunk.Delta.Signature) != "" || nonEmptyJSONPayload(chunk.Delta.Citation) { + if strings.TrimSpace(chunk.Delta.Text) != "" || strings.TrimSpace(chunk.Delta.Thinking) != "" || nonEmptyJSONPayload(chunk.Delta.Citation) { a.hasContent = true } } @@ -1241,8 +1233,8 @@ func (a *emptyCompletionAccum) evalClaudeBlocks(blocks []claudeContentBlock) { a.hasToolCalls = true continue } - if b.Type == "thinking" || b.Type == "redacted_thinking" || b.Type == "reasoning" || 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) != "" { + if b.Type == "thinking" || b.Type == "redacted_thinking" || b.Type == "reasoning" || strings.TrimSpace(b.Thinking) != "" || strings.TrimSpace(b.Data) != "" { + if strings.TrimSpace(b.Thinking) != "" || strings.TrimSpace(b.Data) != "" { a.hasContent = true } continue @@ -1395,9 +1387,6 @@ func (a *emptyCompletionAccum) evalGemini(data []byte) bool { if strings.TrimSpace(part.Text) != "" { a.hasContent = true } - if strings.TrimSpace(part.ThoughtSignature) != "" || strings.TrimSpace(part.Thought_Signature) != "" { - a.hasContent = true - } } } } @@ -1528,11 +1517,6 @@ func (a *emptyCompletionAccum) evalInteractions(data []byte) bool { if chunk.Delta.Content != nil && chunk.Delta.Content.hasMeaningfulContent() { a.hasContent = true } - if strings.TrimSpace(chunk.Delta.Signature) != "" || - strings.TrimSpace(chunk.Delta.ThoughtSignature) != "" || - strings.TrimSpace(chunk.Delta.ThoughtSignatureCamel) != "" { - a.hasContent = true - } if strings.TrimSpace(chunk.Delta.Name) != "" || hasMeaningfulInteractionsArguments(chunk.Delta.Arguments) { a.hasToolCalls = true } @@ -1582,9 +1566,6 @@ func (a *emptyCompletionAccum) evalInteractionsSteps(steps []interactionsStep) { a.hasContent = true } } - if step.hasSignature() { - a.hasContent = true - } for _, content := range step.Content { if content.hasMeaningfulContent() { a.hasContent = true @@ -1708,6 +1689,9 @@ func (s *streamBootstrapState) processLine(line []byte) { line = bytes.TrimSpace(line) if len(line) == 0 { if len(s.dataLines) > 0 && classifyJSONBuffer(bytes.Join(s.dataLines, []byte("\n"))) == jsonBufIncomplete { + // Pretty-printed raw JSON may contain blank lines; keep them in the + // buffer and do not treat them as an SSE event boundary. + s.dataLines = append(s.dataLines, []byte("")) return } s.flushData() @@ -2193,6 +2177,10 @@ func (d *streamPayloadErrorDetector) Observe(chunk []byte) *Error { line := bytes.TrimSpace(d.pending[:newline]) d.pending = d.pending[newline+1:] if len(line) == 0 { + if len(d.dataLines) > 0 && classifyJSONBuffer(bytes.Join(d.dataLines, []byte("\n"))) == jsonBufIncomplete { + d.dataLines = append(d.dataLines, []byte("")) + continue + } d.flushData() if d.err != nil { return d.err diff --git a/sdk/cliproxy/auth/empty_completion_test.go b/sdk/cliproxy/auth/empty_completion_test.go index 84a619d6b68..8718ce59c08 100644 --- a/sdk/cliproxy/auth/empty_completion_test.go +++ b/sdk/cliproxy/auth/empty_completion_test.go @@ -1476,6 +1476,21 @@ func TestStreamPayloadErrorDetectorBuffersPrettyPrintedJSONFrame(t *testing.T) { } } +func TestStreamPayloadErrorDetectorBuffersPrettyPrintedJSONFrameWithBlankLine(t *testing.T) { + var d streamPayloadErrorDetector + frame := "{\n\n \"error\": {\n\n \"message\": \"quota exceeded\",\n \"code\": 429\n }\n}\n" + if streamErr := d.Observe([]byte(frame)); streamErr != nil && !strings.Contains(streamErr.Message, "quota exceeded") { + t.Fatalf("payload detector Observe error = %q", streamErr.Message) + } + streamErr := d.Finish() + if streamErr == nil { + t.Fatal("payload detector did not surface a provider error carried by a pretty-printed raw JSON frame with a blank line") + } + if !strings.Contains(streamErr.Message, "quota exceeded") { + t.Fatalf("payload detector stream error message = %q, want it to carry the provider message", streamErr.Message) + } +} + func TestStreamBootstrapDetector(t *testing.T) { var detector StreamBootstrapDetector if detector.Observe([]byte("data: {\"type\":\"response.created\",\"response\":{\"status\":\"in_progress\"}}\n\n")) { @@ -1630,6 +1645,22 @@ func TestExecuteEmptyCompletionRotatesAuth(t *testing.T) { assertRotatesToContent(t, ids, executor.firstExecute, string(resp.Payload), "real", capture) } +func TestExecuteSignatureOnlyRotatesAuth(t *testing.T) { + executor := &emptyCompletionTestExecutor{ + executePayloads: map[string][]byte{}, + executeCalls: map[string]int{}, + emptyPayload: []byte(`{"candidates":[{"content":{"role":"model","parts":[{"thoughtSignature":"sig_only"}]},"finishReason":"STOP"}]}`), + contentPayload: []byte(`{"candidates":[{"content":{"role":"model","parts":[{"text":"real"}]},"finishReason":"STOP"}]}`), + } + manager, ids, model, capture := newEmptyCompletionTestManager(t, executor) + + resp, err := manager.Execute(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}) + if err != nil { + t.Fatalf("Execute() error = %v", err) + } + assertRotatesToContent(t, ids, executor.firstExecute, string(resp.Payload), "real", capture) +} + func TestExecuteStreamEmptyCompletionRotatesAuth(t *testing.T) { executor := &emptyCompletionTestExecutor{ streamPayloads: map[string][][]byte{}, @@ -2173,15 +2204,22 @@ func assertRotatesToContent(t *testing.T, ids []string, emptyFirst, gotPayload, other = ids[1] } var emptyRecorded bool + var emptySucceeded bool var otherSucceeded bool for _, r := range capture.Results() { if r.AuthID == emptyFirst && !r.Success { emptyRecorded = true } + if r.AuthID == emptyFirst && r.Success { + emptySucceeded = true + } if r.AuthID == other && r.Success { otherSucceeded = true } } + if emptySucceeded { + t.Fatalf("empty auth %q was recorded as success before the empty-completion failure; results=%v", emptyFirst, capture.Results()) + } if !emptyRecorded { t.Fatalf("empty auth %q was not recorded as a failure result; results=%v", emptyFirst, capture.Results()) } @@ -3050,10 +3088,10 @@ func TestReadStreamBootstrapErrorHandling(t *testing.T) { } func TestClaudeSignatureDeltaEmptyCompletion(t *testing.T) { - t.Run("thinking content_block_start followed by signature_delta with signature is not empty", func(t *testing.T) { + t.Run("thinking content_block_start followed by signature_delta with signature is 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") + if !IsEmptyCompletionPayload(payload) { + t.Fatal("IsEmptyCompletionPayload() = false for thinking stream with only signature_delta and 0 tokens, want true") } }) @@ -3194,17 +3232,17 @@ func TestMultiValueJSONMixedUnknownEmptyCompletion(t *testing.T) { } func TestGeminiThoughtSignatureEmptyCompletion(t *testing.T) { - t.Run("gemini STOP with thoughtSignature and omitted candidatesTokenCount is not empty", func(t *testing.T) { + t.Run("gemini STOP with thoughtSignature and omitted candidatesTokenCount is 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") + if !IsEmptyCompletionPayload(payload) { + t.Fatal("IsEmptyCompletionPayload() = false for thoughtSignature-only with omitted token count, want true") } }) - t.Run("gemini STOP with thought_signature and omitted candidatesTokenCount is not empty", func(t *testing.T) { + t.Run("gemini STOP with thought_signature and omitted candidatesTokenCount is 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") + if !IsEmptyCompletionPayload(payload) { + t.Fatal("IsEmptyCompletionPayload() = false for thought_signature-only with omitted token count, want true") } }) @@ -5111,7 +5149,7 @@ func TestReadStreamBootstrapInteractionsNestedFailureSurfacesProviderError(t *te } } -func TestStreamBootstrapDetectorInteractionsSignatureOnlyIsNotEmpty(t *testing.T) { +func TestStreamBootstrapDetectorInteractionsSignatureOnlyIsEmpty(t *testing.T) { createdChunk := []byte("event: interaction.created\ndata: {\"event_type\":\"interaction.created\",\"interaction\":{\"id\":\"int_1\",\"object\":\"interaction\",\"status\":\"in_progress\"}}\n\n") completedChunk := []byte("event: interaction.completed\ndata: {\"event_type\":\"interaction.completed\",\"interaction\":{\"id\":\"int_1\",\"status\":\"completed\",\"usage\":{\"output_tokens\":0}}}\n\n") @@ -5156,11 +5194,11 @@ func TestStreamBootstrapDetectorInteractionsSignatureOnlyIsNotEmpty(t *testing.T detector.Observe([]byte(tc.payload)) detector.Observe(completedChunk) - if detector.IsTerminalEmpty() { - t.Fatalf("IsTerminalEmpty() = true for a signature-carrying interaction, want false; signature would be dropped and the turn retried: %s", tc.payload) + if !detector.IsTerminalEmpty() { + t.Fatalf("IsTerminalEmpty() = false for a signature-only interaction, want true: %s", tc.payload) } - if !detector.HasMeaningfulOutput() { - t.Fatalf("HasMeaningfulOutput() = false for a signature-carrying interaction, want true: %s", tc.payload) + if detector.HasMeaningfulOutput() { + t.Fatalf("HasMeaningfulOutput() = true for a signature-only interaction, want false: %s", tc.payload) } }) } diff --git a/sdk/cliproxy/auth/home_execution_paths_test.go b/sdk/cliproxy/auth/home_execution_paths_test.go index 151e95ccb72..a9ac5373fba 100644 --- a/sdk/cliproxy/auth/home_execution_paths_test.go +++ b/sdk/cliproxy/auth/home_execution_paths_test.go @@ -1642,3 +1642,59 @@ func TestHomePrepareFailureResultPreservesRequestMetadata(t *testing.T) { }) } } + +type emptyThenContentHomeExecutor struct { + mu sync.Mutex + authIDs []string +} + +func (*emptyThenContentHomeExecutor) Identifier() string { return "home-execution" } +func (e *emptyThenContentHomeExecutor) Execute(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + e.mu.Lock() + e.authIDs = append(e.authIDs, auth.ID) + n := len(e.authIDs) + e.mu.Unlock() + if n == 1 { + return cliproxyexecutor.Response{Payload: []byte(`{"choices":[{"message":{"content":""},"finish_reason":"stop"}],"usage":{"completion_tokens":0}}`)}, nil + } + return cliproxyexecutor.Response{Payload: []byte(`{"choices":[{"message":{"content":"ok"},"finish_reason":"stop"}]}`)}, nil +} +func (*emptyThenContentHomeExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + return nil, nil +} +func (*emptyThenContentHomeExecutor) Refresh(context.Context, *Auth) (*Auth, error) { return nil, nil } +func (*emptyThenContentHomeExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} +func (*emptyThenContentHomeExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} +func (e *emptyThenContentHomeExecutor) AuthIDs() []string { + e.mu.Lock() + defer e.mu.Unlock() + return append([]string(nil), e.authIDs...) +} + +func TestHomeExecuteEmptyCompletionRotatesAuth(t *testing.T) { + dispatcher := &freshHomeStreamSelectionDispatcher{} + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetRetryConfig(0, time.Second, 2) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + executor := &emptyThenContentHomeExecutor{} + manager.RegisterExecutor(executor) + + resp, errExecute := manager.Execute(context.Background(), []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{}) + if errExecute != nil { + t.Fatalf("Execute() error = %v", errExecute) + } + if !strings.Contains(string(resp.Payload), "ok") { + t.Fatalf("payload = %q, want content from the second Home auth", resp.Payload) + } + if got := executor.AuthIDs(); len(got) != 2 || got[0] != "home-auth-a" || got[1] != "home-auth-b" { + t.Fatalf("executor auth IDs = %v, want [home-auth-a home-auth-b]", got) + } + if got := dispatcher.calls.Load(); got != 2 { + t.Fatalf("Home RPOP calls = %d, want 2", got) + } +} From 22deb3eb2394eaced4b37b26960f2b0ec048109c Mon Sep 17 00:00:00 2001 From: warelik Date: Sun, 23 Aug 2026 15:44:15 +0300 Subject: [PATCH 3/9] fix(auth): stop stream TTFT at connection, validate count emptiness, redact quoted secret JSON --- sdk/cliproxy/auth/conductor_execution.go | 2 +- sdk/cliproxy/auth/conductor_home_execution.go | 12 +- sdk/cliproxy/auth/conductor_overrides_test.go | 48 +++++++ sdk/cliproxy/auth/conductor_stream.go | 77 +++++++----- sdk/cliproxy/auth/empty_completion.go | 10 ++ .../auth/home_execution_paths_test.go | 63 +++++++++- sdk/cliproxy/auth/stream_ttft_test.go | 118 ++++++++++++++++++ 7 files changed, 296 insertions(+), 34 deletions(-) diff --git a/sdk/cliproxy/auth/conductor_execution.go b/sdk/cliproxy/auth/conductor_execution.go index 773e5aba2f1..f128ce3f3a5 100644 --- a/sdk/cliproxy/auth/conductor_execution.go +++ b/sdk/cliproxy/auth/conductor_execution.go @@ -1723,7 +1723,7 @@ func formatAuthIdentity(auth *Auth, provider string) string { var ( logSecretAPIKeyPattern = regexp.MustCompile(`(?i)\bsk-[A-Za-z0-9_-]{8,}`) logSecretBearerPattern = regexp.MustCompile(`(?i)\b(Bearer|Basic)\s+[A-Za-z0-9._~+/=-]{4,}`) - logSecretLabeledPattern = regexp.MustCompile(`(?i)((?:api[_-]?key|access[_-]?token|secret)\s*[=:]\s*)([^\s"&,;}]+)`) + logSecretLabeledPattern = regexp.MustCompile(`(?i)((?:"?(?:api[_-]?key|access[_-]?token|token|authorization|secret)"?)\s*[=:]\s*"?)([^\s"&,;}]+)`) ) func redactSecretsForLog(msg string) string { diff --git a/sdk/cliproxy/auth/conductor_home_execution.go b/sdk/cliproxy/auth/conductor_home_execution.go index 37035f05f78..ec8b1e6082e 100644 --- a/sdk/cliproxy/auth/conductor_home_execution.go +++ b/sdk/cliproxy/auth/conductor_home_execution.go @@ -278,12 +278,16 @@ func (m *Manager) executeHomeOnce(ctx context.Context, providers []string, req c } result := Result{AuthID: preparedAuth.ID, Provider: selection.Provider, Model: resultModel, Success: errExecute == nil, Options: execOpts} if errExecute == nil { - if !countTokens && isEmptyCompletionPayload(response.Payload) { + if isEmptyCompletionPayload(response.Payload) { result.Success = false - result.Error = errEmptyCompletion + if countTokens { + result.Error = errEmptyCount + } else { + result.Error = errEmptyCompletion + } m.reportHomeResult(execCtx, result, preparedAuth) - tracker.Record(preparedAuth, errEmptyCompletion) - lastErr = errEmptyCompletion + tracker.Record(preparedAuth, result.Error) + lastErr = result.Error break } m.reportHomeResult(execCtx, result, preparedAuth) diff --git a/sdk/cliproxy/auth/conductor_overrides_test.go b/sdk/cliproxy/auth/conductor_overrides_test.go index e9583700085..2615fb5dd78 100644 --- a/sdk/cliproxy/auth/conductor_overrides_test.go +++ b/sdk/cliproxy/auth/conductor_overrides_test.go @@ -2709,6 +2709,54 @@ func TestSummarizeErrorForLogRedactsSecrets(t *testing.T) { } } +func TestRedactSecretsForLog_QuotedJSON(t *testing.T) { + tests := []struct { + name string + in string + leaks []string + }{ + { + name: "apiKey with non-sk secret", + in: `{"apiKey":"AIza-secret"}`, + leaks: []string{"AIza-secret"}, + }, + { + name: "token key", + in: `{"token":"foo"}`, + leaks: []string{"foo"}, + }, + { + name: "authorization key with quoted value", + in: `{"authorization":"AIza-secret"}`, + leaks: []string{"AIza-secret"}, + }, + { + name: "api-key with unquoted key", + in: `api-key:AIza-secret`, + leaks: []string{"AIza-secret"}, + }, + { + name: "sk prefix still redacted", + in: `sk-live-secret`, + leaks: []string{"sk-live-secret"}, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + got := redactSecretsForLog(tc.in) + for _, leak := range tc.leaks { + if strings.Contains(got, leak) { + t.Fatalf("redacted = %q, contains %q", got, leak) + } + } + if !strings.Contains(got, "[REDACTED]") { + t.Fatalf("redacted = %q, want [REDACTED]", got) + } + }) + } +} + func TestManagerExecute_InvalidAPIKeyStopsModelPoolRetries(t *testing.T) { m := NewManager(nil, nil, nil) cfg := &internalconfig.Config{ diff --git a/sdk/cliproxy/auth/conductor_stream.go b/sdk/cliproxy/auth/conductor_stream.go index a1f14ad2ecc..55b0527a84f 100644 --- a/sdk/cliproxy/auth/conductor_stream.go +++ b/sdk/cliproxy/auth/conductor_stream.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "net/http" + "net/http/httptrace" "strings" "sync" "sync/atomic" @@ -323,15 +324,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi var timedOut atomic.Bool var attemptMu sync.Mutex var attemptSeq uint64 - - stopTTFT := func() { - if timer != nil { - timer.Stop() - } - attemptMu.Lock() - attemptSeq++ - attemptMu.Unlock() - } + var stopTTFT func() checkTTFTErr := func(err error) error { if timedOut.Load() { @@ -340,6 +333,53 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi return err } + // Arm the TTFT timer 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. The timer is stopped as soon as the + // HTTP transport reports a connection is obtained (GotConn), which is + // before response headers for standard net/http clients. + armTTFT := func() { + attemptMu.Lock() + if timer != nil { + timer.Stop() + timer = nil + } + var once sync.Once + stopTTFT = func() { + once.Do(func() { + attemptMu.Lock() + if timer != nil { + timer.Stop() + timer = nil + } + attemptSeq++ + attemptMu.Unlock() + }) + } + if ttftTimeout > 0 { + currentSeq := attemptSeq + currentCancel := cancelAttempt + timer = time.AfterFunc(ttftTimeout, func() { + attemptMu.Lock() + defer attemptMu.Unlock() + if currentSeq != attemptSeq { + return + } + timedOut.Store(true) + currentCancel() + }) + trace := &httptrace.ClientTrace{ + GotConn: func(httptrace.GotConnInfo) { + stopTTFT() + }, + } + attemptCtx = httptrace.WithClientTrace(attemptCtx, trace) + } + attemptMu.Unlock() + } + stopTTFT = func() {} + resultModel := m.stateModelForExecution(auth, routeModel, execModel, pooled) execReq := req execReq.Model = execModel @@ -362,25 +402,6 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi cancelAttempt() return nil, errCtx } - // Arm the TTFT timer 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. - armTTFT := func() { - if ttftTimeout > 0 { - currentSeq := attemptSeq - currentCancel := cancelAttempt - timer = time.AfterFunc(ttftTimeout, func() { - attemptMu.Lock() - defer attemptMu.Unlock() - if currentSeq != attemptSeq { - return - } - timedOut.Store(true) - currentCancel() - }) - } - } armTTFT() // The unauthorized-refresh retries below re-execute behind a credential // refresh, which may consume the whole TTFT budget (or fire the timer diff --git a/sdk/cliproxy/auth/empty_completion.go b/sdk/cliproxy/auth/empty_completion.go index 56bd71e202c..d94d252133a 100644 --- a/sdk/cliproxy/auth/empty_completion.go +++ b/sdk/cliproxy/auth/empty_completion.go @@ -70,6 +70,16 @@ var errEmptyCompletion = &Error{ HTTPStatus: http.StatusServiceUnavailable, } +// errEmptyCount indicates the upstream returned an empty count response. It is +// retriable so the conductor marks the auth as failed, cools it down, and +// rotates to the next auth/model. +var errEmptyCount = &Error{ + Code: "empty_count", + Message: "upstream returned an empty count response", + Retryable: true, + HTTPStatus: http.StatusServiceUnavailable, +} + // maxStreamBootstrapBytes bounds how much metadata a stream can accumulate // before the conductor conservatively forwards it. Empty-completion detection // must never create an unbounded pre-output buffer. diff --git a/sdk/cliproxy/auth/home_execution_paths_test.go b/sdk/cliproxy/auth/home_execution_paths_test.go index a9ac5373fba..7814c0e7aaa 100644 --- a/sdk/cliproxy/auth/home_execution_paths_test.go +++ b/sdk/cliproxy/auth/home_execution_paths_test.go @@ -52,7 +52,7 @@ func (*homeExecutionExecutor) ExecuteStream(context.Context, *Auth, cliproxyexec } func (*homeExecutionExecutor) Refresh(context.Context, *Auth) (*Auth, error) { return nil, nil } func (*homeExecutionExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { - return cliproxyexecutor.Response{}, nil + return cliproxyexecutor.Response{Payload: []byte(`{"input_tokens":1}`)}, nil } func (*homeExecutionExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { return nil, nil @@ -1698,3 +1698,64 @@ func TestHomeExecuteEmptyCompletionRotatesAuth(t *testing.T) { t.Fatalf("Home RPOP calls = %d, want 2", got) } } + +// emptyThenContentHomeCountExecutor returns an empty/whitespace count response +// on the first Home auth and a non-empty count on the second, mirroring +// emptyThenContentHomeExecutor for CountTokens. +type emptyThenContentHomeCountExecutor struct { + mu sync.Mutex + authIDs []string +} + +func (*emptyThenContentHomeCountExecutor) Identifier() string { return "home-execution" } +func (*emptyThenContentHomeCountExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} +func (*emptyThenContentHomeCountExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + return nil, nil +} +func (*emptyThenContentHomeCountExecutor) Refresh(context.Context, *Auth) (*Auth, error) { + return nil, nil +} +func (e *emptyThenContentHomeCountExecutor) CountTokens(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + e.mu.Lock() + e.authIDs = append(e.authIDs, auth.ID) + n := len(e.authIDs) + e.mu.Unlock() + if n == 1 { + return cliproxyexecutor.Response{Payload: []byte(` `)}, nil + } + return cliproxyexecutor.Response{Payload: []byte(`{"input_tokens":1}`)}, nil +} +func (*emptyThenContentHomeCountExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} +func (e *emptyThenContentHomeCountExecutor) AuthIDs() []string { + e.mu.Lock() + defer e.mu.Unlock() + return append([]string(nil), e.authIDs...) +} + +func TestHomeCountEmptyBodyRotatesAuth(t *testing.T) { + dispatcher := &freshHomeStreamSelectionDispatcher{} + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetRetryConfig(0, time.Second, 2) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + executor := &emptyThenContentHomeCountExecutor{} + manager.RegisterExecutor(executor) + + resp, errCount := manager.ExecuteCount(context.Background(), []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{}) + if errCount != nil { + t.Fatalf("ExecuteCount() error = %v", errCount) + } + if !strings.Contains(string(resp.Payload), "input_tokens") { + t.Fatalf("payload = %q, want count from the second Home auth", resp.Payload) + } + if got := executor.AuthIDs(); len(got) != 2 || got[0] != "home-auth-a" || got[1] != "home-auth-b" { + t.Fatalf("executor auth IDs = %v, want [home-auth-a home-auth-b]", got) + } + if got := dispatcher.calls.Load(); got != 2 { + t.Fatalf("Home RPOP calls = %d, want 2", got) + } +} diff --git a/sdk/cliproxy/auth/stream_ttft_test.go b/sdk/cliproxy/auth/stream_ttft_test.go index 08285c6c4d2..81373769697 100644 --- a/sdk/cliproxy/auth/stream_ttft_test.go +++ b/sdk/cliproxy/auth/stream_ttft_test.go @@ -2,7 +2,10 @@ package auth import ( "context" + "fmt" + "io" "net/http" + "net/http/httptest" "strings" "sync" "testing" @@ -347,3 +350,118 @@ func TestStreamConnectTimeout_ConfigAndMetadata(t *testing.T) { t.Fatalf("stream_connect_timeout_ms metadata precedence = %v, want 90ms", got) } } + +// httpTTFTStreamExecutor makes an HTTP request using the default client so the +// conductor's httptrace.ClientTrace fires GotConn as soon as the connection is +// obtained. +type httpTTFTStreamExecutor struct { + baseURL string +} + +func (*httpTTFTStreamExecutor) Identifier() string { return "http-ttft" } + +func (*httpTTFTStreamExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusNotImplemented, Message: "not implemented"} +} + +func (e *httpTTFTStreamExecutor) ExecuteStream(ctx context.Context, _ *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + httpClient := &http.Client{} + httpReq, err := http.NewRequestWithContext(ctx, http.MethodGet, e.baseURL, nil) + if err != nil { + return nil, err + } + httpReq.Header.Set("Accept", "text/event-stream") + + resp, err := httpClient.Do(httpReq) + if err != nil { + return nil, err + } + defer resp.Body.Close() + + body, err := io.ReadAll(resp.Body) + if err != nil { + return nil, err + } + + ch := make(chan cliproxyexecutor.StreamChunk, 1) + ch <- cliproxyexecutor.StreamChunk{Payload: body} + close(ch) + return &cliproxyexecutor.StreamResult{Chunks: ch, Headers: resp.Header}, nil +} + +func (*httpTTFTStreamExecutor) Refresh(context.Context, *Auth) (*Auth, error) { return nil, nil } + +func (*httpTTFTStreamExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{Payload: []byte("ok")}, nil +} + +func (*httpTTFTStreamExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +// TestStreamFirstChunkTimeout_StoppedAtConnection verifies that the TTFT timer +// is stopped when the HTTP transport reports a connection, before response +// headers are written. A slow-responding server that waits longer than the TTFT +// window should not trigger a timeout. +func TestStreamFirstChunkTimeout_StoppedAtConnection(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + time.Sleep(150 * time.Millisecond) + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + fmt.Fprint(w, "data: {\"type\":\"content_block_delta\",\"delta\":{\"text\":\"hello\"}}\n\n") + })) + defer server.Close() + + const model = "ttft-http-model" + auth := &Auth{ID: "http-ttft-auth", Provider: "http-ttft", Status: StatusActive} + + exec := &httpTTFTStreamExecutor{baseURL: server.URL} + m := NewManager(nil, nil, nil) + m.RegisterExecutor(exec) + + reg := registry.GetGlobalRegistry() + reg.RegisterClient(auth.ID, "http-ttft", []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { reg.UnregisterClient(auth.ID) }) + + if _, err := m.Register(context.Background(), auth); err != nil { + t.Fatalf("register auth: %v", err) + } + + opts := cliproxyexecutor.Options{ + Metadata: map[string]any{ + "stream_first_chunk_timeout_ms": 50, + }, + } + + start := time.Now() + stream, errStream := m.ExecuteStream(context.Background(), []string{"http-ttft"}, cliproxyexecutor.Request{Model: model}, opts) + elapsed := time.Since(start) + if errStream != nil { + t.Fatalf("ExecuteStream error = %v, want stream after delayed headers", errStream) + } + if elapsed < 100*time.Millisecond { + t.Fatalf("stream returned too quickly (%v), want at least 100ms of delayed headers", elapsed) + } + + done := make(chan struct{}) + var got string + go func() { + defer close(done) + for chunk := range stream.Chunks { + if chunk.Err != nil { + t.Errorf("chunk error = %v", chunk.Err) + return + } + got = string(chunk.Payload) + } + }() + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("timed out reading stream chunks") + } + + if !strings.Contains(got, "hello") { + t.Fatalf("stream payload does not contain expected content: %q", got) + } +} From 89e87a859f1ff6268c8477dfdadc6211388db88d Mon Sep 17 00:00:00 2001 From: warelik Date: Sun, 23 Aug 2026 16:59:45 +0300 Subject: [PATCH 4/9] fix(auth): pass request-local excluded auths to retry decision --- .../conductor_cooldown_retry_reset_test.go | 20 +++++++++++++++++++ sdk/cliproxy/auth/conductor_execution.go | 6 +++--- sdk/cliproxy/auth/conductor_home_execution.go | 2 +- sdk/cliproxy/auth/conductor_overrides_test.go | 2 +- sdk/cliproxy/auth/conductor_selection.go | 15 +++++++++----- sdk/cliproxy/auth/home_retry_contract_test.go | 12 +++++------ 6 files changed, 41 insertions(+), 16 deletions(-) diff --git a/sdk/cliproxy/auth/conductor_cooldown_retry_reset_test.go b/sdk/cliproxy/auth/conductor_cooldown_retry_reset_test.go index c130aa35ae2..3c1e97c487b 100644 --- a/sdk/cliproxy/auth/conductor_cooldown_retry_reset_test.go +++ b/sdk/cliproxy/auth/conductor_cooldown_retry_reset_test.go @@ -216,6 +216,26 @@ func TestCooldownRetryPreservesCallerExclusions(t *testing.T) { } } +func TestShouldRetrySkipsWaitWhenSingleAuthIsExcluded(t *testing.T) { + manager := NewManager(nil, nil, nil) + manager.SetRetryConfig(3, 5*time.Second, 0) + auth := &Auth{ID: "auth-excluded", 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) + + excluded := map[string]struct{}{auth.ID: {}} + err := &retryableRateLimitError{status: http.StatusTooManyRequests, retryAfter: 50 * time.Millisecond} + wait, shouldRetry := manager.shouldRetryAfterErrorWithHomeRetryLimit(context.Background(), cliproxyexecutor.Options{}, err, 0, []string{"gemini"}, "test-model", 5*time.Second, -1, 3, excluded) + if shouldRetry || wait != 0 { + t.Fatalf("shouldRetryAfterErrorWithHomeRetryLimit() = (%v, %t), want (0, false) when the only auth is excluded", wait, shouldRetry) + } +} + func TestCooldownRetryPreservesConfigDisabledCoolingExclusions(t *testing.T) { t.Run("global config disable cooling retains exclusion on retry", func(t *testing.T) { manager := NewManager(nil, nil, nil) diff --git a/sdk/cliproxy/auth/conductor_execution.go b/sdk/cliproxy/auth/conductor_execution.go index f128ce3f3a5..cea021c4a89 100644 --- a/sdk/cliproxy/auth/conductor_execution.go +++ b/sdk/cliproxy/auth/conductor_execution.go @@ -68,7 +68,7 @@ func (m *Manager) Execute(ctx context.Context, providers []string, req cliproxye return cliproxyexecutor.Response{}, wrapRouteExhaustion(unwrapRequestStopError(errExec), tracker) } lastErr = errExec - wait, shouldRetry := m.shouldRetryAfterErrorWithHomeRetryLimit(ctx, opts, errExec, attempt, normalized, retryModel, maxWait, -1, defaultRequestRetry) + wait, shouldRetry := m.shouldRetryAfterErrorWithHomeRetryLimit(ctx, opts, errExec, attempt, normalized, retryModel, maxWait, -1, defaultRequestRetry, tried) if !shouldRetry { break } @@ -120,7 +120,7 @@ func (m *Manager) ExecuteCount(ctx context.Context, providers []string, req clip return cliproxyexecutor.Response{}, wrapRouteExhaustion(unwrapRequestStopError(errExec), tracker) } lastErr = errExec - wait, shouldRetry := m.shouldRetryAfterErrorWithHomeRetryLimit(ctx, opts, errExec, attempt, normalized, retryModel, maxWait, -1, defaultRequestRetry) + wait, shouldRetry := m.shouldRetryAfterErrorWithHomeRetryLimit(ctx, opts, errExec, attempt, normalized, retryModel, maxWait, -1, defaultRequestRetry, tried) if !shouldRetry { break } @@ -184,7 +184,7 @@ func (m *Manager) ExecuteStream(ctx context.Context, providers []string, req cli return nil, wrapRouteExhaustion(unwrapRequestStopError(errStream), tracker) } lastErr = errStream - wait, shouldRetry := m.shouldRetryAfterErrorWithHomeRetryLimit(ctx, opts, errStream, attempt, normalized, retryModel, maxWait, homeRetryLimit, defaultRequestRetry) + wait, shouldRetry := m.shouldRetryAfterErrorWithHomeRetryLimit(ctx, opts, errStream, attempt, normalized, retryModel, maxWait, homeRetryLimit, defaultRequestRetry, tried) if !shouldRetry { break } diff --git a/sdk/cliproxy/auth/conductor_home_execution.go b/sdk/cliproxy/auth/conductor_home_execution.go index ec8b1e6082e..07a1fa2fe7b 100644 --- a/sdk/cliproxy/auth/conductor_home_execution.go +++ b/sdk/cliproxy/auth/conductor_home_execution.go @@ -71,7 +71,7 @@ func (m *Manager) executeHome(ctx context.Context, providers []string, req clipr if isRequestTerminatedError(errExecute) || isRequestStopError(errExecute) { return cliproxyexecutor.Response{}, unwrapRequestStopError(errExecute) } - wait, shouldRetry := m.shouldRetryAfterErrorWithHomeRetryLimit(ctx, opts, errExecute, attempt, providers, retryModel, maxWait, homeRetryLimit, defaultRequestRetry) + wait, shouldRetry := m.shouldRetryAfterErrorWithHomeRetryLimit(ctx, opts, errExecute, attempt, providers, retryModel, maxWait, homeRetryLimit, defaultRequestRetry, nil) if !shouldRetry { return cliproxyexecutor.Response{}, unwrapRequestStopError(errExecute) } diff --git a/sdk/cliproxy/auth/conductor_overrides_test.go b/sdk/cliproxy/auth/conductor_overrides_test.go index 2615fb5dd78..3788faea399 100644 --- a/sdk/cliproxy/auth/conductor_overrides_test.go +++ b/sdk/cliproxy/auth/conductor_overrides_test.go @@ -272,7 +272,7 @@ func TestManager_ShouldRetryAfterError_IgnoresRequestIneligibleOverrides(t *test } } - wait, shouldRetry := manager.shouldRetryAfterErrorWithHomeRetryLimit(test.ctx, test.opts, &Error{HTTPStatus: http.StatusBadGateway, Message: "retryable failure"}, 0, []string{"codex"}, model, 0, -1, 0) + wait, shouldRetry := manager.shouldRetryAfterErrorWithHomeRetryLimit(test.ctx, test.opts, &Error{HTTPStatus: http.StatusBadGateway, Message: "retryable failure"}, 0, []string{"codex"}, model, 0, -1, 0, nil) if shouldRetry || wait != 0 { t.Fatalf("request-ineligible override retry = (%v, %t), want (0, false)", wait, shouldRetry) } diff --git a/sdk/cliproxy/auth/conductor_selection.go b/sdk/cliproxy/auth/conductor_selection.go index e0ed812d49b..22c2245e07e 100644 --- a/sdk/cliproxy/auth/conductor_selection.go +++ b/sdk/cliproxy/auth/conductor_selection.go @@ -918,7 +918,7 @@ func (m *Manager) retryAllowed(attempt int, providers []string, model string, el func (m *Manager) shouldRetryAfterError(err error, attempt int, providers []string, model string, maxWait time.Duration) (time.Duration, bool) { defaultRequestRetry, _, _ := m.retrySettings() - return m.shouldRetryAfterErrorWithHomeRetryLimit(context.Background(), cliproxyexecutor.Options{}, err, attempt, providers, model, maxWait, -1, defaultRequestRetry) + return m.shouldRetryAfterErrorWithHomeRetryLimit(context.Background(), cliproxyexecutor.Options{}, err, attempt, providers, model, maxWait, -1, defaultRequestRetry, nil) } // maxWait limits only positive cooldown waits between credential retry rounds. @@ -926,7 +926,7 @@ func (m *Manager) shouldRetryAfterError(err error, attempt int, providers []stri // credential failover or an additional round that request-retry permits to start // immediately. If every eligible credential still needs a positive cooldown, // retry stops without waiting. -func (m *Manager) shouldRetryAfterErrorWithHomeRetryLimit(ctx context.Context, opts cliproxyexecutor.Options, err error, attempt int, providers []string, model string, maxWait time.Duration, homeRetryLimit int, defaultRequestRetry int) (time.Duration, bool) { +func (m *Manager) shouldRetryAfterErrorWithHomeRetryLimit(ctx context.Context, opts cliproxyexecutor.Options, err error, attempt int, providers []string, model string, maxWait time.Duration, homeRetryLimit int, defaultRequestRetry int, excluded map[string]struct{}) (time.Duration, bool) { if err == nil { return 0, false } @@ -977,11 +977,16 @@ func (m *Manager) shouldRetryAfterErrorWithHomeRetryLimit(ctx context.Context, o } eligibility := authSelectionEligibilityForRequest(ctx, opts) pinnedAuthID := pinnedAuthIDFromMetadata(opts.Metadata) - excluded := extractExcludedAuthIDs(opts.Metadata) - if !isCredentialRetryRoundStatus(status) || !m.retryAllowed(attempt, providers, model, eligibility, pinnedAuthID, defaultRequestRetry, excluded) { + merged := extractExcludedAuthIDs(opts.Metadata) + for authID := range excluded { + if authID = strings.TrimSpace(authID); authID != "" { + merged[authID] = struct{}{} + } + } + if !isCredentialRetryRoundStatus(status) || !m.retryAllowed(attempt, providers, model, eligibility, pinnedAuthID, defaultRequestRetry, merged) { return 0, false } - wait, found := m.closestCooldownWait(providers, model, attempt, eligibility, pinnedAuthID, defaultRequestRetry, excluded) + wait, found := m.closestCooldownWait(providers, model, attempt, eligibility, pinnedAuthID, defaultRequestRetry, merged) if found { if wait > 0 && (maxWait <= 0 || wait > maxWait) { return 0, false diff --git a/sdk/cliproxy/auth/home_retry_contract_test.go b/sdk/cliproxy/auth/home_retry_contract_test.go index 627f6362acd..8ffe3f5ba5c 100644 --- a/sdk/cliproxy/auth/home_retry_contract_test.go +++ b/sdk/cliproxy/auth/home_retry_contract_test.go @@ -563,23 +563,23 @@ func TestHomeRetryPolicyUsesRemoteCredentialOverrideBeforeSelection(t *testing.T hasRequestRetry: true, } - wait, shouldRetry := manager.shouldRetryAfterErrorWithHomeRetryLimit(context.Background(), cliproxyexecutor.Options{}, errRemoteCooldown, 0, []string{"home-retry-contract"}, "gpt", time.Second, -1, 0) + wait, shouldRetry := manager.shouldRetryAfterErrorWithHomeRetryLimit(context.Background(), cliproxyexecutor.Options{}, errRemoteCooldown, 0, []string{"home-retry-contract"}, "gpt", time.Second, -1, 0, nil) if !shouldRetry || wait != 10*time.Millisecond { t.Fatalf("remote credential override retry = (%v, %t), want (10ms, true)", wait, shouldRetry) } - if _, shouldRetry = manager.shouldRetryAfterErrorWithHomeRetryLimit(context.Background(), cliproxyexecutor.Options{}, errRemoteCooldown, 1, []string{"home-retry-contract"}, "gpt", time.Second, -1, 0); shouldRetry { + if _, shouldRetry = manager.shouldRetryAfterErrorWithHomeRetryLimit(context.Background(), cliproxyexecutor.Options{}, errRemoteCooldown, 1, []string{"home-retry-contract"}, "gpt", time.Second, -1, 0, nil); shouldRetry { t.Fatal("remote credential override allowed more than one additional round") } pinnedOpts := cliproxyexecutor.Options{Metadata: map[string]any{ cliproxyexecutor.PinnedAuthMetadataKey: "home-retry-a", }} - if _, shouldRetry = manager.shouldRetryAfterErrorWithHomeRetryLimit(context.Background(), pinnedOpts, errRemoteCooldown, 0, []string{"home-retry-contract"}, "gpt", time.Second, -1, 0); shouldRetry { + if _, shouldRetry = manager.shouldRetryAfterErrorWithHomeRetryLimit(context.Background(), pinnedOpts, errRemoteCooldown, 0, []string{"home-retry-contract"}, "gpt", time.Second, -1, 0, nil); shouldRetry { t.Fatal("aggregate retry limit from unpinned Home credentials affected a pinned request") } errRemoteCooldown.requestRetry = 0 manager.SetRetryConfig(3, time.Second, 0) - if _, shouldRetry = manager.shouldRetryAfterErrorWithHomeRetryLimit(context.Background(), cliproxyexecutor.Options{}, errRemoteCooldown, 0, []string{"home-retry-contract"}, "gpt", time.Second, -1, 0); shouldRetry { + if _, shouldRetry = manager.shouldRetryAfterErrorWithHomeRetryLimit(context.Background(), cliproxyexecutor.Options{}, errRemoteCooldown, 0, []string{"home-retry-contract"}, "gpt", time.Second, -1, 0, nil); shouldRetry { t.Fatal("explicit remote credential override 0 did not suppress the global retry setting") } retryLimit := 3 @@ -860,7 +860,7 @@ func TestHomeCooldownClassificationPreservesNonRetryableRoundStatus(t *testing.T if retryLimit != 2 { t.Fatalf("observed retry limit = %d, want authoritative Home limit 2", retryLimit) } - if wait, shouldRetry := manager.shouldRetryAfterErrorWithHomeRetryLimit(context.Background(), cliproxyexecutor.Options{}, errExecute, 0, []string{"home-retry-contract"}, "gpt", time.Second, retryLimit, 0); shouldRetry || wait != 0 { + if wait, shouldRetry := manager.shouldRetryAfterErrorWithHomeRetryLimit(context.Background(), cliproxyexecutor.Options{}, errExecute, 0, []string{"home-retry-contract"}, "gpt", time.Second, retryLimit, 0, nil); shouldRetry || wait != 0 { t.Fatalf("401 round retry = (%v, %t), want (0, false)", wait, shouldRetry) } }) @@ -911,7 +911,7 @@ func TestHomeRetryRoundStartsImmediatelyWhenHomeReportsAvailableNextRound(t *tes if !isHomeRetryRoundExhausted(errExecute) { t.Fatalf("execution error = %v, want exhausted retry round", errExecute) } - wait, shouldRetry := manager.shouldRetryAfterErrorWithHomeRetryLimit(context.Background(), cliproxyexecutor.Options{}, errExecute, 0, []string{"home-retry-contract"}, "gpt", 10*time.Second, retryLimit, 0) + wait, shouldRetry := manager.shouldRetryAfterErrorWithHomeRetryLimit(context.Background(), cliproxyexecutor.Options{}, errExecute, 0, []string{"home-retry-contract"}, "gpt", 10*time.Second, retryLimit, 0, nil) if !shouldRetry || wait != 0 { t.Fatalf("next-round retry = (%v, %t), want immediate", wait, shouldRetry) } From bbccdfacafe8659c39c9e68d16f4080ff3c0b56e Mon Sep 17 00:00:00 2001 From: warelik Date: Sun, 23 Aug 2026 17:40:38 +0300 Subject: [PATCH 5/9] fix(auth): sanitize all *Error text fields at the upstream boundary --- internal/pluginhost/executor_route.go | 22 +++- .../pluginhost/executor_route_stream_test.go | 30 ++++-- internal/pluginhost/model_router_test.go | 75 +++++++++++++ sdk/cliproxy/auth/conductor_cooldown.go | 37 ++++++- .../conductor_cooldown_retry_reset_test.go | 21 ++++ .../auth/conductor_cooldown_sanitize_test.go | 48 +++++++++ sdk/cliproxy/auth/conductor_stream.go | 14 ++- .../auth/conductor_stream_eof_test.go | 34 ++++++ sdk/cliproxy/auth/conductor_stream_test.go | 100 ++++++++++++++++++ sdk/cliproxy/auth/empty_completion.go | 13 ++- .../auth/empty_completion_sanitize_test.go | 45 ++++++++ sdk/cliproxy/auth/empty_completion_test.go | 26 ++++- sdk/cliproxy/auth/stream_ttft_test.go | 51 +++------ 13 files changed, 455 insertions(+), 61 deletions(-) create mode 100644 sdk/cliproxy/auth/conductor_cooldown_sanitize_test.go create mode 100644 sdk/cliproxy/auth/conductor_stream_test.go create mode 100644 sdk/cliproxy/auth/empty_completion_sanitize_test.go diff --git a/internal/pluginhost/executor_route.go b/internal/pluginhost/executor_route.go index e7fabc0973b..9f951c0a469 100644 --- a/internal/pluginhost/executor_route.go +++ b/internal/pluginhost/executor_route.go @@ -118,7 +118,18 @@ func (h *Host) ExecutePluginExecutorStream(ctx context.Context, pluginID string, // appears or the stream closes; unrecognized streams remain pass-through. func wrapStreamEmptyCompletion(ctx context.Context, streamResult *coreexecutor.StreamResult, requestPayloads ...[]byte) *coreexecutor.StreamResult { if streamResult == nil || streamResult.Chunks == nil { - return streamResult + errChunks := make(chan coreexecutor.StreamChunk, 1) + errChunks <- coreexecutor.StreamChunk{Err: &coreauth.Error{ + Code: "empty_stream", + Message: "upstream stream has no source", + Retryable: true, + }} + close(errChunks) + wrapped := &coreexecutor.StreamResult{Chunks: errChunks} + if streamResult != nil { + wrapped.Headers = streamResult.Headers + } + return wrapped } if ctx == nil { ctx = context.Background() @@ -266,7 +277,14 @@ func (h *Host) CountPluginExecutor(ctx context.Context, pluginID string, req cor if errAdapter != nil { return coreexecutor.Response{}, errAdapter } - return adapter.CountTokens(ctx, (*coreauth.Auth)(nil), req, opts) + resp, err := adapter.CountTokens(ctx, (*coreauth.Auth)(nil), req, opts) + if err != nil { + return coreexecutor.Response{}, err + } + if coreauth.IsEmptyCompletionPayload(resp.Payload) { + return coreexecutor.Response{}, coreauth.EmptyCompletionError() + } + return resp, nil } func (h *Host) executorAdapterForPlugin(pluginID string) (*executorAdapter, error) { diff --git a/internal/pluginhost/executor_route_stream_test.go b/internal/pluginhost/executor_route_stream_test.go index 20d5bf0ba56..1927f204412 100644 --- a/internal/pluginhost/executor_route_stream_test.go +++ b/internal/pluginhost/executor_route_stream_test.go @@ -195,14 +195,28 @@ func TestWrapStreamEmptyCompletionPreservesContentBeforeUpstreamError(t *testing } } -func TestWrapStreamEmptyCompletionPreservesNilResults(t *testing.T) { - if got := wrapStreamEmptyCompletion(context.Background(), nil); got != nil { - t.Fatalf("wrapStreamEmptyCompletion(nil) = %#v, want nil", got) - } - - result := &coreexecutor.StreamResult{Headers: http.Header{"X-Test": []string{"value"}}} - if got := wrapStreamEmptyCompletion(context.Background(), result); got != result { - t.Fatalf("wrapStreamEmptyCompletion(nil chunks) = %#v, want original result", got) +func TestWrapStreamEmptyCompletionRejectsNilSource(t *testing.T) { + for _, tc := range []struct { + name string + result *coreexecutor.StreamResult + }{ + {"nil result", nil}, + {"nil chunks", &coreexecutor.StreamResult{Headers: http.Header{"X-Test": []string{"value"}}}}, + } { + t.Run(tc.name, func(t *testing.T) { + got := wrapStreamEmptyCompletion(context.Background(), tc.result) + if got == nil || got.Chunks == nil { + t.Fatalf("wrapStreamEmptyCompletion(%s) = %#v, want stream with error chunk", tc.name, got) + } + chunk, ok := <-got.Chunks + if !ok || chunk.Err == nil { + t.Fatalf("wrapStreamEmptyCompletion(%s) emitted chunk %v, want error", tc.name, chunk) + } + var authErr *coreauth.Error + if !errors.As(chunk.Err, &authErr) || authErr.Code != "empty_stream" || !authErr.Retryable { + t.Fatalf("error = %v, want retriable empty_stream", chunk.Err) + } + }) } } diff --git a/internal/pluginhost/model_router_test.go b/internal/pluginhost/model_router_test.go index c523a4a45ca..9b90580c391 100644 --- a/internal/pluginhost/model_router_test.go +++ b/internal/pluginhost/model_router_test.go @@ -244,6 +244,81 @@ func TestHostExecutePluginExecutorStreamRejectsEmptyCompletion(t *testing.T) { } } +func TestHostCountPluginExecutorRejectsEmptyCompletion(t *testing.T) { + executor := &fakeExecutor{ + identifier: "plugin-provider", + countTokens: func(ctx context.Context, req pluginapi.ExecutorRequest) (pluginapi.ExecutorResponse, error) { + return pluginapi.ExecutorResponse{Payload: []byte(`{"choices":[{"message":{"content":""},"finish_reason":"stop"}],"usage":{"completion_tokens":0}}`)}, nil + }, + } + host := newRouteModelHostWithRecords(capabilityRecord{ + id: "executor", + plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{ + Executor: executor, + ExecutorInputFormats: []string{"openai"}, + ExecutorOutputFormats: []string{"openai"}, + }}, + }) + + _, errCount := host.CountPluginExecutor(context.Background(), "executor", coreexecutor.Request{Model: "client-model", Payload: []byte(`{"model":"client-model"}`)}, coreexecutor.Options{}) + if errCount == nil { + t.Fatal("CountPluginExecutor() with empty completion = nil, want retriable error") + } + var authErr *coreauth.Error + if !errors.As(errCount, &authErr) { + t.Fatalf("error = %v (%T), want *coreauth.Error", errCount, errCount) + } + if !authErr.Retryable || authErr.Code != "empty_completion" { + t.Fatalf("error = %+v, want retriable empty_completion", authErr) + } + if authErr.StatusCode() != http.StatusServiceUnavailable { + t.Fatalf("error status = %d, want %d", authErr.StatusCode(), http.StatusServiceUnavailable) + } +} + +func TestHostExecutePluginExecutorStreamRejectsNilChunks(t *testing.T) { + executor := &fakeExecutor{ + identifier: "plugin-provider", + executeStream: func(ctx context.Context, req pluginapi.ExecutorRequest) (pluginapi.ExecutorStreamResponse, error) { + return pluginapi.ExecutorStreamResponse{Chunks: nil}, nil + }, + } + host := newRouteModelHostWithRecords(capabilityRecord{ + id: "executor", + plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{ + Executor: executor, + ExecutorInputFormats: []string{"openai"}, + ExecutorOutputFormats: []string{"openai"}, + }}, + }) + + streamResult, errStream := host.ExecutePluginExecutorStream(context.Background(), "executor", coreexecutor.Request{Model: "client-model", Payload: []byte(`{"model":"client-model"}`)}, coreexecutor.Options{}) + if errStream != nil { + t.Fatalf("ExecutePluginExecutorStream() unexpected error = %v", errStream) + } + if streamResult == nil || streamResult.Chunks == nil { + t.Fatal("ExecutePluginExecutorStream() returned nil stream with no source") + } + + var emptyErr error + for chunk := range streamResult.Chunks { + if chunk.Err != nil { + emptyErr = chunk.Err + break + } + } + if emptyErr == nil { + t.Fatal("stream closed clean, want empty_stream error") + } + var authErr *coreauth.Error + if !errors.As(emptyErr, &authErr) { + t.Fatalf("error = %v (%T), want *coreauth.Error", emptyErr, emptyErr) + } + if authErr.Code != "empty_stream" || !authErr.Retryable { + t.Fatalf("error = %+v, want retriable empty_stream", authErr) + } +} + func TestHostRouteModelDefaultsHandledRouterToOwnExecutor(t *testing.T) { host := newRouteModelHostWithRecords(capabilityRecord{ id: "router", diff --git a/sdk/cliproxy/auth/conductor_cooldown.go b/sdk/cliproxy/auth/conductor_cooldown.go index 23696b7596b..05499dd5c3f 100644 --- a/sdk/cliproxy/auth/conductor_cooldown.go +++ b/sdk/cliproxy/auth/conductor_cooldown.go @@ -7,6 +7,7 @@ import ( "fmt" "io" "net/http" + "reflect" "sort" "strings" "sync/atomic" @@ -1410,7 +1411,39 @@ func resultErrorFromError(err error) *Error { resultErr.Code = connectionLifecycleErrorCode } } - return resultErr + // Do not persist or propagate credentials that may be echoed in an + // in-band stream error. sanitizeErrorTextFields redacts every string + // field of the cloned result while the classification above already used + // the original err. + return sanitizeErrorTextFields(resultErr).(*Error) +} + +// sanitizeErrorTextFields redacts secrets from every exported string field of the +// inner *Error. It mutates the value in place when err is an *Error (or wraps +// one), preserving the original error pointer so callers that compare identity +// still work. All untrusted upstream error parsing should return through this. +func sanitizeErrorTextFields(err error) error { + if err == nil { + return nil + } + var authErr *Error + if !errors.As(err, &authErr) || authErr == nil { + return err + } + v := reflect.ValueOf(authErr).Elem() + t := v.Type() + for i := 0; i < v.NumField(); i++ { + field := t.Field(i) + if !field.IsExported() { + continue + } + fv := v.Field(i) + if fv.Kind() == reflect.String { + s := fv.String() + fv.SetString(redactSecretsForLog(s)) + } + } + return err } // shouldSkipCredentialCooldown reports failures that must not mark auth/model cooling. @@ -1516,7 +1549,7 @@ func refreshErrorFromError(err error) *Error { authErr.Code = "unauthorized" authErr.Retryable = false } - return authErr + return sanitizeErrorTextFields(authErr).(*Error) } func retryAfterFromError(err error) *time.Duration { diff --git a/sdk/cliproxy/auth/conductor_cooldown_retry_reset_test.go b/sdk/cliproxy/auth/conductor_cooldown_retry_reset_test.go index 3c1e97c487b..418d34a071f 100644 --- a/sdk/cliproxy/auth/conductor_cooldown_retry_reset_test.go +++ b/sdk/cliproxy/auth/conductor_cooldown_retry_reset_test.go @@ -3,6 +3,7 @@ package auth import ( "context" "net/http" + "strings" "sync" "sync/atomic" "testing" @@ -334,3 +335,23 @@ func TestCooldownRetryPreservesConfigDisabledCoolingExclusions(t *testing.T) { } }) } + +func TestResultErrorFromErrorRedactsSecrets(t *testing.T) { + leaky := &Error{Code: "sk-xyz1234567890", Message: "Bearer sk-abc123def456", HTTPStatus: http.StatusUnauthorized} + result := resultErrorFromError(leaky) + if result == nil { + t.Fatal("resultErrorFromError returned nil") + } + if strings.Contains(result.Message, "sk-abc123def456") { + t.Fatalf("result error message = %q, want secret redacted", result.Message) + } + if !strings.Contains(result.Message, "REDACTED") { + t.Fatalf("result error message = %q, want [REDACTED] placeholder", result.Message) + } + if strings.Contains(result.Code, "sk-xyz1234567890") { + t.Fatalf("result error code = %q, want secret redacted", result.Code) + } + if !strings.Contains(result.Code, "REDACTED") { + t.Fatalf("result error code = %q, want [REDACTED] placeholder", result.Code) + } +} diff --git a/sdk/cliproxy/auth/conductor_cooldown_sanitize_test.go b/sdk/cliproxy/auth/conductor_cooldown_sanitize_test.go new file mode 100644 index 00000000000..70c6425242e --- /dev/null +++ b/sdk/cliproxy/auth/conductor_cooldown_sanitize_test.go @@ -0,0 +1,48 @@ +package auth + +import ( + "reflect" + "strings" + "testing" +) + +// TestSanitizeErrorTextFieldsCoversAllStringFields ensures every exported string +// field of *Error is redacted by sanitizeErrorTextFields. If a new string field +// is added and not handled, this test fails because that field still contains +// the injected secret. +func TestSanitizeErrorTextFieldsCoversAllStringFields(t *testing.T) { + const secret = "sk-live-1234567890" + + e := &Error{} + v := reflect.ValueOf(e).Elem() + stringFieldCount := 0 + for i := 0; i < v.NumField(); i++ { + field := v.Type().Field(i) + if !field.IsExported() || v.Field(i).Kind() != reflect.String { + continue + } + stringFieldCount++ + v.Field(i).SetString(secret) + } + if stringFieldCount == 0 { + t.Fatal("*Error has no string fields to sanitize") + } + + sanitizeErrorTextFields(e) + + sv := reflect.ValueOf(e).Elem() + for i := 0; i < sv.NumField(); i++ { + field := sv.Type().Field(i) + if !field.IsExported() || sv.Field(i).Kind() != reflect.String { + continue + } + got := sv.Field(i).String() + if strings.Contains(got, secret) { + t.Fatalf("field %q not sanitized: %q", field.Name, got) + } + // If the field was a secret-bearing value, redaction should leave a marker. + if got != "" && !strings.Contains(got, "REDACTED") { + t.Fatalf("field %q did not contain redaction marker: %q", field.Name, got) + } + } +} diff --git a/sdk/cliproxy/auth/conductor_stream.go b/sdk/cliproxy/auth/conductor_stream.go index 55b0527a84f..f8cd32ece22 100644 --- a/sdk/cliproxy/auth/conductor_stream.go +++ b/sdk/cliproxy/auth/conductor_stream.go @@ -206,6 +206,7 @@ func (m *Manager) wrapStreamResult(ctx context.Context, auth *Auth, provider, re emit := func(chunk cliproxyexecutor.StreamChunk) bool { if chunk.Err != nil && !failed { failed = true + chunk.Err = sanitizeErrorTextFields(chunk.Err) entry := logEntryWithRequestID(ctx) warnLogUpstreamFailure(ctx, entry, provider, resultModel, auth, time.Since(streamStart), chunk.Err) rerr := resultErrorFromError(chunk.Err) @@ -217,6 +218,7 @@ func (m *Manager) wrapStreamResult(ctx context.Context, auth *Auth, provider, re if !failed && len(chunk.Payload) > 0 { if streamErr := errorDetector.Observe(chunk.Payload); streamErr != nil { failed = true + streamErr = sanitizeErrorTextFields(streamErr).(*Error) entry := logEntryWithRequestID(ctx) warnLogUpstreamFailure(ctx, entry, provider, resultModel, auth, time.Since(streamStart), streamErr) rerr := resultErrorFromError(streamErr) @@ -283,6 +285,7 @@ func (m *Manager) wrapStreamResult(ctx context.Context, auth *Auth, provider, re if !failed { if streamErr := errorDetector.Finish(); streamErr != nil { failed = true + streamErr = sanitizeErrorTextFields(streamErr).(*Error) entry := logEntryWithRequestID(ctx) warnLogUpstreamFailure(ctx, entry, provider, resultModel, auth, time.Since(streamStart), streamErr) rerr := resultErrorFromError(streamErr) @@ -336,9 +339,10 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi // Arm the TTFT timer 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. The timer is stopped as soon as the - // HTTP transport reports a connection is obtained (GotConn), which is - // before response headers for standard net/http clients. + // upstream request was even made. The timer is stopped when the first + // response byte arrives (GotFirstResponseByte), after the actual upstream + // begins responding; this keeps a CONNECT tunnel or slow accept inside + // the budget instead of treating the proxy connection as success. armTTFT := func() { attemptMu.Lock() if timer != nil { @@ -370,7 +374,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi currentCancel() }) trace := &httptrace.ClientTrace{ - GotConn: func(httptrace.GotConnInfo) { + GotFirstResponseByte: func() { stopTTFT() }, } @@ -490,6 +494,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi discardStreamChunks(streamResult.Chunks) } errStream = checkTTFTErr(errStream) + errStream = sanitizeErrorTextFields(errStream) rerr := resultErrorFromError(errStream) action, okAction := matchRequestScopedErrorAction(auth, errStream, m.runtimeConfigSnapshot()) result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: rerr, Options: execOpts} @@ -596,6 +601,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi stopTTFT() cancelAttempt() bootstrapErr = checkTTFTErr(bootstrapErr) + bootstrapErr = sanitizeErrorTextFields(bootstrapErr) action, okAction := matchRequestScopedErrorAction(auth, bootstrapErr, m.runtimeConfigSnapshot()) if okAction { rerr := resultErrorFromError(bootstrapErr) diff --git a/sdk/cliproxy/auth/conductor_stream_eof_test.go b/sdk/cliproxy/auth/conductor_stream_eof_test.go index 54cf0afc1ca..2c1e4a38cc9 100644 --- a/sdk/cliproxy/auth/conductor_stream_eof_test.go +++ b/sdk/cliproxy/auth/conductor_stream_eof_test.go @@ -34,3 +34,37 @@ func TestReadStreamBootstrapFinalizesDetectorAtEOF(t *testing.T) { t.Fatalf("len(buffered) = %d, want 0 when the provider error propagates", len(buffered)) } } + +func TestReadStreamBootstrapWaitsForUsageAfterStop(t *testing.T) { + ch := make(chan cliproxyexecutor.StreamChunk, 2) + ch <- cliproxyexecutor.StreamChunk{Payload: []byte("data: {\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n")} + ch <- cliproxyexecutor.StreamChunk{Payload: []byte("data: {\"choices\":[],\"usage\":{\"completion_tokens\":0}}\n\n")} + close(ch) + + buffered, closed, err := readStreamBootstrap(context.Background(), ch) + if err != nil { + t.Fatalf("readStreamBootstrap() error = %v, want nil", err) + } + if !closed { + t.Fatal("closed = false, want true after terminal usage with zero tokens") + } + if len(buffered) != 2 { + t.Fatalf("len(buffered) = %d, want 2 (stop + usage)", len(buffered)) + } + + ch2 := make(chan cliproxyexecutor.StreamChunk, 2) + ch2 <- cliproxyexecutor.StreamChunk{Payload: []byte("data: {\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n")} + ch2 <- cliproxyexecutor.StreamChunk{Payload: []byte("data: {\"choices\":[],\"usage\":{\"completion_tokens\":3}}\n\n")} + close(ch2) + + buffered2, closed2, err2 := readStreamBootstrap(context.Background(), ch2) + if err2 != nil { + t.Fatalf("readStreamBootstrap() error = %v, want nil", err2) + } + if closed2 { + t.Fatal("closed = true, want false after meaningful usage with positive tokens") + } + if len(buffered2) != 2 { + t.Fatalf("len(buffered2) = %d, want 2 (stop + usage)", len(buffered2)) + } +} diff --git a/sdk/cliproxy/auth/conductor_stream_test.go b/sdk/cliproxy/auth/conductor_stream_test.go new file mode 100644 index 00000000000..b123844f305 --- /dev/null +++ b/sdk/cliproxy/auth/conductor_stream_test.go @@ -0,0 +1,100 @@ +package auth + +import ( + "context" + "net/http" + "strings" + "testing" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +type streamLeakTestExecutor struct { + err error +} + +func (e *streamLeakTestExecutor) Identifier() string { return "stream-leak" } + +func (e *streamLeakTestExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusNotImplemented, Message: "not implemented"} +} + +func (e *streamLeakTestExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + ch := make(chan cliproxyexecutor.StreamChunk, 1) + ch <- cliproxyexecutor.StreamChunk{Err: e.err} + close(ch) + return &cliproxyexecutor.StreamResult{Chunks: ch}, nil +} + +func (e *streamLeakTestExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} + +func (e *streamLeakTestExecutor) Refresh(context.Context, *Auth) (*Auth, error) { return nil, nil } + +func (e *streamLeakTestExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +// TestStreamErrorRedactsQuotedJSONSecretInRecord verifies that an in-band stream +// error containing a non-sk secret in quoted JSON is sanitized before it reaches +// the output stream and the persisted auth result. +func TestStreamErrorRedactsQuotedJSONSecretInRecord(t *testing.T) { + const model = "leak-model" + auth := &Auth{ID: "leak-auth", Provider: "stream-leak", Status: StatusActive} + + exec := &streamLeakTestExecutor{err: &Error{ + Code: "sk-live-test12345", + HTTPStatus: http.StatusUnauthorized, + Message: `{"apiKey":"AIza-secret"}`, + }} + + m := NewManager(nil, nil, nil) + m.RegisterExecutor(exec) + + reg := registry.GetGlobalRegistry() + reg.RegisterClient(auth.ID, "stream-leak", []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { reg.UnregisterClient(auth.ID) }) + + _, err := m.Register(context.Background(), auth) + if err != nil { + t.Fatalf("register auth: %v", err) + } + + stream, errStream := m.ExecuteStream(context.Background(), []string{"stream-leak"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}) + if errStream != nil { + t.Fatalf("ExecuteStream() unexpected error = %v", errStream) + } + + var emittedErr error + for chunk := range stream.Chunks { + if chunk.Err != nil { + emittedErr = chunk.Err + break + } + } + if emittedErr == nil { + t.Fatal("stream emitted no error chunk") + } + if strings.Contains(emittedErr.Error(), "AIza-secret") || strings.Contains(emittedErr.Error(), "sk-live-test12345") { + t.Fatalf("emitted error leaks secrets: %q", emittedErr.Error()) + } + if !strings.Contains(emittedErr.Error(), "REDACTED") { + t.Fatalf("emitted error did not redact secret: %q", emittedErr.Error()) + } + + stored := m.auths[auth.ID] + if stored == nil || stored.LastError == nil { + t.Fatal("auth.LastError is nil, want recorded result") + } + if strings.Contains(stored.LastError.Message, "AIza-secret") { + t.Fatalf("auth.LastError.Message leaks quoted JSON secret: %q", stored.LastError.Message) + } + if strings.Contains(stored.LastError.Code, "sk-live-test12345") { + t.Fatalf("auth.LastError.Code leaks provider code secret: %q", stored.LastError.Code) + } + if !strings.Contains(stored.LastError.Message, "REDACTED") { + t.Fatalf("auth.LastError.Message did not redact secret: %q", stored.LastError.Message) + } +} diff --git a/sdk/cliproxy/auth/empty_completion.go b/sdk/cliproxy/auth/empty_completion.go index d94d252133a..bbd33f43f91 100644 --- a/sdk/cliproxy/auth/empty_completion.go +++ b/sdk/cliproxy/auth/empty_completion.go @@ -1872,6 +1872,11 @@ func (s *streamBootstrapState) isEmptyCompletion() bool { } func (s *streamBootstrapState) isTerminalEmpty() bool { + if s.acc.openAITerminal && !s.acc.sawUsage { + // OpenAI streams may emit finish_reason=stop before the final usage + // frame; do not judge the stream empty until usage arrives. + return false + } return (s.sawDone || s.acc.geminiTerminal || s.acc.claudeTerminal || s.acc.openAITerminal || s.acc.interactionsTerminal) && s.acc.empty() } @@ -2127,7 +2132,9 @@ func parseStreamErrorFromEnvelope(data []byte, envelope streamErrorEnvelope) *Er err.Retryable = true } - return err + // Sanitize every text field before the parsed upstream error reaches + // logging, result recording, or the output stream. + return sanitizeErrorTextFields(err).(*Error) } func evalProviderError(data []byte, sseEvent string) *Error { @@ -2431,11 +2438,11 @@ func couldBeSSEPrefix(payload []byte) bool { // a single non-stream JSON response) represents an empty completion. func isEmptyCompletionPayload(payload []byte) bool { trimmed := bytes.TrimSpace(payload) - if len(trimmed) == 0 { + if len(trimmed) == 0 || bytes.Equal(trimmed, []byte("null")) { // 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. + // credentials. A literal JSON null is equally empty. return true } diff --git a/sdk/cliproxy/auth/empty_completion_sanitize_test.go b/sdk/cliproxy/auth/empty_completion_sanitize_test.go new file mode 100644 index 00000000000..f058fb2622d --- /dev/null +++ b/sdk/cliproxy/auth/empty_completion_sanitize_test.go @@ -0,0 +1,45 @@ +package auth + +import ( + "strings" + "testing" +) + +// TestParseStreamErrorSanitizesCodeAndMessage verifies that evalProviderError +// (which calls parseStreamErrorFromEnvelope) redacts credentials from both the +// parsed code and message, so the secret never reaches logs, results, or +// LastError. +func TestParseStreamErrorSanitizesCodeAndMessage(t *testing.T) { + const codeSecret = "sk-live-test12345" + const msgSecret = "sk-abc123def456" + + data := []byte(`{"error":{"code":"` + codeSecret + `","message":"Bearer ` + msgSecret + `"}}`) + got := evalProviderError(data, "error") + if got == nil { + t.Fatal("evalProviderError returned nil for an error payload") + } + if strings.Contains(got.Code, codeSecret) { + t.Fatalf("parsed error code leaks code secret: %q", got.Code) + } + if !strings.Contains(got.Code, "REDACTED") { + t.Fatalf("parsed error code did not redact code secret: %q", got.Code) + } + if strings.Contains(got.Message, msgSecret) { + t.Fatalf("parsed error message leaks message secret: %q", got.Message) + } + if !strings.Contains(got.Message, "REDACTED") { + t.Fatalf("parsed error message did not redact message secret: %q", got.Message) + } + + // The log summary path must not resurface the raw secret either. + summary := summarizeErrorForLog(got) + if strings.Contains(summary, codeSecret) || strings.Contains(summary, msgSecret) { + t.Fatalf("log summary leaks parsed secret: %q", summary) + } + + // The recorded result path must not resurface the raw secret either. + result := resultErrorFromError(got) + if strings.Contains(result.Code, codeSecret) || strings.Contains(result.Message, msgSecret) { + t.Fatalf("recorded result leaks parsed secret: code=%q message=%q", result.Code, result.Message) + } +} diff --git a/sdk/cliproxy/auth/empty_completion_test.go b/sdk/cliproxy/auth/empty_completion_test.go index 8718ce59c08..8d5fe81c050 100644 --- a/sdk/cliproxy/auth/empty_completion_test.go +++ b/sdk/cliproxy/auth/empty_completion_test.go @@ -4163,15 +4163,19 @@ func TestEmptyCompletion_OpenAIImageGenerationResult(t *testing.T) { } func TestEmptyCompletion_OpenAIFinishReasonStopWithoutDoneIsTerminalEmpty(t *testing.T) { - // Case 1: Single choice finish_reason="stop" without [DONE] is terminal empty. + // Case 1: Single choice finish_reason="stop" without [DONE] is terminal, but + // the bootstrap must not judge it empty until the final usage frame arrives. var detector StreamBootstrapDetector stopChunk := []byte("data: {\"id\":\"1\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n") got := detector.Observe(stopChunk) if got { t.Fatalf("Observe(stopChunk) = %v, want false", got) } - if !detector.IsTerminalEmpty() { - t.Fatal("IsTerminalEmpty() = false, want true for OpenAI finish_reason:stop without [DONE]") + if detector.IsTerminalEmpty() { + t.Fatal("IsTerminalEmpty() = true, want false until usage arrives") + } + if !detector.Finish() { + t.Fatal("Finish() = false, want true for OpenAI finish_reason:stop without usage") } // Case 2: Multi-choice with partial finish_reason (choice 0 "stop", choice 1 nil) is NOT terminal yet. @@ -4236,8 +4240,11 @@ func TestEmptyCompletion_OpenAIFinishReasonStopWithoutDoneIsTerminalEmpty(t *tes if detectorMultiEarly.Observe(frameChoice1Finish) { t.Fatal("Observe(frameChoice1Finish) = true, want false") } - if !detectorMultiEarly.IsTerminalEmpty() { - t.Fatal("IsTerminalEmpty() = false when all n=2 choices finished empty, want true") + if detectorMultiEarly.IsTerminalEmpty() { + t.Fatal("IsTerminalEmpty() = true when all n=2 choices finished empty, want false until usage") + } + if !detectorMultiEarly.Finish() { + t.Fatal("Finish() = false when all n=2 choices finished empty, want true") } } @@ -5214,3 +5221,12 @@ func TestStreamBootstrapDetectorInteractionsWithoutSignatureStaysEmpty(t *testin t.Fatal("IsTerminalEmpty() = false for a genuinely empty interaction, want true") } } + +func TestEmptyCompletionNullPayload(t *testing.T) { + if !IsEmptyCompletionPayload([]byte("null")) { + t.Fatal("IsEmptyCompletionPayload(null) = false, want true") + } + if !IsEmptyCompletionPayload([]byte(" null ")) { + t.Fatal("IsEmptyCompletionPayload(whitespace null) = false, want true") + } +} diff --git a/sdk/cliproxy/auth/stream_ttft_test.go b/sdk/cliproxy/auth/stream_ttft_test.go index 81373769697..89cac69c74c 100644 --- a/sdk/cliproxy/auth/stream_ttft_test.go +++ b/sdk/cliproxy/auth/stream_ttft_test.go @@ -2,6 +2,7 @@ package auth import ( "context" + "errors" "fmt" "io" "net/http" @@ -352,8 +353,8 @@ func TestStreamConnectTimeout_ConfigAndMetadata(t *testing.T) { } // httpTTFTStreamExecutor makes an HTTP request using the default client so the -// conductor's httptrace.ClientTrace fires GotConn as soon as the connection is -// obtained. +// conductor's httptrace.ClientTrace fires GotFirstResponseByte when the +// upstream begins responding. type httpTTFTStreamExecutor struct { baseURL string } @@ -399,13 +400,12 @@ func (*httpTTFTStreamExecutor) HttpRequest(context.Context, *Auth, *http.Request return nil, nil } -// TestStreamFirstChunkTimeout_StoppedAtConnection verifies that the TTFT timer -// is stopped when the HTTP transport reports a connection, before response -// headers are written. A slow-responding server that waits longer than the TTFT -// window should not trigger a timeout. -func TestStreamFirstChunkTimeout_StoppedAtConnection(t *testing.T) { +// TestStreamFirstChunkTimeout_BudgetsSlowHeaders verifies that the TTFT timer +// is not stopped at connection (or a CONNECT tunnel) and still fires when the +// upstream takes longer than the budget to produce the first response byte. +func TestStreamFirstChunkTimeout_BudgetsSlowHeaders(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - time.Sleep(150 * time.Millisecond) + time.Sleep(100 * time.Millisecond) w.Header().Set("Content-Type", "text/event-stream") w.WriteHeader(http.StatusOK) fmt.Fprint(w, "data: {\"type\":\"content_block_delta\",\"delta\":{\"text\":\"hello\"}}\n\n") @@ -433,35 +433,12 @@ func TestStreamFirstChunkTimeout_StoppedAtConnection(t *testing.T) { }, } - start := time.Now() - stream, errStream := m.ExecuteStream(context.Background(), []string{"http-ttft"}, cliproxyexecutor.Request{Model: model}, opts) - elapsed := time.Since(start) - if errStream != nil { - t.Fatalf("ExecuteStream error = %v, want stream after delayed headers", errStream) - } - if elapsed < 100*time.Millisecond { - t.Fatalf("stream returned too quickly (%v), want at least 100ms of delayed headers", elapsed) - } - - done := make(chan struct{}) - var got string - go func() { - defer close(done) - for chunk := range stream.Chunks { - if chunk.Err != nil { - t.Errorf("chunk error = %v", chunk.Err) - return - } - got = string(chunk.Payload) - } - }() - select { - case <-done: - case <-time.After(2 * time.Second): - t.Fatal("timed out reading stream chunks") + _, errStream := m.ExecuteStream(context.Background(), []string{"http-ttft"}, cliproxyexecutor.Request{Model: model}, opts) + if errStream == nil { + t.Fatal("ExecuteStream() = nil, want TTFT timeout when first byte exceeds budget") } - - if !strings.Contains(got, "hello") { - t.Fatalf("stream payload does not contain expected content: %q", got) + var authErr *Error + if !errors.As(errStream, &authErr) || authErr.Code != "stream_first_chunk_timeout" { + t.Fatalf("ExecuteStream error = %v, want TTFT timeout", errStream) } } From 04ae4ef71ed691bb040539e539156b380425c7e8 Mon Sep 17 00:00:00 2001 From: warelik Date: Sun, 23 Aug 2026 18:43:15 +0300 Subject: [PATCH 6/9] fix(auth,pluginhost): stop TTFT at GotConn, use empty_count everywhere --- internal/pluginhost/executor_route.go | 2 +- internal/pluginhost/model_router_test.go | 8 +++--- sdk/cliproxy/auth/conductor_execution.go | 6 ++-- sdk/cliproxy/auth/conductor_stream.go | 11 ++++---- sdk/cliproxy/auth/empty_completion.go | 10 +++++++ sdk/cliproxy/auth/empty_completion_export.go | 7 +++++ sdk/cliproxy/auth/stream_ttft_test.go | 29 +++++++++++--------- 7 files changed, 46 insertions(+), 27 deletions(-) diff --git a/internal/pluginhost/executor_route.go b/internal/pluginhost/executor_route.go index 9f951c0a469..6e3292caf69 100644 --- a/internal/pluginhost/executor_route.go +++ b/internal/pluginhost/executor_route.go @@ -282,7 +282,7 @@ func (h *Host) CountPluginExecutor(ctx context.Context, pluginID string, req cor return coreexecutor.Response{}, err } if coreauth.IsEmptyCompletionPayload(resp.Payload) { - return coreexecutor.Response{}, coreauth.EmptyCompletionError() + return coreexecutor.Response{}, coreauth.EmptyCountError() } return resp, nil } diff --git a/internal/pluginhost/model_router_test.go b/internal/pluginhost/model_router_test.go index 9b90580c391..e5ab16b7298 100644 --- a/internal/pluginhost/model_router_test.go +++ b/internal/pluginhost/model_router_test.go @@ -244,7 +244,7 @@ func TestHostExecutePluginExecutorStreamRejectsEmptyCompletion(t *testing.T) { } } -func TestHostCountPluginExecutorRejectsEmptyCompletion(t *testing.T) { +func TestHostCountPluginExecutorRejectsEmptyCount(t *testing.T) { executor := &fakeExecutor{ identifier: "plugin-provider", countTokens: func(ctx context.Context, req pluginapi.ExecutorRequest) (pluginapi.ExecutorResponse, error) { @@ -262,14 +262,14 @@ func TestHostCountPluginExecutorRejectsEmptyCompletion(t *testing.T) { _, errCount := host.CountPluginExecutor(context.Background(), "executor", coreexecutor.Request{Model: "client-model", Payload: []byte(`{"model":"client-model"}`)}, coreexecutor.Options{}) if errCount == nil { - t.Fatal("CountPluginExecutor() with empty completion = nil, want retriable error") + t.Fatal("CountPluginExecutor() with empty count = nil, want retriable error") } var authErr *coreauth.Error if !errors.As(errCount, &authErr) { t.Fatalf("error = %v (%T), want *coreauth.Error", errCount, errCount) } - if !authErr.Retryable || authErr.Code != "empty_completion" { - t.Fatalf("error = %+v, want retriable empty_completion", authErr) + if !authErr.Retryable || authErr.Code != "empty_count" { + t.Fatalf("error = %+v, want retriable empty_count", authErr) } if authErr.StatusCode() != http.StatusServiceUnavailable { t.Fatalf("error status = %d, want %d", authErr.StatusCode(), http.StatusServiceUnavailable) diff --git a/sdk/cliproxy/auth/conductor_execution.go b/sdk/cliproxy/auth/conductor_execution.go index cea021c4a89..82929ada708 100644 --- a/sdk/cliproxy/auth/conductor_execution.go +++ b/sdk/cliproxy/auth/conductor_execution.go @@ -710,9 +710,9 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, continue } if isEmptyCompletionPayload(resp.Payload) { - lastErr = m.markEmptyCompletion(execCtx, &result) - tracker.Record(auth, errEmptyCompletion) - persistExcludedAuthForRetry(m, auth, errEmptyCompletion, retryRound, defaultRequestRetry, excluded) + lastErr = m.markEmptyCount(execCtx, &result) + tracker.Record(auth, errEmptyCount) + persistExcludedAuthForRetry(m, auth, errEmptyCount, retryRound, defaultRequestRetry, excluded) if homeMode { homeAuthCount++ } diff --git a/sdk/cliproxy/auth/conductor_stream.go b/sdk/cliproxy/auth/conductor_stream.go index f8cd32ece22..394d4bea3ff 100644 --- a/sdk/cliproxy/auth/conductor_stream.go +++ b/sdk/cliproxy/auth/conductor_stream.go @@ -337,12 +337,11 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi } // Arm the TTFT timer only after local interception and request - // preparation: the budget measures upstream responsiveness, so a slow + // preparation: the budget measures connection establishment, so a slow // after-auth interceptor must not cancel the attempt before any - // upstream request was even made. The timer is stopped when the first - // response byte arrives (GotFirstResponseByte), after the actual upstream - // begins responding; this keeps a CONNECT tunnel or slow accept inside - // the budget instead of treating the proxy connection as success. + // upstream request was even made. AGENTS.md:58 permits timeouts only + // until the upstream connection is established, including CONNECT/TLS + // setup for HTTPS proxying, so the timer stops at GotConn. armTTFT := func() { attemptMu.Lock() if timer != nil { @@ -374,7 +373,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi currentCancel() }) trace := &httptrace.ClientTrace{ - GotFirstResponseByte: func() { + GotConn: func(connInfo httptrace.GotConnInfo) { stopTTFT() }, } diff --git a/sdk/cliproxy/auth/empty_completion.go b/sdk/cliproxy/auth/empty_completion.go index bbd33f43f91..8ad673f3a87 100644 --- a/sdk/cliproxy/auth/empty_completion.go +++ b/sdk/cliproxy/auth/empty_completion.go @@ -2587,3 +2587,13 @@ func (m *Manager) markEmptyCompletion(ctx context.Context, result *Result) error m.MarkResult(ctx, *result) return errEmptyCompletion } + +// markEmptyCount records a failed retriable empty count-tokens result and +// returns the error to propagate. It is the count analogue of markEmptyCompletion +// and uses errEmptyCount so count failures have a consistent code everywhere. +func (m *Manager) markEmptyCount(ctx context.Context, result *Result) error { + result.Success = false + result.Error = errEmptyCount + m.MarkResult(ctx, *result) + return errEmptyCount +} diff --git a/sdk/cliproxy/auth/empty_completion_export.go b/sdk/cliproxy/auth/empty_completion_export.go index 60a67f8f249..13805d7b48b 100644 --- a/sdk/cliproxy/auth/empty_completion_export.go +++ b/sdk/cliproxy/auth/empty_completion_export.go @@ -19,6 +19,13 @@ func EmptyCompletionError() error { return errEmptyCompletion } +// EmptyCountError returns the retriable error used when upstream returns an +// empty count-tokens response. The plugin-executor path returns it so count +// failures use the same code as the conductor's count path. +func EmptyCountError() error { + return errEmptyCount +} + type choiceExtractionPayload struct { N *int `json:"n"` CandidateCount *int `json:"candidateCount"` diff --git a/sdk/cliproxy/auth/stream_ttft_test.go b/sdk/cliproxy/auth/stream_ttft_test.go index 89cac69c74c..be83cc345de 100644 --- a/sdk/cliproxy/auth/stream_ttft_test.go +++ b/sdk/cliproxy/auth/stream_ttft_test.go @@ -2,7 +2,6 @@ package auth import ( "context" - "errors" "fmt" "io" "net/http" @@ -353,8 +352,8 @@ func TestStreamConnectTimeout_ConfigAndMetadata(t *testing.T) { } // httpTTFTStreamExecutor makes an HTTP request using the default client so the -// conductor's httptrace.ClientTrace fires GotFirstResponseByte when the -// upstream begins responding. +// conductor's httptrace.ClientTrace fires GotConn when the upstream connection +// is established. type httpTTFTStreamExecutor struct { baseURL string } @@ -400,10 +399,10 @@ func (*httpTTFTStreamExecutor) HttpRequest(context.Context, *Auth, *http.Request return nil, nil } -// TestStreamFirstChunkTimeout_BudgetsSlowHeaders verifies that the TTFT timer -// is not stopped at connection (or a CONNECT tunnel) and still fires when the -// upstream takes longer than the budget to produce the first response byte. -func TestStreamFirstChunkTimeout_BudgetsSlowHeaders(t *testing.T) { +// TestStreamFirstChunkTimeout_SlowHeadersAfterConnectNotTimedOut verifies that +// the TTFT timer stops at connection establishment, so a slow upstream that +// sleeps before writing response headers does not trigger a timeout. +func TestStreamFirstChunkTimeout_SlowHeadersAfterConnectNotTimedOut(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { time.Sleep(100 * time.Millisecond) w.Header().Set("Content-Type", "text/event-stream") @@ -433,12 +432,16 @@ func TestStreamFirstChunkTimeout_BudgetsSlowHeaders(t *testing.T) { }, } - _, errStream := m.ExecuteStream(context.Background(), []string{"http-ttft"}, cliproxyexecutor.Request{Model: model}, opts) - if errStream == nil { - t.Fatal("ExecuteStream() = nil, want TTFT timeout when first byte exceeds budget") + stream, errStream := m.ExecuteStream(context.Background(), []string{"http-ttft"}, cliproxyexecutor.Request{Model: model}, opts) + if errStream != nil { + t.Fatalf("ExecuteStream() = %v, want no timeout for slow headers after connect", errStream) + } + if stream == nil || stream.Chunks == nil { + t.Fatal("ExecuteStream() returned nil stream") } - var authErr *Error - if !errors.As(errStream, &authErr) || authErr.Code != "stream_first_chunk_timeout" { - t.Fatalf("ExecuteStream error = %v, want TTFT timeout", errStream) + for chunk := range stream.Chunks { + if !strings.Contains(string(chunk.Payload), "hello") { + t.Fatalf("unexpected chunk payload: %q", chunk.Payload) + } } } From a995f7b18a308304e9ff57b8672d7cefa17bd957 Mon Sep 17 00:00:00 2001 From: warelik Date: Sun, 23 Aug 2026 19:16:03 +0300 Subject: [PATCH 7/9] fix(pluginhost): exit discardStreamChunks on context cancel and timeout `discardStreamChunks` started a goroutine that ranged over the source channel forever. If a plugin left the channel open after an empty terminal frame, the drainer never exited, causing a cumulative leak in a long-lived proxy. - Pass the request context into `discardStreamChunks`. - Add `streamDrainTimeout` (5s default); the drainer exits on context cancellation, on timeout, or when the source channel closes. - Reset the drain timeout on each received chunk so slow trailing chunks are still drained, but the goroutine is always bounded. - Add `TestDiscardStreamChunksExitsOnContextCancel` and `TestDiscardStreamChunksExitsOnOpenUnclosedChannel` proving the goroutine exits even when the source channel is never closed. --- internal/pluginhost/executor_route.go | 39 ++++++++++++++++--- .../pluginhost/executor_route_stream_test.go | 34 ++++++++++++++++ 2 files changed, 68 insertions(+), 5 deletions(-) diff --git a/internal/pluginhost/executor_route.go b/internal/pluginhost/executor_route.go index 6e3292caf69..db5650fef4b 100644 --- a/internal/pluginhost/executor_route.go +++ b/internal/pluginhost/executor_route.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "strings" + "time" coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" @@ -239,12 +240,12 @@ func wrapStreamEmptyCompletion(ctx context.Context, streamResult *coreexecutor.S } } if streamErr := detector.StreamError(); streamErr != nil { - discardStreamChunks(src) + discardStreamChunks(ctx, src) _ = forward(coreexecutor.StreamChunk{Err: streamErr}) return } if detector.IsTerminalEmpty() { - discardStreamChunks(src) + discardStreamChunks(ctx, src) _ = forward(coreexecutor.StreamChunk{Err: coreauth.EmptyCompletionError()}) return } @@ -253,14 +254,42 @@ func wrapStreamEmptyCompletion(ctx context.Context, streamResult *coreexecutor.S return &coreexecutor.StreamResult{Chunks: wrapped, Headers: streamResult.Headers} } -func discardStreamChunks(ch <-chan coreexecutor.StreamChunk) { +var streamDrainTimeout = 5 * time.Second + +func discardStreamChunks(ctx context.Context, ch <-chan coreexecutor.StreamChunk) <-chan struct{} { + done := make(chan struct{}) if ch == nil { - return + close(done) + return done + } + if ctx == nil { + ctx = context.Background() } go func() { - for range ch { + defer close(done) + timer := time.NewTimer(streamDrainTimeout) + defer timer.Stop() + for { + select { + case <-ctx.Done(): + return + case <-timer.C: + return + case _, ok := <-ch: + if !ok { + return + } + if !timer.Stop() { + select { + case <-timer.C: + default: + } + } + timer.Reset(streamDrainTimeout) + } } }() + return done } func streamChunkPayload(chunks []coreexecutor.StreamChunk) []byte { diff --git a/internal/pluginhost/executor_route_stream_test.go b/internal/pluginhost/executor_route_stream_test.go index 1927f204412..5fcc5a9a8d6 100644 --- a/internal/pluginhost/executor_route_stream_test.go +++ b/internal/pluginhost/executor_route_stream_test.go @@ -312,6 +312,40 @@ func TestWrapStreamEmptyCompletionStopsAtTerminalEmptyMarkersWithoutChannelClose } } +func TestDiscardStreamChunksExitsOnContextCancel(t *testing.T) { + src := make(chan coreexecutor.StreamChunk) + ctx, cancel := context.WithCancel(context.Background()) + done := discardStreamChunks(ctx, src) + + select { + case <-done: + t.Fatal("drain finished too early") + case <-time.After(50 * time.Millisecond): + } + + cancel() + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("discardStreamChunks goroutine did not exit on context cancellation") + } +} + +func TestDiscardStreamChunksExitsOnOpenUnclosedChannel(t *testing.T) { + previous := streamDrainTimeout + streamDrainTimeout = 100 * time.Millisecond + t.Cleanup(func() { streamDrainTimeout = previous }) + + src := make(chan coreexecutor.StreamChunk) + done := discardStreamChunks(context.Background(), src) + + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("discardStreamChunks goroutine did not exit on timeout for open unclosed channel") + } +} + func TestWrapStreamEmptyCompletionDrainsSourceAfterTerminalEmpty(t *testing.T) { src := make(chan coreexecutor.StreamChunk) producerDone := make(chan struct{}) From bc3955f80bb572704dc9eab752751320b9b77fb1 Mon Sep 17 00:00:00 2001 From: warelik Date: Sun, 23 Aug 2026 19:33:48 +0300 Subject: [PATCH 8/9] fix(auth): make sdk/cliproxy/auth discardStreamChunks respect context and timeout `sdk/cliproxy/auth/conductor_stream.go` had its own `discardStreamChunks` variant that started a goroutine ranging over the source channel forever. Without a context, a plugin that left the channel open leaked a goroutine per stream. - Add `ctx` parameter and `streamDrainTimeout` (5s default) to the auth `discardStreamChunks`, mirroring the pluginhost fix. - The drainer exits on context cancellation, on timeout, or when the source channel closes; the timeout resets on each received chunk so slow trailing chunks are still drained. - Update all call sites in `conductor_stream.go` to pass the request context. - Add `TestDiscardStreamChunksExitsOnContextCancel` and `TestDiscardStreamChunksExitsOnOpenUnclosedChannel` in `conductor_stream_test.go`. Also removes the transient-429 changes from the #4881 branch; that finding belongs on #5130 (fix/quota-backoff-hint-floor). --- sdk/cliproxy/auth/conductor_stream.go | 68 +++++++++++++++------- sdk/cliproxy/auth/conductor_stream_test.go | 35 +++++++++++ 2 files changed, 83 insertions(+), 20 deletions(-) diff --git a/sdk/cliproxy/auth/conductor_stream.go b/sdk/cliproxy/auth/conductor_stream.go index 394d4bea3ff..96f459f559b 100644 --- a/sdk/cliproxy/auth/conductor_stream.go +++ b/sdk/cliproxy/auth/conductor_stream.go @@ -54,14 +54,42 @@ func (m *Manager) streamFirstChunkTimeout(opts cliproxyexecutor.Options) time.Du return 0 } -func discardStreamChunks(ch <-chan cliproxyexecutor.StreamChunk) { +var streamDrainTimeout = 5 * time.Second + +func discardStreamChunks(ctx context.Context, ch <-chan cliproxyexecutor.StreamChunk) <-chan struct{} { + done := make(chan struct{}) if ch == nil { - return + close(done) + return done + } + if ctx == nil { + ctx = context.Background() } go func() { - for range ch { + defer close(done) + timer := time.NewTimer(streamDrainTimeout) + defer timer.Stop() + for { + select { + case <-ctx.Done(): + return + case <-timer.C: + return + case _, ok := <-ch: + if !ok { + return + } + if !timer.Stop() { + select { + case <-timer.C: + default: + } + } + timer.Reset(streamDrainTimeout) + } } }() + return done } type streamBootstrapError struct { @@ -266,13 +294,13 @@ func (m *Manager) wrapStreamResult(ctx context.Context, auth *Auth, provider, re } for _, chunk := range buffered { if ok := emit(chunk); !ok { - discardStreamChunks(remaining) + discardStreamChunks(ctx, remaining) return } } for chunk := range remaining { if ok := emit(chunk); !ok { - discardStreamChunks(remaining) + discardStreamChunks(ctx, remaining) return } } @@ -451,7 +479,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi didRefreshOnUnauthorized = true restartAttempt() if streamResult != nil { - discardStreamChunks(streamResult.Chunks) + discardStreamChunks(ctx, streamResult.Chunks) } startRetry := time.Now() streamResult, errStream = executor.ExecuteStream(attemptCtx, auth, execReq, execOpts) @@ -463,7 +491,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi stopTTFT() cancelAttempt() if streamResult != nil { - discardStreamChunks(streamResult.Chunks) + discardStreamChunks(ctx, streamResult.Chunks) } return nil, errCtx } @@ -480,7 +508,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi stopTTFT() cancelAttempt() if streamResult != nil { - discardStreamChunks(streamResult.Chunks) + discardStreamChunks(ctx, streamResult.Chunks) } return nil, errCancel } @@ -490,7 +518,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi stopTTFT() cancelAttempt() if streamResult != nil { - discardStreamChunks(streamResult.Chunks) + discardStreamChunks(ctx, streamResult.Chunks) } errStream = checkTTFTErr(errStream) errStream = sanitizeErrorTextFields(errStream) @@ -529,7 +557,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi if errCtx := ctx.Err(); errCtx != nil { stopTTFT() cancelAttempt() - discardStreamChunks(streamResult.Chunks) + discardStreamChunks(ctx, streamResult.Chunks) return nil, errCtx } bootstrapErr = checkTTFTErr(bootstrapErr) @@ -545,12 +573,12 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi } } if errRefresh != nil { - discardStreamChunks(streamResult.Chunks) + discardStreamChunks(ctx, streamResult.Chunks) bootstrapErr = errRefresh warnLogUpstreamFailure(ctx, entry, provider, execModel, auth, time.Since(startStream), bootstrapErr) streamResult = &cliproxyexecutor.StreamResult{} } else if okRefresh { - discardStreamChunks(streamResult.Chunks) + discardStreamChunks(ctx, streamResult.Chunks) auth = refreshed m.replaceHomeExecutionLifecycleAuth(execOpts.ExecutionLifecycle, auth) publishSelectedAuthMetadata(execOpts.Metadata, auth) @@ -563,7 +591,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi retryErr = checkTTFTErr(retryErr) if retryErr != nil { if retryStream != nil { - discardStreamChunks(retryStream.Chunks) + discardStreamChunks(ctx, retryStream.Chunks) } if errCtx := ctx.Err(); errCtx != nil { stopTTFT() @@ -592,7 +620,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi if errCancel := claudeOAuthRequestCancellation(ctx, auth, bootstrapErr); errCancel != nil { stopTTFT() cancelAttempt() - discardStreamChunks(streamResult.Chunks) + discardStreamChunks(ctx, streamResult.Chunks) return nil, errCancel } } @@ -611,7 +639,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi } applyRequestScopedActionToResult(action, okAction, &result) m.recordExecutionResult(ctx, result, auth, ephemeralResult) - discardStreamChunks(streamResult.Chunks) + discardStreamChunks(ctx, streamResult.Chunks) if isRequestScopedStop(action, okAction) { return nil, wrapRequestStopError(bootstrapErr) } @@ -629,7 +657,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi result.CredentialScope = true } m.recordExecutionResult(ctx, result, auth, ephemeralResult) - discardStreamChunks(streamResult.Chunks) + discardStreamChunks(ctx, streamResult.Chunks) return nil, bootstrapErr } if idx < len(execModels)-1 { @@ -640,7 +668,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi result.CredentialScope = true } m.recordExecutionResult(ctx, result, auth, ephemeralResult) - discardStreamChunks(streamResult.Chunks) + discardStreamChunks(ctx, streamResult.Chunks) lastErr = bootstrapErr if result.CredentialScope { return nil, newStreamBootstrapError(bootstrapErr, streamResult.Headers) @@ -654,7 +682,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi result.CredentialScope = true } m.recordExecutionResult(ctx, result, auth, ephemeralResult) - discardStreamChunks(streamResult.Chunks) + discardStreamChunks(ctx, streamResult.Chunks) return nil, newStreamBootstrapError(bootstrapErr, streamResult.Headers) } @@ -676,7 +704,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi warnLogUpstreamFailure(ctx, entry, provider, execModel, auth, time.Since(startStream), 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) + discardStreamChunks(ctx, streamResult.Chunks) if idx < len(execModels)-1 { lastErr = emptyErr continue @@ -688,7 +716,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi remaining := streamResult.Chunks if closed { - discardStreamChunks(streamResult.Chunks) + discardStreamChunks(ctx, streamResult.Chunks) closedCh := make(chan cliproxyexecutor.StreamChunk) close(closedCh) remaining = closedCh diff --git a/sdk/cliproxy/auth/conductor_stream_test.go b/sdk/cliproxy/auth/conductor_stream_test.go index b123844f305..0571f16d391 100644 --- a/sdk/cliproxy/auth/conductor_stream_test.go +++ b/sdk/cliproxy/auth/conductor_stream_test.go @@ -5,6 +5,7 @@ import ( "net/http" "strings" "testing" + "time" "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" @@ -98,3 +99,37 @@ func TestStreamErrorRedactsQuotedJSONSecretInRecord(t *testing.T) { t.Fatalf("auth.LastError.Message did not redact secret: %q", stored.LastError.Message) } } + +func TestDiscardStreamChunksExitsOnContextCancel(t *testing.T) { + src := make(chan cliproxyexecutor.StreamChunk) + ctx, cancel := context.WithCancel(context.Background()) + done := discardStreamChunks(ctx, src) + + select { + case <-done: + t.Fatal("drain finished too early") + case <-time.After(50 * time.Millisecond): + } + + cancel() + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("discardStreamChunks goroutine did not exit on context cancellation") + } +} + +func TestDiscardStreamChunksExitsOnOpenUnclosedChannel(t *testing.T) { + previous := streamDrainTimeout + streamDrainTimeout = 100 * time.Millisecond + t.Cleanup(func() { streamDrainTimeout = previous }) + + src := make(chan cliproxyexecutor.StreamChunk) + done := discardStreamChunks(context.Background(), src) + + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("discardStreamChunks goroutine did not exit on timeout for open unclosed channel") + } +} From bc92dcc0c50c864b8dbad75375636554bbab5884 Mon Sep 17 00:00:00 2001 From: warelik Date: Mon, 24 Aug 2026 15:06:59 -0400 Subject: [PATCH 9/9] fix(auth): redact in-band error payloads Error-path stream chunks still forwarded raw credentials after the parsed *Error was sanitized. Apply the same recognition set to payloads, and match prose "API key provided" plus vendor prefixes, not only sk-. --- internal/pluginhost/executor_route.go | 13 ++++ .../pluginhost/executor_route_stream_test.go | 30 ++++++++ sdk/cliproxy/auth/conductor_execution.go | 18 ++++- sdk/cliproxy/auth/conductor_overrides_test.go | 15 ++++ sdk/cliproxy/auth/conductor_stream.go | 32 +++++--- sdk/cliproxy/auth/conductor_stream_test.go | 77 +++++++++++++++++++ sdk/cliproxy/auth/empty_completion_export.go | 44 +++++++++++ 7 files changed, 214 insertions(+), 15 deletions(-) diff --git a/internal/pluginhost/executor_route.go b/internal/pluginhost/executor_route.go index db5650fef4b..d226ae70088 100644 --- a/internal/pluginhost/executor_route.go +++ b/internal/pluginhost/executor_route.go @@ -148,7 +148,20 @@ func wrapStreamEmptyCompletion(ctx context.Context, streamResult *coreexecutor.S } } forwarding := false + var payloadErrors coreauth.StreamPayloadErrorDetector forward := func(chunk coreexecutor.StreamChunk) bool { + if chunk.Err != nil { + chunk.Err = coreauth.SanitizeError(chunk.Err) + } + errorPath := chunk.Err != nil || detector.StreamError() != nil + if len(chunk.Payload) > 0 { + if err := payloadErrors.Observe(chunk.Payload); err != nil { + errorPath = true + } + if errorPath { + chunk.Payload = []byte(coreauth.RedactSecrets(string(chunk.Payload))) + } + } select { case <-ctx.Done(): return false diff --git a/internal/pluginhost/executor_route_stream_test.go b/internal/pluginhost/executor_route_stream_test.go index 5fcc5a9a8d6..a2cae1123a8 100644 --- a/internal/pluginhost/executor_route_stream_test.go +++ b/internal/pluginhost/executor_route_stream_test.go @@ -4,6 +4,7 @@ import ( "context" "errors" "net/http" + "strings" "testing" "time" @@ -383,3 +384,32 @@ func TestWrapStreamEmptyCompletionDrainsSourceAfterTerminalEmpty(t *testing.T) { t.Fatal("producer remained blocked after terminal empty return; source was not drained") } } + +// TestWrapStreamEmptyCompletionRedactsInBandErrorPayload is the pluginhost twin of +// P1-A: after meaningful output the wrapper forwards remaining payloads, including +// in-band provider errors, without going through wrapStreamResult. +func TestWrapStreamEmptyCompletionRedactsInBandErrorPayload(t *testing.T) { + const secret = "sk-live-plugin-secret" + src := make(chan coreexecutor.StreamChunk, 2) + src <- coreexecutor.StreamChunk{Payload: []byte("data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hello\"}}]}\n\n")} + src <- coreexecutor.StreamChunk{Payload: []byte(`data: {"error":{"message":"Incorrect API key provided: ` + secret + `","type":"invalid_request_error"}}` + "\n\n")} + close(src) + + wrapped := wrapStreamEmptyCompletion(context.Background(), &coreexecutor.StreamResult{Chunks: src}) + var payloads []string + for chunk := range wrapped.Chunks { + if len(chunk.Payload) > 0 { + payloads = append(payloads, string(chunk.Payload)) + } + } + joined := strings.Join(payloads, "") + if strings.Contains(joined, secret) { + t.Fatalf("plugin stream payload leaks in-band credential: %q", joined) + } + if !strings.Contains(joined, "hello") { + t.Fatalf("meaningful content was dropped: %q", joined) + } + if !strings.Contains(joined, "REDACTED") { + t.Fatalf("in-band error payload was not redacted: %q", joined) + } +} diff --git a/sdk/cliproxy/auth/conductor_execution.go b/sdk/cliproxy/auth/conductor_execution.go index 2a5eae7caaa..3d105eab06b 100644 --- a/sdk/cliproxy/auth/conductor_execution.go +++ b/sdk/cliproxy/auth/conductor_execution.go @@ -1734,9 +1734,14 @@ func formatAuthIdentity(auth *Auth, provider string) string { } var ( - logSecretAPIKeyPattern = regexp.MustCompile(`(?i)\bsk-[A-Za-z0-9_-]{8,}`) - logSecretBearerPattern = regexp.MustCompile(`(?i)\b(Bearer|Basic)\s+[A-Za-z0-9._~+/=-]{4,}`) - logSecretLabeledPattern = regexp.MustCompile(`(?i)((?:"?(?:api[_-]?key|access[_-]?token|token|authorization|secret)"?)\s*[=:]\s*"?)([^\s"&,;}]+)`) + // Vendor API-key prefixes used by providers this proxy talks to, plus the + // historical OpenAI sk- shape. A single extra prefix is not the class. + logSecretAPIKeyPattern = regexp.MustCompile(`(?i)\b(?:sk-[A-Za-z0-9_-]{8,}|xai-[A-Za-z0-9_-]{8,}|gsk_[A-Za-z0-9_-]{8,}|AIza[A-Za-z0-9_-]{8,}|r8_[A-Za-z0-9_-]{8,}|pplx-[A-Za-z0-9_-]{8,}|nvapi-[A-Za-z0-9_-]{8,}|hf_[A-Za-z0-9_-]{8,})`) + logSecretBearerPattern = regexp.MustCompile(`(?i)\b(Bearer|Basic)\s+[A-Za-z0-9._~+/=-]{4,}`) + // Labeled secrets: JSON (`"apiKey":`), `api-key:`, and prose + // (`Incorrect API key provided: `). `[_-]?` does not accept a space, + // so `api[\s_-]*key` plus a short run of filler words before `=`/`:`. + logSecretLabeledPattern = regexp.MustCompile(`(?i)((?:"?(?:api[\s_-]*key|access[\s_-]*token|token|authorization|secret)"?(?:\s+\w+){0,4})\s*[=:]\s*"?)([^\s"&,;}]+)`) ) func redactSecretsForLog(msg string) string { @@ -1746,6 +1751,13 @@ func redactSecretsForLog(msg string) string { return msg } +func redactStreamPayload(payload []byte) []byte { + if len(payload) == 0 { + return payload + } + return []byte(redactSecretsForLog(string(payload))) +} + func summarizeErrorForLog(err error) string { if err == nil { return "" diff --git a/sdk/cliproxy/auth/conductor_overrides_test.go b/sdk/cliproxy/auth/conductor_overrides_test.go index 3788faea399..2a8d75a5257 100644 --- a/sdk/cliproxy/auth/conductor_overrides_test.go +++ b/sdk/cliproxy/auth/conductor_overrides_test.go @@ -2740,6 +2740,21 @@ func TestRedactSecretsForLog_QuotedJSON(t *testing.T) { in: `sk-live-secret`, leaks: []string{"sk-live-secret"}, }, + { + name: "prose API key with xAI prefix", + in: "Incorrect API key provided: xai-live-secret", + leaks: []string{"xai-live-secret"}, + }, + { + name: "prose API key with unlabeled value", + in: "Incorrect API key provided: naked-secret-value12", + leaks: []string{"naked-secret-value12"}, + }, + { + name: "groq prefix without label", + in: "invalid key gsk_live-secret-value", + leaks: []string{"gsk_live-secret-value"}, + }, } for _, tc := range tests { diff --git a/sdk/cliproxy/auth/conductor_stream.go b/sdk/cliproxy/auth/conductor_stream.go index a7cf8917124..e51207d7866 100644 --- a/sdk/cliproxy/auth/conductor_stream.go +++ b/sdk/cliproxy/auth/conductor_stream.go @@ -137,7 +137,7 @@ func (e *streamBootstrapError) Headers() http.Header { func streamErrorResult(headers http.Header, err error) *cliproxyexecutor.StreamResult { ch := make(chan cliproxyexecutor.StreamChunk, 1) - ch <- cliproxyexecutor.StreamChunk{Err: err} + ch <- cliproxyexecutor.StreamChunk{Err: sanitizeErrorTextFields(err)} close(ch) return &cliproxyexecutor.StreamResult{ Headers: cloneHTTPHeader(headers), @@ -243,23 +243,28 @@ func (m *Manager) wrapStreamResult(ctx context.Context, auth *Auth, provider, re applyRequestScopedActionToResult(action, okAction, &result) m.recordExecutionResult(ctx, result, auth, ephemeralResult) } - if !failed && len(chunk.Payload) > 0 { - if streamErr := errorDetector.Observe(chunk.Payload); streamErr != nil { - failed = true - streamErr = sanitizeErrorTextFields(streamErr).(*Error) - entry := logEntryWithRequestID(ctx) - warnLogUpstreamFailure(ctx, entry, provider, resultModel, auth, time.Since(streamStart), streamErr) - rerr := resultErrorFromError(streamErr) - action, okAction := matchRequestScopedErrorAction(auth, streamErr, m.runtimeConfigSnapshot()) - result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: rerr, Options: opts} - applyRequestScopedActionToResult(action, okAction, &result) - m.recordExecutionResult(ctx, result, auth, ephemeralResult) + if len(chunk.Payload) > 0 { + if !failed { + if streamErr := errorDetector.Observe(chunk.Payload); streamErr != nil { + failed = true + streamErr = sanitizeErrorTextFields(streamErr).(*Error) + entry := logEntryWithRequestID(ctx) + warnLogUpstreamFailure(ctx, entry, provider, resultModel, auth, time.Since(streamStart), streamErr) + rerr := resultErrorFromError(streamErr) + action, okAction := matchRequestScopedErrorAction(auth, streamErr, m.runtimeConfigSnapshot()) + result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: rerr, Options: opts} + applyRequestScopedActionToResult(action, okAction, &result) + m.recordExecutionResult(ctx, result, auth, ephemeralResult) + } } } if !forward { return false } if chunk.Err != nil { + if failed { + chunk.Payload = redactStreamPayload(chunk.Payload) + } if ctx == nil { out <- chunk return true @@ -279,6 +284,9 @@ func (m *Manager) wrapStreamResult(ctx context.Context, auth *Auth, provider, re if len(payload) == 0 { return true } + if failed { + payload = redactStreamPayload(payload) + } chunk.Payload = payload if ctx == nil { out <- chunk diff --git a/sdk/cliproxy/auth/conductor_stream_test.go b/sdk/cliproxy/auth/conductor_stream_test.go index 0571f16d391..e5374990605 100644 --- a/sdk/cliproxy/auth/conductor_stream_test.go +++ b/sdk/cliproxy/auth/conductor_stream_test.go @@ -100,6 +100,83 @@ func TestStreamErrorRedactsQuotedJSONSecretInRecord(t *testing.T) { } } +type streamInBandLeakExecutor struct { + chunks []cliproxyexecutor.StreamChunk +} + +func (e *streamInBandLeakExecutor) Identifier() string { return "stream-inband-leak" } + +func (e *streamInBandLeakExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusNotImplemented, Message: "not implemented"} +} + +func (e *streamInBandLeakExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + ch := make(chan cliproxyexecutor.StreamChunk, len(e.chunks)) + for _, chunk := range e.chunks { + ch <- chunk + } + close(ch) + return &cliproxyexecutor.StreamResult{Chunks: ch}, nil +} + +func (e *streamInBandLeakExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} + +func (e *streamInBandLeakExecutor) Refresh(context.Context, *Auth) (*Auth, error) { return nil, nil } + +func (e *streamInBandLeakExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +// TestWrapStreamResultRedactsInBandErrorPayload verifies P1-A: after meaningful +// output, an in-band provider error payload is not forwarded to the caller with +// the raw credential still in chunk.Payload. The *Error object was already +// sanitized; the leak is the payload itself. +func TestWrapStreamResultRedactsInBandErrorPayload(t *testing.T) { + const model = "inband-leak-model" + const secret = "sk-live-inband-secret" + auth := &Auth{ID: "inband-leak-auth", Provider: "stream-inband-leak", Status: StatusActive} + + exec := &streamInBandLeakExecutor{chunks: []cliproxyexecutor.StreamChunk{ + {Payload: []byte("data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hello\"}}]}\n\n")}, + {Payload: []byte(`data: {"error":{"message":"Incorrect API key provided: ` + secret + `","type":"invalid_request_error"}}` + "\n\n")}, + }} + + m := NewManager(nil, nil, nil) + m.RegisterExecutor(exec) + + reg := registry.GetGlobalRegistry() + reg.RegisterClient(auth.ID, "stream-inband-leak", []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { reg.UnregisterClient(auth.ID) }) + + if _, err := m.Register(context.Background(), auth); err != nil { + t.Fatalf("register auth: %v", err) + } + + stream, errStream := m.ExecuteStream(context.Background(), []string{"stream-inband-leak"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}) + if errStream != nil { + t.Fatalf("ExecuteStream() unexpected error = %v", errStream) + } + + var payloads []string + for chunk := range stream.Chunks { + if len(chunk.Payload) > 0 { + payloads = append(payloads, string(chunk.Payload)) + } + } + joined := strings.Join(payloads, "") + if strings.Contains(joined, secret) { + t.Fatalf("caller-visible stream payload leaks in-band credential: %q", joined) + } + if !strings.Contains(joined, "hello") { + t.Fatalf("meaningful content was dropped: %q", joined) + } + if !strings.Contains(joined, "REDACTED") { + t.Fatalf("in-band error payload was not redacted: %q", joined) + } +} + func TestDiscardStreamChunksExitsOnContextCancel(t *testing.T) { src := make(chan cliproxyexecutor.StreamChunk) ctx, cancel := context.WithCancel(context.Background()) diff --git a/sdk/cliproxy/auth/empty_completion_export.go b/sdk/cliproxy/auth/empty_completion_export.go index 13805d7b48b..4ee81900932 100644 --- a/sdk/cliproxy/auth/empty_completion_export.go +++ b/sdk/cliproxy/auth/empty_completion_export.go @@ -163,3 +163,47 @@ func (d *StreamBootstrapDetector) IsTerminalEmpty() bool { } return d.state.isTerminalEmpty() } + +// RedactSecrets redacts credential-shaped values using the same recognition +// set as conductor logging and *Error sanitization. Plugin stream wrappers +// must apply it before emitting an error-path payload to a caller. +func RedactSecrets(s string) string { + return redactSecretsForLog(s) +} + +// SanitizeError redacts every exported string field of an *Error. It is the +// exported form of sanitizeErrorTextFields so pluginhost uses the same +// mechanism as wrapStreamResult rather than a second copy. +func SanitizeError(err error) error { + return sanitizeErrorTextFields(err) +} + +// StreamPayloadErrorDetector incrementally detects in-band provider errors in +// stream bytes after bootstrap has already forwarded meaningful output. +type StreamPayloadErrorDetector struct { + state streamPayloadErrorDetector +} + +// Observe records a stream fragment and returns a sanitized in-band error +// when one has been detected. A nil return is not a successful completion. +func (d *StreamPayloadErrorDetector) Observe(payload []byte) error { + if d == nil { + return nil + } + if err := d.state.Observe(payload); err != nil { + return err + } + return nil +} + +// Finish flushes a trailing unterminated fragment and returns a sanitized +// in-band error when one is present. +func (d *StreamPayloadErrorDetector) Finish() error { + if d == nil { + return nil + } + if err := d.state.Finish(); err != nil { + return err + } + return nil +}