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

- 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
+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()