From 411d7d41eee0ff841b5badba5abacad6c11331ef Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Tue, 14 Jul 2026 16:31:45 +0800 Subject: [PATCH] feat(auth): add `SelectAuthByKind` method and enhance Codex Alpha Search handling - Implemented `SelectAuthByKind` to filter and select credentials by kind with retry logic and error handling. - Updated Codex Alpha Search to enforce OAuth credentials and sanitize upstream request body. - Added unit tests to validate `SelectAuthByKind` behavior across scenarios. - Introduced `sanitizeCodexAlphaSearchBody` for removing sensitive fields from API payloads. Closes #4289 --- internal/api/server.go | 32 +++++- internal/api/server_test.go | 125 ++++++++++++++++++++++++ sdk/cliproxy/auth/conductor.go | 40 ++++++++ sdk/cliproxy/auth/scheduler_test.go | 145 ++++++++++++++++++++++++++++ 4 files changed, 339 insertions(+), 3 deletions(-) diff --git a/internal/api/server.go b/internal/api/server.go index 5893bc0d..73650687 100644 --- a/internal/api/server.go +++ b/internal/api/server.go @@ -553,6 +553,7 @@ func (s *Server) setupRoutes() { codexDirect.GET("/responses", openaiResponsesHandlers.ResponsesWebsocket) codexDirect.POST("/responses", openaiResponsesHandlers.Responses) codexDirect.POST("/responses/compact", openaiResponsesHandlers.Compact) + codexDirect.POST("/alpha/search", s.codexAlphaSearch) } // Gemini compatible API routes @@ -625,6 +626,30 @@ func (s *Server) setupRoutes() { // Management routes are registered lazily by registerManagementRoutes when a secret is configured. } +func sanitizeCodexAlphaSearchBody(body []byte) []byte { + var payload map[string]json.RawMessage + if errUnmarshal := json.Unmarshal(body, &payload); errUnmarshal != nil || payload == nil { + return body + } + + removed := false + for _, field := range []string{"prompt_cache_key", "prompt_cache_retention"} { + if _, exists := payload[field]; exists { + delete(payload, field) + removed = true + } + } + if !removed { + return body + } + + sanitizedBody, errMarshal := json.Marshal(payload) + if errMarshal != nil { + return body + } + return sanitizedBody +} + // codexAlphaSearch forwards the standalone search endpoint used by current // Codex clients. Unlike /responses, this payload is already in Codex search // format and must not pass through a protocol translator. @@ -645,13 +670,14 @@ func (s *Server) codexAlphaSearch(c *gin.Context) { Model string `json:"model"` } _ = json.Unmarshal(body, &routing) + upstreamRequestBody := sanitizeCodexAlphaSearchBody(body) selectionHeaders := c.Request.Header.Clone() if sessionID := strings.TrimSpace(routing.ID); sessionID != "" { selectionHeaders.Set("X-Session-ID", sessionID) } ctx := context.WithValue(c.Request.Context(), "gin", c) - selected, err := s.handlers.AuthManager.SelectAuth(ctx, "codex", strings.TrimSpace(routing.Model), coreexecutor.Options{ + selected, err := s.handlers.AuthManager.SelectAuthByKind(ctx, "codex", strings.TrimSpace(routing.Model), auth.AuthKindOAuth, coreexecutor.Options{ Headers: selectionHeaders, OriginalRequest: body, }) @@ -679,7 +705,7 @@ func (s *Server) codexAlphaSearch(c *gin.Context) { const upstreamURL = "https://chatgpt.com/backend-api/codex/alpha/search" req, err := s.handlers.AuthManager.NewHttpRequest( - ctx, selected, http.MethodPost, upstreamURL, body, headers, + ctx, selected, http.MethodPost, upstreamURL, upstreamRequestBody, headers, ) if err != nil { c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()}) @@ -697,7 +723,7 @@ func (s *Server) codexAlphaSearch(c *gin.Context) { URL: upstreamURL, Method: http.MethodPost, Headers: helpHeaders, - Body: body, + Body: upstreamRequestBody, Provider: "codex", AuthID: authID, AuthLabel: authLabel, diff --git a/internal/api/server_test.go b/internal/api/server_test.go index bb796c2b..1f9c6a3c 100644 --- a/internal/api/server_test.go +++ b/internal/api/server_test.go @@ -67,6 +67,20 @@ func (s *codexSearchGinContextSelector) Pick(ctx context.Context, _ string, _ st return auths[0], nil } +type codexSearchAPIKeyFirstSelector struct{} + +func (s *codexSearchAPIKeyFirstSelector) Pick(_ context.Context, _ string, _ string, _ coreexecutor.Options, auths []*auth.Auth) (*auth.Auth, error) { + for _, candidate := range auths { + if candidate.AuthKind() == auth.AuthKindAPIKey { + return candidate, nil + } + } + if len(auths) == 0 { + return nil, nil + } + return auths[0], nil +} + func (e *codexSearchCaptureExecutor) HttpRequest(_ context.Context, selected *auth.Auth, req *http.Request) (*http.Response, error) { e.request = req.Clone(req.Context()) e.authIDs = append(e.authIDs, selected.ID) @@ -200,6 +214,117 @@ func TestCodexAlphaSearchForwardsRequest(t *testing.T) { } } +func TestCodexAlphaSearchSanitizesResponsesOnlyFields(t *testing.T) { + server := newTestServer(t) + executor := &codexSearchCaptureExecutor{} + server.handlers.AuthManager.RegisterExecutor(executor) + credential := &auth.Auth{ + ID: "codex-auth", + Provider: "codex", + Status: auth.StatusActive, + Metadata: map[string]any{"access_token": "codex-token"}, + } + if _, errRegister := server.handlers.AuthManager.Register(context.Background(), credential); errRegister != nil { + t.Fatalf("register Codex auth: %v", errRegister) + } + registry.GetGlobalRegistry().RegisterClient(credential.ID, credential.Provider, []*registry.ModelInfo{{ID: "gpt-5.6-sol"}}) + t.Cleanup(func() { + registry.GetGlobalRegistry().UnregisterClient(credential.ID) + }) + + payload := `{"id":"session-123","model":"gpt-5.6-sol","commands":{"search_query":[{"q":"golang channels"}]},"prompt_cache_key":"cache-123","prompt_cache_retention":"24h"}` + for _, path := range []string{"/v1/alpha/search", "/backend-api/codex/alpha/search"} { + t.Run(path, func(t *testing.T) { + req := httptest.NewRequest(http.MethodPost, path, strings.NewReader(payload)) + req.Header.Set("Authorization", "Bearer test-key") + req.Header.Set("Content-Type", "application/json") + rr := httptest.NewRecorder() + server.engine.ServeHTTP(rr, req) + + if rr.Code != http.StatusOK { + t.Fatalf("status = %d, want %d; body=%s", rr.Code, http.StatusOK, rr.Body.String()) + } + var upstreamBody map[string]json.RawMessage + if errUnmarshal := json.Unmarshal(executor.body, &upstreamBody); errUnmarshal != nil { + t.Fatalf("unmarshal upstream body: %v; body=%s", errUnmarshal, executor.body) + } + if _, exists := upstreamBody["prompt_cache_key"]; exists { + t.Fatalf("upstream body contains prompt_cache_key: %s", executor.body) + } + if _, exists := upstreamBody["prompt_cache_retention"]; exists { + t.Fatalf("upstream body contains prompt_cache_retention: %s", executor.body) + } + for _, field := range []string{"id", "model", "commands"} { + if _, exists := upstreamBody[field]; !exists { + t.Fatalf("upstream body missing %s: %s", field, executor.body) + } + } + }) + } +} + +func TestCodexAlphaSearchRequiresOAuthCredential(t *testing.T) { + newServer := func(t *testing.T, credentials ...*auth.Auth) (*Server, *codexSearchCaptureExecutor) { + t.Helper() + server := newTestServer(t) + server.handlers.AuthManager.SetSelector(&codexSearchAPIKeyFirstSelector{}) + executor := &codexSearchCaptureExecutor{} + server.handlers.AuthManager.RegisterExecutor(executor) + for _, credential := range credentials { + if _, errRegister := server.handlers.AuthManager.Register(context.Background(), credential); errRegister != nil { + t.Fatalf("register Codex auth %s: %v", credential.ID, errRegister) + } + } + return server, executor + } + apiKeyCredential := func() *auth.Auth { + return &auth.Auth{ + ID: "codex-api-key", + Provider: "codex", + Status: auth.StatusActive, + Attributes: map[string]string{auth.AttributeAPIKey: "codex-key"}, + } + } + oauthCredential := func() *auth.Auth { + return &auth.Auth{ + ID: "codex-oauth", + Provider: "codex", + Status: auth.StatusActive, + Metadata: map[string]any{"access_token": "codex-token"}, + } + } + + t.Run("mixed credentials", func(t *testing.T) { + server, executor := newServer(t, apiKeyCredential(), oauthCredential()) + req := httptest.NewRequest(http.MethodPost, "/v1/alpha/search", strings.NewReader(`{"query":"GPT-5.6"}`)) + req.Header.Set("Authorization", "Bearer test-key") + rr := httptest.NewRecorder() + server.engine.ServeHTTP(rr, req) + + if rr.Code != http.StatusOK { + t.Fatalf("status = %d, want %d; body=%s", rr.Code, http.StatusOK, rr.Body.String()) + } + if got := executor.authIDs; len(got) != 1 || got[0] != "codex-oauth" { + t.Fatalf("selected auth IDs = %v, want [codex-oauth]", got) + } + }) + + t.Run("API key only", func(t *testing.T) { + server, executor := newServer(t, apiKeyCredential()) + req := httptest.NewRequest(http.MethodPost, "/v1/alpha/search", strings.NewReader(`{"query":"GPT-5.6"}`)) + req.Header.Set("Authorization", "Bearer test-key") + rr := httptest.NewRecorder() + server.engine.ServeHTTP(rr, req) + + if rr.Code != http.StatusServiceUnavailable { + t.Fatalf("status = %d, want %d; body=%s", rr.Code, http.StatusServiceUnavailable, rr.Body.String()) + } + if len(executor.authIDs) != 0 { + t.Fatalf("selected auth IDs = %v, want none", executor.authIDs) + } + }) +} + func TestCodexAlphaSearchPassesGinContextToAuthSelection(t *testing.T) { server := newTestServer(t) selector := &codexSearchGinContextSelector{} diff --git a/sdk/cliproxy/auth/conductor.go b/sdk/cliproxy/auth/conductor.go index 042fdb7d..04e87f26 100644 --- a/sdk/cliproxy/auth/conductor.go +++ b/sdk/cliproxy/auth/conductor.go @@ -4655,6 +4655,46 @@ func (m *Manager) SelectAuth(ctx context.Context, provider, model string, opts c return selected, nil } +// SelectAuthByKind selects one credential of the required kind through the +// configured scheduling strategy. Credentials of other kinds are skipped. +func (m *Manager) SelectAuthByKind(ctx context.Context, provider, model, requiredKind string, opts cliproxyexecutor.Options) (*Auth, error) { + requiredKind = normalizeAuthKind(requiredKind) + if requiredKind == "" { + return nil, &Error{Code: "invalid_auth_kind", Message: "required auth kind is invalid", HTTPStatus: http.StatusBadRequest} + } + + homeMode := m.HomeEnabled() + homeAuthCount := homeAuthCountFromMetadata(opts.Metadata) + tried := make(map[string]struct{}) + for { + pickOpts := opts + if homeMode { + pickOpts = withHomeAuthCount(opts, homeAuthCount) + } + selected, _, errPick := m.pickNext(ctx, provider, model, pickOpts, tried) + if errPick != nil { + return nil, errPick + } + if selected == nil { + return nil, &Error{Code: "auth_not_found", Message: "selector returned no auth"} + } + if selected.AuthKind() == requiredKind { + return selected, nil + } + authID := strings.TrimSpace(selected.ID) + if authID == "" { + return nil, &Error{Code: "auth_not_found", Message: "selected auth has no ID"} + } + if _, alreadyTried := tried[authID]; alreadyTried { + return nil, &Error{Code: "auth_not_found", Message: "selector repeatedly returned an ineligible auth"} + } + tried[authID] = struct{}{} + if homeMode { + homeAuthCount++ + } + } +} + func (m *Manager) pickNext(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, tried map[string]struct{}) (*Auth, ProviderExecutor, error) { if m.HomeEnabled() { auth, exec, _, err := m.pickNextViaHome(ctx, model, opts, tried) diff --git a/sdk/cliproxy/auth/scheduler_test.go b/sdk/cliproxy/auth/scheduler_test.go index 99f4f9dc..03296517 100644 --- a/sdk/cliproxy/auth/scheduler_test.go +++ b/sdk/cliproxy/auth/scheduler_test.go @@ -2,12 +2,14 @@ package auth import ( "context" + "encoding/json" "errors" "net/http" "testing" "time" internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/internal/home" "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" @@ -59,6 +61,23 @@ type inactivePluginScheduler struct { fakePluginScheduler } +type authKindHomeDispatcher struct { + auths []Auth + counts []int +} + +func (d *authKindHomeDispatcher) HeartbeatOK() bool { + return true +} + +func (d *authKindHomeDispatcher) RPopAuth(_ context.Context, _ string, _ string, _ http.Header, count int) ([]byte, error) { + d.counts = append(d.counts, count) + if count < 1 || count > len(d.auths) { + return nil, home.ErrAuthNotFound + } + return json.Marshal(homeAuthDispatchResponse{Auth: d.auths[count-1]}) +} + func (s *inactivePluginScheduler) HasScheduler() bool { return false } @@ -427,6 +446,132 @@ func TestManagerPluginSchedulerSelectsAuthID(t *testing.T) { } } +func TestManagerSelectAuthByKindSkipsAPIKey(t *testing.T) { + manager := NewManager(nil, &RoundRobinSelector{}, nil) + manager.executors["codex"] = schedulerTestExecutor{} + for _, candidate := range []*Auth{ + {ID: "codex-api-key", Provider: "codex", Attributes: map[string]string{AttributeAPIKey: "test-key"}}, + {ID: "codex-oauth", Provider: "codex", Metadata: map[string]any{"access_token": "test-token"}}, + } { + if _, errRegister := manager.Register(context.Background(), candidate); errRegister != nil { + t.Fatalf("Register(%s) error = %v", candidate.ID, errRegister) + } + } + + scheduler := &fakePluginScheduler{ + resp: pluginapi.SchedulerPickResponse{Handled: true, AuthID: "codex-api-key"}, + handled: true, + } + manager.SetPluginScheduler(scheduler) + + selected, errSelect := manager.SelectAuthByKind(context.Background(), "codex", "", AuthKindOAuth, cliproxyexecutor.Options{}) + if errSelect != nil { + t.Fatalf("SelectAuthByKind() error = %v", errSelect) + } + if selected == nil || selected.ID != "codex-oauth" { + t.Fatalf("SelectAuthByKind() auth = %#v, want codex-oauth", selected) + } + if scheduler.calls != 2 { + t.Fatalf("scheduler.calls = %d, want 2", scheduler.calls) + } +} + +func TestManagerSelectAuthByKindReturnsErrorWhenUnavailable(t *testing.T) { + manager := NewManager(nil, &RoundRobinSelector{}, nil) + manager.executors["codex"] = schedulerTestExecutor{} + if _, errRegister := manager.Register(context.Background(), &Auth{ + ID: "codex-api-key", + Provider: "codex", + Attributes: map[string]string{AttributeAPIKey: "test-key"}, + }); errRegister != nil { + t.Fatalf("Register(codex-api-key) error = %v", errRegister) + } + + selected, errSelect := manager.SelectAuthByKind(context.Background(), "codex", "", AuthKindOAuth, cliproxyexecutor.Options{}) + if selected != nil { + t.Fatalf("SelectAuthByKind() auth = %#v, want nil", selected) + } + var authErr *Error + if !errors.As(errSelect, &authErr) || authErr.Code != "auth_not_found" { + t.Fatalf("SelectAuthByKind() error = %#v, want auth_not_found", errSelect) + } +} + +func TestManagerSelectAuthByKindRejectsInvalidKind(t *testing.T) { + manager := NewManager(nil, &RoundRobinSelector{}, nil) + selected, errSelect := manager.SelectAuthByKind(context.Background(), "codex", "", "certificate", cliproxyexecutor.Options{}) + if selected != nil { + t.Fatalf("SelectAuthByKind() auth = %#v, want nil", selected) + } + var authErr *Error + if !errors.As(errSelect, &authErr) || authErr.Code != "invalid_auth_kind" || authErr.HTTPStatus != http.StatusBadRequest { + t.Fatalf("SelectAuthByKind() error = %#v, want invalid_auth_kind", errSelect) + } +} + +func TestManagerSelectAuthByKindAdvancesHomeAuthCount(t *testing.T) { + tests := []struct { + name string + auths []Auth + wantAuthID string + wantError string + }{ + { + name: "skips API key for OAuth", + auths: []Auth{ + {ID: "home-api-key", Provider: "test", Attributes: map[string]string{AttributeAPIKey: "test-key"}}, + {ID: "home-oauth", Provider: "test", Metadata: map[string]any{"access_token": "test-token"}}, + }, + wantAuthID: "home-oauth", + }, + { + name: "returns not found without OAuth", + auths: []Auth{ + {ID: "home-api-key", Provider: "test", Attributes: map[string]string{AttributeAPIKey: "test-key"}}, + }, + wantError: "auth_not_found", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + dispatcher := &authKindHomeDispatcher{auths: tt.auths} + oldCurrentHomeDispatcher := currentHomeDispatcher + currentHomeDispatcher = func() homeAuthDispatcher { + return dispatcher + } + t.Cleanup(func() { + currentHomeDispatcher = oldCurrentHomeDispatcher + }) + + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.RegisterExecutor(schedulerTestExecutor{}) + + selected, errSelect := manager.SelectAuthByKind(context.Background(), "codex", "", AuthKindOAuth, cliproxyexecutor.Options{}) + if tt.wantError != "" { + if selected != nil { + t.Fatalf("SelectAuthByKind() auth = %#v, want nil", selected) + } + var authErr *Error + if !errors.As(errSelect, &authErr) || authErr.Code != tt.wantError { + t.Fatalf("SelectAuthByKind() error = %#v, want %s", errSelect, tt.wantError) + } + } else { + if errSelect != nil { + t.Fatalf("SelectAuthByKind() error = %v", errSelect) + } + if selected == nil || selected.ID != tt.wantAuthID { + t.Fatalf("SelectAuthByKind() auth = %#v, want %s", selected, tt.wantAuthID) + } + } + if len(dispatcher.counts) != 2 || dispatcher.counts[0] != 1 || dispatcher.counts[1] != 2 { + t.Fatalf("home auth counts = %v, want [1 2]", dispatcher.counts) + } + }) + } +} + func TestManagerPluginSchedulerSkippedWhenHomeEnabled(t *testing.T) { manager := NewManager(nil, &RoundRobinSelector{}, nil) manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})