diff --git a/config.example.yaml b/config.example.yaml index 027b7acb..627c4f4f 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -332,6 +332,7 @@ nonstream-keepalive-interval: 0 # prefix: "test" # optional: require calls like "test/gpt-5-codex" to target this credential # disable-cooling: false # optional: per-auth override for auth/model cooldown scheduling # base-url: "https://www.example.com" # use the custom codex API endpoint +# alpha-search: false # optional: allow this key to serve /v1/alpha/search via base-url + /alpha/search # headers: # X-Custom-Header: "custom-value" # proxy-url: "socks5://proxy.example.com:1080" # optional: per-key proxy override diff --git a/internal/api/handlers/management/config_codex_alpha_search_test.go b/internal/api/handlers/management/config_codex_alpha_search_test.go new file mode 100644 index 00000000..5c3cc0b7 --- /dev/null +++ b/internal/api/handlers/management/config_codex_alpha_search_test.go @@ -0,0 +1,34 @@ +package management + +import ( + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/gin-gonic/gin" + "github.com/router-for-me/CLIProxyAPI/v7/internal/config" +) + +func TestPatchCodexKeyUpdatesAlphaSearch(t *testing.T) { + h := &Handler{ + cfg: &config.Config{CodexKey: []config.CodexKey{{ + APIKey: "codex-key", + BaseURL: "https://codex.example.com", + }}}, + configFilePath: writeTestConfigFile(t), + } + + rec := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(rec) + ctx.Request = httptest.NewRequest(http.MethodPatch, "/v0/management/codex-api-key", strings.NewReader(`{"index":0,"value":{"alpha-search":true}}`)) + ctx.Request.Header.Set("Content-Type", "application/json") + h.PatchCodexKey(ctx) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want %d; body=%s", rec.Code, http.StatusOK, rec.Body.String()) + } + if !h.cfg.CodexKey[0].AlphaSearch { + t.Fatal("alpha-search = false, want true") + } +} diff --git a/internal/api/handlers/management/config_lists.go b/internal/api/handlers/management/config_lists.go index 541184ff..1fdbc2e9 100644 --- a/internal/api/handlers/management/config_lists.go +++ b/internal/api/handlers/management/config_lists.go @@ -1224,6 +1224,7 @@ func (h *Handler) PatchCodexKey(c *gin.Context) { Prefix *string `json:"prefix"` BaseURL *string `json:"base-url"` ProxyURL *string `json:"proxy-url"` + AlphaSearch *bool `json:"alpha-search"` Models *[]config.CodexModel `json:"models"` Headers *map[string]string `json:"headers"` ExcludedModels *[]string `json:"excluded-models"` @@ -1286,6 +1287,9 @@ func (h *Handler) PatchCodexKey(c *gin.Context) { if body.Value.ProxyURL != nil { entry.ProxyURL = strings.TrimSpace(*body.Value.ProxyURL) } + if body.Value.AlphaSearch != nil { + entry.AlphaSearch = *body.Value.AlphaSearch + } if body.Value.Models != nil { entry.Models = append([]config.CodexModel(nil), (*body.Value.Models)...) } diff --git a/internal/api/server_routes.go b/internal/api/server_routes.go index 98127711..bb35e37e 100644 --- a/internal/api/server_routes.go +++ b/internal/api/server_routes.go @@ -294,12 +294,12 @@ func (s *Server) codexAlphaSearch(c *gin.Context) { var selection *auth.HomeDispatchSelection var selected *auth.Auth if s.handlers.AuthManager.HomeEnabled() { - selection, err = s.handlers.AuthManager.SelectHomeAuthByKind(ctx, "codex", selectionModel, auth.AuthKindOAuth, selectionOpts) + selection, err = s.handlers.AuthManager.SelectHomeAuthWithCredentialPolicy(ctx, "codex", selectionModel, auth.CredentialPolicyCodexAlphaSearchV1, selectionOpts) if selection != nil { selected = selection.CloneAuth() } } else { - selected, err = s.handlers.AuthManager.SelectAuthByKind(ctx, "codex", selectionModel, auth.AuthKindOAuth, selectionOpts) + selected, err = s.handlers.AuthManager.SelectAuthWithCredentialPolicy(ctx, "codex", selectionModel, auth.CredentialPolicyCodexAlphaSearchV1, selectionOpts) } if err != nil { status := http.StatusServiceUnavailable @@ -346,7 +346,21 @@ func (s *Server) codexAlphaSearch(c *gin.Context) { headers.Set("Chatgpt-Account-Id", accountID) } - const upstreamURL = "https://chatgpt.com/backend-api/codex/alpha/search" + upstreamURL := "https://chatgpt.com/backend-api/codex/alpha/search" + if selected.AuthKind() == auth.AuthKindAPIKey { + baseURL := "" + if selected.Attributes != nil { + baseURL = strings.TrimSpace(selected.Attributes["base_url"]) + } + if baseURL == "" { + if selection != nil { + selection.End("missing_base_url") + } + c.JSON(http.StatusServiceUnavailable, gin.H{"error": "Codex Alpha Search API key base URL unavailable"}) + return + } + upstreamURL = strings.TrimRight(baseURL, "/") + "/alpha/search" + } req, err := s.handlers.AuthManager.NewHttpRequest( ctx, selected, http.MethodPost, upstreamURL, upstreamRequestBody, headers, ) diff --git a/internal/api/server_test.go b/internal/api/server_test.go index fb1c202a..511a94f7 100644 --- a/internal/api/server_test.go +++ b/internal/api/server_test.go @@ -63,7 +63,12 @@ func (e *codexSearchCaptureExecutor) PrepareRequest(req *http.Request, a *auth.A return e.prepareErr } token, _ := a.Metadata["access_token"].(string) - req.Header.Set("Authorization", "Bearer "+token) + if strings.TrimSpace(token) == "" && a.Attributes != nil { + token = a.Attributes[auth.AttributeAPIKey] + } + if strings.TrimSpace(token) != "" { + req.Header.Set("Authorization", "Bearer "+token) + } return nil } @@ -127,7 +132,8 @@ func (e *codexSearchCaptureExecutor) HttpRequest(_ context.Context, selected *au } type codexSearchHomeDispatcher struct { - calls atomic.Int32 + calls atomic.Int32 + policy atomic.Value } func (*codexSearchHomeDispatcher) HeartbeatOK() bool { return true } @@ -151,6 +157,11 @@ func (d *codexSearchHomeDispatcher) RPopAuth(_ context.Context, model string, _ }) } +func (d *codexSearchHomeDispatcher) RPopAuthWithPolicy(ctx context.Context, model string, sessionID string, headers http.Header, count int, policy string) ([]byte, error) { + d.policy.Store(policy) + return d.RPopAuth(ctx, model, sessionID, headers, count) +} + func (*codexSearchHomeDispatcher) AbortAmbiguousDispatch() {} type codexSearchBusyHomeDispatcher struct{} @@ -159,6 +170,9 @@ func (*codexSearchBusyHomeDispatcher) HeartbeatOK() bool { return true } func (*codexSearchBusyHomeDispatcher) RPopAuth(context.Context, string, string, http.Header, int) ([]byte, error) { return []byte(`{"error":{"type":"credential_concurrency_exceeded","message":"busy","retry_after_ms":750}}`), nil } +func (d *codexSearchBusyHomeDispatcher) RPopAuthWithPolicy(ctx context.Context, model string, sessionID string, headers http.Header, count int, _ string) ([]byte, error) { + return d.RPopAuth(ctx, model, sessionID, headers, count) +} func (*codexSearchBusyHomeDispatcher) AbortAmbiguousDispatch() {} type trackedSearchResponseBody struct { @@ -349,6 +363,9 @@ func TestHomeCodexAlphaSearchEndsSelectionAcrossDirectHTTPPaths(t *testing.T) { if got := dispatcher.calls.Load(); got != 1 { t.Fatalf("Home RPOP calls = %d, want 1", got) } + if got, _ := dispatcher.policy.Load().(string); got != auth.CredentialPolicyCodexAlphaSearchV1 { + t.Fatalf("Home credential policy = %q, want %q", got, auth.CredentialPolicyCodexAlphaSearchV1) + } if got := body.closed.Load(); got != test.wantClosed { t.Fatalf("response body closed = %t, want %t", got, test.wantClosed) } @@ -737,7 +754,7 @@ func TestCodexAlphaSearchSanitizesResponsesOnlyFields(t *testing.T) { } } -func TestCodexAlphaSearchRequiresOAuthCredential(t *testing.T) { +func TestCodexAlphaSearchCredentialPolicy(t *testing.T) { newServer := func(t *testing.T, credentials ...*auth.Auth) (*Server, *codexSearchCaptureExecutor) { t.Helper() server := newTestServer(t) @@ -783,7 +800,7 @@ func TestCodexAlphaSearchRequiresOAuthCredential(t *testing.T) { } }) - t.Run("API key only", func(t *testing.T) { + t.Run("ordinary 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") @@ -799,6 +816,82 @@ func TestCodexAlphaSearchRequiresOAuthCredential(t *testing.T) { }) } +func TestCodexAlphaSearchOptInAPIKeyUsesConfiguredEndpoint(t *testing.T) { + server := newTestServer(t) + executor := &codexSearchCaptureExecutor{} + server.handlers.AuthManager.RegisterExecutor(executor) + credential := &auth.Auth{ + ID: "codex-alpha-api-key", + Provider: "codex", + Status: auth.StatusActive, + Attributes: map[string]string{ + auth.AttributeAPIKey: "codex-alpha-key", + auth.AttributeCodexAlphaSearch: "true", + "base_url": "https://codex.example.com/v1/", + }, + } + if _, errRegister := server.handlers.AuthManager.Register(context.Background(), credential); errRegister != nil { + t.Fatalf("register Codex API key: %v", errRegister) + } + + payload := `{"query":"golang","prompt_cache_key":"cache","prompt_cache_retention":"24h"}` + req := httptest.NewRequest(http.MethodPost, "/v1/alpha/search", strings.NewReader(payload)) + 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 executor.request == nil { + t.Fatal("Codex executor did not receive a request") + } + if got, want := executor.request.URL.String(), "https://codex.example.com/v1/alpha/search"; got != want { + t.Fatalf("upstream URL = %q, want %q", got, want) + } + if got := executor.request.Header.Get("Authorization"); got != "Bearer codex-alpha-key" { + t.Fatalf("Authorization = %q, want API key bearer", got) + } + var upstreamBody map[string]json.RawMessage + if errUnmarshal := json.Unmarshal(executor.body, &upstreamBody); errUnmarshal != nil { + t.Fatalf("unmarshal upstream body: %v", errUnmarshal) + } + for _, field := range []string{"prompt_cache_key", "prompt_cache_retention"} { + if _, exists := upstreamBody[field]; exists { + t.Fatalf("upstream body contains %s: %s", field, executor.body) + } + } +} + +func TestCodexAlphaSearchOptInAPIKeyWithoutBaseURLFailsClosed(t *testing.T) { + server := newTestServer(t) + executor := &codexSearchCaptureExecutor{} + server.handlers.AuthManager.RegisterExecutor(executor) + if _, errRegister := server.handlers.AuthManager.Register(context.Background(), &auth.Auth{ + ID: "codex-alpha-api-key", + Provider: "codex", + Status: auth.StatusActive, + Attributes: map[string]string{ + auth.AttributeAPIKey: "codex-alpha-key", + auth.AttributeCodexAlphaSearch: "true", + }, + }); errRegister != nil { + t.Fatalf("register Codex API key: %v", errRegister) + } + + 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 executor.request != nil { + t.Fatal("request was sent without an API key base URL") + } +} + func TestCodexAlphaSearchPassesGinContextToAuthSelection(t *testing.T) { server := newTestServer(t) selector := &codexSearchGinContextSelector{} diff --git a/internal/config/config_normalization.go b/internal/config/config_normalization.go index 697f6c3d..7f9af128 100644 --- a/internal/config/config_normalization.go +++ b/internal/config/config_normalization.go @@ -139,6 +139,9 @@ func (cfg *Config) SanitizeXAIKeys() { return } cfg.XAIKey = sanitizeCodexKeyEntries(cfg.XAIKey) + for i := range cfg.XAIKey { + cfg.XAIKey[i].AlphaSearch = false + } } func sanitizeCodexKeyEntries(entries []CodexKey) []CodexKey { diff --git a/internal/config/config_types.go b/internal/config/config_types.go index a75a4c70..1c9e7f96 100644 --- a/internal/config/config_types.go +++ b/internal/config/config_types.go @@ -417,6 +417,9 @@ type CodexKey struct { // Websockets enables the Responses API websocket transport for this credential. Websockets bool `yaml:"websockets,omitempty" json:"websockets,omitempty"` + // AlphaSearch allows this Codex API key to serve the Alpha Search endpoint. + AlphaSearch bool `yaml:"alpha-search,omitempty" json:"alpha-search,omitempty"` + // ProxyURL overrides the global proxy setting for this API key if provided. ProxyURL string `yaml:"proxy-url" json:"proxy-url"` diff --git a/internal/config/xai_alpha_search_test.go b/internal/config/xai_alpha_search_test.go new file mode 100644 index 00000000..8a0c10d0 --- /dev/null +++ b/internal/config/xai_alpha_search_test.go @@ -0,0 +1,20 @@ +package config + +import "testing" + +func TestSanitizeXAIKeysClearsCodexAlphaSearchCapability(t *testing.T) { + cfg := &Config{XAIKey: []XAIKey{{ + APIKey: "xai-key", + BaseURL: "https://api.x.ai/v1", + AlphaSearch: true, + }}} + + cfg.SanitizeXAIKeys() + + if len(cfg.XAIKey) != 1 { + t.Fatalf("XAI key count = %d, want 1", len(cfg.XAIKey)) + } + if cfg.XAIKey[0].AlphaSearch { + t.Fatal("SanitizeXAIKeys() retained the Codex-only alpha-search capability") + } +} diff --git a/internal/home/client.go b/internal/home/client.go index f4a295c6..d07d4d42 100644 --- a/internal/home/client.go +++ b/internal/home/client.go @@ -1248,7 +1248,7 @@ func queryToLowerMap(query url.Values) map[string]string { return out } -func newAuthDispatchRequest(requestedModel string, sessionID string, headers http.Header, count int) authDispatchRequest { +func newAuthDispatchRequest(requestedModel string, sessionID string, headers http.Header, count int, credentialPolicy string) authDispatchRequest { if count <= 0 { count = 1 } @@ -1259,10 +1259,20 @@ func newAuthDispatchRequest(requestedModel string, sessionID string, headers htt ConcurrencyProtocol: 1, SessionID: strings.TrimSpace(sessionID), Headers: headersToLowerMap(headers), + CredentialPolicy: strings.TrimSpace(credentialPolicy), } } func (c *Client) RPopAuth(ctx context.Context, requestedModel string, sessionID string, headers http.Header, count int) ([]byte, error) { + return c.rPopAuth(ctx, requestedModel, sessionID, headers, count, "") +} + +// RPopAuthWithPolicy requests a Home credential constrained by the supplied fixed policy. +func (c *Client) RPopAuthWithPolicy(ctx context.Context, requestedModel string, sessionID string, headers http.Header, count int, credentialPolicy string) ([]byte, error) { + return c.rPopAuth(ctx, requestedModel, sessionID, headers, count, credentialPolicy) +} + +func (c *Client) rPopAuth(ctx context.Context, requestedModel string, sessionID string, headers http.Header, count int, credentialPolicy string) ([]byte, error) { if c == nil || c.dispatchFenced.Load() { return nil, ErrDispatchFenced } @@ -1276,7 +1286,7 @@ func (c *Client) RPopAuth(ctx context.Context, requestedModel string, sessionID if requestedModel == "" { return nil, fmt.Errorf("home: requested model is empty") } - req := newAuthDispatchRequest(requestedModel, sessionID, headers, count) + req := newAuthDispatchRequest(requestedModel, sessionID, headers, count, credentialPolicy) keyBytes, errMarshal := json.Marshal(&req) if errMarshal != nil { return nil, errMarshal diff --git a/internal/home/client_test.go b/internal/home/client_test.go index 66f87d75..ccf2db3f 100644 --- a/internal/home/client_test.go +++ b/internal/home/client_test.go @@ -26,7 +26,7 @@ import ( ) func TestAuthDispatchRequestIncludesCount(t *testing.T) { - req := newAuthDispatchRequest("gpt-5.4", "session-1", http.Header{"Authorization": {"Bearer test"}}, 2) + req := newAuthDispatchRequest("gpt-5.4", "session-1", http.Header{"Authorization": {"Bearer test"}}, 2, "") raw, err := json.Marshal(&req) if err != nil { @@ -46,11 +46,29 @@ func TestAuthDispatchRequestIncludesCount(t *testing.T) { } func TestAuthDispatchRequestDefaultsCountToOne(t *testing.T) { - req := newAuthDispatchRequest("gpt-5.4", "", nil, 0) + req := newAuthDispatchRequest("gpt-5.4", "", nil, 0, "") if req.Count != 1 { t.Fatalf("count = %d, want 1", req.Count) } + if req.CredentialPolicy != "" { + t.Fatalf("credential policy = %q, want empty", req.CredentialPolicy) + } +} + +func TestAuthDispatchRequestIncludesCredentialPolicy(t *testing.T) { + req := newAuthDispatchRequest("gpt-5.4", "", nil, 1, "codex_alpha_search_v1") + raw, errMarshal := json.Marshal(&req) + if errMarshal != nil { + t.Fatalf("marshal auth dispatch request: %v", errMarshal) + } + var payload map[string]any + if errUnmarshal := json.Unmarshal(raw, &payload); errUnmarshal != nil { + t.Fatalf("unmarshal auth dispatch request: %v", errUnmarshal) + } + if got := payload["credential_policy"]; got != "codex_alpha_search_v1" { + t.Fatalf("credential_policy = %#v, want codex_alpha_search_v1", got) + } } func TestRedisOptionsHomeTLSDisabled(t *testing.T) { diff --git a/internal/home/requests.go b/internal/home/requests.go index eca63742..655fc601 100644 --- a/internal/home/requests.go +++ b/internal/home/requests.go @@ -9,6 +9,7 @@ type authDispatchRequest struct { ConcurrencyProtocol int `json:"concurrency_protocol,omitempty"` SessionID string `json:"session_id,omitempty"` Headers map[string]string `json:"headers,omitempty"` + CredentialPolicy string `json:"credential_policy,omitempty"` } type modelsRequest struct { diff --git a/internal/watcher/diff/config_diff.go b/internal/watcher/diff/config_diff.go index cc3a0507..4c01a076 100644 --- a/internal/watcher/diff/config_diff.go +++ b/internal/watcher/diff/config_diff.go @@ -289,6 +289,9 @@ func BuildConfigChangeDetails(oldCfg, newCfg *config.Config) []string { if o.Websockets != n.Websockets { changes = append(changes, fmt.Sprintf("codex[%d].websockets: %t -> %t", i, o.Websockets, n.Websockets)) } + if o.AlphaSearch != n.AlphaSearch { + changes = append(changes, fmt.Sprintf("codex[%d].alpha-search: %t -> %t", i, o.AlphaSearch, n.AlphaSearch)) + } if strings.TrimSpace(o.APIKey) != strings.TrimSpace(n.APIKey) { changes = append(changes, fmt.Sprintf("codex[%d].api-key: updated", i)) } diff --git a/internal/watcher/diff/config_diff_test.go b/internal/watcher/diff/config_diff_test.go index 2fe86540..f13a3658 100644 --- a/internal/watcher/diff/config_diff_test.go +++ b/internal/watcher/diff/config_diff_test.go @@ -196,6 +196,14 @@ func TestBuildConfigChangeDetails_ModelPrefixes(t *testing.T) { expectContains(t, changes, "vertex[0].prefix: old-v -> new-v") } +func TestBuildConfigChangeDetails_CodexAlphaSearch(t *testing.T) { + oldCfg := &config.Config{CodexKey: []config.CodexKey{{APIKey: "key", BaseURL: "https://codex.example.com"}}} + newCfg := &config.Config{CodexKey: []config.CodexKey{{APIKey: "key", BaseURL: "https://codex.example.com", AlphaSearch: true}}} + + changes := BuildConfigChangeDetails(oldCfg, newCfg) + expectContains(t, changes, "codex[0].alpha-search: false -> true") +} + func TestBuildConfigChangeDetails_XAIKeys(t *testing.T) { oldCfg := &config.Config{XAIKey: []config.XAIKey{{ APIKey: "old-key", diff --git a/internal/watcher/synthesizer/config.go b/internal/watcher/synthesizer/config.go index f15284bd..27f20c8d 100644 --- a/internal/watcher/synthesizer/config.go +++ b/internal/watcher/synthesizer/config.go @@ -232,6 +232,9 @@ func (s *ConfigSynthesizer) synthesizeCodexStyleKeys(ctx *SynthesisContext, entr if entry.Websockets { attrs["websockets"] = "true" } + if provider == "codex" && entry.AlphaSearch { + attrs[coreauth.AttributeCodexAlphaSearch] = "true" + } if hash := diff.ComputeCodexModelsHash(entry.Models); hash != "" { attrs["models_hash"] = hash } diff --git a/internal/watcher/synthesizer/config_test.go b/internal/watcher/synthesizer/config_test.go index ac5d1e1c..6ab64811 100644 --- a/internal/watcher/synthesizer/config_test.go +++ b/internal/watcher/synthesizer/config_test.go @@ -311,6 +311,7 @@ func TestConfigSynthesizer_CodexKeys(t *testing.T) { BaseURL: "https://api.openai.com", ProxyURL: "http://proxy.local", Websockets: true, + AlphaSearch: true, DisableCooling: true, }, }, @@ -339,6 +340,9 @@ func TestConfigSynthesizer_CodexKeys(t *testing.T) { if auths[0].Attributes["websockets"] != "true" { t.Errorf("expected websockets=true, got %s", auths[0].Attributes["websockets"]) } + if auths[0].Attributes[coreauth.AttributeCodexAlphaSearch] != "true" { + t.Errorf("expected codex_alpha_search=true, got %s", auths[0].Attributes[coreauth.AttributeCodexAlphaSearch]) + } if v, ok := auths[0].Metadata["disable_cooling"].(bool); !ok || !v { t.Errorf("expected disable_cooling=true, got %v", auths[0].Metadata["disable_cooling"]) } @@ -354,6 +358,7 @@ func TestConfigSynthesizer_XAIKeys(t *testing.T) { BaseURL: "https://api.x.ai/v1", ProxyURL: "http://proxy.local", Websockets: true, + AlphaSearch: true, DisableCooling: true, Headers: map[string]string{"X-Custom": "value"}, Models: []config.XAIModel{{Name: "grok-4.5", Alias: "grok-latest"}}, @@ -380,6 +385,9 @@ func TestConfigSynthesizer_XAIKeys(t *testing.T) { if auth.Attributes["websockets"] != "true" { t.Fatalf("websockets = %q, want true", auth.Attributes["websockets"]) } + if _, exists := auth.Attributes[coreauth.AttributeCodexAlphaSearch]; exists { + t.Fatal("xAI auth unexpectedly contains codex_alpha_search") + } if auth.Attributes["base_url"] != "https://api.x.ai/v1" { t.Fatalf("base_url = %q, want https://api.x.ai/v1", auth.Attributes["base_url"]) } @@ -421,6 +429,9 @@ func TestConfigSynthesizer_CodexKeys_SkipsEmptyAndHeaders(t *testing.T) { if auths[0].Attributes["header:Authorization"] != "Bearer xyz" { t.Errorf("expected header:Authorization=Bearer xyz, got %s", auths[0].Attributes["header:Authorization"]) } + if _, exists := auths[0].Attributes[coreauth.AttributeCodexAlphaSearch]; exists { + t.Fatal("default alpha-search=false unexpectedly generated codex_alpha_search") + } } func TestConfigSynthesizer_OpenAICompat(t *testing.T) { diff --git a/sdk/cliproxy/auth/classification.go b/sdk/cliproxy/auth/classification.go index f1344fa9..2a9059a8 100644 --- a/sdk/cliproxy/auth/classification.go +++ b/sdk/cliproxy/auth/classification.go @@ -13,14 +13,15 @@ const ( AuthSourceObjectStore = "objectstore" AuthSourcePostgres = "postgres" - AttributeAPIKey = "api_key" - AttributeAuthKind = "auth_kind" - AttributeConfigIndex = "config_index" - AttributePath = "path" - AttributeRuntimeOnly = "runtime_only" - AttributeSource = "source" - AttributeSourceBackend = "source_backend" - AttributeWeight = "weight" + AttributeAPIKey = "api_key" + AttributeAuthKind = "auth_kind" + AttributeCodexAlphaSearch = "codex_alpha_search" + AttributeConfigIndex = "config_index" + AttributePath = "path" + AttributeRuntimeOnly = "runtime_only" + AttributeSource = "source" + AttributeSourceBackend = "source_backend" + AttributeWeight = "weight" ) // AuthKind returns the credential kind using explicit metadata first and legacy diff --git a/sdk/cliproxy/auth/conductor_home.go b/sdk/cliproxy/auth/conductor_home.go index 324d2f7e..fb576f93 100644 --- a/sdk/cliproxy/auth/conductor_home.go +++ b/sdk/cliproxy/auth/conductor_home.go @@ -180,6 +180,10 @@ type homeAuthDispatcher interface { AbortAmbiguousDispatch() } +type homeCredentialPolicyDispatcher interface { + RPopAuthWithPolicy(ctx context.Context, requestedModel string, sessionID string, headers http.Header, count int, credentialPolicy string) ([]byte, error) +} + var currentHomeDispatcher = func() homeAuthDispatcher { return home.Current() } @@ -719,7 +723,17 @@ func (m *Manager) pickHomeDispatchSelection(ctx context.Context, model string, o sessionID := m.homeDispatchSessionID(opts) dispatchHeaders := homeDispatchHeaders(ctx, opts.Headers) - raw, errRPop := client.RPopAuth(ctx, requestedModel, sessionID, dispatchHeaders, homeAuthCountFromMetadata(opts.Metadata)) + credentialPolicy := credentialPolicyFromContext(ctx) + var raw []byte + var errRPop error + if credentialPolicy == "" { + raw, errRPop = client.RPopAuth(ctx, requestedModel, sessionID, dispatchHeaders, homeAuthCountFromMetadata(opts.Metadata)) + } else if policyClient, okPolicy := client.(homeCredentialPolicyDispatcher); okPolicy { + raw, errRPop = policyClient.RPopAuthWithPolicy(ctx, requestedModel, sessionID, dispatchHeaders, homeAuthCountFromMetadata(opts.Metadata), credentialPolicy) + } else { + pending.End() + return nil, &Error{Code: "home_unavailable", Message: "home dispatcher does not support credential policies", HTTPStatus: http.StatusServiceUnavailable} + } if errRPop != nil { if home.IsAmbiguousDispatchError(errRPop) { client.AbortAmbiguousDispatch() diff --git a/sdk/cliproxy/auth/conductor_selection.go b/sdk/cliproxy/auth/conductor_selection.go index 81e41b38..9578e178 100644 --- a/sdk/cliproxy/auth/conductor_selection.go +++ b/sdk/cliproxy/auth/conductor_selection.go @@ -50,9 +50,11 @@ func isBuiltInSelector(selector Selector) bool { } type requiredAuthKindContextKey struct{} +type credentialPolicyContextKey struct{} type authSelectionEligibility struct { requiredKind string + credentialPolicy string disallowFreeAuth bool } @@ -60,10 +62,23 @@ func withRequiredAuthKind(ctx context.Context, requiredKind string) context.Cont return context.WithValue(ctx, requiredAuthKindContextKey{}, requiredKind) } +func withCredentialPolicy(ctx context.Context, policy string) context.Context { + return context.WithValue(ctx, credentialPolicyContextKey{}, policy) +} + +func credentialPolicyFromContext(ctx context.Context) string { + if ctx == nil { + return "" + } + policy, _ := ctx.Value(credentialPolicyContextKey{}).(string) + return policy +} + func authSelectionEligibilityForRequest(ctx context.Context, opts cliproxyexecutor.Options) authSelectionEligibility { eligibility := authSelectionEligibility{disallowFreeAuth: disallowFreeAuthFromMetadata(opts.Metadata)} if ctx != nil { eligibility.requiredKind, _ = ctx.Value(requiredAuthKindContextKey{}).(string) + eligibility.credentialPolicy, _ = ctx.Value(credentialPolicyContextKey{}).(string) } return eligibility } @@ -75,6 +90,9 @@ func (e authSelectionEligibility) allows(auth *Auth) bool { if e.requiredKind != "" && auth.AuthKind() != e.requiredKind { return false } + if e.credentialPolicy != "" && !credentialPolicyAllows(e.credentialPolicy, auth) { + return false + } return !e.disallowFreeAuth || !isFreeCodexAuth(auth) } @@ -1068,6 +1086,81 @@ func (m *Manager) SelectAuthByKind(ctx context.Context, provider, model, require return selected, nil } +// SelectAuthWithCredentialPolicy selects one local credential allowed by a fixed policy. +func (m *Manager) SelectAuthWithCredentialPolicy(ctx context.Context, provider, model, policy string, opts cliproxyexecutor.Options) (*Auth, error) { + if m != nil && m.HomeEnabled() { + return nil, &Error{Code: "home_unavailable", Message: "legacy auth selection is unavailable while Home is enabled", HTTPStatus: http.StatusServiceUnavailable} + } + policy = normalizeCredentialPolicy(policy) + if policy == "" { + return nil, &Error{Code: "invalid_credential_policy", Message: "credential policy is invalid", HTTPStatus: http.StatusBadRequest} + } + if ctx == nil { + ctx = context.Background() + } + selectionCtx := withCredentialPolicy(ctx, policy) + selected, _, errPick := m.pickNextLegacy(selectionCtx, provider, model, opts, nil) + if errPick != nil { + return nil, errPick + } + if selected == nil || !credentialPolicyAllows(policy, selected) { + return nil, &Error{Code: "auth_not_found", Message: "selector returned no eligible auth"} + } + if m.HomeEnabled() { + return nil, &Error{Code: "home_unavailable", Message: "legacy auth selection is unavailable while Home is enabled", HTTPStatus: http.StatusServiceUnavailable} + } + return selected, nil +} + +// SelectHomeAuthWithCredentialPolicy selects a policy-constrained Home dispatch while retaining its execution scope. +func (m *Manager) SelectHomeAuthWithCredentialPolicy(ctx context.Context, provider, model, policy string, opts cliproxyexecutor.Options) (*HomeDispatchSelection, error) { + policy = normalizeCredentialPolicy(policy) + if policy == "" { + return nil, &Error{Code: "invalid_credential_policy", Message: "credential policy is invalid", HTTPStatus: http.StatusBadRequest} + } + if m == nil || !m.HomeEnabled() { + return nil, &Error{Code: "home_unavailable", Message: "home control center unavailable", HTTPStatus: http.StatusServiceUnavailable} + } + if ctx == nil { + ctx = context.Background() + } + selectionCtx := withCredentialPolicy(ctx, policy) + homeAuthCount := homeAuthCountFromMetadata(opts.Metadata) + tried := make(map[string]struct{}) + for { + selectionOpts := withHomeAuthCount(opts, homeAuthCount) + selection, errSelection := m.pickHomeDispatchSelection(selectionCtx, model, selectionOpts) + if errSelection != nil { + return nil, errSelection + } + providerMatches := strings.TrimSpace(provider) == "" || strings.EqualFold(strings.TrimSpace(selection.Provider), strings.TrimSpace(provider)) + policyMatches := credentialPolicyAllows(policy, selection.Auth) + if providerMatches && policyMatches { + return selection, nil + } + + authID := "" + if selection.Auth != nil { + authID = strings.TrimSpace(selection.Auth.ID) + } + reason := "credential_policy_mismatch" + if !providerMatches { + reason = "provider_mismatch" + } + if errEnd := m.endHomeSelectionBeforeRedispatch(selectionCtx, selection, reason); errEnd != nil { + return nil, errEnd + } + 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{}{} + homeAuthCount++ + } +} + // SelectHomeAuthByKind selects a Home dispatch while retaining its execution scope. func (m *Manager) SelectHomeAuthByKind(ctx context.Context, provider string, model string, requiredKind string, opts cliproxyexecutor.Options) (*HomeDispatchSelection, error) { requiredKind = normalizeAuthKind(requiredKind) diff --git a/sdk/cliproxy/auth/credential_policy.go b/sdk/cliproxy/auth/credential_policy.go new file mode 100644 index 00000000..a290759f --- /dev/null +++ b/sdk/cliproxy/auth/credential_policy.go @@ -0,0 +1,39 @@ +package auth + +import "strings" + +const ( + // CredentialPolicyCodexAlphaSearchV1 selects credentials supported by Codex Alpha Search. + CredentialPolicyCodexAlphaSearchV1 = "codex_alpha_search_v1" +) + +func normalizeCredentialPolicy(policy string) string { + switch strings.ToLower(strings.TrimSpace(policy)) { + case CredentialPolicyCodexAlphaSearchV1: + return CredentialPolicyCodexAlphaSearchV1 + default: + return "" + } +} + +func credentialPolicyAllows(policy string, auth *Auth) bool { + if auth == nil { + return false + } + switch policy { + case CredentialPolicyCodexAlphaSearchV1: + if !strings.EqualFold(strings.TrimSpace(auth.Provider), "codex") { + return false + } + switch auth.AuthKind() { + case AuthKindOAuth: + return true + case AuthKindAPIKey: + return strings.EqualFold(authAttribute(auth, AttributeCodexAlphaSearch), "true") + default: + return false + } + default: + return false + } +} diff --git a/sdk/cliproxy/auth/scheduler_test.go b/sdk/cliproxy/auth/scheduler_test.go index 3e3156d1..b0c2aba5 100644 --- a/sdk/cliproxy/auth/scheduler_test.go +++ b/sdk/cliproxy/auth/scheduler_test.go @@ -86,8 +86,9 @@ type inactivePluginScheduler struct { } type authKindHomeDispatcher struct { - auths []Auth - counts []int + auths []Auth + counts []int + policies []string } func (d *authKindHomeDispatcher) HeartbeatOK() bool { @@ -102,6 +103,11 @@ func (d *authKindHomeDispatcher) RPopAuth(_ context.Context, _ string, _ string, return json.Marshal(homeAuthDispatchResponse{Auth: d.auths[count-1]}) } +func (d *authKindHomeDispatcher) RPopAuthWithPolicy(ctx context.Context, model string, sessionID string, headers http.Header, count int, policy string) ([]byte, error) { + d.policies = append(d.policies, policy) + return d.RPopAuth(ctx, model, sessionID, headers, count) +} + func (*authKindHomeDispatcher) AbortAmbiguousDispatch() {} func (s *inactivePluginScheduler) HasScheduler() bool { @@ -721,6 +727,57 @@ func TestManagerSelectAuthByKindSkipsAPIKey(t *testing.T) { } } +func TestManagerCodexAlphaSearchPolicyFiltersBeforePluginScheduler(t *testing.T) { + manager := NewManager(nil, &RoundRobinSelector{}, nil) + manager.executors["codex"] = schedulerTestExecutor{} + for _, candidate := range []*Auth{ + {ID: "ordinary-api-key", Provider: "codex", Attributes: map[string]string{AttributeAPIKey: "ordinary"}}, + {ID: "alpha-api-key", Provider: "codex", Attributes: map[string]string{AttributeAPIKey: "alpha", AttributeCodexAlphaSearch: "true", "base_url": "https://codex.example.com"}}, + } { + 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: "alpha-api-key"}, + handled: true, + } + manager.SetPluginScheduler(scheduler) + + selected, errSelect := manager.SelectAuthWithCredentialPolicy(context.Background(), "codex", "", CredentialPolicyCodexAlphaSearchV1, cliproxyexecutor.Options{}) + if errSelect != nil { + t.Fatalf("SelectAuthWithCredentialPolicy() error = %v", errSelect) + } + if selected == nil || selected.ID != "alpha-api-key" { + t.Fatalf("SelectAuthWithCredentialPolicy() auth = %#v, want alpha-api-key", selected) + } + if len(scheduler.requests) != 1 || len(scheduler.requests[0].Candidates) != 1 || scheduler.requests[0].Candidates[0].ID != "alpha-api-key" { + t.Fatalf("scheduler candidates = %#v, want only alpha-api-key", scheduler.requests) + } +} + +func TestManagerCodexAlphaSearchPolicyRejectsOrdinaryAPIKey(t *testing.T) { + manager := NewManager(nil, &RoundRobinSelector{}, nil) + manager.executors["codex"] = schedulerTestExecutor{} + if _, errRegister := manager.Register(context.Background(), &Auth{ + ID: "ordinary-api-key", + Provider: "codex", + Attributes: map[string]string{AttributeAPIKey: "ordinary"}, + }); errRegister != nil { + t.Fatalf("Register() error = %v", errRegister) + } + + selected, errSelect := manager.SelectAuthWithCredentialPolicy(context.Background(), "codex", "", CredentialPolicyCodexAlphaSearchV1, cliproxyexecutor.Options{}) + if selected != nil { + t.Fatalf("SelectAuthWithCredentialPolicy() auth = %#v, want nil", selected) + } + var authErr *Error + if !errors.As(errSelect, &authErr) || authErr.Code != "auth_not_found" { + t.Fatalf("SelectAuthWithCredentialPolicy() error = %#v, want auth_not_found", errSelect) + } +} + func TestManagerSelectAuthByKindWeightedRoundRobinIgnoresIneligibleAPIKeyWeight(t *testing.T) { manager := NewManager(nil, &WeightedRoundRobinSelector{}, nil) manager.executors["codex"] = schedulerTestExecutor{} @@ -959,6 +1016,36 @@ func TestSelectHomeAuthByKindSkipsProviderMismatch(t *testing.T) { selection.End("test_complete") } +func TestSelectHomeAuthWithCredentialPolicyTransportsAndValidatesPolicy(t *testing.T) { + dispatcher := &authKindHomeDispatcher{auths: []Auth{ + {ID: "ordinary-api-key", Provider: "codex", Attributes: map[string]string{AttributeAPIKey: "ordinary", "base_url": "https://ordinary.example.com"}}, + {ID: "alpha-api-key", Provider: "codex", Attributes: map[string]string{AttributeAPIKey: "alpha", AttributeCodexAlphaSearch: "true", "base_url": "https://alpha.example.com"}}, + }} + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + registry := executionregistry.New() + manager.PublishHomeDispatch(dispatcher, registry, 1) + manager.RegisterExecutor(schedulerTestExecutor{provider: "codex"}) + + selection, errSelect := manager.SelectHomeAuthWithCredentialPolicy(context.Background(), "codex", "gpt-5.4", CredentialPolicyCodexAlphaSearchV1, cliproxyexecutor.Options{}) + if errSelect != nil { + t.Fatalf("SelectHomeAuthWithCredentialPolicy() error = %v", errSelect) + } + if selection == nil || selection.Auth == nil || selection.Auth.ID != "alpha-api-key" { + t.Fatalf("SelectHomeAuthWithCredentialPolicy() = %#v, want alpha-api-key", selection) + } + if got := dispatcher.counts; len(got) != 2 || got[0] != 1 || got[1] != 2 { + t.Fatalf("Home auth counts = %v, want [1 2]", got) + } + if got := dispatcher.policies; len(got) != 2 || got[0] != CredentialPolicyCodexAlphaSearchV1 || got[1] != CredentialPolicyCodexAlphaSearchV1 { + t.Fatalf("Home credential policies = %v", got) + } + selection.End("test_complete") + if errDrain := registry.Drain(context.Background()); errDrain != nil { + t.Fatalf("Drain() error = %v", errDrain) + } +} + func TestSelectHomeAuthByKindKeepsLogicalProviderWhenUsingCompatibilityExecutor(t *testing.T) { dispatcher := &authKindHomeDispatcher{auths: []Auth{{ ID: "compat-auth",