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