package handlers import ( "context" "encoding/json" "errors" "net/http" "net/http/httptest" "strings" "sync" "sync/atomic" "testing" "time" "github.com/gin-gonic/gin" internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" "github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces" "github.com/router-for-me/CLIProxyAPI/v7/internal/logging" "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config" "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" ) type failOnceStreamExecutor struct { mu sync.Mutex calls int } func (e *failOnceStreamExecutor) Identifier() string { return "codex" } func (e *failOnceStreamExecutor) Execute(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) { return coreexecutor.Response{}, &coreauth.Error{Code: "not_implemented", Message: "Execute not implemented"} } func (e *failOnceStreamExecutor) ExecuteStream(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (*coreexecutor.StreamResult, error) { e.mu.Lock() e.calls++ call := e.calls e.mu.Unlock() ch := make(chan coreexecutor.StreamChunk, 1) if call == 1 { ch <- coreexecutor.StreamChunk{ Err: &coreauth.Error{ Code: "unauthorized", Message: "unauthorized", Retryable: false, HTTPStatus: http.StatusUnauthorized, }, } close(ch) return &coreexecutor.StreamResult{ Headers: http.Header{"X-Upstream-Attempt": {"1"}}, Chunks: ch, }, nil } ch <- coreexecutor.StreamChunk{Payload: []byte("ok")} close(ch) return &coreexecutor.StreamResult{ Headers: http.Header{"X-Upstream-Attempt": {"2"}}, Chunks: ch, }, nil } func (e *failOnceStreamExecutor) Refresh(ctx context.Context, auth *coreauth.Auth) (*coreauth.Auth, error) { return auth, nil } func (e *failOnceStreamExecutor) CountTokens(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) { return coreexecutor.Response{}, &coreauth.Error{Code: "not_implemented", Message: "CountTokens not implemented"} } func (e *failOnceStreamExecutor) HttpRequest(ctx context.Context, auth *coreauth.Auth, req *http.Request) (*http.Response, error) { return nil, &coreauth.Error{ Code: "not_implemented", Message: "HttpRequest not implemented", HTTPStatus: http.StatusNotImplemented, } } func (e *failOnceStreamExecutor) Calls() int { e.mu.Lock() defer e.mu.Unlock() return e.calls } type blockingRetryStreamExecutor struct { mu sync.Mutex calls int retryStarted chan struct{} allowRetry chan struct{} } func (e *blockingRetryStreamExecutor) Identifier() string { return "codex" } func (e *blockingRetryStreamExecutor) Execute(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) { return coreexecutor.Response{}, &coreauth.Error{Code: "not_implemented", Message: "Execute not implemented"} } func (e *blockingRetryStreamExecutor) ExecuteStream(ctx context.Context, auth *coreauth.Auth, req coreexecutor.Request, opts coreexecutor.Options) (*coreexecutor.StreamResult, error) { e.mu.Lock() e.calls++ call := e.calls e.mu.Unlock() if call == 1 { chunks := make(chan coreexecutor.StreamChunk, 1) chunks <- coreexecutor.StreamChunk{Err: &coreauth.Error{Code: "unauthorized", Message: "unauthorized", HTTPStatus: http.StatusUnauthorized}} close(chunks) return &coreexecutor.StreamResult{Headers: http.Header{"X-Upstream-Attempt": {"1"}}, Chunks: chunks}, nil } close(e.retryStarted) select { case <-ctx.Done(): return nil, ctx.Err() case <-e.allowRetry: } chunks := make(chan coreexecutor.StreamChunk, 1) chunks <- coreexecutor.StreamChunk{Payload: []byte("ok")} close(chunks) return &coreexecutor.StreamResult{Headers: http.Header{"X-Upstream-Attempt": {"2"}}, Chunks: chunks}, nil } func (e *blockingRetryStreamExecutor) Refresh(ctx context.Context, auth *coreauth.Auth) (*coreauth.Auth, error) { return auth, nil } func (e *blockingRetryStreamExecutor) CountTokens(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) { return coreexecutor.Response{}, &coreauth.Error{Code: "not_implemented", Message: "CountTokens not implemented"} } func (e *blockingRetryStreamExecutor) HttpRequest(ctx context.Context, auth *coreauth.Auth, req *http.Request) (*http.Response, error) { return nil, &coreauth.Error{Code: "not_implemented", Message: "HttpRequest not implemented", HTTPStatus: http.StatusNotImplemented} } type payloadThenErrorStreamExecutor struct { mu sync.Mutex calls int } func (e *payloadThenErrorStreamExecutor) Identifier() string { return "codex" } func (e *payloadThenErrorStreamExecutor) Execute(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) { return coreexecutor.Response{}, &coreauth.Error{Code: "not_implemented", Message: "Execute not implemented"} } func (e *payloadThenErrorStreamExecutor) ExecuteStream(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (*coreexecutor.StreamResult, error) { e.mu.Lock() e.calls++ e.mu.Unlock() ch := make(chan coreexecutor.StreamChunk, 2) ch <- coreexecutor.StreamChunk{Payload: []byte("partial")} ch <- coreexecutor.StreamChunk{ Err: &coreauth.Error{ Code: "upstream_closed", Message: "upstream closed", Retryable: false, HTTPStatus: http.StatusBadGateway, }, } close(ch) return &coreexecutor.StreamResult{Chunks: ch}, nil } func (e *payloadThenErrorStreamExecutor) Refresh(ctx context.Context, auth *coreauth.Auth) (*coreauth.Auth, error) { return auth, nil } func (e *payloadThenErrorStreamExecutor) CountTokens(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) { return coreexecutor.Response{}, &coreauth.Error{Code: "not_implemented", Message: "CountTokens not implemented"} } func (e *payloadThenErrorStreamExecutor) HttpRequest(ctx context.Context, auth *coreauth.Auth, req *http.Request) (*http.Response, error) { return nil, &coreauth.Error{ Code: "not_implemented", Message: "HttpRequest not implemented", HTTPStatus: http.StatusNotImplemented, } } func (e *payloadThenErrorStreamExecutor) Calls() int { e.mu.Lock() defer e.mu.Unlock() return e.calls } type authAwareStreamExecutor struct { mu sync.Mutex calls int authIDs []string } type invalidJSONStreamExecutor struct{} type splitResponsesEventStreamExecutor struct{} func (e *invalidJSONStreamExecutor) Identifier() string { return "codex" } func (e *invalidJSONStreamExecutor) Execute(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) { return coreexecutor.Response{}, &coreauth.Error{Code: "not_implemented", Message: "Execute not implemented"} } func (e *invalidJSONStreamExecutor) ExecuteStream(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (*coreexecutor.StreamResult, error) { ch := make(chan coreexecutor.StreamChunk, 1) ch <- coreexecutor.StreamChunk{Payload: []byte("event: response.completed\ndata: {\"type\"")} close(ch) return &coreexecutor.StreamResult{Chunks: ch}, nil } func (e *invalidJSONStreamExecutor) Refresh(ctx context.Context, auth *coreauth.Auth) (*coreauth.Auth, error) { return auth, nil } func (e *invalidJSONStreamExecutor) CountTokens(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) { return coreexecutor.Response{}, &coreauth.Error{Code: "not_implemented", Message: "CountTokens not implemented"} } func (e *invalidJSONStreamExecutor) HttpRequest(ctx context.Context, auth *coreauth.Auth, req *http.Request) (*http.Response, error) { return nil, &coreauth.Error{ Code: "not_implemented", Message: "HttpRequest not implemented", HTTPStatus: http.StatusNotImplemented, } } func (e *splitResponsesEventStreamExecutor) Identifier() string { return "split-sse" } func (e *splitResponsesEventStreamExecutor) Execute(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) { return coreexecutor.Response{}, &coreauth.Error{Code: "not_implemented", Message: "Execute not implemented"} } 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\":[]}}")} close(ch) return &coreexecutor.StreamResult{Chunks: ch}, nil } func (e *splitResponsesEventStreamExecutor) Refresh(ctx context.Context, auth *coreauth.Auth) (*coreauth.Auth, error) { return auth, nil } func (e *splitResponsesEventStreamExecutor) CountTokens(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) { return coreexecutor.Response{}, &coreauth.Error{Code: "not_implemented", Message: "CountTokens not implemented"} } func (e *splitResponsesEventStreamExecutor) HttpRequest(ctx context.Context, auth *coreauth.Auth, req *http.Request) (*http.Response, error) { return nil, &coreauth.Error{ Code: "not_implemented", Message: "HttpRequest not implemented", HTTPStatus: http.StatusNotImplemented, } } func (e *authAwareStreamExecutor) Identifier() string { return "codex" } func (e *authAwareStreamExecutor) Execute(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) { return coreexecutor.Response{}, &coreauth.Error{Code: "not_implemented", Message: "Execute not implemented"} } func (e *authAwareStreamExecutor) ExecuteStream(ctx context.Context, auth *coreauth.Auth, req coreexecutor.Request, opts coreexecutor.Options) (*coreexecutor.StreamResult, error) { _ = ctx _ = req _ = opts ch := make(chan coreexecutor.StreamChunk, 1) authID := "" if auth != nil { authID = auth.ID } e.mu.Lock() e.calls++ e.authIDs = append(e.authIDs, authID) e.mu.Unlock() if authID == "auth1" { ch <- coreexecutor.StreamChunk{ Err: &coreauth.Error{ Code: "unauthorized", Message: "unauthorized", Retryable: false, HTTPStatus: http.StatusUnauthorized, }, } close(ch) return &coreexecutor.StreamResult{Chunks: ch}, nil } ch <- coreexecutor.StreamChunk{Payload: []byte("ok")} close(ch) return &coreexecutor.StreamResult{Chunks: ch}, nil } func (e *authAwareStreamExecutor) Refresh(ctx context.Context, auth *coreauth.Auth) (*coreauth.Auth, error) { return auth, nil } func (e *authAwareStreamExecutor) CountTokens(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) { return coreexecutor.Response{}, &coreauth.Error{Code: "not_implemented", Message: "CountTokens not implemented"} } func (e *authAwareStreamExecutor) HttpRequest(ctx context.Context, auth *coreauth.Auth, req *http.Request) (*http.Response, error) { return nil, &coreauth.Error{ Code: "not_implemented", Message: "HttpRequest not implemented", HTTPStatus: http.StatusNotImplemented, } } func (e *authAwareStreamExecutor) Calls() int { e.mu.Lock() defer e.mu.Unlock() return e.calls } func (e *authAwareStreamExecutor) AuthIDs() []string { e.mu.Lock() defer e.mu.Unlock() out := make([]string, len(e.authIDs)) copy(out, e.authIDs) return out } func TestExecuteStreamWithAuthManager_RetriesBeforeFirstByte(t *testing.T) { executor := &failOnceStreamExecutor{} manager := coreauth.NewManager(nil, nil, nil) manager.RegisterExecutor(executor) auth1 := &coreauth.Auth{ ID: "auth1", Provider: "codex", Status: coreauth.StatusActive, Metadata: map[string]any{"email": "test1@example.com"}, } if _, err := manager.Register(context.Background(), auth1); err != nil { t.Fatalf("manager.Register(auth1): %v", err) } auth2 := &coreauth.Auth{ ID: "auth2", Provider: "codex", Status: coreauth.StatusActive, Metadata: map[string]any{"email": "test2@example.com"}, } if _, err := manager.Register(context.Background(), auth2); err != nil { t.Fatalf("manager.Register(auth2): %v", err) } registry.GetGlobalRegistry().RegisterClient(auth1.ID, auth1.Provider, []*registry.ModelInfo{{ID: "test-model"}}) registry.GetGlobalRegistry().RegisterClient(auth2.ID, auth2.Provider, []*registry.ModelInfo{{ID: "test-model"}}) t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth1.ID) registry.GetGlobalRegistry().UnregisterClient(auth2.ID) }) handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{ PassthroughHeaders: true, Streaming: sdkconfig.StreamingConfig{ BootstrapRetries: 1, }, }, manager) dataChan, upstreamHeaders, errChan := handler.ExecuteStreamWithAuthManager(context.Background(), "openai", "test-model", []byte(`{"model":"test-model"}`), "") if dataChan == nil || errChan == nil { t.Fatalf("expected non-nil channels") } var got []byte for chunk := range dataChan { got = append(got, chunk...) } for msg := range errChan { if msg != nil { t.Fatalf("unexpected error: %+v", msg) } } if string(got) != "ok" { t.Fatalf("expected payload ok, got %q", string(got)) } if executor.Calls() != 2 { t.Fatalf("expected 2 stream attempts, got %d", executor.Calls()) } upstreamAttemptHeader := upstreamHeaders.Get("X-Upstream-Attempt") if upstreamAttemptHeader != "2" { t.Fatalf("expected upstream header from retry attempt, got %q", upstreamAttemptHeader) } } func TestExecuteStreamWithAuthManager_ResolvesBootstrapRetryHeadersBeforeReturn(t *testing.T) { executor := &blockingRetryStreamExecutor{ retryStarted: make(chan struct{}), allowRetry: make(chan struct{}), } manager := coreauth.NewManager(nil, nil, nil) manager.RegisterExecutor(executor) auth1 := &coreauth.Auth{ID: "auth1", Provider: "codex", Status: coreauth.StatusActive, Metadata: map[string]any{"email": "test1@example.com"}} if _, err := manager.Register(context.Background(), auth1); err != nil { t.Fatalf("manager.Register(auth1): %v", err) } auth2 := &coreauth.Auth{ID: "auth2", Provider: "codex", Status: coreauth.StatusActive, Metadata: map[string]any{"email": "test2@example.com"}} if _, err := manager.Register(context.Background(), auth2); err != nil { t.Fatalf("manager.Register(auth2): %v", err) } registry.GetGlobalRegistry().RegisterClient(auth1.ID, auth1.Provider, []*registry.ModelInfo{{ID: "test-model"}}) registry.GetGlobalRegistry().RegisterClient(auth2.ID, auth2.Provider, []*registry.ModelInfo{{ID: "test-model"}}) t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth1.ID) registry.GetGlobalRegistry().UnregisterClient(auth2.ID) }) handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{PassthroughHeaders: true, Streaming: sdkconfig.StreamingConfig{BootstrapRetries: 1}}, manager) type streamResult struct { dataChan <-chan []byte upstreamHeaders http.Header errChan <-chan *interfaces.ErrorMessage } resultChan := make(chan streamResult, 1) go func() { dataChan, upstreamHeaders, errChan := handler.ExecuteStreamWithAuthManager(context.Background(), "openai", "test-model", []byte(`{"model":"test-model"}`), "") resultChan <- streamResult{dataChan: dataChan, upstreamHeaders: upstreamHeaders, errChan: errChan} }() select { case result := <-resultChan: t.Fatalf("ExecuteStreamWithAuthManager returned before bootstrap retry completed: %#v", result.upstreamHeaders) case <-executor.retryStarted: } select { case result := <-resultChan: t.Fatalf("ExecuteStreamWithAuthManager returned while bootstrap retry was blocked: %#v", result.upstreamHeaders) default: } close(executor.allowRetry) result := <-resultChan if result.upstreamHeaders.Get("X-Upstream-Attempt") != "2" { t.Fatalf("upstream headers = %#v, want retry attempt headers", result.upstreamHeaders) } for range result.dataChan { } for msg := range result.errChan { if msg != nil { t.Fatalf("unexpected stream error: %+v", msg) } } } type bootstrapStreamExecutor struct { mu sync.Mutex calls int stream func(context.Context, int) (*coreexecutor.StreamResult, error) } func (*bootstrapStreamExecutor) Identifier() string { return "bootstrap-test" } func (e *bootstrapStreamExecutor) Execute(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) { return coreexecutor.Response{}, &coreauth.Error{Code: "not_implemented", Message: "Execute not implemented"} } func (e *bootstrapStreamExecutor) ExecuteStream(ctx context.Context, _ *coreauth.Auth, _ coreexecutor.Request, _ coreexecutor.Options) (*coreexecutor.StreamResult, error) { e.mu.Lock() e.calls++ call := e.calls e.mu.Unlock() return e.stream(ctx, call) } func (e *bootstrapStreamExecutor) Refresh(context.Context, *coreauth.Auth) (*coreauth.Auth, error) { return nil, nil } func (e *bootstrapStreamExecutor) CountTokens(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) { return coreexecutor.Response{}, &coreauth.Error{Code: "not_implemented", Message: "CountTokens not implemented"} } func (e *bootstrapStreamExecutor) HttpRequest(context.Context, *coreauth.Auth, *http.Request) (*http.Response, error) { return nil, &coreauth.Error{Code: "not_implemented", Message: "HttpRequest not implemented", HTTPStatus: http.StatusNotImplemented} } func (e *bootstrapStreamExecutor) Calls() int { e.mu.Lock() defer e.mu.Unlock() return e.calls } func registerBootstrapExecutor(t *testing.T, executor *bootstrapStreamExecutor) (*BaseAPIHandler, *coreauth.Manager) { t.Helper() manager := coreauth.NewManager(nil, nil, nil) manager.RegisterExecutor(executor) auth := &coreauth.Auth{ID: "bootstrap-auth", Provider: executor.Identifier(), Status: coreauth.StatusActive, Metadata: map[string]any{"email": "bootstrap@example.com"}} if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { t.Fatalf("manager.Register(): %v", errRegister) } registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: "bootstrap-model"}}) authRetry := &coreauth.Auth{ID: "bootstrap-auth-retry", Provider: executor.Identifier(), Status: coreauth.StatusActive, Metadata: map[string]any{"email": "bootstrap-retry@example.com"}} if _, errRegister := manager.Register(context.Background(), authRetry); errRegister != nil { t.Fatalf("manager.Register(retry): %v", errRegister) } registry.GetGlobalRegistry().RegisterClient(authRetry.ID, authRetry.Provider, []*registry.ModelInfo{{ID: "bootstrap-model"}}) t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth.ID) registry.GetGlobalRegistry().UnregisterClient(authRetry.ID) }) return NewBaseAPIHandlers(&sdkconfig.SDKConfig{Streaming: sdkconfig.StreamingConfig{BootstrapRetries: 1}}, manager), manager } func TestExecuteStreamWithAuthManager_RetriesAfterDroppedBootstrapPayload(t *testing.T) { executor := &bootstrapStreamExecutor{stream: func(_ context.Context, call int) (*coreexecutor.StreamResult, error) { chunks := make(chan coreexecutor.StreamChunk, 2) if call == 1 { chunks <- coreexecutor.StreamChunk{Payload: []byte("drop")} chunks <- coreexecutor.StreamChunk{Err: &coreauth.Error{HTTPStatus: http.StatusUnauthorized, Message: "unauthorized"}} } else { chunks <- coreexecutor.StreamChunk{Payload: []byte("ok")} } close(chunks) return &coreexecutor.StreamResult{Chunks: chunks}, nil }} handler, _ := registerBootstrapExecutor(t, executor) var intercepted []string handler.SetPluginHost(&handlerInterceptorTestHost{interceptStreamChunk: func(_ context.Context, req pluginapi.StreamChunkInterceptRequest) pluginapi.StreamChunkInterceptResponse { if req.ChunkIndex >= 0 { intercepted = append(intercepted, string(req.Body)) } return pluginapi.StreamChunkInterceptResponse{Body: cloneBytes(req.Body), DropChunk: string(req.Body) == "drop"} }}) dataChan, _, errChan := handler.ExecuteStreamWithAuthManager(context.Background(), "openai", "bootstrap-model", []byte(`{"model":"bootstrap-model"}`), "") var got []byte for chunk := range dataChan { got = append(got, chunk...) } for msg := range errChan { if msg != nil { t.Fatalf("unexpected stream error: %+v", msg) } } if string(got) != "ok" { t.Fatalf("stream payload = %q, want ok", got) } if executor.Calls() != 2 { t.Fatalf("stream attempts = %d, want 2", executor.Calls()) } if strings.Join(intercepted, ",") != "drop,ok" { t.Fatalf("intercepted payloads = %v, want [drop ok] without double interception", intercepted) } } func TestExecuteStreamWithAuthManager_ResetsResponsesValidatorOnBootstrapRetry(t *testing.T) { executor := &bootstrapStreamExecutor{stream: func(_ context.Context, call int) (*coreexecutor.StreamResult, error) { chunks := make(chan coreexecutor.StreamChunk, 2) if call == 1 { chunks <- coreexecutor.StreamChunk{Payload: []byte("event: response.completed\ndata: {\"type\":\"response.completed\",")} chunks <- coreexecutor.StreamChunk{Err: &coreauth.Error{HTTPStatus: http.StatusUnauthorized, Message: "unauthorized"}} } else { chunks <- coreexecutor.StreamChunk{Payload: []byte("event: response.completed\ndata: {\"type\":\"response.completed\",\"response\":{\"status\":\"completed\"}}\n\n")} } close(chunks) return &coreexecutor.StreamResult{Chunks: chunks}, nil }} handler, _ := registerBootstrapExecutor(t, executor) dataChan, _, errChan := handler.ExecuteStreamWithAuthManager(context.Background(), "openai-response", "bootstrap-model", []byte(`{"model":"bootstrap-model"}`), "") var got []byte for chunk := range dataChan { got = append(got, chunk...) } for msg := range errChan { if msg != nil { t.Fatalf("unexpected stream error after retry: %+v", msg) } } if executor.Calls() != 2 || !strings.Contains(string(got), "response.completed") { t.Fatalf("retry calls=%d payload=%q", executor.Calls(), got) } } func TestExecuteStreamWithAuthManager_CancelDuringSynchronousBootstrap(t *testing.T) { started := make(chan struct{}) executor := &bootstrapStreamExecutor{stream: func(_ context.Context, _ int) (*coreexecutor.StreamResult, error) { close(started) return &coreexecutor.StreamResult{Chunks: make(chan coreexecutor.StreamChunk)}, nil }} handler, _ := registerBootstrapExecutor(t, executor) ctx, cancel := context.WithCancel(context.Background()) defer cancel() type result struct { data <-chan []byte errs <-chan *interfaces.ErrorMessage } results := make(chan result, 1) go func() { dataChan, _, errChan := handler.ExecuteStreamWithAuthManager(ctx, "openai", "bootstrap-model", []byte(`{"model":"bootstrap-model"}`), "") results <- result{data: dataChan, errs: errChan} }() <-started cancel() select { case got := <-results: if got.data != nil { if _, ok := <-got.data; ok { t.Fatal("data channel remains open after bootstrap cancellation") } } if got.errs != nil { for range got.errs { } } case <-time.After(time.Second): t.Fatal("bootstrap cancellation did not return") } } func TestExecuteStreamWithAuthManager_EmptyClosedStream(t *testing.T) { executor := &bootstrapStreamExecutor{stream: func(_ context.Context, _ int) (*coreexecutor.StreamResult, error) { chunks := make(chan coreexecutor.StreamChunk) close(chunks) return &coreexecutor.StreamResult{Chunks: chunks}, nil }} handler, _ := registerBootstrapExecutor(t, executor) dataChan, _, errChan := handler.ExecuteStreamWithAuthManager(context.Background(), "openai", "bootstrap-model", []byte(`{"model":"bootstrap-model"}`), "") if _, ok := <-dataChan; ok { t.Fatal("empty stream produced data") } var streamErr *interfaces.ErrorMessage for msg := range errChan { if msg != nil { streamErr = msg } } if streamErr == nil || streamErr.StatusCode != http.StatusInternalServerError { t.Fatalf("empty stream error = %+v, want terminal internal-server error", streamErr) } } type handlerReleaseNotification struct { group executionregistry.ReleaseGroup sequence int64 } type handlerReleaseSink struct { mu sync.Mutex notifications []handlerReleaseNotification notified chan struct{} } func newHandlerReleaseSink() *handlerReleaseSink { return &handlerReleaseSink{notified: make(chan struct{}, 1)} } func (s *handlerReleaseSink) MarkDirty(group executionregistry.ReleaseGroup, sequence int64) { s.mu.Lock() s.notifications = append(s.notifications, handlerReleaseNotification{group: group, sequence: sequence}) s.mu.Unlock() select { case s.notified <- struct{}{}: default: } } func (s *handlerReleaseSink) Notifications() []handlerReleaseNotification { s.mu.Lock() defer s.mu.Unlock() return append([]handlerReleaseNotification(nil), s.notifications...) } type handlerAccountedHomeDispatcher struct { calls atomic.Int32 } func (*handlerAccountedHomeDispatcher) HeartbeatOK() bool { return true } func (d *handlerAccountedHomeDispatcher) RPopAuth(_ context.Context, model string, _ string, _ http.Header, _ int) ([]byte, error) { d.calls.Add(1) return json.Marshal(map[string]any{ "concurrency": map[string]any{"accounted": true, "credential_id": "handler-cred", "model": model}, "model": model, "auth_index": "handler-cred", "auth": map[string]any{"id": "handler-cred", "provider": "bootstrap-test", "status": coreauth.StatusActive}, }) } func (*handlerAccountedHomeDispatcher) AbortAmbiguousDispatch() {} func TestExecuteStreamWithAuthManager_HomeBootstrapFailureDoesNotRedispatch(t *testing.T) { executor := &bootstrapStreamExecutor{stream: func(_ context.Context, _ int) (*coreexecutor.StreamResult, error) { chunks := make(chan coreexecutor.StreamChunk, 2) chunks <- coreexecutor.StreamChunk{Payload: []byte("drop")} chunks <- coreexecutor.StreamChunk{Err: &coreauth.Error{HTTPStatus: http.StatusUnauthorized, Message: "unauthorized"}} close(chunks) return &coreexecutor.StreamResult{Chunks: chunks}, nil }} manager := coreauth.NewManager(nil, nil, nil) manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) manager.RegisterExecutor(executor) registry := executionregistry.New() releaseSink := newHandlerReleaseSink() registry.SetReleaseSink(releaseSink.MarkDirty) dispatcher := &handlerAccountedHomeDispatcher{} manager.PublishHomeDispatch(dispatcher, registry, 1) handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{Streaming: sdkconfig.StreamingConfig{BootstrapRetries: 1}}, manager) handler.SetPluginHost(&handlerInterceptorTestHost{interceptStreamChunk: func(_ context.Context, req pluginapi.StreamChunkInterceptRequest) pluginapi.StreamChunkInterceptResponse { return pluginapi.StreamChunkInterceptResponse{Body: cloneBytes(req.Body), DropChunk: string(req.Body) == "drop"} }}) dataChan, _, errChan := handler.ExecuteStreamWithAuthManager(context.Background(), "openai", "home-model", []byte(`{"model":"home-model"}`), "") for range dataChan { t.Fatal("Home bootstrap failure produced data") } var streamErr *interfaces.ErrorMessage for msg := range errChan { if msg != nil { streamErr = msg } } if streamErr == nil || streamErr.StatusCode != http.StatusUnauthorized { t.Fatalf("stream error = %+v, want unauthorized terminal error", streamErr) } if got := dispatcher.calls.Load(); got != 1 { t.Fatalf("Home RPOP calls = %d, want 1", got) } select { case <-releaseSink.notified: case <-time.After(time.Second): t.Fatal("accounted Home selection was not released") } wantRelease := handlerReleaseNotification{ group: executionregistry.ReleaseGroup{CredentialID: "handler-cred", Model: "home-model"}, sequence: 1, } if got := releaseSink.Notifications(); len(got) != 1 || got[0] != wantRelease { t.Fatalf("release notifications = %#v, want [%#v]", got, wantRelease) } if errDrain := registry.Drain(context.Background()); errDrain != nil { t.Fatalf("registry.Drain(): %v", errDrain) } if got := releaseSink.Notifications(); len(got) != 1 || got[0] != wantRelease { t.Fatalf("release notifications after drain = %#v, want [%#v]", got, wantRelease) } } func TestExecuteStreamWithAuthManager_HeaderPassthroughDisabledByDefault(t *testing.T) { executor := &failOnceStreamExecutor{} manager := coreauth.NewManager(nil, nil, nil) manager.RegisterExecutor(executor) auth1 := &coreauth.Auth{ ID: "auth1", Provider: "codex", Status: coreauth.StatusActive, Metadata: map[string]any{"email": "test1@example.com"}, } if _, err := manager.Register(context.Background(), auth1); err != nil { t.Fatalf("manager.Register(auth1): %v", err) } auth2 := &coreauth.Auth{ ID: "auth2", Provider: "codex", Status: coreauth.StatusActive, Metadata: map[string]any{"email": "test2@example.com"}, } if _, err := manager.Register(context.Background(), auth2); err != nil { t.Fatalf("manager.Register(auth2): %v", err) } registry.GetGlobalRegistry().RegisterClient(auth1.ID, auth1.Provider, []*registry.ModelInfo{{ID: "test-model"}}) registry.GetGlobalRegistry().RegisterClient(auth2.ID, auth2.Provider, []*registry.ModelInfo{{ID: "test-model"}}) t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth1.ID) registry.GetGlobalRegistry().UnregisterClient(auth2.ID) }) handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{ Streaming: sdkconfig.StreamingConfig{ BootstrapRetries: 1, }, }, manager) dataChan, upstreamHeaders, errChan := handler.ExecuteStreamWithAuthManager(context.Background(), "openai", "test-model", []byte(`{"model":"test-model"}`), "") if dataChan == nil || errChan == nil { t.Fatalf("expected non-nil channels") } var got []byte for chunk := range dataChan { got = append(got, chunk...) } for msg := range errChan { if msg != nil { t.Fatalf("unexpected error: %+v", msg) } } if string(got) != "ok" { t.Fatalf("expected payload ok, got %q", string(got)) } if upstreamHeaders != nil { t.Fatalf("expected nil upstream headers when passthrough is disabled, got %#v", upstreamHeaders) } } func TestExecuteStreamWithAuthManager_DoesNotRetryAfterFirstByte(t *testing.T) { executor := &payloadThenErrorStreamExecutor{} manager := coreauth.NewManager(nil, nil, nil) manager.RegisterExecutor(executor) auth1 := &coreauth.Auth{ ID: "auth1", Provider: "codex", Status: coreauth.StatusActive, Metadata: map[string]any{"email": "test1@example.com"}, } if _, err := manager.Register(context.Background(), auth1); err != nil { t.Fatalf("manager.Register(auth1): %v", err) } auth2 := &coreauth.Auth{ ID: "auth2", Provider: "codex", Status: coreauth.StatusActive, Metadata: map[string]any{"email": "test2@example.com"}, } if _, err := manager.Register(context.Background(), auth2); err != nil { t.Fatalf("manager.Register(auth2): %v", err) } registry.GetGlobalRegistry().RegisterClient(auth1.ID, auth1.Provider, []*registry.ModelInfo{{ID: "test-model"}}) registry.GetGlobalRegistry().RegisterClient(auth2.ID, auth2.Provider, []*registry.ModelInfo{{ID: "test-model"}}) t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth1.ID) registry.GetGlobalRegistry().UnregisterClient(auth2.ID) }) handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{ Streaming: sdkconfig.StreamingConfig{ BootstrapRetries: 1, }, }, manager) dataChan, _, errChan := handler.ExecuteStreamWithAuthManager(context.Background(), "openai", "test-model", []byte(`{"model":"test-model"}`), "") if dataChan == nil || errChan == nil { t.Fatalf("expected non-nil channels") } var got []byte for chunk := range dataChan { got = append(got, chunk...) } var gotErr error var gotStatus int for msg := range errChan { if msg != nil && msg.Error != nil { gotErr = msg.Error gotStatus = msg.StatusCode } } if string(got) != "partial" { t.Fatalf("expected payload partial, got %q", string(got)) } if gotErr == nil { t.Fatalf("expected terminal error, got nil") } if gotStatus != http.StatusBadGateway { t.Fatalf("expected status %d, got %d", http.StatusBadGateway, gotStatus) } if executor.Calls() != 1 { t.Fatalf("expected 1 stream attempt, got %d", executor.Calls()) } } func TestExecuteStreamWithAuthManager_EnrichesBootstrapRetryAuthUnavailableError(t *testing.T) { executor := &failOnceStreamExecutor{} manager := coreauth.NewManager(nil, nil, nil) manager.RegisterExecutor(executor) auth1 := &coreauth.Auth{ ID: "auth1", Provider: "codex", Status: coreauth.StatusActive, Metadata: map[string]any{"email": "test1@example.com"}, } if _, err := manager.Register(context.Background(), auth1); err != nil { t.Fatalf("manager.Register(auth1): %v", err) } registry.GetGlobalRegistry().RegisterClient(auth1.ID, auth1.Provider, []*registry.ModelInfo{{ID: "test-model"}}) t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth1.ID) }) handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{ Streaming: sdkconfig.StreamingConfig{ BootstrapRetries: 1, }, }, manager) dataChan, _, errChan := handler.ExecuteStreamWithAuthManager(context.Background(), "openai", "test-model", []byte(`{"model":"test-model"}`), "") if dataChan == nil || errChan == nil { t.Fatalf("expected non-nil channels") } var got []byte for chunk := range dataChan { got = append(got, chunk...) } if len(got) != 0 { t.Fatalf("expected empty payload, got %q", string(got)) } var gotErr *interfaces.ErrorMessage for msg := range errChan { if msg != nil { gotErr = msg } } if gotErr == nil { t.Fatalf("expected terminal error") } if gotErr.StatusCode != http.StatusServiceUnavailable { t.Fatalf("status = %d, want %d", gotErr.StatusCode, http.StatusServiceUnavailable) } var authErr *coreauth.Error if !errors.As(gotErr.Error, &authErr) || authErr == nil { t.Fatalf("expected coreauth.Error, got %T", gotErr.Error) } if authErr.Code != "auth_unavailable" { t.Fatalf("code = %q, want %q", authErr.Code, "auth_unavailable") } if !strings.Contains(authErr.Message, "providers=codex") { t.Fatalf("message missing provider context: %q", authErr.Message) } if !strings.Contains(authErr.Message, "model=test-model") { t.Fatalf("message missing model context: %q", authErr.Message) } if executor.Calls() != 1 { t.Fatalf("expected exactly one upstream call before retry path selection failure, got %d", executor.Calls()) } } func TestExecuteStreamWithAuthManager_PinnedAuthKeepsSameUpstream(t *testing.T) { executor := &authAwareStreamExecutor{} manager := coreauth.NewManager(nil, nil, nil) manager.RegisterExecutor(executor) auth1 := &coreauth.Auth{ ID: "auth1", Provider: "codex", Status: coreauth.StatusActive, Metadata: map[string]any{"email": "test1@example.com"}, } if _, err := manager.Register(context.Background(), auth1); err != nil { t.Fatalf("manager.Register(auth1): %v", err) } auth2 := &coreauth.Auth{ ID: "auth2", Provider: "codex", Status: coreauth.StatusActive, Metadata: map[string]any{"email": "test2@example.com"}, } if _, err := manager.Register(context.Background(), auth2); err != nil { t.Fatalf("manager.Register(auth2): %v", err) } registry.GetGlobalRegistry().RegisterClient(auth1.ID, auth1.Provider, []*registry.ModelInfo{{ID: "test-model"}}) registry.GetGlobalRegistry().RegisterClient(auth2.ID, auth2.Provider, []*registry.ModelInfo{{ID: "test-model"}}) t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth1.ID) registry.GetGlobalRegistry().UnregisterClient(auth2.ID) }) handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{ Streaming: sdkconfig.StreamingConfig{ BootstrapRetries: 1, }, }, manager) ctx := WithPinnedAuthID(context.Background(), "auth1") dataChan, _, errChan := handler.ExecuteStreamWithAuthManager(ctx, "openai", "test-model", []byte(`{"model":"test-model"}`), "") if dataChan == nil || errChan == nil { t.Fatalf("expected non-nil channels") } var got []byte for chunk := range dataChan { got = append(got, chunk...) } var gotErr error for msg := range errChan { if msg != nil && msg.Error != nil { gotErr = msg.Error } } if len(got) != 0 { t.Fatalf("expected empty payload, got %q", string(got)) } if gotErr == nil { t.Fatalf("expected terminal error, got nil") } authIDs := executor.AuthIDs() if len(authIDs) == 0 { t.Fatalf("expected at least one upstream attempt") } for _, authID := range authIDs { if authID != "auth1" { t.Fatalf("expected all attempts on auth1, got sequence %v", authIDs) } } } func TestExecuteStreamWithAuthManager_SelectedAuthCallbackReceivesAuthID(t *testing.T) { executor := &authAwareStreamExecutor{} manager := coreauth.NewManager(nil, nil, nil) manager.RegisterExecutor(executor) auth2 := &coreauth.Auth{ ID: "auth2", Provider: "codex", Status: coreauth.StatusActive, Metadata: map[string]any{"email": "test2@example.com"}, } if _, err := manager.Register(context.Background(), auth2); err != nil { t.Fatalf("manager.Register(auth2): %v", err) } registry.GetGlobalRegistry().RegisterClient(auth2.ID, auth2.Provider, []*registry.ModelInfo{{ID: "test-model"}}) t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth2.ID) }) handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{ Streaming: sdkconfig.StreamingConfig{ BootstrapRetries: 0, }, }, manager) gin.SetMode(gin.TestMode) recorder := httptest.NewRecorder() ginCtx, _ := gin.CreateTestContext(recorder) ginCtx.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) logging.SetGinRequestID(ginCtx, "1234abcd") selectedAuthID := "" ctx := context.WithValue(context.Background(), "gin", ginCtx) ctx = WithSelectedAuthIDCallback(ctx, func(authID string) { selectedAuthID = authID }) dataChan, _, errChan := handler.ExecuteStreamWithAuthManager(ctx, "openai", "test-model", []byte(`{"model":"test-model"}`), "") if dataChan == nil || errChan == nil { t.Fatalf("expected non-nil channels") } var got []byte for chunk := range dataChan { got = append(got, chunk...) } for msg := range errChan { if msg != nil { t.Fatalf("unexpected error: %+v", msg) } } if string(got) != "ok" { t.Fatalf("expected payload ok, got %q", string(got)) } if selectedAuthID != "auth2" { t.Fatalf("selectedAuthID = %q, want %q", selectedAuthID, "auth2") } traceID := logging.GetGinCPATraceID(ginCtx) parts := strings.Split(traceID, "-") if len(parts) != 3 || parts[1] != auth2.Index || parts[2] != "1234abcd" { t.Fatalf("trace ID = %q, want timestamp-%s-1234abcd", traceID, auth2.Index) } if _, errParse := time.Parse("20060102150405", parts[0]); errParse != nil { t.Fatalf("trace timestamp = %q: %v", parts[0], errParse) } } func TestExecuteStreamWithAuthManager_ValidatesOpenAIResponsesStreamDataJSON(t *testing.T) { executor := &invalidJSONStreamExecutor{} manager := coreauth.NewManager(nil, nil, nil) manager.RegisterExecutor(executor) auth1 := &coreauth.Auth{ ID: "auth1", Provider: "codex", Status: coreauth.StatusActive, Metadata: map[string]any{"email": "test1@example.com"}, } if _, err := manager.Register(context.Background(), auth1); err != nil { t.Fatalf("manager.Register(auth1): %v", err) } registry.GetGlobalRegistry().RegisterClient(auth1.ID, auth1.Provider, []*registry.ModelInfo{{ID: "test-model"}}) t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth1.ID) }) handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager) dataChan, _, errChan := handler.ExecuteStreamWithAuthManager(context.Background(), "openai-response", "test-model", []byte(`{"model":"test-model"}`), "") if dataChan == nil || errChan == nil { t.Fatalf("expected non-nil channels") } var got []byte for chunk := range dataChan { got = append(got, chunk...) } if len(got) != 0 { t.Fatalf("expected empty payload, got %q", string(got)) } gotErr := false for msg := range errChan { if msg == nil { continue } if msg.StatusCode != http.StatusBadGateway { t.Fatalf("expected status %d, got %d", http.StatusBadGateway, msg.StatusCode) } if msg.Error == nil { t.Fatalf("expected error") } gotErr = true } if !gotErr { t.Fatalf("expected terminal error") } } func TestExecuteStreamWithAuthManager_AllowsSplitOpenAIResponsesSSEEventLines(t *testing.T) { executor := &splitResponsesEventStreamExecutor{} manager := coreauth.NewManager(nil, nil, nil) manager.RegisterExecutor(executor) auth1 := &coreauth.Auth{ ID: "auth1", Provider: "split-sse", Status: coreauth.StatusActive, Metadata: map[string]any{"email": "test1@example.com"}, } if _, err := manager.Register(context.Background(), auth1); err != nil { t.Fatalf("manager.Register(auth1): %v", err) } registry.GetGlobalRegistry().RegisterClient(auth1.ID, auth1.Provider, []*registry.ModelInfo{{ID: "test-model"}}) t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth1.ID) }) handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager) dataChan, _, errChan := handler.ExecuteStreamWithAuthManager(context.Background(), "openai-response", "test-model", []byte(`{"model":"test-model"}`), "") if dataChan == nil || errChan == nil { t.Fatalf("expected non-nil channels") } var got []string for chunk := range dataChan { got = append(got, string(chunk)) } for msg := range errChan { if msg != nil { t.Fatalf("unexpected error: %+v", msg) } } if len(got) != 2 { t.Fatalf("expected 2 forwarded chunks, got %d: %#v", len(got), got) } if got[0] != "event: response.completed" { t.Fatalf("unexpected first chunk: %q", got[0]) } expectedData := "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp-1\",\"output\":[]}}" if got[1] != expectedData { t.Fatalf("unexpected second chunk.\nGot: %q\nWant: %q", got[1], expectedData) } }