oadmin/xinghuoapi
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files

57
backend/internal/service/gemini_messages_compat_service.go
View File

@@ -386,12 +386,21 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex
sleepGeminiBackoff(attempt)
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) {
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
_ = 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 {
// Mark as rate-limited early so concurrent requests avoid this account.
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)
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
@@ -633,12 +648,21 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
FirstTokenMs: 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) {
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
_ = 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 {
s.handleGeminiUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
}
@@ -659,7 +683,13 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
FirstTokenMs: 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
@@ -752,7 +782,15 @@ func (s *GeminiMessagesCompatService) shouldRetryGeminiUpstreamError(account *mo
return true
case 403:
// 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:
return false
}
@@ -774,6 +812,15 @@ func sleepGeminiBackoff(attempt int) {
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 {
var statusCode int
var errType, errMsg string

4
backend/internal/service/gemini_oauth.go
View File

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

110
backend/internal/service/gemini_oauth_service.go
View File

@@ -25,6 +25,11 @@ type GeminiOAuthService struct {
cfg *config.Config
}
type GeminiOAuthCapabilities struct {
AIStudioOAuthEnabled bool `json:"ai_studio_oauth_enabled"`
RequiredRedirectURIs []string `json:"required_redirect_uris"`
}
func NewGeminiOAuthService(
proxyRepo ProxyRepository,
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 {
AuthURL string `json:"auth_url"`
SessionID string `json:"session_id"`
@@ -69,13 +87,18 @@ func (s *GeminiOAuthService) GenerateAuthURL(ctx context.Context, proxyID *int64
}
}
// 两种 OAuth 模式都使用相同的配置，只是 scopes 不同
// scopes 会在 EffectiveOAuthConfig 中根据 oauthType 自动选择
// OAuth client selection:
// - 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{
ClientID: s.cfg.Gemini.OAuth.ClientID,
ClientSecret: s.cfg.Gemini.OAuth.ClientSecret,
Scopes: s.cfg.Gemini.OAuth.Scopes,
}
if oauthType == "code_assist" {
oauthCfg.ClientID = ""
oauthCfg.ClientSecret = ""
}
session := &geminicli.OAuthSession{
State: state,
@@ -93,12 +116,24 @@ func (s *GeminiOAuthService) GenerateAuthURL(ctx context.Context, proxyID *int64
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" {
redirectURI = geminicli.GeminiCLIRedirectURI
session.RedirectURI = redirectURI
s.sessionStore.Set(sessionID, session)
} else {
redirectURI = geminicli.AIStudioOAuthRedirectURI
}
session.RedirectURI = redirectURI
s.sessionStore.Set(sessionID, session)
authURL, err := geminicli.BuildAuthorizationURL(effectiveCfg, state, codeChallenge, redirectURI, session.ProjectID, oauthType)
if err != nil {
@@ -150,15 +185,39 @@ func (s *GeminiOAuthService) ExchangeCode(ctx context.Context, input *GeminiExch
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 {
return nil, fmt.Errorf("failed to exchange code: %w", err)
}
sessionProjectID := strings.TrimSpace(session.ProjectID)
oauthType := session.OAuthType
if oauthType == "" {
oauthType = "code_assist" // 默认为 code_assist 以兼容旧 session
}
s.sessionStore.Delete(input.SessionID)
// 计算过期时间时减去 5 分钟安全时间窗口,考虑网络延迟和时钟偏差
@@ -194,7 +253,7 @@ func (s *GeminiOAuthService) ExchangeCode(ctx context.Context, input *GeminiExch
}, 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
for attempt := 0; attempt <= 3; attempt++ {
@@ -206,7 +265,7 @@ func (s *GeminiOAuthService) RefreshToken(ctx context.Context, refreshToken, pro
time.Sleep(backoff)
}
tokenResp, err := s.oauthClient.RefreshToken(ctx, refreshToken, proxyURL)
tokenResp, err := s.oauthClient.RefreshToken(ctx, oauthType, refreshToken, proxyURL)
if err == nil {
// 计算过期时间时减去 5 分钟安全时间窗口,考虑网络延迟和时钟偏差
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")
}
// 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
if account.ProxyID != nil {
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 {
// 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
}
// 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
// Preserve account's project_id when present.