diff --git a/internal/runtime/executor/aistudio_executor.go b/internal/runtime/executor/aistudio_executor.go index d37cd2c2..5517d970 100644 --- a/internal/runtime/executor/aistudio_executor.go +++ b/internal/runtime/executor/aistudio_executor.go @@ -1,3 +1,6 @@ +// Package executor provides runtime execution capabilities for various AI service providers. +// This file implements the AI Studio executor that routes requests through a websocket-backed +// transport for the AI Studio provider. package executor import ( @@ -26,15 +29,23 @@ type AIStudioExecutor struct { cfg *config.Config } -// NewAIStudioExecutor constructs a websocket executor for the provider name. +// NewAIStudioExecutor creates a new AI Studio executor instance. +// +// Parameters: +// - cfg: The application configuration +// - provider: The provider name +// - relay: The websocket relay manager +// +// Returns: +// - *AIStudioExecutor: A new AI Studio executor instance func NewAIStudioExecutor(cfg *config.Config, provider string, relay *wsrelay.Manager) *AIStudioExecutor { return &AIStudioExecutor{provider: strings.ToLower(provider), relay: relay, cfg: cfg} } -// Identifier returns the logical provider key for routing. +// Identifier returns the executor identifier. func (e *AIStudioExecutor) Identifier() string { return "aistudio" } -// PrepareRequest is a no-op because websocket transport already injects headers. +// PrepareRequest prepares the HTTP request for execution (no-op for AI Studio). func (e *AIStudioExecutor) PrepareRequest(_ *http.Request, _ *cliproxyauth.Auth) error { return nil } @@ -293,8 +304,8 @@ func (e *AIStudioExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.A return cliproxyexecutor.Response{Payload: []byte(translated)}, nil } -func (e *AIStudioExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) { - _ = ctx +// Refresh refreshes the authentication credentials (no-op for AI Studio). +func (e *AIStudioExecutor) Refresh(_ context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) { return auth, nil } diff --git a/internal/runtime/executor/antigravity_executor.go b/internal/runtime/executor/antigravity_executor.go index f5b5ef06..53f024bf 100644 --- a/internal/runtime/executor/antigravity_executor.go +++ b/internal/runtime/executor/antigravity_executor.go @@ -1,3 +1,6 @@ +// Package executor provides runtime execution capabilities for various AI service providers. +// This file implements the Antigravity executor that proxies requests to the antigravity +// upstream using OAuth credentials. package executor import ( @@ -29,16 +32,15 @@ import ( const ( antigravityBaseURLDaily = "https://daily-cloudcode-pa.sandbox.googleapis.com" // antigravityBaseURLAutopush = "https://autopush-cloudcode-pa.sandbox.googleapis.com" - antigravityBaseURLProd = "https://cloudcode-pa.googleapis.com" - antigravityStreamPath = "/v1internal:streamGenerateContent" - antigravityGeneratePath = "/v1internal:generateContent" - antigravityModelsPath = "/v1internal:fetchAvailableModels" - antigravityClientID = "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com" - antigravityClientSecret = "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" - defaultAntigravityAgent = "antigravity/1.11.5 windows/amd64" - antigravityAuthType = "antigravity" - refreshSkew = 3000 * time.Second - streamScannerBuffer int = 52_428_800 // 50MB + antigravityBaseURLProd = "https://cloudcode-pa.googleapis.com" + antigravityStreamPath = "/v1internal:streamGenerateContent" + antigravityGeneratePath = "/v1internal:generateContent" + antigravityModelsPath = "/v1internal:fetchAvailableModels" + antigravityClientID = "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com" + antigravityClientSecret = "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" + defaultAntigravityAgent = "antigravity/1.11.5 windows/amd64" + antigravityAuthType = "antigravity" + refreshSkew = 3000 * time.Second ) var randSource = rand.New(rand.NewSource(time.Now().UnixNano())) @@ -48,15 +50,21 @@ type AntigravityExecutor struct { cfg *config.Config } -// NewAntigravityExecutor constructs a new executor instance. +// NewAntigravityExecutor creates a new Antigravity executor instance. +// +// Parameters: +// - cfg: The application configuration +// +// Returns: +// - *AntigravityExecutor: A new Antigravity executor instance func NewAntigravityExecutor(cfg *config.Config) *AntigravityExecutor { return &AntigravityExecutor{cfg: cfg} } -// Identifier implements ProviderExecutor. +// Identifier returns the executor identifier. func (e *AntigravityExecutor) Identifier() string { return antigravityAuthType } -// PrepareRequest implements ProviderExecutor. +// PrepareRequest prepares the HTTP request for execution (no-op for Antigravity). func (e *AntigravityExecutor) PrepareRequest(_ *http.Request, _ *cliproxyauth.Auth) error { return nil } // Execute handles non-streaming requests via the antigravity generate endpoint. @@ -292,7 +300,7 @@ func (e *AntigravityExecutor) ExecuteStream(ctx context.Context, auth *cliproxya return nil, err } -// Refresh refreshes the OAuth token using the refresh token. +// Refresh refreshes the authentication credentials using the refresh token. func (e *AntigravityExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) { if auth == nil { return auth, nil diff --git a/internal/runtime/executor/gemini_cli_executor.go b/internal/runtime/executor/gemini_cli_executor.go index 2c4f3f88..35a5f03d 100644 --- a/internal/runtime/executor/gemini_cli_executor.go +++ b/internal/runtime/executor/gemini_cli_executor.go @@ -1,3 +1,6 @@ +// Package executor provides runtime execution capabilities for various AI service providers. +// This file implements the Gemini CLI executor that talks to Cloud Code Assist endpoints +// using OAuth credentials from auth metadata. package executor import ( @@ -44,12 +47,21 @@ type GeminiCLIExecutor struct { cfg *config.Config } +// NewGeminiCLIExecutor creates a new Gemini CLI executor instance. +// +// Parameters: +// - cfg: The application configuration +// +// Returns: +// - *GeminiCLIExecutor: A new Gemini CLI executor instance func NewGeminiCLIExecutor(cfg *config.Config) *GeminiCLIExecutor { return &GeminiCLIExecutor{cfg: cfg} } +// Identifier returns the executor identifier. func (e *GeminiCLIExecutor) Identifier() string { return "gemini-cli" } +// PrepareRequest prepares the HTTP request for execution (no-op for Gemini CLI). func (e *GeminiCLIExecutor) PrepareRequest(_ *http.Request, _ *cliproxyauth.Auth) error { return nil } func (e *GeminiCLIExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) { @@ -309,7 +321,7 @@ func (e *GeminiCLIExecutor) ExecuteStream(ctx context.Context, auth *cliproxyaut }() if opts.Alt == "" { scanner := bufio.NewScanner(resp.Body) - scanner.Buffer(nil, 52_428_800) // 50MB + scanner.Buffer(nil, streamScannerBuffer) var param any for scanner.Scan() { line := scanner.Bytes() @@ -471,9 +483,8 @@ func (e *GeminiCLIExecutor) CountTokens(ctx context.Context, auth *cliproxyauth. return cliproxyexecutor.Response{}, newGeminiStatusErr(lastStatus, lastBody) } -func (e *GeminiCLIExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) { - log.Debugf("gemini cli executor: refresh called") - _ = ctx +// Refresh refreshes the authentication credentials (no-op for Gemini CLI). +func (e *GeminiCLIExecutor) Refresh(_ context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) { return auth, nil } diff --git a/internal/runtime/executor/gemini_executor.go b/internal/runtime/executor/gemini_executor.go index 7b94b145..a6d76c58 100644 --- a/internal/runtime/executor/gemini_executor.go +++ b/internal/runtime/executor/gemini_executor.go @@ -11,7 +11,6 @@ import ( "io" "net/http" "strings" - "time" "github.com/router-for-me/CLIProxyAPI/v6/internal/config" "github.com/router-for-me/CLIProxyAPI/v6/internal/util" @@ -21,8 +20,6 @@ import ( log "github.com/sirupsen/logrus" "github.com/tidwall/gjson" "github.com/tidwall/sjson" - "golang.org/x/oauth2" - "golang.org/x/oauth2/google" ) const ( @@ -31,6 +28,9 @@ const ( // glAPIVersion is the API version used for Gemini requests. glAPIVersion = "v1beta" + + // streamScannerBuffer is the buffer size for SSE stream scanning. + streamScannerBuffer = 52_428_800 ) // GeminiExecutor is a stateless executor for the official Gemini API using API keys. @@ -48,9 +48,11 @@ type GeminiExecutor struct { // // Returns: // - *GeminiExecutor: A new Gemini executor instance -func NewGeminiExecutor(cfg *config.Config) *GeminiExecutor { return &GeminiExecutor{cfg: cfg} } +func NewGeminiExecutor(cfg *config.Config) *GeminiExecutor { + return &GeminiExecutor{cfg: cfg} +} -// Identifier returns the executor identifier for Gemini. +// Identifier returns the executor identifier. func (e *GeminiExecutor) Identifier() string { return "gemini" } // PrepareRequest prepares the HTTP request for execution (no-op for Gemini). @@ -249,7 +251,7 @@ func (e *GeminiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A } }() scanner := bufio.NewScanner(httpResp.Body) - scanner.Buffer(nil, 52_428_800) // 50MB + scanner.Buffer(nil, streamScannerBuffer) var param any for scanner.Scan() { line := scanner.Bytes() @@ -353,106 +355,8 @@ func (e *GeminiExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Aut return cliproxyexecutor.Response{Payload: []byte(translated)}, nil } -func (e *GeminiExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) { - log.Debugf("gemini executor: refresh called") - // OAuth bearer token refresh for official Gemini API. - if auth == nil { - return nil, fmt.Errorf("gemini executor: auth is nil") - } - if auth.Metadata == nil { - return auth, nil - } - // Token data is typically nested under "token" map in Gemini files. - tokenMap, _ := auth.Metadata["token"].(map[string]any) - var refreshToken, accessToken, clientID, clientSecret, tokenURI, expiryStr string - if tokenMap != nil { - if v, ok := tokenMap["refresh_token"].(string); ok { - refreshToken = v - } - if v, ok := tokenMap["access_token"].(string); ok { - accessToken = v - } - if v, ok := tokenMap["client_id"].(string); ok { - clientID = v - } - if v, ok := tokenMap["client_secret"].(string); ok { - clientSecret = v - } - if v, ok := tokenMap["token_uri"].(string); ok { - tokenURI = v - } - if v, ok := tokenMap["expiry"].(string); ok { - expiryStr = v - } - } else { - // Fallback to top-level keys if present - if v, ok := auth.Metadata["refresh_token"].(string); ok { - refreshToken = v - } - if v, ok := auth.Metadata["access_token"].(string); ok { - accessToken = v - } - if v, ok := auth.Metadata["client_id"].(string); ok { - clientID = v - } - if v, ok := auth.Metadata["client_secret"].(string); ok { - clientSecret = v - } - if v, ok := auth.Metadata["token_uri"].(string); ok { - tokenURI = v - } - if v, ok := auth.Metadata["expiry"].(string); ok { - expiryStr = v - } - } - if refreshToken == "" { - // Nothing to do for API key or cookie based entries - return auth, nil - } - - // Prepare oauth2 config; default to Google endpoints - endpoint := google.Endpoint - if tokenURI != "" { - endpoint.TokenURL = tokenURI - } - conf := &oauth2.Config{ClientID: clientID, ClientSecret: clientSecret, Endpoint: endpoint} - - // Ensure proxy-aware HTTP client for token refresh - httpClient := util.SetProxy(&e.cfg.SDKConfig, &http.Client{}) - ctx = context.WithValue(ctx, oauth2.HTTPClient, httpClient) - - // Build base token - tok := &oauth2.Token{AccessToken: accessToken, RefreshToken: refreshToken} - if t, err := time.Parse(time.RFC3339, expiryStr); err == nil { - tok.Expiry = t - } - newTok, err := conf.TokenSource(ctx, tok).Token() - if err != nil { - return nil, err - } - - // Persist back to metadata; prefer nested token map if present - if tokenMap == nil { - tokenMap = make(map[string]any) - } - tokenMap["access_token"] = newTok.AccessToken - tokenMap["refresh_token"] = newTok.RefreshToken - tokenMap["expiry"] = newTok.Expiry.Format(time.RFC3339) - if clientID != "" { - tokenMap["client_id"] = clientID - } - if clientSecret != "" { - tokenMap["client_secret"] = clientSecret - } - if tokenURI != "" { - tokenMap["token_uri"] = tokenURI - } - auth.Metadata["token"] = tokenMap - - // Also mirror top-level access_token for compatibility if previously present - if _, ok := auth.Metadata["access_token"]; ok { - auth.Metadata["access_token"] = newTok.AccessToken - } +// Refresh refreshes the authentication credentials (no-op for Gemini API key). +func (e *GeminiExecutor) Refresh(_ context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) { return auth, nil } diff --git a/internal/runtime/executor/gemini_vertex_executor.go b/internal/runtime/executor/gemini_vertex_executor.go index 51a6118c..80a110ee 100644 --- a/internal/runtime/executor/gemini_vertex_executor.go +++ b/internal/runtime/executor/gemini_vertex_executor.go @@ -1,6 +1,6 @@ -// Package executor contains provider executors. This file implements the Vertex AI -// Gemini executor that talks to Google Vertex AI endpoints using service account -// credentials imported by the CLI. +// Package executor provides runtime execution capabilities for various AI service providers. +// This file implements the Vertex AI Gemini executor that talks to Google Vertex AI +// endpoints using service account credentials or API keys. package executor import ( @@ -36,15 +36,21 @@ type GeminiVertexExecutor struct { cfg *config.Config } -// NewGeminiVertexExecutor constructs the Vertex executor. +// NewGeminiVertexExecutor creates a new Vertex AI Gemini executor instance. +// +// Parameters: +// - cfg: The application configuration +// +// Returns: +// - *GeminiVertexExecutor: A new Vertex AI Gemini executor instance func NewGeminiVertexExecutor(cfg *config.Config) *GeminiVertexExecutor { return &GeminiVertexExecutor{cfg: cfg} } -// Identifier returns provider key for manager routing. +// Identifier returns the executor identifier. func (e *GeminiVertexExecutor) Identifier() string { return "vertex" } -// PrepareRequest is a no-op for Vertex. +// PrepareRequest prepares the HTTP request for execution (no-op for Vertex). func (e *GeminiVertexExecutor) PrepareRequest(_ *http.Request, _ *cliproxyauth.Auth) error { return nil } @@ -281,7 +287,7 @@ func (e *GeminiVertexExecutor) countTokensWithAPIKey(ctx context.Context, auth * return cliproxyexecutor.Response{Payload: []byte(out)}, nil } -// Refresh is a no-op for service account based credentials. +// Refresh refreshes the authentication credentials (no-op for Vertex). func (e *GeminiVertexExecutor) Refresh(_ context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) { return auth, nil } @@ -579,7 +585,7 @@ func (e *GeminiVertexExecutor) executeStreamWithServiceAccount(ctx context.Conte } }() scanner := bufio.NewScanner(httpResp.Body) - scanner.Buffer(nil, 52_428_800) // 50MB + scanner.Buffer(nil, streamScannerBuffer) var param any for scanner.Scan() { line := scanner.Bytes() @@ -696,7 +702,7 @@ func (e *GeminiVertexExecutor) executeStreamWithAPIKey(ctx context.Context, auth } }() scanner := bufio.NewScanner(httpResp.Body) - scanner.Buffer(nil, 52_428_800) // 50MB + scanner.Buffer(nil, streamScannerBuffer) var param any for scanner.Scan() { line := scanner.Bytes()