674 lines
23 KiB
Go
674 lines
23 KiB
Go
package handler
|
||
|
||
import (
|
||
"bytes"
|
||
"context"
|
||
"crypto/sha256"
|
||
"encoding/hex"
|
||
"encoding/json"
|
||
"errors"
|
||
"io"
|
||
"log"
|
||
"net/http"
|
||
"regexp"
|
||
"strings"
|
||
"time"
|
||
|
||
"github.com/Wei-Shaw/sub2api/internal/domain"
|
||
"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
|
||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||
"github.com/Wei-Shaw/sub2api/internal/pkg/gemini"
|
||
"github.com/Wei-Shaw/sub2api/internal/pkg/googleapi"
|
||
"github.com/Wei-Shaw/sub2api/internal/pkg/ip"
|
||
"github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||
"github.com/google/uuid"
|
||
|
||
"github.com/gin-gonic/gin"
|
||
)
|
||
|
||
// geminiCLITmpDirRegex 用于从 Gemini CLI 请求体中提取 tmp 目录的哈希值
|
||
// 匹配格式: /Users/xxx/.gemini/tmp/[64位十六进制哈希]
|
||
var geminiCLITmpDirRegex = regexp.MustCompile(`/\.gemini/tmp/([A-Fa-f0-9]{64})`)
|
||
|
||
// GeminiV1BetaListModels proxies:
|
||
// GET /v1beta/models
|
||
func (h *GatewayHandler) GeminiV1BetaListModels(c *gin.Context) {
|
||
apiKey, ok := middleware.GetAPIKeyFromContext(c)
|
||
if !ok || apiKey == nil {
|
||
googleError(c, http.StatusUnauthorized, "Invalid API key")
|
||
return
|
||
}
|
||
// 检查平台：优先使用强制平台（/antigravity 路由），否则要求 gemini 分组
|
||
forcePlatform, hasForcePlatform := middleware.GetForcePlatformFromContext(c)
|
||
if !hasForcePlatform && (apiKey.Group == nil || apiKey.Group.Platform != service.PlatformGemini) {
|
||
googleError(c, http.StatusBadRequest, "API key group platform is not gemini")
|
||
return
|
||
}
|
||
|
||
// 强制 antigravity 模式：返回 antigravity 支持的模型列表
|
||
if forcePlatform == service.PlatformAntigravity {
|
||
c.JSON(http.StatusOK, antigravity.FallbackGeminiModelsList())
|
||
return
|
||
}
|
||
|
||
account, err := h.geminiCompatService.SelectAccountForAIStudioEndpoints(c.Request.Context(), apiKey.GroupID)
|
||
if err != nil {
|
||
// 没有 gemini 账户，检查是否有 antigravity 账户可用
|
||
hasAntigravity, _ := h.geminiCompatService.HasAntigravityAccounts(c.Request.Context(), apiKey.GroupID)
|
||
if hasAntigravity {
|
||
// antigravity 账户使用静态模型列表
|
||
c.JSON(http.StatusOK, gemini.FallbackModelsList())
|
||
return
|
||
}
|
||
googleError(c, http.StatusServiceUnavailable, "No available Gemini accounts: "+err.Error())
|
||
return
|
||
}
|
||
|
||
res, err := h.geminiCompatService.ForwardAIStudioGET(c.Request.Context(), account, "/v1beta/models")
|
||
if err != nil {
|
||
googleError(c, http.StatusBadGateway, err.Error())
|
||
return
|
||
}
|
||
if shouldFallbackGeminiModels(res) {
|
||
c.JSON(http.StatusOK, gemini.FallbackModelsList())
|
||
return
|
||
}
|
||
writeUpstreamResponse(c, res)
|
||
}
|
||
|
||
// GeminiV1BetaGetModel proxies:
|
||
// GET /v1beta/models/{model}
|
||
func (h *GatewayHandler) GeminiV1BetaGetModel(c *gin.Context) {
|
||
apiKey, ok := middleware.GetAPIKeyFromContext(c)
|
||
if !ok || apiKey == nil {
|
||
googleError(c, http.StatusUnauthorized, "Invalid API key")
|
||
return
|
||
}
|
||
// 检查平台：优先使用强制平台（/antigravity 路由），否则要求 gemini 分组
|
||
forcePlatform, hasForcePlatform := middleware.GetForcePlatformFromContext(c)
|
||
if !hasForcePlatform && (apiKey.Group == nil || apiKey.Group.Platform != service.PlatformGemini) {
|
||
googleError(c, http.StatusBadRequest, "API key group platform is not gemini")
|
||
return
|
||
}
|
||
|
||
modelName := strings.TrimSpace(c.Param("model"))
|
||
if modelName == "" {
|
||
googleError(c, http.StatusBadRequest, "Missing model in URL")
|
||
return
|
||
}
|
||
|
||
// 强制 antigravity 模式：返回 antigravity 模型信息
|
||
if forcePlatform == service.PlatformAntigravity {
|
||
c.JSON(http.StatusOK, antigravity.FallbackGeminiModel(modelName))
|
||
return
|
||
}
|
||
|
||
account, err := h.geminiCompatService.SelectAccountForAIStudioEndpoints(c.Request.Context(), apiKey.GroupID)
|
||
if err != nil {
|
||
// 没有 gemini 账户，检查是否有 antigravity 账户可用
|
||
hasAntigravity, _ := h.geminiCompatService.HasAntigravityAccounts(c.Request.Context(), apiKey.GroupID)
|
||
if hasAntigravity {
|
||
// antigravity 账户使用静态模型信息
|
||
c.JSON(http.StatusOK, gemini.FallbackModel(modelName))
|
||
return
|
||
}
|
||
googleError(c, http.StatusServiceUnavailable, "No available Gemini accounts: "+err.Error())
|
||
return
|
||
}
|
||
|
||
res, err := h.geminiCompatService.ForwardAIStudioGET(c.Request.Context(), account, "/v1beta/models/"+modelName)
|
||
if err != nil {
|
||
googleError(c, http.StatusBadGateway, err.Error())
|
||
return
|
||
}
|
||
if shouldFallbackGeminiModels(res) {
|
||
c.JSON(http.StatusOK, gemini.FallbackModel(modelName))
|
||
return
|
||
}
|
||
writeUpstreamResponse(c, res)
|
||
}
|
||
|
||
// GeminiV1BetaModels proxies Gemini native REST endpoints like:
|
||
// POST /v1beta/models/{model}:generateContent
|
||
// POST /v1beta/models/{model}:streamGenerateContent?alt=sse
|
||
func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
|
||
apiKey, ok := middleware.GetAPIKeyFromContext(c)
|
||
if !ok || apiKey == nil {
|
||
googleError(c, http.StatusUnauthorized, "Invalid API key")
|
||
return
|
||
}
|
||
authSubject, ok := middleware.GetAuthSubjectFromContext(c)
|
||
if !ok {
|
||
googleError(c, http.StatusInternalServerError, "User context not found")
|
||
return
|
||
}
|
||
|
||
// 检查平台：优先使用强制平台（/antigravity 路由，中间件已设置 request.Context），否则要求 gemini 分组
|
||
if !middleware.HasForcePlatform(c) {
|
||
if apiKey.Group == nil || apiKey.Group.Platform != service.PlatformGemini {
|
||
googleError(c, http.StatusBadRequest, "API key group platform is not gemini")
|
||
return
|
||
}
|
||
}
|
||
|
||
modelName, action, err := parseGeminiModelAction(strings.TrimPrefix(c.Param("modelAction"), "/"))
|
||
if err != nil {
|
||
googleError(c, http.StatusNotFound, err.Error())
|
||
return
|
||
}
|
||
|
||
stream := action == "streamGenerateContent"
|
||
|
||
body, err := io.ReadAll(c.Request.Body)
|
||
if err != nil {
|
||
if maxErr, ok := extractMaxBytesError(err); ok {
|
||
googleError(c, http.StatusRequestEntityTooLarge, buildBodyTooLargeMessage(maxErr.Limit))
|
||
return
|
||
}
|
||
googleError(c, http.StatusBadRequest, "Failed to read request body")
|
||
return
|
||
}
|
||
if len(body) == 0 {
|
||
googleError(c, http.StatusBadRequest, "Request body is empty")
|
||
return
|
||
}
|
||
|
||
setOpsRequestContext(c, modelName, stream, body)
|
||
|
||
// Get subscription (may be nil)
|
||
subscription, _ := middleware.GetSubscriptionFromContext(c)
|
||
|
||
// For Gemini native API, do not send Claude-style ping frames.
|
||
geminiConcurrency := NewConcurrencyHelper(h.concurrencyHelper.concurrencyService, SSEPingFormatNone, 0)
|
||
|
||
// 0) wait queue check
|
||
maxWait := service.CalculateMaxWait(authSubject.Concurrency)
|
||
canWait, err := geminiConcurrency.IncrementWaitCount(c.Request.Context(), authSubject.UserID, maxWait)
|
||
waitCounted := false
|
||
if err != nil {
|
||
log.Printf("Increment wait count failed: %v", err)
|
||
} else if !canWait {
|
||
googleError(c, http.StatusTooManyRequests, "Too many pending requests, please retry later")
|
||
return
|
||
}
|
||
if err == nil && canWait {
|
||
waitCounted = true
|
||
}
|
||
defer func() {
|
||
if waitCounted {
|
||
geminiConcurrency.DecrementWaitCount(c.Request.Context(), authSubject.UserID)
|
||
}
|
||
}()
|
||
|
||
// 1) user concurrency slot
|
||
streamStarted := false
|
||
if h.errorPassthroughService != nil {
|
||
service.BindErrorPassthroughService(c, h.errorPassthroughService)
|
||
}
|
||
userReleaseFunc, err := geminiConcurrency.AcquireUserSlotWithWait(c, authSubject.UserID, authSubject.Concurrency, stream, &streamStarted)
|
||
if err != nil {
|
||
googleError(c, http.StatusTooManyRequests, err.Error())
|
||
return
|
||
}
|
||
if waitCounted {
|
||
geminiConcurrency.DecrementWaitCount(c.Request.Context(), authSubject.UserID)
|
||
waitCounted = false
|
||
}
|
||
// 确保请求取消时也会释放槽位，避免长连接被动中断造成泄漏
|
||
userReleaseFunc = wrapReleaseOnDone(c.Request.Context(), userReleaseFunc)
|
||
if userReleaseFunc != nil {
|
||
defer userReleaseFunc()
|
||
}
|
||
|
||
// 2) billing eligibility check (after wait)
|
||
if err := h.billingCacheService.CheckBillingEligibility(c.Request.Context(), apiKey.User, apiKey, apiKey.Group, subscription); err != nil {
|
||
status, _, message := billingErrorDetails(err)
|
||
googleError(c, status, message)
|
||
return
|
||
}
|
||
|
||
// 3) select account (sticky session based on request body)
|
||
// 优先使用 Gemini CLI 的会话标识（privileged-user-id + tmp 目录哈希）
|
||
sessionHash := extractGeminiCLISessionHash(c, body)
|
||
if sessionHash == "" {
|
||
// Fallback: 使用通用的会话哈希生成逻辑（适用于其他客户端）
|
||
parsedReq, _ := service.ParseGatewayRequest(body, domain.PlatformGemini)
|
||
if parsedReq != nil {
|
||
parsedReq.SessionContext = &service.SessionContext{
|
||
ClientIP: ip.GetClientIP(c),
|
||
UserAgent: c.GetHeader("User-Agent"),
|
||
APIKeyID: apiKey.ID,
|
||
}
|
||
}
|
||
sessionHash = h.gatewayService.GenerateSessionHash(parsedReq)
|
||
}
|
||
sessionKey := sessionHash
|
||
if sessionHash != "" {
|
||
sessionKey = "gemini:" + sessionHash
|
||
}
|
||
|
||
// 查询粘性会话绑定的账号 ID（用于检测账号切换）
|
||
var sessionBoundAccountID int64
|
||
if sessionKey != "" {
|
||
sessionBoundAccountID, _ = h.gatewayService.GetCachedSessionAccountID(c.Request.Context(), apiKey.GroupID, sessionKey)
|
||
}
|
||
|
||
// === Gemini 内容摘要会话 Fallback 逻辑 ===
|
||
// 当原有会话标识无效时（sessionBoundAccountID == 0），尝试基于内容摘要链匹配
|
||
var geminiDigestChain string
|
||
var geminiPrefixHash string
|
||
var geminiSessionUUID string
|
||
useDigestFallback := sessionBoundAccountID == 0
|
||
|
||
if useDigestFallback {
|
||
// 解析 Gemini 请求体
|
||
var geminiReq antigravity.GeminiRequest
|
||
if err := json.Unmarshal(body, &geminiReq); err == nil && len(geminiReq.Contents) > 0 {
|
||
// 生成摘要链
|
||
geminiDigestChain = service.BuildGeminiDigestChain(&geminiReq)
|
||
if geminiDigestChain != "" {
|
||
// 生成前缀 hash
|
||
userAgent := c.GetHeader("User-Agent")
|
||
clientIP := ip.GetClientIP(c)
|
||
platform := ""
|
||
if apiKey.Group != nil {
|
||
platform = apiKey.Group.Platform
|
||
}
|
||
geminiPrefixHash = service.GenerateGeminiPrefixHash(
|
||
authSubject.UserID,
|
||
apiKey.ID,
|
||
clientIP,
|
||
userAgent,
|
||
platform,
|
||
modelName,
|
||
)
|
||
|
||
// 查找会话
|
||
foundUUID, foundAccountID, found := h.gatewayService.FindGeminiSession(
|
||
c.Request.Context(),
|
||
derefGroupID(apiKey.GroupID),
|
||
geminiPrefixHash,
|
||
geminiDigestChain,
|
||
)
|
||
if found {
|
||
sessionBoundAccountID = foundAccountID
|
||
geminiSessionUUID = foundUUID
|
||
log.Printf("[Gemini] Digest fallback matched: uuid=%s, accountID=%d, chain=%s",
|
||
safeShortPrefix(foundUUID, 8), foundAccountID, truncateDigestChain(geminiDigestChain))
|
||
|
||
// 关键：如果原 sessionKey 为空，使用 prefixHash + uuid 作为 sessionKey
|
||
// 这样 SelectAccountWithLoadAwareness 的粘性会话逻辑会优先使用匹配到的账号
|
||
if sessionKey == "" {
|
||
sessionKey = service.GenerateGeminiDigestSessionKey(geminiPrefixHash, foundUUID)
|
||
}
|
||
_ = h.gatewayService.BindStickySession(c.Request.Context(), apiKey.GroupID, sessionKey, foundAccountID)
|
||
} else {
|
||
// 生成新的会话 UUID
|
||
geminiSessionUUID = uuid.New().String()
|
||
// 为新会话也生成 sessionKey（用于后续请求的粘性会话）
|
||
if sessionKey == "" {
|
||
sessionKey = service.GenerateGeminiDigestSessionKey(geminiPrefixHash, geminiSessionUUID)
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// 判断是否真的绑定了粘性会话：有 sessionKey 且已经绑定到某个账号
|
||
hasBoundSession := sessionKey != "" && sessionBoundAccountID > 0
|
||
cleanedForUnknownBinding := false
|
||
|
||
maxAccountSwitches := h.maxAccountSwitchesGemini
|
||
switchCount := 0
|
||
failedAccountIDs := make(map[int64]struct{})
|
||
var lastFailoverErr *service.UpstreamFailoverError
|
||
var forceCacheBilling bool // 粘性会话切换时的缓存计费标记
|
||
|
||
for {
|
||
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, sessionKey, modelName, failedAccountIDs, "") // Gemini 不使用会话限制
|
||
if err != nil {
|
||
if len(failedAccountIDs) == 0 {
|
||
googleError(c, http.StatusServiceUnavailable, "No available Gemini accounts: "+err.Error())
|
||
return
|
||
}
|
||
h.handleGeminiFailoverExhausted(c, lastFailoverErr)
|
||
return
|
||
}
|
||
account := selection.Account
|
||
setOpsSelectedAccount(c, account.ID)
|
||
|
||
// 检测账号切换：如果粘性会话绑定的账号与当前选择的账号不同，清除 thoughtSignature
|
||
// 注意：Gemini 原生 API 的 thoughtSignature 与具体上游账号强相关；跨账号透传会导致 400。
|
||
if sessionBoundAccountID > 0 && sessionBoundAccountID != account.ID {
|
||
log.Printf("[Gemini] Sticky session account switched: %d -> %d, cleaning thoughtSignature", sessionBoundAccountID, account.ID)
|
||
body = service.CleanGeminiNativeThoughtSignatures(body)
|
||
sessionBoundAccountID = account.ID
|
||
} else if sessionKey != "" && sessionBoundAccountID == 0 && !cleanedForUnknownBinding && bytes.Contains(body, []byte(`"thoughtSignature"`)) {
|
||
// 无缓存绑定但请求里已有 thoughtSignature：常见于缓存丢失/TTL 过期后，客户端继续携带旧签名。
|
||
// 为避免第一次转发就 400，这里做一次确定性清理，让新账号重新生成签名链路。
|
||
log.Printf("[Gemini] Sticky session binding missing, cleaning thoughtSignature proactively")
|
||
body = service.CleanGeminiNativeThoughtSignatures(body)
|
||
cleanedForUnknownBinding = true
|
||
sessionBoundAccountID = account.ID
|
||
} else if sessionBoundAccountID == 0 {
|
||
// 记录本次请求中首次选择到的账号，便于同一请求内 failover 时检测切换。
|
||
sessionBoundAccountID = account.ID
|
||
}
|
||
|
||
// 4) account concurrency slot
|
||
accountReleaseFunc := selection.ReleaseFunc
|
||
if !selection.Acquired {
|
||
if selection.WaitPlan == nil {
|
||
googleError(c, http.StatusServiceUnavailable, "No available Gemini accounts")
|
||
return
|
||
}
|
||
accountWaitCounted := false
|
||
canWait, err := geminiConcurrency.IncrementAccountWaitCount(c.Request.Context(), account.ID, selection.WaitPlan.MaxWaiting)
|
||
if err != nil {
|
||
log.Printf("Increment account wait count failed: %v", err)
|
||
} else if !canWait {
|
||
log.Printf("Account wait queue full: account=%d", account.ID)
|
||
googleError(c, http.StatusTooManyRequests, "Too many pending requests, please retry later")
|
||
return
|
||
}
|
||
if err == nil && canWait {
|
||
accountWaitCounted = true
|
||
}
|
||
defer func() {
|
||
if accountWaitCounted {
|
||
geminiConcurrency.DecrementAccountWaitCount(c.Request.Context(), account.ID)
|
||
}
|
||
}()
|
||
|
||
accountReleaseFunc, err = geminiConcurrency.AcquireAccountSlotWithWaitTimeout(
|
||
c,
|
||
account.ID,
|
||
selection.WaitPlan.MaxConcurrency,
|
||
selection.WaitPlan.Timeout,
|
||
stream,
|
||
&streamStarted,
|
||
)
|
||
if err != nil {
|
||
googleError(c, http.StatusTooManyRequests, err.Error())
|
||
return
|
||
}
|
||
if accountWaitCounted {
|
||
geminiConcurrency.DecrementAccountWaitCount(c.Request.Context(), account.ID)
|
||
accountWaitCounted = false
|
||
}
|
||
if err := h.gatewayService.BindStickySession(c.Request.Context(), apiKey.GroupID, sessionKey, account.ID); err != nil {
|
||
log.Printf("Bind sticky session failed: %v", err)
|
||
}
|
||
}
|
||
// 账号槽位/等待计数需要在超时或断开时安全回收
|
||
accountReleaseFunc = wrapReleaseOnDone(c.Request.Context(), accountReleaseFunc)
|
||
|
||
// 5) forward (根据平台分流)
|
||
var result *service.ForwardResult
|
||
requestCtx := c.Request.Context()
|
||
if switchCount > 0 {
|
||
requestCtx = context.WithValue(requestCtx, ctxkey.AccountSwitchCount, switchCount)
|
||
}
|
||
if account.Platform == service.PlatformAntigravity && account.Type != service.AccountTypeAPIKey {
|
||
result, err = h.antigravityGatewayService.ForwardGemini(requestCtx, c, account, modelName, action, stream, body, hasBoundSession)
|
||
} else {
|
||
result, err = h.geminiCompatService.ForwardNative(requestCtx, c, account, modelName, action, stream, body)
|
||
}
|
||
if accountReleaseFunc != nil {
|
||
accountReleaseFunc()
|
||
}
|
||
if err != nil {
|
||
var failoverErr *service.UpstreamFailoverError
|
||
if errors.As(err, &failoverErr) {
|
||
failedAccountIDs[account.ID] = struct{}{}
|
||
if failoverErr.ForceCacheBilling {
|
||
forceCacheBilling = true
|
||
}
|
||
if switchCount >= maxAccountSwitches {
|
||
lastFailoverErr = failoverErr
|
||
h.handleGeminiFailoverExhausted(c, lastFailoverErr)
|
||
return
|
||
}
|
||
lastFailoverErr = failoverErr
|
||
switchCount++
|
||
log.Printf("Gemini account %d: upstream error %d, switching account %d/%d", account.ID, failoverErr.StatusCode, switchCount, maxAccountSwitches)
|
||
if account.Platform == service.PlatformAntigravity {
|
||
if !sleepFailoverDelay(c.Request.Context(), switchCount) {
|
||
return
|
||
}
|
||
}
|
||
continue
|
||
}
|
||
// ForwardNative already wrote the response
|
||
log.Printf("Gemini native forward failed: %v", err)
|
||
return
|
||
}
|
||
|
||
// 捕获请求信息（用于异步记录，避免在 goroutine 中访问 gin.Context）
|
||
userAgent := c.GetHeader("User-Agent")
|
||
clientIP := ip.GetClientIP(c)
|
||
|
||
// 保存 Gemini 内容摘要会话（用于 Fallback 匹配）
|
||
if useDigestFallback && geminiDigestChain != "" && geminiPrefixHash != "" {
|
||
if err := h.gatewayService.SaveGeminiSession(
|
||
c.Request.Context(),
|
||
derefGroupID(apiKey.GroupID),
|
||
geminiPrefixHash,
|
||
geminiDigestChain,
|
||
geminiSessionUUID,
|
||
account.ID,
|
||
); err != nil {
|
||
log.Printf("[Gemini] Failed to save digest session: %v", err)
|
||
}
|
||
}
|
||
|
||
// 6) record usage async (Gemini 使用长上下文双倍计费)
|
||
go func(result *service.ForwardResult, usedAccount *service.Account, ua, ip string, fcb bool) {
|
||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||
defer cancel()
|
||
|
||
if err := h.gatewayService.RecordUsageWithLongContext(ctx, &service.RecordUsageLongContextInput{
|
||
Result: result,
|
||
APIKey: apiKey,
|
||
User: apiKey.User,
|
||
Account: usedAccount,
|
||
Subscription: subscription,
|
||
UserAgent: ua,
|
||
IPAddress: ip,
|
||
LongContextThreshold: 200000, // Gemini 200K 阈值
|
||
LongContextMultiplier: 2.0, // 超出部分双倍计费
|
||
ForceCacheBilling: fcb,
|
||
APIKeyService: h.apiKeyService,
|
||
}); err != nil {
|
||
log.Printf("Record usage failed: %v", err)
|
||
}
|
||
}(result, account, userAgent, clientIP, forceCacheBilling)
|
||
return
|
||
}
|
||
}
|
||
|
||
func parseGeminiModelAction(rest string) (model string, action string, err error) {
|
||
rest = strings.TrimSpace(rest)
|
||
if rest == "" {
|
||
return "", "", &pathParseError{"missing path"}
|
||
}
|
||
|
||
// Standard: {model}:{action}
|
||
if i := strings.Index(rest, ":"); i > 0 && i < len(rest)-1 {
|
||
return rest[:i], rest[i+1:], nil
|
||
}
|
||
|
||
// Fallback: {model}/{action}
|
||
if i := strings.Index(rest, "/"); i > 0 && i < len(rest)-1 {
|
||
return rest[:i], rest[i+1:], nil
|
||
}
|
||
|
||
return "", "", &pathParseError{"invalid model action path"}
|
||
}
|
||
|
||
func (h *GatewayHandler) handleGeminiFailoverExhausted(c *gin.Context, failoverErr *service.UpstreamFailoverError) {
|
||
if failoverErr == nil {
|
||
googleError(c, http.StatusBadGateway, "Upstream request failed")
|
||
return
|
||
}
|
||
|
||
statusCode := failoverErr.StatusCode
|
||
responseBody := failoverErr.ResponseBody
|
||
|
||
// 先检查透传规则
|
||
if h.errorPassthroughService != nil && len(responseBody) > 0 {
|
||
if rule := h.errorPassthroughService.MatchRule(service.PlatformGemini, statusCode, responseBody); rule != nil {
|
||
// 确定响应状态码
|
||
respCode := statusCode
|
||
if !rule.PassthroughCode && rule.ResponseCode != nil {
|
||
respCode = *rule.ResponseCode
|
||
}
|
||
|
||
// 确定响应消息
|
||
msg := service.ExtractUpstreamErrorMessage(responseBody)
|
||
if !rule.PassthroughBody && rule.CustomMessage != nil {
|
||
msg = *rule.CustomMessage
|
||
}
|
||
|
||
googleError(c, respCode, msg)
|
||
return
|
||
}
|
||
}
|
||
|
||
// 使用默认的错误映射
|
||
status, message := mapGeminiUpstreamError(statusCode)
|
||
googleError(c, status, message)
|
||
}
|
||
|
||
func mapGeminiUpstreamError(statusCode int) (int, string) {
|
||
switch statusCode {
|
||
case 401:
|
||
return http.StatusBadGateway, "Upstream authentication failed, please contact administrator"
|
||
case 403:
|
||
return http.StatusBadGateway, "Upstream access forbidden, please contact administrator"
|
||
case 429:
|
||
return http.StatusTooManyRequests, "Upstream rate limit exceeded, please retry later"
|
||
case 529:
|
||
return http.StatusServiceUnavailable, "Upstream service overloaded, please retry later"
|
||
case 500, 502, 503, 504:
|
||
return http.StatusBadGateway, "Upstream service temporarily unavailable"
|
||
default:
|
||
return http.StatusBadGateway, "Upstream request failed"
|
||
}
|
||
}
|
||
|
||
type pathParseError struct{ msg string }
|
||
|
||
func (e *pathParseError) Error() string { return e.msg }
|
||
|
||
func googleError(c *gin.Context, status int, message string) {
|
||
c.JSON(status, gin.H{
|
||
"error": gin.H{
|
||
"code": status,
|
||
"message": message,
|
||
"status": googleapi.HTTPStatusToGoogleStatus(status),
|
||
},
|
||
})
|
||
}
|
||
|
||
func writeUpstreamResponse(c *gin.Context, res *service.UpstreamHTTPResult) {
|
||
if res == nil {
|
||
googleError(c, http.StatusBadGateway, "Empty upstream response")
|
||
return
|
||
}
|
||
for k, vv := range res.Headers {
|
||
// Avoid overriding content-length and hop-by-hop headers.
|
||
if strings.EqualFold(k, "Content-Length") || strings.EqualFold(k, "Transfer-Encoding") || strings.EqualFold(k, "Connection") {
|
||
continue
|
||
}
|
||
for _, v := range vv {
|
||
c.Writer.Header().Add(k, v)
|
||
}
|
||
}
|
||
contentType := res.Headers.Get("Content-Type")
|
||
if contentType == "" {
|
||
contentType = "application/json"
|
||
}
|
||
c.Data(res.StatusCode, contentType, res.Body)
|
||
}
|
||
|
||
func shouldFallbackGeminiModels(res *service.UpstreamHTTPResult) bool {
|
||
if res == nil {
|
||
return true
|
||
}
|
||
if res.StatusCode != http.StatusUnauthorized && res.StatusCode != http.StatusForbidden {
|
||
return false
|
||
}
|
||
if strings.Contains(strings.ToLower(res.Headers.Get("Www-Authenticate")), "insufficient_scope") {
|
||
return true
|
||
}
|
||
if strings.Contains(strings.ToLower(string(res.Body)), "insufficient authentication scopes") {
|
||
return true
|
||
}
|
||
if strings.Contains(strings.ToLower(string(res.Body)), "access_token_scope_insufficient") {
|
||
return true
|
||
}
|
||
return false
|
||
}
|
||
|
||
// extractGeminiCLISessionHash 从 Gemini CLI 请求中提取会话标识。
|
||
// 组合 x-gemini-api-privileged-user-id header 和请求体中的 tmp 目录哈希。
|
||
//
|
||
// 会话标识生成策略：
|
||
// 1. 从请求体中提取 tmp 目录哈希（64位十六进制）
|
||
// 2. 从 header 中提取 privileged-user-id（UUID）
|
||
// 3. 组合两者生成 SHA256 哈希作为最终的会话标识
|
||
//
|
||
// 如果找不到 tmp 目录哈希，返回空字符串（不使用粘性会话）。
|
||
//
|
||
// extractGeminiCLISessionHash extracts session identifier from Gemini CLI requests.
|
||
// Combines x-gemini-api-privileged-user-id header with tmp directory hash from request body.
|
||
func extractGeminiCLISessionHash(c *gin.Context, body []byte) string {
|
||
// 1. 从请求体中提取 tmp 目录哈希
|
||
match := geminiCLITmpDirRegex.FindSubmatch(body)
|
||
if len(match) < 2 {
|
||
return "" // 没有找到 tmp 目录，不使用粘性会话
|
||
}
|
||
tmpDirHash := string(match[1])
|
||
|
||
// 2. 提取 privileged-user-id
|
||
privilegedUserID := strings.TrimSpace(c.GetHeader("x-gemini-api-privileged-user-id"))
|
||
|
||
// 3. 组合生成最终的 session hash
|
||
if privilegedUserID != "" {
|
||
// 组合两个标识符：privileged-user-id + tmp 目录哈希
|
||
combined := privilegedUserID + ":" + tmpDirHash
|
||
hash := sha256.Sum256([]byte(combined))
|
||
return hex.EncodeToString(hash[:])
|
||
}
|
||
|
||
// 如果没有 privileged-user-id，直接使用 tmp 目录哈希
|
||
return tmpDirHash
|
||
}
|
||
|
||
// truncateDigestChain 截断摘要链用于日志显示
|
||
func truncateDigestChain(chain string) string {
|
||
if len(chain) <= 50 {
|
||
return chain
|
||
}
|
||
return chain[:50] + "..."
|
||
}
|
||
|
||
// safeShortPrefix 返回字符串前 n 个字符；长度不足时返回原字符串。
|
||
// 用于日志展示，避免切片越界。
|
||
func safeShortPrefix(value string, n int) string {
|
||
if n <= 0 || len(value) <= n {
|
||
return value
|
||
}
|
||
return value[:n]
|
||
}
|
||
|
||
// derefGroupID 安全解引用 *int64，nil 返回 0
|
||
func derefGroupID(groupID *int64) int64 {
|
||
if groupID == nil {
|
||
return 0
|
||
}
|
||
return *groupID
|
||
}
|