gemini_v1beta_handler.go 32 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889
  1. package handler
  2. import (
  3. "bytes"
  4. "context"
  5. "crypto/sha256"
  6. "encoding/hex"
  7. "encoding/json"
  8. "errors"
  9. "net/http"
  10. "regexp"
  11. "strconv"
  12. "strings"
  13. "github.com/Wei-Shaw/sub2api/internal/domain"
  14. "github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
  15. "github.com/Wei-Shaw/sub2api/internal/pkg/gemini"
  16. "github.com/Wei-Shaw/sub2api/internal/pkg/googleapi"
  17. pkghttputil "github.com/Wei-Shaw/sub2api/internal/pkg/httputil"
  18. "github.com/Wei-Shaw/sub2api/internal/pkg/ip"
  19. "github.com/Wei-Shaw/sub2api/internal/pkg/logger"
  20. "github.com/Wei-Shaw/sub2api/internal/server/middleware"
  21. "github.com/Wei-Shaw/sub2api/internal/service"
  22. "github.com/google/uuid"
  23. "github.com/gin-gonic/gin"
  24. "go.uber.org/zap"
  25. )
  26. // geminiCLITmpDirRegex 用于从 Gemini CLI 请求体中提取 tmp 目录的哈希值
  27. // 匹配格式: /Users/xxx/.gemini/tmp/[64位十六进制哈希]
  28. var geminiCLITmpDirRegex = regexp.MustCompile(`/\.gemini/tmp/([A-Fa-f0-9]{64})`)
  29. // GeminiV1BetaListModels proxies:
  30. // GET /v1beta/models
  31. func (h *GatewayHandler) GeminiV1BetaListModels(c *gin.Context) {
  32. apiKey, ok := middleware.GetAPIKeyFromContext(c)
  33. if !ok || apiKey == nil {
  34. googleError(c, http.StatusUnauthorized, "Invalid API key")
  35. return
  36. }
  37. // 检查平台:优先使用强制平台(/antigravity 路由),否则要求 gemini 分组
  38. forcePlatform, hasForcePlatform := middleware.GetForcePlatformFromContext(c)
  39. if !hasForcePlatform && effectiveAPIKeyPlatform(c, apiKey) != service.PlatformGemini {
  40. googleError(c, http.StatusBadRequest, "API key group platform is not gemini")
  41. return
  42. }
  43. // 分组级模型白名单开启时过滤 models[].name(名字形如 models/xxx)。
  44. filterGeminiModels := func(models []gemini.Model) []gemini.Model {
  45. if apiKey.Group == nil || !apiKey.Group.ModelAllowlistEnabled() {
  46. return models
  47. }
  48. filtered := make([]gemini.Model, 0, len(models))
  49. for _, model := range models {
  50. if apiKey.Group.ModelAllowlist.Allows(model.Name) {
  51. filtered = append(filtered, model)
  52. }
  53. }
  54. return filtered
  55. }
  56. // 强制 antigravity 模式:返回 antigravity 支持的模型列表
  57. if forcePlatform == service.PlatformAntigravity {
  58. if apiKey.Group != nil && apiKey.Group.ModelAllowlistEnabled() {
  59. agModels := antigravity.DefaultGeminiModels()
  60. filtered := make([]antigravity.GeminiModel, 0, len(agModels))
  61. for _, model := range agModels {
  62. if apiKey.Group.ModelAllowlist.Allows(model.Name) {
  63. filtered = append(filtered, model)
  64. }
  65. }
  66. c.JSON(http.StatusOK, antigravity.GeminiModelsListResponse{Models: filtered})
  67. return
  68. }
  69. c.JSON(http.StatusOK, antigravity.FallbackGeminiModelsList())
  70. return
  71. }
  72. account, err := h.geminiCompatService.SelectAccountForAIStudioEndpoints(c.Request.Context(), apiKey.GroupID)
  73. if err != nil {
  74. // 没有 gemini 账户,检查是否有 antigravity 账户可用
  75. hasAntigravity, _ := h.geminiCompatService.HasAntigravityAccounts(c.Request.Context(), apiKey.GroupID)
  76. if hasAntigravity {
  77. // antigravity 账户使用静态模型列表
  78. c.JSON(http.StatusOK, gemini.ModelsListResponse{Models: filterGeminiModels(gemini.DefaultModels())})
  79. return
  80. }
  81. markOpsRoutingCapacityLimitedIfNoAvailable(c, err)
  82. googleError(c, http.StatusServiceUnavailable, "No available Gemini accounts: "+err.Error())
  83. return
  84. }
  85. res, err := h.geminiCompatService.ForwardAIStudioGET(c.Request.Context(), account, "/v1beta/models")
  86. if err != nil {
  87. googleError(c, http.StatusBadGateway, err.Error())
  88. return
  89. }
  90. if shouldFallbackGeminiModels(res) {
  91. c.JSON(http.StatusOK, gemini.ModelsListResponse{Models: filterGeminiModels(gemini.DefaultModels())})
  92. return
  93. }
  94. if apiKey.Group != nil && apiKey.Group.ModelAllowlistEnabled() {
  95. if filtered, dropped, ok := filterUpstreamGeminiModelsBody(res.Body, apiKey.Group.ModelAllowlist); ok && dropped {
  96. // 只在确有条目被过滤时替换响应体;全命中或解析失败时保持原始响应,
  97. // 统一经 writeUpstreamResponse 写出(保留全部上游响应头)。
  98. res.Body = filtered
  99. }
  100. }
  101. writeUpstreamResponse(c, res)
  102. }
  103. // filterUpstreamGeminiModelsBody 按白名单过滤上游 /v1beta/models 响应中的
  104. // models[].name,其余信封字段(如 nextPageToken)原样保留。
  105. // 返回值:filtered 为过滤后的响应体;dropped 表示是否有条目被移除(全命中时
  106. // 为 false,调用方应保持原始响应以完整透传上游头);ok=false 表示解析失败,
  107. // 调用方同样应透传原始响应。
  108. func filterUpstreamGeminiModelsBody(body []byte, allowlist service.GroupModelAllowlist) (filtered []byte, dropped bool, ok bool) {
  109. var envelope map[string]json.RawMessage
  110. if err := json.Unmarshal(body, &envelope); err != nil {
  111. return nil, false, false
  112. }
  113. rawModels, hasModels := envelope["models"]
  114. if !hasModels {
  115. return body, false, true
  116. }
  117. type geminiModelName struct {
  118. Name string `json:"name"`
  119. }
  120. var models []json.RawMessage
  121. if err := json.Unmarshal(rawModels, &models); err != nil {
  122. return nil, false, false
  123. }
  124. kept := make([]json.RawMessage, 0, len(models))
  125. for _, raw := range models {
  126. var model geminiModelName
  127. if err := json.Unmarshal(raw, &model); err != nil {
  128. return nil, false, false
  129. }
  130. if allowlist.Allows(model.Name) {
  131. kept = append(kept, raw)
  132. }
  133. }
  134. if len(kept) == len(models) {
  135. // 全部命中时直接透传原始响应体。
  136. return body, false, true
  137. }
  138. mergedModels, err := json.Marshal(kept)
  139. if err != nil {
  140. return nil, false, false
  141. }
  142. envelope["models"] = mergedModels
  143. merged, err := json.Marshal(envelope)
  144. if err != nil {
  145. return nil, false, false
  146. }
  147. return merged, true, true
  148. }
  149. // GeminiV1BetaGetModel proxies:
  150. // GET /v1beta/models/{model}
  151. func (h *GatewayHandler) GeminiV1BetaGetModel(c *gin.Context) {
  152. apiKey, ok := middleware.GetAPIKeyFromContext(c)
  153. if !ok || apiKey == nil {
  154. googleError(c, http.StatusUnauthorized, "Invalid API key")
  155. return
  156. }
  157. // 检查平台:优先使用强制平台(/antigravity 路由),否则要求 gemini 分组
  158. forcePlatform, hasForcePlatform := middleware.GetForcePlatformFromContext(c)
  159. if !hasForcePlatform && effectiveAPIKeyPlatform(c, apiKey) != service.PlatformGemini {
  160. googleError(c, http.StatusBadRequest, "API key group platform is not gemini")
  161. return
  162. }
  163. modelName := strings.TrimSpace(c.Param("model"))
  164. if modelName == "" {
  165. googleError(c, http.StatusBadRequest, "Missing model in URL")
  166. return
  167. }
  168. // 模型名会被拼进上游 URL 的 path,先在入口校验片段合规性,
  169. // 见 service/upstream_path_guard.go。
  170. if !service.IsSafeGeminiModelPathSegment(modelName) {
  171. googleError(c, http.StatusBadRequest, "Invalid model in URL")
  172. return
  173. }
  174. if resolvedModel, ok := service.ResolvedUpstreamModelFromContext(c.Request.Context()); ok && strings.TrimSpace(resolvedModel) != "" {
  175. modelName = strings.TrimSpace(resolvedModel)
  176. }
  177. // 强制 antigravity 模式:返回 antigravity 模型信息
  178. if forcePlatform == service.PlatformAntigravity {
  179. c.JSON(http.StatusOK, antigravity.FallbackGeminiModel(modelName))
  180. return
  181. }
  182. account, err := h.geminiCompatService.SelectAccountForAIStudioEndpoints(c.Request.Context(), apiKey.GroupID)
  183. if err != nil {
  184. // 没有 gemini 账户,检查是否有 antigravity 账户可用
  185. hasAntigravity, _ := h.geminiCompatService.HasAntigravityAccounts(c.Request.Context(), apiKey.GroupID)
  186. if hasAntigravity {
  187. // antigravity 账户使用静态模型信息
  188. c.JSON(http.StatusOK, gemini.FallbackModel(modelName))
  189. return
  190. }
  191. markOpsRoutingCapacityLimitedIfNoAvailable(c, err)
  192. googleError(c, http.StatusServiceUnavailable, "No available Gemini accounts: "+err.Error())
  193. return
  194. }
  195. res, err := h.geminiCompatService.ForwardAIStudioGET(c.Request.Context(), account, "/v1beta/models/"+modelName)
  196. if err != nil {
  197. googleError(c, http.StatusBadGateway, err.Error())
  198. return
  199. }
  200. if shouldFallbackGeminiModel(modelName, res) {
  201. c.JSON(http.StatusOK, gemini.FallbackModel(modelName))
  202. return
  203. }
  204. writeUpstreamResponse(c, res)
  205. }
  206. // GeminiV1BetaModels proxies Gemini native REST endpoints like:
  207. // POST /v1beta/models/{model}:generateContent
  208. // POST /v1beta/models/{model}:streamGenerateContent?alt=sse
  209. func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
  210. apiKey, ok := middleware.GetAPIKeyFromContext(c)
  211. if !ok || apiKey == nil {
  212. googleError(c, http.StatusUnauthorized, "Invalid API key")
  213. return
  214. }
  215. authSubject, ok := middleware.GetAuthSubjectFromContext(c)
  216. if !ok {
  217. googleError(c, http.StatusInternalServerError, "User context not found")
  218. return
  219. }
  220. reqLog := requestLogger(
  221. c,
  222. "handler.gemini_v1beta.models",
  223. zap.Int64("user_id", authSubject.UserID),
  224. zap.Int64("api_key_id", apiKey.ID),
  225. zap.Any("group_id", apiKey.GroupID),
  226. )
  227. // 检查平台:优先使用强制平台(/antigravity 路由,中间件已设置 request.Context),否则要求 gemini 分组
  228. if !middleware.HasForcePlatform(c) {
  229. if effectiveAPIKeyPlatform(c, apiKey) != service.PlatformGemini {
  230. googleError(c, http.StatusBadRequest, "API key group platform is not gemini")
  231. return
  232. }
  233. }
  234. modelName, action, err := parseGeminiModelAction(strings.TrimPrefix(c.Param("modelAction"), "/"))
  235. if err != nil {
  236. googleError(c, http.StatusNotFound, err.Error())
  237. return
  238. }
  239. // URL 里的模型名最终会被拼进上游 /v1beta/models/{model}:{action},
  240. // 先在入口校验片段合规性,见 service/upstream_path_guard.go。
  241. if !service.IsSafeGeminiModelPathSegment(modelName) {
  242. googleError(c, http.StatusBadRequest, "Invalid model in URL")
  243. return
  244. }
  245. if resolvedModel, ok := service.ResolvedUpstreamModelFromContext(c.Request.Context()); ok && strings.TrimSpace(resolvedModel) != "" {
  246. modelName = strings.TrimSpace(resolvedModel)
  247. }
  248. stream := action == "streamGenerateContent"
  249. reqLog = reqLog.With(zap.String("model", modelName), zap.String("action", action), zap.Bool("stream", stream))
  250. body, err := pkghttputil.ReadRequestBodyWithPrealloc(c.Request)
  251. if err != nil {
  252. if maxErr, ok := extractMaxBytesError(err); ok {
  253. googleError(c, http.StatusRequestEntityTooLarge, buildBodyTooLargeMessage(maxErr.Limit))
  254. return
  255. }
  256. googleError(c, http.StatusBadRequest, "Failed to read request body")
  257. return
  258. }
  259. if len(body) == 0 {
  260. googleError(c, http.StatusBadRequest, "Request body is empty")
  261. return
  262. }
  263. setOpsRequestContext(c, modelName, stream)
  264. setOpsEndpointContext(c, "", int16(service.RequestTypeFromLegacy(stream, false)))
  265. pricingCtx, pricingAt := service.WithGatewayTokenRequestPricing(c.Request.Context())
  266. c.Request = c.Request.WithContext(pricingCtx)
  267. if decision := h.checkSecurityAudit(c, reqLog, apiKey, authSubject, service.ContentModerationProtocolGemini, modelName, body); decision != nil && !decision.AllowNextStage {
  268. googleSecurityAuditError(c, decision)
  269. return
  270. }
  271. // 解析渠道级模型映射
  272. channelMapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(c.Request.Context(), apiKey.GroupID, modelName)
  273. reqModel := modelName // 保存映射前的原始模型名
  274. if channelMapping.Mapped {
  275. modelName = channelMapping.MappedModel
  276. }
  277. // Get subscription (may be nil)
  278. subscription, _ := middleware.GetSubscriptionFromContext(c)
  279. // For Gemini native API, do not send Claude-style ping frames.
  280. geminiConcurrency := NewConcurrencyHelper(h.concurrencyHelper.concurrencyService, SSEPingFormatNone, 0)
  281. // 1) user concurrency slot
  282. streamStarted := false
  283. if h.errorPassthroughService != nil {
  284. service.BindErrorPassthroughService(c, h.errorPassthroughService)
  285. }
  286. userReleaseFunc, err := geminiConcurrency.AcquireUserSlotWithWait(c, authSubject.UserID, authSubject.Concurrency, stream, &streamStarted)
  287. if err != nil {
  288. reqLog.Warn("gemini.user_slot_acquire_failed", zap.Error(err))
  289. googleError(c, http.StatusTooManyRequests, err.Error())
  290. return
  291. }
  292. // 确保请求取消时也会释放槽位,避免长连接被动中断造成泄漏
  293. userReleaseFunc = wrapReleaseOnDone(c.Request.Context(), userReleaseFunc)
  294. if userReleaseFunc != nil {
  295. defer userReleaseFunc()
  296. }
  297. // 2) billing eligibility check (after wait)
  298. if err := h.billingCacheService.CheckBillingEligibility(c.Request.Context(), apiKey.User, apiKey, apiKey.Group, subscription, service.QuotaPlatform(c.Request.Context(), apiKey)); err != nil {
  299. reqLog.Info("gemini.billing_eligibility_check_failed", zap.Error(err))
  300. status, _, message, retryAfter := billingErrorDetails(err)
  301. if retryAfter > 0 {
  302. c.Header("Retry-After", strconv.Itoa(retryAfter))
  303. }
  304. googleError(c, status, message)
  305. return
  306. }
  307. // 3) select account (sticky session based on request body)
  308. // 优先使用 Gemini CLI 的会话标识(privileged-user-id + tmp 目录哈希)
  309. sessionHash := extractGeminiCLISessionHash(c, body)
  310. if sessionHash == "" {
  311. // Fallback: 使用通用的会话哈希生成逻辑(适用于其他客户端)
  312. parsedReq, _ := service.ParseGatewayRequest(service.NewRequestBodyRef(body), domain.PlatformGemini)
  313. if parsedReq != nil {
  314. parsedReq.SessionContext = &service.SessionContext{
  315. ClientIP: ip.GetClientIP(c),
  316. UserAgent: c.GetHeader("User-Agent"),
  317. APIKeyID: apiKey.ID,
  318. }
  319. }
  320. sessionHash = h.gatewayService.GenerateSessionHash(parsedReq)
  321. }
  322. sessionKey := sessionHash
  323. if sessionHash != "" {
  324. sessionKey = "gemini:" + sessionHash
  325. }
  326. // 查询粘性会话绑定的账号 ID(用于检测账号切换)
  327. var sessionBoundAccountID int64
  328. if sessionKey != "" {
  329. sessionBoundAccountID, _ = h.gatewayService.GetCachedSessionAccountID(c.Request.Context(), apiKey.GroupID, sessionKey)
  330. if sessionBoundAccountID > 0 {
  331. prefetchedGroupID := int64(0)
  332. if apiKey.GroupID != nil {
  333. prefetchedGroupID = *apiKey.GroupID
  334. }
  335. ctx := service.WithPrefetchedStickySession(c.Request.Context(), sessionBoundAccountID, prefetchedGroupID, h.metadataBridgeEnabled())
  336. c.Request = c.Request.WithContext(ctx)
  337. }
  338. }
  339. // === Gemini 内容摘要会话 Fallback 逻辑 ===
  340. // 当原有会话标识无效时(sessionBoundAccountID == 0),尝试基于内容摘要链匹配
  341. var geminiDigestChain string
  342. var geminiPrefixHash string
  343. var geminiSessionUUID string
  344. var matchedDigestChain string
  345. useDigestFallback := sessionBoundAccountID == 0
  346. if useDigestFallback {
  347. // 解析 Gemini 请求体
  348. var geminiReq antigravity.GeminiRequest
  349. if err := json.Unmarshal(body, &geminiReq); err == nil && len(geminiReq.Contents) > 0 {
  350. // 生成摘要链
  351. geminiDigestChain = service.BuildGeminiDigestChain(&geminiReq)
  352. if geminiDigestChain != "" {
  353. // 生成前缀 hash
  354. userAgent := c.GetHeader("User-Agent")
  355. clientIP := ip.GetClientIP(c)
  356. platform := ""
  357. if apiKey.Group != nil {
  358. platform = apiKey.Group.Platform
  359. }
  360. geminiPrefixHash = service.GenerateGeminiPrefixHash(
  361. authSubject.UserID,
  362. apiKey.ID,
  363. clientIP,
  364. userAgent,
  365. platform,
  366. modelName,
  367. )
  368. // 查找会话
  369. foundUUID, foundAccountID, foundMatchedChain, found := h.gatewayService.FindGeminiSession(
  370. c.Request.Context(),
  371. derefGroupID(apiKey.GroupID),
  372. geminiPrefixHash,
  373. geminiDigestChain,
  374. )
  375. if found {
  376. matchedDigestChain = foundMatchedChain
  377. sessionBoundAccountID = foundAccountID
  378. geminiSessionUUID = foundUUID
  379. reqLog.Info("gemini.digest_fallback_matched",
  380. zap.String("session_uuid_prefix", safeShortPrefix(foundUUID, 8)),
  381. zap.Int64("account_id", foundAccountID),
  382. zap.String("digest_chain", truncateDigestChain(geminiDigestChain)),
  383. )
  384. // 关键:如果原 sessionKey 为空,使用 prefixHash + uuid 作为 sessionKey
  385. // 这样 SelectAccountWithLoadAwareness 的粘性会话逻辑会优先使用匹配到的账号
  386. if sessionKey == "" {
  387. sessionKey = service.GenerateGeminiDigestSessionKey(geminiPrefixHash, foundUUID)
  388. }
  389. _ = h.gatewayService.BindStickySession(c.Request.Context(), apiKey.GroupID, sessionKey, foundAccountID)
  390. } else {
  391. // 生成新的会话 UUID
  392. geminiSessionUUID = uuid.New().String()
  393. // 为新会话也生成 sessionKey(用于后续请求的粘性会话)
  394. if sessionKey == "" {
  395. sessionKey = service.GenerateGeminiDigestSessionKey(geminiPrefixHash, geminiSessionUUID)
  396. }
  397. }
  398. }
  399. }
  400. }
  401. // 判断是否真的绑定了粘性会话:有 sessionKey 且已经绑定到某个账号
  402. hasBoundSession := sessionKey != "" && sessionBoundAccountID > 0
  403. cleanedForUnknownBinding := false
  404. fs := NewFailoverState(h.maxAccountSwitchesGemini, hasBoundSession)
  405. // 单账号分组提前设置 SingleAccountRetry 标记,让 Service 层首次 503 就不设模型限流标记。
  406. // 避免单账号分组收到 503 (MODEL_CAPACITY_EXHAUSTED) 时设 29s 限流,导致后续请求连续快速失败。
  407. if h.gatewayService.IsSingleAntigravityAccountGroup(c.Request.Context(), apiKey.GroupID) {
  408. ctx := service.WithSingleAccountRetry(c.Request.Context(), true, h.metadataBridgeEnabled())
  409. c.Request = c.Request.WithContext(ctx)
  410. }
  411. for {
  412. selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, sessionKey, modelName, fs.FailedAccountIDs, "", int64(0)) // Gemini 不使用会话限制
  413. if err != nil {
  414. if len(fs.FailedAccountIDs) == 0 {
  415. cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, modelName, modelName, service.PlatformGemini)
  416. if !cls.ModelNotFound {
  417. markOpsRoutingCapacityLimitedIfNoAvailable(c, err)
  418. }
  419. message := cls.Message
  420. if !cls.ModelNotFound {
  421. message = "No available Gemini accounts: " + err.Error()
  422. }
  423. googleError(c, cls.Status, message)
  424. return
  425. }
  426. action := fs.HandleSelectionExhausted(c.Request.Context())
  427. switch action {
  428. case FailoverContinue:
  429. ctx := service.WithSingleAccountRetry(c.Request.Context(), true, h.metadataBridgeEnabled())
  430. c.Request = c.Request.WithContext(ctx)
  431. continue
  432. case FailoverCanceled:
  433. failoverClientGone(c)
  434. return
  435. default: // FailoverExhausted
  436. h.handleGeminiFailoverExhausted(c, fs.LastFailoverErr)
  437. return
  438. }
  439. }
  440. account := selection.Account
  441. setOpsSelectedAccount(c, account.ID, account.Platform)
  442. // 检测账号切换:如果粘性会话绑定的账号与当前选择的账号不同,清除 thoughtSignature
  443. // 注意:Gemini 原生 API 的 thoughtSignature 与具体上游账号强相关;跨账号透传会导致 400。
  444. if sessionBoundAccountID > 0 && sessionBoundAccountID != account.ID {
  445. reqLog.Info("gemini.sticky_session_account_switched",
  446. zap.Int64("from_account_id", sessionBoundAccountID),
  447. zap.Int64("to_account_id", account.ID),
  448. zap.Bool("clean_thought_signature", true),
  449. )
  450. body = service.CleanGeminiNativeThoughtSignatures(body)
  451. sessionBoundAccountID = account.ID
  452. } else if sessionKey != "" && sessionBoundAccountID == 0 && !cleanedForUnknownBinding && bytes.Contains(body, []byte(`"thoughtSignature"`)) {
  453. // 无缓存绑定但请求里已有 thoughtSignature:常见于缓存丢失/TTL 过期后,客户端继续携带旧签名。
  454. // 为避免第一次转发就 400,这里做一次确定性清理,让新账号重新生成签名链路。
  455. reqLog.Info("gemini.sticky_session_binding_missing",
  456. zap.Bool("clean_thought_signature", true),
  457. )
  458. body = service.CleanGeminiNativeThoughtSignatures(body)
  459. cleanedForUnknownBinding = true
  460. sessionBoundAccountID = account.ID
  461. } else if sessionBoundAccountID == 0 {
  462. // 记录本次请求中首次选择到的账号,便于同一请求内 failover 时检测切换。
  463. sessionBoundAccountID = account.ID
  464. }
  465. // 4) account concurrency slot
  466. accountReleaseFunc := selection.ReleaseFunc
  467. if !selection.Acquired {
  468. if selection.WaitPlan == nil {
  469. markOpsRoutingCapacityLimited(c)
  470. googleError(c, http.StatusServiceUnavailable, "No available Gemini accounts")
  471. return
  472. }
  473. accountWaitCounted := false
  474. canWait, err := geminiConcurrency.IncrementAccountWaitCount(c.Request.Context(), account.ID, selection.WaitPlan.MaxWaiting)
  475. if err != nil {
  476. reqLog.Warn("gemini.account_wait_counter_increment_failed", zap.Int64("account_id", account.ID), zap.Error(err))
  477. } else if !canWait {
  478. reqLog.Info("gemini.account_wait_queue_full",
  479. zap.Int64("account_id", account.ID),
  480. zap.Int("max_waiting", selection.WaitPlan.MaxWaiting),
  481. )
  482. googleError(c, http.StatusTooManyRequests, "Too many pending requests, please retry later")
  483. return
  484. }
  485. if err == nil && canWait {
  486. accountWaitCounted = true
  487. }
  488. defer func() {
  489. if accountWaitCounted {
  490. geminiConcurrency.DecrementAccountWaitCount(c.Request.Context(), account.ID)
  491. }
  492. }()
  493. accountReleaseFunc, err = geminiConcurrency.AcquireAccountSlotWithWaitTimeout(
  494. c,
  495. account.ID,
  496. selection.WaitPlan.MaxConcurrency,
  497. selection.WaitPlan.Timeout,
  498. stream,
  499. &streamStarted,
  500. )
  501. if err != nil {
  502. reqLog.Warn("gemini.account_slot_acquire_failed", zap.Int64("account_id", account.ID), zap.Error(err))
  503. googleError(c, http.StatusTooManyRequests, err.Error())
  504. return
  505. }
  506. if accountWaitCounted {
  507. geminiConcurrency.DecrementAccountWaitCount(c.Request.Context(), account.ID)
  508. accountWaitCounted = false
  509. }
  510. }
  511. // 终检与准入后绑定使用选号结果携带的门(见 responses 同名注释)。
  512. admissionCtx := service.ContextWithSelectionProfitGate(c.Request.Context(), selection)
  513. latest, vetoed, reason := h.gatewayService.GatewayProfitControlVetoLatest(admissionCtx, account)
  514. if vetoed {
  515. if accountReleaseFunc != nil {
  516. accountReleaseFunc()
  517. }
  518. reqLog.Debug("gemini.account_slot_profit_vetoed", zap.Int64("account_id", account.ID), zap.String("reason", reason))
  519. if fs.RecordProfitVeto(account.ID) == FailoverExhausted {
  520. reqLog.Warn("gemini.profit_veto_attempts_exhausted", zap.Int("profit_veto_count", fs.ProfitVetoCount()))
  521. markOpsRoutingCapacityLimited(c)
  522. googleError(c, http.StatusServiceUnavailable, profitVetoExhaustedMessage)
  523. return
  524. }
  525. continue
  526. }
  527. account = latest
  528. selection.Account = latest
  529. // 等待路径保持既有 eager 绑定(无门时 helper 直接绑定);调度器已抢槽
  530. // 的直达路径无门时由选号内部绑定,这里只在门下补准入后绑定。
  531. if selection.ProfitGateActive() || !selection.Acquired {
  532. if err := h.gatewayService.BindStickySessionAfterProfitAdmission(admissionCtx, apiKey.GroupID, sessionKey, account.ID); err != nil {
  533. reqLog.Warn("gemini.bind_sticky_session_after_profit_admission_failed", zap.Int64("account_id", account.ID), zap.Error(err))
  534. }
  535. }
  536. // 账号槽位/等待计数需要在超时或断开时安全回收
  537. accountReleaseFunc = wrapReleaseOnDone(c.Request.Context(), accountReleaseFunc)
  538. // 5) forward (根据平台分流)
  539. var result *service.ForwardResult
  540. requestCtx := c.Request.Context()
  541. if fs.SwitchCount > 0 {
  542. requestCtx = service.WithAccountSwitchCount(requestCtx, fs.SwitchCount, h.metadataBridgeEnabled())
  543. }
  544. sessionGroupID := derefGroupID(apiKey.GroupID)
  545. if account.Platform == service.PlatformAntigravity && account.Type != service.AccountTypeAPIKey {
  546. result, err = h.antigravityGatewayService.ForwardGemini(
  547. requestCtx,
  548. c,
  549. account,
  550. modelName,
  551. action,
  552. stream,
  553. body,
  554. hasBoundSession,
  555. service.WithForwardGeminiSession(sessionGroupID, sessionKey),
  556. )
  557. } else {
  558. result, err = h.geminiCompatService.ForwardNative(requestCtx, c, account, modelName, action, stream, body)
  559. }
  560. if accountReleaseFunc != nil {
  561. accountReleaseFunc()
  562. }
  563. if err != nil {
  564. var failoverErr *service.UpstreamFailoverError
  565. if errors.As(err, &failoverErr) {
  566. failoverAction := fs.HandleFailoverError(c.Request.Context(), h.gatewayService, account.ID, account.Platform, account.GetPoolModeRetryCount(), failoverErr)
  567. switch failoverAction {
  568. case FailoverContinue:
  569. continue
  570. case FailoverExhausted:
  571. h.handleGeminiFailoverExhausted(c, fs.LastFailoverErr)
  572. return
  573. case FailoverCanceled:
  574. failoverClientGone(c)
  575. return
  576. }
  577. }
  578. // ForwardNative already wrote the response
  579. reqLog.Error("gemini.forward_failed", zap.Int64("account_id", account.ID), zap.Error(err))
  580. return
  581. }
  582. // 捕获请求信息(用于异步记录,避免在 goroutine 中访问 gin.Context)
  583. userAgent := c.GetHeader("User-Agent")
  584. clientIP := ip.GetClientIP(c)
  585. // 保存 Gemini 内容摘要会话(用于 Fallback 匹配)
  586. if useDigestFallback && geminiDigestChain != "" && geminiPrefixHash != "" {
  587. if err := h.gatewayService.SaveGeminiSession(
  588. c.Request.Context(),
  589. derefGroupID(apiKey.GroupID),
  590. geminiPrefixHash,
  591. geminiDigestChain,
  592. geminiSessionUUID,
  593. account.ID,
  594. matchedDigestChain,
  595. ); err != nil {
  596. reqLog.Warn("gemini.digest_session_save_failed", zap.Int64("account_id", account.ID), zap.Error(err))
  597. }
  598. }
  599. // 使用量记录通过有界 worker 池提交,避免请求热路径创建无界 goroutine。
  600. requestPayloadHash := service.HashUsageRequestPayload(body)
  601. inboundEndpoint := GetInboundEndpoint(c)
  602. upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform)
  603. // ForceCacheBilling 提前拍成标量,避免 worker 闭包保活 failover 状态里的响应体。
  604. forceCacheBilling := fs.ForceCacheBilling
  605. quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
  606. sessionID := service.ExtractClientSessionID(c)
  607. // 长上下文阶梯由目录数据驱动,统一在计费路径内生效,入口无需声明。
  608. h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) {
  609. if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{
  610. Result: result,
  611. QuotaPlatform: quotaPlatform,
  612. APIKey: apiKey,
  613. User: apiKey.User,
  614. Account: account,
  615. Subscription: subscription,
  616. PricingAt: pricingAt,
  617. InboundEndpoint: inboundEndpoint,
  618. UpstreamEndpoint: upstreamEndpoint,
  619. UserAgent: userAgent,
  620. IPAddress: clientIP,
  621. RequestPayloadHash: requestPayloadHash,
  622. ForceCacheBilling: forceCacheBilling,
  623. APIKeyService: h.apiKeyService,
  624. SessionID: sessionID,
  625. ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, result.UpstreamModel),
  626. }); err != nil {
  627. logger.L().With(
  628. zap.String("component", "handler.gemini_v1beta.models"),
  629. zap.Int64("user_id", authSubject.UserID),
  630. zap.Int64("api_key_id", apiKey.ID),
  631. zap.Any("group_id", apiKey.GroupID),
  632. zap.String("model", modelName),
  633. zap.Int64("account_id", account.ID),
  634. ).Error("gemini.record_usage_failed", zap.Error(err))
  635. }
  636. })
  637. reqLog.Debug("gemini.request_completed",
  638. zap.Int64("account_id", account.ID),
  639. zap.Int("switch_count", fs.SwitchCount),
  640. )
  641. return
  642. }
  643. }
  644. func parseGeminiModelAction(rest string) (model string, action string, err error) {
  645. rest = strings.TrimSpace(rest)
  646. if rest == "" {
  647. return "", "", &pathParseError{"missing path"}
  648. }
  649. // Standard: {model}:{action}
  650. if i := strings.Index(rest, ":"); i > 0 && i < len(rest)-1 {
  651. return rest[:i], rest[i+1:], nil
  652. }
  653. // Fallback: {model}/{action}
  654. if i := strings.Index(rest, "/"); i > 0 && i < len(rest)-1 {
  655. return rest[:i], rest[i+1:], nil
  656. }
  657. return "", "", &pathParseError{"invalid model action path"}
  658. }
  659. func (h *GatewayHandler) handleGeminiFailoverExhausted(c *gin.Context, failoverErr *service.UpstreamFailoverError) {
  660. if failoverErr == nil {
  661. googleError(c, http.StatusBadGateway, "Upstream request failed")
  662. return
  663. }
  664. statusCode := failoverErr.StatusCode
  665. responseBody := failoverErr.ResponseBody
  666. // 先检查透传规则
  667. if h.errorPassthroughService != nil && len(responseBody) > 0 {
  668. if rule := h.errorPassthroughService.MatchRule(service.PlatformGemini, statusCode, responseBody); rule != nil {
  669. // 确定响应状态码
  670. respCode := statusCode
  671. if !rule.PassthroughCode && rule.ResponseCode != nil {
  672. respCode = *rule.ResponseCode
  673. }
  674. // 确定响应消息
  675. msg := service.ExtractUpstreamErrorMessage(responseBody)
  676. if !rule.PassthroughBody && rule.CustomMessage != nil {
  677. msg = *rule.CustomMessage
  678. }
  679. if rule.SkipMonitoring {
  680. c.Set(service.OpsSkipPassthroughKey, true)
  681. }
  682. googleError(c, respCode, msg)
  683. return
  684. }
  685. }
  686. // 记录原始上游状态码,以便 ops 错误日志捕获真实的上游错误
  687. upstreamMsg := service.ExtractUpstreamErrorMessage(responseBody)
  688. service.SetOpsUpstreamError(c, statusCode, upstreamMsg, "")
  689. // 使用默认的错误映射
  690. status, message := mapGeminiUpstreamError(statusCode)
  691. googleError(c, status, message)
  692. }
  693. func mapGeminiUpstreamError(statusCode int) (int, string) {
  694. switch statusCode {
  695. case 401:
  696. return http.StatusBadGateway, "Upstream authentication failed, please contact administrator"
  697. case 403:
  698. return http.StatusBadGateway, "Upstream access forbidden, please contact administrator"
  699. case 429:
  700. return http.StatusTooManyRequests, "Upstream rate limit exceeded, please retry later"
  701. case 529:
  702. return http.StatusServiceUnavailable, "Upstream service overloaded, please retry later"
  703. case 500, 502, 503, 504:
  704. return http.StatusBadGateway, "Upstream service temporarily unavailable"
  705. default:
  706. return http.StatusBadGateway, "Upstream request failed"
  707. }
  708. }
  709. type pathParseError struct{ msg string }
  710. func (e *pathParseError) Error() string { return e.msg }
  711. func googleError(c *gin.Context, status int, message string) {
  712. c.JSON(status, gin.H{
  713. "error": gin.H{
  714. "code": status,
  715. "message": message,
  716. "status": googleapi.HTTPStatusToGoogleStatus(status),
  717. },
  718. })
  719. }
  720. func writeUpstreamResponse(c *gin.Context, res *service.UpstreamHTTPResult) {
  721. if res == nil {
  722. googleError(c, http.StatusBadGateway, "Empty upstream response")
  723. return
  724. }
  725. for k, vv := range res.Headers {
  726. // Avoid overriding content-length and hop-by-hop headers.
  727. if strings.EqualFold(k, "Content-Length") || strings.EqualFold(k, "Transfer-Encoding") || strings.EqualFold(k, "Connection") {
  728. continue
  729. }
  730. for _, v := range vv {
  731. c.Writer.Header().Add(k, v)
  732. }
  733. }
  734. contentType := res.Headers.Get("Content-Type")
  735. if contentType == "" {
  736. contentType = "application/json"
  737. }
  738. c.Data(res.StatusCode, contentType, res.Body)
  739. }
  740. func shouldFallbackGeminiModels(res *service.UpstreamHTTPResult) bool {
  741. if res == nil {
  742. return true
  743. }
  744. if res.StatusCode != http.StatusUnauthorized && res.StatusCode != http.StatusForbidden {
  745. return false
  746. }
  747. if strings.Contains(strings.ToLower(res.Headers.Get("Www-Authenticate")), "insufficient_scope") {
  748. return true
  749. }
  750. if strings.Contains(strings.ToLower(string(res.Body)), "insufficient authentication scopes") {
  751. return true
  752. }
  753. if strings.Contains(strings.ToLower(string(res.Body)), "access_token_scope_insufficient") {
  754. return true
  755. }
  756. return false
  757. }
  758. func shouldFallbackGeminiModel(modelName string, res *service.UpstreamHTTPResult) bool {
  759. if shouldFallbackGeminiModels(res) {
  760. return true
  761. }
  762. if res == nil || res.StatusCode != http.StatusNotFound {
  763. return false
  764. }
  765. return gemini.HasFallbackModel(modelName)
  766. }
  767. // extractGeminiCLISessionHash 从 Gemini CLI 请求中提取会话标识。
  768. // 组合 x-gemini-api-privileged-user-id header 和请求体中的 tmp 目录哈希。
  769. //
  770. // 会话标识生成策略:
  771. // 1. 从请求体中提取 tmp 目录哈希(64位十六进制)
  772. // 2. 从 header 中提取 privileged-user-id(UUID)
  773. // 3. 组合两者生成 SHA256 哈希作为最终的会话标识
  774. //
  775. // 如果找不到 tmp 目录哈希,返回空字符串(不使用粘性会话)。
  776. //
  777. // extractGeminiCLISessionHash extracts session identifier from Gemini CLI requests.
  778. // Combines x-gemini-api-privileged-user-id header with tmp directory hash from request body.
  779. func extractGeminiCLISessionHash(c *gin.Context, body []byte) string {
  780. // 1. 从请求体中提取 tmp 目录哈希
  781. match := geminiCLITmpDirRegex.FindSubmatch(body)
  782. if len(match) < 2 {
  783. return "" // 没有找到 tmp 目录,不使用粘性会话
  784. }
  785. tmpDirHash := string(match[1])
  786. // 2. 提取 privileged-user-id
  787. privilegedUserID := strings.TrimSpace(c.GetHeader("x-gemini-api-privileged-user-id"))
  788. // 3. 组合生成最终的 session hash
  789. if privilegedUserID != "" {
  790. // 组合两个标识符:privileged-user-id + tmp 目录哈希
  791. combined := privilegedUserID + ":" + tmpDirHash
  792. hash := sha256.Sum256([]byte(combined))
  793. return hex.EncodeToString(hash[:])
  794. }
  795. // 如果没有 privileged-user-id,直接使用 tmp 目录哈希
  796. return tmpDirHash
  797. }
  798. // truncateDigestChain 截断摘要链用于日志显示
  799. func truncateDigestChain(chain string) string {
  800. if len(chain) <= 50 {
  801. return chain
  802. }
  803. return chain[:50] + "..."
  804. }
  805. // safeShortPrefix 返回字符串前 n 个字符;长度不足时返回原字符串。
  806. // 用于日志展示,避免切片越界。
  807. func safeShortPrefix(value string, n int) string {
  808. if n <= 0 || len(value) <= n {
  809. return value
  810. }
  811. return value[:n]
  812. }
  813. // derefGroupID 安全解引用 *int64,nil 返回 0
  814. func derefGroupID(groupID *int64) int64 {
  815. if groupID == nil {
  816. return 0
  817. }
  818. return *groupID
  819. }