feat(service): 改进 Gemini OAuth 服务层,区分 Code Assist 和 AI Studio 客户端

OAuth 服务改进:
- 添加 GetOAuthConfig 返回 AI Studio OAuth 可用性
- Code Assist 强制使用内置 Gemini CLI 客户端
- AI Studio OAuth 要求用户配置自定义 OAuth 客户端
- ExchangeCode/RefreshToken 接口添加 oauthType 参数
- 添加 unauthorized_client 错误的向后兼容重试逻辑

兼容层改进:
- 403 重试逻辑仅对 Code Assist OAuth 生效
- 添加 insufficient-scope 错误检测,避免无效重试
- 上游错误消息脱敏处理(隐藏 API key 等敏感信息)
- 改进错误提示,显示更多上游错误详情
This commit is contained in:
ianshaw
2025-12-25 23:51:38 -08:00
parent eea949853a
commit 456e8984b0
3 changed files with 146 additions and 25 deletions

View File

@@ -386,12 +386,21 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex
sleepGeminiBackoff(attempt) sleepGeminiBackoff(attempt)
continue continue
} }
return nil, s.writeClaudeError(c, http.StatusBadGateway, "upstream_error", "Upstream request failed after retries") return nil, s.writeClaudeError(c, http.StatusBadGateway, "upstream_error", "Upstream request failed after retries: "+sanitizeUpstreamErrorMessage(err.Error()))
} }
if resp.StatusCode >= 400 && s.shouldRetryGeminiUpstreamError(account, resp.StatusCode) { if resp.StatusCode >= 400 && s.shouldRetryGeminiUpstreamError(account, resp.StatusCode) {
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
_ = resp.Body.Close() _ = resp.Body.Close()
// Don't treat insufficient-scope as transient.
if resp.StatusCode == 403 && isGeminiInsufficientScope(resp.Header, respBody) {
resp = &http.Response{
StatusCode: resp.StatusCode,
Header: resp.Header.Clone(),
Body: io.NopCloser(bytes.NewReader(respBody)),
}
break
}
if resp.StatusCode == 429 { if resp.StatusCode == 429 {
// Mark as rate-limited early so concurrent requests avoid this account. // Mark as rate-limited early so concurrent requests avoid this account.
s.handleGeminiUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody) s.handleGeminiUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
@@ -401,7 +410,13 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex
sleepGeminiBackoff(attempt) sleepGeminiBackoff(attempt)
continue continue
} }
return nil, s.writeClaudeError(c, http.StatusBadGateway, "upstream_error", "Upstream request failed after retries") // Final attempt: surface the upstream error body (mapped below) instead of a generic retry error.
resp = &http.Response{
StatusCode: resp.StatusCode,
Header: resp.Header.Clone(),
Body: io.NopCloser(bytes.NewReader(respBody)),
}
break
} }
break break
@@ -633,12 +648,21 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
FirstTokenMs: nil, FirstTokenMs: nil,
}, nil }, nil
} }
return nil, s.writeGoogleError(c, http.StatusBadGateway, "Upstream request failed after retries") return nil, s.writeGoogleError(c, http.StatusBadGateway, "Upstream request failed after retries: "+sanitizeUpstreamErrorMessage(err.Error()))
} }
if resp.StatusCode >= 400 && s.shouldRetryGeminiUpstreamError(account, resp.StatusCode) { if resp.StatusCode >= 400 && s.shouldRetryGeminiUpstreamError(account, resp.StatusCode) {
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
_ = resp.Body.Close() _ = resp.Body.Close()
// Don't treat insufficient-scope as transient.
if resp.StatusCode == 403 && isGeminiInsufficientScope(resp.Header, respBody) {
resp = &http.Response{
StatusCode: resp.StatusCode,
Header: resp.Header.Clone(),
Body: io.NopCloser(bytes.NewReader(respBody)),
}
break
}
if resp.StatusCode == 429 { if resp.StatusCode == 429 {
s.handleGeminiUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody) s.handleGeminiUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
} }
@@ -659,7 +683,13 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
FirstTokenMs: nil, FirstTokenMs: nil,
}, nil }, nil
} }
return nil, s.writeGoogleError(c, http.StatusBadGateway, "Upstream request failed after retries") // Final attempt: surface the upstream error body (passed through below) instead of a generic retry error.
resp = &http.Response{
StatusCode: resp.StatusCode,
Header: resp.Header.Clone(),
Body: io.NopCloser(bytes.NewReader(respBody)),
}
break
} }
break break
@@ -752,7 +782,15 @@ func (s *GeminiMessagesCompatService) shouldRetryGeminiUpstreamError(account *mo
return true return true
case 403: case 403:
// GeminiCli OAuth occasionally returns 403 transiently (activation/quota propagation); allow retry. // GeminiCli OAuth occasionally returns 403 transiently (activation/quota propagation); allow retry.
return account != nil && account.Type == model.AccountTypeOAuth if account == nil || account.Type != model.AccountTypeOAuth {
return false
}
oauthType := strings.ToLower(strings.TrimSpace(account.GetCredential("oauth_type")))
if oauthType == "" && strings.TrimSpace(account.GetCredential("project_id")) != "" {
// Legacy/implicit Code Assist OAuth accounts.
oauthType = "code_assist"
}
return oauthType == "code_assist"
default: default:
return false return false
} }
@@ -774,6 +812,15 @@ func sleepGeminiBackoff(attempt int) {
time.Sleep(sleepFor) time.Sleep(sleepFor)
} }
var sensitiveQueryParamRegex = regexp.MustCompile(`(?i)([?&](?:key|client_secret|access_token|refresh_token)=)[^&"\s]+`)
func sanitizeUpstreamErrorMessage(msg string) string {
if msg == "" {
return msg
}
return sensitiveQueryParamRegex.ReplaceAllString(msg, `$1***`)
}
func (s *GeminiMessagesCompatService) writeGeminiMappedError(c *gin.Context, upstreamStatus int, body []byte) error { func (s *GeminiMessagesCompatService) writeGeminiMappedError(c *gin.Context, upstreamStatus int, body []byte) error {
var statusCode int var statusCode int
var errType, errMsg string var errType, errMsg string

View File

@@ -8,6 +8,6 @@ import (
// GeminiOAuthClient performs Google OAuth token exchange/refresh for Gemini integration. // GeminiOAuthClient performs Google OAuth token exchange/refresh for Gemini integration.
type GeminiOAuthClient interface { type GeminiOAuthClient interface {
ExchangeCode(ctx context.Context, code, codeVerifier, redirectURI, proxyURL string) (*geminicli.TokenResponse, error) ExchangeCode(ctx context.Context, oauthType, code, codeVerifier, redirectURI, proxyURL string) (*geminicli.TokenResponse, error)
RefreshToken(ctx context.Context, refreshToken, proxyURL string) (*geminicli.TokenResponse, error) RefreshToken(ctx context.Context, oauthType, refreshToken, proxyURL string) (*geminicli.TokenResponse, error)
} }

View File

@@ -25,6 +25,11 @@ type GeminiOAuthService struct {
cfg *config.Config cfg *config.Config
} }
type GeminiOAuthCapabilities struct {
AIStudioOAuthEnabled bool `json:"ai_studio_oauth_enabled"`
RequiredRedirectURIs []string `json:"required_redirect_uris"`
}
func NewGeminiOAuthService( func NewGeminiOAuthService(
proxyRepo ProxyRepository, proxyRepo ProxyRepository,
oauthClient GeminiOAuthClient, oauthClient GeminiOAuthClient,
@@ -40,6 +45,19 @@ func NewGeminiOAuthService(
} }
} }
func (s *GeminiOAuthService) GetOAuthConfig() *GeminiOAuthCapabilities {
// AI Studio OAuth is only enabled when the operator configures a custom OAuth client.
clientID := strings.TrimSpace(s.cfg.Gemini.OAuth.ClientID)
clientSecret := strings.TrimSpace(s.cfg.Gemini.OAuth.ClientSecret)
enabled := clientID != "" && clientSecret != "" &&
!(clientID == geminicli.GeminiCLIOAuthClientID && clientSecret == geminicli.GeminiCLIOAuthClientSecret)
return &GeminiOAuthCapabilities{
AIStudioOAuthEnabled: enabled,
RequiredRedirectURIs: []string{geminicli.AIStudioOAuthRedirectURI},
}
}
type GeminiAuthURLResult struct { type GeminiAuthURLResult struct {
AuthURL string `json:"auth_url"` AuthURL string `json:"auth_url"`
SessionID string `json:"session_id"` SessionID string `json:"session_id"`
@@ -69,13 +87,18 @@ func (s *GeminiOAuthService) GenerateAuthURL(ctx context.Context, proxyID *int64
} }
} }
// 两种 OAuth 模式都使用相同的配置,只是 scopes 不同 // OAuth client selection:
// scopes 会在 EffectiveOAuthConfig 中根据 oauthType 自动选择 // - code_assist: always use built-in Gemini CLI OAuth client (public), regardless of configured client_id/secret.
// - ai_studio: requires a user-provided OAuth client.
oauthCfg := geminicli.OAuthConfig{ oauthCfg := geminicli.OAuthConfig{
ClientID: s.cfg.Gemini.OAuth.ClientID, ClientID: s.cfg.Gemini.OAuth.ClientID,
ClientSecret: s.cfg.Gemini.OAuth.ClientSecret, ClientSecret: s.cfg.Gemini.OAuth.ClientSecret,
Scopes: s.cfg.Gemini.OAuth.Scopes, Scopes: s.cfg.Gemini.OAuth.Scopes,
} }
if oauthType == "code_assist" {
oauthCfg.ClientID = ""
oauthCfg.ClientSecret = ""
}
session := &geminicli.OAuthSession{ session := &geminicli.OAuthSession{
State: state, State: state,
@@ -93,12 +116,24 @@ func (s *GeminiOAuthService) GenerateAuthURL(ctx context.Context, proxyID *int64
return nil, err return nil, err
} }
// For Code Assist with Gemini CLI credentials, use the CLI's redirect URI isBuiltinClient := effectiveCfg.ClientID == geminicli.GeminiCLIOAuthClientID &&
effectiveCfg.ClientSecret == geminicli.GeminiCLIOAuthClientSecret
// AI Studio OAuth requires a user-provided OAuth client (built-in Gemini CLI client is scope-restricted).
if oauthType == "ai_studio" && isBuiltinClient {
return nil, fmt.Errorf("AI Studio OAuth requires a custom OAuth Client (GEMINI_OAUTH_CLIENT_ID / GEMINI_OAUTH_CLIENT_SECRET). If you don't want to configure an OAuth client, please use an AI Studio API Key account instead")
}
// Redirect URI strategy:
// - code_assist: use Gemini CLI redirect URI (codeassist.google.com/authcode)
// - ai_studio: use localhost callback for manual copy/paste flow
if oauthType == "code_assist" { if oauthType == "code_assist" {
redirectURI = geminicli.GeminiCLIRedirectURI redirectURI = geminicli.GeminiCLIRedirectURI
session.RedirectURI = redirectURI } else {
s.sessionStore.Set(sessionID, session) redirectURI = geminicli.AIStudioOAuthRedirectURI
} }
session.RedirectURI = redirectURI
s.sessionStore.Set(sessionID, session)
authURL, err := geminicli.BuildAuthorizationURL(effectiveCfg, state, codeChallenge, redirectURI, session.ProjectID, oauthType) authURL, err := geminicli.BuildAuthorizationURL(effectiveCfg, state, codeChallenge, redirectURI, session.ProjectID, oauthType)
if err != nil { if err != nil {
@@ -150,15 +185,39 @@ func (s *GeminiOAuthService) ExchangeCode(ctx context.Context, input *GeminiExch
redirectURI := session.RedirectURI redirectURI := session.RedirectURI
tokenResp, err := s.oauthClient.ExchangeCode(ctx, input.Code, session.CodeVerifier, redirectURI, proxyURL) // Resolve oauth_type early (defaults to code_assist for backward compatibility).
oauthType := session.OAuthType
if oauthType == "" {
oauthType = "code_assist"
}
// If the session was created for AI Studio OAuth, ensure a custom OAuth client is configured.
if oauthType == "ai_studio" {
effectiveCfg, err := geminicli.EffectiveOAuthConfig(geminicli.OAuthConfig{
ClientID: s.cfg.Gemini.OAuth.ClientID,
ClientSecret: s.cfg.Gemini.OAuth.ClientSecret,
Scopes: s.cfg.Gemini.OAuth.Scopes,
}, "ai_studio")
if err != nil {
return nil, err
}
isBuiltinClient := effectiveCfg.ClientID == geminicli.GeminiCLIOAuthClientID &&
effectiveCfg.ClientSecret == geminicli.GeminiCLIOAuthClientSecret
if isBuiltinClient {
return nil, fmt.Errorf("AI Studio OAuth requires a custom OAuth Client. Please use an AI Studio API Key account, or configure GEMINI_OAUTH_CLIENT_ID / GEMINI_OAUTH_CLIENT_SECRET and re-authorize")
}
}
// code_assist always uses the built-in client and its fixed redirect URI.
if oauthType == "code_assist" {
redirectURI = geminicli.GeminiCLIRedirectURI
}
tokenResp, err := s.oauthClient.ExchangeCode(ctx, oauthType, input.Code, session.CodeVerifier, redirectURI, proxyURL)
if err != nil { if err != nil {
return nil, fmt.Errorf("failed to exchange code: %w", err) return nil, fmt.Errorf("failed to exchange code: %w", err)
} }
sessionProjectID := strings.TrimSpace(session.ProjectID) sessionProjectID := strings.TrimSpace(session.ProjectID)
oauthType := session.OAuthType
if oauthType == "" {
oauthType = "code_assist" // 默认为 code_assist 以兼容旧 session
}
s.sessionStore.Delete(input.SessionID) s.sessionStore.Delete(input.SessionID)
// 计算过期时间时减去 5 分钟安全时间窗口,考虑网络延迟和时钟偏差 // 计算过期时间时减去 5 分钟安全时间窗口,考虑网络延迟和时钟偏差
@@ -194,7 +253,7 @@ func (s *GeminiOAuthService) ExchangeCode(ctx context.Context, input *GeminiExch
}, nil }, nil
} }
func (s *GeminiOAuthService) RefreshToken(ctx context.Context, refreshToken, proxyURL string) (*GeminiTokenInfo, error) { func (s *GeminiOAuthService) RefreshToken(ctx context.Context, oauthType, refreshToken, proxyURL string) (*GeminiTokenInfo, error) {
var lastErr error var lastErr error
for attempt := 0; attempt <= 3; attempt++ { for attempt := 0; attempt <= 3; attempt++ {
@@ -206,7 +265,7 @@ func (s *GeminiOAuthService) RefreshToken(ctx context.Context, refreshToken, pro
time.Sleep(backoff) time.Sleep(backoff)
} }
tokenResp, err := s.oauthClient.RefreshToken(ctx, refreshToken, proxyURL) tokenResp, err := s.oauthClient.RefreshToken(ctx, oauthType, refreshToken, proxyURL)
if err == nil { if err == nil {
// 计算过期时间时减去 5 分钟安全时间窗口,考虑网络延迟和时钟偏差 // 计算过期时间时减去 5 分钟安全时间窗口,考虑网络延迟和时钟偏差
expiresAt := time.Now().Unix() + tokenResp.ExpiresIn - 300 expiresAt := time.Now().Unix() + tokenResp.ExpiresIn - 300
@@ -255,6 +314,12 @@ func (s *GeminiOAuthService) RefreshAccountToken(ctx context.Context, account *m
return nil, fmt.Errorf("no refresh token available") return nil, fmt.Errorf("no refresh token available")
} }
// Preserve oauth_type from the account (defaults to code_assist for backward compatibility).
oauthType := strings.TrimSpace(account.GetCredential("oauth_type"))
if oauthType == "" {
oauthType = "code_assist"
}
var proxyURL string var proxyURL string
if account.ProxyID != nil { if account.ProxyID != nil {
proxy, err := s.proxyRepo.GetByID(ctx, *account.ProxyID) proxy, err := s.proxyRepo.GetByID(ctx, *account.ProxyID)
@@ -263,16 +328,25 @@ func (s *GeminiOAuthService) RefreshAccountToken(ctx context.Context, account *m
} }
} }
tokenInfo, err := s.RefreshToken(ctx, refreshToken, proxyURL) tokenInfo, err := s.RefreshToken(ctx, oauthType, refreshToken, proxyURL)
// Backward compatibility:
// Older versions could refresh Code Assist tokens using a user-provided OAuth client when configured.
// If the refresh token was originally issued to that custom client, forcing the built-in client will
// fail with "unauthorized_client". In that case, retry with the custom client (ai_studio path) when available.
if err != nil && oauthType == "code_assist" && strings.Contains(err.Error(), "unauthorized_client") && s.GetOAuthConfig().AIStudioOAuthEnabled {
if alt, altErr := s.RefreshToken(ctx, "ai_studio", refreshToken, proxyURL); altErr == nil {
tokenInfo = alt
err = nil
}
}
if err != nil { if err != nil {
// Provide a more actionable error for common OAuth client mismatch issues.
if strings.Contains(err.Error(), "unauthorized_client") {
return nil, fmt.Errorf("%w (OAuth client mismatch: the refresh_token is bound to the OAuth client used during authorization; please re-authorize this account or restore the original GEMINI_OAUTH_CLIENT_ID/SECRET)", err)
}
return nil, err return nil, err
} }
// Preserve oauth_type from the account (defaults to code_assist for backward compatibility).
oauthType := strings.TrimSpace(account.GetCredential("oauth_type"))
if oauthType == "" {
oauthType = "code_assist"
}
tokenInfo.OAuthType = oauthType tokenInfo.OAuthType = oauthType
// Preserve account's project_id when present. // Preserve account's project_id when present.