From 0607e52767c1cf779fd19072f1020c2ef636d736 Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Mon, 22 Sep 2025 09:27:03 +0800 Subject: [PATCH] feat: implement token refresh support for executors - Added `Refresh` method implementations for Codex, Claude, Gemini, and Qwen executors. - Introduced OAuth-based token handling for Gemini and Qwen with support for refresh tokens. - Updated Codex and Claude to use new internal auth services. - Enhanced metadata structure and consistency for token storage across all executors. --- internal/runtime/executor/claude_executor.go | 29 +++++- internal/runtime/executor/codex_executor.go | 34 +++++- internal/runtime/executor/gemini_executor.go | 104 ++++++++++++++++++- internal/runtime/executor/qwen_executor.go | 35 ++++++- 4 files changed, 197 insertions(+), 5 deletions(-) diff --git a/internal/runtime/executor/claude_executor.go b/internal/runtime/executor/claude_executor.go index 07bc64f2..5beaca65 100644 --- a/internal/runtime/executor/claude_executor.go +++ b/internal/runtime/executor/claude_executor.go @@ -9,6 +9,7 @@ import ( "net/http" "strings" + claudeauth "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/claude" "github.com/router-for-me/CLIProxyAPI/v6/internal/config" "github.com/router-for-me/CLIProxyAPI/v6/internal/misc" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth" @@ -141,7 +142,33 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A } func (e *ClaudeExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) { - _ = ctx + if auth == nil { + return nil, fmt.Errorf("claude executor: auth is nil") + } + var refreshToken string + if auth.Metadata != nil { + if v, ok := auth.Metadata["refresh_token"].(string); ok && v != "" { + refreshToken = v + } + } + if refreshToken == "" { + return auth, nil + } + svc := claudeauth.NewClaudeAuth(e.cfg) + td, err := svc.RefreshTokens(ctx, refreshToken) + if err != nil { + return nil, err + } + if auth.Metadata == nil { + auth.Metadata = make(map[string]any) + } + auth.Metadata["access_token"] = td.AccessToken + if td.RefreshToken != "" { + auth.Metadata["refresh_token"] = td.RefreshToken + } + auth.Metadata["email"] = td.Email + auth.Metadata["expired"] = td.Expire + auth.Metadata["type"] = "claude" return auth, nil } diff --git a/internal/runtime/executor/codex_executor.go b/internal/runtime/executor/codex_executor.go index f41bb704..caaad9a3 100644 --- a/internal/runtime/executor/codex_executor.go +++ b/internal/runtime/executor/codex_executor.go @@ -8,6 +8,7 @@ import ( "net/http" "strings" + codexauth "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/codex" "github.com/router-for-me/CLIProxyAPI/v6/internal/config" "github.com/router-for-me/CLIProxyAPI/v6/internal/util" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth" @@ -187,7 +188,38 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au } func (e *CodexExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) { - _ = ctx + if auth == nil { + return nil, statusErr{code: 500, msg: "codex executor: auth is nil"} + } + var refreshToken string + if auth.Metadata != nil { + if v, ok := auth.Metadata["refresh_token"].(string); ok && v != "" { + refreshToken = v + } + } + if refreshToken == "" { + return auth, nil + } + svc := codexauth.NewCodexAuth(e.cfg) + td, err := svc.RefreshTokensWithRetry(ctx, refreshToken, 3) + if err != nil { + return nil, err + } + if auth.Metadata == nil { + auth.Metadata = make(map[string]any) + } + auth.Metadata["id_token"] = td.IDToken + auth.Metadata["access_token"] = td.AccessToken + if td.RefreshToken != "" { + auth.Metadata["refresh_token"] = td.RefreshToken + } + if td.AccountID != "" { + auth.Metadata["account_id"] = td.AccountID + } + auth.Metadata["email"] = td.Email + // Use unified key in files + auth.Metadata["expired"] = td.Expire + auth.Metadata["type"] = "codex" return auth, nil } diff --git a/internal/runtime/executor/gemini_executor.go b/internal/runtime/executor/gemini_executor.go index f9e86b21..d8243027 100644 --- a/internal/runtime/executor/gemini_executor.go +++ b/internal/runtime/executor/gemini_executor.go @@ -7,11 +7,15 @@ import ( "fmt" "io" "net/http" + "time" "github.com/router-for-me/CLIProxyAPI/v6/internal/config" + "github.com/router-for-me/CLIProxyAPI/v6/internal/util" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor" sdktranslator "github.com/router-for-me/CLIProxyAPI/v6/sdk/translator" + "golang.org/x/oauth2" + "golang.org/x/oauth2/google" ) const ( @@ -161,8 +165,104 @@ func (e *GeminiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A } func (e *GeminiExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) { - // API-key based: no-op; cookie-based handled by legacy fallback when used. - _ = ctx + // 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, &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 + } return auth, nil } diff --git a/internal/runtime/executor/qwen_executor.go b/internal/runtime/executor/qwen_executor.go index d9d0a8e7..7eef2c3e 100644 --- a/internal/runtime/executor/qwen_executor.go +++ b/internal/runtime/executor/qwen_executor.go @@ -9,6 +9,7 @@ import ( "net/http" "strings" + qwenauth "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/qwen" "github.com/router-for-me/CLIProxyAPI/v6/internal/config" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor" @@ -143,7 +144,39 @@ func (e *QwenExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Aut } func (e *QwenExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) { - _ = ctx + if auth == nil { + return nil, fmt.Errorf("qwen executor: auth is nil") + } + // Expect refresh_token in metadata for OAuth-based accounts + var refreshToken string + if auth.Metadata != nil { + if v, ok := auth.Metadata["refresh_token"].(string); ok && strings.TrimSpace(v) != "" { + refreshToken = v + } + } + if strings.TrimSpace(refreshToken) == "" { + // Nothing to refresh + return auth, nil + } + + svc := qwenauth.NewQwenAuth(e.cfg) + td, err := svc.RefreshTokens(ctx, refreshToken) + if err != nil { + return nil, err + } + if auth.Metadata == nil { + auth.Metadata = make(map[string]any) + } + auth.Metadata["access_token"] = td.AccessToken + if td.RefreshToken != "" { + auth.Metadata["refresh_token"] = td.RefreshToken + } + if td.ResourceURL != "" { + auth.Metadata["resource_url"] = td.ResourceURL + } + // Use "expired" for consistency with existing file format + auth.Metadata["expired"] = td.Expire + auth.Metadata["type"] = "qwen" return auth, nil }