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
+20
View File
@@ -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.