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:
@@ -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()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user