diff --git a/internal/runtime/executor/codex_executor.go b/internal/runtime/executor/codex_executor.go index 38667231a..dc3254a76 100644 --- a/internal/runtime/executor/codex_executor.go +++ b/internal/runtime/executor/codex_executor.go @@ -180,7 +180,7 @@ func (e *CodexExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, re body, _ = sjson.DeleteBytes(body, "safety_identifier") body, _ = sjson.DeleteBytes(body, "stream_options") body = normalizeCodexInstructions(body) - body = ensureImageGenerationTool(body, baseModel) + body = ensureImageGenerationTool(body, baseModel, auth) url := strings.TrimSuffix(baseURL, "/") + "/responses" httpReq, err := e.cacheHelper(ctx, from, url, req, body) @@ -327,7 +327,7 @@ func (e *CodexExecutor) executeCompact(ctx context.Context, auth *cliproxyauth.A body, _ = sjson.SetBytes(body, "model", baseModel) body, _ = sjson.DeleteBytes(body, "stream") body = normalizeCodexInstructions(body) - body = ensureImageGenerationTool(body, baseModel) + body = ensureImageGenerationTool(body, baseModel, auth) url := strings.TrimSuffix(baseURL, "/") + "/responses/compact" httpReq, err := e.cacheHelper(ctx, from, url, req, body) @@ -422,7 +422,7 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au body, _ = sjson.DeleteBytes(body, "stream_options") body, _ = sjson.SetBytes(body, "model", baseModel) body = normalizeCodexInstructions(body) - body = ensureImageGenerationTool(body, baseModel) + body = ensureImageGenerationTool(body, baseModel, auth) url := strings.TrimSuffix(baseURL, "/") + "/responses" httpReq, err := e.cacheHelper(ctx, from, url, req, body) @@ -827,10 +827,23 @@ func normalizeCodexInstructions(body []byte) []byte { var imageGenToolJSON = []byte(`{"type":"image_generation","output_format":"png"}`) var imageGenToolArrayJSON = []byte(`[{"type":"image_generation","output_format":"png"}]`) -func ensureImageGenerationTool(body []byte, baseModel string) []byte { +func isCodexFreePlanAuth(auth *cliproxyauth.Auth) bool { + if auth == nil || auth.Attributes == nil { + return false + } + if !strings.EqualFold(strings.TrimSpace(auth.Provider), "codex") { + return false + } + return strings.EqualFold(strings.TrimSpace(auth.Attributes["plan_type"]), "free") +} + +func ensureImageGenerationTool(body []byte, baseModel string, auth *cliproxyauth.Auth) []byte { if strings.HasSuffix(baseModel, "spark") { return body } + if isCodexFreePlanAuth(auth) { + return body + } tools := gjson.GetBytes(body, "tools") if !tools.Exists() || !tools.IsArray() { diff --git a/internal/runtime/executor/codex_executor_imagegen_test.go b/internal/runtime/executor/codex_executor_imagegen_test.go index 5e67c598a..1657209a9 100644 --- a/internal/runtime/executor/codex_executor_imagegen_test.go +++ b/internal/runtime/executor/codex_executor_imagegen_test.go @@ -3,12 +3,13 @@ package executor import ( "testing" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth" "github.com/tidwall/gjson" ) func TestEnsureImageGenerationTool_NoTools(t *testing.T) { body := []byte(`{"model":"gpt-5.4","input":"draw a cat"}`) - result := ensureImageGenerationTool(body, "gpt-5.4") + result := ensureImageGenerationTool(body, "gpt-5.4", nil) tools := gjson.GetBytes(result, "tools") if !tools.IsArray() { @@ -28,7 +29,7 @@ func TestEnsureImageGenerationTool_NoTools(t *testing.T) { func TestEnsureImageGenerationTool_ExistingToolsWithoutImageGen(t *testing.T) { body := []byte(`{"model":"gpt-5.4","tools":[{"type":"function","name":"get_weather","parameters":{}}]}`) - result := ensureImageGenerationTool(body, "gpt-5.4") + result := ensureImageGenerationTool(body, "gpt-5.4", nil) tools := gjson.GetBytes(result, "tools") arr := tools.Array() @@ -45,7 +46,7 @@ func TestEnsureImageGenerationTool_ExistingToolsWithoutImageGen(t *testing.T) { func TestEnsureImageGenerationTool_AlreadyPresent(t *testing.T) { body := []byte(`{"model":"gpt-5.4","tools":[{"type":"image_generation","output_format":"webp"},{"type":"function","name":"f1"}]}`) - result := ensureImageGenerationTool(body, "gpt-5.4") + result := ensureImageGenerationTool(body, "gpt-5.4", nil) tools := gjson.GetBytes(result, "tools") arr := tools.Array() @@ -59,7 +60,7 @@ func TestEnsureImageGenerationTool_AlreadyPresent(t *testing.T) { func TestEnsureImageGenerationTool_EmptyToolsArray(t *testing.T) { body := []byte(`{"model":"gpt-5.4","tools":[]}`) - result := ensureImageGenerationTool(body, "gpt-5.4") + result := ensureImageGenerationTool(body, "gpt-5.4", nil) tools := gjson.GetBytes(result, "tools") arr := tools.Array() @@ -73,7 +74,7 @@ func TestEnsureImageGenerationTool_EmptyToolsArray(t *testing.T) { func TestEnsureImageGenerationTool_WebSearchAndImageGen(t *testing.T) { body := []byte(`{"model":"gpt-5.4","tools":[{"type":"web_search"}]}`) - result := ensureImageGenerationTool(body, "gpt-5.4") + result := ensureImageGenerationTool(body, "gpt-5.4", nil) tools := gjson.GetBytes(result, "tools") arr := tools.Array() @@ -90,7 +91,7 @@ func TestEnsureImageGenerationTool_WebSearchAndImageGen(t *testing.T) { func TestEnsureImageGenerationTool_GPT53CodexSparkDoesNotInjectTool(t *testing.T) { body := []byte(`{"model":"gpt-5.3-codex-spark","input":"draw a cat"}`) - result := ensureImageGenerationTool(body, "gpt-5.3-codex-spark") + result := ensureImageGenerationTool(body, "gpt-5.3-codex-spark", nil) if string(result) != string(body) { t.Fatalf("expected body to be unchanged, got %s", string(result)) @@ -99,3 +100,19 @@ func TestEnsureImageGenerationTool_GPT53CodexSparkDoesNotInjectTool(t *testing.T t.Fatalf("expected no tools for gpt-5.3-codex-spark, got %s", gjson.GetBytes(result, "tools").Raw) } } + +func TestEnsureImageGenerationTool_FreeCodexAuthDoesNotInjectTool(t *testing.T) { + body := []byte(`{"model":"gpt-5.4","input":"draw a cat"}`) + freeAuth := &cliproxyauth.Auth{ + Provider: "codex", + Attributes: map[string]string{"plan_type": "free"}, + } + result := ensureImageGenerationTool(body, "gpt-5.4", freeAuth) + + if string(result) != string(body) { + t.Fatalf("expected body to be unchanged, got %s", string(result)) + } + if gjson.GetBytes(result, "tools").Exists() { + t.Fatalf("expected no tools for free codex auth, got %s", gjson.GetBytes(result, "tools").Raw) + } +} diff --git a/sdk/api/handlers/handlers.go b/sdk/api/handlers/handlers.go index 1fda8f49f..5f0ea7b81 100644 --- a/sdk/api/handlers/handlers.go +++ b/sdk/api/handlers/handlers.go @@ -55,6 +55,7 @@ const ( type pinnedAuthContextKey struct{} type selectedAuthCallbackContextKey struct{} type executionSessionContextKey struct{} +type disallowFreeAuthContextKey struct{} // WithPinnedAuthID returns a child context that requests execution on a specific auth ID. func WithPinnedAuthID(ctx context.Context, authID string) context.Context { @@ -91,6 +92,14 @@ func WithExecutionSessionID(ctx context.Context, sessionID string) context.Conte return context.WithValue(ctx, executionSessionContextKey{}, sessionID) } +// WithDisallowFreeAuth returns a child context that requests skipping known free-tier credentials. +func WithDisallowFreeAuth(ctx context.Context) context.Context { + if ctx == nil { + ctx = context.Background() + } + return context.WithValue(ctx, disallowFreeAuthContextKey{}, true) +} + // BuildErrorResponseBody builds an OpenAI-compatible JSON error response body. // If errText is already valid JSON, it is returned as-is to preserve upstream error payloads. func BuildErrorResponseBody(status int, errText string) []byte { @@ -208,6 +217,9 @@ func requestExecutionMetadata(ctx context.Context) map[string]any { if executionSessionID := executionSessionIDFromContext(ctx); executionSessionID != "" { meta[coreexecutor.ExecutionSessionMetadataKey] = executionSessionID } + if disallowFreeAuthFromContext(ctx) { + meta[coreexecutor.DisallowFreeAuthMetadataKey] = true + } return meta } @@ -252,6 +264,14 @@ func executionSessionIDFromContext(ctx context.Context) string { } } +func disallowFreeAuthFromContext(ctx context.Context) bool { + if ctx == nil { + return false + } + raw, ok := ctx.Value(disallowFreeAuthContextKey{}).(bool) + return ok && raw +} + // BaseAPIHandler contains the handlers for API endpoints. // It holds a pool of clients to interact with the backend service and manages // load balancing, client selection, and configuration. diff --git a/sdk/api/handlers/openai/openai_images_handlers.go b/sdk/api/handlers/openai/openai_images_handlers.go index 93d45460d..17243314f 100644 --- a/sdk/api/handlers/openai/openai_images_handlers.go +++ b/sdk/api/handlers/openai/openai_images_handlers.go @@ -527,6 +527,7 @@ func (h *OpenAIAPIHandler) collectImagesFromResponses(c *gin.Context, responsesR c.Header("Content-Type", "application/json") cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background()) + cliCtx = handlers.WithDisallowFreeAuth(cliCtx) stopKeepAlive := h.StartNonStreamingKeepAlive(c, cliCtx) dataChan, upstreamHeaders, errChan := h.ExecuteStreamWithAuthManager(cliCtx, "openai-response", defaultImagesMainModel, responsesReq, "") @@ -716,6 +717,7 @@ func (h *OpenAIAPIHandler) streamImagesFromResponses(c *gin.Context, responsesRe } cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background()) + cliCtx = handlers.WithDisallowFreeAuth(cliCtx) dataChan, upstreamHeaders, errChan := h.ExecuteStreamWithAuthManager(cliCtx, "openai-response", defaultImagesMainModel, responsesReq, "") setSSEHeaders := func() { diff --git a/sdk/cliproxy/auth/conductor.go b/sdk/cliproxy/auth/conductor.go index 05a32ceb2..2091f669a 100644 --- a/sdk/cliproxy/auth/conductor.go +++ b/sdk/cliproxy/auth/conductor.go @@ -1549,6 +1549,38 @@ func pinnedAuthIDFromMetadata(meta map[string]any) string { } } +func disallowFreeAuthFromMetadata(meta map[string]any) bool { + if len(meta) == 0 { + return false + } + raw, ok := meta[cliproxyexecutor.DisallowFreeAuthMetadataKey] + if !ok || raw == nil { + return false + } + switch val := raw.(type) { + case bool: + return val + case string: + parsed, err := strconv.ParseBool(strings.TrimSpace(val)) + return err == nil && parsed + case []byte: + parsed, err := strconv.ParseBool(strings.TrimSpace(string(val))) + return err == nil && parsed + default: + return false + } +} + +func isFreeCodexAuth(auth *Auth) bool { + if auth == nil || auth.Attributes == nil { + return false + } + if !strings.EqualFold(strings.TrimSpace(auth.Provider), "codex") { + return false + } + return strings.EqualFold(strings.TrimSpace(auth.Attributes["plan_type"]), "free") +} + func publishSelectedAuthMetadata(meta map[string]any, authID string) { if len(meta) == 0 { return @@ -2633,6 +2665,7 @@ func (m *Manager) routeAwareSelectionRequired(auth *Auth, routeModel string) boo func (m *Manager) pickNextLegacy(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, tried map[string]struct{}) (*Auth, ProviderExecutor, error) { pinnedAuthID := pinnedAuthIDFromMetadata(opts.Metadata) + disallowFreeAuth := disallowFreeAuthFromMetadata(opts.Metadata) m.mu.RLock() executor, okExecutor := m.executors[provider] @@ -2657,6 +2690,9 @@ func (m *Manager) pickNextLegacy(ctx context.Context, provider, model string, op if pinnedAuthID != "" && candidate.ID != pinnedAuthID { continue } + if disallowFreeAuth && isFreeCodexAuth(candidate) { + continue + } if _, used := tried[candidate.ID]; used { continue } @@ -2720,31 +2756,42 @@ func (m *Manager) pickNext(ctx context.Context, provider, model string, opts cli if !okExecutor { return nil, nil, &Error{Code: "executor_not_found", Message: "executor not registered"} } - selected, errPick := m.scheduler.pickSingle(ctx, provider, model, opts, tried) - if errPick != nil && model != "" && shouldRetrySchedulerPick(errPick) { - m.syncScheduler() - selected, errPick = m.scheduler.pickSingle(ctx, provider, model, opts, tried) - } - if errPick != nil { - return nil, nil, errPick - } - if selected == nil { - return nil, nil, &Error{Code: "auth_not_found", Message: "selector returned no auth"} - } - authCopy := selected.Clone() - if !selected.indexAssigned { - m.mu.Lock() - if current := m.auths[authCopy.ID]; current != nil && !current.indexAssigned { - current.EnsureIndex() - authCopy = current.Clone() + disallowFreeAuth := disallowFreeAuthFromMetadata(opts.Metadata) + for { + selected, errPick := m.scheduler.pickSingle(ctx, provider, model, opts, tried) + if errPick != nil && model != "" && shouldRetrySchedulerPick(errPick) { + m.syncScheduler() + selected, errPick = m.scheduler.pickSingle(ctx, provider, model, opts, tried) } - m.mu.Unlock() + if errPick != nil { + return nil, nil, errPick + } + if selected == nil { + return nil, nil, &Error{Code: "auth_not_found", Message: "selector returned no auth"} + } + if disallowFreeAuth && isFreeCodexAuth(selected) { + if tried == nil { + tried = make(map[string]struct{}) + } + tried[selected.ID] = struct{}{} + continue + } + authCopy := selected.Clone() + if !selected.indexAssigned { + m.mu.Lock() + if current := m.auths[authCopy.ID]; current != nil && !current.indexAssigned { + current.EnsureIndex() + authCopy = current.Clone() + } + m.mu.Unlock() + } + return authCopy, executor, nil } - return authCopy, executor, nil } func (m *Manager) pickNextMixedLegacy(ctx context.Context, providers []string, model string, opts cliproxyexecutor.Options, tried map[string]struct{}) (*Auth, ProviderExecutor, string, error) { pinnedAuthID := pinnedAuthIDFromMetadata(opts.Metadata) + disallowFreeAuth := disallowFreeAuthFromMetadata(opts.Metadata) providerSet := make(map[string]struct{}, len(providers)) for _, provider := range providers { @@ -2776,6 +2823,9 @@ func (m *Manager) pickNextMixedLegacy(ctx context.Context, providers []string, m if pinnedAuthID != "" && candidate.ID != pinnedAuthID { continue } + if disallowFreeAuth && isFreeCodexAuth(candidate) { + continue + } providerKey := strings.TrimSpace(strings.ToLower(candidate.Provider)) if providerKey == "" { continue @@ -2879,31 +2929,41 @@ func (m *Manager) pickNextMixed(ctx context.Context, providers []string, model s m.mu.RUnlock() } - selected, providerKey, errPick := m.scheduler.pickMixed(ctx, eligibleProviders, model, opts, tried) - if errPick != nil && model != "" && shouldRetrySchedulerPick(errPick) { - m.syncScheduler() - selected, providerKey, errPick = m.scheduler.pickMixed(ctx, eligibleProviders, model, opts, tried) - } - if errPick != nil { - return nil, nil, "", errPick - } - if selected == nil { - return nil, nil, "", &Error{Code: "auth_not_found", Message: "selector returned no auth"} - } - executor, okExecutor := m.Executor(providerKey) - if !okExecutor { - return nil, nil, "", &Error{Code: "executor_not_found", Message: "executor not registered"} - } - authCopy := selected.Clone() - if !selected.indexAssigned { - m.mu.Lock() - if current := m.auths[authCopy.ID]; current != nil && !current.indexAssigned { - current.EnsureIndex() - authCopy = current.Clone() + disallowFreeAuth := disallowFreeAuthFromMetadata(opts.Metadata) + for { + selected, providerKey, errPick := m.scheduler.pickMixed(ctx, eligibleProviders, model, opts, tried) + if errPick != nil && model != "" && shouldRetrySchedulerPick(errPick) { + m.syncScheduler() + selected, providerKey, errPick = m.scheduler.pickMixed(ctx, eligibleProviders, model, opts, tried) } - m.mu.Unlock() + if errPick != nil { + return nil, nil, "", errPick + } + if selected == nil { + return nil, nil, "", &Error{Code: "auth_not_found", Message: "selector returned no auth"} + } + if disallowFreeAuth && isFreeCodexAuth(selected) { + if tried == nil { + tried = make(map[string]struct{}) + } + tried[selected.ID] = struct{}{} + continue + } + executor, okExecutor := m.Executor(providerKey) + if !okExecutor { + return nil, nil, "", &Error{Code: "executor_not_found", Message: "executor not registered"} + } + authCopy := selected.Clone() + if !selected.indexAssigned { + m.mu.Lock() + if current := m.auths[authCopy.ID]; current != nil && !current.indexAssigned { + current.EnsureIndex() + authCopy = current.Clone() + } + m.mu.Unlock() + } + return authCopy, executor, providerKey, nil } - return authCopy, executor, providerKey, nil } func (m *Manager) findAllAntigravityCreditsCandidateAuths(routeModel string, opts cliproxyexecutor.Options) []creditsCandidateEntry { diff --git a/sdk/cliproxy/auth/scheduler_test.go b/sdk/cliproxy/auth/scheduler_test.go index d744ec32d..8caaa4735 100644 --- a/sdk/cliproxy/auth/scheduler_test.go +++ b/sdk/cliproxy/auth/scheduler_test.go @@ -333,6 +333,39 @@ func TestManager_PickNextMixed_UsesWeightedProviderRotationBeforeCredentialRotat } } +func TestManager_PickNextMixed_DisallowFreeAuthSkipsCodexFreePlan(t *testing.T) { + t.Parallel() + + model := "gpt-5.4-mini" + registerSchedulerModels(t, "codex", model, "codex-a-free", "codex-b-plus") + + manager := NewManager(nil, &RoundRobinSelector{}, nil) + manager.executors["codex"] = schedulerTestExecutor{} + if _, errRegister := manager.Register(context.Background(), &Auth{ID: "codex-a-free", Provider: "codex", Attributes: map[string]string{"plan_type": "free"}}); errRegister != nil { + t.Fatalf("Register(codex-a-free) error = %v", errRegister) + } + if _, errRegister := manager.Register(context.Background(), &Auth{ID: "codex-b-plus", Provider: "codex", Attributes: map[string]string{"plan_type": "plus"}}); errRegister != nil { + t.Fatalf("Register(codex-b-plus) error = %v", errRegister) + } + + opts := cliproxyexecutor.Options{ + Metadata: map[string]any{cliproxyexecutor.DisallowFreeAuthMetadataKey: true}, + } + got, _, provider, errPick := manager.pickNextMixed(context.Background(), []string{"codex"}, model, opts, map[string]struct{}{}) + if errPick != nil { + t.Fatalf("pickNextMixed() error = %v", errPick) + } + if got == nil { + t.Fatalf("pickNextMixed() auth = nil") + } + if provider != "codex" { + t.Fatalf("pickNextMixed() provider = %q, want %q", provider, "codex") + } + if got.ID != "codex-b-plus" { + t.Fatalf("pickNextMixed() auth.ID = %q, want %q", got.ID, "codex-b-plus") + } +} + func TestManagerCustomSelector_FallsBackToLegacyPath(t *testing.T) { t.Parallel() diff --git a/sdk/cliproxy/executor/types.go b/sdk/cliproxy/executor/types.go index 4ea810394..ac58286fd 100644 --- a/sdk/cliproxy/executor/types.go +++ b/sdk/cliproxy/executor/types.go @@ -10,6 +10,9 @@ import ( // RequestedModelMetadataKey stores the client-requested model name in Options.Metadata. const RequestedModelMetadataKey = "requested_model" +// DisallowFreeAuthMetadataKey instructs auth selection to skip known free-tier credentials. +const DisallowFreeAuthMetadataKey = "disallow_free_auth" + const ( // PinnedAuthMetadataKey locks execution to a specific auth ID. PinnedAuthMetadataKey = "pinned_auth_id"