airkjw/CLIProxyAPI
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(auth): disallow free-tier Codex auth during selection process

Browse Source 
- Introduced `disallowFreeAuthFromMetadata` and `isFreeCodexAuth` to enforce skipping free-tier credentials.
- Modified scheduler logic to honor `DisallowFreeAuthMetadataKey` during auth selection.
- Updated `ensureImageGenerationTool` to skip tool injection for free-tier Codex auth.
- Added context utility `WithDisallowFreeAuth` and integrated with image handlers.
- Augmented relevant tests to cover free-tier exclusion scenarios.
This commit is contained in:
Luis Pater
2026-04-24 23:18:56 +08:00
parent 36cc762fc9
commit a7e92e2639
7 changed files with 200 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sdk/cliproxy/auth/conductor.go
+102 -42
View File
@@ -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 {
sdk/cliproxy/auth/scheduler_test.go
+33
View File
@@ -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()