From 3bfd40f5afd820ff7fae32f3f26c75fe3e0dd504 Mon Sep 17 00:00:00 2001 From: feeeei Date: Fri, 9 Oct 2026 17:20:11 +0800 Subject: [PATCH 01/15] =?UTF-8?q?fix(gemini):=20=E5=BC=95=E7=94=A8?= =?UTF-8?q?=E7=BC=93=E5=AD=98=E4=B8=8D=E5=AD=98=E5=9C=A8=E7=9A=84=20403=20?= =?UTF-8?q?=E4=B8=8D=E5=86=8D=E5=81=9C=E7=94=A8=E8=B4=A6=E5=8F=B7=E5=B9=B6?= =?UTF-8?q?=E6=8D=A2=E5=8F=B7?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 上游对不存在、已过期、已删除或不属于当前项目的 cachedContent 一律返回 403 "CachedContent not found (or permission denied)"。原生转发路径在错误 策略之前识别该响应,原样返回客户端,不标记账号错误、不临时停调度、 不重试、不换号;其余 403 处理不变。 --- .../gemini_cached_content_not_found_test.go | 129 ++++++++++++++++++ .../service/gemini_messages_compat_service.go | 36 +++++ 2 files changed, 165 insertions(+) create mode 100644 backend/internal/service/gemini_cached_content_not_found_test.go diff --git a/backend/internal/service/gemini_cached_content_not_found_test.go b/backend/internal/service/gemini_cached_content_not_found_test.go new file mode 100644 index 0000000000..0bfade0855 --- /dev/null +++ b/backend/internal/service/gemini_cached_content_not_found_test.go @@ -0,0 +1,129 @@ +//go:build unit + +package service + +import ( + "context" + "errors" + "io" + "net/http" + "strings" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +// 上游对不存在、已过期、已删除或不属于当前项目的显式缓存一律返回这个 403。 +const geminiCachedContentNotFoundBody = `{"error":{"code":403,"message":"CachedContent not found (or permission denied)","status":"PERMISSION_DENIED"}}` + +func newGeminiCachedContentTestService(status int, body string) (*GeminiMessagesCompatService, *geminiCompatHTTPUpstreamStub, *errorPolicyRepoStub) { + repo := &errorPolicyRepoStub{} + httpStub := &geminiCompatHTTPUpstreamStub{ + response: &http.Response{ + StatusCode: status, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(body)), + }, + } + svc := &GeminiMessagesCompatService{ + httpUpstream: httpStub, + cfg: &config.Config{}, + rateLimitService: NewRateLimitService(repo, nil, &config.Config{}, nil, nil), + } + return svc, httpStub, repo +} + +func geminiDefaultPolicyAPIKeyAccount() *Account { + return &Account{ + ID: 710, + Platform: PlatformGemini, + Type: AccountTypeAPIKey, + Credentials: map[string]any{ + "api_key": "test-key", + }, + } +} + +const geminiCachedContentRequestBody = `{"cachedContent":"cachedContents/doesnotexist","contents":[{"role":"user","parts":[{"text":"hi"}]}]}` + +func TestGeminiForwardNative_CachedContentNotFoundPassthroughWithoutAccountPenalty(t *testing.T) { + gin.SetMode(gin.TestMode) + for _, tc := range []struct { + name string + action string + stream bool + }{ + {name: "generateContent", action: "generateContent"}, + {name: "streamGenerateContent", action: "streamGenerateContent", stream: true}, + {name: "countTokens", action: "countTokens"}, + } { + t.Run(tc.name, func(t *testing.T) { + svc, httpStub, repo := newGeminiCachedContentTestService(http.StatusForbidden, geminiCachedContentNotFoundBody) + c, rec := newGeminiNativeTestContext(t) + + result, err := svc.ForwardNative(context.Background(), c, geminiDefaultPolicyAPIKeyAccount(), + "gemini-3.8-flash", tc.action, tc.stream, []byte(geminiCachedContentRequestBody)) + + require.Nil(t, result) + require.Error(t, err) + var failoverErr *UpstreamFailoverError + require.False(t, errors.As(err, &failoverErr), "缓存不存在是请求问题，不应换号") + require.Equal(t, 0, repo.setErrCalls, "不应把账号标记为错误") + require.Equal(t, 0, repo.tempCalls, "不应临时停调度") + require.Equal(t, 1, httpStub.calls, "不应重试") + require.Equal(t, http.StatusForbidden, rec.Code) + require.Equal(t, geminiCachedContentNotFoundBody, rec.Body.String(), "响应体应原样透传") + }) + } +} + +func TestGeminiForwardNative_CachedContentNotFoundBypassesTempUnschedulableRules(t *testing.T) { + gin.SetMode(gin.TestMode) + svc, _, repo := newGeminiCachedContentTestService(http.StatusForbidden, geminiCachedContentNotFoundBody) + c, rec := newGeminiNativeTestContext(t) + account := geminiDefaultPolicyAPIKeyAccount() + account.Credentials["temp_unschedulable_enabled"] = true + account.Credentials["temp_unschedulable_rules"] = []any{ + map[string]any{ + "error_code": float64(403), + "keywords": []any{"permission denied"}, + "duration_minutes": float64(10), + "description": "403 permission denied", + }, + } + + _, err := svc.ForwardNative(context.Background(), c, account, + "gemini-3.8-flash", "generateContent", false, []byte(geminiCachedContentRequestBody)) + + require.Error(t, err) + var failoverErr *UpstreamFailoverError + require.False(t, errors.As(err, &failoverErr)) + require.Equal(t, 0, repo.tempCalls, "临时不可调度规则不应命中") + require.Equal(t, 0, repo.setErrCalls) + require.Equal(t, http.StatusForbidden, rec.Code) +} + +func TestGeminiForwardNative_OtherForbiddenStillPenalizesAccount(t *testing.T) { + gin.SetMode(gin.TestMode) + body := `{"error":{"code":403,"message":"Your project has been denied access. Please contact support.","status":"PERMISSION_DENIED"}}` + svc, _, repo := newGeminiCachedContentTestService(http.StatusForbidden, body) + c, _ := newGeminiNativeTestContext(t) + + _, err := svc.ForwardNative(context.Background(), c, geminiDefaultPolicyAPIKeyAccount(), + "gemini-3.8-flash", "generateContent", false, []byte(geminiCachedContentRequestBody)) + + var failoverErr *UpstreamFailoverError + require.True(t, errors.As(err, &failoverErr), "其他 403 仍按账号问题换号") + require.Equal(t, 1, repo.setErrCalls, "其他 403 仍标记账号错误") +} + +func TestIsGeminiCachedContentNotFound(t *testing.T) { + require.True(t, isGeminiCachedContentNotFound(http.StatusForbidden, []byte(geminiCachedContentNotFoundBody))) + require.True(t, isGeminiCachedContentNotFound(http.StatusForbidden, + []byte(`[{"error":{"code":403,"message":"CachedContent not found (or permission denied)","status":"PERMISSION_DENIED"}}]`))) + require.False(t, isGeminiCachedContentNotFound(http.StatusBadRequest, []byte(geminiCachedContentNotFoundBody))) + require.False(t, isGeminiCachedContentNotFound(http.StatusForbidden, + []byte(`{"error":{"code":403,"message":"Permission denied: Consumer has been suspended.","status":"PERMISSION_DENIED"}}`))) +} diff --git a/backend/internal/service/gemini_messages_compat_service.go b/backend/internal/service/gemini_messages_compat_service.go index 278cbfa6fa..01c3b2cc07 100644 --- a/backend/internal/service/gemini_messages_compat_service.go +++ b/backend/internal/service/gemini_messages_compat_service.go @@ -1374,6 +1374,14 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin. return nil, transportErr } + // 引用缓存不存在先于错误策略识别：错误策略会写临时不可调度，重试与换号也无法恢复。 + if notFound, rebuilt := s.cachedContentNotFoundInLoop(resp); notFound { + resp = rebuilt + break + } else { + resp = rebuilt + } + // 错误策略优先：匹配则跳过重试直接处理。 if matched, rebuilt := s.checkErrorPolicyInLoop(ctx, account, resp, mappedModel); matched { resp = rebuilt @@ -1467,6 +1475,10 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin. if resp.StatusCode >= 400 { respBody := s.readUpstreamErrorBody(resp) + // 引用缓存不存在是请求问题：原样返回客户端，不改账号状态、不换号。 + if isGeminiCachedContentNotFound(resp.StatusCode, respBody) { + return nil, s.writeGeminiNativeUpstreamError(c, account, resp, respBody, requestID, isOAuth) + } // Best-effort fallback for OAuth tokens missing AI Studio scopes when calling countTokens. // This avoids Gemini SDKs failing hard during preflight token counting. // Checked before error policy so it always works regardless of custom error codes. @@ -1676,6 +1688,21 @@ func (s *GeminiMessagesCompatService) checkErrorPolicyInLoop( return policy != ErrorPolicyNone, rebuilt } +// cachedContentNotFoundInLoop 判断上游响应是否为引用缓存不存在的 403；读取过的响应体以内存副本重建。 +func (s *GeminiMessagesCompatService) cachedContentNotFoundInLoop(resp *http.Response) (notFound bool, rebuilt *http.Response) { + if resp.StatusCode != http.StatusForbidden { + return false, resp + } + body := s.readUpstreamErrorBody(resp) + _ = resp.Body.Close() + rebuilt = &http.Response{ + StatusCode: resp.StatusCode, + Header: resp.Header.Clone(), + Body: io.NopCloser(bytes.NewReader(body)), + } + return isGeminiCachedContentNotFound(resp.StatusCode, body), rebuilt +} + func (s *GeminiMessagesCompatService) shouldRetryGeminiUpstreamError(account *Account, statusCode int) bool { switch statusCode { case 429, 500, 502, 503, 504, 529: @@ -2635,6 +2662,15 @@ func isGeminiInsufficientScope(headers http.Header, body []byte) bool { return strings.Contains(lower, "insufficient authentication scopes") || strings.Contains(lower, "access_token_scope_insufficient") } +// isGeminiCachedContentNotFound 识别请求引用的显式缓存不存在、已过期、已删除或不属于当前项目时 +// 上游返回的 403。四种情况上游响应完全相同，均由请求的 cachedContent 决定，与账号凭据状态无关。 +func isGeminiCachedContentNotFound(statusCode int, body []byte) bool { + if statusCode != http.StatusForbidden { + return false + } + return strings.Contains(strings.ToLower(string(body)), "cachedcontent not found") +} + func estimateGeminiCountTokens(reqBody []byte) int { total := 0 From d5a802f84a6ae3405bee2a6e0160675785abf237 Mon Sep 17 00:00:00 2001 From: feeeei Date: Sat, 10 Oct 2026 22:21:28 +0800 Subject: [PATCH 02/15] =?UTF-8?q?feat(gemini):=20=E6=94=AF=E6=8C=81?= =?UTF-8?q?=E6=98=BE=E5=BC=8F=E4=B8=8A=E4=B8=8B=E6=96=87=E7=BC=93=E5=AD=98?= =?UTF-8?q?=EF=BC=88AI=20Studio=20/=20Vertex=20/=20=E8=87=AA=E5=BB=BA?= =?UTF-8?q?=E7=BD=91=E5=85=B3=E4=B8=8A=E6=B8=B8=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 新增 /v1beta/cachedContents 的创建、列表、查询、更新有效期与删除接口。 参与账号：官方地址的 AI Studio API Key 账号、Vertex service account 账号， 以及开启"上游支持显式缓存"的自定义地址 API Key 账号。 - 网关自发缓存 ID，绑定上游资源名、持有账号与创建它的 API Key （迁移 243 gemini_cached_contents）；list/get 只读本地记录 - 引用缓存的 generateContent / streamGenerateContent / countTokens 改写为上游资源名，只发往持有缓存的账号，不参与会话粘性、不换号； 拒绝重复键与 snake_case 写法；在持有账号上重试耗尽后， 上游 429 / 5xx 原样透传，401 / 403 映射为 502 - 计费：创建时一次性收取缓存 token 输入费与整个有效期的存储费； 延长按库中实际推进的时长收费，纯存储费不算一次请求；缩短与删除不退费 - 存储单价随计费模型取价，来自定价覆盖文件字段 cache_storage_cost_per_token_per_hour，未配置的模型拒绝创建； 存储费同时计入账号统计成本，渠道按 token 重算成本时叠加且不重复计入 - 扣费前校验：创建在上游建成后按确定金额校验余额 / 订阅额度 / API Key 额度并登记不封顶的在途预留，不足时删除上游缓存、只收输入费； 延长在调用上游前按增量时长校验 - 有效期上限 gateway.gemini_cached_content_max_ttl_seconds，默认 86400， 有效期向上取整到秒；创建请求保留数字原文、兼容 snake_case 元数据字段 - 上游调用与落库同客户端连接脱钩；2xx 不可用响应不换号并删除已建缓存 - 缓存不存在：AI Studio 的 403、Vertex 生成请求的 400 与缓存资源接口的 404 统一识别；原生转发（含流式）不重试、不换号、不改账号状态， 向客户端返回标准 403，不暴露上游缓存资源 ID，上游原始错误登记到 ops； PATCH 遇上游 404 移除本地记录并返回 403，DELETE 视为成功 - 自建网关上游：账号开关 explicit_cache_upstream，缓存接口发往账号自身 base_url；可选 explicit_cache_upstream_max_ttl_seconds 对齐上游 TTL 上限， 显式 TTL 超限时创建跳过该账号（全部超限返回 400），未显式时按上限创建， PATCH 取本站与账号上限的较小值；绑定缓存的请求在网关账号上收到 429 / 5xx 时不重试、不限流该账号；前端创建 / 编辑账号弹窗新增开关与上限输入框 - 用量与 ops 端点统计：带缓存 ID 的路径归一化为 /v1beta/cachedContents， 缓存操作的上游端点记为 /v1beta/cachedContents --- backend/cmd/server/wire_gen.go | 4 +- backend/internal/config/config.go | 3 + backend/internal/handler/endpoint.go | 9 +- backend/internal/handler/endpoint_test.go | 3 + backend/internal/handler/gateway_handler.go | 14 +- .../gemini_cached_contents_billing_test.go | 413 ++++++++ .../handler/gemini_cached_contents_handler.go | 806 +++++++++++++++ .../gemini_cached_contents_handler_test.go | 220 ++++ .../internal/handler/gemini_v1beta_handler.go | 62 +- .../handler/security_audit_order_test.go | 1 + backend/internal/handler/wire.go | 2 + .../repository/gemini_cached_content_repo.go | 124 +++ ...ni_cached_content_repo_integration_test.go | 82 ++ backend/internal/repository/wire.go | 1 + backend/internal/server/routes/gateway.go | 6 + .../prompt_audit_route_coverage_test.go | 1 + .../internal/service/billing_cache_service.go | 69 +- .../service/billing_inflight_reservation.go | 14 +- backend/internal/service/billing_service.go | 13 + backend/internal/service/gateway_service.go | 3 + .../internal/service/gateway_usage_billing.go | 172 +++- .../internal/service/gemini_cached_content.go | 957 ++++++++++++++++++ .../gemini_cached_content_hardening_test.go | 368 +++++++ .../gemini_cached_content_not_found_test.go | 55 +- .../service/gemini_cached_content_test.go | 506 +++++++++ .../service/gemini_messages_compat_service.go | 97 +- backend/internal/service/pricing_service.go | 6 + backend/internal/service/wire.go | 1 + .../migrations/243_gemini_cached_contents.sql | 30 + deploy/config.example.yaml | 9 + .../components/account/CreateAccountModal.vue | 53 + .../components/account/EditAccountModal.vue | 61 ++ .../__tests__/CreateAccountModal.spec.ts | 37 + .../__tests__/EditAccountModal.spec.ts | 72 ++ .../src/i18n/locales/en/admin/accounts.ts | 6 + .../src/i18n/locales/zh/admin/accounts.ts | 6 + .../__tests__/geminiExplicitCache.spec.ts | 50 + frontend/src/utils/geminiExplicitCache.ts | 47 + 38 files changed, 4286 insertions(+), 97 deletions(-) create mode 100644 backend/internal/handler/gemini_cached_contents_billing_test.go create mode 100644 backend/internal/handler/gemini_cached_contents_handler.go create mode 100644 backend/internal/handler/gemini_cached_contents_handler_test.go create mode 100644 backend/internal/repository/gemini_cached_content_repo.go create mode 100644 backend/internal/repository/gemini_cached_content_repo_integration_test.go create mode 100644 backend/internal/service/gemini_cached_content.go create mode 100644 backend/internal/service/gemini_cached_content_hardening_test.go create mode 100644 backend/internal/service/gemini_cached_content_test.go create mode 100644 backend/migrations/243_gemini_cached_contents.sql create mode 100644 frontend/src/utils/__tests__/geminiExplicitCache.spec.ts create mode 100644 frontend/src/utils/geminiExplicitCache.ts diff --git a/backend/cmd/server/wire_gen.go b/backend/cmd/server/wire_gen.go index 1eda60da76..2c2f0cfbc9 100644 --- a/backend/cmd/server/wire_gen.go +++ b/backend/cmd/server/wire_gen.go @@ -296,7 +296,9 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { userMessageQueueService := service.ProvideUserMessageQueueService(userMsgQueueCache, rpmCache, configConfig) legacyEngine := securityaudit.NewLegacyModerationAdapter(contentModerationService) coordinator := securityaudit.NewCoordinator(legacyEngine, promptService) - gatewayHandler := handler.ProvideGatewayHandler(gatewayService, openAIGatewayService, geminiMessagesCompatService, antigravityGatewayService, userService, concurrencyService, billingCacheService, usageService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, userMessageQueueService, configConfig, settingService, coordinator) + geminiCachedContentRepository := repository.NewGeminiCachedContentRepository(db) + geminiCachedContentService := service.NewGeminiCachedContentService(geminiCachedContentRepository, geminiMessagesCompatService, configConfig) + gatewayHandler := handler.ProvideGatewayHandler(gatewayService, openAIGatewayService, geminiMessagesCompatService, antigravityGatewayService, userService, concurrencyService, billingCacheService, usageService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, userMessageQueueService, configConfig, settingService, coordinator, geminiCachedContentService) openAIGatewayHandler := handler.ProvideOpenAIGatewayHandler(openAIGatewayService, pluginManager, concurrencyService, billingCacheService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, opsService, grokQuotaService, configConfig, coordinator, compositeRouteResolver) handlerSettingHandler := handler.ProvideSettingHandler(settingService, buildInfo, notificationEmailService) totpHandler := handler.NewTotpHandler(totpService) diff --git a/backend/internal/config/config.go b/backend/internal/config/config.go index fe0224aed0..bdb8086582 100644 --- a/backend/internal/config/config.go +++ b/backend/internal/config/config.go @@ -1005,6 +1005,8 @@ type GatewayConfig struct { ProxyProbeResponseReadMaxBytes int64 `mapstructure:"proxy_probe_response_read_max_bytes"` // Gemini 上游响应头调试日志开关（默认关闭，避免高频日志开销） GeminiDebugResponseHeaders bool `mapstructure:"gemini_debug_response_headers"` + // GeminiCachedContentMaxTTLSeconds: Gemini 显式缓存单次可设置的最长有效期（秒），<=0 时按 86400 处理 + GeminiCachedContentMaxTTLSeconds int `mapstructure:"gemini_cached_content_max_ttl_seconds"` // ConnectionPoolIsolation: 上游连接池隔离策略（proxy/account/account_proxy） ConnectionPoolIsolation string `mapstructure:"connection_pool_isolation"` // ForceCodexCLI: 强制将 OpenAI `/v1/responses` 请求按 Codex CLI 处理。 @@ -2527,6 +2529,7 @@ func setDefaults() { viper.SetDefault("gateway.models_list_read_max_bytes", DefaultModelsListReadMaxBytes) viper.SetDefault("gateway.proxy_probe_response_read_max_bytes", int64(1024*1024)) viper.SetDefault("gateway.gemini_debug_response_headers", false) + viper.SetDefault("gateway.gemini_cached_content_max_ttl_seconds", 86400) viper.SetDefault("gateway.connection_pool_isolation", ConnectionPoolIsolationAccountProxy) // HTTP 上游连接池配置（针对 5000+ 并发用户优化） viper.SetDefault("gateway.max_idle_conns", 2560) // 最大空闲连接总数（高并发场景可调大） diff --git a/backend/internal/handler/endpoint.go b/backend/internal/handler/endpoint.go index 82d8279e90..11ddc065c4 100644 --- a/backend/internal/handler/endpoint.go +++ b/backend/internal/handler/endpoint.go @@ -32,6 +32,7 @@ const ( EndpointVideos = "/v1/videos" EndpointSeedanceTasks = "/api/v3/contents/generations/tasks" EndpointGeminiModels = "/v1beta/models" + EndpointGeminiCachedContents = "/v1beta/cachedContents" ) const EndpointAntigravityGenerateContent = "/v1internal:streamGenerateContent" @@ -53,6 +54,7 @@ const ( // "/v1/chat/completions" → "/v1/chat/completions" // "/openai/v1/responses/foo" → "/v1/responses" // "/v1beta/models/gemini:gen" → "/v1beta/models" +// "/v1beta/cachedContents/abc" → "/v1beta/cachedContents" // // The OpenAI Responses API is also exposed via a few bare/alias // routes that do not carry a "/v1/" prefix (top-level bare route and @@ -117,6 +119,8 @@ func NormalizeInboundEndpoint(path string) string { return EndpointResponses case strings.Contains(path, EndpointGeminiModels): return EndpointGeminiModels + case strings.Contains(path, EndpointGeminiCachedContents): + return EndpointGeminiCachedContents default: return path } @@ -194,7 +198,7 @@ func isBareOrSubpathOf(path, root string) bool { // retain their paths. Grok raw Chat requests override this through the // forwarding result consumed by resolveOpenAIUpstreamEndpoint. // - Anthropic → /v1/messages -// - Gemini → /v1beta/models +// - Gemini → /v1beta/models（显式缓存管理接口保留 /v1beta/cachedContents） // - Antigravity → /v1/messages (Claude) or gemini (Gemini) // - Antigravity routes may target either Claude or Gemini, so the // inbound endpoint is used to distinguish. @@ -230,6 +234,9 @@ func DeriveUpstreamEndpoint(inbound, rawRequestPath, platform string) string { return EndpointSystemOne case service.PlatformGemini: + if inbound == EndpointGeminiCachedContents { + return EndpointGeminiCachedContents + } return EndpointGeminiModels case service.PlatformAntigravity: diff --git a/backend/internal/handler/endpoint_test.go b/backend/internal/handler/endpoint_test.go index 6941e56374..7fbb4aae10 100644 --- a/backend/internal/handler/endpoint_test.go +++ b/backend/internal/handler/endpoint_test.go @@ -37,6 +37,8 @@ func TestNormalizeInboundEndpoint(t *testing.T) { {"/v1/videos/generations", EndpointVideosGenerations}, {"/v1/videos/req_123", EndpointVideos}, {"/v1beta/models", EndpointGeminiModels}, + {"/v1beta/cachedContents", EndpointGeminiCachedContents}, + {"/v1beta/cachedContents/gfo2i2haszm3wbwiyivyok5cyddubfwsr77efg7h", EndpointGeminiCachedContents}, // Prefixed paths (antigravity, openai) — root Responses. {"/antigravity/v1/messages", EndpointMessages}, @@ -102,6 +104,7 @@ func TestDeriveUpstreamEndpoint(t *testing.T) { // Gemini. {"gemini models", EndpointGeminiModels, "/v1beta/models/gemini:gen", service.PlatformGemini, EndpointGeminiModels}, + {"gemini cached contents", EndpointGeminiCachedContents, "/v1beta/cachedContents/abc", service.PlatformGemini, EndpointGeminiCachedContents}, // OpenAI — root Responses. {"openai responses root", EndpointResponses, "/v1/responses", service.PlatformOpenAI, EndpointResponses}, diff --git a/backend/internal/handler/gateway_handler.go b/backend/internal/handler/gateway_handler.go index 0b58876d8d..4ed38abb10 100644 --- a/backend/internal/handler/gateway_handler.go +++ b/backend/internal/handler/gateway_handler.go @@ -52,12 +52,14 @@ type GatewayHandler struct { errorPassthroughService *service.ErrorPassthroughService contentModerationService *service.ContentModerationService securityAuditCoordinator *securityaudit.Coordinator - concurrencyHelper *ConcurrencyHelper - userMsgQueueHelper *UserMsgQueueHelper - maxAccountSwitches int - maxAccountSwitchesGemini int - cfg *config.Config - settingService *service.SettingService + // geminiCachedContentService 管理 Gemini 显式缓存；为 nil 时相关接口不可用。 + geminiCachedContentService *service.GeminiCachedContentService + concurrencyHelper *ConcurrencyHelper + userMsgQueueHelper *UserMsgQueueHelper + maxAccountSwitches int + maxAccountSwitchesGemini int + cfg *config.Config + settingService *service.SettingService } // NewGatewayHandler creates a new GatewayHandler diff --git a/backend/internal/handler/gemini_cached_contents_billing_test.go b/backend/internal/handler/gemini_cached_contents_billing_test.go new file mode 100644 index 0000000000..f567c2cb01 --- /dev/null +++ b/backend/internal/handler/gemini_cached_contents_billing_test.go @@ -0,0 +1,413 @@ +//go:build unit + +package handler + +import ( + "context" + "errors" + "fmt" + "io" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "sync" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey" + middleware "github.com/Wei-Shaw/sub2api/internal/server/middleware" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +const ( + cachedBillingInputPrice = 0.00000125 + cachedBillingStoragePrice = 0.0000045 + cachedBillingTokens = 1_000_000 +) + +type cachedBillingRepo struct { + mu sync.Mutex + records map[string]*service.GeminiCachedContent + failUpdate bool + nextID int64 + deleted []int64 +} + +func (r *cachedBillingRepo) Create(_ context.Context, record *service.GeminiCachedContent) error { + r.mu.Lock() + defer r.mu.Unlock() + r.nextID++ + record.ID = r.nextID + cp := *record + r.records[record.PublicID] = &cp + return nil +} + +func (r *cachedBillingRepo) GetForOwner(_ context.Context, apiKeyID int64, publicID string) (*service.GeminiCachedContent, error) { + r.mu.Lock() + defer r.mu.Unlock() + record, ok := r.records[publicID] + if !ok || record.APIKeyID != apiKeyID { + return nil, service.ErrGeminiCachedContentNotFound + } + cp := *record + return &cp, nil +} + +func (r *cachedBillingRepo) ListForOwner(context.Context, int64, int64, time.Time, int64, int) ([]*service.GeminiCachedContent, error) { + return nil, nil +} + +func (r *cachedBillingRepo) UpdateExpireTime(_ context.Context, id int64, expireTime time.Time) (time.Time, error) { + r.mu.Lock() + defer r.mu.Unlock() + if r.failUpdate { + return time.Time{}, errors.New("db unavailable") + } + for _, record := range r.records { + if record.ID == id { + previous := record.ExpireTime + record.ExpireTime = expireTime + return previous, nil + } + } + return time.Time{}, service.ErrGeminiCachedContentNotFound +} + +func (r *cachedBillingRepo) SoftDelete(_ context.Context, id int64) error { + r.mu.Lock() + defer r.mu.Unlock() + r.deleted = append(r.deleted, id) + return nil +} + +type cachedBillingAccountRepo struct { + service.AccountRepository + accounts map[int64]*service.Account +} + +func (r *cachedBillingAccountRepo) GetByID(_ context.Context, id int64) (*service.Account, error) { + account, ok := r.accounts[id] + if !ok { + return nil, service.ErrAccountNotFound + } + cp := *account + return &cp, nil +} + +type cachedBillingUpstream struct { + service.HTTPUpstream + mu sync.Mutex + createBody string + calls []string + beforePost func(n int) + // failures 按请求方法返回上游错误响应。 + failures map[string]cachedBillingUpstreamFailure +} + +type cachedBillingUpstreamFailure struct { + status int + body string +} + +func (u *cachedBillingUpstream) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) { + u.mu.Lock() + u.calls = append(u.calls, req.Method+" "+req.URL.Path) + n := len(u.calls) + failure, failed := u.failures[req.Method] + u.mu.Unlock() + if failed { + return &http.Response{StatusCode: failure.status, Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(failure.body))}, nil + } + body := "{}" + switch req.Method { + case http.MethodPost: + body = u.createBody + case http.MethodPatch: + if u.beforePost != nil { + u.beforePost(n) + } + body = fmt.Sprintf(`{"name":"cachedContents/up1","expireTime":%q}`, time.Now().Add(24*time.Hour).UTC().Format(time.RFC3339Nano)) + } + return &http.Response{StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(body))}, nil +} + +func (u *cachedBillingUpstream) methodCalls(method string) int { + u.mu.Lock() + defer u.mu.Unlock() + n := 0 + for _, call := range u.calls { + if strings.HasPrefix(call, method+" ") { + n++ + } + } + return n +} + +type cachedBillingUserRepo struct { + service.UserRepository + balance float64 +} + +func (r *cachedBillingUserRepo) GetByID(_ context.Context, id int64) (*service.User, error) { + return &service.User{ID: id, Balance: r.balance}, nil +} + +func (r *cachedBillingUserRepo) DeductBalance(context.Context, int64, float64) error { return nil } + +type cachedBillingUsageRepo struct { + service.UsageLogRepository + mu sync.Mutex + logs []*service.UsageLog +} + +func (r *cachedBillingUsageRepo) Create(_ context.Context, log *service.UsageLog) (bool, error) { + r.mu.Lock() + defer r.mu.Unlock() + r.logs = append(r.logs, log) + return true, nil +} + +func (r *cachedBillingUsageRepo) totals() []float64 { + r.mu.Lock() + defer r.mu.Unlock() + out := make([]float64, 0, len(r.logs)) + for _, log := range r.logs { + out = append(out, log.TotalCost) + } + return out +} + +type cachedBillingEnv struct { + h *GatewayHandler + repo *cachedBillingRepo + upstream *cachedBillingUpstream + usage *cachedBillingUsageRepo + apiKey *service.APIKey + group *service.Group +} + +func newCachedBillingEnv(t *testing.T, balance float64, accountIDs ...int64) *cachedBillingEnv { + t.Helper() + dir := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(dir, "model_pricing.json"), []byte(fmt.Sprintf(`{"gemini-2.5-pro":{ + "litellm_provider":"gemini","mode":"chat","input_cost_per_token":%g,"output_cost_per_token":0.00001, + "cache_storage_cost_per_token_per_hour":%g}}`, cachedBillingInputPrice, cachedBillingStoragePrice)), 0o644)) + cfg := &config.Config{} + cfg.Pricing.DataDir = dir + cfg.Pricing.UpdateIntervalHours = 100000 + cfg.Default.RateMultiplier = 1 + pricing := service.NewPricingService(cfg, nil) + require.NoError(t, pricing.Initialize()) + billing := service.NewBillingService(cfg, pricing) + + groupID := int64(42) + group := &service.Group{ID: groupID, Hydrated: true, Platform: service.PlatformGemini, Status: service.StatusActive, RateMultiplier: 1} + accounts := map[int64]*service.Account{} + var scheduled []*service.Account + for _, id := range accountIDs { + account := &service.Account{ID: id, Platform: service.PlatformGemini, Type: service.AccountTypeAPIKey, + Credentials: map[string]any{"api_key": "k"}, Status: service.StatusActive, Schedulable: true, + Concurrency: 5, Priority: 1, AccountGroups: []service.AccountGroup{{AccountID: id, GroupID: groupID}}} + accounts[id] = account + scheduled = append(scheduled, account) + } + userRepo := &cachedBillingUserRepo{balance: balance} + usage := &cachedBillingUsageRepo{} + bcs := service.NewBillingCacheService(nil, userRepo, nil, nil, nil, nil, cfg, nil) + t.Cleanup(bcs.Stop) + snapshot := service.NewSchedulerSnapshotService(&fakeSchedulerCache{accounts: scheduled}, nil, nil, nil, nil) + gw := service.NewGatewayService(nil, &fakeGroupRepo{group: group}, usage, nil, userRepo, nil, nil, nil, cfg, snapshot, nil, + billing, nil, &service.BillingCacheService{}, nil, nil, &service.DeferredService{}, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil) + upstream := &cachedBillingUpstream{} + compat := service.NewGeminiMessagesCompatService(&cachedBillingAccountRepo{accounts: accounts}, nil, nil, nil, nil, nil, upstream, nil, cfg) + repo := &cachedBillingRepo{records: map[string]*service.GeminiCachedContent{}} + h := &GatewayHandler{ + gatewayService: gw, + billingCacheService: bcs, + concurrencyHelper: NewConcurrencyHelper(service.NewConcurrencyService(&fakeConcurrencyCache{}), SSEPingFormatNone, 0), + geminiCachedContentService: service.NewGeminiCachedContentService(repo, compat, cfg), + } + apiKey := &service.APIKey{ID: 7, UserID: 3, GroupID: &groupID, Group: group, Status: service.StatusActive, + User: &service.User{ID: 3, Concurrency: 10, Balance: balance}} + return &cachedBillingEnv{h: h, repo: repo, upstream: upstream, usage: usage, apiKey: apiKey, group: group} +} + +func (e *cachedBillingEnv) serve(method, target, body string, params gin.Params, call func(*gin.Context)) *httptest.ResponseRecorder { + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + req := httptest.NewRequest(method, target, strings.NewReader(body)) + c.Request = req.WithContext(context.WithValue(req.Context(), ctxkey.Group, e.group)) + c.Params = params + c.Set(string(middleware.ContextKeyAPIKey), e.apiKey) + c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: e.apiKey.UserID, Concurrency: 10}) + call(c) + return rec +} + +// waitForUsage 等待异步计费任务落库（无 worker 池时同步执行，直接返回）。 +func (e *cachedBillingEnv) waitForUsage(t *testing.T, n int) []float64 { + t.Helper() + require.Eventually(t, func() bool { return len(e.usage.totals()) >= n }, time.Second, 10*time.Millisecond) + return e.usage.totals() +} + +func (e *cachedBillingEnv) create() *httptest.ResponseRecorder { + return e.serve(http.MethodPost, "/v1beta/cachedContents", + `{"model":"models/gemini-2.5-pro","ttl":"86400s","contents":[{"role":"user","parts":[{"text":"doc"}]}]}`, + nil, e.h.GeminiCachedContentsCreate) +} + +func (e *cachedBillingEnv) seed(expire time.Time) { + e.repo.records["abc"] = &service.GeminiCachedContent{ + ID: 100, PublicID: "abc", UserID: 3, APIKeyID: 7, GroupID: 42, AccountID: 900, + UpstreamName: "cachedContents/up1", Model: "gemini-2.5-pro", UpstreamModel: "models/gemini-2.5-pro", + TotalTokenCount: cachedBillingTokens, ExpireTime: expire, + } +} + +func (e *cachedBillingEnv) patch() *httptest.ResponseRecorder { + return e.serve(http.MethodPatch, "/v1beta/cachedContents/abc?updateMask=ttl", `{"ttl":"86400s"}`, + gin.Params{{Key: "cacheID", Value: "abc"}}, e.h.GeminiCachedContentsPatch) +} + +func cachedBillingCreateBody(expire time.Time) string { + return fmt.Sprintf(`{"name":"cachedContents/up1","usageMetadata":{"totalTokenCount":%d},"expireTime":%q}`, + cachedBillingTokens, expire.UTC().Format(time.RFC3339Nano)) +} + +func TestGeminiCachedContentsCreate_ChargesInputAndStorage(t *testing.T) { + env := newCachedBillingEnv(t, 500, 900) + env.upstream.createBody = cachedBillingCreateBody(time.Now().Add(100 * time.Hour)) + + rec := env.create() + require.Equal(t, http.StatusOK, rec.Code, rec.Body.String()) + totals := env.waitForUsage(t, 1) + require.Len(t, totals, 1) + require.InDelta(t, cachedBillingTokens*cachedBillingInputPrice+cachedBillingTokens*cachedBillingStoragePrice*24, totals[0], 1e-9) + require.Len(t, env.repo.records, 1) + for _, record := range env.repo.records { + require.False(t, record.ExpireTime.After(time.Now().Add(24*time.Hour)), "本地到期时间不超过请求的有效期") + } +} + +func TestGeminiCachedContentsCreate_UnaffordableStorageDeletesUpstreamAndBillsInputOnly(t *testing.T) { + env := newCachedBillingEnv(t, 1, 900) + env.upstream.createBody = cachedBillingCreateBody(time.Now().Add(24 * time.Hour)) + + rec := env.create() + require.NotEqual(t, http.StatusOK, rec.Code) + require.Contains(t, rec.Body.String(), "error") + require.Equal(t, 1, env.upstream.methodCalls(http.MethodDelete), "删除已创建的上游缓存，避免产生存储费") + require.Empty(t, env.repo.records, "不登记缓存") + totals := env.waitForUsage(t, 1) + require.Len(t, totals, 1) + require.InDelta(t, cachedBillingTokens*cachedBillingInputPrice, totals[0], 1e-9, "只收已发生的缓存 token 输入费") +} + +func TestGeminiCachedContentsCreate_UnusableUpstreamResponseCleansUpWithoutRetry(t *testing.T) { + env := newCachedBillingEnv(t, 500, 900, 901) + env.upstream.createBody = `{"name":"cachedContents/up1","expireTime":"2099-01-01T00:00:00Z"}` + + rec := env.create() + require.Equal(t, http.StatusBadGateway, rec.Code) + require.Equal(t, 1, env.upstream.methodCalls(http.MethodPost), "不换号重建") + require.Equal(t, 1, env.upstream.methodCalls(http.MethodDelete), "删除已创建的上游缓存") + require.Empty(t, env.repo.records) + require.Empty(t, env.usage.totals()) +} + +func TestGeminiCachedContentsPatch_ConcurrentExtensionsChargeOnlyActualProgress(t *testing.T) { + env := newCachedBillingEnv(t, 1000, 900) + original := time.Now().Add(time.Hour) + env.seed(original) + env.upstream.beforePost = func(n int) { + if n == 1 { + // 第一个 PATCH 的上游调用返回前，第二个 PATCH 完整执行（读到的仍是旧到期时间）。 + require.Equal(t, http.StatusOK, env.patch().Code) + } + } + require.Equal(t, http.StatusOK, env.patch().Code) + + totals := env.waitForUsage(t, 1) + final := env.repo.records["abc"].ExpireTime + fair := cachedBillingTokens * cachedBillingStoragePrice * final.Sub(original).Hours() + sum := 0.0 + for _, v := range totals { + sum += v + } + require.InDelta(t, fair, sum, cachedBillingTokens*cachedBillingStoragePrice*(time.Minute.Hours()), "两次延长合计只收实际推进的时长") +} + +func TestGeminiCachedContentsPatch_PersistFailureDoesNotCharge(t *testing.T) { + env := newCachedBillingEnv(t, 1000, 900) + env.seed(time.Now().Add(time.Hour)) + env.repo.failUpdate = true + + rec := env.patch() + require.Equal(t, http.StatusInternalServerError, rec.Code) + require.Equal(t, 1, env.upstream.methodCalls(http.MethodPatch)) + require.Empty(t, env.usage.totals(), "落库失败不收费") +} + +func TestGeminiCachedContentsPatch_UnaffordableExtensionRejectedBeforeUpstream(t *testing.T) { + env := newCachedBillingEnv(t, 1, 900) + env.seed(time.Now().Add(time.Hour)) + + rec := env.patch() + require.NotEqual(t, http.StatusOK, rec.Code) + require.Zero(t, env.upstream.methodCalls(http.MethodPatch), "余额不足以支付延长部分时不调用上游") + require.Empty(t, env.usage.totals()) +} + +// 上游对缓存不存在的响应：AI Studio 403；Vertex 缓存资源接口 404，文案含上游缓存数字 ID。 +var cachedBillingUpstreamNotFoundCases = []struct { + name string + failure cachedBillingUpstreamFailure +}{ + {name: "ai studio 403", failure: cachedBillingUpstreamFailure{status: http.StatusForbidden, + body: `{"error":{"code":403,"message":"CachedContent not found (or permission denied)","status":"PERMISSION_DENIED"}}`}}, + {name: "vertex 404", failure: cachedBillingUpstreamFailure{status: http.StatusNotFound, + body: `{"error":{"code":404,"message":"Cached content 8079767076522164224 is not found.","status":"NOT_FOUND"}}`}}, +} + +func TestGeminiCachedContentsPatch_UpstreamNotFoundForgetsAndReturnsStandardForbidden(t *testing.T) { + for _, tc := range cachedBillingUpstreamNotFoundCases { + t.Run(tc.name, func(t *testing.T) { + env := newCachedBillingEnv(t, 1000, 900) + env.seed(time.Now().Add(time.Hour)) + env.upstream.failures = map[string]cachedBillingUpstreamFailure{http.MethodPatch: tc.failure} + + rec := env.patch() + require.Equal(t, http.StatusForbidden, rec.Code) + require.Equal(t, service.GeminiCachedContentNotFoundResponse, rec.Body.String()) + require.NotContains(t, rec.Body.String(), "8079767076522164224", "不应暴露上游缓存 ID") + require.Equal(t, []int64{100}, env.repo.deleted, "上游已不存在的缓存应从本地移除") + require.Empty(t, env.usage.totals(), "上游未延长不收费") + }) + } +} + +func TestGeminiCachedContentsDelete_UpstreamNotFoundStillForgets(t *testing.T) { + for _, tc := range cachedBillingUpstreamNotFoundCases { + t.Run(tc.name, func(t *testing.T) { + env := newCachedBillingEnv(t, 1000, 900) + env.seed(time.Now().Add(time.Hour)) + env.upstream.failures = map[string]cachedBillingUpstreamFailure{http.MethodDelete: tc.failure} + + rec := env.serve(http.MethodDelete, "/v1beta/cachedContents/abc", "", + gin.Params{{Key: "cacheID", Value: "abc"}}, env.h.GeminiCachedContentsDelete) + require.Equal(t, http.StatusOK, rec.Code, rec.Body.String()) + require.JSONEq(t, `{}`, rec.Body.String()) + require.Equal(t, []int64{100}, env.repo.deleted) + }) + } +} diff --git a/backend/internal/handler/gemini_cached_contents_handler.go b/backend/internal/handler/gemini_cached_contents_handler.go new file mode 100644 index 0000000000..82b3741fad --- /dev/null +++ b/backend/internal/handler/gemini_cached_contents_handler.go @@ -0,0 +1,806 @@ +package handler + +import ( + "context" + "encoding/base64" + "encoding/json" + "errors" + "fmt" + "net/http" + "strconv" + "strings" + "time" + + pkghttputil "github.com/Wei-Shaw/sub2api/internal/pkg/httputil" + "github.com/Wei-Shaw/sub2api/internal/pkg/ip" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "github.com/Wei-Shaw/sub2api/internal/server/middleware" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/gin-gonic/gin" + "github.com/google/uuid" + "go.uber.org/zap" +) + +// geminiCachedContentCreateMaxAttempts 限制创建缓存时实际发往上游的账号数。 +const geminiCachedContentCreateMaxAttempts = 3 + +// geminiCachedContentPersistTimeout 是上游操作完成后落库的超时；落库与客户端连接脱钩。 +const geminiCachedContentPersistTimeout = 15 * time.Second + +const geminiCachedContentUnavailableMessage = "The account holding this cached content is temporarily unavailable" + +// geminiCachedContentsPreamble 校验显式缓存管理接口的公共前置条件：gemini 分组、非 composite 分组、功能已装配。 +func (h *GatewayHandler) geminiCachedContentsPreamble(c *gin.Context, component string) (*service.APIKey, middleware.AuthSubject, *zap.Logger, bool) { + apiKey, ok := middleware.GetAPIKeyFromContext(c) + if !ok || apiKey == nil { + googleError(c, http.StatusUnauthorized, "Invalid API key") + return nil, middleware.AuthSubject{}, nil, false + } + authSubject, ok := middleware.GetAuthSubjectFromContext(c) + if !ok { + googleError(c, http.StatusInternalServerError, "User context not found") + return nil, middleware.AuthSubject{}, nil, false + } + if h.geminiCachedContentService == nil { + googleError(c, http.StatusNotFound, "Explicit context caching is not enabled") + return nil, middleware.AuthSubject{}, nil, false + } + if apiKey.Group == nil || apiKey.GroupID == nil || apiKey.Group.Platform == service.PlatformComposite { + googleError(c, http.StatusBadRequest, "Explicit context caching is not supported for this API key group") + return nil, middleware.AuthSubject{}, nil, false + } + if effectiveAPIKeyPlatform(c, apiKey) != service.PlatformGemini { + googleError(c, http.StatusBadRequest, "API key group platform is not gemini") + return nil, middleware.AuthSubject{}, nil, false + } + reqLog := requestLogger( + c, + component, + zap.Int64("user_id", authSubject.UserID), + zap.Int64("api_key_id", apiKey.ID), + zap.Any("group_id", apiKey.GroupID), + ) + return apiKey, authSubject, reqLog, true +} + +func readGeminiCachedContentBody(c *gin.Context) ([]byte, bool) { + body, err := pkghttputil.ReadRequestBodyWithPrealloc(c.Request) + if err != nil { + if maxErr, ok := extractMaxBytesError(err); ok { + googleError(c, http.StatusRequestEntityTooLarge, buildBodyTooLargeMessage(maxErr.Limit)) + return nil, false + } + googleError(c, http.StatusBadRequest, "Failed to read request body") + return nil, false + } + if len(body) == 0 { + googleError(c, http.StatusBadRequest, "Request body is empty") + return nil, false + } + return body, true +} + +func (h *GatewayHandler) checkGeminiCachedContentBilling(c *gin.Context, reqLog *zap.Logger, apiKey *service.APIKey, subscription *service.UserSubscription) bool { + if err := h.billingCacheService.CheckBillingEligibility(c.Request.Context(), apiKey.User, apiKey, apiKey.Group, subscription, service.QuotaPlatform(c.Request.Context(), apiKey)); err != nil { + reqLog.Info("gemini.cached_content.billing_eligibility_check_failed", zap.Error(err)) + status, _, message, retryAfter := billingErrorDetails(err) + if retryAfter > 0 { + c.Header("Retry-After", strconv.Itoa(retryAfter)) + } + googleError(c, status, message) + return false + } + return true +} + +// GeminiCachedContentsCreate POST /v1beta/cachedContents +// 创建请求内一次性收取缓存 token 的输入费与整个有效期的存储费。 +func (h *GatewayHandler) GeminiCachedContentsCreate(c *gin.Context) { + apiKey, authSubject, reqLog, ok := h.geminiCachedContentsPreamble(c, "handler.gemini_v1beta.cached_contents.create") + if !ok { + return + } + svc := h.geminiCachedContentService + + body, ok := readGeminiCachedContentBody(c) + if !ok { + return + } + req, err := service.ParseGeminiCachedContentCreateRequest(body, svc.Now(), svc.MaxTTL()) + if err != nil { + googleError(c, http.StatusBadRequest, err.Error()) + return + } + modelName := req.Model + reqLog = reqLog.With(zap.String("model", modelName)) + + setOpsRequestContext(c, modelName, false) + setOpsEndpointContext(c, "", int16(service.RequestTypeFromLegacy(false, false))) + pricingCtx, pricingAt := service.WithGatewayTokenRequestPricing(c.Request.Context()) + c.Request = c.Request.WithContext(pricingCtx) + + if decision := h.checkSecurityAudit(c, reqLog, apiKey, authSubject, service.ContentModerationProtocolGemini, modelName, body); decision != nil && !decision.AllowNextStage { + googleSecurityAuditError(c, decision) + return + } + + channelMapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(c.Request.Context(), apiKey.GroupID, modelName) + reqModel := modelName + if channelMapping.Mapped { + modelName = channelMapping.MappedModel + } + + subscription, _ := middleware.GetSubscriptionFromContext(c) + geminiConcurrency := NewConcurrencyHelper(h.concurrencyHelper.concurrencyService, SSEPingFormatNone, 0) + streamStarted := false + userReleaseFunc, err := geminiConcurrency.AcquireUserSlotWithWait(c, authSubject.UserID, authSubject.Concurrency, false, &streamStarted) + if err != nil { + reqLog.Warn("gemini.cached_content.user_slot_acquire_failed", zap.Error(err)) + googleConcurrencyError(c, err, "user") + return + } + userReleaseFunc = wrapReleaseOnDone(c.Request.Context(), userReleaseFunc) + if userReleaseFunc != nil { + defer userReleaseFunc() + } + + if !h.checkGeminiCachedContentBilling(c, reqLog, apiKey, subscription) { + return + } + inflightRelease, err := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(modelName, body)) + if err != nil { + reqLog.Info("gemini.cached_content.inflight_reservation_rejected", zap.Error(err)) + status, _, message, retryAfter := billingErrorDetails(err) + if retryAfter > 0 { + c.Header("Retry-After", strconv.Itoa(retryAfter)) + } + googleError(c, status, message) + return + } + defer inflightRelease() + + excluded, err := h.gatewayService.GeminiExplicitCacheExclusions(c.Request.Context(), apiKey.GroupID, 0) + if err != nil { + reqLog.Warn("gemini.cached_content.list_accounts_failed", zap.Error(err)) + googleError(c, http.StatusServiceUnavailable, "No available Gemini accounts support explicit context caching") + return + } + + var lastUpstreamErr *service.GeminiCachedContentUpstreamError + var ttlRejectedLimit time.Duration + for attempts := 0; attempts < geminiCachedContentCreateMaxAttempts; { + selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, "", modelName, excluded, "", int64(0)) + if err != nil { + if failoverClientGone(c) { + return + } + reqLog.Info("gemini.cached_content.account_select_exhausted", zap.Error(err)) + break + } + account := selection.Account + if !service.IsGeminiExplicitCacheAccount(account) { + if selection.Acquired && selection.ReleaseFunc != nil { + selection.ReleaseFunc() + } + excluded[account.ID] = struct{}{} + continue + } + ttl, ttlAllowed := req.TTLForAccount(account) + if !ttlAllowed { + if selection.Acquired && selection.ReleaseFunc != nil { + selection.ReleaseFunc() + } + excluded[account.ID] = struct{}{} + ttlRejectedLimit = max(ttlRejectedLimit, account.GeminiExplicitCacheUpstreamMaxTTL()) + continue + } + accountReleaseFunc, ok := h.acquireGeminiCachedContentAccountSlot(c, reqLog, geminiConcurrency, selection) + if !ok { + return + } + admissionCtx := service.ContextWithSelectionProfitGate(c.Request.Context(), selection) + latest, vetoed, reason := h.gatewayService.GatewayProfitControlVetoLatest(admissionCtx, account) + if vetoed { + accountReleaseFunc() + reqLog.Debug("gemini.cached_content.account_profit_vetoed", zap.Int64("account_id", account.ID), zap.String("reason", reason)) + excluded[account.ID] = struct{}{} + continue + } + account = latest + attempts++ + setOpsSelectedAccount(c, account.ID, account.Platform) + + mappedModel := account.GetMappedModel(modelName) + usageFields := clientRequestedUsageFields(c, channelMapping, reqModel, mappedModel) + if !h.gatewayService.CacheStoragePriced(c.Request.Context(), &service.RecordUsageInput{ + Result: &service.ForwardResult{Model: modelName, UpstreamModel: mappedModel, CacheStorageTokenHours: 1}, + APIKey: apiKey, + ChannelUsageFields: usageFields, + }) { + accountReleaseFunc() + googleError(c, http.StatusBadRequest, fmt.Sprintf("Explicit context caching is not available for model %s: storage pricing is not configured", modelName)) + return + } + upstreamModel, err := service.GeminiCachedContentUpstreamModelResource(account, mappedModel) + if err != nil { + accountReleaseFunc() + reqLog.Warn("gemini.cached_content.upstream_model_invalid", zap.Int64("account_id", account.ID), zap.Error(err)) + excluded[account.ID] = struct{}{} + continue + } + publicID, err := service.NewGeminiCachedContentPublicID() + if err != nil { + accountReleaseFunc() + googleError(c, http.StatusInternalServerError, "Failed to allocate cached content id") + return + } + + upstream, err := svc.CreateUpstream(c.Request.Context(), account, req, mappedModel, publicID, ttl) + accountReleaseFunc() + if err != nil { + var upErr *service.GeminiCachedContentUpstreamError + if errors.As(err, &upErr) { + reqLog.Info("gemini.cached_content.upstream_create_failed", zap.Int64("account_id", account.ID), zap.Int("status", upErr.StatusCode)) + if upErr.RetryableOnOtherAccount() { + lastUpstreamErr = upErr + excluded[account.ID] = struct{}{} + continue + } + writeGeminiCachedContentUpstreamError(c, upErr) + return + } + // 上游已返回成功但响应不可用：不换号重建，能识别出资源名时删除上游缓存。 + var malformed *service.GeminiCachedContentMalformedResponseError + if errors.As(err, &malformed) { + reqLog.Error("gemini.cached_content.upstream_create_malformed", zap.Int64("account_id", account.ID), zap.Error(err)) + if malformed.Name != "" { + h.cleanupGeminiCachedContentUpstream(reqLog, account, malformed.Name) + } + googleError(c, http.StatusBadGateway, "Upstream returned an invalid cached content response") + return + } + if failoverClientGone(c) { + return + } + reqLog.Warn("gemini.cached_content.upstream_create_error", zap.Int64("account_id", account.ID), zap.Error(err)) + excluded[account.ID] = struct{}{} + continue + } + + // 缓存 token 数在上游创建后才确定：按确定金额校验余额 / 额度并预留，不足时删除上游缓存， + // 只收取已发生的缓存 token 输入费。 + tokens := upstream.TotalTokenCount + storageTokenHours := service.GeminiCachedContentStorageTokenHours(tokens, ttl) + amount := h.gatewayService.QuoteUsageCost(c.Request.Context(), &service.RecordUsageInput{ + Result: &service.ForwardResult{ + Model: modelName, UpstreamModel: mappedModel, + Usage: service.ClaudeUsage{InputTokens: int(tokens)}, CacheStorageTokenHours: storageTokenHours, + }, + APIKey: apiKey, User: apiKey.User, Subscription: subscription, PricingAt: pricingAt, ChannelUsageFields: usageFields, + }) + inflightRelease() + chargeDone, err := h.reserveGeminiCachedContentCharge(c, apiKey, subscription, amount) + if err != nil { + reqLog.Info("gemini.cached_content.storage_charge_rejected", zap.Float64("amount", amount), zap.Error(err)) + h.cleanupGeminiCachedContentUpstream(reqLog, account, upstream.Name) + h.submitGeminiCachedContentUsage(c, apiKey, account, subscription, pricingAt, + "create:"+publicID, modelName, mappedModel, int(tokens), 0, usageFields, body) + writeGeminiCachedContentBillingError(c, err) + return + } + defer chargeDone() + + expireTime := service.ClampGeminiCachedContentExpire(upstream.ExpireTime, svc.Now().Add(ttl)) + record := &service.GeminiCachedContent{ + PublicID: publicID, + UserID: authSubject.UserID, + APIKeyID: apiKey.ID, + GroupID: *apiKey.GroupID, + AccountID: account.ID, + UpstreamName: upstream.Name, + Model: modelName, + UpstreamModel: upstreamModel, + DisplayName: req.DisplayName, + TotalTokenCount: tokens, + ExpireTime: expireTime, + } + persistCtx, cancelPersist := geminiCachedContentPersistContext(c) + err = svc.Save(persistCtx, record) + cancelPersist() + if err != nil { + reqLog.Error("gemini.cached_content.save_failed", zap.Int64("account_id", account.ID), zap.Error(err)) + h.cleanupGeminiCachedContentUpstream(reqLog, account, record.UpstreamName) + googleError(c, http.StatusInternalServerError, "Failed to persist cached content") + return + } + + h.submitGeminiCachedContentUsage(c, apiKey, account, subscription, pricingAt, + "create:"+publicID, modelName, mappedModel, int(tokens), storageTokenHours, usageFields, body) + c.JSON(http.StatusOK, service.GeminiCachedContentView(record)) + return + } + + if lastUpstreamErr != nil { + writeGeminiCachedContentUpstreamError(c, lastUpstreamErr) + return + } + if ttlRejectedLimit > 0 { + googleError(c, http.StatusBadRequest, fmt.Sprintf("ttl exceeds the maximum allowed %ds by the upstream accounts", int64(ttlRejectedLimit.Seconds()))) + return + } + markOpsRoutingCapacityLimited(c) + googleError(c, http.StatusServiceUnavailable, "No available Gemini accounts support explicit context caching") +} + +// GeminiCachedContentsList GET /v1beta/cachedContents +// 只返回当前 API Key 在当前分组下仍有效的缓存，不透传上游 list（同一账号上其他租户的缓存不可见）。 +func (h *GatewayHandler) GeminiCachedContentsList(c *gin.Context) { + apiKey, _, reqLog, ok := h.geminiCachedContentsPreamble(c, "handler.gemini_v1beta.cached_contents.list") + if !ok { + return + } + pageSize := service.GeminiCachedContentListDefaultPageSize + if raw := strings.TrimSpace(c.Query("pageSize")); raw != "" { + n, err := strconv.Atoi(raw) + if err != nil || n < 0 { + googleError(c, http.StatusBadRequest, "pageSize must be a non-negative integer") + return + } + if n > 0 { + pageSize = min(n, service.GeminiCachedContentListMaxPageSize) + } + } + beforeID := int64(0) + if raw := strings.TrimSpace(c.Query("pageToken")); raw != "" { + id, ok := decodeGeminiCachedContentPageToken(raw) + if !ok { + googleError(c, http.StatusBadRequest, "Invalid pageToken") + return + } + beforeID = id + } + records, err := h.geminiCachedContentService.List(c.Request.Context(), apiKey.ID, *apiKey.GroupID, beforeID, pageSize+1) + if err != nil { + reqLog.Error("gemini.cached_content.list_failed", zap.Error(err)) + googleError(c, http.StatusInternalServerError, "Failed to list cached contents") + return + } + resp := gin.H{} + if len(records) > pageSize { + records = records[:pageSize] + resp["nextPageToken"] = encodeGeminiCachedContentPageToken(records[len(records)-1].ID) + } + if len(records) > 0 { + views := make([]map[string]any, 0, len(records)) + for _, record := range records { + views = append(views, service.GeminiCachedContentView(record)) + } + resp["cachedContents"] = views + } + c.JSON(http.StatusOK, resp) +} + +// GeminiCachedContentsGet GET /v1beta/cachedContents/:cacheID +func (h *GatewayHandler) GeminiCachedContentsGet(c *gin.Context) { + apiKey, _, reqLog, ok := h.geminiCachedContentsPreamble(c, "handler.gemini_v1beta.cached_contents.get") + if !ok { + return + } + record, ok := h.resolveGeminiCachedContentOrWrite(c, reqLog, apiKey, c.Param("cacheID")) + if !ok { + return + } + c.JSON(http.StatusOK, service.GeminiCachedContentView(record)) +} + +// GeminiCachedContentsPatch PATCH /v1beta/cachedContents/:cacheID +// 只允许修改有效期；延长部分在本次请求内按增量时长收取存储费，缩短不退费。 +func (h *GatewayHandler) GeminiCachedContentsPatch(c *gin.Context) { + apiKey, _, reqLog, ok := h.geminiCachedContentsPreamble(c, "handler.gemini_v1beta.cached_contents.patch") + if !ok { + return + } + svc := h.geminiCachedContentService + body, ok := readGeminiCachedContentBody(c) + if !ok { + return + } + record, ok := h.resolveGeminiCachedContentOrWrite(c, reqLog, apiKey, c.Param("cacheID")) + if !ok { + return + } + account, err := svc.BoundAccount(c.Request.Context(), record) + if err != nil { + reqLog.Error("gemini.cached_content.load_account_failed", zap.Int64("account_id", record.AccountID), zap.Error(err)) + googleError(c, http.StatusInternalServerError, "Failed to load cached content account") + return + } + if account == nil { + googleError(c, http.StatusServiceUnavailable, geminiCachedContentUnavailableMessage) + return + } + ttl, err := service.ParseGeminiCachedContentPatchRequest(body, c.Query("updateMask"), svc.Now(), svc.MaxTTLForAccount(account)) + if err != nil { + googleError(c, http.StatusBadRequest, err.Error()) + return + } + setOpsRequestContext(c, record.Model, false) + pricingCtx, pricingAt := service.WithGatewayTokenRequestPricing(c.Request.Context()) + c.Request = c.Request.WithContext(pricingCtx) + + subscription, _ := middleware.GetSubscriptionFromContext(c) + if extension := svc.Now().Add(ttl).Sub(record.ExpireTime); extension > 0 { + usage := &service.RecordUsageInput{ + Result: &service.ForwardResult{ + Model: record.Model, UpstreamModel: record.UpstreamModelID(), + CacheStorageTokenHours: service.GeminiCachedContentStorageTokenHours(record.TotalTokenCount, extension), + }, + APIKey: apiKey, User: apiKey.User, Subscription: subscription, PricingAt: pricingAt, + } + if !h.gatewayService.CacheStoragePriced(c.Request.Context(), usage) { + googleError(c, http.StatusBadRequest, fmt.Sprintf("Explicit context caching is not available for model %s: storage pricing is not configured", record.Model)) + return + } + if !h.checkGeminiCachedContentBilling(c, reqLog, apiKey, subscription) { + return + } + chargeDone, err := h.reserveGeminiCachedContentCharge(c, apiKey, subscription, h.gatewayService.QuoteUsageCost(c.Request.Context(), usage)) + if err != nil { + reqLog.Info("gemini.cached_content.storage_charge_rejected", zap.Error(err)) + writeGeminiCachedContentBillingError(c, err) + return + } + defer chargeDone() + } + + setOpsSelectedAccount(c, account.ID, account.Platform) + upstream, err := svc.PatchUpstream(c.Request.Context(), account, record, ttl) + if err != nil { + var upErr *service.GeminiCachedContentUpstreamError + if errors.As(err, &upErr) { + if upErr.NotFound() { + if forgetErr := svc.Forget(c.Request.Context(), record); forgetErr != nil { + reqLog.Warn("gemini.cached_content.forget_failed", zap.Error(forgetErr)) + } + } + writeGeminiCachedContentUpstreamError(c, upErr) + return + } + reqLog.Warn("gemini.cached_content.upstream_patch_error", zap.Int64("account_id", account.ID), zap.Error(err)) + googleError(c, http.StatusBadGateway, "Upstream request failed") + return + } + expireTime := service.ClampGeminiCachedContentExpire(upstream.ExpireTime, svc.Now().Add(ttl)) + persistCtx, cancelPersist := geminiCachedContentPersistContext(c) + previousExpire, err := svc.UpdateExpireTime(persistCtx, record, expireTime) + cancelPersist() + if err != nil { + reqLog.Error("gemini.cached_content.update_expire_failed", zap.Int64("account_id", account.ID), zap.Error(err)) + if errors.Is(err, service.ErrGeminiCachedContentNotFound) { + googleError(c, http.StatusForbidden, service.GeminiCachedContentNotFoundMessage) + return + } + googleError(c, http.StatusInternalServerError, "Failed to persist cached content expiration") + return + } + // 按库中实际推进的时长收费：并发延长各自只为自己推进的部分付费。 + if extended := expireTime.Sub(previousExpire); extended > 0 { + h.submitGeminiCachedContentUsage(c, apiKey, account, subscription, pricingAt, + "patch:"+record.PublicID+":"+uuid.NewString(), record.Model, record.UpstreamModelID(), 0, + service.GeminiCachedContentStorageTokenHours(record.TotalTokenCount, extended), service.ChannelUsageFields{}, body) + } + c.JSON(http.StatusOK, service.GeminiCachedContentView(record)) +} + +// GeminiCachedContentsDelete DELETE /v1beta/cachedContents/:cacheID +// 删除不退还已收取的存储费。 +func (h *GatewayHandler) GeminiCachedContentsDelete(c *gin.Context) { + apiKey, _, reqLog, ok := h.geminiCachedContentsPreamble(c, "handler.gemini_v1beta.cached_contents.delete") + if !ok { + return + } + svc := h.geminiCachedContentService + record, ok := h.resolveGeminiCachedContentOrWrite(c, reqLog, apiKey, c.Param("cacheID")) + if !ok { + return + } + account, err := svc.BoundAccount(c.Request.Context(), record) + if err != nil { + reqLog.Error("gemini.cached_content.load_account_failed", zap.Int64("account_id", record.AccountID), zap.Error(err)) + googleError(c, http.StatusInternalServerError, "Failed to load cached content account") + return + } + if account != nil { + setOpsSelectedAccount(c, account.ID, account.Platform) + if err := svc.DeleteUpstream(c.Request.Context(), account, record); err != nil { + var upErr *service.GeminiCachedContentUpstreamError + if !errors.As(err, &upErr) { + reqLog.Warn("gemini.cached_content.upstream_delete_error", zap.Int64("account_id", account.ID), zap.Error(err)) + googleError(c, http.StatusBadGateway, "Upstream request failed") + return + } + if !upErr.NotFound() { + writeGeminiCachedContentUpstreamError(c, upErr) + return + } + } + } + persistCtx, cancelPersist := geminiCachedContentPersistContext(c) + err = svc.Forget(persistCtx, record) + cancelPersist() + if err != nil { + reqLog.Error("gemini.cached_content.forget_failed", zap.Error(err)) + googleError(c, http.StatusInternalServerError, "Failed to delete cached content") + return + } + c.JSON(http.StatusOK, gin.H{}) +} + +func (h *GatewayHandler) resolveGeminiCachedContentOrWrite(c *gin.Context, reqLog *zap.Logger, apiKey *service.APIKey, cacheID string) (*service.GeminiCachedContent, bool) { + record, err := h.geminiCachedContentService.Resolve(c.Request.Context(), apiKey.ID, apiKey.GroupID, cacheID) + if err != nil { + if errors.Is(err, service.ErrGeminiCachedContentNotFound) { + googleError(c, http.StatusForbidden, service.GeminiCachedContentNotFoundMessage) + return nil, false + } + reqLog.Error("gemini.cached_content.lookup_failed", zap.Error(err)) + googleError(c, http.StatusInternalServerError, "Failed to load cached content") + return nil, false + } + return record, true +} + +// resolveGeminiCachedContentBinding 处理 generateContent / countTokens 请求对显式缓存的引用： +// 校验归属与模型后把引用改写为上游资源名，返回绑定记录；请求未引用缓存时返回 nil。 +// 返回 false 表示已写出错误响应。 +func (h *GatewayHandler) resolveGeminiCachedContentBinding(c *gin.Context, reqLog *zap.Logger, apiKey *service.APIKey, modelName string, body []byte) (*service.GeminiCachedContent, []byte, bool) { + ref, err := service.FindGeminiCachedContentReference(body) + if err != nil { + googleError(c, http.StatusBadRequest, err.Error()) + return nil, nil, false + } + if ref == nil { + return nil, body, true + } + if middleware.HasForcePlatform(c) || h.geminiCachedContentService == nil || + apiKey.Group == nil || apiKey.Group.Platform == service.PlatformComposite { + googleError(c, http.StatusBadRequest, "cachedContent is not supported on this route") + return nil, nil, false + } + record, err := h.geminiCachedContentService.Resolve(c.Request.Context(), apiKey.ID, apiKey.GroupID, ref.Name) + if err != nil { + if errors.Is(err, service.ErrGeminiCachedContentNotFound) { + googleError(c, http.StatusForbidden, service.GeminiCachedContentNotFoundMessage) + return nil, nil, false + } + reqLog.Error("gemini.cached_content.lookup_failed", zap.Error(err)) + googleError(c, http.StatusInternalServerError, "Failed to load cached content") + return nil, nil, false + } + if record.Model != modelName { + googleError(c, http.StatusBadRequest, fmt.Sprintf( + "Model used by GenerateContent request (models/%s) and CachedContent (models/%s) has to be the same.", modelName, record.Model)) + return nil, nil, false + } + rewritten, err := service.RewriteGeminiCachedContentReference(body, ref, record.UpstreamName) + if err != nil { + googleError(c, http.StatusBadRequest, "Invalid cachedContent reference") + return nil, nil, false + } + return record, rewritten, true +} + +// acquireGeminiCachedContentAccountSlot 获取选中账号的并发槽（含等待计划）；返回 false 表示已写出错误响应。 +func (h *GatewayHandler) acquireGeminiCachedContentAccountSlot(c *gin.Context, reqLog *zap.Logger, helper *ConcurrencyHelper, selection *service.AccountSelectionResult) (func(), bool) { + if selection.Acquired { + release := selection.ReleaseFunc + if release == nil { + release = func() {} + } + return wrapReleaseOnDone(c.Request.Context(), release), true + } + account := selection.Account + if selection.WaitPlan == nil { + markOpsRoutingCapacityLimited(c) + googleError(c, http.StatusServiceUnavailable, "No available Gemini accounts") + return nil, false + } + canWait, err := helper.IncrementAccountWaitCount(c.Request.Context(), account.ID, selection.WaitPlan.MaxWaiting) + if err != nil { + reqLog.Warn("gemini.cached_content.account_wait_counter_increment_failed", zap.Int64("account_id", account.ID), zap.Error(err)) + } else if !canWait { + googleError(c, http.StatusTooManyRequests, "Too many pending requests, please retry later") + return nil, false + } + streamStarted := false + release, acquireErr := helper.AcquireAccountSlotWithWaitTimeout(c, account.ID, selection.WaitPlan.MaxConcurrency, selection.WaitPlan.Timeout, false, &streamStarted) + if err == nil { + helper.DecrementAccountWaitCount(c.Request.Context(), account.ID) + } + if acquireErr != nil { + reqLog.Warn("gemini.cached_content.account_slot_acquire_failed", zap.Int64("account_id", account.ID), zap.Error(acquireErr)) + googleConcurrencyError(c, acquireErr, "account") + return nil, false + } + if release == nil { + release = func() {} + } + return wrapReleaseOnDone(c.Request.Context(), release), true +} + +// submitGeminiCachedContentUsage 记录缓存创建 / 延长的计费：inputTokens 按输入价计，storageTokenHours 按计费模型的存储单价计。 +func (h *GatewayHandler) submitGeminiCachedContentUsage( + c *gin.Context, + apiKey *service.APIKey, + account *service.Account, + subscription *service.UserSubscription, + pricingAt time.Time, + requestSuffix string, + model string, + upstreamModel string, + inputTokens int, + storageTokenHours float64, + usageFields service.ChannelUsageFields, + payload []byte, +) { + if inputTokens <= 0 && storageTokenHours <= 0 { + return + } + if upstreamModel == model { + upstreamModel = "" + } + result := &service.ForwardResult{ + RequestID: "gemini-cached-content:" + requestSuffix, + Model: model, + UpstreamModel: upstreamModel, + Usage: service.ClaudeUsage{InputTokens: inputTokens}, + CacheStorageTokenHours: storageTokenHours, + } + inboundEndpoint := GetInboundEndpoint(c) + upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform) + userAgent := c.GetHeader("User-Agent") + clientIP := ip.GetClientIP(c) + requestPayloadHash := service.HashUsageRequestPayload(payload) + quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey) + h.submitMandatoryUsageRecordTask(c.Request.Context(), func(ctx context.Context) { + if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{ + Result: result, + APIKey: apiKey, + User: apiKey.User, + Account: account, + Subscription: subscription, + PricingAt: pricingAt, + InboundEndpoint: inboundEndpoint, + UpstreamEndpoint: upstreamEndpoint, + UserAgent: userAgent, + IPAddress: clientIP, + RequestPayloadHash: requestPayloadHash, + APIKeyService: h.apiKeyService, + QuotaPlatform: quotaPlatform, + ChannelUsageFields: usageFields, + }); err != nil { + logger.L().With( + zap.String("component", "handler.gemini_v1beta.cached_contents"), + zap.Int64("api_key_id", apiKey.ID), + zap.Int64("account_id", account.ID), + zap.String("request_id", result.RequestID), + ).Error("gemini.cached_content.record_usage_failed", zap.Error(err)) + } + }) +} + +// geminiCachedContentPersistContext 返回与客户端连接脱钩的落库 context：上游操作已生效时，客户端断开不应让本地记录缺失。 +func geminiCachedContentPersistContext(c *gin.Context) (context.Context, context.CancelFunc) { + return context.WithTimeout(context.WithoutCancel(c.Request.Context()), geminiCachedContentPersistTimeout) +} + +// cleanupGeminiCachedContentUpstream 尽力删除已在上游创建、但网关不会登记的缓存，避免其持续产生存储费。 +func (h *GatewayHandler) cleanupGeminiCachedContentUpstream(reqLog *zap.Logger, account *service.Account, upstreamName string) { + ctx, cancel := context.WithTimeout(context.Background(), geminiCachedContentPersistTimeout) + defer cancel() + if err := h.geminiCachedContentService.DeleteUpstream(ctx, account, &service.GeminiCachedContent{UpstreamName: upstreamName}); err != nil { + reqLog.Error("gemini.cached_content.orphan_cleanup_failed", zap.Int64("account_id", account.ID), zap.String("upstream_name", upstreamName), zap.Error(err)) + } +} + +// reserveGeminiCachedContentCharge 校验余额 / 订阅额度 / API Key 额度足以支付确定金额 amount（倍率后）， +// 并按该金额登记在途预留（不受预留上限封顶）。返回的 done 须在 handler 结束时调用。 +func (h *GatewayHandler) reserveGeminiCachedContentCharge(c *gin.Context, apiKey *service.APIKey, subscription *service.UserSubscription, amount float64) (func(), error) { + ctx := c.Request.Context() + if err := h.billingCacheService.CheckChargeAffordable(ctx, apiKey.User, apiKey, apiKey.Group, subscription, amount); err != nil { + return inflightNoop, err + } + res, err := h.billingCacheService.ReserveInflightCharge(ctx, apiKey.User, apiKey.Group, subscription, amount) + if err != nil { + return inflightNoop, err + } + if res == nil { + return inflightNoop, nil + } + c.Request = c.Request.WithContext(service.WithInflightReservation(ctx, res)) + return res.HandlerDone, nil +} + +func writeGeminiCachedContentBillingError(c *gin.Context, err error) { + status, _, message, retryAfter := billingErrorDetails(err) + if retryAfter > 0 { + c.Header("Retry-After", strconv.Itoa(retryAfter)) + } + googleError(c, status, message) +} + +// writeGeminiCachedContentUpstreamError 透传上游错误：缓存不存在统一返回标准 403（各上游形态不同且可能含上游资源 ID）； +// 账号凭据或权限类错误不是客户端问题，统一改为 502。 +func writeGeminiCachedContentUpstreamError(c *gin.Context, upErr *service.GeminiCachedContentUpstreamError) { + status := upErr.StatusCode + message := strings.TrimSpace(extractGoogleErrorMessage(upErr.Body)) + service.SetOpsUpstreamError(c, status, message, "") + if upErr.NotFound() { + c.Data(http.StatusForbidden, "application/json", []byte(service.GeminiCachedContentNotFoundResponse)) + return + } + if status == http.StatusUnauthorized || status == http.StatusForbidden { + googleError(c, http.StatusBadGateway, "Upstream account rejected the cached content request") + return + } + if len(upErr.Body) == 0 || !json.Valid(upErr.Body) { + if message == "" { + message = http.StatusText(status) + } + googleError(c, status, message) + return + } + c.Data(status, "application/json", upErr.Body) +} + +func extractGoogleErrorMessage(body []byte) string { + var envelope struct { + Error struct { + Message string `json:"message"` + } `json:"error"` + } + if err := json.Unmarshal(body, &envelope); err != nil { + return "" + } + return envelope.Error.Message +} + +func encodeGeminiCachedContentPageToken(id int64) string { + return base64.RawURLEncoding.EncodeToString([]byte(strconv.FormatInt(id, 10))) +} + +func decodeGeminiCachedContentPageToken(raw string) (int64, bool) { + decoded, err := base64.RawURLEncoding.DecodeString(raw) + if err != nil { + return 0, false + } + id, err := strconv.ParseInt(string(decoded), 10, 64) + if err != nil || id <= 0 { + return 0, false + } + return id, true +} + +// mergeAccountIDSets 返回两个账号 ID 集合的并集（不修改入参）。 +func mergeAccountIDSets(a, b map[int64]struct{}) map[int64]struct{} { + out := make(map[int64]struct{}, len(a)+len(b)) + for id := range a { + out[id] = struct{}{} + } + for id := range b { + out[id] = struct{}{} + } + return out +} + +// handleGeminiCachedContentFailoverExhausted 处理绑定缓存的请求在持有账号上重试耗尽：没有其他账号可换， +// 上游的限流 / 不可用响应原样返回，客户端据此判断缓存所在账号的状态；账号凭据类错误仍按网关映射。 +func (h *GatewayHandler) handleGeminiCachedContentFailoverExhausted(c *gin.Context, failoverErr *service.UpstreamFailoverError) { + if failoverErr == nil || len(failoverErr.ResponseBody) == 0 || !json.Valid(failoverErr.ResponseBody) || + (failoverErr.StatusCode != http.StatusTooManyRequests && failoverErr.StatusCode < http.StatusInternalServerError) { + h.handleGeminiFailoverExhausted(c, failoverErr) + return + } + service.SetOpsUpstreamError(c, failoverErr.StatusCode, service.ExtractUpstreamErrorMessage(failoverErr.ResponseBody), "") + c.Data(failoverErr.StatusCode, "application/json", failoverErr.ResponseBody) +} diff --git a/backend/internal/handler/gemini_cached_contents_handler_test.go b/backend/internal/handler/gemini_cached_contents_handler_test.go new file mode 100644 index 0000000000..8edbcc86e9 --- /dev/null +++ b/backend/internal/handler/gemini_cached_contents_handler_test.go @@ -0,0 +1,220 @@ +//go:build unit + +package handler + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + middleware "github.com/Wei-Shaw/sub2api/internal/server/middleware" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "go.uber.org/zap" +) + +type geminiCachedContentHandlerRepoStub struct { + records []*service.GeminiCachedContent +} + +func (r *geminiCachedContentHandlerRepoStub) Create(_ context.Context, record *service.GeminiCachedContent) error { + record.ID = int64(len(r.records) + 1) + r.records = append(r.records, record) + return nil +} + +func (r *geminiCachedContentHandlerRepoStub) GetForOwner(_ context.Context, apiKeyID int64, publicID string) (*service.GeminiCachedContent, error) { + for _, record := range r.records { + if record.PublicID == publicID && record.APIKeyID == apiKeyID { + return record, nil + } + } + return nil, service.ErrGeminiCachedContentNotFound +} + +func (r *geminiCachedContentHandlerRepoStub) ListForOwner(_ context.Context, apiKeyID, groupID int64, activeAt time.Time, beforeID int64, limit int) ([]*service.GeminiCachedContent, error) { + var out []*service.GeminiCachedContent + for i := len(r.records) - 1; i >= 0; i-- { + record := r.records[i] + if record.APIKeyID == apiKeyID && record.GroupID == groupID && record.ExpireTime.After(activeAt) && (beforeID <= 0 || record.ID < beforeID) { + out = append(out, record) + } + if len(out) == limit { + break + } + } + return out, nil +} + +func (r *geminiCachedContentHandlerRepoStub) UpdateExpireTime(context.Context, int64, time.Time) (time.Time, error) { + return time.Time{}, nil +} + +func (r *geminiCachedContentHandlerRepoStub) SoftDelete(context.Context, int64) error { + return nil +} + +func newGeminiCachedContentHandlerForTest(repo *geminiCachedContentHandlerRepoStub) *GatewayHandler { + return &GatewayHandler{geminiCachedContentService: service.NewGeminiCachedContentService(repo, nil, &config.Config{})} +} + +func geminiCachedContentTestAPIKey(platform string) *service.APIKey { + groupID := int64(42) + return &service.APIKey{ID: 7, UserID: 3, GroupID: &groupID, Group: &service.Group{ID: groupID, Platform: platform}} +} + +func newGeminiCachedContentTestContext(method, target string, apiKey *service.APIKey) (*gin.Context, *httptest.ResponseRecorder) { + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(method, target, nil) + c.Set(string(middleware.ContextKeyAPIKey), apiKey) + c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: apiKey.UserID, Concurrency: 10}) + return c, rec +} + +func seedGeminiCachedContent(repo *geminiCachedContentHandlerRepoStub, publicID string, apiKeyID int64, expire time.Time) *service.GeminiCachedContent { + record := &service.GeminiCachedContent{ + PublicID: publicID, + UserID: 3, + APIKeyID: apiKeyID, + GroupID: 42, + AccountID: 900, + UpstreamName: "cachedContents/upstream" + publicID, + Model: "gemini-3.8-flash", + UpstreamModel: "models/gemini-3.8-flash", + TotalTokenCount: 6273, + ExpireTime: expire, + CreatedAt: time.Date(2026, 10, 9, 8, 0, 0, 0, time.UTC), + UpdatedAt: time.Date(2026, 10, 9, 8, 0, 0, 0, time.UTC), + } + _ = repo.Create(context.Background(), record) + return record +} + +func TestResolveGeminiCachedContentBinding(t *testing.T) { + repo := &geminiCachedContentHandlerRepoStub{} + seedGeminiCachedContent(repo, "abc", 7, time.Now().Add(time.Hour)) + seedGeminiCachedContent(repo, "foreign", 8, time.Now().Add(time.Hour)) + h := newGeminiCachedContentHandlerForTest(repo) + apiKey := geminiCachedContentTestAPIKey(service.PlatformGemini) + + t.Run("request without cache passes through", func(t *testing.T) { + c, _ := newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", apiKey) + body := []byte(`{"contents":[]}`) + record, out, ok := h.resolveGeminiCachedContentBinding(c, zap.NewNop(), apiKey, "gemini-3.8-flash", body) + require.True(t, ok) + require.Nil(t, record) + require.Equal(t, body, out) + }) + + t.Run("owned cache rewrites reference and binds account", func(t *testing.T) { + c, _ := newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", apiKey) + record, out, ok := h.resolveGeminiCachedContentBinding(c, zap.NewNop(), apiKey, "gemini-3.8-flash", + []byte(`{"cachedContent":"cachedContents/abc","contents":[]}`)) + require.True(t, ok) + require.Equal(t, int64(900), record.AccountID) + require.JSONEq(t, `{"cachedContent":"cachedContents/upstreamabc","contents":[]}`, string(out)) + }) + + t.Run("cache owned by another key looks missing", func(t *testing.T) { + c, rec := newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", apiKey) + _, _, ok := h.resolveGeminiCachedContentBinding(c, zap.NewNop(), apiKey, "gemini-3.8-flash", + []byte(`{"cachedContent":"cachedContents/foreign","contents":[]}`)) + require.False(t, ok) + require.Equal(t, http.StatusForbidden, rec.Code) + require.Contains(t, rec.Body.String(), service.GeminiCachedContentNotFoundMessage) + }) + + t.Run("model must match the cache literally", func(t *testing.T) { + c, rec := newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-flash-latest:generateContent", apiKey) + _, _, ok := h.resolveGeminiCachedContentBinding(c, zap.NewNop(), apiKey, "gemini-flash-latest", + []byte(`{"cachedContent":"cachedContents/abc","contents":[]}`)) + require.False(t, ok) + require.Equal(t, http.StatusBadRequest, rec.Code) + require.Contains(t, rec.Body.String(), "has to be the same") + }) + + t.Run("composite groups are rejected", func(t *testing.T) { + composite := geminiCachedContentTestAPIKey(service.PlatformComposite) + c, rec := newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", composite) + _, _, ok := h.resolveGeminiCachedContentBinding(c, zap.NewNop(), composite, "gemini-3.8-flash", + []byte(`{"cachedContent":"cachedContents/abc","contents":[]}`)) + require.False(t, ok) + require.Equal(t, http.StatusBadRequest, rec.Code) + }) +} + +func TestGeminiCachedContentsListAndGet(t *testing.T) { + repo := &geminiCachedContentHandlerRepoStub{} + seedGeminiCachedContent(repo, "first", 7, time.Now().Add(time.Hour)) + seedGeminiCachedContent(repo, "second", 7, time.Now().Add(time.Hour)) + seedGeminiCachedContent(repo, "expired", 7, time.Now().Add(-time.Minute)) + seedGeminiCachedContent(repo, "foreign", 8, time.Now().Add(time.Hour)) + h := newGeminiCachedContentHandlerForTest(repo) + apiKey := geminiCachedContentTestAPIKey(service.PlatformGemini) + + c, rec := newGeminiCachedContentTestContext(http.MethodGet, "/v1beta/cachedContents?pageSize=1", apiKey) + h.GeminiCachedContentsList(c) + require.Equal(t, http.StatusOK, rec.Code) + var page struct { + CachedContents []map[string]any `json:"cachedContents"` + NextPageToken string `json:"nextPageToken"` + } + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &page)) + require.Len(t, page.CachedContents, 1) + require.Equal(t, "cachedContents/second", page.CachedContents[0]["name"]) + require.NotEmpty(t, page.NextPageToken) + + c, rec = newGeminiCachedContentTestContext(http.MethodGet, "/v1beta/cachedContents?pageSize=1&pageToken="+page.NextPageToken, apiKey) + h.GeminiCachedContentsList(c) + require.Equal(t, http.StatusOK, rec.Code) + page.NextPageToken = "" + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &page)) + require.Len(t, page.CachedContents, 1) + require.Equal(t, "cachedContents/first", page.CachedContents[0]["name"]) + require.Empty(t, page.NextPageToken, "过期与其他 Key 的缓存不出现在列表里") + + c, rec = newGeminiCachedContentTestContext(http.MethodGet, "/v1beta/cachedContents/first", apiKey) + c.Params = gin.Params{{Key: "cacheID", Value: "first"}} + h.GeminiCachedContentsGet(c) + require.Equal(t, http.StatusOK, rec.Code) + var view map[string]any + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &view)) + require.Equal(t, "cachedContents/first", view["name"]) + require.Equal(t, "models/gemini-3.8-flash", view["model"]) + require.NotContains(t, rec.Body.String(), "upstreamfirst", "上游资源名不暴露给客户端") + + c, rec = newGeminiCachedContentTestContext(http.MethodGet, "/v1beta/cachedContents/foreign", apiKey) + c.Params = gin.Params{{Key: "cacheID", Value: "foreign"}} + h.GeminiCachedContentsGet(c) + require.Equal(t, http.StatusForbidden, rec.Code) + require.True(t, strings.Contains(rec.Body.String(), service.GeminiCachedContentNotFoundMessage)) +} + +func TestHandleGeminiCachedContentFailoverExhausted(t *testing.T) { + h := &GatewayHandler{} + apiKey := geminiCachedContentTestAPIKey(service.PlatformGemini) + unavailable := `{"error":{"code":503,"message":"The account holding this cached content is temporarily unavailable","status":"UNAVAILABLE"}}` + + c, rec := newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", apiKey) + h.handleGeminiCachedContentFailoverExhausted(c, &service.UpstreamFailoverError{StatusCode: http.StatusServiceUnavailable, ResponseBody: []byte(unavailable)}) + require.Equal(t, http.StatusServiceUnavailable, rec.Code, "绑定缓存的请求无账号可换，上游 503 原样返回") + require.JSONEq(t, unavailable, rec.Body.String()) + + c, rec = newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", apiKey) + h.handleGeminiCachedContentFailoverExhausted(c, &service.UpstreamFailoverError{StatusCode: http.StatusTooManyRequests, ResponseBody: []byte(`{"error":{"code":429,"message":"Too many pending requests, please retry later","status":"RESOURCE_EXHAUSTED"}}`)}) + require.Equal(t, http.StatusTooManyRequests, rec.Code) + require.Contains(t, rec.Body.String(), "Too many pending requests") + + c, rec = newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", apiKey) + h.handleGeminiCachedContentFailoverExhausted(c, &service.UpstreamFailoverError{StatusCode: http.StatusUnauthorized, ResponseBody: []byte(`{"error":{"code":401,"message":"bad key"}}`)}) + require.Equal(t, http.StatusBadGateway, rec.Code, "账号凭据类错误仍按网关映射，不透传给客户端") + require.NotContains(t, rec.Body.String(), "bad key") +} diff --git a/backend/internal/handler/gemini_v1beta_handler.go b/backend/internal/handler/gemini_v1beta_handler.go index 4e56c37232..18ac84084f 100644 --- a/backend/internal/handler/gemini_v1beta_handler.go +++ b/backend/internal/handler/gemini_v1beta_handler.go @@ -356,6 +356,12 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) { modelName = channelMapping.MappedModel } + // 引用显式缓存的请求只能发往持有该缓存的账号，且不参与会话粘性与换号。 + cachedContentBinding, body, ok := h.resolveGeminiCachedContentBinding(c, reqLog, apiKey, modelName, body) + if !ok { + return + } + // Get subscription (may be nil) subscription, _ := middleware.GetSubscriptionFromContext(c) @@ -405,8 +411,11 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) { // 3) select account (sticky session based on request body) // 优先使用 Gemini CLI 的会话标识（privileged-user-id + tmp 目录哈希） - sessionHash := extractGeminiCLISessionHash(c, body) - if sessionHash == "" { + sessionHash := "" + if cachedContentBinding == nil { + sessionHash = extractGeminiCLISessionHash(c, body) + } + if sessionHash == "" && cachedContentBinding == nil { // Fallback: 使用通用的会话哈希生成逻辑（适用于其他客户端） parsedReq, _ := service.ParseGatewayRequest(service.NewRequestBodyRef(body), domain.PlatformGemini) if parsedReq != nil { @@ -443,7 +452,7 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) { var geminiPrefixHash string var geminiSessionUUID string var matchedDigestChain string - useDigestFallback := sessionBoundAccountID == 0 + useDigestFallback := sessionBoundAccountID == 0 && cachedContentBinding == nil if useDigestFallback { // 解析 Gemini 请求体 @@ -507,7 +516,21 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) { hasBoundSession := sessionKey != "" && sessionBoundAccountID > 0 cleanedForUnknownBinding := false - fs := NewFailoverState(h.maxAccountSwitchesGemini, hasBoundSession) + maxAccountSwitches := h.maxAccountSwitchesGemini + var cachedContentExclusions map[int64]struct{} + if cachedContentBinding != nil { + // 缓存只存在于绑定账号：排除分组内其余账号，失败不换号。 + maxAccountSwitches = 0 + sessionBoundAccountID = cachedContentBinding.AccountID + c.Request = c.Request.WithContext(service.WithGeminiCachedContentBound(c.Request.Context())) + cachedContentExclusions, err = h.gatewayService.GeminiExplicitCacheExclusions(c.Request.Context(), apiKey.GroupID, cachedContentBinding.AccountID) + if err != nil { + reqLog.Warn("gemini.cached_content.list_accounts_failed", zap.Error(err)) + googleError(c, http.StatusServiceUnavailable, geminiCachedContentUnavailableMessage) + return + } + } + fs := NewFailoverState(maxAccountSwitches, hasBoundSession) // 单账号分组提前设置 SingleAccountRetry 标记，让 Service 层首次 503 就不设模型限流标记。 // 避免单账号分组收到 503 (MODEL_CAPACITY_EXHAUSTED) 时设 29s 限流，导致后续请求连续快速失败。 @@ -517,12 +540,32 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) { } for { - selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, sessionKey, modelName, fs.FailedAccountIDs, "", int64(0)) // Gemini 不使用会话限制 + excludedAccountIDs := fs.FailedAccountIDs + if cachedContentExclusions != nil { + excludedAccountIDs = mergeAccountIDSets(cachedContentExclusions, fs.FailedAccountIDs) + } + selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, sessionKey, modelName, excludedAccountIDs, "", int64(0)) // Gemini 不使用会话限制 + if err == nil && cachedContentBinding != nil && selection.Account.ID != cachedContentBinding.AccountID { + if selection.Acquired && selection.ReleaseFunc != nil { + selection.ReleaseFunc() + } + err = service.ErrNoAvailableAccounts + } if err != nil { if failoverClientGone(c) { reqLog.Info("gemini.account_select_aborted_client_disconnected", zap.Error(err)) return } + if cachedContentBinding != nil { + if fs.LastFailoverErr != nil { + h.handleGeminiCachedContentFailoverExhausted(c, fs.LastFailoverErr) + return + } + reqLog.Info("gemini.cached_content.bound_account_unavailable", zap.Int64("account_id", cachedContentBinding.AccountID), zap.Error(err)) + markOpsRoutingCapacityLimited(c) + googleError(c, http.StatusServiceUnavailable, geminiCachedContentUnavailableMessage) + return + } if len(fs.FailedAccountIDs) == 0 { cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, modelName, modelName, service.PlatformGemini) if !cls.ModelNotFound { @@ -631,6 +674,11 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) { accountReleaseFunc() } reqLog.Debug("gemini.account_slot_profit_vetoed", zap.Int64("account_id", account.ID), zap.String("reason", reason)) + if cachedContentBinding != nil { + markOpsRoutingCapacityLimited(c) + googleError(c, http.StatusServiceUnavailable, geminiCachedContentUnavailableMessage) + return + } if fs.RecordProfitVeto(account.ID) == FailoverExhausted { reqLog.Warn("gemini.profit_veto_attempts_exhausted", zap.Int("profit_veto_count", fs.ProfitVetoCount())) markOpsRoutingCapacityLimited(c) @@ -684,6 +732,10 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) { case FailoverContinue: continue case FailoverExhausted: + if cachedContentBinding != nil { + h.handleGeminiCachedContentFailoverExhausted(c, fs.LastFailoverErr) + return + } h.handleGeminiFailoverExhausted(c, fs.LastFailoverErr) return case FailoverCanceled: diff --git a/backend/internal/handler/security_audit_order_test.go b/backend/internal/handler/security_audit_order_test.go index e9cb4e1f91..2e01829307 100644 --- a/backend/internal/handler/security_audit_order_test.go +++ b/backend/internal/handler/security_audit_order_test.go @@ -24,6 +24,7 @@ func TestPromptAuditGatePrecedesAccountBillingAndUpstreamSideEffects(t *testing. {file: "gateway_handler_chat_completions.go", function: "ChatCompletions", auditToken: "checkSecurityAudit"}, {file: "gateway_handler_responses.go", function: "Responses", auditToken: "checkSecurityAudit"}, {file: "gemini_v1beta_handler.go", function: "GeminiV1BetaModels", auditToken: "checkSecurityAudit"}, + {file: "gemini_cached_contents_handler.go", function: "GeminiCachedContentsCreate", auditToken: "checkSecurityAudit"}, {file: "openai_gateway_handler.go", function: "Responses", auditToken: "checkSecurityAudit"}, {file: "openai_gateway_handler.go", function: "Messages", auditToken: "checkSecurityAudit"}, {file: "openai_chat_completions.go", function: "ChatCompletions", auditToken: "checkSecurityAudit"}, diff --git a/backend/internal/handler/wire.go b/backend/internal/handler/wire.go index e87d440c33..bdd47f4265 100644 --- a/backend/internal/handler/wire.go +++ b/backend/internal/handler/wire.go @@ -113,11 +113,13 @@ func ProvideGatewayHandler( cfg *config.Config, settingService *service.SettingService, coordinator *securityaudit.Coordinator, + geminiCachedContentService *service.GeminiCachedContentService, ) *GatewayHandler { h := NewGatewayHandler(gatewayService, openAIGatewayService, geminiCompatService, antigravityGatewayService, userService, concurrencyService, billingCacheService, usageService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, userMsgQueueService, cfg, settingService) h.securityAuditCoordinator = coordinator + h.geminiCachedContentService = geminiCachedContentService return h } diff --git a/backend/internal/repository/gemini_cached_content_repo.go b/backend/internal/repository/gemini_cached_content_repo.go new file mode 100644 index 0000000000..d63929c3e6 --- /dev/null +++ b/backend/internal/repository/gemini_cached_content_repo.go @@ -0,0 +1,124 @@ +package repository + +import ( + "context" + "database/sql" + "time" + + "github.com/Wei-Shaw/sub2api/internal/service" +) + +type geminiCachedContentSQLExecutor interface { + ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) + QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error) + QueryRowContext(ctx context.Context, query string, args ...any) *sql.Row +} + +type geminiCachedContentRepository struct { + sql geminiCachedContentSQLExecutor +} + +func NewGeminiCachedContentRepository(db *sql.DB) service.GeminiCachedContentRepository { + return &geminiCachedContentRepository{sql: db} +} + +const geminiCachedContentColumns = `id, public_id, user_id, api_key_id, group_id, account_id, upstream_name, model, + upstream_model, display_name, total_token_count, expire_time, created_at, updated_at` + +func (r *geminiCachedContentRepository) Create(ctx context.Context, record *service.GeminiCachedContent) error { + row := r.sql.QueryRowContext(ctx, ` +INSERT INTO gemini_cached_contents (public_id, user_id, api_key_id, group_id, account_id, upstream_name, model, + upstream_model, display_name, total_token_count, expire_time) +VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11) +RETURNING id, created_at, updated_at`, + record.PublicID, record.UserID, record.APIKeyID, record.GroupID, record.AccountID, record.UpstreamName, + record.Model, record.UpstreamModel, record.DisplayName, record.TotalTokenCount, record.ExpireTime) + if err := row.Scan(&record.ID, &record.CreatedAt, &record.UpdatedAt); err != nil { + return translatePersistenceError(err, nil, service.ErrGeminiCachedContentExists) + } + return nil +} + +func (r *geminiCachedContentRepository) GetForOwner(ctx context.Context, apiKeyID int64, publicID string) (*service.GeminiCachedContent, error) { + record, err := scanGeminiCachedContent(r.sql.QueryRowContext(ctx, `SELECT `+geminiCachedContentColumns+` + FROM gemini_cached_contents + WHERE public_id = $1 AND api_key_id = $2 AND deleted_at IS NULL`, publicID, apiKeyID)) + if err != nil { + return nil, translatePersistenceError(err, service.ErrGeminiCachedContentNotFound, nil) + } + return record, nil +} + +func (r *geminiCachedContentRepository) ListForOwner(ctx context.Context, apiKeyID, groupID int64, activeAt time.Time, beforeID int64, limit int) ([]*service.GeminiCachedContent, error) { + if limit <= 0 { + limit = service.GeminiCachedContentListDefaultPageSize + } + rows, err := r.sql.QueryContext(ctx, `SELECT `+geminiCachedContentColumns+` + FROM gemini_cached_contents + WHERE api_key_id = $1 AND group_id = $2 AND deleted_at IS NULL AND expire_time > $3 AND ($4 <= 0 OR id < $4) + ORDER BY id DESC + LIMIT $5`, apiKeyID, groupID, activeAt, beforeID, limit) + if err != nil { + return nil, err + } + defer func() { _ = rows.Close() }() + var out []*service.GeminiCachedContent + for rows.Next() { + record, err := scanGeminiCachedContent(rows) + if err != nil { + return nil, err + } + out = append(out, record) + } + return out, rows.Err() +} + +func (r *geminiCachedContentRepository) UpdateExpireTime(ctx context.Context, id int64, expireTime time.Time) (time.Time, error) { + var previous time.Time + err := r.sql.QueryRowContext(ctx, `WITH prev AS ( + SELECT id, expire_time FROM gemini_cached_contents WHERE id = $1 AND deleted_at IS NULL FOR UPDATE +) +UPDATE gemini_cached_contents AS g + SET expire_time = $2, updated_at = NOW() + FROM prev + WHERE g.id = prev.id + RETURNING prev.expire_time`, id, expireTime).Scan(&previous) + if err != nil { + return time.Time{}, translatePersistenceError(err, service.ErrGeminiCachedContentNotFound, nil) + } + return previous, nil +} + +func (r *geminiCachedContentRepository) SoftDelete(ctx context.Context, id int64) error { + _, err := r.sql.ExecContext(ctx, `UPDATE gemini_cached_contents + SET deleted_at = NOW(), updated_at = NOW() + WHERE id = $1 AND deleted_at IS NULL`, id) + return err +} + +type geminiCachedContentScanner interface { + Scan(dest ...any) error +} + +func scanGeminiCachedContent(row geminiCachedContentScanner) (*service.GeminiCachedContent, error) { + var record service.GeminiCachedContent + if err := row.Scan( + &record.ID, + &record.PublicID, + &record.UserID, + &record.APIKeyID, + &record.GroupID, + &record.AccountID, + &record.UpstreamName, + &record.Model, + &record.UpstreamModel, + &record.DisplayName, + &record.TotalTokenCount, + &record.ExpireTime, + &record.CreatedAt, + &record.UpdatedAt, + ); err != nil { + return nil, err + } + return &record, nil +} diff --git a/backend/internal/repository/gemini_cached_content_repo_integration_test.go b/backend/internal/repository/gemini_cached_content_repo_integration_test.go new file mode 100644 index 0000000000..6f6f3cfe9d --- /dev/null +++ b/backend/internal/repository/gemini_cached_content_repo_integration_test.go @@ -0,0 +1,82 @@ +//go:build integration + +package repository + +import ( + "context" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/stretchr/testify/require" +) + +func TestGeminiCachedContentRepository_Lifecycle(t *testing.T) { + ctx := context.Background() + tx := testTx(t) + repo := &geminiCachedContentRepository{sql: tx} + now := time.Now().UTC().Truncate(time.Microsecond) + suffix := time.Now().Format("150405.000000") + apiKeyID := int64(910001) + + newRecord := func(publicID string, keyID, groupID int64, expire time.Time) *service.GeminiCachedContent { + return &service.GeminiCachedContent{ + PublicID: publicID + suffix, UserID: 1, APIKeyID: keyID, GroupID: groupID, AccountID: 99, + UpstreamName: "cachedContents/up" + publicID, Model: "gemini-3.8-flash", UpstreamModel: "models/gemini-3.8-flash", + DisplayName: "label", TotalTokenCount: 2730, ExpireTime: expire, + } + } + first := newRecord("a", apiKeyID, 42, now.Add(time.Hour)) + require.NoError(t, repo.Create(ctx, first)) + require.Positive(t, first.ID) + require.False(t, first.CreatedAt.IsZero()) + _, err := tx.ExecContext(ctx, `SAVEPOINT gemini_cached_content_duplicate`) + require.NoError(t, err) + require.ErrorIs(t, repo.Create(ctx, newRecord("a", apiKeyID, 42, now.Add(time.Hour))), service.ErrGeminiCachedContentExists) + _, err = tx.ExecContext(ctx, `ROLLBACK TO SAVEPOINT gemini_cached_content_duplicate`) + require.NoError(t, err) + + second := newRecord("b", apiKeyID, 42, now.Add(time.Hour)) + require.NoError(t, repo.Create(ctx, second)) + require.NoError(t, repo.Create(ctx, newRecord("expired", apiKeyID, 42, now.Add(-time.Minute)))) + require.NoError(t, repo.Create(ctx, newRecord("othergroup", apiKeyID, 43, now.Add(time.Hour)))) + require.NoError(t, repo.Create(ctx, newRecord("otherkey", apiKeyID+1, 42, now.Add(time.Hour)))) + + got, err := repo.GetForOwner(ctx, apiKeyID, first.PublicID) + require.NoError(t, err) + require.Equal(t, "cachedContents/upa", got.UpstreamName) + require.Equal(t, int64(2730), got.TotalTokenCount) + require.True(t, got.ExpireTime.Equal(first.ExpireTime)) + _, err = repo.GetForOwner(ctx, apiKeyID+1, first.PublicID) + require.ErrorIs(t, err, service.ErrGeminiCachedContentNotFound) + + list, err := repo.ListForOwner(ctx, apiKeyID, 42, now, 0, 10) + require.NoError(t, err) + require.Len(t, list, 2) + require.Equal(t, second.PublicID, list[0].PublicID) + page, err := repo.ListForOwner(ctx, apiKeyID, 42, now, list[0].ID, 10) + require.NoError(t, err) + require.Len(t, page, 1) + require.Equal(t, first.PublicID, page[0].PublicID) + + newExpire := now.Add(3 * time.Hour) + previous, err := repo.UpdateExpireTime(ctx, first.ID, newExpire) + require.NoError(t, err) + require.True(t, previous.Equal(first.ExpireTime), "返回更新前的到期时间") + got, err = repo.GetForOwner(ctx, apiKeyID, first.PublicID) + require.NoError(t, err) + require.True(t, got.ExpireTime.Equal(newExpire)) + previous, err = repo.UpdateExpireTime(ctx, first.ID, newExpire.Add(time.Hour)) + require.NoError(t, err) + require.True(t, previous.Equal(newExpire), "连续更新各自拿到自己的前值") + newExpire = newExpire.Add(time.Hour) + + require.NoError(t, repo.SoftDelete(ctx, first.ID)) + _, err = repo.GetForOwner(ctx, apiKeyID, first.PublicID) + require.ErrorIs(t, err, service.ErrGeminiCachedContentNotFound) + _, err = repo.UpdateExpireTime(ctx, first.ID, newExpire) + require.ErrorIs(t, err, service.ErrGeminiCachedContentNotFound) + list, err = repo.ListForOwner(ctx, apiKeyID, 42, now, 0, 10) + require.NoError(t, err) + require.Len(t, list, 1) +} diff --git a/backend/internal/repository/wire.go b/backend/internal/repository/wire.go index f56fee1a58..ba13b0f767 100644 --- a/backend/internal/repository/wire.go +++ b/backend/internal/repository/wire.go @@ -82,6 +82,7 @@ var ProviderSet = wire.NewSet( NewUsageLogRepository, NewUsageBillingRepository, NewBatchImageRepository, + NewGeminiCachedContentRepository, NewIdempotencyRepository, NewUsageCleanupRepository, NewDashboardAggregationRepository, diff --git a/backend/internal/server/routes/gateway.go b/backend/internal/server/routes/gateway.go index a921a42d44..0af1c4cb23 100644 --- a/backend/internal/server/routes/gateway.go +++ b/backend/internal/server/routes/gateway.go @@ -349,6 +349,12 @@ func RegisterGatewayRoutes( gemini.GET("/models/:model", h.Gateway.GeminiV1BetaGetModel) // Gin treats ":" as a param marker, but Gemini uses "{model}:{action}" in the same segment. gemini.POST("/models/*modelAction", h.Gateway.GeminiV1BetaModels) + // 显式上下文缓存：网关自管资源 ID 与账号绑定，list/get 只读本地记录。 + gemini.POST("/cachedContents", h.Gateway.GeminiCachedContentsCreate) + gemini.GET("/cachedContents", h.Gateway.GeminiCachedContentsList) + gemini.GET("/cachedContents/:cacheID", h.Gateway.GeminiCachedContentsGet) + gemini.PATCH("/cachedContents/:cacheID", h.Gateway.GeminiCachedContentsPatch) + gemini.DELETE("/cachedContents/:cacheID", h.Gateway.GeminiCachedContentsDelete) } // OpenAI Responses API（不带v1前缀的别名）— auto-route based on group platform diff --git a/backend/internal/server/routes/prompt_audit_route_coverage_test.go b/backend/internal/server/routes/prompt_audit_route_coverage_test.go index 2d82f7c1da..64b33d92b9 100644 --- a/backend/internal/server/routes/prompt_audit_route_coverage_test.go +++ b/backend/internal/server/routes/prompt_audit_route_coverage_test.go @@ -47,6 +47,7 @@ func TestEveryGatewayPOSTRouteIsClassifiedForPromptAuditCoverage(t *testing.T) { "/videos/edits": {"grok_media.go"}, "/videos/extensions": {"grok_media.go"}, "/models/*modelAction": {"gemini_v1beta_handler.go"}, + "/cachedContents": {"gemini_cached_contents_handler.go"}, "/tts": {"grok_audio.go"}, "/web_search": {"gateway_web_search.go"}, "/x_search": {"gateway_web_search.go"}, diff --git a/backend/internal/service/billing_cache_service.go b/backend/internal/service/billing_cache_service.go index 691d2558f7..70f75a7ae8 100644 --- a/backend/internal/service/billing_cache_service.go +++ b/backend/internal/service/billing_cache_service.go @@ -3,6 +3,7 @@ package service import ( "context" "fmt" + "math" "strconv" "sync" "sync/atomic" @@ -573,6 +574,11 @@ func (s *BillingCacheService) InvalidateAPIKeyRateLimit(ctx context.Context, key // resets expired windows in-memory and triggers async DB reset, // and returns an error if any window limit is exceeded. func (s *BillingCacheService) checkAPIKeyRateLimits(ctx context.Context, apiKey *APIKey) error { + return s.checkAPIKeyRateLimitsWithPending(ctx, apiKey, 0) +} + +// checkAPIKeyRateLimitsWithPending 在窗口用量上叠加 pending（即将扣除的金额）后检查限速窗口。 +func (s *BillingCacheService) checkAPIKeyRateLimitsWithPending(ctx context.Context, apiKey *APIKey, pending float64) error { if s.cache == nil { // No cache: fall back to reading from DB directly if s.apiKeyRateLimitLoader == nil { @@ -583,7 +589,7 @@ func (s *BillingCacheService) checkAPIKeyRateLimits(ctx context.Context, apiKey return nil // Don't block requests on DB errors } return s.evaluateRateLimits(ctx, apiKey, data.Usage5h, data.Usage1d, data.Usage7d, - data.Window5hStart, data.Window1dStart, data.Window7dStart) + data.Window5hStart, data.Window1dStart, data.Window7dStart, pending) } cacheData, err := s.cache.GetAPIKeyRateLimit(ctx, apiKey.ID) @@ -628,11 +634,11 @@ func (s *BillingCacheService) checkAPIKeyRateLimits(ctx context.Context, apiKey t := time.Unix(cacheData.Window7d, 0) w7d = &t } - return s.evaluateRateLimits(ctx, apiKey, cacheData.Usage5h, cacheData.Usage1d, cacheData.Usage7d, w5h, w1d, w7d) + return s.evaluateRateLimits(ctx, apiKey, cacheData.Usage5h, cacheData.Usage1d, cacheData.Usage7d, w5h, w1d, w7d, pending) } -// evaluateRateLimits checks usage against limits, triggering async resets for expired windows. -func (s *BillingCacheService) evaluateRateLimits(ctx context.Context, apiKey *APIKey, usage5h, usage1d, usage7d float64, w5h, w1d, w7d *time.Time) error { +// evaluateRateLimits checks usage (plus pending) against limits, triggering async resets for expired windows. +func (s *BillingCacheService) evaluateRateLimits(ctx context.Context, apiKey *APIKey, usage5h, usage1d, usage7d float64, w5h, w1d, w7d *time.Time, pending float64) error { needsReset := false // Reset expired windows in-memory for check purposes @@ -675,13 +681,16 @@ func (s *BillingCacheService) evaluateRateLimits(ctx context.Context, apiKey *AP } // Check limits - if apiKey.RateLimit5h > 0 && usage5h >= apiKey.RateLimit5h { + exceeds := func(usage, limit float64) bool { + return limit > 0 && (usage >= limit || usage+pending > limit) + } + if exceeds(usage5h, apiKey.RateLimit5h) { return ErrAPIKeyRateLimit5hExceeded } - if apiKey.RateLimit1d > 0 && usage1d >= apiKey.RateLimit1d { + if exceeds(usage1d, apiKey.RateLimit1d) { return ErrAPIKeyRateLimit1dExceeded } - if apiKey.RateLimit7d > 0 && usage7d >= apiKey.RateLimit7d { + if exceeds(usage7d, apiKey.RateLimit7d) { return ErrAPIKeyRateLimit7dExceeded } return nil @@ -781,6 +790,52 @@ func (s *BillingCacheService) CheckBillingEligibility(ctx context.Context, user return nil } +// CheckChargeAffordable 校验扣费前即可确定金额的费用（倍率后 USD，如显式缓存存储费）不会透支： +// 余额模式要求缓存余额覆盖 amount，订阅模式要求日 / 周 / 月剩余额度覆盖 amount， +// API Key 总额度与限速窗口同样叠加 amount 检查。须在 CheckBillingEligibility 通过后调用；简易模式不校验。 +func (s *BillingCacheService) CheckChargeAffordable(ctx context.Context, user *User, apiKey *APIKey, group *Group, subscription *UserSubscription, amount float64) error { + if user == nil || amount <= 0 || math.IsNaN(amount) || math.IsInf(amount, 0) { + return nil + } + if s.cfg != nil && s.cfg.RunMode == config.RunModeSimple { + return nil + } + if group != nil && group.IsSubscriptionType() && subscription != nil { + subData, err := s.GetSubscriptionStatus(ctx, user.ID, group.ID) + if err != nil { + return ErrBillingServiceUnavailable.WithCause(err) + } + if group.HasDailyLimit() && subData.DailyUsage+amount > *group.DailyLimitUSD { + return ErrDailyLimitExceeded + } + if group.HasWeeklyLimit() && subData.WeeklyUsage+amount > *group.WeeklyLimitUSD { + return ErrWeeklyLimitExceeded + } + if group.HasMonthlyLimit() && subData.MonthlyUsage+amount > *group.MonthlyLimitUSD { + return ErrMonthlyLimitExceeded + } + } else { + balance, err := s.GetUserBalance(ctx, user.ID) + if err != nil { + return ErrBillingServiceUnavailable.WithCause(err) + } + if balance < amount { + return ErrInsufficientBalance + } + } + if apiKey != nil { + if apiKey.Quota > 0 && apiKey.QuotaUsed+amount > apiKey.Quota { + return ErrAPIKeyQuotaExhausted + } + if apiKey.HasRateLimits() { + if err := s.checkAPIKeyRateLimitsWithPending(ctx, apiKey, amount); err != nil { + return err + } + } + } + return nil +} + // checkSimpleModeAPIKeyRateLimits is deliberately DB-authoritative. Redis // updates are asynchronous and can be dropped or missed after a committed // transaction, so using the cache here could let a limited simple-mode key diff --git a/backend/internal/service/billing_inflight_reservation.go b/backend/internal/service/billing_inflight_reservation.go index 7c6c7b45fc..714143d554 100644 --- a/backend/internal/service/billing_inflight_reservation.go +++ b/backend/internal/service/billing_inflight_reservation.go @@ -213,7 +213,7 @@ func (s *BillingCacheService) InflightReservationFailClosedOnUnpriced() bool { // ReserveInflightBalance 简化封装：返回释放函数，不续期、不做计费交接（预留最长存活 TTL）。 func (s *BillingCacheService) ReserveInflightBalance(ctx context.Context, user *User, group *Group, subscription *UserSubscription, estimate float64) (func(), error) { - r, err := s.reserveInflight(ctx, user, group, subscription, estimate, false) + r, err := s.reserveInflight(ctx, user, group, subscription, estimate, false, true) if err != nil { return noopRelease, err } @@ -232,10 +232,16 @@ func (s *BillingCacheService) ReserveInflightBalance(ctx context.Context, user * // 开关关闭 / 简易模式 / 订阅模式 / estimate <= 0 / 缓存不支持 / 余额读取失败 / Redis 执行失败。 // 仅当 Redis 明确判定 缓存余额 - 在途合计 < estimate（且已有在途请求）时返回 ErrInsufficientBalance。 func (s *BillingCacheService) ReserveInflight(ctx context.Context, user *User, group *Group, subscription *UserSubscription, estimate float64) (*InflightReservation, error) { - return s.reserveInflight(ctx, user, group, subscription, estimate, true) + return s.reserveInflight(ctx, user, group, subscription, estimate, true, true) } -func (s *BillingCacheService) reserveInflight(ctx context.Context, user *User, group *Group, subscription *UserSubscription, estimate float64, renew bool) (*InflightReservation, error) { +// ReserveInflightCharge 为扣费前即可确定金额的费用登记在途预留，金额不受 max_reservation_usd 封顶， +// 其余语义与 ReserveInflight 相同。 +func (s *BillingCacheService) ReserveInflightCharge(ctx context.Context, user *User, group *Group, subscription *UserSubscription, amount float64) (*InflightReservation, error) { + return s.reserveInflight(ctx, user, group, subscription, amount, true, false) +} + +func (s *BillingCacheService) reserveInflight(ctx context.Context, user *User, group *Group, subscription *UserSubscription, estimate float64, renew, capped bool) (*InflightReservation, error) { cfg, ok := s.inflightReservationConfig() if !ok || user == nil { return nil, nil @@ -243,7 +249,7 @@ func (s *BillingCacheService) reserveInflight(ctx context.Context, user *User, g if group != nil && group.IsSubscriptionType() && subscription != nil { return nil, nil } - if cfg.MaxReservationUSD > 0 && estimate > cfg.MaxReservationUSD { + if capped && cfg.MaxReservationUSD > 0 && estimate > cfg.MaxReservationUSD { estimate = cfg.MaxReservationUSD } if estimate <= 0 || math.IsNaN(estimate) || math.IsInf(estimate, 0) { diff --git a/backend/internal/service/billing_service.go b/backend/internal/service/billing_service.go index 449acf1597..6dfb9a0399 100644 --- a/backend/internal/service/billing_service.go +++ b/backend/internal/service/billing_service.go @@ -1328,6 +1328,19 @@ func (s *BillingService) GetModelPricing(model string) (*ModelPricing, error) { // getModelPricingAt 是 GetModelPricing 的带计费时点内部变体：pricingAt 显式 // 驱动 DeepSeek pro→Flash 切换判定（切换点前 Pro 价、之后 Flash 价），使 // 展示/估算路径可与历史补账同刻复算，测试也能用固定时点钉住断言。 +// CacheStoragePricePerTokenHour 返回模型的显式上下文缓存存储单价（每 token 每小时，USD）。 +// 只做确定性识别，不按系列兜底；价格表未配置时返回 0。 +func (s *BillingService) CacheStoragePricePerTokenHour(model string) float64 { + if s == nil || s.pricingService == nil { + return 0 + } + pricing := s.pricingService.GetIdentifiedModelPricing(model) + if pricing == nil || pricing.CacheStorageCostPerTokenPerHour <= 0 { + return 0 + } + return pricing.CacheStorageCostPerTokenPerHour +} + func (s *BillingService) getModelPricingAt(model string, pricingAt time.Time) (*ModelPricing, error) { // 标准化模型名称（转小写） model = strings.ToLower(model) diff --git a/backend/internal/service/gateway_service.go b/backend/internal/service/gateway_service.go index c22097ede1..bdb90cc193 100644 --- a/backend/internal/service/gateway_service.go +++ b/backend/internal/service/gateway_service.go @@ -649,6 +649,9 @@ type ForwardResult struct { ImageSizeBreakdown map[string]int SearchCount int AudioUsage *AudioUsage + // CacheStorageTokenHours 显式上下文缓存的存储量（缓存 token 数 × 存储小时数），按计费模型的 + // 存储单价折算为费用后叠加在 token 费用之上。 + CacheStorageTokenHours float64 } // GatewayFailureStage identifies which request stage failed. The zero value is diff --git a/backend/internal/service/gateway_usage_billing.go b/backend/internal/service/gateway_usage_billing.go index 3a612d17d0..4c164f2ed0 100644 --- a/backend/internal/service/gateway_usage_billing.go +++ b/backend/internal/service/gateway_usage_billing.go @@ -787,42 +787,9 @@ func (s *GatewayService) recordUsageCore(ctx context.Context, input *recordUsage cacheTTLOverridden = (result.Usage.CacheCreation5mTokens + result.Usage.CacheCreation1hTokens) > 0 } - // 获取费率倍数（优先级：用户专属 > 分组默认 > 系统默认） - multiplier := 1.0 - if s.cfg != nil { - multiplier = s.cfg.Default.RateMultiplier - } - if apiKey.GroupID != nil && apiKey.Group != nil { - groupDefault := apiKey.Group.RateMultiplier - multiplier = s.ResolveUserGroupRateMultiplier(ctx, user.ID, *apiKey.GroupID, groupDefault) - } - // token 倍率叠加高峰因子（token 计费含图片 token，图片按次倍率不受影响）。高峰因子按请求时刻现算， - // 不并入上面的 getUserGroupRateMultiplier，以免污染 user:group 倍率缓存。 - pricingAt := input.PricingAt - if pricingAt.IsZero() { - pricingAt = timezone.Now() - } - multiplier, imageMultiplier := computePeakAwareMultipliers(apiKey, multiplier, pricingAt) - - // 确定计费模型 - concreteBillingModel := forwardResultBillingModel(result.Model, result.UpstreamModel) - billingModel := concreteBillingModel - if input.BillingModelSource == BillingModelSourceChannelMapped && input.ChannelMappedModel != "" { - billingModel = input.ChannelMappedModel - } - if input.BillingModelSource == BillingModelSourceRequested && input.OriginalModel != "" { - billingModel = input.OriginalModel - } - // composite 分组的公开别名（如 all/claude）会经 OriginalModel/ChannelMappedModel - // 进入上面的来源覆盖：任意别名查无价会静默落 $0，含家族词的别名则被价格表的 - // 家族模糊匹配错计（如 Opus 流量按 Sonnet 兜底价）。除非管理员为别名显式配置了 - // 渠道定价（OpenRouter 式自定价），composite 请求一律按实际转发的具体模型计费。 - if apiKey.Group != nil && apiKey.Group.Platform == PlatformComposite { - billingModel = s.compositeBillableModel(ctx, apiKey, billingModel, concreteBillingModel) - } - // 通用兜底（与 OpenAI 路径的 usageBillingModelCandidates 语义对齐）： - // 选定模型查不到任何价格时回退到实际转发的具体模型。已定价流量不受影响。 - billingModel = s.billableModelWithFallback(ctx, apiKey, billingModel, result.UpstreamModel, result.Model) + multiplier, imageMultiplier, pricingAt := s.recordUsageMultipliers(ctx, apiKey, user, input.PricingAt) + billingModel := s.recordUsageBillingModel(ctx, apiKey, result, input.ChannelUsageFields) + cacheStorageCost := s.cacheStorageCost(result, billingModel) // 确定 RequestedModel（渠道映射前的原始模型） requestedModel := result.Model @@ -840,7 +807,7 @@ func (s *GatewayService) recordUsageCore(ctx context.Context, input *recordUsage input.BillingModelSource, result.UpstreamResponseModel, result.UpstreamResponseModelConflict, - result.ImageCount > 0 || result.AudioUsage != nil || result.SearchCount > 0, + result.ImageCount > 0 || result.AudioUsage != nil || result.SearchCount > 0 || result.CacheStorageTokenHours > 0, ); responseModel != "" && !strings.EqualFold(responseModel, strings.TrimSpace(billingModel)) { if identified, responseChannelPriced := s.hasIdentifiedResponseModelPricing(ctx, responseModel, apiKey); identified { responseCost := s.calculateRecordUsageCost(ctx, result, apiKey, responseModel, multiplier, imageMultiplier, pricingAt) @@ -866,8 +833,9 @@ func (s *GatewayService) recordUsageCore(ctx context.Context, input *recordUsage usageLog := s.buildRecordUsageLog(ctx, input, result, apiKey, user, account, subscription, requestedModel, multiplier, imageMultiplier, accountRateMultiplier, billingType, cacheTTLOverridden, cost) - // 计算账号统计定价费用（使用最终上游模型匹配自定义规则） - if apiKey.GroupID != nil { + // 计算账号统计定价费用（使用最终上游模型匹配自定义规则）。纯存储费不是一次模型请求， + // 账号统计成本回落 total_cost（官方存储价）。 + if apiKey.GroupID != nil && !isCacheStorageOnlyResult(result) { applyAccountStatsCost(ctx, usageLog, s.channelService, s.billingService, account.ID, *apiKey.GroupID, result.UpstreamModel, result.Model, // Anthropic's input_tokens excludes cache_read and cache_creation (billed separately); @@ -879,9 +847,10 @@ func (s *GatewayService) recordUsageCore(ctx context.Context, input *recordUsage CacheReadTokens: result.Usage.CacheReadInputTokens, ImageOutputTokens: result.Usage.ImageOutputTokens, }, - cost.TotalCost, pricingAt, + cost.TotalCost-cacheStorageCost, pricingAt, accountStatsLongContextPricingEnabled(nil), ) + addCacheStorageToAccountStatsCost(usageLog, cacheStorageCost) } simpleModeKeyRateLimitOnly := simpleModeKeyRateLimitBillingEnabled(s.cfg, apiKey) @@ -927,6 +896,112 @@ func (s *GatewayService) recordUsageCore(ctx context.Context, input *recordUsage return nil } +// recordUsageMultipliers 返回 token / 图片售价倍率（用户专属 > 分组默认 > 系统默认，token 倍率叠加高峰因子） +// 与计价时刻；pricingAt 为零值时取当前时刻。 +func (s *GatewayService) recordUsageMultipliers(ctx context.Context, apiKey *APIKey, user *User, pricingAt time.Time) (float64, float64, time.Time) { + multiplier := 1.0 + if s.cfg != nil { + multiplier = s.cfg.Default.RateMultiplier + } + if apiKey.GroupID != nil && apiKey.Group != nil { + multiplier = apiKey.Group.RateMultiplier + if user != nil { + multiplier = s.ResolveUserGroupRateMultiplier(ctx, user.ID, *apiKey.GroupID, multiplier) + } + } + // 高峰因子按请求时刻现算，不并入 getUserGroupRateMultiplier，以免污染 user:group 倍率缓存。 + if pricingAt.IsZero() { + pricingAt = timezone.Now() + } + multiplier, imageMultiplier := computePeakAwareMultipliers(apiKey, multiplier, pricingAt) + return multiplier, imageMultiplier, pricingAt +} + +// recordUsageBillingModel 按计费模型来源、composite 规则与无价兜底确定计费模型。 +func (s *GatewayService) recordUsageBillingModel(ctx context.Context, apiKey *APIKey, result *ForwardResult, fields ChannelUsageFields) string { + concreteBillingModel := forwardResultBillingModel(result.Model, result.UpstreamModel) + billingModel := concreteBillingModel + if fields.BillingModelSource == BillingModelSourceChannelMapped && fields.ChannelMappedModel != "" { + billingModel = fields.ChannelMappedModel + } + if fields.BillingModelSource == BillingModelSourceRequested && fields.OriginalModel != "" { + billingModel = fields.OriginalModel + } + // composite 分组的公开别名（如 all/claude）会经 OriginalModel/ChannelMappedModel + // 进入上面的来源覆盖：任意别名查无价会静默落 $0，含家族词的别名则被价格表的 + // 家族模糊匹配错计（如 Opus 流量按 Sonnet 兜底价）。除非管理员为别名显式配置了 + // 渠道定价（OpenRouter 式自定价），composite 请求一律按实际转发的具体模型计费。 + if apiKey.Group != nil && apiKey.Group.Platform == PlatformComposite { + billingModel = s.compositeBillableModel(ctx, apiKey, billingModel, concreteBillingModel) + } + // 通用兜底（与 OpenAI 路径的 usageBillingModelCandidates 语义对齐）： + // 选定模型查不到任何价格时回退到实际转发的具体模型。已定价流量不受影响。 + return s.billableModelWithFallback(ctx, apiKey, billingModel, result.UpstreamModel, result.Model) +} + +// cacheStoragePrice 返回计费模型的显式缓存存储单价；计费模型未配置时依次回退到实际转发模型与请求模型。 +func (s *GatewayService) cacheStoragePrice(result *ForwardResult, billingModel string) float64 { + if s.billingService == nil { + return 0 + } + for _, model := range []string{billingModel, result.UpstreamModel, result.Model} { + if strings.TrimSpace(model) == "" { + continue + } + if price := s.billingService.CacheStoragePricePerTokenHour(model); price > 0 { + return price + } + } + return 0 +} + +// cacheStorageCost 返回 result 的显式缓存存储费（USD，倍率前）。 +func (s *GatewayService) cacheStorageCost(result *ForwardResult, billingModel string) float64 { + if result.CacheStorageTokenHours <= 0 { + return 0 + } + return result.CacheStorageTokenHours * s.cacheStoragePrice(result, billingModel) +} + +// isCacheStorageOnlyResult 判断 result 是否只有显式缓存存储费（延长有效期）：不含 token 用量，不算一次模型请求。 +func isCacheStorageOnlyResult(result *ForwardResult) bool { + u := result.Usage + return result.CacheStorageTokenHours > 0 && u.InputTokens == 0 && u.OutputTokens == 0 && + u.CacheCreationInputTokens == 0 && u.CacheReadInputTokens == 0 && u.ImageOutputTokens == 0 && + result.ImageCount == 0 && result.AudioUsage == nil && result.SearchCount == 0 +} + +// CacheStoragePriced 判断 RecordUsage 对 input.Result 采用的计费模型（含回退候选）是否配置了显式缓存存储单价。 +func (s *GatewayService) CacheStoragePriced(ctx context.Context, input *RecordUsageInput) bool { + billingModel := s.recordUsageBillingModel(ctx, input.APIKey, input.Result, input.ChannelUsageFields) + return s.cacheStoragePrice(input.Result, billingModel) > 0 +} + +// QuoteUsageCost 按与 RecordUsage 相同的倍率、计费模型与定价计算 input.Result 的应付金额（倍率后 USD），不记录也不扣费。 +func (s *GatewayService) QuoteUsageCost(ctx context.Context, input *RecordUsageInput) float64 { + user := input.User + if user == nil { + user = input.APIKey.User + } + multiplier, imageMultiplier, pricingAt := s.recordUsageMultipliers(ctx, input.APIKey, user, input.PricingAt) + billingModel := s.recordUsageBillingModel(ctx, input.APIKey, input.Result, input.ChannelUsageFields) + cost := s.calculateRecordUsageCost(ctx, input.Result, input.APIKey, billingModel, multiplier, imageMultiplier, pricingAt) + if cost == nil { + return 0 + } + return cost.ActualCost +} + +// addCacheStorageToAccountStatsCost 把显式缓存存储费按官方价叠加到账号统计成本。 +// 账号统计定价只覆盖 token 部分；统计成本为空时成本统计回落 total_cost，其中已含存储费。 +func addCacheStorageToAccountStatsCost(usageLog *UsageLog, storageCost float64) { + if usageLog == nil || usageLog.AccountStatsCost == nil || storageCost <= 0 { + return + } + cost := *usageLog.AccountStatsCost + storageCost + usageLog.AccountStatsCost = &cost +} + // calculateRecordUsageCost 根据请求类型计算费用。 func (s *GatewayService) calculateRecordUsageCost( ctx context.Context, @@ -964,8 +1039,13 @@ func (s *GatewayService) calculateRecordUsageCost( return s.billingService.CalculateAudioCost(result.AudioUsage.Mode, result.AudioUsage.DurationOrUnits, cfg, multiplier) } - // Token 计费；SearchCount 为叠加 surcharge（不替代 token）。 - tokenCost := s.calculateTokenCost(ctx, result, apiKey, billingModel, multiplier, pricingAt) + // Token 计费；SearchCount 为叠加 surcharge（不替代 token）。纯存储费不是一次模型请求，不计 token / 按次费用。 + var tokenCost *CostBreakdown + if isCacheStorageOnlyResult(result) { + tokenCost = &CostBreakdown{} + } else { + tokenCost = s.calculateTokenCost(ctx, result, apiKey, billingModel, multiplier, pricingAt) + } if result.SearchCount > 0 { price := groupSearchPricePer1kFromAPIKey(apiKey) if price != nil && *price == 0 { @@ -980,6 +1060,14 @@ func (s *GatewayService) calculateRecordUsageCost( tokenCost.ActualCost += searchCost.ActualCost } } + // 显式缓存存储费同为叠加 surcharge，与 token 费用使用同一计费模型与倍率。 + if storageCost := s.cacheStorageCost(result, billingModel); storageCost > 0 { + if tokenCost == nil { + tokenCost = &CostBreakdown{} + } + tokenCost.TotalCost += storageCost + tokenCost.ActualCost += storageCost * multiplier + } return tokenCost } diff --git a/backend/internal/service/gemini_cached_content.go b/backend/internal/service/gemini_cached_content.go new file mode 100644 index 0000000000..7c3c6b14a6 --- /dev/null +++ b/backend/internal/service/gemini_cached_content.go @@ -0,0 +1,957 @@ +package service + +import ( + "bytes" + "context" + "crypto/rand" + "encoding/base32" + "encoding/json" + "errors" + "fmt" + "io" + "math" + "net/http" + "net/url" + "regexp" + "slices" + "strconv" + "strings" + "time" + "unicode/utf8" + + "github.com/Wei-Shaw/sub2api/internal/config" + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + "github.com/Wei-Shaw/sub2api/internal/pkg/geminicli" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" +) + +const ( + // GeminiCachedContentNamePrefix 是缓存资源名前缀，网关与上游 AI Studio 一致。 + GeminiCachedContentNamePrefix = "cachedContents/" + // GeminiCachedContentNotFoundMessage 与上游对不存在、已过期或无权访问的缓存返回的文案一致。 + GeminiCachedContentNotFoundMessage = "CachedContent not found (or permission denied)" + // GeminiCachedContentNotFoundResponse 是网关对引用不存在缓存统一返回的 403 响应体（与 AI Studio 上游一致）。 + GeminiCachedContentNotFoundResponse = `{"error":{"code":403,"message":"CachedContent not found (or permission denied)","status":"PERMISSION_DENIED"}}` + + geminiCachedContentDefaultTTL = time.Hour + geminiCachedContentDefaultMaxTTL = 24 * time.Hour + geminiCachedContentDisplayNameMaxRunes = 128 + geminiCachedContentUpstreamTagPrefix = "s2a-" + geminiCachedContentResponseMaxBytes = 4 << 20 + geminiCachedContentUpstreamTimeout = 5 * time.Minute + geminiCachedContentOfficialHost = "generativelanguage.googleapis.com" + + // GeminiCachedContentListDefaultPageSize / GeminiCachedContentListMaxPageSize 约束 list 分页。 + GeminiCachedContentListDefaultPageSize = 50 + GeminiCachedContentListMaxPageSize = 1000 +) + +var ( + ErrGeminiCachedContentNotFound = infraerrors.New(http.StatusNotFound, "GEMINI_CACHED_CONTENT_NOT_FOUND", "gemini cached content not found") + ErrGeminiCachedContentExists = infraerrors.New(http.StatusConflict, "GEMINI_CACHED_CONTENT_EXISTS", "gemini cached content already exists") + + geminiCachedContentPublicIDPattern = regexp.MustCompile(`^[a-z0-9]{1,64}$`) + geminiCachedContentAIStudioName = regexp.MustCompile(`^cachedContents/[A-Za-z0-9_-]+$`) + geminiCachedContentVertexName = regexp.MustCompile(`^projects/[A-Za-z0-9_.:-]+/locations/([a-z0-9-]+)/cachedContents/[A-Za-z0-9_-]+$`) + geminiCachedContentPublicIDEncoding = base32.StdEncoding.WithPadding(base32.NoPadding) +) + +// GeminiCachedContent 是网关对外暴露的显式缓存与上游资源、账号、归属的绑定。 +type GeminiCachedContent struct { + ID int64 + PublicID string + UserID int64 + APIKeyID int64 + GroupID int64 + // AccountID 是持有上游缓存的账号；引用缓存的请求只能发往该账号。 + AccountID int64 + UpstreamName string + // Model 是客户端视角的模型名（渠道映射后，不含 models/ 前缀），引用缓存的请求模型必须与之字面一致。 + Model string + // UpstreamModel 是创建时发给上游的模型资源名。 + UpstreamModel string + DisplayName string + TotalTokenCount int64 + ExpireTime time.Time + CreatedAt time.Time + UpdatedAt time.Time +} + +// PublicName 返回客户端可见的缓存资源名。 +func (c *GeminiCachedContent) PublicName() string { + return GeminiCachedContentNamePrefix + c.PublicID +} + +// UpstreamModelID 返回上游模型资源名的末段（裸模型名），用于计费定价查找。 +func (c *GeminiCachedContent) UpstreamModelID() string { + model := c.UpstreamModel + if i := strings.LastIndex(model, "/"); i >= 0 { + model = model[i+1:] + } + return model +} + +// GeminiCachedContentRepository 持久化显式缓存绑定。所有读取都过滤已软删除的记录。 +type GeminiCachedContentRepository interface { + Create(ctx context.Context, record *GeminiCachedContent) error + GetForOwner(ctx context.Context, apiKeyID int64, publicID string) (*GeminiCachedContent, error) + // ListForOwner 按 id 倒序返回分组内 expire_time 晚于 activeAt 的记录；beforeID>0 时只返回 id 更小的记录。 + ListForOwner(ctx context.Context, apiKeyID, groupID int64, activeAt time.Time, beforeID int64, limit int) ([]*GeminiCachedContent, error) + // UpdateExpireTime 原子地写入新的到期时间并返回写入前的值。 + UpdateExpireTime(ctx context.Context, id int64, expireTime time.Time) (time.Time, error) + SoftDelete(ctx context.Context, id int64) error +} + +// IsGeminiExplicitCacheAccountType 只按平台与账号类型判断是否可能支持显式缓存，适用于调度快照账号。 +func IsGeminiExplicitCacheAccountType(account *Account) bool { + if account == nil || account.Platform != PlatformGemini { + return false + } + return account.Type == AccountTypeAPIKey || account.Type == AccountTypeServiceAccount +} + +// IsGeminiExplicitCacheAccount 判断完整账号是否支持显式缓存：官方地址的 AI Studio API Key 账号、 +// 声明上游支持显式缓存的自定义地址 API Key 账号，或 Vertex service account 账号。 +func IsGeminiExplicitCacheAccount(account *Account) bool { + if !IsGeminiExplicitCacheAccountType(account) { + return false + } + if account.Type == AccountTypeServiceAccount { + return true + } + if strings.TrimSpace(account.GetCredential("api_key")) == "" { + return false + } + return isOfficialGeminiAPIBaseURL(account.GetCredential("base_url")) || account.GeminiExplicitCacheUpstreamEnabled() +} + +const ( + geminiExplicitCacheUpstreamKey = "explicit_cache_upstream" + geminiExplicitCacheUpstreamMaxTTLKey = "explicit_cache_upstream_max_ttl_seconds" +) + +// GeminiExplicitCacheUpstreamEnabled 判断 API Key 账号是否声明其上游实现了显式缓存接口， +// 并保证缓存只在创建它的上游账号上使用（本系统同版本网关）。 +func (a *Account) GeminiExplicitCacheUpstreamEnabled() bool { + if a == nil || a.Type != AccountTypeAPIKey || a.Credentials == nil { + return false + } + enabled, _ := a.Credentials[geminiExplicitCacheUpstreamKey].(bool) + return enabled +} + +// GeminiExplicitCacheUpstreamMaxTTL 返回上游网关允许的缓存最长有效期；0 表示未配置。 +func (a *Account) GeminiExplicitCacheUpstreamMaxTTL() time.Duration { + if !a.GeminiExplicitCacheUpstreamEnabled() { + return 0 + } + var seconds float64 + switch v := a.Credentials[geminiExplicitCacheUpstreamMaxTTLKey].(type) { + case float64: + seconds = v + case int: + seconds = float64(v) + case int64: + seconds = float64(v) + case json.Number: + seconds, _ = v.Float64() + case string: + seconds, _ = strconv.ParseFloat(strings.TrimSpace(v), 64) + } + if seconds <= 0 || math.IsNaN(seconds) || math.IsInf(seconds, 0) { + return 0 + } + return time.Duration(seconds) * time.Second +} + +// geminiCachedContentAPIBase 返回 API Key 账号缓存接口的上游根地址。 +func geminiCachedContentAPIBase(account *Account) string { + if isOfficialGeminiAPIBaseURL(account.GetCredential("base_url")) { + return geminicli.AIStudioBaseURL + } + return account.GetGeminiBaseURL(geminicli.AIStudioBaseURL) +} + +func isOfficialGeminiAPIBaseURL(raw string) bool { + raw = strings.TrimSpace(raw) + if raw == "" { + return true + } + u, err := url.Parse(raw) + if err != nil { + return false + } + path := strings.TrimRight(u.Path, "/") + return strings.EqualFold(u.Scheme, "https") && + strings.EqualFold(u.Hostname(), geminiCachedContentOfficialHost) && + u.Port() == "" && path == "" && u.RawQuery == "" +} + +// NewGeminiCachedContentPublicID 生成 40 位小写字母数字的缓存 ID。 +func NewGeminiCachedContentPublicID() (string, error) { + var b [25]byte + if _, err := rand.Read(b[:]); err != nil { + return "", err + } + return strings.ToLower(geminiCachedContentPublicIDEncoding.EncodeToString(b[:])), nil +} + +// ParseGeminiCachedContentPublicID 从 "cachedContents/{id}" 或裸 id 中解析缓存 ID。 +func ParseGeminiCachedContentPublicID(name string) (string, bool) { + id := strings.TrimPrefix(strings.TrimSpace(name), GeminiCachedContentNamePrefix) + if !geminiCachedContentPublicIDPattern.MatchString(id) { + return "", false + } + return id, true +} + +type geminiCachedContentBoundCtxKey struct{} + +// WithGeminiCachedContentBound 标记请求引用了显式缓存并已硬绑定到持有账号。 +func WithGeminiCachedContentBound(ctx context.Context) context.Context { + return context.WithValue(ctx, geminiCachedContentBoundCtxKey{}, true) +} + +func geminiCachedContentBoundFromContext(ctx context.Context) bool { + if ctx == nil { + return false + } + bound, _ := ctx.Value(geminiCachedContentBoundCtxKey{}).(bool) + return bound +} + +// isGeminiBoundCacheUpstreamPassthrough 判断绑定缓存的请求在上游网关账号上收到的错误是否只反映 +// 上游持有缓存的那个账号的状态（429 / 5xx），不应据此限流或重试本账号。 +func isGeminiBoundCacheUpstreamPassthrough(ctx context.Context, account *Account, statusCode int) bool { + if statusCode != http.StatusTooManyRequests && statusCode < http.StatusInternalServerError { + return false + } + return account.GeminiExplicitCacheUpstreamEnabled() && geminiCachedContentBoundFromContext(ctx) +} + +// GeminiCachedContentReference 是请求体中引用显式缓存的位置。 +type GeminiCachedContentReference struct { + Name string + Paths []string +} + +var geminiCachedContentReferencePaths = []string{"cachedContent", "generateContentRequest.cachedContent"} + +// FindGeminiCachedContentReference 返回 generateContent / countTokens 请求体中引用的缓存名；未引用时返回 nil。 +// 引用键重复、使用 snake_case 写法或多个位置引用不同缓存时返回错误，保证校验与改写的引用就是发往上游的引用。 +func FindGeminiCachedContentReference(body []byte) (*GeminiCachedContentReference, error) { + root := gjson.ParseBytes(body) + if !root.IsObject() { + return nil, nil + } + if err := checkGeminiCachedContentReferenceKeys(root, "cachedContent", "generateContentRequest"); err != nil { + return nil, err + } + if nested := root.Get("generateContentRequest"); nested.IsObject() { + if err := checkGeminiCachedContentReferenceKeys(nested, "cachedContent"); err != nil { + return nil, err + } + } + var ref *GeminiCachedContentReference + for _, path := range geminiCachedContentReferencePaths { + value := gjson.GetBytes(body, path) + if !value.Exists() || value.Type == gjson.Null { + continue + } + if value.Type != gjson.String { + return nil, errors.New("cachedContent must be a string") + } + name := strings.TrimSpace(value.String()) + if name == "" { + continue + } + if ref == nil { + ref = &GeminiCachedContentReference{Name: name} + } else if ref.Name != name { + return nil, errors.New("conflicting cachedContent references") + } + ref.Paths = append(ref.Paths, path) + } + return ref, nil +} + +// checkGeminiCachedContentReferenceKeys 拒绝对象中重复出现的 uniqueKeys 与 snake_case 写法的缓存引用键。 +func checkGeminiCachedContentReferenceKeys(obj gjson.Result, uniqueKeys ...string) error { + seen := make(map[string]bool, len(uniqueKeys)) + var err error + obj.ForEach(func(key, value gjson.Result) bool { + name := key.String() + switch { + case name == "cached_content": + err = errors.New("cached_content is not supported, use cachedContent") + case name == "generate_content_request" && (value.Get("cachedContent").Exists() || value.Get("cached_content").Exists()): + err = errors.New("generate_content_request is not supported with cachedContent, use generateContentRequest") + case slices.Contains(uniqueKeys, name): + if seen[name] { + err = fmt.Errorf("duplicate %s field", name) + } + seen[name] = true + } + return err == nil + }) + return err +} + +// RewriteGeminiCachedContentReference 把请求体中的缓存引用改写为上游资源名。 +func RewriteGeminiCachedContentReference(body []byte, ref *GeminiCachedContentReference, upstreamName string) ([]byte, error) { + if ref == nil { + return body, nil + } + out := body + for _, path := range ref.Paths { + next, err := sjson.SetBytes(out, path, upstreamName) + if err != nil { + return nil, err + } + out = next + } + return out, nil +} + +// GeminiCachedContentCreateRequest 是校验后的创建请求。 +type GeminiCachedContentCreateRequest struct { + // Model 为客户端请求的模型名（不含 models/ 前缀）。 + Model string + TTL time.Duration + DisplayName string + // ttlExplicit 表示客户端显式指定了 ttl 或 expireTime。 + ttlExplicit bool + payload map[string]any +} + +// TTLSeconds 返回发往上游的整秒 TTL。 +func (r *GeminiCachedContentCreateRequest) TTLSeconds() int64 { + return int64(math.Ceil(r.TTL.Seconds())) +} + +// TTLForAccount 返回在账号上创建缓存使用的整秒有效期。账号上限小于请求有效期时： +// 未显式指定有效期的请求按账号上限创建，显式指定的返回 false（该账号不可用）。 +func (r *GeminiCachedContentCreateRequest) TTLForAccount(account *Account) (time.Duration, bool) { + ttl := time.Duration(r.TTLSeconds()) * time.Second + limit := account.GeminiExplicitCacheUpstreamMaxTTL() + if limit <= 0 || ttl <= limit { + return ttl, true + } + if r.ttlExplicit { + return 0, false + } + return limit, true +} + +// ParseGeminiCachedContentCreateRequest 校验创建请求。未指定有效期时使用上游默认的 1 小时（不超过 maxTTL）。 +// 数字按原文保留，重新编码后发往上游的内容与请求一致。 +func ParseGeminiCachedContentCreateRequest(body []byte, now time.Time, maxTTL time.Duration) (*GeminiCachedContentCreateRequest, error) { + payload, err := decodeGeminiCachedContentPayload(body) + if err != nil { + return nil, err + } + rawModel, _ := payload["model"].(string) + model := strings.TrimPrefix(strings.TrimSpace(rawModel), "models/") + if model == "" { + return nil, errors.New("model is required") + } + if !IsSafeGeminiModelPathSegment(model) { + return nil, errors.New("model name must be in the format 'models/{model_name}'") + } + + ttl, err := parseGeminiCachedContentExpiration(payload, now) + if err != nil { + return nil, err + } + ttlExplicit := ttl > 0 + if ttl == 0 { + ttl = geminiCachedContentDefaultTTL + if maxTTL > 0 && ttl > maxTTL { + ttl = maxTTL + } + } + if maxTTL > 0 && ttl > maxTTL { + return nil, fmt.Errorf("ttl exceeds the maximum allowed %ds", int64(maxTTL.Seconds())) + } + + displayName := "" + if raw, ok := payload["displayName"]; ok && raw != nil { + s, ok := raw.(string) + if !ok { + return nil, errors.New("displayName must be a string") + } + if utf8.RuneCountInString(s) > geminiCachedContentDisplayNameMaxRunes { + return nil, fmt.Errorf("displayName must be at most %d characters", geminiCachedContentDisplayNameMaxRunes) + } + displayName = s + } + + for _, key := range []string{"name", "model", "displayName", "ttl", "expireTime", "createTime", "updateTime", "usageMetadata"} { + delete(payload, key) + } + return &GeminiCachedContentCreateRequest{Model: model, TTL: ttl, DisplayName: displayName, ttlExplicit: ttlExplicit, payload: payload}, nil +} + +// ParseGeminiCachedContentPatchRequest 校验更新请求，只允许修改有效期，返回新的 TTL（自 now 起算）。 +func ParseGeminiCachedContentPatchRequest(body []byte, updateMask string, now time.Time, maxTTL time.Duration) (time.Duration, error) { + payload, err := decodeGeminiCachedContentPayload(body) + if err != nil { + return 0, err + } + for key := range payload { + switch key { + case "ttl", "expireTime", "name": + default: + return 0, errors.New("can not update immutable fields (model or display_name)") + } + } + for _, field := range strings.Split(updateMask, ",") { + switch strings.TrimSpace(field) { + case "", "ttl", "expireTime", "expire_time": + default: + return 0, errors.New("can not update immutable fields (model or display_name)") + } + } + ttl, err := parseGeminiCachedContentExpiration(payload, now) + if err != nil { + return 0, err + } + if ttl == 0 { + return 0, errors.New("ttl or expireTime is required") + } + if maxTTL > 0 && ttl > maxTTL { + return 0, fmt.Errorf("ttl exceeds the maximum allowed %ds", int64(maxTTL.Seconds())) + } + return ttl, nil +} + +// geminiCachedContentFieldAliases 是缓存元数据字段的 snake_case 写法到 camelCase 的映射。 +var geminiCachedContentFieldAliases = map[string]string{ + "display_name": "displayName", + "expire_time": "expireTime", + "create_time": "createTime", + "update_time": "updateTime", + "usage_metadata": "usageMetadata", +} + +// decodeGeminiCachedContentPayload 解码缓存请求体（数字保留原文），并把 snake_case 写法的元数据字段改为 camelCase； +// 同一字段两种写法同时出现时报错。 +func decodeGeminiCachedContentPayload(body []byte) (map[string]any, error) { + dec := json.NewDecoder(bytes.NewReader(body)) + dec.UseNumber() + var payload map[string]any + if err := dec.Decode(&payload); err != nil || payload == nil { + return nil, errors.New("request body must be a JSON object") + } + if _, err := dec.Token(); !errors.Is(err, io.EOF) { + return nil, errors.New("request body must be a JSON object") + } + for alias, canonical := range geminiCachedContentFieldAliases { + value, ok := payload[alias] + if !ok { + continue + } + if _, dup := payload[canonical]; dup { + return nil, fmt.Errorf("only one of %s or %s may be set", canonical, alias) + } + delete(payload, alias) + payload[canonical] = value + } + return payload, nil +} + +// parseGeminiCachedContentExpiration 解析 ttl / expireTime（二选一），未设置时返回 0；结果向上取整到秒，与发往上游的 ttl 一致。 +func parseGeminiCachedContentExpiration(payload map[string]any, now time.Time) (time.Duration, error) { + rawTTL, hasTTL := payload["ttl"] + rawExpire, hasExpire := payload["expireTime"] + hasTTL = hasTTL && rawTTL != nil + hasExpire = hasExpire && rawExpire != nil + if hasTTL && hasExpire { + return 0, errors.New("only one of ttl or expireTime may be set") + } + if hasTTL { + s, ok := rawTTL.(string) + if !ok { + return 0, errors.New("ttl must be a duration string such as \"3600s\"") + } + return parseGoogleDurationSeconds(s) + } + if hasExpire { + s, ok := rawExpire.(string) + if !ok { + return 0, errors.New("expireTime must be an RFC 3339 timestamp") + } + expire, err := time.Parse(time.RFC3339Nano, strings.TrimSpace(s)) + if err != nil { + return 0, errors.New("expireTime must be an RFC 3339 timestamp") + } + ttl := expire.Sub(now) + if ttl <= 0 { + return 0, errors.New("expireTime must be in the future") + } + return time.Duration(math.Ceil(ttl.Seconds())) * time.Second, nil + } + return 0, nil +} + +func parseGoogleDurationSeconds(raw string) (time.Duration, error) { + raw = strings.TrimSpace(raw) + if !strings.HasSuffix(raw, "s") { + return 0, errors.New("ttl must be a duration string such as \"3600s\"") + } + seconds, err := strconv.ParseFloat(strings.TrimSuffix(raw, "s"), 64) + if err != nil || math.IsNaN(seconds) || math.IsInf(seconds, 0) { + return 0, errors.New("ttl must be a duration string such as \"3600s\"") + } + if seconds <= 0 { + return 0, errors.New("ttl must be positive") + } + if seconds > float64(math.MaxInt64/int64(time.Second)) { + return 0, errors.New("ttl is too large") + } + return time.Duration(math.Ceil(seconds)) * time.Second, nil +} + +// GeminiCachedContentUpstream 是上游缓存资源的元数据。 +type GeminiCachedContentUpstream struct { + Name string + TotalTokenCount int64 + ExpireTime time.Time +} + +// GeminiCachedContentUpstreamError 是上游返回的非 2xx 响应。 +type GeminiCachedContentUpstreamError struct { + StatusCode int + Body []byte +} + +func (e *GeminiCachedContentUpstreamError) Error() string { + return fmt.Sprintf("gemini cached content upstream error: %d", e.StatusCode) +} + +// NotFound 判断上游是否报告缓存不存在（含已过期、已删除）。 +func (e *GeminiCachedContentUpstreamError) NotFound() bool { + return e != nil && isGeminiCachedContentNotFound(e.StatusCode, e.Body) +} + +// GeminiCachedContentMalformedResponseError 表示上游返回 2xx 但响应体不可用；Name 非空时上游资源已经创建。 +type GeminiCachedContentMalformedResponseError struct { + Name string + Err error +} + +func (e *GeminiCachedContentMalformedResponseError) Error() string { + return "gemini cached content upstream returned an unusable response: " + e.Err.Error() +} + +func (e *GeminiCachedContentMalformedResponseError) Unwrap() error { + return e.Err +} + +// RetryableOnOtherAccount 判断创建失败是否值得换一个账号重试（账号凭据、配额或上游可用性问题）。 +func (e *GeminiCachedContentUpstreamError) RetryableOnOtherAccount() bool { + if e == nil { + return false + } + switch e.StatusCode { + case http.StatusUnauthorized, http.StatusForbidden, http.StatusTooManyRequests: + return true + default: + return e.StatusCode >= 500 + } +} + +// GeminiCachedContentService 管理显式缓存的上游资源与网关绑定。 +type GeminiCachedContentService struct { + repo GeminiCachedContentRepository + compat *GeminiMessagesCompatService + cfg *config.Config + now func() time.Time +} + +func NewGeminiCachedContentService(repo GeminiCachedContentRepository, compat *GeminiMessagesCompatService, cfg *config.Config) *GeminiCachedContentService { + return &GeminiCachedContentService{repo: repo, compat: compat, cfg: cfg, now: time.Now} +} + +// Now 返回服务时钟。 +func (s *GeminiCachedContentService) Now() time.Time { + return s.now() +} + +// MaxTTL 返回单次创建或延长允许的最长有效期。 +func (s *GeminiCachedContentService) MaxTTL() time.Duration { + if s.cfg != nil && s.cfg.Gateway.GeminiCachedContentMaxTTLSeconds > 0 { + return time.Duration(s.cfg.Gateway.GeminiCachedContentMaxTTLSeconds) * time.Second + } + return geminiCachedContentDefaultMaxTTL +} + +// MaxTTLForAccount 返回在账号上可设置的最长有效期：全局上限与上游网关上限取较小值。 +func (s *GeminiCachedContentService) MaxTTLForAccount(account *Account) time.Duration { + limit := s.MaxTTL() + if upstream := account.GeminiExplicitCacheUpstreamMaxTTL(); upstream > 0 && upstream < limit { + return upstream + } + return limit +} + +// GeminiCachedContentStorageTokenHours 返回 tokens 个缓存 token 存储 duration 的存储量（token·小时）。 +func GeminiCachedContentStorageTokenHours(tokens int64, duration time.Duration) float64 { + if tokens <= 0 || duration <= 0 { + return 0 + } + return float64(tokens) * duration.Hours() +} + +// ClampGeminiCachedContentExpire 返回本地记录与计费采用的到期时间：上游回报的到期时间不晚于 limit，缺失时取 limit。 +func ClampGeminiCachedContentExpire(upstream, limit time.Time) time.Time { + if upstream.IsZero() || upstream.After(limit) { + return limit + } + return upstream +} + +// Resolve 返回 API Key 名下仍有效的缓存；不存在、不在当前分组或已过期时返回 ErrGeminiCachedContentNotFound。 +func (s *GeminiCachedContentService) Resolve(ctx context.Context, apiKeyID int64, groupID *int64, name string) (*GeminiCachedContent, error) { + publicID, ok := ParseGeminiCachedContentPublicID(name) + if !ok { + return nil, ErrGeminiCachedContentNotFound + } + record, err := s.repo.GetForOwner(ctx, apiKeyID, publicID) + if err != nil { + return nil, err + } + if groupID == nil || record.GroupID != *groupID || !record.ExpireTime.After(s.now()) { + return nil, ErrGeminiCachedContentNotFound + } + return record, nil +} + +func (s *GeminiCachedContentService) Save(ctx context.Context, record *GeminiCachedContent) error { + return s.repo.Create(ctx, record) +} + +func (s *GeminiCachedContentService) List(ctx context.Context, apiKeyID, groupID int64, beforeID int64, limit int) ([]*GeminiCachedContent, error) { + return s.repo.ListForOwner(ctx, apiKeyID, groupID, s.now(), beforeID, limit) +} + +// UpdateExpireTime 更新到期时间并返回更新前库中的到期时间（并发更新按提交顺序各自拿到自己的前值）。 +func (s *GeminiCachedContentService) UpdateExpireTime(ctx context.Context, record *GeminiCachedContent, expireTime time.Time) (time.Time, error) { + previous, err := s.repo.UpdateExpireTime(ctx, record.ID, expireTime) + if err != nil { + return time.Time{}, err + } + record.ExpireTime = expireTime + record.UpdatedAt = s.now() + return previous, nil +} + +func (s *GeminiCachedContentService) Forget(ctx context.Context, record *GeminiCachedContent) error { + return s.repo.SoftDelete(ctx, record.ID) +} + +// BoundAccount 返回持有缓存的账号；账号已删除或不再支持显式缓存时返回 nil。 +func (s *GeminiCachedContentService) BoundAccount(ctx context.Context, record *GeminiCachedContent) (*Account, error) { + if s.compat == nil || s.compat.accountRepo == nil { + return nil, errors.New("account repository not configured") + } + account, err := s.compat.accountRepo.GetByID(ctx, record.AccountID) + if err != nil { + if errors.Is(err, ErrAccountNotFound) { + return nil, nil + } + return nil, err + } + if account == nil || !IsGeminiExplicitCacheAccount(account) { + return nil, nil + } + return account, nil +} + +// GeminiCachedContentUpstreamModelResource 返回创建缓存时发往上游的模型资源名。 +func GeminiCachedContentUpstreamModelResource(account *Account, mappedModel string) (string, error) { + switch account.Type { + case AccountTypeAPIKey: + return "models/" + mappedModel, nil + case AccountTypeServiceAccount: + projectID := strings.TrimSpace(account.VertexProjectID()) + if projectID == "" { + return "", errors.New("vertex project_id is required") + } + location := strings.TrimSpace(account.VertexLocation(mappedModel)) + if location == "" { + location = vertexDefaultLocation + } + if !vertexLocationPattern.MatchString(location) { + return "", fmt.Errorf("invalid vertex location: %s", location) + } + return fmt.Sprintf("projects/%s/locations/%s/publishers/google/models/%s", projectID, location, mappedModel), nil + default: + return "", fmt.Errorf("account type %s does not support explicit context caching", account.Type) + } +} + +// CreateUpstream 在账号上创建上游缓存。 +func (s *GeminiCachedContentService) CreateUpstream(ctx context.Context, account *Account, req *GeminiCachedContentCreateRequest, mappedModel, publicID string, ttl time.Duration) (*GeminiCachedContentUpstream, error) { + upstreamModel, err := GeminiCachedContentUpstreamModelResource(account, mappedModel) + if err != nil { + return nil, err + } + payload := make(map[string]any, len(req.payload)+3) + for k, v := range req.payload { + payload[k] = v + } + payload["model"] = upstreamModel + payload["ttl"] = fmt.Sprintf("%ds", int64(math.Ceil(ttl.Seconds()))) + payload["displayName"] = geminiCachedContentUpstreamTagPrefix + publicID + body, err := json.Marshal(payload) + if err != nil { + return nil, err + } + base, err := s.apiBase(account) + if err != nil { + return nil, err + } + collectionURL, err := geminiCachedContentCollectionURL(account, base, mappedModel) + if err != nil { + return nil, err + } + respBody, err := s.doUpstream(ctx, account, http.MethodPost, collectionURL, body) + if err != nil { + return nil, err + } + upstream, err := parseGeminiCachedContentUpstream(respBody) + if err != nil { + return nil, err + } + if upstream.TotalTokenCount <= 0 { + return nil, &GeminiCachedContentMalformedResponseError{Name: upstream.Name, Err: errors.New("missing usageMetadata.totalTokenCount")} + } + return upstream, nil +} + +// PatchUpstream 把上游缓存的有效期改为自现在起 ttl。 +func (s *GeminiCachedContentService) PatchUpstream(ctx context.Context, account *Account, record *GeminiCachedContent, ttl time.Duration) (*GeminiCachedContentUpstream, error) { + base, err := s.apiBase(account) + if err != nil { + return nil, err + } + resourceURL, err := geminiCachedContentResourceURL(account, base, record.UpstreamName) + if err != nil { + return nil, err + } + seconds := int64(math.Ceil(ttl.Seconds())) + body, _ := json.Marshal(map[string]any{"ttl": fmt.Sprintf("%ds", seconds)}) + respBody, err := s.doUpstream(ctx, account, http.MethodPatch, resourceURL+"?updateMask=ttl", body) + if err != nil { + return nil, err + } + out := &GeminiCachedContentUpstream{Name: record.UpstreamName} + if raw := gjson.GetBytes(respBody, "expireTime").String(); raw != "" { + if t, err := time.Parse(time.RFC3339Nano, raw); err == nil { + out.ExpireTime = t + } + } + return out, nil +} + +// DeleteUpstream 删除上游缓存。 +func (s *GeminiCachedContentService) DeleteUpstream(ctx context.Context, account *Account, record *GeminiCachedContent) error { + base, err := s.apiBase(account) + if err != nil { + return err + } + resourceURL, err := geminiCachedContentResourceURL(account, base, record.UpstreamName) + if err != nil { + return err + } + _, err = s.doUpstream(ctx, account, http.MethodDelete, resourceURL, nil) + return err +} + +// doUpstream 调用上游缓存接口并返回 2xx 响应体。上游调用与客户端连接脱钩：客户端断开不会中途取消 +// 已发出的创建 / 修改，调用方据结果完成落库、计费或清理。 +func (s *GeminiCachedContentService) doUpstream(ctx context.Context, account *Account, method, target string, body []byte) ([]byte, error) { + if s.compat == nil || s.compat.httpUpstream == nil { + return nil, errors.New("gemini upstream client not configured") + } + ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), geminiCachedContentUpstreamTimeout) + defer cancel() + var reader io.Reader + if body != nil { + reader = bytes.NewReader(body) + } + req, err := http.NewRequestWithContext(ctx, method, target, reader) + if err != nil { + return nil, err + } + if body != nil { + req.Header.Set("Content-Type", "application/json") + } + switch account.Type { + case AccountTypeAPIKey: + apiKey := strings.TrimSpace(account.GetCredential("api_key")) + if apiKey == "" { + return nil, errors.New("gemini api_key not configured") + } + req.Header.Set("x-goog-api-key", apiKey) + case AccountTypeServiceAccount: + if s.compat.tokenProvider == nil { + return nil, errors.New("gemini token provider not configured") + } + accessToken, err := s.compat.tokenProvider.GetAccessToken(ctx, account) + if err != nil { + return nil, err + } + req.Header.Set("Authorization", "Bearer "+accessToken) + default: + return nil, fmt.Errorf("account type %s does not support explicit context caching", account.Type) + } + proxyURL := "" + if account.ProxyID != nil && account.Proxy != nil { + proxyURL = account.Proxy.URL() + } + resp, err := s.compat.httpUpstream.Do(req, proxyURL, account.ID, account.Concurrency) + if err != nil { + return nil, err + } + defer func() { _ = resp.Body.Close() }() + respBody, err := io.ReadAll(io.LimitReader(resp.Body, geminiCachedContentResponseMaxBytes)) + success := resp.StatusCode >= 200 && resp.StatusCode < 300 + if err != nil { + if success { + return nil, &GeminiCachedContentMalformedResponseError{Err: err} + } + return nil, err + } + if !success { + return nil, &GeminiCachedContentUpstreamError{StatusCode: resp.StatusCode, Body: respBody} + } + return respBody, nil +} + +// apiBase 返回 API Key 账号经地址校验后的上游根地址；Vertex 账号返回空串（地址由区域决定）。 +func (s *GeminiCachedContentService) apiBase(account *Account) (string, error) { + if account.Type != AccountTypeAPIKey { + return "", nil + } + if s.compat == nil { + return "", errors.New("gemini upstream client not configured") + } + return s.compat.validateUpstreamBaseURL(geminiCachedContentAPIBase(account)) +} + +func parseGeminiCachedContentUpstream(body []byte) (*GeminiCachedContentUpstream, error) { + name := strings.TrimSpace(gjson.GetBytes(body, "name").String()) + if name == "" { + return nil, &GeminiCachedContentMalformedResponseError{Err: errors.New("missing name")} + } + if !geminiCachedContentAIStudioName.MatchString(name) && !geminiCachedContentVertexName.MatchString(name) { + return nil, &GeminiCachedContentMalformedResponseError{Err: fmt.Errorf("unexpected name %q", name)} + } + out := &GeminiCachedContentUpstream{ + Name: name, + TotalTokenCount: gjson.GetBytes(body, "usageMetadata.totalTokenCount").Int(), + } + if raw := gjson.GetBytes(body, "expireTime").String(); raw != "" { + if t, err := time.Parse(time.RFC3339Nano, raw); err == nil { + out.ExpireTime = t + } + } + return out, nil +} + +func geminiCachedContentCollectionURL(account *Account, base, mappedModel string) (string, error) { + switch account.Type { + case AccountTypeAPIKey: + base = strings.TrimRight(strings.TrimSpace(base), "/") + if base == "" { + return "", errors.New("gemini base url is required") + } + return base + "/v1beta/cachedContents", nil + case AccountTypeServiceAccount: + projectID := strings.TrimSpace(account.VertexProjectID()) + if projectID == "" { + return "", errors.New("vertex project_id is required") + } + location := strings.TrimSpace(account.VertexLocation(mappedModel)) + if location == "" { + location = vertexDefaultLocation + } + if !vertexLocationPattern.MatchString(location) { + return "", fmt.Errorf("invalid vertex location: %s", location) + } + return fmt.Sprintf("https://%s/v1/projects/%s/locations/%s/cachedContents", + geminiCachedContentVertexHost(location), url.PathEscape(projectID), url.PathEscape(location)), nil + default: + return "", fmt.Errorf("account type %s does not support explicit context caching", account.Type) + } +} + +func geminiCachedContentResourceURL(account *Account, base, upstreamName string) (string, error) { + switch account.Type { + case AccountTypeAPIKey: + if !geminiCachedContentAIStudioName.MatchString(upstreamName) { + return "", fmt.Errorf("invalid upstream cached content name: %s", upstreamName) + } + base = strings.TrimRight(strings.TrimSpace(base), "/") + if base == "" { + return "", errors.New("gemini base url is required") + } + return base + "/v1beta/" + upstreamName, nil + case AccountTypeServiceAccount: + m := geminiCachedContentVertexName.FindStringSubmatch(upstreamName) + if m == nil { + return "", fmt.Errorf("invalid upstream cached content name: %s", upstreamName) + } + return fmt.Sprintf("https://%s/v1/%s", geminiCachedContentVertexHost(m[1]), upstreamName), nil + default: + return "", fmt.Errorf("account type %s does not support explicit context caching", account.Type) + } +} + +func geminiCachedContentVertexHost(location string) string { + if location == "global" { + return "aiplatform.googleapis.com" + } + return location + "-aiplatform.googleapis.com" +} + +// GeminiCachedContentView 返回客户端可见的缓存元数据。 +func GeminiCachedContentView(record *GeminiCachedContent) map[string]any { + view := map[string]any{ + "name": record.PublicName(), + "model": "models/" + record.Model, + "usageMetadata": map[string]any{"totalTokenCount": record.TotalTokenCount}, + "createTime": record.CreatedAt.UTC().Format(time.RFC3339Nano), + "updateTime": record.UpdatedAt.UTC().Format(time.RFC3339Nano), + "expireTime": record.ExpireTime.UTC().Format(time.RFC3339Nano), + } + if record.DisplayName != "" { + view["displayName"] = record.DisplayName + } + return view +} + +// GeminiExplicitCacheExclusions 返回显式缓存请求在 gemini 分组内需要排除的可调度账号。 +// boundAccountID > 0 时只保留该账号；否则排除平台或类型不支持显式缓存的账号 +// （调度快照不含 base_url，官方地址判定在选中后对完整账号进行）。 +func (s *GatewayService) GeminiExplicitCacheExclusions(ctx context.Context, groupID *int64, boundAccountID int64) (map[int64]struct{}, error) { + accounts, _, err := s.listSchedulableAccounts(ctx, groupID, PlatformGemini, false) + if err != nil { + return nil, err + } + excluded := make(map[int64]struct{}, len(accounts)) + for i := range accounts { + account := &accounts[i] + keep := IsGeminiExplicitCacheAccountType(account) + if boundAccountID > 0 { + keep = account.ID == boundAccountID + } + if !keep { + excluded[account.ID] = struct{}{} + } + } + return excluded, nil +} diff --git a/backend/internal/service/gemini_cached_content_hardening_test.go b/backend/internal/service/gemini_cached_content_hardening_test.go new file mode 100644 index 0000000000..cd9c6fa57a --- /dev/null +++ b/backend/internal/service/gemini_cached_content_hardening_test.go @@ -0,0 +1,368 @@ +//go:build unit + +package service + +import ( + "context" + "errors" + "fmt" + "io" + "net/http" + "strings" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +const geminiCacheHardeningPricing = `{ + "gemini-3.8-flash":{"litellm_provider":"gemini","input_cost_per_token":0.00000075,"output_cost_per_token":0.00000375, + "cache_storage_cost_per_token_per_hour":0.0000005}, + "gemini-2.5-pro":{"litellm_provider":"gemini","input_cost_per_token":0.00000125,"output_cost_per_token":0.00001, + "cache_storage_cost_per_token_per_hour":0.0000045}}` + +func newGeminiCacheBillingGatewayForTest(t *testing.T) (*GatewayService, *openAIRecordUsageLogRepoStub) { + t.Helper() + usageRepo := &openAIRecordUsageLogRepoStub{inserted: true} + svc := newGatewayRecordUsageServiceForTest(usageRepo, &openAIRecordUsageUserRepoStub{}, &openAIRecordUsageSubRepoStub{}) + pricing := &PricingService{} + data, err := pricing.parsePricingData([]byte(geminiCacheHardeningPricing)) + require.NoError(t, err) + pricing.pricingData = data + svc.billingService = NewBillingService(svc.cfg, pricing) + return svc, usageRepo +} + +func TestGeminiCachedContentStoragePriceFollowsBillingModel(t *testing.T) { + ctx := context.Background() + svc, usageRepo := newGeminiCacheBillingGatewayForTest(t) + groupID := int64(33) + apiKey := &APIKey{ID: 821, GroupID: &groupID, User: &User{ID: 621}, + Group: &Group{ID: groupID, Platform: PlatformGemini, RateMultiplier: 2, Hydrated: true}} + tokenHours := GeminiCachedContentStorageTokenHours(1000, 24*time.Hour) + + mapped := &RecordUsageInput{ + Result: &ForwardResult{RequestID: "cc-mapped", Model: "gemini-2.5-pro", UpstreamModel: "gemini-3.8-flash", + Usage: ClaudeUsage{InputTokens: 1000}, CacheStorageTokenHours: tokenHours}, + APIKey: apiKey, User: apiKey.User, Account: &Account{ID: 713, Platform: PlatformGemini}, + } + require.True(t, svc.CacheStoragePriced(ctx, mapped)) + want := (1000*0.00000125 + tokenHours*0.0000045) * 2 + require.InDelta(t, want, svc.QuoteUsageCost(ctx, mapped), 1e-12, "存储与 token 同按计费模型（请求模型）定价") + require.NoError(t, svc.RecordUsage(ctx, mapped)) + require.InDelta(t, want, usageRepo.lastLog.ActualCost, 1e-12, "报价与实际计费同口径") + + alias := &RecordUsageInput{ + Result: &ForwardResult{Model: "customer-alias", UpstreamModel: "gemini-3.8-flash", CacheStorageTokenHours: tokenHours}, + APIKey: apiKey, + } + require.True(t, svc.CacheStoragePriced(ctx, alias)) + require.InDelta(t, tokenHours*0.0000005*2, svc.QuoteUsageCost(ctx, alias), 1e-12, "计费模型无存储价时回退到实际转发模型") + + unpriced := &RecordUsageInput{ + Result: &ForwardResult{Model: "customer-alias", UpstreamModel: "unknown-upstream", CacheStorageTokenHours: tokenHours}, + APIKey: apiKey, + } + require.False(t, svc.CacheStoragePriced(ctx, unpriced)) +} + +func TestGeminiCachedContentStorageOnlyIsNotARequest(t *testing.T) { + ctx := context.Background() + svc, usageRepo := newGeminiCacheBillingGatewayForTest(t) + groupID := int64(34) + perRequest := 0.05 + pricing := []ChannelModelPricing{{Models: []string{"gemini-3.8-flash"}, BillingMode: BillingModePerRequest, PerRequestPrice: &perRequest}} + svc.channelService = newTestChannelServiceForStats(t, &Channel{ID: 1, Status: StatusActive, + AccountStatsPricingRules: []AccountStatsPricingRule{{GroupIDs: []int64{groupID}, Pricing: pricing}}}, groupID, PlatformGemini) + svc.resolver = NewModelPricingResolver(svc.channelService, svc.billingService) + apiKey := &APIKey{ID: 822, GroupID: &groupID, User: &User{ID: 622}, + Group: &Group{ID: groupID, Platform: PlatformGemini, RateMultiplier: 1, Hydrated: true, ModelPricing: pricing}} + tokenHours := 20000.0 + storage := tokenHours * 0.0000005 + + record := func(tokens int) *UsageLog { + require.NoError(t, svc.RecordUsage(ctx, &RecordUsageInput{ + Result: &ForwardResult{RequestID: fmt.Sprintf("cc-%d", tokens), Model: "gemini-3.8-flash", + Usage: ClaudeUsage{InputTokens: tokens}, CacheStorageTokenHours: tokenHours}, + APIKey: apiKey, User: apiKey.User, Account: &Account{ID: 712, Platform: PlatformGemini}, + })) + return usageRepo.lastLog + } + + created := record(10000) + require.InDelta(t, perRequest+storage, created.TotalCost, 1e-12, "创建是一次请求：按次价 + 存储费") + require.NotNil(t, created.AccountStatsCost) + require.InDelta(t, perRequest+storage, *created.AccountStatsCost, 1e-12) + + extended := record(0) + require.InDelta(t, storage, extended.TotalCost, 1e-12, "延长有效期不算一次请求，只收存储费") + require.Nil(t, extended.AccountStatsCost, "账号统计成本回落 total_cost，不套按次规则") +} + +func TestFindGeminiCachedContentReferenceRejectsAmbiguousKeys(t *testing.T) { + for name, body := range map[string]string{ + "duplicate top-level": `{"cachedContent":"cachedContents/a","contents":[],"cachedContent":"cachedContents/b"}`, + "snake top-level": `{"cached_content":"cachedContents/a"}`, + "snake beside camel": `{"cachedContent":"cachedContents/a","cached_content":"cachedContents/b"}`, + "duplicate nested": `{"generateContentRequest":{"cachedContent":"cachedContents/a","cachedContent":"cachedContents/b"}}`, + "snake nested": `{"generateContentRequest":{"cached_content":"cachedContents/a"}}`, + "snake parent with cache": `{"generate_content_request":{"cachedContent":"cachedContents/a"}}`, + "duplicate nested parent": `{"generateContentRequest":{},"generateContentRequest":{"cachedContent":"cachedContents/a"}}`, + "conflicting two locations": `{"cachedContent":"cachedContents/a","generateContentRequest":{"cachedContent":"cachedContents/b"}}`, + } { + t.Run(name, func(t *testing.T) { + _, err := FindGeminiCachedContentReference([]byte(body)) + require.Error(t, err) + }) + } + + for name, body := range map[string]string{ + "snake parent without cache": `{"generate_content_request":{"contents":[]}}`, + "key name inside content": `{"contents":[{"role":"user","cached_content":"x","parts":[{"text":"cached_content"}]}]}`, + "not an object": `[{"cachedContent":"cachedContents/a"}]`, + } { + t.Run("allows "+name, func(t *testing.T) { + ref, err := FindGeminiCachedContentReference([]byte(body)) + require.NoError(t, err) + require.Nil(t, ref) + }) + } + + ref, err := FindGeminiCachedContentReference([]byte(`{"cachedContent":"cachedContents/a","generateContentRequest":{"cachedContent":"cachedContents/a"}}`)) + require.NoError(t, err) + require.Equal(t, []string{"cachedContent", "generateContentRequest.cachedContent"}, ref.Paths) +} + +func TestGeminiCachedContentPayloadNormalization(t *testing.T) { + now := time.Date(2026, 10, 9, 8, 0, 0, 0, time.UTC) + maxTTL := 24 * time.Hour + + req, err := ParseGeminiCachedContentCreateRequest([]byte(`{"model":"gemini-3.8-flash","display_name":"docs","expire_time":"2026-10-09T10:00:00Z"}`), now, maxTTL) + require.NoError(t, err) + require.Equal(t, "docs", req.DisplayName) + require.Equal(t, 2*time.Hour, req.TTL) + for _, key := range []string{"display_name", "expire_time", "displayName", "expireTime"} { + require.NotContains(t, req.payload, key, "元数据字段不随内容发往上游") + } + + for name, body := range map[string]string{ + "both spellings": `{"model":"gemini-3.8-flash","displayName":"a","display_name":"b"}`, + "trailing data": `{"model":"gemini-3.8-flash"} {"model":"x"}`, + } { + t.Run("rejects "+name, func(t *testing.T) { + _, err := ParseGeminiCachedContentCreateRequest([]byte(body), now, maxTTL) + require.Error(t, err) + }) + } + + ttl, err := ParseGeminiCachedContentPatchRequest([]byte(`{"expire_time":"2026-10-09T09:00:00Z"}`), "", now, maxTTL) + require.NoError(t, err) + require.Equal(t, time.Hour, ttl) + _, err = ParseGeminiCachedContentPatchRequest([]byte(`{"display_name":"x"}`), "", now, maxTTL) + require.ErrorContains(t, err, "immutable") +} + +func TestGeminiCachedContentTTLRoundsUpToWholeSeconds(t *testing.T) { + now := time.Date(2026, 10, 9, 8, 0, 0, 0, time.UTC) + req, err := ParseGeminiCachedContentCreateRequest([]byte(`{"model":"gemini-3.8-flash","ttl":"1e-10s"}`), now, 24*time.Hour) + require.NoError(t, err) + require.Equal(t, time.Second, req.TTL, "极小的显式有效期按 1 秒处理，不回落到默认 1 小时") + require.True(t, req.ttlExplicit) + + ttl, err := ParseGeminiCachedContentPatchRequest([]byte(`{"ttl":"1e-10s"}`), "", now, 24*time.Hour) + require.NoError(t, err) + require.Equal(t, time.Second, ttl) + + req, err = ParseGeminiCachedContentCreateRequest([]byte(`{"model":"gemini-3.8-flash","expireTime":"2026-10-09T08:00:10.2Z"}`), now, 24*time.Hour) + require.NoError(t, err) + require.Equal(t, 11*time.Second, req.TTL, "expireTime 折算的有效期与发往上游的整秒 ttl 一致") +} + +func TestGeminiCachedContentCreateKeepsNumbersVerbatim(t *testing.T) { + svc, httpStub := newGeminiCachedContentServiceForTest(t, http.StatusOK, + `{"name":"cachedContents/up1","usageMetadata":{"totalTokenCount":2000}}`) + req, err := ParseGeminiCachedContentCreateRequest([]byte(`{"model":"models/gemini-3.8-flash","contents":[{"role":"user","parts":[ + {"functionResponse":{"name":"lookup","response":{"order_id":1790000000000000001,"ratio":0.10}}}]}]}`), time.Now(), 24*time.Hour) + require.NoError(t, err) + account := &Account{ID: 9, Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "secret"}} + _, err = svc.CreateUpstream(context.Background(), account, req, "gemini-3.8-flash", "pub1", time.Hour) + require.NoError(t, err) + raw, err := io.ReadAll(httpStub.lastReq.Body) + require.NoError(t, err) + require.Contains(t, string(raw), `"order_id":1790000000000000001`) + require.Contains(t, string(raw), `"ratio":0.10`) +} + +func TestGeminiCachedContentCreateUpstreamUnusableResponse(t *testing.T) { + account := &Account{ID: 9, Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "secret"}} + for name, tc := range map[string]struct { + body string + wantName string + }{ + "missing token count": {`{"name":"cachedContents/up9","expireTime":"2099-01-01T00:00:00Z"}`, "cachedContents/up9"}, + "zero token count": {`{"name":"cachedContents/up9","usageMetadata":{"totalTokenCount":0}}`, "cachedContents/up9"}, + "unexpected name": {`{"name":"cachedContents/a/../b","usageMetadata":{"totalTokenCount":5}}`, ""}, + "not json": {``, ""}, + } { + t.Run(name, func(t *testing.T) { + svc, _ := newGeminiCachedContentServiceForTest(t, http.StatusOK, tc.body) + req, err := ParseGeminiCachedContentCreateRequest([]byte(`{"model":"gemini-3.8-flash"}`), time.Now(), 24*time.Hour) + require.NoError(t, err) + _, err = svc.CreateUpstream(context.Background(), account, req, "gemini-3.8-flash", "pub1", time.Hour) + var malformed *GeminiCachedContentMalformedResponseError + require.ErrorAs(t, err, &malformed) + require.Equal(t, tc.wantName, malformed.Name) + }) + } +} + +type geminiCachedContentCtxProbeUpstream struct { + geminiCompatHTTPUpstreamStub + ctxErrAtSend error +} + +func (s *geminiCachedContentCtxProbeUpstream) Do(req *http.Request, proxyURL string, accountID int64, accountConcurrency int) (*http.Response, error) { + s.ctxErrAtSend = req.Context().Err() + return s.geminiCompatHTTPUpstreamStub.Do(req, proxyURL, accountID, accountConcurrency) +} + +func TestGeminiCachedContentUpstreamCallIgnoresClientCancel(t *testing.T) { + up := &geminiCachedContentCtxProbeUpstream{geminiCompatHTTPUpstreamStub: geminiCompatHTTPUpstreamStub{response: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"name":"cachedContents/up1","usageMetadata":{"totalTokenCount":2000}}`)), + }}} + svc := NewGeminiCachedContentService(newGeminiCachedContentRepoStub(), + &GeminiMessagesCompatService{httpUpstream: up, cfg: &config.Config{}}, &config.Config{}) + req, err := ParseGeminiCachedContentCreateRequest([]byte(`{"model":"gemini-3.8-flash"}`), time.Now(), 24*time.Hour) + require.NoError(t, err) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + account := &Account{ID: 9, Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "secret"}} + _, err = svc.CreateUpstream(ctx, account, req, "gemini-3.8-flash", "pub1", time.Hour) + require.NoError(t, err) + require.NoError(t, up.ctxErrAtSend, "客户端断开不取消发往上游的创建") +} + +func TestClampGeminiCachedContentExpire(t *testing.T) { + limit := time.Date(2026, 10, 10, 8, 0, 0, 0, time.UTC) + require.Equal(t, limit, ClampGeminiCachedContentExpire(time.Time{}, limit)) + require.Equal(t, limit, ClampGeminiCachedContentExpire(limit.Add(100*time.Hour), limit), "上游回报的到期时间不超过请求的有效期") + earlier := limit.Add(-time.Second) + require.Equal(t, earlier, ClampGeminiCachedContentExpire(earlier, limit)) +} + +type geminiBoundCacheRateLimitRepo struct { + errorPolicyRepoStub + rateLimitedCalls int +} + +func (r *geminiBoundCacheRateLimitRepo) SetRateLimited(context.Context, int64, time.Time) error { + r.rateLimitedCalls++ + return nil +} + +func TestGeminiForwardNative_BoundCacheUpstreamGatewayErrorsSkipAccountPenalty(t *testing.T) { + gin.SetMode(gin.TestMode) + for _, status := range []int{http.StatusTooManyRequests, http.StatusServiceUnavailable} { + t.Run(http.StatusText(status), func(t *testing.T) { + body := fmt.Sprintf(`{"error":{"code":%d,"message":"The account holding this cached content is temporarily unavailable","status":"UNAVAILABLE"}}`, status) + svc, httpStub, _ := newGeminiCachedContentTestService(status, body) + repo := &geminiBoundCacheRateLimitRepo{} + svc.accountRepo = repo + svc.rateLimitService = NewRateLimitService(repo, nil, &config.Config{}, nil, nil) + account := &Account{ID: 711, Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: map[string]any{ + "api_key": "k", "base_url": "https://gw.example.com", "explicit_cache_upstream": true, + }} + c, _ := newGeminiNativeTestContext(t) + + _, err := svc.ForwardNative(WithGeminiCachedContentBound(context.Background()), c, account, + "gemini-3.8-flash", "generateContent", false, []byte(geminiCachedContentRequestBody)) + + var failoverErr *UpstreamFailoverError + require.True(t, errors.As(err, &failoverErr), "交由绑定请求的耗尽处理透传") + require.Equal(t, status, failoverErr.StatusCode) + require.JSONEq(t, body, string(failoverErr.ResponseBody)) + require.False(t, failoverErr.RetryableOnSameAccount) + require.Equal(t, 1, httpStub.calls, "不在同一账号上退避重试") + require.Zero(t, repo.rateLimitedCalls, "不限流整个网关账号") + require.Zero(t, repo.tempCalls) + require.Zero(t, repo.setErrCalls) + }) + } +} + +type chargeAffordableCacheStub struct { + billingCacheWorkerStub + balance float64 + sub *SubscriptionCacheData + rate *APIKeyRateLimitCacheData +} + +func (s *chargeAffordableCacheStub) GetUserBalance(context.Context, int64) (float64, error) { + return s.balance, nil +} + +func (s *chargeAffordableCacheStub) GetSubscriptionCache(context.Context, int64, int64) (*SubscriptionCacheData, error) { + if s.sub == nil { + return nil, errors.New("cache miss") + } + return s.sub, nil +} + +func (s *chargeAffordableCacheStub) GetAPIKeyRateLimit(context.Context, int64) (*APIKeyRateLimitCacheData, error) { + if s.rate == nil { + return nil, errors.New("cache miss") + } + return s.rate, nil +} + +func TestCheckChargeAffordable(t *testing.T) { + ctx := context.Background() + user := &User{ID: 1} + newSvc := func(cache *chargeAffordableCacheStub) *BillingCacheService { + svc := NewBillingCacheService(cache, nil, nil, nil, nil, nil, &config.Config{}, nil) + t.Cleanup(svc.Stop) + return svc + } + + t.Run("balance must cover the amount", func(t *testing.T) { + svc := newSvc(&chargeAffordableCacheStub{balance: 10}) + require.NoError(t, svc.CheckChargeAffordable(ctx, user, nil, nil, nil, 10)) + require.ErrorIs(t, svc.CheckChargeAffordable(ctx, user, nil, nil, nil, 10.01), ErrInsufficientBalance) + require.NoError(t, svc.CheckChargeAffordable(ctx, user, nil, nil, nil, 0)) + }) + + t.Run("subscription windows must cover the amount", func(t *testing.T) { + daily := 5.0 + group := &Group{ID: 2, SubscriptionType: SubscriptionTypeSubscription, DailyLimitUSD: &daily} + svc := newSvc(&chargeAffordableCacheStub{sub: &SubscriptionCacheData{ + Status: SubscriptionStatusActive, ExpiresAt: time.Now().Add(time.Hour), DailyUsage: 3, + }}) + sub := &UserSubscription{ID: 1} + require.NoError(t, svc.CheckChargeAffordable(ctx, user, nil, group, sub, 2)) + require.ErrorIs(t, svc.CheckChargeAffordable(ctx, user, nil, group, sub, 2.5), ErrDailyLimitExceeded) + }) + + t.Run("api key quota and rate windows include the amount", func(t *testing.T) { + svc := newSvc(&chargeAffordableCacheStub{balance: 100, rate: &APIKeyRateLimitCacheData{ + Usage1d: 8, Window1d: time.Now().Unix(), + }}) + quotaKey := &APIKey{ID: 5, Quota: 10, QuotaUsed: 9} + require.NoError(t, svc.CheckChargeAffordable(ctx, user, quotaKey, nil, nil, 1)) + require.ErrorIs(t, svc.CheckChargeAffordable(ctx, user, quotaKey, nil, nil, 1.5), ErrAPIKeyQuotaExhausted) + + windowKey := &APIKey{ID: 6, RateLimit1d: 10} + require.NoError(t, svc.CheckChargeAffordable(ctx, user, windowKey, nil, nil, 2)) + require.ErrorIs(t, svc.CheckChargeAffordable(ctx, user, windowKey, nil, nil, 2.5), ErrAPIKeyRateLimit1dExceeded) + }) + + t.Run("simple mode skips the check", func(t *testing.T) { + cfg := &config.Config{RunMode: config.RunModeSimple} + svc := NewBillingCacheService(&chargeAffordableCacheStub{}, nil, nil, nil, nil, nil, cfg, nil) + t.Cleanup(svc.Stop) + require.NoError(t, svc.CheckChargeAffordable(ctx, user, nil, nil, nil, 1000)) + }) +} diff --git a/backend/internal/service/gemini_cached_content_not_found_test.go b/backend/internal/service/gemini_cached_content_not_found_test.go index 0bfade0855..783ba71c72 100644 --- a/backend/internal/service/gemini_cached_content_not_found_test.go +++ b/backend/internal/service/gemini_cached_content_not_found_test.go @@ -74,7 +74,7 @@ func TestGeminiForwardNative_CachedContentNotFoundPassthroughWithoutAccountPenal require.Equal(t, 0, repo.tempCalls, "不应临时停调度") require.Equal(t, 1, httpStub.calls, "不应重试") require.Equal(t, http.StatusForbidden, rec.Code) - require.Equal(t, geminiCachedContentNotFoundBody, rec.Body.String(), "响应体应原样透传") + require.Equal(t, GeminiCachedContentNotFoundResponse, rec.Body.String(), "应返回标准缓存不存在响应") }) } } @@ -119,11 +119,64 @@ func TestGeminiForwardNative_OtherForbiddenStillPenalizesAccount(t *testing.T) { require.Equal(t, 1, repo.setErrCalls, "其他 403 仍标记账号错误") } +// Vertex 对已删除 / 已过期缓存的响应（2026-10 实测）：生成请求 400，缓存资源接口 404，文案含上游缓存数字 ID。 +const ( + geminiVertexCachedContentInvalidStateBody = `{"error":{"code":400,"message":"Invalid resource state for cache content 8079767076522164224.","status":"FAILED_PRECONDITION"}}` + geminiVertexCachedContentNotFoundBody = `{"error":{"code":404,"message":"Cached content 8079767076522164224 is not found.","status":"NOT_FOUND"}}` +) + func TestIsGeminiCachedContentNotFound(t *testing.T) { require.True(t, isGeminiCachedContentNotFound(http.StatusForbidden, []byte(geminiCachedContentNotFoundBody))) require.True(t, isGeminiCachedContentNotFound(http.StatusForbidden, []byte(`[{"error":{"code":403,"message":"CachedContent not found (or permission denied)","status":"PERMISSION_DENIED"}}]`))) + require.True(t, isGeminiCachedContentNotFound(http.StatusBadRequest, []byte(geminiVertexCachedContentInvalidStateBody))) + require.True(t, isGeminiCachedContentNotFound(http.StatusNotFound, []byte(geminiVertexCachedContentNotFoundBody))) + require.False(t, isGeminiCachedContentNotFound(http.StatusBadRequest, []byte(geminiCachedContentNotFoundBody))) require.False(t, isGeminiCachedContentNotFound(http.StatusForbidden, []byte(`{"error":{"code":403,"message":"Permission denied: Consumer has been suspended.","status":"PERMISSION_DENIED"}}`))) + require.False(t, isGeminiCachedContentNotFound(http.StatusNotFound, + []byte(`{"error":{"code":404,"message":"Publisher Model `+"`projects/p/locations/global/publishers/google/models/gemini-9`"+` was not found.","status":"NOT_FOUND"}}`))) + require.False(t, isGeminiCachedContentNotFound(http.StatusBadRequest, + []byte(`{"error":{"code":400,"message":"Tool config, tools and system instruction should not be set in the request when using cached content.","status":"INVALID_ARGUMENT"}}`))) + require.False(t, isGeminiCachedContentNotFound(http.StatusInternalServerError, []byte(geminiVertexCachedContentNotFoundBody))) +} + +func TestGeminiForwardNative_VertexCachedContentNotFoundReturnsStandardForbidden(t *testing.T) { + gin.SetMode(gin.TestMode) + for _, tc := range []struct { + name string + status int + body string + action string + stream bool + }{ + {name: "generate invalid resource state", status: http.StatusBadRequest, body: geminiVertexCachedContentInvalidStateBody, action: "generateContent"}, + {name: "stream invalid resource state", status: http.StatusBadRequest, body: geminiVertexCachedContentInvalidStateBody, action: "streamGenerateContent", stream: true}, + {name: "generate not found", status: http.StatusNotFound, body: geminiVertexCachedContentNotFoundBody, action: "generateContent"}, + } { + t.Run(tc.name, func(t *testing.T) { + svc, httpStub, repo := newGeminiCachedContentTestService(tc.status, tc.body) + svc.tokenProvider = NewGeminiTokenProvider(nil, &fakeGeminiTokenCache{token: "ya29.test-token"}, nil) + c, rec := newGeminiNativeTestContext(t) + account := vertexServiceAccount() + account.ID = 811 + + result, err := svc.ForwardNative(context.Background(), c, account, + "gemini-3.8-flash", tc.action, tc.stream, + []byte(`{"cachedContent":"projects/proj-1/locations/global/cachedContents/8079767076522164224","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`)) + + require.Nil(t, result) + require.Error(t, err) + var failoverErr *UpstreamFailoverError + require.False(t, errors.As(err, &failoverErr), "缓存不存在是请求问题，不应换号") + require.Equal(t, 0, repo.setErrCalls) + require.Equal(t, 0, repo.tempCalls) + require.Empty(t, repo.modelRateLimitCalls, "404 不应按模型不存在记模型限流") + require.Equal(t, 1, httpStub.calls, "不应重试") + require.Equal(t, http.StatusForbidden, rec.Code) + require.Equal(t, GeminiCachedContentNotFoundResponse, rec.Body.String()) + require.NotContains(t, rec.Body.String(), "8079767076522164224", "不应暴露上游缓存 ID") + }) + } } diff --git a/backend/internal/service/gemini_cached_content_test.go b/backend/internal/service/gemini_cached_content_test.go new file mode 100644 index 0000000000..f53ad6c94d --- /dev/null +++ b/backend/internal/service/gemini_cached_content_test.go @@ -0,0 +1,506 @@ +//go:build unit + +package service + +import ( + "context" + "encoding/json" + "io" + "net/http" + "strings" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/stretchr/testify/require" +) + +type geminiCachedContentRepoStub struct { + records map[string]*GeminiCachedContent + nextID int64 + deleted []int64 +} + +func newGeminiCachedContentRepoStub() *geminiCachedContentRepoStub { + return &geminiCachedContentRepoStub{records: map[string]*GeminiCachedContent{}} +} + +func (r *geminiCachedContentRepoStub) Create(_ context.Context, record *GeminiCachedContent) error { + r.nextID++ + record.ID = r.nextID + r.records[record.PublicID] = record + return nil +} + +func (r *geminiCachedContentRepoStub) GetForOwner(_ context.Context, apiKeyID int64, publicID string) (*GeminiCachedContent, error) { + record, ok := r.records[publicID] + if !ok || record.APIKeyID != apiKeyID { + return nil, ErrGeminiCachedContentNotFound + } + return record, nil +} + +func (r *geminiCachedContentRepoStub) ListForOwner(_ context.Context, apiKeyID, groupID int64, activeAt time.Time, beforeID int64, limit int) ([]*GeminiCachedContent, error) { + var out []*GeminiCachedContent + for _, record := range r.records { + if record.APIKeyID == apiKeyID && record.GroupID == groupID && record.ExpireTime.After(activeAt) && (beforeID <= 0 || record.ID < beforeID) { + out = append(out, record) + } + } + if len(out) > limit { + out = out[:limit] + } + return out, nil +} + +func (r *geminiCachedContentRepoStub) UpdateExpireTime(_ context.Context, id int64, expireTime time.Time) (time.Time, error) { + for _, record := range r.records { + if record.ID == id { + previous := record.ExpireTime + record.ExpireTime = expireTime + return previous, nil + } + } + return time.Time{}, ErrGeminiCachedContentNotFound +} + +func (r *geminiCachedContentRepoStub) SoftDelete(_ context.Context, id int64) error { + r.deleted = append(r.deleted, id) + for key, record := range r.records { + if record.ID == id { + delete(r.records, key) + } + } + return nil +} + +func TestParseGeminiCachedContentCreateRequest(t *testing.T) { + now := time.Date(2026, 10, 9, 8, 0, 0, 0, time.UTC) + maxTTL := 24 * time.Hour + + t.Run("defaults to one hour and strips output-only fields", func(t *testing.T) { + req, err := ParseGeminiCachedContentCreateRequest([]byte(`{"model":"models/gemini-3.8-flash","displayName":"docs", + "contents":[{"role":"user","parts":[{"text":"x"}]}],"name":"cachedContents/forged","usageMetadata":{"totalTokenCount":1}}`), now, maxTTL) + require.NoError(t, err) + require.Equal(t, "gemini-3.8-flash", req.Model) + require.Equal(t, time.Hour, req.TTL) + require.Equal(t, int64(3600), req.TTLSeconds()) + require.Equal(t, "docs", req.DisplayName) + require.NotContains(t, req.payload, "name") + require.NotContains(t, req.payload, "usageMetadata") + require.Contains(t, req.payload, "contents") + }) + + t.Run("default is capped by a shorter max ttl", func(t *testing.T) { + req, err := ParseGeminiCachedContentCreateRequest([]byte(`{"model":"gemini-3.8-flash"}`), now, 30*time.Minute) + require.NoError(t, err) + require.Equal(t, 30*time.Minute, req.TTL) + }) + + t.Run("fractional ttl rounds up to whole seconds", func(t *testing.T) { + req, err := ParseGeminiCachedContentCreateRequest([]byte(`{"model":"gemini-3.8-flash","ttl":"3.5s"}`), now, maxTTL) + require.NoError(t, err) + require.Equal(t, int64(4), req.TTLSeconds()) + }) + + t.Run("expireTime converts to ttl", func(t *testing.T) { + req, err := ParseGeminiCachedContentCreateRequest([]byte(`{"model":"gemini-3.8-flash","expireTime":"2026-10-09T10:00:00Z"}`), now, maxTTL) + require.NoError(t, err) + require.Equal(t, 2*time.Hour, req.TTL) + }) + + for name, body := range map[string]string{ + "missing model": `{"ttl":"60s"}`, + "vertex style model": `{"model":"projects/p/locations/l/publishers/google/models/x"}`, + "ttl and expireTime": `{"model":"m","ttl":"60s","expireTime":"2026-10-09T10:00:00Z"}`, + "ttl over max": `{"model":"m","ttl":"86401s"}`, + "non positive ttl": `{"model":"m","ttl":"0s"}`, + "ttl without unit": `{"model":"m","ttl":"60"}`, + "expireTime in past": `{"model":"m","expireTime":"2026-10-09T07:00:00Z"}`, + "displayName too long": `{"model":"m","displayName":"` + strings.Repeat("字", 129) + `"}`, + "not an object": `[]`, + } { + t.Run("rejects "+name, func(t *testing.T) { + _, err := ParseGeminiCachedContentCreateRequest([]byte(body), now, maxTTL) + require.Error(t, err) + }) + } +} + +func TestParseGeminiCachedContentPatchRequest(t *testing.T) { + now := time.Date(2026, 10, 9, 8, 0, 0, 0, time.UTC) + ttl, err := ParseGeminiCachedContentPatchRequest([]byte(`{"ttl":"7200s"}`), "ttl", now, 24*time.Hour) + require.NoError(t, err) + require.Equal(t, 2*time.Hour, ttl) + + ttl, err = ParseGeminiCachedContentPatchRequest([]byte(`{"expireTime":"2026-10-09T09:00:00Z"}`), "", now, 24*time.Hour) + require.NoError(t, err) + require.Equal(t, time.Hour, ttl) + + _, err = ParseGeminiCachedContentPatchRequest([]byte(`{"displayName":"x"}`), "displayName", now, 24*time.Hour) + require.ErrorContains(t, err, "immutable") + _, err = ParseGeminiCachedContentPatchRequest([]byte(`{"ttl":"60s"}`), "ttl,displayName", now, 24*time.Hour) + require.ErrorContains(t, err, "immutable") + _, err = ParseGeminiCachedContentPatchRequest([]byte(`{}`), "", now, 24*time.Hour) + require.Error(t, err) + _, err = ParseGeminiCachedContentPatchRequest([]byte(`{"ttl":"90000s"}`), "", now, 24*time.Hour) + require.ErrorContains(t, err, "maximum") +} + +func TestGeminiCachedContentReferenceRewrite(t *testing.T) { + ref, err := FindGeminiCachedContentReference([]byte(`{"contents":[]}`)) + require.NoError(t, err) + require.Nil(t, ref) + + body := []byte(`{"cachedContent":"cachedContents/abc","contents":[]}`) + ref, err = FindGeminiCachedContentReference(body) + require.NoError(t, err) + require.Equal(t, "cachedContents/abc", ref.Name) + out, err := RewriteGeminiCachedContentReference(body, ref, "cachedContents/upstream1") + require.NoError(t, err) + require.JSONEq(t, `{"cachedContent":"cachedContents/upstream1","contents":[]}`, string(out)) + + countBody := []byte(`{"generateContentRequest":{"model":"models/m","cachedContent":"cachedContents/abc","contents":[]}}`) + ref, err = FindGeminiCachedContentReference(countBody) + require.NoError(t, err) + out, err = RewriteGeminiCachedContentReference(countBody, ref, "projects/1/locations/us-central1/cachedContents/9") + require.NoError(t, err) + require.JSONEq(t, `{"generateContentRequest":{"model":"models/m","cachedContent":"projects/1/locations/us-central1/cachedContents/9","contents":[]}}`, string(out)) + + _, err = FindGeminiCachedContentReference([]byte(`{"cachedContent":"cachedContents/a","generateContentRequest":{"cachedContent":"cachedContents/b"}}`)) + require.Error(t, err) + _, err = FindGeminiCachedContentReference([]byte(`{"cachedContent":1}`)) + require.Error(t, err) +} + +func TestIsGeminiExplicitCacheAccount(t *testing.T) { + cases := []struct { + name string + account *Account + want bool + }{ + {"official api key", &Account{Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "k"}}, true}, + {"explicit official base url", &Account{Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "k", "base_url": "https://generativelanguage.googleapis.com/"}}, true}, + {"reseller base url", &Account{Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "k", "base_url": "https://relay.example.com"}}, false}, + {"upstream gateway with explicit cache switch", &Account{Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "k", "base_url": "https://relay.example.com", "explicit_cache_upstream": true}}, true}, + {"official host with path prefix", &Account{Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "k", "base_url": "https://generativelanguage.googleapis.com/antigravity"}}, false}, + {"missing api key", &Account{Platform: PlatformGemini, Type: AccountTypeAPIKey}, false}, + {"vertex service account", &Account{Platform: PlatformGemini, Type: AccountTypeServiceAccount}, true}, + {"oauth", &Account{Platform: PlatformGemini, Type: AccountTypeOAuth}, false}, + {"antigravity api key", &Account{Platform: PlatformAntigravity, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "k"}}, false}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + require.Equal(t, tc.want, IsGeminiExplicitCacheAccount(tc.account)) + }) + } +} + +func TestGeminiCachedContentVertexURLs(t *testing.T) { + account := &Account{Platform: PlatformGemini, Type: AccountTypeServiceAccount, Credentials: map[string]any{ + "project_id": "proj-1", + "location": "us-central1", + "vertex_model_locations": map[string]any{"gemini-3.8-flash": "global"}, + }} + model, err := GeminiCachedContentUpstreamModelResource(account, "gemini-3.8-flash") + require.NoError(t, err) + require.Equal(t, "projects/proj-1/locations/global/publishers/google/models/gemini-3.8-flash", model) + + collection, err := geminiCachedContentCollectionURL(account, "", "gemini-3.8-flash") + require.NoError(t, err) + require.Equal(t, "https://aiplatform.googleapis.com/v1/projects/proj-1/locations/global/cachedContents", collection) + + collection, err = geminiCachedContentCollectionURL(account, "", "gemini-3.1-pro-preview") + require.NoError(t, err) + require.Equal(t, "https://us-central1-aiplatform.googleapis.com/v1/projects/proj-1/locations/us-central1/cachedContents", collection) + + resource, err := geminiCachedContentResourceURL(account, "", "projects/123/locations/europe-west4/cachedContents/456") + require.NoError(t, err) + require.Equal(t, "https://europe-west4-aiplatform.googleapis.com/v1/projects/123/locations/europe-west4/cachedContents/456", resource) + + _, err = geminiCachedContentResourceURL(account, "", "projects/123/locations/europe-west4/cachedContents/../x") + require.Error(t, err) + _, err = geminiCachedContentResourceURL(&Account{Type: AccountTypeAPIKey}, "https://generativelanguage.googleapis.com", "cachedContents/a/../b") + require.Error(t, err) +} + +func newGeminiCachedContentServiceForTest(t *testing.T, status int, respBody string) (*GeminiCachedContentService, *geminiCompatHTTPUpstreamStub) { + t.Helper() + httpStub := &geminiCompatHTTPUpstreamStub{response: &http.Response{ + StatusCode: status, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(respBody)), + }} + compat := &GeminiMessagesCompatService{httpUpstream: httpStub, cfg: &config.Config{}} + return NewGeminiCachedContentService(newGeminiCachedContentRepoStub(), compat, &config.Config{}), httpStub +} + +func TestGeminiCachedContentCreateUpstreamAIStudio(t *testing.T) { + svc, httpStub := newGeminiCachedContentServiceForTest(t, http.StatusOK, + `{"name":"cachedContents/up123","model":"models/gemini-3.8-flash","usageMetadata":{"totalTokenCount":6273},"expireTime":"2026-10-09T09:50:26.591575799Z"}`) + account := &Account{ID: 9, Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "secret"}} + req, err := ParseGeminiCachedContentCreateRequest([]byte(`{"model":"models/gemini-3.8-flash","displayName":"user label", + "systemInstruction":{"parts":[{"text":"sys"}]},"contents":[{"role":"user","parts":[{"text":"doc"}]}]}`), time.Now(), 24*time.Hour) + require.NoError(t, err) + + upstream, err := svc.CreateUpstream(context.Background(), account, req, "gemini-3.8-flash", "pub1", req.TTL) + require.NoError(t, err) + require.Equal(t, "cachedContents/up123", upstream.Name) + require.Equal(t, int64(6273), upstream.TotalTokenCount) + require.Equal(t, 2026, upstream.ExpireTime.Year()) + + sent := httpStub.lastReq + require.Equal(t, http.MethodPost, sent.Method) + require.Equal(t, "https://generativelanguage.googleapis.com/v1beta/cachedContents", sent.URL.String()) + require.Equal(t, "secret", sent.Header.Get("x-goog-api-key")) + raw, err := io.ReadAll(sent.Body) + require.NoError(t, err) + var payload map[string]any + require.NoError(t, json.Unmarshal(raw, &payload)) + require.Equal(t, "models/gemini-3.8-flash", payload["model"]) + require.Equal(t, "3600s", payload["ttl"]) + require.Equal(t, "s2a-pub1", payload["displayName"], "上游 displayName 写网关标记，用户标签只存本地") + require.Contains(t, payload, "systemInstruction") + require.Contains(t, payload, "contents") +} + +func TestGeminiCachedContentUpstreamErrorClassification(t *testing.T) { + notFoundBody := `{"error":{"code":403,"message":"CachedContent not found (or permission denied)","status":"PERMISSION_DENIED"}}` + svc, _ := newGeminiCachedContentServiceForTest(t, http.StatusForbidden, notFoundBody) + account := &Account{ID: 9, Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "secret"}} + record := &GeminiCachedContent{UpstreamName: "cachedContents/up123"} + + err := svc.DeleteUpstream(context.Background(), account, record) + var upErr *GeminiCachedContentUpstreamError + require.ErrorAs(t, err, &upErr) + require.True(t, upErr.NotFound()) + require.True(t, upErr.RetryableOnOtherAccount()) + + quota := &GeminiCachedContentUpstreamError{StatusCode: http.StatusTooManyRequests} + require.True(t, quota.RetryableOnOtherAccount()) + badRequest := &GeminiCachedContentUpstreamError{StatusCode: http.StatusBadRequest} + require.False(t, badRequest.RetryableOnOtherAccount()) + require.False(t, badRequest.NotFound()) +} + +func TestGeminiCachedContentResolve(t *testing.T) { + repo := newGeminiCachedContentRepoStub() + now := time.Date(2026, 10, 9, 8, 0, 0, 0, time.UTC) + svc := NewGeminiCachedContentService(repo, nil, &config.Config{}) + svc.now = func() time.Time { return now } + groupID := int64(5) + require.NoError(t, repo.Create(context.Background(), &GeminiCachedContent{ + PublicID: "abc123", APIKeyID: 1, GroupID: groupID, ExpireTime: now.Add(time.Hour), + })) + require.NoError(t, repo.Create(context.Background(), &GeminiCachedContent{ + PublicID: "expired1", APIKeyID: 1, GroupID: groupID, ExpireTime: now.Add(-time.Second), + })) + + record, err := svc.Resolve(context.Background(), 1, &groupID, "cachedContents/abc123") + require.NoError(t, err) + require.Equal(t, "abc123", record.PublicID) + + otherGroup := int64(6) + for name, call := range map[string]func() error{ + "other api key": func() error { + _, err := svc.Resolve(context.Background(), 2, &groupID, "cachedContents/abc123") + return err + }, + "other group": func() error { + _, err := svc.Resolve(context.Background(), 1, &otherGroup, "cachedContents/abc123") + return err + }, + "expired": func() error { + _, err := svc.Resolve(context.Background(), 1, &groupID, "cachedContents/expired1") + return err + }, + "malformed": func() error { + _, err := svc.Resolve(context.Background(), 1, &groupID, "cachedContents/../x") + return err + }, + } { + t.Run(name, func(t *testing.T) { + require.ErrorIs(t, call(), ErrGeminiCachedContentNotFound) + }) + } +} + +func TestGeminiCachedContentPublicID(t *testing.T) { + id, err := NewGeminiCachedContentPublicID() + require.NoError(t, err) + require.Len(t, id, 40) + parsed, ok := ParseGeminiCachedContentPublicID(GeminiCachedContentNamePrefix + id) + require.True(t, ok) + require.Equal(t, id, parsed) +} + +func TestGeminiCachedContentStoragePricingAndSurcharge(t *testing.T) { + pricing := &PricingService{} + data, err := pricing.parsePricingData([]byte(`{"gemini-3.8-flash":{"litellm_provider":"gemini","input_cost_per_token":0.00000075, + "output_cost_per_token":0.00000375,"cache_read_input_token_cost":0.000000075,"cache_storage_cost_per_token_per_hour":0.0000005}}`)) + require.NoError(t, err) + pricing.pricingData = data + billing := NewBillingService(&config.Config{}, pricing) + + price := billing.CacheStoragePricePerTokenHour("models/gemini-3.8-flash") + require.InDelta(t, 0.0000005, price, 1e-15) + require.Zero(t, billing.CacheStoragePricePerTokenHour("gemini-unknown-model")) + + tokenHours := GeminiCachedContentStorageTokenHours(10000, 2*time.Hour) + require.InDelta(t, 20000, tokenHours, 1e-9) + storage := tokenHours * price + require.InDelta(t, 0.01, storage, 1e-12) + + gateway := &GatewayService{billingService: billing} + groupID := int64(3) + apiKey := &APIKey{GroupID: &groupID, Group: &Group{ID: groupID, Hydrated: true}} + cost := gateway.calculateRecordUsageCost(context.Background(), &ForwardResult{ + Model: "gemini-3.8-flash", + Usage: ClaudeUsage{InputTokens: 10000}, + CacheStorageTokenHours: tokenHours, + }, apiKey, "gemini-3.8-flash", 2, 2, time.Now()) + inputCost := 10000 * 0.00000075 + require.InDelta(t, inputCost+storage, cost.TotalCost, 1e-12) + require.InDelta(t, (inputCost+storage)*2, cost.ActualCost, 1e-12) + + storageOnly := gateway.calculateRecordUsageCost(context.Background(), &ForwardResult{ + Model: "gemini-3.8-flash", + CacheStorageTokenHours: tokenHours, + }, apiKey, "gemini-3.8-flash", 1.5, 1.5, time.Now()) + require.InDelta(t, storage, storageOnly.TotalCost, 1e-12) + require.InDelta(t, storage*1.5, storageOnly.ActualCost, 1e-12) +} + +func TestGeminiCachedContentStorageCostInAccountStats(t *testing.T) { + const ( + model = "gemini-3.8-flash" + groupID = int64(31) + inputTokens = 10000 + storage = 0.01 + inputPrice = 0.00000075 + rulePrice = 0.0000002 + ) + newPricing := func() *PricingService { + pricing := &PricingService{} + data, err := pricing.parsePricingData([]byte(`{"gemini-3.8-flash":{"litellm_provider":"gemini","input_cost_per_token":0.00000075, + "output_cost_per_token":0.00000375,"cache_storage_cost_per_token_per_hour":0.0000005}}`)) + require.NoError(t, err) + pricing.pricingData = data + return pricing + } + officialInput := inputTokens * inputPrice + + for _, tc := range []struct { + name string + channel *Channel + tokens int + wantStats float64 + }{ + {name: "no_channel", tokens: inputTokens, wantStats: officialInput + storage}, + {name: "model_file_pricing", channel: &Channel{ID: 1, Status: StatusActive}, tokens: inputTokens, wantStats: officialInput + storage}, + {name: "apply_pricing_to_account_stats", channel: &Channel{ID: 1, Status: StatusActive, ApplyPricingToAccountStats: true}, + tokens: inputTokens, wantStats: officialInput + storage}, + {name: "custom_rule", channel: &Channel{ID: 1, Status: StatusActive, AccountStatsPricingRules: []AccountStatsPricingRule{{ + GroupIDs: []int64{groupID}, + Pricing: []ChannelModelPricing{{Models: []string{model}, InputPrice: testPtrFloat64(rulePrice)}}, + }}}, tokens: inputTokens, wantStats: inputTokens*rulePrice + storage}, + {name: "storage_only_with_custom_rule", channel: &Channel{ID: 1, Status: StatusActive, AccountStatsPricingRules: []AccountStatsPricingRule{{ + GroupIDs: []int64{groupID}, + Pricing: []ChannelModelPricing{{Models: []string{model}, InputPrice: testPtrFloat64(rulePrice)}}, + }}}, wantStats: storage}, + } { + t.Run(tc.name, func(t *testing.T) { + usageRepo := &openAIRecordUsageLogRepoStub{inserted: true} + svc := newGatewayRecordUsageServiceForTest(usageRepo, &openAIRecordUsageUserRepoStub{}, &openAIRecordUsageSubRepoStub{}) + svc.billingService = NewBillingService(svc.cfg, newPricing()) + if tc.channel != nil { + svc.channelService = newTestChannelServiceForStats(t, tc.channel, groupID, PlatformGemini) + } + groupRate, accountRate := 1.5, 0.6 + gid := groupID + err := svc.RecordUsage(context.Background(), &RecordUsageInput{ + Result: &ForwardResult{ + RequestID: "gemini-cached-content:create:" + tc.name, + Model: model, + Usage: ClaudeUsage{InputTokens: tc.tokens}, + CacheStorageTokenHours: storage / 0.0000005, + }, + APIKey: &APIKey{ID: 811, GroupID: &gid, Group: &Group{ID: groupID, Platform: PlatformGemini, RateMultiplier: groupRate, Hydrated: true}}, + User: &User{ID: 611}, + Account: &Account{ID: 711, Platform: PlatformGemini, RateMultiplier: &accountRate}, + }) + require.NoError(t, err) + log := usageRepo.lastLog + require.NotNil(t, log) + + wantTotal := float64(tc.tokens)*inputPrice + storage + require.InDelta(t, wantTotal, log.TotalCost, 1e-12) + require.InDelta(t, wantTotal*groupRate, log.ActualCost, 1e-12, "售价 = 官方价 × 分组倍率") + + statsCost := log.TotalCost + if log.AccountStatsCost != nil { + statsCost = *log.AccountStatsCost + } + require.InDelta(t, tc.wantStats, statsCost, 1e-12, "账号统计成本须含存储费且只计一次") + require.NotNil(t, log.AccountRateMultiplier) + require.InDelta(t, tc.wantStats*accountRate, statsCost**log.AccountRateMultiplier, 1e-12, "成本 = 统计成本 × 账号倍率") + }) + } +} + +func TestGeminiExplicitCacheUpstreamTTL(t *testing.T) { + now := time.Date(2026, 10, 9, 8, 0, 0, 0, time.UTC) + upstream := &Account{Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: map[string]any{ + "api_key": "k", "base_url": "https://relay.example.com", "explicit_cache_upstream": true, + "explicit_cache_upstream_max_ttl_seconds": float64(1800), + }} + official := &Account{Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: map[string]any{ + "api_key": "k", "explicit_cache_upstream_max_ttl_seconds": float64(1800), + }} + require.Equal(t, 30*time.Minute, upstream.GeminiExplicitCacheUpstreamMaxTTL()) + require.Zero(t, official.GeminiExplicitCacheUpstreamMaxTTL(), "未开启开关时上限字段不生效") + + svc := NewGeminiCachedContentService(newGeminiCachedContentRepoStub(), nil, &config.Config{}) + require.Equal(t, 30*time.Minute, svc.MaxTTLForAccount(upstream)) + require.Equal(t, 24*time.Hour, svc.MaxTTLForAccount(official)) + + implicit, err := ParseGeminiCachedContentCreateRequest([]byte(`{"model":"gemini-3.8-flash"}`), now, 24*time.Hour) + require.NoError(t, err) + ttl, ok := implicit.TTLForAccount(upstream) + require.True(t, ok) + require.Equal(t, 30*time.Minute, ttl, "未显式指定有效期时按上游上限创建") + ttl, ok = implicit.TTLForAccount(official) + require.True(t, ok) + require.Equal(t, time.Hour, ttl) + + explicit, err := ParseGeminiCachedContentCreateRequest([]byte(`{"model":"gemini-3.8-flash","ttl":"3600s"}`), now, 24*time.Hour) + require.NoError(t, err) + _, ok = explicit.TTLForAccount(upstream) + require.False(t, ok, "显式有效期超过上游上限时该账号不可用") + ttl, ok = explicit.TTLForAccount(official) + require.True(t, ok) + require.Equal(t, time.Hour, ttl) +} + +func TestGeminiCachedContentCreateUpstreamGateway(t *testing.T) { + svc, httpStub := newGeminiCachedContentServiceForTest(t, http.StatusOK, + `{"name":"cachedContents/gatewayid0001","model":"models/gemini-3.8-flash","usageMetadata":{"totalTokenCount":2730},"expireTime":"2026-10-09T09:00:00Z"}`) + account := &Account{ID: 11, Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: map[string]any{ + "api_key": "downstream-key", "base_url": "https://upstream.example.com/gw/", "explicit_cache_upstream": true, + }} + req, err := ParseGeminiCachedContentCreateRequest([]byte(`{"model":"gemini-3.8-flash","contents":[{"role":"user","parts":[{"text":"doc"}]}]}`), time.Now(), 24*time.Hour) + require.NoError(t, err) + + upstream, err := svc.CreateUpstream(context.Background(), account, req, "gemini-3.8-flash", "pub2", 30*time.Minute) + require.NoError(t, err) + require.Equal(t, "cachedContents/gatewayid0001", upstream.Name) + require.Equal(t, "https://upstream.example.com/gw/v1beta/cachedContents", httpStub.lastReq.URL.String()) + require.Equal(t, "downstream-key", httpStub.lastReq.Header.Get("x-goog-api-key")) + raw, err := io.ReadAll(httpStub.lastReq.Body) + require.NoError(t, err) + require.Contains(t, string(raw), `"ttl":"1800s"`) + + resource, err := geminiCachedContentResourceURL(account, geminiCachedContentAPIBase(account), upstream.Name) + require.NoError(t, err) + require.Equal(t, "https://upstream.example.com/gw/v1beta/cachedContents/gatewayid0001", resource) +} diff --git a/backend/internal/service/gemini_messages_compat_service.go b/backend/internal/service/gemini_messages_compat_service.go index 01c3b2cc07..60421d1f7f 100644 --- a/backend/internal/service/gemini_messages_compat_service.go +++ b/backend/internal/service/gemini_messages_compat_service.go @@ -1374,8 +1374,9 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin. return nil, transportErr } - // 引用缓存不存在先于错误策略识别：错误策略会写临时不可调度，重试与换号也无法恢复。 - if notFound, rebuilt := s.cachedContentNotFoundInLoop(resp); notFound { + // 引用缓存不存在、上游网关持有缓存的账号不可用先于错误策略识别：错误策略会写临时不可调度， + // 重试与换号也无法恢复。 + if passthrough, rebuilt := s.cachedContentPassthroughInLoop(ctx, account, resp); passthrough { resp = rebuilt break } else { @@ -1475,9 +1476,13 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin. if resp.StatusCode >= 400 { respBody := s.readUpstreamErrorBody(resp) - // 引用缓存不存在是请求问题：原样返回客户端，不改账号状态、不换号。 + // 引用缓存不存在是请求问题：统一返回 403，不改账号状态、不换号。 if isGeminiCachedContentNotFound(resp.StatusCode, respBody) { - return nil, s.writeGeminiNativeUpstreamError(c, account, resp, respBody, requestID, isOAuth) + return nil, s.writeGeminiCachedContentNotFound(c, account, resp, respBody, requestID, isOAuth) + } + // 上游网关的限流 / 不可用只针对其持有缓存的账号：不改本账号状态，交由绑定请求的耗尽处理透传。 + if isGeminiBoundCacheUpstreamPassthrough(ctx, account, resp.StatusCode) { + return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: respBody} } // Best-effort fallback for OAuth tokens missing AI Studio scopes when calling countTokens. // This avoids Gemini SDKs failing hard during preflight token counting. @@ -1688,9 +1693,15 @@ func (s *GeminiMessagesCompatService) checkErrorPolicyInLoop( return policy != ErrorPolicyNone, rebuilt } -// cachedContentNotFoundInLoop 判断上游响应是否为引用缓存不存在的 403；读取过的响应体以内存副本重建。 -func (s *GeminiMessagesCompatService) cachedContentNotFoundInLoop(resp *http.Response) (notFound bool, rebuilt *http.Response) { - if resp.StatusCode != http.StatusForbidden { +// cachedContentPassthroughInLoop 判断上游响应是否跳过同号重试与错误策略直接处理：引用缓存不存在， +// 或绑定缓存的请求在上游网关账号上收到的 429 / 5xx。读取过的响应体以内存副本重建。 +func (s *GeminiMessagesCompatService) cachedContentPassthroughInLoop(ctx context.Context, account *Account, resp *http.Response) (passthrough bool, rebuilt *http.Response) { + if isGeminiBoundCacheUpstreamPassthrough(ctx, account, resp.StatusCode) { + return true, resp + } + switch resp.StatusCode { + case http.StatusBadRequest, http.StatusForbidden, http.StatusNotFound: + default: return false, resp } body := s.readUpstreamErrorBody(resp) @@ -1808,25 +1819,7 @@ func (s *GeminiMessagesCompatService) writeGeminiCustomCodeSkippedError(c *gin.C // 并记录 ops 错误事件。状态码保真：下游据此区分请求级错误与可重试的链路故障。 func (s *GeminiMessagesCompatService) writeGeminiNativeUpstreamError(c *gin.Context, account *Account, resp *http.Response, respBody []byte, requestID string, isOAuth bool) error { respBody = unwrapIfNeeded(isOAuth, respBody) - upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody)) - upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) - upstreamDetail := s.upstreamErrorDetail(respBody) - if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody { - logger.LegacyPrintf("service.gemini_messages_compat", "[Gemini] native upstream error %d: %s", resp.StatusCode, truncateForLog(respBody, s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes)) - } - setOpsUpstreamError(c, resp.StatusCode, upstreamMsg, upstreamDetail) - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - ProxyID: opsUpstreamProxyID(account), - ProxyName: opsUpstreamProxyName(account), - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: resp.StatusCode, - UpstreamRequestID: requestID, - Kind: "http_error", - Message: upstreamMsg, - Detail: upstreamDetail, - }) + upstreamMsg := s.recordGeminiNativeUpstreamError(c, account, resp.StatusCode, respBody, requestID) contentType := resp.Header.Get("Content-Type") if contentType == "" { @@ -1840,6 +1833,40 @@ func (s *GeminiMessagesCompatService) writeGeminiNativeUpstreamError(c *gin.Cont return fmt.Errorf("gemini upstream error: %d message=%s", resp.StatusCode, upstreamMsg) } +// writeGeminiCachedContentNotFound 记录上游原始错误，并向客户端统一返回缓存不存在的 403： +// 各上游的响应形态不同，且 Vertex 的文案含上游缓存资源 ID。 +func (s *GeminiMessagesCompatService) writeGeminiCachedContentNotFound(c *gin.Context, account *Account, resp *http.Response, respBody []byte, requestID string, isOAuth bool) error { + respBody = unwrapIfNeeded(isOAuth, respBody) + upstreamMsg := s.recordGeminiNativeUpstreamError(c, account, resp.StatusCode, respBody, requestID) + MarkResponseCommitted(c) + c.Data(http.StatusForbidden, "application/json", []byte(GeminiCachedContentNotFoundResponse)) + return fmt.Errorf("gemini upstream error: %d message=%s", resp.StatusCode, upstreamMsg) +} + +// recordGeminiNativeUpstreamError 把上游错误响应登记到 ops 上下文，返回脱敏后的上游错误文案。 +func (s *GeminiMessagesCompatService) recordGeminiNativeUpstreamError(c *gin.Context, account *Account, statusCode int, respBody []byte, requestID string) string { + upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody)) + upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) + upstreamDetail := s.upstreamErrorDetail(respBody) + if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody { + logger.LegacyPrintf("service.gemini_messages_compat", "[Gemini] native upstream error %d: %s", statusCode, truncateForLog(respBody, s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes)) + } + setOpsUpstreamError(c, statusCode, upstreamMsg, upstreamDetail) + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + ProxyID: opsUpstreamProxyID(account), + ProxyName: opsUpstreamProxyName(account), + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: statusCode, + UpstreamRequestID: requestID, + Kind: "http_error", + Message: upstreamMsg, + Detail: upstreamDetail, + }) + return upstreamMsg +} + func sleepGeminiBackoff(attempt int) { delay := geminiRetryBaseDelay * time.Duration(1< geminiRetryMaxDelay { @@ -2662,13 +2689,23 @@ func isGeminiInsufficientScope(headers http.Header, body []byte) bool { return strings.Contains(lower, "insufficient authentication scopes") || strings.Contains(lower, "access_token_scope_insufficient") } -// isGeminiCachedContentNotFound 识别请求引用的显式缓存不存在、已过期、已删除或不属于当前项目时 -// 上游返回的 403。四种情况上游响应完全相同，均由请求的 cachedContent 决定，与账号凭据状态无关。 +// isGeminiCachedContentNotFound 识别请求引用的显式缓存不存在（含已过期、已删除、不属于当前项目）时上游的响应， +// 这类响应由请求的 cachedContent 决定，与账号凭据状态无关： +// - AI Studio 及本系统网关：一律 403 "CachedContent not found (or permission denied)"； +// - Vertex 生成请求：400 "Invalid resource state for cache content {id}."； +// - Vertex 缓存资源 GET / PATCH / DELETE：404 "Cached content {id} is not found."。 func isGeminiCachedContentNotFound(statusCode int, body []byte) bool { - if statusCode != http.StatusForbidden { + lower := strings.ToLower(string(body)) + switch statusCode { + case http.StatusForbidden: + return strings.Contains(lower, "cachedcontent not found") + case http.StatusBadRequest: + return strings.Contains(lower, "invalid resource state for cache content") + case http.StatusNotFound: + return strings.Contains(lower, "cached content") && strings.Contains(lower, "not found") + default: return false } - return strings.Contains(strings.ToLower(string(body)), "cachedcontent not found") } func estimateGeminiCountTokens(reqBody []byte) int { diff --git a/backend/internal/service/pricing_service.go b/backend/internal/service/pricing_service.go index f2a1eee9f7..45b530fd6d 100644 --- a/backend/internal/service/pricing_service.go +++ b/backend/internal/service/pricing_service.go @@ -223,6 +223,8 @@ type LiteLLMModelPricing struct { OutputCostPerImageToken float64 `json:"output_cost_per_image_token"` // 图片输出 token 价格 InputCostPerImageToken float64 `json:"input_cost_per_image_token"` // 图片输入 token 价格（如 gpt-image-2 图片编辑） CacheReadInputImageTokenCost float64 `json:"cache_read_input_image_token_cost"` + // CacheStorageCostPerTokenPerHour 显式上下文缓存存储单价（每 token 每小时，USD），仅由定价覆盖文件提供。 + CacheStorageCostPerTokenPerHour float64 `json:"cache_storage_cost_per_token_per_hour"` // TokenPricingAbsent 表示源数据中 input/output token 价格均缺失（仅有图片价）。 // 此类条目只可用于图片计费，token 计费必须回退到 fallback 或 fail-closed， @@ -258,6 +260,7 @@ type LiteLLMRawEntry struct { OutputCostPerImageToken *float64 `json:"output_cost_per_image_token"` InputCostPerImageToken *float64 `json:"input_cost_per_image_token"` CacheReadInputImageTokenCost *float64 `json:"cache_read_input_image_token_cost"` + CacheStorageCostPerTokenPerHour *float64 `json:"cache_storage_cost_per_token_per_hour"` } // PricingService 动态价格服务 @@ -730,6 +733,9 @@ func (s *PricingService) parsePricingData(body []byte) (map[string]*LiteLLMModel if entry.CacheReadInputImageTokenCost != nil { pricing.CacheReadInputImageTokenCost = *entry.CacheReadInputImageTokenCost } + if entry.CacheStorageCostPerTokenPerHour != nil { + pricing.CacheStorageCostPerTokenPerHour = *entry.CacheStorageCostPerTokenPerHour + } hasExplicitLongContext := entry.LongContextInputTokenThreshold != nil || entry.LongContextInputCostMultiplier != nil || diff --git a/backend/internal/service/wire.go b/backend/internal/service/wire.go index ec3058dc66..56fe96f52d 100644 --- a/backend/internal/service/wire.go +++ b/backend/internal/service/wire.go @@ -885,6 +885,7 @@ var ProviderSet = wire.NewSet( ProvideGrokOAuthService, wire.Bind(new(GrokOAuthTokenService), new(*GrokOAuthService)), NewGeminiOAuthService, + NewGeminiCachedContentService, NewGeminiQuotaService, NewCompositeTokenCacheInvalidator, wire.Bind(new(TokenCacheInvalidator), new(*CompositeTokenCacheInvalidator)), diff --git a/backend/migrations/243_gemini_cached_contents.sql b/backend/migrations/243_gemini_cached_contents.sql new file mode 100644 index 0000000000..d3cbda27df --- /dev/null +++ b/backend/migrations/243_gemini_cached_contents.sql @@ -0,0 +1,30 @@ +-- Gemini 显式缓存（cachedContents）的网关资源绑定。 +-- public_id 是网关对客户端暴露的缓存 ID；upstream_name 是上游资源名， +-- 引用缓存的请求只能路由到 account_id 对应的账号。 +CREATE TABLE IF NOT EXISTS gemini_cached_contents ( + id BIGSERIAL PRIMARY KEY, + public_id VARCHAR(64) NOT NULL, + user_id BIGINT NOT NULL, + api_key_id BIGINT NOT NULL, + group_id BIGINT NOT NULL, + account_id BIGINT NOT NULL, + upstream_name VARCHAR(512) NOT NULL, + model VARCHAR(255) NOT NULL, + upstream_model VARCHAR(512) NOT NULL, + display_name VARCHAR(512) NOT NULL DEFAULT '', + total_token_count BIGINT NOT NULL DEFAULT 0, + expire_time TIMESTAMPTZ NOT NULL, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + deleted_at TIMESTAMPTZ NULL +); + +CREATE UNIQUE INDEX IF NOT EXISTS idx_gemini_cached_contents_public_id + ON gemini_cached_contents (public_id); + +CREATE INDEX IF NOT EXISTS idx_gemini_cached_contents_api_key_active + ON gemini_cached_contents (api_key_id, id DESC) + WHERE deleted_at IS NULL; + +CREATE INDEX IF NOT EXISTS idx_gemini_cached_contents_account + ON gemini_cached_contents (account_id); diff --git a/deploy/config.example.yaml b/deploy/config.example.yaml index 7066f17208..174d5a4437 100644 --- a/deploy/config.example.yaml +++ b/deploy/config.example.yaml @@ -253,6 +253,9 @@ gateway: # Enable Gemini upstream response header debug logs (default: false) # 是否开启 Gemini 上游响应头调试日志（默认 false） gemini_debug_response_headers: false + # Maximum TTL (seconds) a Gemini explicit context cache may be created or extended to (default: 86400) + # Gemini 显式上下文缓存单次可设置的最长有效期（秒，默认 86400）；存储费在创建/延长请求中一次性收取 + gemini_cached_content_max_ttl_seconds: 86400 # Connection pool isolation strategy: # 连接池隔离策略： # - proxy: Isolate by proxy, same proxy shares connection pool (suitable for few proxies, many accounts) @@ -1102,6 +1105,12 @@ pricing: # 覆盖补丁文件（可选，优先级最高）：条目按字段浅合并覆盖目录/回退数据， # 字段值为 null 表示删除该字段。模型名须与目录键完全一致。 # 修改后在下一次哈希检查时按内容哈希自动重载，无需重启。 + # Gemini explicit context caching requires a storage price per model, e.g. + # {"gemini-3.8-flash": {"cache_storage_cost_per_token_per_hour": 0.0000005}}; + # models without it reject POST /v1beta/cachedContents. + # Gemini 显式缓存需按模型配置存储单价（每 token 每小时，USD）字段 + # cache_storage_cost_per_token_per_hour；未配置的模型拒绝创建缓存。 + # 目录中不存在的模型条目还须同时提供 input_cost_per_token / output_cost_per_token。 # override_file: "./data/model_pricing_overrides.json" # Update interval in hours # 更新间隔（小时） diff --git a/frontend/src/components/account/CreateAccountModal.vue b/frontend/src/components/account/CreateAccountModal.vue index 2ea44ee11c..28ca1b2b00 100644 --- a/frontend/src/components/account/CreateAccountModal.vue +++ b/frontend/src/components/account/CreateAccountModal.vue @@ -1767,6 +1767,48 @@ + +
+
+
+ +

+ {{ t('admin.accounts.explicitCacheUpstreamHint') }} +

+
+ +
+
+

+ + {{ t('admin.accounts.explicitCacheUpstreamInfo') }} +

+
+
+ + +

+ {{ t('admin.accounts.explicitCacheUpstreamMaxTTLHint') }} +

+
+
+
@@ -4005,6 +4047,7 @@ import QuotaLimitCard from '@/components/account/QuotaLimitCard.vue' import Toggle from '@/components/common/Toggle.vue' import GrokBaseUrlPresets from '@/components/account/GrokBaseUrlPresets.vue' import CnBaseUrlPresets from '@/components/account/CnBaseUrlPresets.vue' +import { applyExplicitCacheUpstreamCredentials, isOfficialGeminiBaseUrl } from '@/utils/geminiExplicitCache' import OpenCodeGoProtocolRulesEditor from '@/components/account/OpenCodeGoProtocolRulesEditor.vue' import HeaderOverrideEditor from '@/components/account/HeaderOverrideEditor.vue' import { allSelectedGroupsEnableLongContextPricing } from '@/components/account/longContextBilling' @@ -4473,6 +4516,8 @@ const DEFAULT_POOL_MODE_RETRY_COUNT = 3 const MAX_POOL_MODE_RETRY_COUNT = 10 const DEFAULT_POOL_MODE_RETRY_STATUS_CODES = [401, 403, 429] const poolModeEnabled = ref(false) +const explicitCacheUpstreamEnabled = ref(false) +const explicitCacheUpstreamMaxTTL = ref('') const poolModeRetryCount = ref(DEFAULT_POOL_MODE_RETRY_COUNT) const poolModeRetryStatusCodesInput = ref('') @@ -5458,6 +5503,8 @@ const resetForm = () => { poolModeEnabled.value = false poolModeRetryCount.value = DEFAULT_POOL_MODE_RETRY_COUNT poolModeRetryStatusCodesInput.value = '' + explicitCacheUpstreamEnabled.value = false + explicitCacheUpstreamMaxTTL.value = '' customErrorCodesEnabled.value = false selectedErrorCodes.value = [] customErrorCodeInput.value = null @@ -5963,6 +6010,12 @@ const handleSubmit = async () => { } } + applyExplicitCacheUpstreamCredentials( + credentials, + form.platform === 'gemini' && !isOfficialGeminiBaseUrl(apiKeyBaseUrl.value) && explicitCacheUpstreamEnabled.value, + explicitCacheUpstreamMaxTTL.value + ) + // Add custom error codes if enabled if (customErrorCodesEnabled.value) { credentials.custom_error_codes_enabled = true diff --git a/frontend/src/components/account/EditAccountModal.vue b/frontend/src/components/account/EditAccountModal.vue index 949ac2386b..d5aa3fdaea 100644 --- a/frontend/src/components/account/EditAccountModal.vue +++ b/frontend/src/components/account/EditAccountModal.vue @@ -500,6 +500,48 @@
+ +
+
+
+ +

+ {{ t('admin.accounts.explicitCacheUpstreamHint') }} +

+
+ +
+
+

+ + {{ t('admin.accounts.explicitCacheUpstreamInfo') }} +

+
+
+ + +

+ {{ t('admin.accounts.explicitCacheUpstreamMaxTTLHint') }} +

+
+
+
@@ -3160,6 +3202,11 @@ import ModelWhitelistSelector from '@/components/account/ModelWhitelistSelector. import QuotaLimitCard from '@/components/account/QuotaLimitCard.vue' import GrokBaseUrlPresets from '@/components/account/GrokBaseUrlPresets.vue' import CnBaseUrlPresets from '@/components/account/CnBaseUrlPresets.vue' +import { + applyExplicitCacheUpstreamCredentials, + isOfficialGeminiBaseUrl, + normalizeExplicitCacheUpstreamMaxTTL +} from '@/utils/geminiExplicitCache' import OpenCodeGoProtocolRulesEditor from '@/components/account/OpenCodeGoProtocolRulesEditor.vue' import HeaderOverrideEditor from '@/components/account/HeaderOverrideEditor.vue' import OllamaCloudUsageSettings from '@/components/account/OllamaCloudUsageSettings.vue' @@ -3561,6 +3608,8 @@ const MAX_POOL_MODE_RETRY_COUNT = 10 const DEFAULT_POOL_MODE_RETRY_STATUS_CODES = [401, 403, 429] const GROK_CLIENT_TOOL_CACHE_EXTRA_KEY = 'grok_client_tool_cache_enabled' const poolModeEnabled = ref(false) +const explicitCacheUpstreamEnabled = ref(false) +const explicitCacheUpstreamMaxTTL = ref('') const poolModeRetryCount = ref(DEFAULT_POOL_MODE_RETRY_COUNT) const poolModeRetryStatusCodesInput = ref('') @@ -4500,6 +4549,10 @@ const syncFormFromAccount = (newAccount: Account | null) => { ) poolModeRetryStatusCodesInput.value = formatPoolModeRetryStatusCodes(credentials.pool_mode_retry_status_codes) + explicitCacheUpstreamEnabled.value = credentials.explicit_cache_upstream === true + explicitCacheUpstreamMaxTTL.value = + normalizeExplicitCacheUpstreamMaxTTL(credentials.explicit_cache_upstream_max_ttl_seconds) ?? '' + // Load custom error codes customErrorCodesEnabled.value = credentials.custom_error_codes_enabled === true const existingErrorCodes = credentials.custom_error_codes as number[] | undefined @@ -4573,6 +4626,8 @@ const syncFormFromAccount = (newAccount: Account | null) => { poolModeEnabled.value = false poolModeRetryCount.value = DEFAULT_POOL_MODE_RETRY_COUNT poolModeRetryStatusCodesInput.value = '' + explicitCacheUpstreamEnabled.value = false + explicitCacheUpstreamMaxTTL.value = '' customErrorCodesEnabled.value = false selectedErrorCodes.value = [] } @@ -5299,6 +5354,12 @@ const handleSubmit = async () => { delete newCredentials.pool_mode_retry_status_codes } + applyExplicitCacheUpstreamCredentials( + newCredentials, + props.account.platform === 'gemini' && !isOfficialGeminiBaseUrl(editBaseUrl.value) && explicitCacheUpstreamEnabled.value, + explicitCacheUpstreamMaxTTL.value + ) + // Add custom error codes if enabled if (customErrorCodesEnabled.value) { newCredentials.custom_error_codes_enabled = true diff --git a/frontend/src/components/account/__tests__/CreateAccountModal.spec.ts b/frontend/src/components/account/__tests__/CreateAccountModal.spec.ts index edd67a1d9f..4a317eb174 100644 --- a/frontend/src/components/account/__tests__/CreateAccountModal.spec.ts +++ b/frontend/src/components/account/__tests__/CreateAccountModal.spec.ts @@ -895,3 +895,40 @@ describe('CreateAccountModal OpenAI long-context billing', () => { expect(createOpenAICodexPATMock.mock.calls[0]?.[0]?.extra?.openai_long_context_billing_enabled).toBe(false) }) }) + +describe('CreateAccountModal gemini explicit cache upstream switch', () => { + beforeEach(() => { + authIsSimpleMode.value = true + createAccountMock.mockReset().mockResolvedValue({ id: 43, platform: 'gemini', type: 'apikey' }) + probeUpstreamBillingMock.mockReset().mockResolvedValue({}) + syncUpstreamModelsMock.mockReset().mockResolvedValue({ models: [], metadata: {} }) + }) + + it('only offers the switch for a custom base url and sends it on create', async () => { + const wrapper = mountModal() + await selectButtonByText(wrapper, 'Gemini') + await selectButtonByText(wrapper, 'admin.accounts.gemini.accountType.apiKeyTitle') + expect(wrapper.find('[data-testid="explicit-cache-upstream-section"]').exists()).toBe(false) + + const baseUrlInput = wrapper + .findAll('form#create-account-form input[type="text"]') + .find((input) => (input.element as HTMLInputElement).value === 'https://generativelanguage.googleapis.com') + expect(baseUrlInput).toBeDefined() + await baseUrlInput!.setValue('https://upstream.example.com') + expect(wrapper.find('[data-testid="explicit-cache-upstream-section"]').exists()).toBe(true) + + await wrapper.get('[data-testid="explicit-cache-upstream-toggle"]').trigger('click') + await wrapper.get('[data-testid="explicit-cache-upstream-max-ttl"]').setValue('1800') + await wrapper.get('form#create-account-form input[type="text"]').setValue('gemini upstream gateway') + await wrapper.get('form#create-account-form input[type="password"]').setValue('sk-upstream') + await wrapper.get('form#create-account-form').trigger('submit.prevent') + await flushPromises() + + expect(createAccountMock).toHaveBeenCalledTimes(1) + expect(createAccountMock.mock.calls[0]?.[0]?.credentials).toMatchObject({ + base_url: 'https://upstream.example.com', + explicit_cache_upstream: true, + explicit_cache_upstream_max_ttl_seconds: 1800 + }) + }) +}) diff --git a/frontend/src/components/account/__tests__/EditAccountModal.spec.ts b/frontend/src/components/account/__tests__/EditAccountModal.spec.ts index 2281bacd4d..19402e8389 100644 --- a/frontend/src/components/account/__tests__/EditAccountModal.spec.ts +++ b/frontend/src/components/account/__tests__/EditAccountModal.spec.ts @@ -1834,3 +1834,75 @@ describe('EditAccountModal OpenAI 自动使用重置卡', () => { wrapper.unmount() }) }) + +describe('EditAccountModal gemini explicit cache upstream switch', () => { + function buildGeminiApiKeyAccount(credentials: Record) { + return { + id: 21, + name: 'Gemini upstream gateway', + notes: '', + platform: 'gemini', + type: 'apikey', + credentials: { api_key: 'sk-upstream', ...credentials }, + credentials_status: { has_api_key: true }, + extra: {}, + proxy_id: null, + concurrency: 1, + priority: 1, + rate_multiplier: 1, + status: 'active', + group_ids: [], + expires_at: null, + auto_pause_on_expired: false + } as any + } + + beforeEach(() => { + authIsSimpleMode.value = true + updateAccountMock.mockReset() + checkMixedChannelRiskMock.mockReset().mockResolvedValue({ has_risk: false }) + }) + + it('hides the switch for the official AI Studio address', () => { + const wrapper = mountModal(buildGeminiApiKeyAccount({ base_url: 'https://generativelanguage.googleapis.com' })) + expect(wrapper.find('[data-testid="explicit-cache-upstream-section"]').exists()).toBe(false) + }) + + it('loads and saves the switch and ttl limit for a custom upstream', async () => { + const account = buildGeminiApiKeyAccount({ + base_url: 'https://upstream.example.com', + explicit_cache_upstream: true, + explicit_cache_upstream_max_ttl_seconds: 1800 + }) + updateAccountMock.mockResolvedValue(account) + const wrapper = mountModal(account) + + expect(wrapper.find('[data-testid="explicit-cache-upstream-section"]').exists()).toBe(true) + expect(wrapper.get('[data-testid="explicit-cache-upstream-toggle"]').attributes('aria-checked')).toBe('true') + const ttlInput = wrapper.get('[data-testid="explicit-cache-upstream-max-ttl"]') + expect((ttlInput.element as HTMLInputElement).value).toBe('1800') + + await ttlInput.setValue('3600') + await wrapper.get('form#edit-account-form').trigger('submit.prevent') + expect(updateAccountMock.mock.calls[0]?.[1]?.credentials).toMatchObject({ + explicit_cache_upstream: true, + explicit_cache_upstream_max_ttl_seconds: 3600 + }) + }) + + it('removes both fields when the switch is turned off', async () => { + const account = buildGeminiApiKeyAccount({ + base_url: 'https://upstream.example.com', + explicit_cache_upstream: true, + explicit_cache_upstream_max_ttl_seconds: 1800 + }) + updateAccountMock.mockResolvedValue(account) + const wrapper = mountModal(account) + + await wrapper.get('[data-testid="explicit-cache-upstream-toggle"]').trigger('click') + await wrapper.get('form#edit-account-form').trigger('submit.prevent') + const credentials = updateAccountMock.mock.calls[0]?.[1]?.credentials + expect(credentials).not.toHaveProperty('explicit_cache_upstream') + expect(credentials).not.toHaveProperty('explicit_cache_upstream_max_ttl_seconds') + }) +}) diff --git a/frontend/src/i18n/locales/en/admin/accounts.ts b/frontend/src/i18n/locales/en/admin/accounts.ts index 9af550436d..8a512037b3 100644 --- a/frontend/src/i18n/locales/en/admin/accounts.ts +++ b/frontend/src/i18n/locales/en/admin/accounts.ts @@ -853,6 +853,12 @@ export default { poolModeRetryStatusCodes: 'Retry Status Codes', poolModeRetryStatusCodesHint: 'Comma-separated HTTP status codes (100-599) that trigger same-account retry in pool mode. Leave blank to use defaults ({default}).', + explicitCacheUpstream: 'Upstream Supports Explicit Caching', + explicitCacheUpstreamHint: 'Enable only when the upstream is this system with explicit caching support', + explicitCacheUpstreamInfo: + 'When enabled, this account can create Gemini explicit caches (cachedContents). A cache is only used on this account and never fails over. The upstream must route each cache to the account that created it; account-pool gateways or upstreams without explicit caching support will make caches unreliable.', + explicitCacheUpstreamMaxTTL: 'Upstream Cache TTL Limit (seconds)', + explicitCacheUpstreamMaxTTLHint: 'Maximum TTL allowed by the upstream gateway; requests above it are not sent to this account. Leave blank to use this site\'s limit.', customErrorCodes: 'Custom Error Codes', customErrorCodesHint: 'Only stop scheduling for selected error codes', customErrorCodesWarning: diff --git a/frontend/src/i18n/locales/zh/admin/accounts.ts b/frontend/src/i18n/locales/zh/admin/accounts.ts index 18a9288f73..6991f01c7e 100644 --- a/frontend/src/i18n/locales/zh/admin/accounts.ts +++ b/frontend/src/i18n/locales/zh/admin/accounts.ts @@ -956,6 +956,12 @@ export default { poolModeRetryCountHint: '仅在池模式下生效。0 表示不原地重试；默认 {default}，最大 {max}。', poolModeRetryStatusCodes: '同账号重试状态码', poolModeRetryStatusCodesHint: '仅在池模式下生效。以英文逗号分隔的 HTTP 状态码（100-599），命中时触发同账号重试。留空使用默认值（{default}）。', + explicitCacheUpstream: '上游支持显式缓存', + explicitCacheUpstreamHint: '仅当上游是已支持显式缓存的本系统时开启', + explicitCacheUpstreamInfo: + '开启后该账号可创建 Gemini 显式缓存（cachedContents）：缓存只在本账号上使用，失败不换号。上游必须保证缓存只发往创建它的账号；号池类网关或未支持显式缓存的上游开启后，缓存会时有时无。', + explicitCacheUpstreamMaxTTL: '上游缓存有效期上限（秒）', + explicitCacheUpstreamMaxTTLHint: '填写上游网关允许的最长有效期，超出的请求不会发往该账号；留空沿用本站上限。', customErrorCodes: '自定义错误码', customErrorCodesHint: '仅对选中的错误码停止调度', customErrorCodesWarning: diff --git a/frontend/src/utils/__tests__/geminiExplicitCache.spec.ts b/frontend/src/utils/__tests__/geminiExplicitCache.spec.ts new file mode 100644 index 0000000000..ef9f80158b --- /dev/null +++ b/frontend/src/utils/__tests__/geminiExplicitCache.spec.ts @@ -0,0 +1,50 @@ +import { describe, expect, it } from 'vitest' +import { + applyExplicitCacheUpstreamCredentials, + isOfficialGeminiBaseUrl, + normalizeExplicitCacheUpstreamMaxTTL +} from '../geminiExplicitCache' + +describe('isOfficialGeminiBaseUrl', () => { + it('treats empty and the official host as official', () => { + expect(isOfficialGeminiBaseUrl('')).toBe(true) + expect(isOfficialGeminiBaseUrl(undefined)).toBe(true) + expect(isOfficialGeminiBaseUrl('https://generativelanguage.googleapis.com')).toBe(true) + expect(isOfficialGeminiBaseUrl('https://generativelanguage.googleapis.com/')).toBe(true) + }) + + it('treats relays, path prefixes and invalid urls as custom', () => { + expect(isOfficialGeminiBaseUrl('https://relay.example.com')).toBe(false) + expect(isOfficialGeminiBaseUrl('https://generativelanguage.googleapis.com/gw')).toBe(false) + expect(isOfficialGeminiBaseUrl('http://generativelanguage.googleapis.com')).toBe(false) + expect(isOfficialGeminiBaseUrl('not a url')).toBe(false) + }) +}) + +describe('applyExplicitCacheUpstreamCredentials', () => { + it('writes the switch and an optional ttl cap', () => { + const credentials: Record = {} + applyExplicitCacheUpstreamCredentials(credentials, true, '1800') + expect(credentials).toEqual({ explicit_cache_upstream: true, explicit_cache_upstream_max_ttl_seconds: 1800 }) + + applyExplicitCacheUpstreamCredentials(credentials, true, '') + expect(credentials).toEqual({ explicit_cache_upstream: true }) + }) + + it('removes both fields when disabled', () => { + const credentials: Record = { + api_key: 'k', + explicit_cache_upstream: true, + explicit_cache_upstream_max_ttl_seconds: 600 + } + applyExplicitCacheUpstreamCredentials(credentials, false, 600) + expect(credentials).toEqual({ api_key: 'k' }) + }) + + it('ignores non-positive caps', () => { + expect(normalizeExplicitCacheUpstreamMaxTTL(0)).toBeNull() + expect(normalizeExplicitCacheUpstreamMaxTTL(-5)).toBeNull() + expect(normalizeExplicitCacheUpstreamMaxTTL('abc')).toBeNull() + expect(normalizeExplicitCacheUpstreamMaxTTL(90.7)).toBe(90) + }) +}) diff --git a/frontend/src/utils/geminiExplicitCache.ts b/frontend/src/utils/geminiExplicitCache.ts new file mode 100644 index 0000000000..72fc31170e --- /dev/null +++ b/frontend/src/utils/geminiExplicitCache.ts @@ -0,0 +1,47 @@ +const OFFICIAL_GEMINI_API_HOST = 'generativelanguage.googleapis.com' + +/** Gemini API Key 账号的 base_url 是否为官方地址；留空视为官方地址。 */ +export function isOfficialGeminiBaseUrl(raw: string | null | undefined): boolean { + const value = (raw ?? '').trim() + if (!value) return true + try { + const url = new URL(value) + const path = url.pathname.replace(/\/+$/, '') + return ( + url.protocol === 'https:' && + url.hostname.toLowerCase() === OFFICIAL_GEMINI_API_HOST && + url.port === '' && + path === '' && + url.search === '' + ) + } catch { + return false + } +} + +/** 上游缓存有效期上限（秒）：非正数或非数字视为未配置。 */ +export function normalizeExplicitCacheUpstreamMaxTTL(value: unknown): number | null { + if (value === null || value === undefined || value === '') return null + const seconds = Number(value) + return Number.isFinite(seconds) && seconds > 0 ? Math.floor(seconds) : null +} + +/** 把上游显式缓存开关写入（或移出）账号凭据。 */ +export function applyExplicitCacheUpstreamCredentials( + credentials: Record, + enabled: boolean, + maxTTL: unknown +): void { + if (!enabled) { + delete credentials.explicit_cache_upstream + delete credentials.explicit_cache_upstream_max_ttl_seconds + return + } + credentials.explicit_cache_upstream = true + const seconds = normalizeExplicitCacheUpstreamMaxTTL(maxTTL) + if (seconds === null) { + delete credentials.explicit_cache_upstream_max_ttl_seconds + } else { + credentials.explicit_cache_upstream_max_ttl_seconds = seconds + } +} From 3abd64658297f5ec47ad09120a0a26d6081d42de Mon Sep 17 00:00:00 2001 From: feeeei Date: Sat, 10 Oct 2026 22:09:14 +0800 Subject: [PATCH 03/15] =?UTF-8?q?fix(ops):=20Gemini=20=E5=8E=9F=E7=94=9F?= =?UTF-8?q?=E8=B7=AF=E7=94=B1=E7=9A=84=E7=BD=91=E5=85=B3=E6=9C=AC=E5=9C=B0?= =?UTF-8?q?=E9=94=99=E8=AF=AF=E6=8C=89=E9=94=99=E8=AF=AF=E7=B1=BB=E5=9E=8B?= =?UTF-8?q?=E5=BD=92=E7=B1=BB?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Google 格式错误体没有 type 字段，ops 一律记为 api_error，客户端参数错误、 引用不存在的缓存等都被归为平台内部错误。 - googleError 写出时按状态码向 ops 登记错误类型：400/413 invalid_request_error、 401 authentication_error、403 permission_error、404 not_found_error， 429 与 5xx 不登记 - 计费、并发与风控拦截沿用 Claude / OpenAI 格式错误的类型与错误码，归类口径一致 - 存储价未配置、显式缓存未启用改用 googlePlatformError，仍按平台错误统计； 错误透传规则改写的上游错误不登记本地类型 - 上游返回的缓存不存在标记为请求级失败：按客户端请求错误分类， 上游原始错误保留为明细但不参与上游归因 --- .../handler/gemini_cached_contents_handler.go | 19 +-- .../internal/handler/gemini_v1beta_handler.go | 45 +++++- backend/internal/handler/ops_error_logger.go | 13 ++ .../handler/ops_error_logger_google_test.go | 135 ++++++++++++++++++ .../internal/handler/security_audit_errors.go | 7 +- .../gemini_cached_content_not_found_test.go | 4 + .../service/gemini_messages_compat_service.go | 3 +- .../internal/service/ops_upstream_context.go | 44 ++++++ 8 files changed, 252 insertions(+), 18 deletions(-) create mode 100644 backend/internal/handler/ops_error_logger_google_test.go diff --git a/backend/internal/handler/gemini_cached_contents_handler.go b/backend/internal/handler/gemini_cached_contents_handler.go index 82b3741fad..00404329d5 100644 --- a/backend/internal/handler/gemini_cached_contents_handler.go +++ b/backend/internal/handler/gemini_cached_contents_handler.go @@ -42,7 +42,7 @@ func (h *GatewayHandler) geminiCachedContentsPreamble(c *gin.Context, component return nil, middleware.AuthSubject{}, nil, false } if h.geminiCachedContentService == nil { - googleError(c, http.StatusNotFound, "Explicit context caching is not enabled") + googlePlatformError(c, http.StatusNotFound, "Explicit context caching is not enabled") return nil, middleware.AuthSubject{}, nil, false } if apiKey.Group == nil || apiKey.GroupID == nil || apiKey.Group.Platform == service.PlatformComposite { @@ -83,11 +83,11 @@ func readGeminiCachedContentBody(c *gin.Context) ([]byte, bool) { func (h *GatewayHandler) checkGeminiCachedContentBilling(c *gin.Context, reqLog *zap.Logger, apiKey *service.APIKey, subscription *service.UserSubscription) bool { if err := h.billingCacheService.CheckBillingEligibility(c.Request.Context(), apiKey.User, apiKey, apiKey.Group, subscription, service.QuotaPlatform(c.Request.Context(), apiKey)); err != nil { reqLog.Info("gemini.cached_content.billing_eligibility_check_failed", zap.Error(err)) - status, _, message, retryAfter := billingErrorDetails(err) + status, code, message, retryAfter := billingErrorDetails(err) if retryAfter > 0 { c.Header("Retry-After", strconv.Itoa(retryAfter)) } - googleError(c, status, message) + googleErrorWithType(c, status, code, "", message) return false } return true @@ -150,11 +150,11 @@ func (h *GatewayHandler) GeminiCachedContentsCreate(c *gin.Context) { inflightRelease, err := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(modelName, body)) if err != nil { reqLog.Info("gemini.cached_content.inflight_reservation_rejected", zap.Error(err)) - status, _, message, retryAfter := billingErrorDetails(err) + status, code, message, retryAfter := billingErrorDetails(err) if retryAfter > 0 { c.Header("Retry-After", strconv.Itoa(retryAfter)) } - googleError(c, status, message) + googleErrorWithType(c, status, code, "", message) return } defer inflightRelease() @@ -218,7 +218,7 @@ func (h *GatewayHandler) GeminiCachedContentsCreate(c *gin.Context) { ChannelUsageFields: usageFields, }) { accountReleaseFunc() - googleError(c, http.StatusBadRequest, fmt.Sprintf("Explicit context caching is not available for model %s: storage pricing is not configured", modelName)) + googlePlatformError(c, http.StatusBadRequest, fmt.Sprintf("Explicit context caching is not available for model %s: storage pricing is not configured", modelName)) return } upstreamModel, err := service.GeminiCachedContentUpstreamModelResource(account, mappedModel) @@ -438,7 +438,7 @@ func (h *GatewayHandler) GeminiCachedContentsPatch(c *gin.Context) { APIKey: apiKey, User: apiKey.User, Subscription: subscription, PricingAt: pricingAt, } if !h.gatewayService.CacheStoragePriced(c.Request.Context(), usage) { - googleError(c, http.StatusBadRequest, fmt.Sprintf("Explicit context caching is not available for model %s: storage pricing is not configured", record.Model)) + googlePlatformError(c, http.StatusBadRequest, fmt.Sprintf("Explicit context caching is not available for model %s: storage pricing is not configured", record.Model)) return } if !h.checkGeminiCachedContentBilling(c, reqLog, apiKey, subscription) { @@ -722,11 +722,11 @@ func (h *GatewayHandler) reserveGeminiCachedContentCharge(c *gin.Context, apiKey } func writeGeminiCachedContentBillingError(c *gin.Context, err error) { - status, _, message, retryAfter := billingErrorDetails(err) + status, code, message, retryAfter := billingErrorDetails(err) if retryAfter > 0 { c.Header("Retry-After", strconv.Itoa(retryAfter)) } - googleError(c, status, message) + googleErrorWithType(c, status, code, "", message) } // writeGeminiCachedContentUpstreamError 透传上游错误：缓存不存在统一返回标准 403（各上游形态不同且可能含上游资源 ID）； @@ -736,6 +736,7 @@ func writeGeminiCachedContentUpstreamError(c *gin.Context, upErr *service.Gemini message := strings.TrimSpace(extractGoogleErrorMessage(upErr.Body)) service.SetOpsUpstreamError(c, status, message, "") if upErr.NotFound() { + service.MarkOpsRequestScopedError(c, "permission_error") c.Data(http.StatusForbidden, "application/json", []byte(service.GeminiCachedContentNotFoundResponse)) return } diff --git a/backend/internal/handler/gemini_v1beta_handler.go b/backend/internal/handler/gemini_v1beta_handler.go index 18ac84084f..63e944cb5e 100644 --- a/backend/internal/handler/gemini_v1beta_handler.go +++ b/backend/internal/handler/gemini_v1beta_handler.go @@ -388,11 +388,11 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) { // 2) billing eligibility check (after wait) if err := h.billingCacheService.CheckBillingEligibility(c.Request.Context(), apiKey.User, apiKey, apiKey.Group, subscription, service.QuotaPlatform(c.Request.Context(), apiKey)); err != nil { reqLog.Info("gemini.billing_eligibility_check_failed", zap.Error(err)) - status, _, message, retryAfter := billingErrorDetails(err) + status, code, message, retryAfter := billingErrorDetails(err) if retryAfter > 0 { c.Header("Retry-After", strconv.Itoa(retryAfter)) } - googleError(c, status, message) + googleErrorWithType(c, status, code, "", message) return } @@ -400,11 +400,11 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) { inflightRelease, err := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(modelName, body)) if err != nil { reqLog.Info("gemini.inflight_reservation_rejected", zap.Error(err)) - status, _, message, retryAfter := billingErrorDetails(err) + status, code, message, retryAfter := billingErrorDetails(err) if retryAfter > 0 { c.Header("Retry-After", strconv.Itoa(retryAfter)) } - googleError(c, status, message) + googleErrorWithType(c, status, code, "", message) return } defer inflightRelease() @@ -861,7 +861,8 @@ func (h *GatewayHandler) handleGeminiFailoverExhausted(c *gin.Context, failoverE c.Set(service.OpsSkipPassthroughKey, true) } - googleError(c, respCode, msg) + // 按透传规则改写的是上游错误，不登记本地错误类型。 + googleErrorWithType(c, respCode, "", "", msg) return } } @@ -896,7 +897,21 @@ type pathParseError struct{ msg string } func (e *pathParseError) Error() string { return e.msg } +// googleError 以 Google 错误格式回写网关产生的错误，并按状态码向 ops 登记错误类型： +// 4xx 为客户端请求问题，其余状态码不登记（按平台错误统计）。 func googleError(c *gin.Context, status int, message string) { + googleErrorWithType(c, status, googleClientErrorType(status), "", message) +} + +// googlePlatformError 以 Google 错误格式回写由平台侧原因（如管理员未完成配置）导致的错误，ops 按平台错误统计。 +func googlePlatformError(c *gin.Context, status int, message string) { + googleErrorWithType(c, status, "", "", message) +} + +// googleErrorWithType 以 Google 错误格式回写错误，并向 ops 登记错误类型与错误码（Google 错误体没有 type 字段）， +// 与 Claude / OpenAI 格式错误同口径分类；errType 为空时不登记。 +func googleErrorWithType(c *gin.Context, status int, errType, code, message string) { + service.SetOpsLocalErrorType(c, errType, code) c.JSON(status, gin.H{ "error": gin.H{ "code": status, @@ -906,11 +921,27 @@ func googleError(c *gin.Context, status int, message string) { }) } +// googleClientErrorType 返回客户端请求类 4xx 状态码对应的 ops 错误类型；其余状态码返回空串。 +func googleClientErrorType(status int) string { + switch status { + case http.StatusBadRequest, http.StatusRequestEntityTooLarge: + return "invalid_request_error" + case http.StatusUnauthorized: + return "authentication_error" + case http.StatusForbidden: + return "permission_error" + case http.StatusNotFound: + return "not_found_error" + default: + return "" + } +} + // googleConcurrencyError 以 Google 错误格式回写并发槽获取失败，状态码与文案 // 沿用 concurrencyErrorResponse 的统一映射（客户端断开为 499）。 func googleConcurrencyError(c *gin.Context, err error, slotType string) { - status, _, _, message := concurrencyErrorResponse(err, slotType) - googleError(c, status, message) + status, errType, code, message := concurrencyErrorResponse(err, slotType) + googleErrorWithType(c, status, errType, code, message) } func writeUpstreamResponse(c *gin.Context, res *service.UpstreamHTTPResult) { diff --git a/backend/internal/handler/ops_error_logger.go b/backend/internal/handler/ops_error_logger.go index a4e30814ca..679cf4e8a2 100644 --- a/backend/internal/handler/ops_error_logger.go +++ b/backend/internal/handler/ops_error_logger.go @@ -1183,6 +1183,14 @@ func OpsErrorLoggerMiddleware(ops *service.OpsService) gin.HandlerFunc { } } + // Google 格式错误体没有 type 字段：按网关回写时登记的错误类型补齐。 + if errType, code := service.OpsLocalErrorType(c); errType != "" && !parsed.StreamFailure && + (parsed.ErrorType == "" || parsed.ErrorType == "api_error") { + parsed.ErrorType = errType + if code != "" { + parsed.Code = code + } + } normalizedType := normalizeOpsErrorType(parsed.ErrorType, parsed.Code) phase, isBusinessLimited, errorOwner, errorSource := classifyOpsErrorLog(c, normalizedType, parsed.Message, parsed.Code, status) @@ -2214,6 +2222,11 @@ func classifyOpsSeverity(errType string, status int) string { func classifyOpsErrorLog(c *gin.Context, errType, message, code string, status int) (phase string, isBusinessLimited bool, errorOwner string, errorSource string) { phase = classifyOpsPhase(errType, message, code) + if service.IsOpsRequestScopedError(c) { + // 失败由请求内容决定：只按错误类型分类，上游错误上下文不参与归因。 + isBusinessLimited = classifyOpsIsBusinessLimited(errType, phase, code, status, message) + return phase, isBusinessLimited, classifyOpsErrorOwner(phase, message), classifyOpsErrorSource(phase, message) + } routingCapacityLimited := isOpsRoutingCapacityLimited(c) clientBusinessLimited := service.HasOpsClientBusinessLimited(c) localModelConfiguration := clientBusinessLimited && service.OpsClientBusinessLimitedReason(c) == service.OpsClientBusinessLimitedReasonLocalModelConfiguration diff --git a/backend/internal/handler/ops_error_logger_google_test.go b/backend/internal/handler/ops_error_logger_google_test.go new file mode 100644 index 0000000000..392421dffd --- /dev/null +++ b/backend/internal/handler/ops_error_logger_google_test.go @@ -0,0 +1,135 @@ +package handler + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +// recordGoogleFormatOpsError 让 handle 以 Google 错误格式回写错误，返回 ops 中间件记录的错误日志。 +func recordGoogleFormatOpsError(t *testing.T, handle gin.HandlerFunc) *service.OpsInsertErrorLogInput { + t.Helper() + setupOpsErrorLogTestQueue(t, 1) + gin.SetMode(gin.TestMode) + ops := service.NewOpsService(nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil) + router := gin.New() + router.Use(OpsErrorLoggerMiddleware(ops)) + router.POST("/v1beta/cachedContents", handle) + + router.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodPost, "/v1beta/cachedContents", nil)) + + require.Equal(t, int64(1), OpsErrorLogQueueLength()) + return (<-opsErrorLogQueue).entry +} + +func TestOpsErrorLoggerMiddleware_GoogleClientErrorsClassifiedAsClient(t *testing.T) { + for _, tc := range []struct { + status int + errType string + phase string + }{ + {status: http.StatusBadRequest, errType: "invalid_request_error", phase: "request"}, + {status: http.StatusRequestEntityTooLarge, errType: "invalid_request_error", phase: "request"}, + {status: http.StatusUnauthorized, errType: "authentication_error", phase: "auth"}, + {status: http.StatusForbidden, errType: "permission_error", phase: "request"}, + {status: http.StatusNotFound, errType: "not_found_error", phase: "request"}, + } { + t.Run(http.StatusText(tc.status), func(t *testing.T) { + entry := recordGoogleFormatOpsError(t, func(c *gin.Context) { + googleError(c, tc.status, "client request problem") + }) + require.Equal(t, tc.status, entry.StatusCode) + require.Equal(t, tc.errType, entry.ErrorType) + require.Equal(t, tc.phase, entry.ErrorPhase) + require.Equal(t, "client", entry.ErrorOwner) + require.Equal(t, "client_request", entry.ErrorSource) + require.Equal(t, "P3", entry.Severity) + }) + } +} + +func TestOpsErrorLoggerMiddleware_GooglePlatformErrorStaysPlatform(t *testing.T) { + entry := recordGoogleFormatOpsError(t, func(c *gin.Context) { + googlePlatformError(c, http.StatusBadRequest, "Explicit context caching is not available for model m: storage pricing is not configured") + }) + require.Equal(t, "api_error", entry.ErrorType) + require.Equal(t, "internal", entry.ErrorPhase) + require.Equal(t, "platform", entry.ErrorOwner) + require.Equal(t, "gateway", entry.ErrorSource) +} + +func TestOpsErrorLoggerMiddleware_GoogleServerErrorsNotTypedByStatus(t *testing.T) { + entry := recordGoogleFormatOpsError(t, func(c *gin.Context) { + googleError(c, http.StatusInternalServerError, "Failed to persist cached content") + }) + require.Equal(t, "api_error", entry.ErrorType) + require.Equal(t, "platform", entry.ErrorOwner) +} + +func TestOpsErrorLoggerMiddleware_GoogleBillingErrorMatchesClaudeFormat(t *testing.T) { + entry := recordGoogleFormatOpsError(t, func(c *gin.Context) { + status, code, message, _ := billingErrorDetails(service.ErrInsufficientBalance) + googleErrorWithType(c, status, code, "", message) + }) + require.Equal(t, "billing_error", entry.ErrorType) + require.Equal(t, "request", entry.ErrorPhase) + require.Equal(t, "client", entry.ErrorOwner) + require.True(t, entry.IsBusinessLimited) +} + +func TestOpsErrorLoggerMiddleware_GoogleConcurrencyErrorClassifiedAsRequest(t *testing.T) { + entry := recordGoogleFormatOpsError(t, func(c *gin.Context) { + googleConcurrencyError(c, &ConcurrencyError{SlotType: "user"}, "user") + }) + require.Equal(t, http.StatusTooManyRequests, entry.StatusCode) + require.Equal(t, "rate_limit_error", entry.ErrorType) + require.Equal(t, "request", entry.ErrorPhase) + require.Equal(t, "client", entry.ErrorOwner) +} + +func TestOpsErrorLoggerMiddleware_GoogleErrorAfterUpstreamFailureKeepsUpstreamAttribution(t *testing.T) { + entry := recordGoogleFormatOpsError(t, func(c *gin.Context) { + service.SetOpsUpstreamError(c, http.StatusBadRequest, "upstream rejected", "") + googleError(c, http.StatusBadRequest, "upstream rejected") + }) + require.Equal(t, "upstream", entry.ErrorPhase) + require.Equal(t, "provider", entry.ErrorOwner) + require.Equal(t, "upstream_http", entry.ErrorSource) +} + +func TestOpsErrorLoggerMiddleware_RequestScopedUpstreamErrorClassifiedAsClient(t *testing.T) { + entry := recordGoogleFormatOpsError(t, func(c *gin.Context) { + service.SetOpsUpstreamError(c, http.StatusBadRequest, "Invalid resource state for cache content 8079767076522164224.", "") + c.Set(service.OpsUpstreamErrorsKey, []*service.OpsUpstreamErrorEvent{{ + Platform: service.PlatformGemini, AccountID: 173, UpstreamStatusCode: http.StatusBadRequest, Kind: "http_error", + Message: "Invalid resource state for cache content 8079767076522164224.", + }}) + service.MarkOpsRequestScopedError(c, "permission_error") + c.Data(http.StatusForbidden, "application/json", []byte(service.GeminiCachedContentNotFoundResponse)) + }) + require.Equal(t, http.StatusForbidden, entry.StatusCode) + require.Equal(t, "permission_error", entry.ErrorType) + require.Equal(t, "request", entry.ErrorPhase) + require.Equal(t, "client", entry.ErrorOwner) + require.Equal(t, "client_request", entry.ErrorSource) + require.False(t, entry.IsBusinessLimited) + require.NotNil(t, entry.UpstreamStatusCode, "上游原始错误保留为明细") + require.Equal(t, http.StatusBadRequest, *entry.UpstreamStatusCode) +} + +func TestSetOpsLocalErrorTypeClearsRequestScopedMark(t *testing.T) { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + service.MarkOpsRequestScopedError(c, "permission_error") + require.True(t, service.IsOpsRequestScopedError(c)) + + service.SetOpsLocalErrorType(c, "", "") + require.False(t, service.IsOpsRequestScopedError(c)) + errType, code := service.OpsLocalErrorType(c) + require.Empty(t, errType) + require.Empty(t, code) +} diff --git a/backend/internal/handler/security_audit_errors.go b/backend/internal/handler/security_audit_errors.go index b6f42c7111..7c55e9bfde 100644 --- a/backend/internal/handler/security_audit_errors.go +++ b/backend/internal/handler/security_audit_errors.go @@ -100,9 +100,14 @@ func googleSecurityAuditError(c *gin.Context, decision *securityaudit.Decision) return } if decision.Legacy != nil && decision.Legacy.Blocked { - googleError(c, securityAuditStatus(decision), securityAuditMessage(decision)) + googleErrorWithType(c, securityAuditStatus(decision), securityAuditErrorCode(decision), "", securityAuditMessage(decision)) return } + errType := "api_error" + if decision.Kind == securityaudit.DecisionBlock { + errType = "permission_error" + } + service.SetOpsLocalErrorType(c, errType, securityAuditErrorCode(decision)) status := securityAuditStatus(decision) googleStatus := googleapi.HTTPStatusToGoogleStatus(status) if status == http.StatusServiceUnavailable { diff --git a/backend/internal/service/gemini_cached_content_not_found_test.go b/backend/internal/service/gemini_cached_content_not_found_test.go index 783ba71c72..fc4daa1181 100644 --- a/backend/internal/service/gemini_cached_content_not_found_test.go +++ b/backend/internal/service/gemini_cached_content_not_found_test.go @@ -75,6 +75,7 @@ func TestGeminiForwardNative_CachedContentNotFoundPassthroughWithoutAccountPenal require.Equal(t, 1, httpStub.calls, "不应重试") require.Equal(t, http.StatusForbidden, rec.Code) require.Equal(t, GeminiCachedContentNotFoundResponse, rec.Body.String(), "应返回标准缓存不存在响应") + require.True(t, IsOpsRequestScopedError(c), "缓存不存在按客户端请求错误统计") }) } } @@ -177,6 +178,9 @@ func TestGeminiForwardNative_VertexCachedContentNotFoundReturnsStandardForbidden require.Equal(t, http.StatusForbidden, rec.Code) require.Equal(t, GeminiCachedContentNotFoundResponse, rec.Body.String()) require.NotContains(t, rec.Body.String(), "8079767076522164224", "不应暴露上游缓存 ID") + require.True(t, IsOpsRequestScopedError(c), "缓存不存在按客户端请求错误统计") + errType, _ := OpsLocalErrorType(c) + require.Equal(t, "permission_error", errType) }) } } diff --git a/backend/internal/service/gemini_messages_compat_service.go b/backend/internal/service/gemini_messages_compat_service.go index 60421d1f7f..c3edbd2fd4 100644 --- a/backend/internal/service/gemini_messages_compat_service.go +++ b/backend/internal/service/gemini_messages_compat_service.go @@ -1834,10 +1834,11 @@ func (s *GeminiMessagesCompatService) writeGeminiNativeUpstreamError(c *gin.Cont } // writeGeminiCachedContentNotFound 记录上游原始错误，并向客户端统一返回缓存不存在的 403： -// 各上游的响应形态不同，且 Vertex 的文案含上游缓存资源 ID。 +// 各上游的响应形态不同，且 Vertex 的文案含上游缓存资源 ID。失败由请求引用的缓存决定，ops 按客户端请求错误统计。 func (s *GeminiMessagesCompatService) writeGeminiCachedContentNotFound(c *gin.Context, account *Account, resp *http.Response, respBody []byte, requestID string, isOAuth bool) error { respBody = unwrapIfNeeded(isOAuth, respBody) upstreamMsg := s.recordGeminiNativeUpstreamError(c, account, resp.StatusCode, respBody, requestID) + MarkOpsRequestScopedError(c, "permission_error") MarkResponseCommitted(c) c.Data(http.StatusForbidden, "application/json", []byte(GeminiCachedContentNotFoundResponse)) return fmt.Errorf("gemini upstream error: %d message=%s", resp.StatusCode, upstreamMsg) diff --git a/backend/internal/service/ops_upstream_context.go b/backend/internal/service/ops_upstream_context.go index 4d09a43ae6..a8d0625fd8 100644 --- a/backend/internal/service/ops_upstream_context.go +++ b/backend/internal/service/ops_upstream_context.go @@ -52,6 +52,14 @@ const ( // ensureForwardErrorResponse 检查此 key，为 true 时跳过兜底写入，避免在已完成的 JSON 后追加 SSE。 ResponseCommittedKey = "response_committed" + // OpsLocalErrorTypeKey / OpsLocalErrorCodeKey 记录网关以 Google 错误格式回写的错误类型与错误码： + // Google 错误体没有 type 字段，ops_error_logger 据此按与 Claude / OpenAI 格式错误相同的口径分类。 + OpsLocalErrorTypeKey = "ops_local_error_type" + OpsLocalErrorCodeKey = "ops_local_error_code" + // OpsRequestScopedErrorKey 标记失败由请求内容决定（如引用的显式缓存不存在）：只按错误类型分类， + // 上游错误上下文保留为明细但不参与归因。 + OpsRequestScopedErrorKey = "ops_request_scoped_error" + OpsClientBusinessLimitedKey = "ops_client_business_limited" OpsClientBusinessLimitedReasonKey = "ops_client_business_limited_reason" OpsClientBusinessLimitedReasonIPRestriction = "api_key_ip_restriction" @@ -134,6 +142,42 @@ func OpsClientBusinessLimitedReason(c *gin.Context) string { return strings.TrimSpace(reason) } +// SetOpsLocalErrorType 登记响应错误体缺少的错误类型与错误码，并清除请求级失败标记；errType 为空表示不登记。 +func SetOpsLocalErrorType(c *gin.Context, errType, code string) { + if c == nil { + return + } + c.Set(OpsLocalErrorTypeKey, strings.TrimSpace(errType)) + c.Set(OpsLocalErrorCodeKey, strings.TrimSpace(code)) + c.Set(OpsRequestScopedErrorKey, false) +} + +// OpsLocalErrorType 返回 SetOpsLocalErrorType 登记的错误类型与错误码；未登记时返回空串。 +func OpsLocalErrorType(c *gin.Context) (errType, code string) { + if c == nil { + return "", "" + } + errType = c.GetString(OpsLocalErrorTypeKey) + if errType == "" { + return "", "" + } + return errType, c.GetString(OpsLocalErrorCodeKey) +} + +// MarkOpsRequestScopedError 登记错误类型，并标记失败由请求内容决定：上游错误上下文不参与 ops 归因。 +func MarkOpsRequestScopedError(c *gin.Context, errType string) { + if c == nil { + return + } + SetOpsLocalErrorType(c, errType, "") + c.Set(OpsRequestScopedErrorKey, true) +} + +// IsOpsRequestScopedError 判断 MarkOpsRequestScopedError 是否标记了本次请求的失败。 +func IsOpsRequestScopedError(c *gin.Context) bool { + return c != nil && c.GetBool(OpsRequestScopedErrorKey) +} + // OpsStreamError 描述承载在 2xx 响应上的带内错误：网关在响应状态已固化为 200 之后 // 就地以 SSE error 帧返回的错误（并发限流回退、Wait 后二次计费校验失败、流开始后才无 // 可用账号等），以及上游 2xx 正文或事件里携带的错误结果（NonStream 标记非流式正文）。 From 64331a97921bd5a43e7b54ba6a517476c985e8e5 Mon Sep 17 00:00:00 2001 From: feeeei Date: Sat, 10 Oct 2026 22:22:12 +0800 Subject: [PATCH 04/15] =?UTF-8?q?ui(gemini):=20Vertex=20=E6=98=BE=E7=A4=BA?= =?UTF-8?q?=E5=90=8D=E6=94=B9=E4=B8=BA=20Agent=20Platform=20(Vertex)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Google 已将 Vertex AI 更名为 Gemini Enterprise Agent Platform。 账号类型按钮、平台徽章、Vertex 相关提示、配额文档链接与批量生图报错文案 统一显示为 Agent Platform (Vertex)；创建账号弹窗的 Vertex 按钮改用 i18n 文案。 --- .../src/components/account/CreateAccountModal.vue | 8 ++++---- frontend/src/components/common/PlatformTypeBadge.vue | 2 +- frontend/src/i18n/locales/en/admin/accounts.ts | 12 ++++++------ frontend/src/i18n/locales/en/batchImage.ts | 6 +++--- frontend/src/i18n/locales/zh/admin/accounts.ts | 12 ++++++------ frontend/src/i18n/locales/zh/batchImage.ts | 6 +++--- 6 files changed, 23 insertions(+), 23 deletions(-) diff --git a/frontend/src/components/account/CreateAccountModal.vue b/frontend/src/components/account/CreateAccountModal.vue index 28ca1b2b00..e4aa3dfc4c 100644 --- a/frontend/src/components/account/CreateAccountModal.vue +++ b/frontend/src/components/account/CreateAccountModal.vue @@ -381,8 +381,8 @@
- Vertex - Service Account + {{ t('admin.accounts.vertexLabel') }} + {{ t('admin.accounts.vertexDesc') }}
@@ -830,10 +830,10 @@
- Vertex + {{ t('admin.accounts.vertexLabel') }} - Service Account + {{ t('admin.accounts.vertexDesc') }}
diff --git a/frontend/src/components/common/PlatformTypeBadge.vue b/frontend/src/components/common/PlatformTypeBadge.vue index f0831c593a..c726ff9ce8 100644 --- a/frontend/src/components/common/PlatformTypeBadge.vue +++ b/frontend/src/components/common/PlatformTypeBadge.vue @@ -108,7 +108,7 @@ const typeLabel = computed(() => { case 'bedrock': return 'AWS' case 'service_account': - return 'Vertex' + return 'Agent Platform (Vertex)' default: return props.type } diff --git a/frontend/src/i18n/locales/en/admin/accounts.ts b/frontend/src/i18n/locales/en/admin/accounts.ts index 8a512037b3..350551b3c7 100644 --- a/frontend/src/i18n/locales/en/admin/accounts.ts +++ b/frontend/src/i18n/locales/en/admin/accounts.ts @@ -589,10 +589,10 @@ export default { claudeConsole: 'Claude Console', bedrockLabel: 'AWS Bedrock', bedrockDesc: 'SigV4 / API Key', - vertexLabel: 'Vertex', + vertexLabel: 'Agent Platform (Vertex)', vertexDesc: 'Service Account', - vertexAnthropicHint: 'Use a Google Cloud Service Account JSON to call Anthropic Claude via Vertex AI. It is recommended to configure model mapping to map client Claude model names to Vertex model IDs.', - vertexGeminiHint: 'Use a Google Cloud Service Account JSON to access Vertex AI Gemini. It is recommended to place Vertex accounts in a separate group to avoid mixing with AI Studio/Gemini OAuth on the same models.', + vertexAnthropicHint: 'Use a Google Cloud Service Account JSON to call Anthropic Claude via Agent Platform (Vertex). It is recommended to configure model mapping to map client Claude model names to Agent Platform (Vertex) model IDs.', + vertexGeminiHint: 'Use a Google Cloud Service Account JSON to call Gemini via Agent Platform (Vertex). It is recommended to place these accounts in a separate group to avoid mixing with AI Studio/Gemini OAuth on the same models.', vertexSaJsonLabel: 'Service Account JSON', vertexSaJsonLoaded: 'Service Account JSON loaded', vertexSaJsonDrop: 'Drop Service Account JSON here', @@ -602,8 +602,8 @@ export default { vertexSaJsonUploadHint: 'After uploading or dropping a JSON file, the project_id will be auto-extracted. Key content is only used for account creation.', vertexSaJsonEditHint: 'Service Account JSON is not shown on the edit page; to change the JSON, delete the account and recreate it.', vertexProjectIdPlaceholder: 'Auto-extracted from JSON', - vertexLocationHint: 'Available locations vary by Vertex model. Select the default endpoint location for this account.', - vertexLocationRequired: 'Please enter a Vertex location', + vertexLocationHint: 'Available locations vary by model. Select the default endpoint location for this account.', + vertexLocationRequired: 'Please select a location', vertexSaJsonMissingFields: 'Service Account JSON is missing project_id, client_email, or private_key', vertexSaJsonMissingProjectId: 'Service Account JSON is missing project_id', vertexSaJsonMissingClientEmail: 'Service Account JSON is missing client_email', @@ -1499,7 +1499,7 @@ export default { docs: { codeAssist: 'Code Assist Quotas', aiStudio: 'AI Studio Pricing', - vertex: 'Vertex AI Quotas' + vertex: 'Agent Platform (Vertex) Quotas' }, simulatedNote: 'Simulated quota, for reference only', rows: { diff --git a/frontend/src/i18n/locales/en/batchImage.ts b/frontend/src/i18n/locales/en/batchImage.ts index d7b2274871..bfc0dc55d8 100644 --- a/frontend/src/i18n/locales/en/batchImage.ts +++ b/frontend/src/i18n/locales/en/batchImage.ts @@ -96,7 +96,7 @@ export default { noUsableImage: 'No usable image was generated.', cancelled: 'Job cancelled.', waiting: 'Waiting for results.', - emptyImageOutput: 'The upstream returned a result, but this item has no image content. This usually means a single Gemini/Vertex generation failed or was blocked by safety policies.', + emptyImageOutput: 'The upstream returned a result, but this item has no image content. This usually means a single Gemini/Agent Platform (Vertex) generation failed or was blocked by safety policies.', providerItemFailed: 'The upstream result for this item has no usable image.', }, imagePreview: { @@ -174,10 +174,10 @@ export default { loadingModels: 'Loading available models...', noModels: 'No available models', noModelsHint: 'This key’s group has no models configured for batch image generation.', - noCompatibleAccount: 'No usable upstream batch image account is available for this key’s group. Contact an administrator to check the group’s schedulable Gemini API key or Vertex service account and model support.', + noCompatibleAccount: 'No usable upstream batch image account is available for this key’s group. Contact an administrator to check the group’s schedulable Gemini API key or Agent Platform (Vertex) service account and model support.', unsupportedProvider: 'The batch image provider for this job is not available. Contact an administrator to check the batch image provider configuration.', providerSubmitFailed: 'The upstream batch image job failed to submit. Contact an administrator to check the upstream account, model permission, or provider status.', - vertexGcsBucketMissing: 'Vertex batch image generation is missing the managed GCS bucket configuration. Contact an administrator to configure BATCH_IMAGE_VERTEX_MANAGED_GCS_BUCKET before submitting again.', + vertexGcsBucketMissing: 'Agent Platform (Vertex) batch image generation is missing the managed GCS bucket configuration. Contact an administrator to configure BATCH_IMAGE_VERTEX_MANAGED_GCS_BUCKET before submitting again.', queueFailed: 'The task queue is temporarily unavailable, so the batch job was not queued. Contact an administrator to check the queue service.', billingHoldFailed: 'The cost hold failed, so the batch job was not submitted. Contact an administrator to check billing or balance hold service.', groupDisabled: 'Batch image generation is not enabled for this key’s group. Choose another enabled key or contact an administrator.', diff --git a/frontend/src/i18n/locales/zh/admin/accounts.ts b/frontend/src/i18n/locales/zh/admin/accounts.ts index 6991f01c7e..94f71ff5dd 100644 --- a/frontend/src/i18n/locales/zh/admin/accounts.ts +++ b/frontend/src/i18n/locales/zh/admin/accounts.ts @@ -707,10 +707,10 @@ export default { claudeConsole: 'Claude Console', bedrockLabel: 'AWS Bedrock', bedrockDesc: 'SigV4 / API Key', - vertexLabel: 'Vertex', + vertexLabel: 'Agent Platform (Vertex)', vertexDesc: 'Service Account', - vertexAnthropicHint: '使用 Google Cloud Service Account JSON 通过 Vertex AI 调用 Anthropic Claude。建议配置模型映射，将客户端 Claude 模型名映射到 Vertex 模型 ID。', - vertexGeminiHint: '使用 Google Cloud Service Account JSON 访问 Vertex AI Gemini。建议将 Vertex 账号放入独立分组，避免和 AI Studio/Gemini OAuth 同模型混调。', + vertexAnthropicHint: '使用 Google Cloud Service Account JSON 通过 Agent Platform (Vertex) 调用 Anthropic Claude。建议配置模型映射，将客户端 Claude 模型名映射到 Agent Platform (Vertex) 模型 ID。', + vertexGeminiHint: '使用 Google Cloud Service Account JSON 通过 Agent Platform (Vertex) 调用 Gemini。建议将此类账号放入独立分组，避免和 AI Studio/Gemini OAuth 同模型混调。', vertexSaJsonLabel: 'Service Account JSON', vertexSaJsonLoaded: '已读取 Service Account JSON', vertexSaJsonDrop: '拖入 Service Account JSON', @@ -720,8 +720,8 @@ export default { vertexSaJsonUploadHint: '上传或拖入 JSON 后会自动读取 project_id，密钥内容仅用于创建账号提交。', vertexSaJsonEditHint: 'Service Account JSON 不在编辑页显示；需要更换 JSON 时请删除账号后重新创建。', vertexProjectIdPlaceholder: '从 JSON 自动读取', - vertexLocationHint: '不同 Vertex 模型可用 location 可能不同，这里选择账号默认 endpoint location。', - vertexLocationRequired: '请填写 Vertex location', + vertexLocationHint: '不同模型可用 location 可能不同，这里选择账号默认 endpoint location。', + vertexLocationRequired: '请选择 Location', vertexSaJsonMissingFields: 'Service Account JSON 缺少 project_id、client_email 或 private_key', vertexSaJsonMissingProjectId: 'Service Account JSON 缺少 project_id', vertexSaJsonMissingClientEmail: 'Service Account JSON 缺少 client_email', @@ -1589,7 +1589,7 @@ export default { docs: { codeAssist: 'Code Assist 配额', aiStudio: 'AI Studio 定价', - vertex: 'Vertex AI 配额' + vertex: 'Agent Platform (Vertex) 配额' }, simulatedNote: '本地模拟配额，仅供参考', rows: { diff --git a/frontend/src/i18n/locales/zh/batchImage.ts b/frontend/src/i18n/locales/zh/batchImage.ts index cf018cef89..5aafcae77f 100644 --- a/frontend/src/i18n/locales/zh/batchImage.ts +++ b/frontend/src/i18n/locales/zh/batchImage.ts @@ -96,7 +96,7 @@ export default { noUsableImage: '未生成可用图片', cancelled: '任务已取消', waiting: '等待生成结果', - emptyImageOutput: '上游返回了结果，但这条没有图片内容。通常是 Gemini/Vertex 单条生成失败或被安全策略拦截。', + emptyImageOutput: '上游返回了结果，但这条没有图片内容。通常是 Gemini/Agent Platform (Vertex) 单条生成失败或被安全策略拦截。', providerItemFailed: '上游返回的这条结果没有可用图片。', }, imagePreview: { @@ -174,10 +174,10 @@ export default { loadingModels: '加载可用模型中...', noModels: '无可用模型', noModelsHint: '当前密钥所属分组没有配置可用于批量生图的模型。', - noCompatibleAccount: '当前密钥所属分组没有可用的批量生图上游账号。请联系管理员检查：该分组是否绑定了可调度的 Gemini API Key 或 Vertex 服务账号，以及账号是否支持所选模型。', + noCompatibleAccount: '当前密钥所属分组没有可用的批量生图上游账号。请联系管理员检查：该分组是否绑定了可调度的 Gemini API Key 或 Agent Platform (Vertex) Service Account 账号，以及账号是否支持所选模型。', unsupportedProvider: '这个任务使用的批量生图通道当前不可用。请联系管理员检查批量生图通道配置。', providerSubmitFailed: '上游批量生图任务提交失败。请联系管理员检查上游账号状态、模型权限或服务状态。', - vertexGcsBucketMissing: 'Vertex 批量生图缺少托管 GCS 存储桶配置。请联系管理员配置 BATCH_IMAGE_VERTEX_MANAGED_GCS_BUCKET 后再提交。', + vertexGcsBucketMissing: 'Agent Platform (Vertex) 批量生图缺少托管 GCS 存储桶配置。请联系管理员配置 BATCH_IMAGE_VERTEX_MANAGED_GCS_BUCKET 后再提交。', queueFailed: '任务队列暂时不可用，批量任务没有成功入队。请联系管理员检查队列服务。', billingHoldFailed: '费用冻结失败，批量任务没有成功提交。请联系管理员检查余额冻结或计费服务。', groupDisabled: '当前密钥所属分组没有开启批量生图。你可以换一个已开启批量生图的密钥，或联系管理员开启。', From 81e6ec5346c645fa03d794495198e32e4b910598 Mon Sep 17 00:00:00 2001 From: feeeei Date: Sat, 10 Oct 2026 22:22:12 +0800 Subject: [PATCH 05/15] =?UTF-8?q?feat(gemini):=20Vertex=20=E8=B4=A6?= =?UTF-8?q?=E5=8F=B7=E6=94=AF=E6=8C=81=20API=20Key=20=E8=AE=A4=E8=AF=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Gemini 平台的 Vertex 账号新增 API Key 认证（credentials.auth_mode=apikey）， 不填 Project ID 走全局端点，填写后走项目级端点，Key 通过 x-goog-api-key 传递； Anthropic 平台的 Vertex 账号仍只支持 Service Account JSON - 转发（Claude 兼容 / 原生 / Chat Completions）与测试连接统一走 buildVertexGeminiAccountURL / setVertexGeminiAuth - API Key 账号不参与显式缓存（Vertex 缓存接口不接受 API Key）与批量生图 - 创建弹窗新增认证方式选择；编辑弹窗支持轮换 Key 与修改 Project ID， Key 留空表示保留 --- .../internal/service/account_test_service.go | 13 +- .../internal/service/gemini_cached_content.go | 5 +- .../service/gemini_cached_content_test.go | 1 + .../gemini_chat_completions_compat_service.go | 14 +- .../service/gemini_messages_compat_service.go | 28 +- .../gemini_vertex_api_key_forward_test.go | 101 +++++++ .../service/vertex_service_account.go | 74 ++++- .../service/vertex_service_account_test.go | 56 ++++ .../components/account/CreateAccountModal.vue | 272 +++++++++++++----- .../components/account/EditAccountModal.vue | 114 ++++++-- .../__tests__/CreateAccountModal.spec.ts | 57 ++++ .../__tests__/EditAccountModal.spec.ts | 56 ++++ .../src/i18n/locales/en/admin/accounts.ts | 10 +- .../src/i18n/locales/zh/admin/accounts.ts | 10 +- 14 files changed, 655 insertions(+), 156 deletions(-) create mode 100644 backend/internal/service/gemini_vertex_api_key_forward_test.go diff --git a/backend/internal/service/account_test_service.go b/backend/internal/service/account_test_service.go index 5ae0859a8f..26060efc78 100644 --- a/backend/internal/service/account_test_service.go +++ b/backend/internal/service/account_test_service.go @@ -2585,14 +2585,7 @@ func (s *AccountTestService) buildGeminiOAuthRequest(ctx context.Context, accoun } func (s *AccountTestService) buildGeminiServiceAccountRequest(ctx context.Context, account *Account, modelID string, payload []byte) (*http.Request, error) { - if s.geminiTokenProvider == nil { - return nil, fmt.Errorf("gemini token provider not configured") - } - accessToken, err := s.geminiTokenProvider.GetAccessToken(ctx, account) - if err != nil { - return nil, fmt.Errorf("failed to get service account access token: %w", err) - } - fullURL, err := buildVertexGeminiURL(account.VertexProjectID(), account.VertexLocation(modelID), modelID, "streamGenerateContent", true) + fullURL, err := buildVertexGeminiAccountURL(account, modelID, "streamGenerateContent", true) if err != nil { return nil, err } @@ -2601,7 +2594,9 @@ func (s *AccountTestService) buildGeminiServiceAccountRequest(ctx context.Contex return nil, err } req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", "Bearer "+accessToken) + if err := setVertexGeminiAuth(ctx, req, account, s.geminiTokenProvider); err != nil { + return nil, fmt.Errorf("failed to authenticate vertex request: %w", err) + } return req, nil } diff --git a/backend/internal/service/gemini_cached_content.go b/backend/internal/service/gemini_cached_content.go index 7c3c6b14a6..0fe24249c3 100644 --- a/backend/internal/service/gemini_cached_content.go +++ b/backend/internal/service/gemini_cached_content.go @@ -112,13 +112,14 @@ func IsGeminiExplicitCacheAccountType(account *Account) bool { } // IsGeminiExplicitCacheAccount 判断完整账号是否支持显式缓存：官方地址的 AI Studio API Key 账号、 -// 声明上游支持显式缓存的自定义地址 API Key 账号，或 Vertex service account 账号。 +// 声明上游支持显式缓存的自定义地址 API Key 账号，或 Service Account JSON 认证的 Vertex 账号。 +// API Key 认证的 Vertex 账号不参与显式缓存：Vertex 缓存接口只接受 OAuth access token，不接受 API Key。 func IsGeminiExplicitCacheAccount(account *Account) bool { if !IsGeminiExplicitCacheAccountType(account) { return false } if account.Type == AccountTypeServiceAccount { - return true + return !account.IsVertexAPIKey() } if strings.TrimSpace(account.GetCredential("api_key")) == "" { return false diff --git a/backend/internal/service/gemini_cached_content_test.go b/backend/internal/service/gemini_cached_content_test.go index f53ad6c94d..0aec766105 100644 --- a/backend/internal/service/gemini_cached_content_test.go +++ b/backend/internal/service/gemini_cached_content_test.go @@ -186,6 +186,7 @@ func TestIsGeminiExplicitCacheAccount(t *testing.T) { {"official host with path prefix", &Account{Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "k", "base_url": "https://generativelanguage.googleapis.com/antigravity"}}, false}, {"missing api key", &Account{Platform: PlatformGemini, Type: AccountTypeAPIKey}, false}, {"vertex service account", &Account{Platform: PlatformGemini, Type: AccountTypeServiceAccount}, true}, + {"vertex api key", &Account{Platform: PlatformGemini, Type: AccountTypeServiceAccount, Credentials: map[string]any{"auth_mode": "apikey", "api_key": "k", "project_id": "my-project"}}, false}, {"oauth", &Account{Platform: PlatformGemini, Type: AccountTypeOAuth}, false}, {"antigravity api key", &Account{Platform: PlatformAntigravity, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "k"}}, false}, } diff --git a/backend/internal/service/gemini_chat_completions_compat_service.go b/backend/internal/service/gemini_chat_completions_compat_service.go index a98c94cf93..1ad2b487e8 100644 --- a/backend/internal/service/gemini_chat_completions_compat_service.go +++ b/backend/internal/service/gemini_chat_completions_compat_service.go @@ -399,19 +399,11 @@ func (s *GeminiMessagesCompatService) buildGeminiChatCompletionsUpstreamRequestF case AccountTypeServiceAccount: return func(ctx context.Context) (*http.Request, string, error) { - if s.tokenProvider == nil { - return nil, "", errors.New("gemini token provider not configured") - } - accessToken, err := s.tokenProvider.GetAccessToken(ctx, account) - if err != nil { - return nil, "", err - } - action := "generateContent" if clientStream { action = "streamGenerateContent" } - fullURL, err := buildVertexGeminiURL(account.VertexProjectID(), account.VertexLocation(mappedModel), mappedModel, action, clientStream) + fullURL, err := buildVertexGeminiAccountURL(account, mappedModel, action, clientStream) if err != nil { return nil, "", err } @@ -422,7 +414,9 @@ func (s *GeminiMessagesCompatService) buildGeminiChatCompletionsUpstreamRequestF return nil, "", err } upstreamReq.Header.Set("Content-Type", "application/json") - upstreamReq.Header.Set("Authorization", "Bearer "+accessToken) + if err := setVertexGeminiAuth(ctx, upstreamReq, account, s.tokenProvider); err != nil { + return nil, "", err + } return upstreamReq, "x-request-id", nil }, "x-request-id" diff --git a/backend/internal/service/gemini_messages_compat_service.go b/backend/internal/service/gemini_messages_compat_service.go index c3edbd2fd4..dec439ea8b 100644 --- a/backend/internal/service/gemini_messages_compat_service.go +++ b/backend/internal/service/gemini_messages_compat_service.go @@ -770,19 +770,11 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex case AccountTypeServiceAccount: buildReq = func(ctx context.Context) (*http.Request, string, error) { - if s.tokenProvider == nil { - return nil, "", errors.New("gemini token provider not configured") - } - accessToken, err := s.tokenProvider.GetAccessToken(ctx, account) - if err != nil { - return nil, "", err - } - action := "generateContent" if req.Stream { action = "streamGenerateContent" } - fullURL, err := buildVertexGeminiURL(account.VertexProjectID(), account.VertexLocation(mappedModel), mappedModel, action, req.Stream) + fullURL, err := buildVertexGeminiAccountURL(account, mappedModel, action, req.Stream) if err != nil { return nil, "", err } @@ -793,7 +785,9 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex return nil, "", err } upstreamReq.Header.Set("Content-Type", "application/json") - upstreamReq.Header.Set("Authorization", "Bearer "+accessToken) + if err := setVertexGeminiAuth(ctx, upstreamReq, account, s.tokenProvider); err != nil { + return nil, "", err + } return upstreamReq, "x-request-id", nil } requestIDHeader = "x-request-id" @@ -1311,15 +1305,7 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin. case AccountTypeServiceAccount: buildReq = func(ctx context.Context) (*http.Request, string, error) { - if s.tokenProvider == nil { - return nil, "", errors.New("gemini token provider not configured") - } - accessToken, err := s.tokenProvider.GetAccessToken(ctx, account) - if err != nil { - return nil, "", err - } - - fullURL, err := buildVertexGeminiURL(account.VertexProjectID(), account.VertexLocation(mappedModel), mappedModel, upstreamAction, useUpstreamStream) + fullURL, err := buildVertexGeminiAccountURL(account, mappedModel, upstreamAction, useUpstreamStream) if err != nil { return nil, "", err } @@ -1329,7 +1315,9 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin. return nil, "", err } upstreamReq.Header.Set("Content-Type", "application/json") - upstreamReq.Header.Set("Authorization", "Bearer "+accessToken) + if err := setVertexGeminiAuth(ctx, upstreamReq, account, s.tokenProvider); err != nil { + return nil, "", err + } return upstreamReq, "x-request-id", nil } requestIDHeader = "x-request-id" diff --git a/backend/internal/service/gemini_vertex_api_key_forward_test.go b/backend/internal/service/gemini_vertex_api_key_forward_test.go new file mode 100644 index 0000000000..bf8be6c1c0 --- /dev/null +++ b/backend/internal/service/gemini_vertex_api_key_forward_test.go @@ -0,0 +1,101 @@ +//go:build unit + +package service + +import ( + "bytes" + "context" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +const geminiVertexAPIKeyTestResponse = `{"candidates":[{"content":{"parts":[{"text":"hello"}],"role":"model"},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":5,"totalTokenCount":15}}` + +func newGeminiVertexAPIKeyTestService() (*GeminiMessagesCompatService, *geminiCompatHTTPUpstreamStub) { + httpStub := &geminiCompatHTTPUpstreamStub{ + response: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}, "X-Request-Id": []string{"vertex-req-1"}}, + Body: io.NopCloser(strings.NewReader(geminiVertexAPIKeyTestResponse)), + }, + } + // tokenProvider 留空：API Key 认证不依赖 access token。 + return &GeminiMessagesCompatService{httpUpstream: httpStub, cfg: &config.Config{}}, httpStub +} + +func geminiVertexAPIKeyTestAccount(projectID string) *Account { + credentials := map[string]any{ + "auth_mode": "apikey", + "api_key": "vertex-test-key", + "location": "us-central1", + "tier_id": "vertex", + } + if projectID != "" { + credentials["project_id"] = projectID + } + return &Account{ + ID: 811, + Platform: PlatformGemini, + Type: AccountTypeServiceAccount, + Credentials: credentials, + Concurrency: 1, + } +} + +func requireVertexAPIKeyAuth(t *testing.T, req *http.Request) { + t.Helper() + require.NotNil(t, req) + require.Equal(t, "vertex-test-key", req.Header.Get("x-goog-api-key")) + require.Empty(t, req.Header.Get("Authorization")) +} + +func TestGeminiForwardNative_VertexAPIKeyWithoutProjectUsesExpressEndpoint(t *testing.T) { + gin.SetMode(gin.TestMode) + svc, httpStub := newGeminiVertexAPIKeyTestService() + c, _ := newGeminiNativeTestContext(t) + + result, err := svc.ForwardNative(context.Background(), c, geminiVertexAPIKeyTestAccount(""), + "gemini-3-flash-preview", "generateContent", false, []byte(`{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`)) + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, 1, httpStub.calls) + require.Equal(t, "https://aiplatform.googleapis.com/v1/publishers/google/models/gemini-3-flash-preview:generateContent", httpStub.lastReq.URL.String()) + requireVertexAPIKeyAuth(t, httpStub.lastReq) +} + +func TestGeminiMessagesCompatServiceForward_VertexAPIKeyWithProjectUsesProjectEndpoint(t *testing.T) { + gin.SetMode(gin.TestMode) + svc, httpStub := newGeminiVertexAPIKeyTestService() + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) + body := []byte(`{"model":"gemini-3-flash-preview","max_tokens":16,"messages":[{"role":"user","content":"hello"}]}`) + + result, err := svc.Forward(context.Background(), c, geminiVertexAPIKeyTestAccount("my-project"), body) + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, "https://us-central1-aiplatform.googleapis.com/v1/projects/my-project/locations/us-central1/publishers/google/models/gemini-3-flash-preview:generateContent", httpStub.lastReq.URL.String()) + requireVertexAPIKeyAuth(t, httpStub.lastReq) +} + +func TestGeminiForwardAsChatCompletions_VertexAPIKeyUsesAPIKeyHeader(t *testing.T) { + gin.SetMode(gin.TestMode) + svc, httpStub := newGeminiVertexAPIKeyTestService() + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + body := []byte(`{"model":"gemini-3-flash-preview","messages":[{"role":"user","content":"hi"}]}`) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body)) + + result, err := svc.ForwardAsChatCompletions(context.Background(), c, geminiVertexAPIKeyTestAccount(""), body) + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, http.StatusOK, rec.Code) + require.Equal(t, "https://aiplatform.googleapis.com/v1/publishers/google/models/gemini-3-flash-preview:generateContent", httpStub.lastReq.URL.String()) + requireVertexAPIKeyAuth(t, httpStub.lastReq) +} diff --git a/backend/internal/service/vertex_service_account.go b/backend/internal/service/vertex_service_account.go index 7ccbeee43c..2a19a50e21 100644 --- a/backend/internal/service/vertex_service_account.go +++ b/backend/internal/service/vertex_service_account.go @@ -29,6 +29,8 @@ const ( vertexServiceAccountCacheSkew = 5 * time.Minute vertexLockWaitTime = 200 * time.Millisecond vertexAnthropicVersion = "vertex-2023-10-16" + vertexGlobalHost = "aiplatform.googleapis.com" + vertexAuthModeAPIKey = "apikey" ) var ( @@ -58,6 +60,13 @@ func (a *Account) IsVertexServiceAccount() bool { return a != nil && a.Type == AccountTypeServiceAccount } +// IsVertexAPIKey 判断 Gemini 平台的 Vertex 账号是否使用 API Key 认证（credentials.auth_mode=apikey）。 +// Anthropic 平台的 Vertex 账号只支持 Service Account JSON。 +func (a *Account) IsVertexAPIKey() bool { + return a.IsVertexServiceAccount() && a.Platform == PlatformGemini && + strings.TrimSpace(a.GetCredential("auth_mode")) == vertexAuthModeAPIKey +} + func (a *Account) VertexProjectID() string { if a == nil { return "" @@ -283,6 +292,63 @@ func exchangeVertexServiceAccountToken(ctx context.Context, key *vertexServiceAc return parsed.AccessToken, ttl, nil } +// buildVertexGeminiAccountURL 按账号构造 Vertex Gemini 请求地址：未配置 project_id 的 API Key 账号 +// 走 Express mode 全局端点，其余账号走项目级端点。 +func buildVertexGeminiAccountURL(account *Account, model, action string, stream bool) (string, error) { + if account.IsVertexAPIKey() && account.VertexProjectID() == "" { + return buildVertexGeminiExpressURL(model, action, stream) + } + return buildVertexGeminiURL(account.VertexProjectID(), account.VertexLocation(model), model, action, stream) +} + +// buildVertexGeminiExpressURL 构造不带项目与区域的全局端点地址，仅 API Key 认证可用。 +func buildVertexGeminiExpressURL(model, action string, stream bool) (string, error) { + model = strings.TrimSpace(model) + action = strings.TrimSpace(action) + if model == "" { + return "", errors.New("vertex model is required") + } + if err := validateVertexGeminiAction(action); err != nil { + return "", err + } + u := fmt.Sprintf("https://%s/v1/publishers/google/models/%s:%s", vertexGlobalHost, url.PathEscape(model), action) + if stream { + u += "?alt=sse" + } + return u, nil +} + +func validateVertexGeminiAction(action string) error { + switch action { + case "generateContent", "streamGenerateContent", "countTokens": + return nil + default: + return fmt.Errorf("unsupported vertex gemini action: %s", action) + } +} + +// setVertexGeminiAuth 给 Vertex Gemini 请求设置认证头：API Key 账号用 x-goog-api-key， +// Service Account JSON 账号用换取的 Bearer access token。 +func setVertexGeminiAuth(ctx context.Context, req *http.Request, account *Account, tokenProvider *GeminiTokenProvider) error { + if account.IsVertexAPIKey() { + apiKey := strings.TrimSpace(account.GetCredential("api_key")) + if apiKey == "" { + return errors.New("vertex api_key not configured") + } + req.Header.Set("x-goog-api-key", apiKey) + return nil + } + if tokenProvider == nil { + return errors.New("gemini token provider not configured") + } + accessToken, err := tokenProvider.GetAccessToken(ctx, account) + if err != nil { + return err + } + req.Header.Set("Authorization", "Bearer "+accessToken) + return nil +} + func buildVertexGeminiURL(projectID, location, model, action string, stream bool) (string, error) { projectID = strings.TrimSpace(projectID) location = strings.TrimSpace(location) @@ -300,14 +366,12 @@ func buildVertexGeminiURL(projectID, location, model, action string, stream bool if model == "" { return "", errors.New("vertex model is required") } - switch action { - case "generateContent", "streamGenerateContent", "countTokens": - default: - return "", fmt.Errorf("unsupported vertex gemini action: %s", action) + if err := validateVertexGeminiAction(action); err != nil { + return "", err } host := fmt.Sprintf("%s-aiplatform.googleapis.com", location) if location == "global" { - host = "aiplatform.googleapis.com" + host = vertexGlobalHost } u := fmt.Sprintf( "https://%s/v1/projects/%s/locations/%s/publishers/google/models/%s:%s", diff --git a/backend/internal/service/vertex_service_account_test.go b/backend/internal/service/vertex_service_account_test.go index 68c756eaa2..77cb6617c7 100644 --- a/backend/internal/service/vertex_service_account_test.go +++ b/backend/internal/service/vertex_service_account_test.go @@ -30,6 +30,62 @@ func TestBuildVertexGeminiURLUsesGlobalEndpointHost(t *testing.T) { require.Equal(t, "https://aiplatform.googleapis.com/v1/projects/my-project/locations/global/publishers/google/models/gemini-3-flash-preview:streamGenerateContent?alt=sse", got) } +func TestIsVertexAPIKey(t *testing.T) { + apiKeyCreds := map[string]any{"auth_mode": "apikey", "api_key": "k"} + require.True(t, (&Account{Platform: PlatformGemini, Type: AccountTypeServiceAccount, Credentials: apiKeyCreds}).IsVertexAPIKey()) + require.False(t, (&Account{Platform: PlatformGemini, Type: AccountTypeServiceAccount, Credentials: map[string]any{"service_account_json": "{}"}}).IsVertexAPIKey()) + require.False(t, (&Account{Platform: PlatformAnthropic, Type: AccountTypeServiceAccount, Credentials: apiKeyCreds}).IsVertexAPIKey()) + require.False(t, (&Account{Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: apiKeyCreds}).IsVertexAPIKey()) +} + +func TestBuildVertexGeminiAccountURL(t *testing.T) { + cases := []struct { + name string + credentials map[string]any + want string + }{ + { + name: "api key without project uses express global endpoint", + credentials: map[string]any{"auth_mode": "apikey", "api_key": "k", "location": "us-central1"}, + want: "https://aiplatform.googleapis.com/v1/publishers/google/models/gemini-3-flash-preview:streamGenerateContent?alt=sse", + }, + { + name: "api key with project uses project endpoint", + credentials: map[string]any{"auth_mode": "apikey", "api_key": "k", "project_id": "my-project", "location": "us-central1"}, + want: "https://us-central1-aiplatform.googleapis.com/v1/projects/my-project/locations/us-central1/publishers/google/models/gemini-3-flash-preview:streamGenerateContent?alt=sse", + }, + { + name: "service account uses project endpoint", + credentials: map[string]any{"project_id": "sa-project", "location": "global"}, + want: "https://aiplatform.googleapis.com/v1/projects/sa-project/locations/global/publishers/google/models/gemini-3-flash-preview:streamGenerateContent?alt=sse", + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + account := &Account{Platform: PlatformGemini, Type: AccountTypeServiceAccount, Credentials: tc.credentials} + got, err := buildVertexGeminiAccountURL(account, "gemini-3-flash-preview", "streamGenerateContent", true) + require.NoError(t, err) + require.Equal(t, tc.want, got) + }) + } + + account := &Account{Platform: PlatformGemini, Type: AccountTypeServiceAccount, Credentials: map[string]any{"auth_mode": "apikey", "api_key": "k"}} + _, err := buildVertexGeminiAccountURL(account, "gemini-3-flash-preview", "rawPredict", false) + require.ErrorContains(t, err, "unsupported vertex gemini action") +} + +func TestSetVertexGeminiAuthUsesAPIKeyHeader(t *testing.T) { + account := &Account{Platform: PlatformGemini, Type: AccountTypeServiceAccount, Credentials: map[string]any{"auth_mode": "apikey", "api_key": " vertex-key "}} + req := httptest.NewRequest(http.MethodPost, "https://aiplatform.googleapis.com/v1/publishers/google/models/gemini-3-flash-preview:generateContent", nil) + + require.NoError(t, setVertexGeminiAuth(context.Background(), req, account, nil)) + require.Equal(t, "vertex-key", req.Header.Get("x-goog-api-key")) + require.Empty(t, req.Header.Get("Authorization")) + + account.Credentials["api_key"] = "" + require.ErrorContains(t, setVertexGeminiAuth(context.Background(), req, account, nil), "vertex api_key not configured") +} + func TestBuildVertexAnthropicURL(t *testing.T) { got, err := buildVertexAnthropicURL("my-project", "us-east5", "claude-sonnet-4-5@20250929", false) require.NoError(t, err) diff --git a/frontend/src/components/account/CreateAccountModal.vue b/frontend/src/components/account/CreateAccountModal.vue index e4aa3dfc4c..5ea21bbb87 100644 --- a/frontend/src/components/account/CreateAccountModal.vue +++ b/frontend/src/components/account/CreateAccountModal.vue @@ -833,7 +833,7 @@ {{ t('admin.accounts.vertexLabel') }} - {{ t('admin.accounts.vertexDesc') }} + {{ t('admin.accounts.vertexGeminiDesc') }} @@ -1215,92 +1215,169 @@
-
- - -
-
-
-
- - {{ vertexClientEmail ? t('admin.accounts.vertexSaJsonLoaded') : t('admin.accounts.vertexSaJsonDrop') }} -
-

- {{ vertexClientEmail ? t('admin.accounts.vertexSaJsonKeyHidden') : t('admin.accounts.vertexSaJsonDropHint') }} -

-
- -
-
-
Project ID: {{ vertexProjectId }}
-
Client Email: {{ vertexClientEmail }}
-
+
+ +
+ +
-

{{ t('admin.accounts.vertexSaJsonUploadHint') }}

-
+ + +
@@ -4651,6 +4728,17 @@ const vertexProjectId = ref('') const vertexClientEmail = ref('') const vertexLocation = ref('global') const vertexServiceAccountDragActive = ref(false) +const vertexAuthMode = ref<'service_account' | 'apikey'>('service_account') +const vertexApiKey = ref('') +// API Key 认证只对 Gemini 平台开放，Anthropic 平台的 Vertex 账号始终使用 Service Account JSON。 +const vertexUsesApiKey = computed(() => form.platform === 'gemini' && vertexAuthMode.value === 'apikey') +// 两种认证方式的 Project ID 来源不同（JSON 读取 / 手填），切换时清空凭证字段避免串用。 +watch(vertexAuthMode, () => { + vertexServiceAccountJson.value = '' + vertexProjectId.value = '' + vertexClientEmail.value = '' + vertexApiKey.value = '' +}) const tempUnschedEnabled = ref(false) const tempUnschedRules = ref([]) const getModelMappingKey = createStableObjectKeyResolver('create-model-mapping') @@ -5055,6 +5143,8 @@ watch( vertexProjectId.value = '' vertexClientEmail.value = '' vertexLocation.value = 'global' + vertexAuthMode.value = 'service_account' + vertexApiKey.value = '' // Reset Anthropic/Antigravity-specific settings when switching to other platforms if (newPlatform !== 'anthropic' && newPlatform !== 'antigravity') { interceptWarmupRequests.value = false @@ -5558,6 +5648,8 @@ const resetForm = () => { vertexProjectId.value = '' vertexClientEmail.value = '' vertexLocation.value = 'global' + vertexAuthMode.value = 'service_account' + vertexApiKey.value = '' tempUnschedEnabled.value = false tempUnschedRules.value = [] geminiOAuthType.value = 'code_assist' @@ -5906,6 +5998,28 @@ const handleSubmit = async () => { appStore.showError(t('admin.accounts.pleaseEnterAccountName')) return } + if (vertexUsesApiKey.value) { + if (!vertexApiKey.value.trim()) { + appStore.showError(t('admin.accounts.vertexApiKeyRequired')) + return + } + const projectId = vertexProjectId.value.trim() + if (projectId && !vertexLocation.value.trim()) { + appStore.showError(t('admin.accounts.vertexLocationRequired')) + return + } + const credentials: Record = { + auth_mode: 'apikey', + api_key: vertexApiKey.value.trim(), + location: vertexLocation.value.trim() || 'global', + tier_id: 'vertex' + } + if (projectId) { + credentials.project_id = projectId + } + await createAccountAndFinish(form.platform, 'service_account' as AccountType, credentials) + return + } if (!parseVertexServiceAccountJson()) { return } diff --git a/frontend/src/components/account/EditAccountModal.vue b/frontend/src/components/account/EditAccountModal.vue index d5aa3fdaea..1fdd3727f2 100644 --- a/frontend/src/components/account/EditAccountModal.vue +++ b/frontend/src/components/account/EditAccountModal.vue @@ -957,8 +957,30 @@
+
+ + +

{{ t('admin.accounts.vertexApiKeyHint') }}

+
-
+
+ + +

{{ t('admin.accounts.vertexApiKeyProjectIdHint') }}

+
+

{{ t('admin.accounts.vertexSaJsonEditHint') }}

-
+
-
+
- @@ -6061,7 +6062,7 @@ const handleSubmit = async () => { base_url: apiKeyBaseUrl.value.trim() || defaultBaseUrl, api_key: apiKeyValue.value.trim() } - if (form.platform === 'gemini') { + if (form.platform === 'gemini' && isOfficialGeminiBaseUrl(apiKeyBaseUrl.value)) { credentials.tier_id = geminiTierAIStudio.value } diff --git a/frontend/src/components/account/__tests__/AccountUsageCell.spec.ts b/frontend/src/components/account/__tests__/AccountUsageCell.spec.ts index 71c5f5d7a5..d595102f66 100644 --- a/frontend/src/components/account/__tests__/AccountUsageCell.spec.ts +++ b/frontend/src/components/account/__tests__/AccountUsageCell.spec.ts @@ -1709,6 +1709,101 @@ describe('AccountUsageCell', () => { expect(wrapper.text()).toContain('0') expect(wrapper.text()).toContain('A $0.00') expect(wrapper.text()).toContain('U $0.00') + expect(wrapper.text()).toContain('vertex') + expect(wrapper.text()).not.toContain('admin.accounts.gemini.quotaPolicy.title') + expect(wrapper.text()).not.toContain('admin.accounts.gemini.rateLimit.unlimited') + expect(getUsage).not.toHaveBeenCalled() + }) + + it('Gemini 第三方上游 API Key 按普通 Key 账号展示今日统计与额度，不显示 Gemini 徽章', async () => { + const wrapper = mount(AccountUsageCell, { + props: { + account: makeAccount({ + id: 4002, + platform: 'gemini', + type: 'apikey', + credentials: { + base_url: 'https://relay.example.com/gemini', + api_key: 'sk-relay', + tier_id: 'aistudio_free' + }, + extra: {}, + quota_daily_limit: 10, + quota_daily_used: 2.5 + }), + todayStats: { + requests: 123, + tokens: 2300000, + cost: 1.79, + standard_cost: 1.79, + user_cost: 0 + } + }, + global: { + stubs: { + UsageProgressBar: true, + AccountQuotaInfo: true + } + } + }) + + await flushPromises() + + const text = wrapper.text() + expect(text).not.toContain('api key') + expect(text).not.toContain('ai studio') + expect(text).not.toContain('free') + expect(text).not.toContain('relay.example.com') + expect(wrapper.find('account-quota-info-stub').exists()).toBe(false) + expect(text).toContain('123 req') + expect(text).toContain('A $1.79') + expect(text).toContain('U $0.00') + const bars = wrapper.findAll('usage-progress-bar-stub') + expect(bars.map((bar) => bar.attributes('label'))).toEqual(['1d']) + expect(text).not.toContain('admin.accounts.gemini.quotaPolicy.title') + expect(text).not.toContain('admin.accounts.gemini.quotaPolicy.simulatedNote') + expect(text).not.toContain('admin.accounts.gemini.rateLimit.unlimited') + expect(getUsage).not.toHaveBeenCalled() + }) + + it('Gemini 官方 API Key 保留 AI Studio 等级与模拟配额', async () => { + getUsage.mockResolvedValue({ + gemini_pro_daily: { + utilization: 20, + resets_at: '2026-10-11T07:00:00Z', + remaining_seconds: 3600 + } + }) + + const wrapper = mount(AccountUsageCell, { + props: { + account: makeAccount({ + id: 4003, + platform: 'gemini', + type: 'apikey', + credentials: { + base_url: 'https://generativelanguage.googleapis.com', + api_key: 'AIza-official', + tier_id: 'aistudio_paid' + }, + extra: {} + }) + }, + global: { + stubs: { + UsageProgressBar: true, + AccountQuotaInfo: true + } + } + }) + + await flushPromises() + + const text = wrapper.text() + expect(text).toContain('ai studio paid') + expect(text).toContain('admin.accounts.gemini.quotaPolicy.title') + expect(text).toContain('admin.accounts.gemini.quotaPolicy.simulatedNote') + expect(getUsage).toHaveBeenCalledWith(4003) }) it('Anthropic OAuth 会渲染 7d F (Fable) 进度条，且 7d S 逻辑保留', async () => { diff --git a/frontend/src/components/account/__tests__/CreateAccountModal.spec.ts b/frontend/src/components/account/__tests__/CreateAccountModal.spec.ts index 3d74b506d6..84cfa81123 100644 --- a/frontend/src/components/account/__tests__/CreateAccountModal.spec.ts +++ b/frontend/src/components/account/__tests__/CreateAccountModal.spec.ts @@ -931,6 +931,58 @@ describe('CreateAccountModal gemini explicit cache upstream switch', () => { explicit_cache_upstream_max_ttl_seconds: 1800 }) }) + + it('only offers the AI Studio tier for the official base url', async () => { + const wrapper = mountModal() + await selectButtonByText(wrapper, 'Gemini') + await selectButtonByText(wrapper, 'admin.accounts.gemini.accountType.apiKeyTitle') + expect(wrapper.find('[data-testid="gemini-oauth-tier"]').exists()).toBe(false) + expect(wrapper.find('[data-testid="gemini-apikey-note"]').exists()).toBe(true) + await wrapper.get('[data-testid="gemini-apikey-tier"]').setValue('aistudio_paid') + + const baseUrlInput = wrapper + .findAll('form#create-account-form input[type="text"]') + .find((input) => (input.element as HTMLInputElement).value === 'https://generativelanguage.googleapis.com') + expect(baseUrlInput).toBeDefined() + await baseUrlInput!.setValue('https://upstream.example.com') + expect(wrapper.find('[data-testid="gemini-apikey-tier"]').exists()).toBe(false) + expect(wrapper.find('[data-testid="gemini-apikey-note"]').exists()).toBe(false) + + await wrapper.get('form#create-account-form input[type="text"]').setValue('gemini relay') + await wrapper.get('form#create-account-form input[type="password"]').setValue('sk-upstream') + await wrapper.get('form#create-account-form').trigger('submit.prevent') + await flushPromises() + + expect(createAccountMock).toHaveBeenCalledTimes(1) + const credentials = createAccountMock.mock.calls[0]?.[0]?.credentials + expect(credentials).toMatchObject({ base_url: 'https://upstream.example.com' }) + expect(credentials).not.toHaveProperty('tier_id') + }) + + it('shows the OAuth tier fallback only on the OAuth form', async () => { + const wrapper = mountModal() + await selectButtonByText(wrapper, 'Gemini') + await selectButtonByText(wrapper, 'admin.accounts.gemini.accountType.oauthTitle') + expect(wrapper.find('[data-testid="gemini-oauth-tier"]').exists()).toBe(true) + expect(wrapper.find('[data-testid="gemini-apikey-note"]').exists()).toBe(false) + }) + + it('keeps the selected AI Studio tier for the official base url', async () => { + const wrapper = mountModal() + await selectButtonByText(wrapper, 'Gemini') + await selectButtonByText(wrapper, 'admin.accounts.gemini.accountType.apiKeyTitle') + await wrapper.get('[data-testid="gemini-apikey-tier"]').setValue('aistudio_paid') + await wrapper.get('form#create-account-form input[type="text"]').setValue('gemini official') + await wrapper.get('form#create-account-form input[type="password"]').setValue('AIza-official') + await wrapper.get('form#create-account-form').trigger('submit.prevent') + await flushPromises() + + expect(createAccountMock).toHaveBeenCalledTimes(1) + expect(createAccountMock.mock.calls[0]?.[0]?.credentials).toMatchObject({ + base_url: 'https://generativelanguage.googleapis.com', + tier_id: 'aistudio_paid' + }) + }) }) describe('CreateAccountModal Vertex API key auth', () => { diff --git a/frontend/src/i18n/locales/en/admin/accounts.ts b/frontend/src/i18n/locales/en/admin/accounts.ts index 95e8042471..317e8ec933 100644 --- a/frontend/src/i18n/locales/en/admin/accounts.ts +++ b/frontend/src/i18n/locales/en/admin/accounts.ts @@ -1445,8 +1445,8 @@ export default { accountType: { oauthTitle: 'OAuth (Gemini)', oauthDesc: 'Authorize with your Google account and choose an OAuth type.', - apiKeyTitle: 'API Key (AI Studio)', - apiKeyDesc: 'Fastest setup. Use an AIza API key.', + apiKeyTitle: 'API Key', + apiKeyDesc: 'An official AI Studio key, or a Gemini-compatible third-party upstream.', apiKeyNote: 'Best for light testing. Free tier has strict rate limits and data may be used for training.', apiKeyLink: 'Get API Key', diff --git a/frontend/src/i18n/locales/zh/admin/accounts.ts b/frontend/src/i18n/locales/zh/admin/accounts.ts index 129225846e..3e9c06fefa 100644 --- a/frontend/src/i18n/locales/zh/admin/accounts.ts +++ b/frontend/src/i18n/locales/zh/admin/accounts.ts @@ -1536,8 +1536,8 @@ export default { accountType: { oauthTitle: 'OAuth 授权（Gemini）', oauthDesc: '使用 Google 账号授权，并选择 OAuth 子类型。', - apiKeyTitle: 'API 密钥（AI Studio）', - apiKeyDesc: '最快接入方式，使用 AIza API Key。', + apiKeyTitle: 'API 密钥', + apiKeyDesc: 'AI Studio 官方 Key，或兼容 Gemini 协议的第三方上游。', apiKeyNote: '适合轻量测试。免费层限流严格，数据可能用于训练。', apiKeyLink: '获取 API Key', quotaLink: '配额说明' diff --git a/frontend/src/types/index.ts b/frontend/src/types/index.ts index cfd3f2f167..3d5dd452f7 100644 --- a/frontend/src/types/index.ts +++ b/frontend/src/types/index.ts @@ -1010,6 +1010,7 @@ export interface ProxyQualityCheckResult { export interface GeminiCredentials { // API Key authentication api_key?: string + base_url?: string // OAuth authentication access_token?: string From 1c68f9996df8e0bd7933d94d4ca826bc3dc570c7 Mon Sep 17 00:00:00 2001 From: feeeei Date: Sat, 10 Oct 2026 22:22:29 +0800 Subject: [PATCH 07/15] =?UTF-8?q?feat(gemini):=20=E9=80=82=E9=85=8D=20AI?= =?UTF-8?q?=20Studio=20AQ.=20=E6=A0=BC=E5=BC=8F=20Key?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - AI Studio 新建 Key 为 AQ.Ab 前缀的 auth key：logredact.RedactText 与 渠道监控 sanitizeErrorMessage 增加 AQ.Ab 规则，长度按 40 位以上放宽以覆盖格式变体 - 创建 / 编辑账号的 Key 占位符改为 AQ.Ab...，提示文案说明新旧两种前缀 与第三方上游 Key --- .../service/channel_monitor_checker.go | 4 +++- .../channel_monitor_checker_body_test.go | 11 ++++++++++ backend/internal/util/logredact/redact.go | 3 +++ .../internal/util/logredact/redact_test.go | 22 +++++++++++++++++++ .../components/account/CreateAccountModal.vue | 2 +- .../components/account/EditAccountModal.vue | 2 +- .../src/i18n/locales/en/admin/accounts.ts | 2 +- .../src/i18n/locales/zh/admin/accounts.ts | 2 +- 8 files changed, 43 insertions(+), 5 deletions(-) diff --git a/backend/internal/service/channel_monitor_checker.go b/backend/internal/service/channel_monitor_checker.go index ddb5e78317..c8b4fdf069 100644 --- a/backend/internal/service/channel_monitor_checker.go +++ b/backend/internal/service/channel_monitor_checker.go @@ -636,6 +636,8 @@ var monitorAPIKeyPatterns = []struct { {regexp.MustCompile(`xai-[A-Za-z0-9_-]{6,}`), "xai-***REDACTED***"}, // Gemini / Google API Key：固定前缀 + 35 位 {regexp.MustCompile(`AIza[A-Za-z0-9_-]{35}`), "AIza***REDACTED***"}, + // Gemini auth key：AQ.Ab 前缀 + 48 位，长度放宽以覆盖格式变体 + {regexp.MustCompile(`AQ\.Ab[A-Za-z0-9_-]{40,}`), "AQ.***REDACTED***"}, // JWT 三段式（Bearer 后常出现）：eyJxxx.eyJxxx.signature {regexp.MustCompile(`eyJ[A-Za-z0-9_-]{8,}\.eyJ[A-Za-z0-9_-]{8,}\.[A-Za-z0-9_-]{8,}`), "eyJ***REDACTED.JWT***"}, } @@ -643,7 +645,7 @@ var monitorAPIKeyPatterns = []struct { // sanitizeErrorMessage 擦除错误/响应文本中可能泄露的 API key。 // 处理两类来源： // 1. URL query 中的 ?key= / ?api_key= 等（Go *url.Error 会回填完整 URL） -// 2. 上游 HTTP body 文本里直接出现的 sk-* / xai-* / AIza* / JWT 等密钥碎片 +// 2. 上游 HTTP body 文本里直接出现的 sk-* / xai-* / AIza* / AQ.* / JWT 等密钥碎片 // // 注意：与 gemini_messages_compat_service.go 的 sanitizeUpstreamErrorMessage 关注点类似但参数集更广， // 监控模块独立维护，避免互相耦合。 diff --git a/backend/internal/service/channel_monitor_checker_body_test.go b/backend/internal/service/channel_monitor_checker_body_test.go index 922eb86732..5c347cae89 100644 --- a/backend/internal/service/channel_monitor_checker_body_test.go +++ b/backend/internal/service/channel_monitor_checker_body_test.go @@ -279,6 +279,17 @@ func TestRunCheckForModel_Grok_RedactsXAIKeyFromUpstreamBody(t *testing.T) { } } +func TestSanitizeErrorMessage_RedactsGeminiAuthKey(t *testing.T) { + authKey := "AQ." + "Ab" + strings.Repeat("Xy9_-z", 8) + got := sanitizeErrorMessage(`{"error":{"message":"API key not valid: ` + authKey + `"}}`) + if strings.Contains(got, authKey) { + t.Fatalf("Gemini auth key leaked: %q", got) + } + if !strings.Contains(got, "AQ.***REDACTED***") { + t.Fatalf("expected redaction marker, got %q", got) + } +} + func TestRunCheckForModel_OpenAIResponses_DefaultRequest(t *testing.T) { h := &openAICaptureHandler{} endpoint := setupFakeOpenAI(t, h) diff --git a/backend/internal/util/logredact/redact.go b/backend/internal/util/logredact/redact.go index 9249b761c7..8be6448a18 100644 --- a/backend/internal/util/logredact/redact.go +++ b/backend/internal/util/logredact/redact.go @@ -42,6 +42,8 @@ type textRedactPatterns struct { var ( reGOCSPX = regexp.MustCompile(`GOCSPX-[0-9A-Za-z_-]{24,}`) reAIza = regexp.MustCompile(`AIza[0-9A-Za-z_-]{35}`) + // Gemini auth key：AQ.Ab 前缀 + 48 位，长度放宽以覆盖格式变体。 + reGeminiAuthKey = regexp.MustCompile(`AQ\.Ab[0-9A-Za-z_-]{40,}`) defaultTextRedactPatterns = compileTextRedactPatterns(nil) extraTextPatternCache sync.Map // map[string]*textRedactPatterns @@ -99,6 +101,7 @@ func RedactText(input string, extraKeys ...string) string { out := input out = reGOCSPX.ReplaceAllString(out, "GOCSPX-***") out = reAIza.ReplaceAllString(out, "AIza***") + out = reGeminiAuthKey.ReplaceAllString(out, "AQ.***") out = patterns.reJSONLike.ReplaceAllString(out, `$1***$3`) out = patterns.reQueryLike.ReplaceAllString(out, `$1=***`) out = patterns.rePlain.ReplaceAllString(out, `$1$2***`) diff --git a/backend/internal/util/logredact/redact_test.go b/backend/internal/util/logredact/redact_test.go index 266db69dbd..e615f33e0e 100644 --- a/backend/internal/util/logredact/redact_test.go +++ b/backend/internal/util/logredact/redact_test.go @@ -38,6 +38,28 @@ func TestRedactText_GOCSPX(t *testing.T) { } } +func TestRedactText_GeminiAPIKeys(t *testing.T) { + legacyKey := "AI" + "za" + strings.Repeat("Ab1_-", 7) + authKey := "AQ." + "Ab" + strings.Repeat("Xy9_-z", 8) + in := "upstream rejected " + legacyKey + " and " + authKey + " for model gemini-3.1-pro" + out := RedactText(in) + if strings.Contains(out, legacyKey) || strings.Contains(out, authKey) { + t.Fatalf("expected Gemini keys redacted, got %q", out) + } + for _, want := range []string{"AIza***", "AQ.***", "for model gemini-3.1-pro"} { + if !strings.Contains(out, want) { + t.Fatalf("expected %q in %q", want, out) + } + } +} + +func TestRedactText_GeminiAuthKeyPrefixAloneUnchanged(t *testing.T) { + in := "key format AQ.Ab... is documented" + if out := RedactText(in); out != in { + t.Fatalf("expected short AQ. prefix unchanged, got %q", out) + } +} + func TestRedactText_ExtraKeyCacheUsesNormalizedSortedKey(t *testing.T) { clearExtraTextPatternCache() diff --git a/frontend/src/components/account/CreateAccountModal.vue b/frontend/src/components/account/CreateAccountModal.vue index 2dccc59363..315feae303 100644 --- a/frontend/src/components/account/CreateAccountModal.vue +++ b/frontend/src/components/account/CreateAccountModal.vue @@ -4253,7 +4253,7 @@ const apiKeyValuePlaceholder = computed(() => { case 'openai': return 'sk-proj-...' case 'gemini': - return 'AIza...' + return 'AQ.Ab...' case 'grok': return 'xai-...' case 'kimi': diff --git a/frontend/src/components/account/EditAccountModal.vue b/frontend/src/components/account/EditAccountModal.vue index 1fdd3727f2..9bc4f9167a 100644 --- a/frontend/src/components/account/EditAccountModal.vue +++ b/frontend/src/components/account/EditAccountModal.vue @@ -239,7 +239,7 @@ account.platform === 'openai' ? 'sk-proj-...' : account.platform === 'gemini' - ? 'AIza...' + ? 'AQ.Ab...' : account.platform === 'antigravity' ? 'sk-...' : account.platform === 'grok' diff --git a/frontend/src/i18n/locales/en/admin/accounts.ts b/frontend/src/i18n/locales/en/admin/accounts.ts index 317e8ec933..bb6fa68552 100644 --- a/frontend/src/i18n/locales/en/admin/accounts.ts +++ b/frontend/src/i18n/locales/en/admin/accounts.ts @@ -1422,7 +1422,7 @@ export default { modelPassthroughDesc: 'All model requests are forwarded directly to the Gemini API without model restrictions or mappings.', baseUrlHint: 'Leave default for official Gemini API', - apiKeyHint: 'Your Gemini API Key (starts with AIza)', + apiKeyHint: 'New AI Studio keys start with AQ. (legacy keys start with AIza); for a third-party upstream, use the key it provides', tier: { label: 'Account Tier', hint: 'Tip: The system will try to auto-detect the tier first; if auto-detection is unavailable or fails, your selected tier is used as a fallback (simulated quota).', diff --git a/frontend/src/i18n/locales/zh/admin/accounts.ts b/frontend/src/i18n/locales/zh/admin/accounts.ts index 3e9c06fefa..73a2f3f21b 100644 --- a/frontend/src/i18n/locales/zh/admin/accounts.ts +++ b/frontend/src/i18n/locales/zh/admin/accounts.ts @@ -1513,7 +1513,7 @@ export default { modelPassthrough: 'Gemini 直接转发模型', modelPassthroughDesc: '所有模型请求将直接转发至 Gemini API，不进行模型限制或映射。', baseUrlHint: '留空使用官方 Gemini API', - apiKeyHint: '您的 Gemini API Key（以 AIza 开头）', + apiKeyHint: 'AI Studio 新建的 Key 以 AQ. 开头（旧版以 AIza 开头）；第三方上游填写其提供的 Key', tier: { label: '账号等级', hint: '提示：系统会优先尝试自动识别账号等级；若自动识别不可用或失败，则使用你选择的等级作为回退（本地模拟配额）。', From 6f9b02f12c3d414ed324460a0e82d7793a471f72 Mon Sep 17 00:00:00 2001 From: feeeei Date: Sat, 10 Oct 2026 19:12:10 +0800 Subject: [PATCH 08/15] =?UTF-8?q?feat(gemini):=20Vertex=20=E6=9C=8D?= =?UTF-8?q?=E5=8A=A1=E8=B4=A6=E5=8F=B7=E6=94=AF=E6=8C=81=E5=90=8C=E6=AD=A5?= =?UTF-8?q?=E4=B8=8A=E6=B8=B8=E6=A8=A1=E5=9E=8B=E5=88=97=E8=A1=A8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Gemini 平台 Service Account 账号改用 v1beta1 publishers/google/models 分页拉取模型（pageSize 300，最多 20 页），按账号默认区域选端点； 只保留 publishers/google/models/gemini-* 条目，排除区域端点里的 Gemma、Veo、AutoML 等模型花园条目 - 凭据缺失或 location 非法返回配置错误（400），且不发起 token 交换； API Key 认证的 Vertex 账号返回不支持 - 翻页途中的 404/405 不再触发配置模型回退 - 模型列表的请求、读体与状态码检查抽成 readUpstreamModelsResponse 供各平台共用 --- backend/internal/service/upstream_models.go | 165 +++++++++-- .../service/upstream_models_vertex_test.go | 276 ++++++++++++++++++ 2 files changed, 419 insertions(+), 22 deletions(-) create mode 100644 backend/internal/service/upstream_models_vertex_test.go diff --git a/backend/internal/service/upstream_models.go b/backend/internal/service/upstream_models.go index 3c5eb2fcd2..c7a392c550 100644 --- a/backend/internal/service/upstream_models.go +++ b/backend/internal/service/upstream_models.go @@ -10,6 +10,7 @@ import ( "net/http" "net/url" "sort" + "strconv" "strings" "time" @@ -25,6 +26,10 @@ const ( UpstreamModelMetadataExtraKey = "upstream_model_metadata" UpstreamModelMetadataIncompleteCode = "upstream_model_metadata_incomplete" UpstreamModelMetadataPartialCode = "upstream_model_metadata_partial" + + vertexPublisherModelsPageSize = 300 + vertexPublisherModelsMaxPages = 20 + vertexGooglePublisherModelPrefix = "publishers/google/models/" ) type UpstreamModelMetadata struct { @@ -741,34 +746,19 @@ func (s *AccountTestService) fetchUpstreamModelList(ctx context.Context, account return nil, nil, newUpstreamModelSyncConfigError("Upstream HTTP client is not configured", nil) } + if account.IsGemini() && account.IsVertexServiceAccount() { + models, err := s.fetchVertexGeminiUpstreamModels(ctx, account) + return models, nil, err + } + req, err := s.buildUpstreamModelsRequest(ctx, account) if err != nil { return nil, nil, err } - proxyURL := upstreamModelsProxyURL(account) - resp, err := s.doUpstreamModelsRequest(req, proxyURL, account) + body, err := s.readUpstreamModelsResponse(req, account) if err != nil { - return nil, nil, newUpstreamModelSyncUpstreamError("Failed to request upstream model list", err) - } - defer func() { _ = resp.Body.Close() }() - - bodyLimit := resolveModelsListReadLimit(s.cfg) - body, err := io.ReadAll(io.LimitReader(resp.Body, bodyLimit+1)) - if err != nil { - return nil, nil, newUpstreamModelSyncUpstreamError("Failed to read upstream model list", err) - } - if int64(len(body)) > bodyLimit { - return nil, nil, newUpstreamModelSyncUpstreamError("Upstream model list response is too large", fmt.Errorf("response exceeds %d bytes", bodyLimit)) - } - - if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices { - return nil, nil, &UpstreamModelSyncError{ - Kind: UpstreamModelSyncErrorUpstream, - Message: fmt.Sprintf("Upstream model list request failed with HTTP %d", resp.StatusCode), - StatusCode: resp.StatusCode, - Err: fmt.Errorf("upstream model list returned HTTP %d", resp.StatusCode), - } + return nil, nil, err } extractModels := extractUpstreamModelIDs @@ -786,6 +776,117 @@ func (s *AccountTestService) fetchUpstreamModelList(ctx context.Context, account return models, body, nil } +// readUpstreamModelsResponse 发送模型列表请求并读取受限大小的响应体；非 2xx 返回带状态码的同步错误。 +func (s *AccountTestService) readUpstreamModelsResponse(req *http.Request, account *Account) ([]byte, error) { + resp, err := s.doUpstreamModelsRequest(req, upstreamModelsProxyURL(account), account) + if err != nil { + return nil, newUpstreamModelSyncUpstreamError("Failed to request upstream model list", err) + } + defer func() { _ = resp.Body.Close() }() + + bodyLimit := resolveModelsListReadLimit(s.cfg) + body, err := io.ReadAll(io.LimitReader(resp.Body, bodyLimit+1)) + if err != nil { + return nil, newUpstreamModelSyncUpstreamError("Failed to read upstream model list", err) + } + if int64(len(body)) > bodyLimit { + return nil, newUpstreamModelSyncUpstreamError("Upstream model list response is too large", fmt.Errorf("response exceeds %d bytes", bodyLimit)) + } + + if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices { + return nil, &UpstreamModelSyncError{ + Kind: UpstreamModelSyncErrorUpstream, + Message: fmt.Sprintf("Upstream model list request failed with HTTP %d", resp.StatusCode), + StatusCode: resp.StatusCode, + Err: fmt.Errorf("upstream model list returned HTTP %d", resp.StatusCode), + } + } + return body, nil +} + +// fetchVertexGeminiUpstreamModels 按账号默认区域分页读取 Vertex 的 Google 发布方模型列表，只保留 gemini- 前缀的模型。 +// 区域端点还会列出 Gemma、Veo、AutoML 等非 Gemini 的模型花园条目；global 端点只含 Gemini。 +// 列表接口不接受 API Key。 +func (s *AccountTestService) fetchVertexGeminiUpstreamModels(ctx context.Context, account *Account) ([]string, error) { + if account.IsVertexAPIKey() { + return nil, newUpstreamModelSyncUnsupportedError("Vertex API key accounts do not support upstream model listing", nil) + } + if _, err := parseVertexServiceAccountKey(account); err != nil { + return nil, newUpstreamModelSyncConfigError("Invalid Vertex service account credentials", err) + } + location := account.VertexLocation("") + if _, err := buildVertexPublisherModelsURL(location, ""); err != nil { + return nil, newUpstreamModelSyncConfigError("Invalid Vertex location", err) + } + if s.geminiTokenProvider == nil { + return nil, newUpstreamModelSyncConfigError("Gemini token provider is not configured", nil) + } + accessToken, err := s.geminiTokenProvider.GetAccessToken(ctx, account) + if err != nil { + return nil, newUpstreamModelSyncUpstreamError("Failed to get Vertex access token", err) + } + accessToken = strings.TrimSpace(accessToken) + if accessToken == "" { + return nil, newUpstreamModelSyncConfigError("No Vertex access token is available", nil) + } + + models := make([]string, 0) + pageToken := "" + for page := 0; page < vertexPublisherModelsMaxPages; page++ { + listURL, err := buildVertexPublisherModelsURL(location, pageToken) + if err != nil { + return nil, newUpstreamModelSyncConfigError("Invalid Vertex location", err) + } + req, err := http.NewRequestWithContext(ctx, http.MethodGet, listURL, nil) + if err != nil { + return nil, newUpstreamModelSyncConfigError("Invalid Vertex model list URL", err) + } + req.Header.Set("Accept", "application/json") + req.Header.Set("Authorization", "Bearer "+accessToken) + + body, err := s.readUpstreamModelsResponse(req, account) + if err != nil { + // 只有首页的 404/405 表示列表端点不可用；翻页途中的失败不触发配置模型回退。 + var syncErr *UpstreamModelSyncError + if page > 0 && errors.As(err, &syncErr) { + syncErr.StatusCode = 0 + } + return nil, err + } + var response struct { + PublisherModels []struct { + Name string `json:"name"` + } `json:"publisherModels"` + NextPageToken string `json:"nextPageToken"` + } + if err := json.Unmarshal(body, &response); err != nil { + return nil, newUpstreamModelSyncUpstreamError("Upstream model list response was not valid JSON", err) + } + for _, entry := range response.PublisherModels { + name := strings.TrimSpace(entry.Name) + if !strings.HasPrefix(name, vertexGooglePublisherModelPrefix) { + continue + } + modelID := strings.TrimPrefix(name, vertexGooglePublisherModelPrefix) + if strings.HasPrefix(modelID, "gemini-") && len(modelID) > len("gemini-") { + models = append(models, modelID) + } + } + + pageToken = strings.TrimSpace(response.NextPageToken) + if pageToken == "" { + if len(models) == 0 { + return nil, newUpstreamModelSyncUpstreamError("Upstream returned no supported models", nil) + } + return dedupeAndSortModelIDs(models), nil + } + } + return nil, newUpstreamModelSyncUpstreamError( + "Upstream model list pagination did not finish", + fmt.Errorf("more than %d pages", vertexPublisherModelsMaxPages), + ) +} + func (s *AccountTestService) buildUpstreamModelsRequest(ctx context.Context, account *Account) (*http.Request, error) { switch { case account.Platform == PlatformAntigravity: @@ -1218,6 +1319,26 @@ func buildOpenAIModelsURL(base string) string { return buildOpenAIEndpointURL(base, "/v1/models") } +// buildVertexPublisherModelsURL 构造 Google 发布方模型列表地址。该接口只有 v1beta1 版本，pageSize 上限 300。 +func buildVertexPublisherModelsURL(location, pageToken string) (string, error) { + location = strings.TrimSpace(location) + if location == "" { + location = vertexDefaultLocation + } + if !vertexLocationPattern.MatchString(location) { + return "", fmt.Errorf("invalid vertex location: %s", location) + } + host := location + "-aiplatform.googleapis.com" + if location == "global" { + host = vertexGlobalHost + } + query := url.Values{"pageSize": []string{strconv.Itoa(vertexPublisherModelsPageSize)}} + if pageToken != "" { + query.Set("pageToken", pageToken) + } + return fmt.Sprintf("https://%s/v1beta1/publishers/google/models?%s", host, query.Encode()), nil +} + func buildGeminiModelsURL(base string) string { normalized := strings.TrimRight(strings.TrimSpace(base), "/") if strings.HasSuffix(normalized, "/v1beta/models") { diff --git a/backend/internal/service/upstream_models_vertex_test.go b/backend/internal/service/upstream_models_vertex_test.go new file mode 100644 index 0000000000..bd4678614e --- /dev/null +++ b/backend/internal/service/upstream_models_vertex_test.go @@ -0,0 +1,276 @@ +package service + +import ( + "context" + "errors" + "io" + "net/http" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +type vertexModelsTokenCacheStub struct { + token string + calls int +} + +func (c *vertexModelsTokenCacheStub) GetAccessToken(context.Context, string) (string, error) { + c.calls++ + if c.token == "" { + return "", errors.New("missing token") + } + return c.token, nil +} + +func (c *vertexModelsTokenCacheStub) SetAccessToken(context.Context, string, string, time.Duration) error { + return nil +} + +func (c *vertexModelsTokenCacheStub) DeleteAccessToken(context.Context, string) error { + return nil +} + +func (c *vertexModelsTokenCacheStub) AcquireRefreshLock(context.Context, string, time.Duration) (bool, error) { + return false, nil +} + +func (c *vertexModelsTokenCacheStub) ReleaseRefreshLock(context.Context, string) error { + return nil +} + +func vertexModelSyncTestAccount(credentials map[string]any) *Account { + merged := map[string]any{ + "service_account_json": map[string]any{ + "type": "service_account", + "project_id": "proj", + "private_key": "-----BEGIN PRIVATE KEY-----\nabc\n-----END PRIVATE KEY-----\n", + "client_email": "svc@proj.iam.gserviceaccount.com", + }, + } + for key, value := range credentials { + merged[key] = value + } + return &Account{ + ID: 173, + Platform: PlatformGemini, + Type: AccountTypeServiceAccount, + Credentials: merged, + } +} + +func vertexModelSyncTestService(upstream *httpUpstreamRecorder) *AccountTestService { + return &AccountTestService{ + httpUpstream: upstream, + cfg: upstreamModelSyncTestConfig(), + geminiTokenProvider: NewGeminiTokenProvider(nil, &vertexModelsTokenCacheStub{token: "ya29.vertex-token"}, nil), + } +} + +func vertexModelsJSONResponse(status int, body string) *http.Response { + return &http.Response{ + StatusCode: status, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(body)), + } +} + +func TestBuildVertexPublisherModelsURL(t *testing.T) { + t.Parallel() + + u, err := buildVertexPublisherModelsURL("global", "") + require.NoError(t, err) + require.Equal(t, "https://aiplatform.googleapis.com/v1beta1/publishers/google/models?pageSize=300", u) + + u, err = buildVertexPublisherModelsURL("us-central1", "next/page+token") + require.NoError(t, err) + require.Equal(t, "https://us-central1-aiplatform.googleapis.com/v1beta1/publishers/google/models?pageSize=300&pageToken=next%2Fpage%2Btoken", u) + + u, err = buildVertexPublisherModelsURL("", "") + require.NoError(t, err) + require.Equal(t, "https://us-central1-aiplatform.googleapis.com/v1beta1/publishers/google/models?pageSize=300", u) + + _, err = buildVertexPublisherModelsURL("evil.com/x", "") + require.Error(t, err) +} + +// Scenario: Vertex SA 账号按账号区域分页读取发布方模型列表，只保留 gemini- 前缀模型。 +func TestFetchUpstreamSupportedModelsVertexServiceAccountPaginatesAndKeepsGemini(t *testing.T) { + t.Parallel() + + upstream := &httpUpstreamRecorder{responses: []*http.Response{ + vertexModelsJSONResponse(http.StatusOK, `{"publisherModels":[ + {"name":"publishers/google/models/gemini-3.8-flash"}, + {"name":"publishers/google/models/gemma4"}, + {"name":"publishers/google/models/veo-3.1-generate-001"} + ],"nextPageToken":"page-2"}`), + vertexModelsJSONResponse(http.StatusOK, `{"publisherModels":[ + {"name":" publishers/google/models/gemini-2.5-flash "}, + {"name":"publishers/google/models/gemini-3.8-flash"}, + {"name":"gemini-without-publisher-prefix"}, + {"name":"publishers/anthropic/models/gemini-other-publisher"}, + {"name":"publishers/google/models/gemini-"} + ]}`), + }} + svc := vertexModelSyncTestService(upstream) + + models, err := svc.FetchUpstreamSupportedModels(context.Background(), vertexModelSyncTestAccount(map[string]any{ + "location": "europe-west4", + })) + require.NoError(t, err) + require.Equal(t, []string{"gemini-2.5-flash", "gemini-3.8-flash"}, models) + + require.Len(t, upstream.requests, 2) + require.Equal(t, "https://europe-west4-aiplatform.googleapis.com/v1beta1/publishers/google/models?pageSize=300", upstream.requests[0].URL.String()) + require.Equal(t, "https://europe-west4-aiplatform.googleapis.com/v1beta1/publishers/google/models?pageSize=300&pageToken=page-2", upstream.requests[1].URL.String()) + for _, req := range upstream.requests { + require.Equal(t, http.MethodGet, req.Method) + require.Equal(t, "Bearer ya29.vertex-token", req.Header.Get("Authorization")) + } +} + +func TestFetchUpstreamSupportedModelsVertexUpstreamErrorKeepsStatus(t *testing.T) { + t.Parallel() + + upstream := &httpUpstreamRecorder{resp: vertexModelsJSONResponse(http.StatusForbidden, `{"error":{"code":403,"status":"PERMISSION_DENIED"}}`)} + svc := vertexModelSyncTestService(upstream) + + _, err := svc.FetchUpstreamSupportedModels(context.Background(), vertexModelSyncTestAccount(map[string]any{"location": "global"})) + var syncErr *UpstreamModelSyncError + require.ErrorAs(t, err, &syncErr) + require.Equal(t, UpstreamModelSyncErrorUpstream, syncErr.Kind) + require.Equal(t, http.StatusForbidden, syncErr.StatusCode) + require.Equal(t, "https://aiplatform.googleapis.com/v1beta1/publishers/google/models?pageSize=300", upstream.lastReq.URL.String()) +} + +func TestFetchUpstreamSupportedModelsVertexWithoutGeminiModelsFails(t *testing.T) { + t.Parallel() + + upstream := &httpUpstreamRecorder{resp: vertexModelsJSONResponse(http.StatusOK, `{"publisherModels":[{"name":"publishers/google/models/gemma4"}]}`)} + svc := vertexModelSyncTestService(upstream) + + _, err := svc.FetchUpstreamSupportedModels(context.Background(), vertexModelSyncTestAccount(nil)) + var syncErr *UpstreamModelSyncError + require.ErrorAs(t, err, &syncErr) + require.Equal(t, UpstreamModelSyncErrorUpstream, syncErr.Kind) + require.Equal(t, "Upstream returned no supported models", syncErr.SafeMessage()) +} + +func TestFetchUpstreamSupportedModelsVertexStopsEndlessPagination(t *testing.T) { + t.Parallel() + + responses := make([]*http.Response, 0, vertexPublisherModelsMaxPages) + for i := 0; i < vertexPublisherModelsMaxPages; i++ { + responses = append(responses, vertexModelsJSONResponse(http.StatusOK, `{"publisherModels":[{"name":"publishers/google/models/gemini-3.8-flash"}],"nextPageToken":"again"}`)) + } + upstream := &httpUpstreamRecorder{responses: responses} + svc := vertexModelSyncTestService(upstream) + + _, err := svc.FetchUpstreamSupportedModels(context.Background(), vertexModelSyncTestAccount(nil)) + var syncErr *UpstreamModelSyncError + require.ErrorAs(t, err, &syncErr) + require.Equal(t, "Upstream model list pagination did not finish", syncErr.SafeMessage()) + require.Len(t, upstream.requests, vertexPublisherModelsMaxPages) +} + +func TestFetchUpstreamSupportedModelsVertexAPIKeyUnsupported(t *testing.T) { + t.Parallel() + + upstream := &httpUpstreamRecorder{} + svc := vertexModelSyncTestService(upstream) + account := &Account{ + ID: 174, + Platform: PlatformGemini, + Type: AccountTypeServiceAccount, + Credentials: map[string]any{ + "auth_mode": "apikey", + "api_key": "AQ.test-key", + }, + } + + _, err := svc.FetchUpstreamSupportedModels(context.Background(), account) + var syncErr *UpstreamModelSyncError + require.ErrorAs(t, err, &syncErr) + require.Equal(t, UpstreamModelSyncErrorUnsupported, syncErr.Kind) + require.Empty(t, upstream.requests) +} + +func TestFetchUpstreamSupportedModelsVertexInvalidCredentialsIsConfigError(t *testing.T) { + t.Parallel() + + cache := &vertexModelsTokenCacheStub{token: "ya29.vertex-token"} + upstream := &httpUpstreamRecorder{} + svc := &AccountTestService{ + httpUpstream: upstream, + cfg: upstreamModelSyncTestConfig(), + geminiTokenProvider: NewGeminiTokenProvider(nil, cache, nil), + } + account := &Account{ + ID: 175, + Platform: PlatformGemini, + Type: AccountTypeServiceAccount, + Credentials: map[string]any{"api_key": "not-a-service-account"}, + } + + _, err := svc.FetchUpstreamSupportedModels(context.Background(), account) + var syncErr *UpstreamModelSyncError + require.ErrorAs(t, err, &syncErr) + require.Equal(t, UpstreamModelSyncErrorConfiguration, syncErr.Kind) + require.Equal(t, "Invalid Vertex service account credentials", syncErr.SafeMessage()) + require.Zero(t, cache.calls) + require.Empty(t, upstream.requests) +} + +func TestFetchUpstreamSupportedModelsVertexInvalidLocationSkipsTokenExchange(t *testing.T) { + t.Parallel() + + cache := &vertexModelsTokenCacheStub{token: "ya29.vertex-token"} + upstream := &httpUpstreamRecorder{} + svc := &AccountTestService{ + httpUpstream: upstream, + cfg: upstreamModelSyncTestConfig(), + geminiTokenProvider: NewGeminiTokenProvider(nil, cache, nil), + } + + _, err := svc.FetchUpstreamSupportedModels(context.Background(), vertexModelSyncTestAccount(map[string]any{ + "location": "us-central1.evil.com", + })) + var syncErr *UpstreamModelSyncError + require.ErrorAs(t, err, &syncErr) + require.Equal(t, UpstreamModelSyncErrorConfiguration, syncErr.Kind) + require.Equal(t, "Invalid Vertex location", syncErr.SafeMessage()) + require.Zero(t, cache.calls) + require.Empty(t, upstream.requests) +} + +// Scenario: 翻页途中的 404 不能被当成列表端点不可用而回退到配置模型。 +func TestSyncUpstreamModelCatalogVertexMidPagination404DoesNotFallback(t *testing.T) { + t.Parallel() + + upstream := &httpUpstreamRecorder{responses: []*http.Response{ + vertexModelsJSONResponse(http.StatusOK, `{"publisherModels":[{"name":"publishers/google/models/gemini-a"}],"nextPageToken":"page-2"}`), + vertexModelsJSONResponse(http.StatusNotFound, `{"error":{"code":404}}`), + }} + svc := vertexModelSyncTestService(upstream) + + _, err := svc.SyncUpstreamModelCatalog(context.Background(), vertexModelSyncTestAccount(map[string]any{ + "model_mapping": map[string]any{"x": "gemini-mapped"}, + })) + var syncErr *UpstreamModelSyncError + require.ErrorAs(t, err, &syncErr) + require.Equal(t, UpstreamModelSyncErrorUpstream, syncErr.Kind) + require.Zero(t, syncErr.StatusCode) + require.Len(t, upstream.requests, 2) +} + +func TestFetchUpstreamSupportedModelsVertexFirstPage404KeepsStatus(t *testing.T) { + t.Parallel() + + upstream := &httpUpstreamRecorder{resp: vertexModelsJSONResponse(http.StatusNotFound, `{"error":{"code":404}}`)} + svc := vertexModelSyncTestService(upstream) + + _, err := svc.FetchUpstreamSupportedModels(context.Background(), vertexModelSyncTestAccount(nil)) + require.True(t, upstreamModelListEndpointUnsupported(err)) +} From c9e4b267b6fac1f00060eab983e606f7a2b2d4fc Mon Sep 17 00:00:00 2001 From: feeeei Date: Sat, 10 Oct 2026 19:12:10 +0800 Subject: [PATCH 09/15] =?UTF-8?q?fix(account):=20=E5=90=8C=E6=AD=A5?= =?UTF-8?q?=E4=B8=8A=E6=B8=B8=E6=A8=A1=E5=9E=8B=E5=A4=B1=E8=B4=A5=E6=97=B6?= =?UTF-8?q?=E6=98=BE=E7=A4=BA=E5=90=8E=E7=AB=AF=E8=BF=94=E5=9B=9E=E7=9A=84?= =?UTF-8?q?=E9=94=99=E8=AF=AF=E5=8E=9F=E5=9B=A0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - apiClient 拒绝的是普通对象而非 Error，原写法 instanceof Error 恒为假， 错误提示被替换成通用文案"同步上游模型失败：同步上游模型失败" - 白名单选择器与 Antigravity 同步按钮改用 extractApiErrorMessage 取消息 --- .../components/account/EditAccountModal.vue | 2 +- .../account/ModelWhitelistSelector.vue | 3 ++- .../__tests__/ModelWhitelistSelector.spec.ts | 22 ++++++++++++++++++- 3 files changed, 24 insertions(+), 3 deletions(-) diff --git a/frontend/src/components/account/EditAccountModal.vue b/frontend/src/components/account/EditAccountModal.vue index 9bc4f9167a..04b1c9f68f 100644 --- a/frontend/src/components/account/EditAccountModal.vue +++ b/frontend/src/components/account/EditAccountModal.vue @@ -4771,7 +4771,7 @@ const syncAntigravityUpstreamModels = async () => { appStore.showWarning(t('admin.accounts.syncUpstreamModelsMetadataPartial')) } } catch (error) { - const message = error instanceof Error ? error.message : t('admin.accounts.syncUpstreamModelsFailed') + const message = extractApiErrorMessage(error, t('admin.accounts.syncUpstreamModelsFailed')) appStore.showError(t('admin.accounts.syncUpstreamModelsError', { message })) } finally { isSyncingAntigravityUpstream.value = false diff --git a/frontend/src/components/account/ModelWhitelistSelector.vue b/frontend/src/components/account/ModelWhitelistSelector.vue index d953c21ce9..aa09de7bbb 100644 --- a/frontend/src/components/account/ModelWhitelistSelector.vue +++ b/frontend/src/components/account/ModelWhitelistSelector.vue @@ -155,6 +155,7 @@ import ModelIcon from '@/components/common/ModelIcon.vue' import Icon from '@/components/icons/Icon.vue' import { allModels, getModelsByPlatform } from '@/composables/useModelWhitelist' import { supportsUpstreamModelSync } from '@/constants/platformCatalog' +import { extractApiErrorMessage } from '@/utils/apiError' const { t } = useI18n() @@ -343,7 +344,7 @@ const syncUpstreamModels = async () => { appStore.showWarning(t('admin.accounts.syncUpstreamModelsMetadataPartial')) } } catch (error) { - const message = error instanceof Error ? error.message : t('admin.accounts.syncUpstreamModelsFailed') + const message = extractApiErrorMessage(error, t('admin.accounts.syncUpstreamModelsFailed')) appStore.showError(t('admin.accounts.syncUpstreamModelsError', { message })) } finally { isSyncingUpstream.value = false diff --git a/frontend/src/components/account/__tests__/ModelWhitelistSelector.spec.ts b/frontend/src/components/account/__tests__/ModelWhitelistSelector.spec.ts index 9e27d29db9..0b5f10a015 100644 --- a/frontend/src/components/account/__tests__/ModelWhitelistSelector.spec.ts +++ b/frontend/src/components/account/__tests__/ModelWhitelistSelector.spec.ts @@ -29,7 +29,7 @@ vi.mock('vue-i18n', async () => { return { ...actual, useI18n: () => ({ - t: (key: string, params?: Record) => key === 'common.copy' ? '复制' : key === 'admin.accounts.modelMappingConflict' ? `Model mapping conflict: ${params?.from} → ${params?.to}` : key + t: (key: string, params?: Record) => key === 'common.copy' ? '复制' : key === 'admin.accounts.modelMappingConflict' ? `Model mapping conflict: ${params?.from} → ${params?.to}` : key === 'admin.accounts.syncUpstreamModelsError' ? `${key}: ${params?.message}` : key }) } }) @@ -244,6 +244,26 @@ describe('ModelWhitelistSelector', () => { expect(showSuccess).not.toHaveBeenCalled() }) + it('shows the backend message when upstream sync is rejected', async () => { + syncUpstreamModels.mockRejectedValue({ + status: 400, + code: 'BAD_REQUEST', + message: 'Vertex API key accounts do not support upstream model listing' + }) + const wrapper = mountSelector({ platform: 'gemini', accountId: 174 }) + + const syncButton = wrapper + .findAll('button') + .find(button => button.text() === 'admin.accounts.syncUpstreamModels') + expect(syncButton).toBeDefined() + await syncButton!.trigger('click') + await flushPromises() + + expect(showError).toHaveBeenCalledWith( + 'admin.accounts.syncUpstreamModelsError: Vertex API key accounts do not support upstream model listing' + ) + }) + it('shows success and a partial warning when some capabilities were saved', async () => { syncUpstreamModels.mockResolvedValue({ models: ['gpt-6-astra', 'gpt-image-2'], From 05d7946af84228589e5a6a8d8b6bc08d686ba778 Mon Sep 17 00:00:00 2001 From: feeeei Date: Sun, 11 Oct 2026 01:37:04 +0800 Subject: [PATCH 10/15] =?UTF-8?q?fix(gemini):=20=E6=98=BE=E5=BC=8F?= =?UTF-8?q?=E7=BC=93=E5=AD=98=E9=94=99=E8=AF=AF=E9=80=8F=E4=BC=A0=E6=94=B6?= =?UTF-8?q?=E7=AA=84=E5=88=B0=E4=B8=8A=E6=B8=B8=E7=BD=91=E5=85=B3=E8=B4=A6?= =?UTF-8?q?=E5=8F=B7=EF=BC=8C=E5=BB=B6=E9=95=BF=E6=B2=BF=E7=94=A8=E5=88=9B?= =?UTF-8?q?=E5=BB=BA=E6=97=B6=E8=AE=A1=E8=B4=B9=E5=8F=A3=E5=BE=84?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 绑定缓存的请求与缓存管理接口收到 429 / 5xx 时，只有上游网关账号的响应原样回写 （UpstreamFailoverError.BoundUpstreamPassthrough / GeminiCachedContentUpstreamError.UpstreamGateway）， 官方账号仍走错误透传规则与通用映射，不向客户端暴露配额明细 - 记录新增 channel_usage（迁移 244）保存创建时的渠道计费口径，PATCH 延长按同一口径预检、报价与记账 - 创建 / 延长费用行的 request_id 改为固定前缀 gcache:{create|patch}:{id}，不被客户端或网关请求 id 覆盖， 按缓存 ID 即可检索一个缓存的全部费用行 - 调度快照凭据投影补齐 base_url / auth_mode / 上游显式缓存声明，显式缓存候选过滤一次算全， 不再逐个选号后剔除 --- .../gemini_cached_contents_billing_test.go | 62 +++++++++++++++++-- .../handler/gemini_cached_contents_handler.go | 45 ++++++++------ .../gemini_cached_contents_handler_test.go | 34 +++++++++- .../repository/gemini_cached_content_repo.go | 20 ++++-- ...ni_cached_content_repo_integration_test.go | 2 + .../internal/repository/scheduler_cache.go | 6 +- .../repository/scheduler_cache_test.go | 17 +++++ backend/internal/service/gateway_service.go | 1 + .../internal/service/gateway_usage_billing.go | 3 +- .../gateway_usage_billing_request_id_test.go | 18 ++++++ .../internal/service/gemini_cached_content.go | 46 +++++++++++--- .../gemini_cached_content_hardening_test.go | 1 + .../service/gemini_cached_content_test.go | 2 +- .../service/gemini_messages_compat_service.go | 2 +- ...4_gemini_cached_contents_channel_usage.sql | 4 ++ 15 files changed, 218 insertions(+), 45 deletions(-) create mode 100644 backend/migrations/244_gemini_cached_contents_channel_usage.sql diff --git a/backend/internal/handler/gemini_cached_contents_billing_test.go b/backend/internal/handler/gemini_cached_contents_billing_test.go index f567c2cb01..a1fd414519 100644 --- a/backend/internal/handler/gemini_cached_contents_billing_test.go +++ b/backend/internal/handler/gemini_cached_contents_billing_test.go @@ -25,9 +25,10 @@ import ( ) const ( - cachedBillingInputPrice = 0.00000125 - cachedBillingStoragePrice = 0.0000045 - cachedBillingTokens = 1_000_000 + cachedBillingInputPrice = 0.00000125 + cachedBillingStoragePrice = 0.0000045 + cachedBillingAliasStoragePrice = 0.000009 + cachedBillingTokens = 1_000_000 ) type cachedBillingRepo struct { @@ -175,6 +176,12 @@ func (r *cachedBillingUsageRepo) Create(_ context.Context, log *service.UsageLog return true, nil } +func (r *cachedBillingUsageRepo) first() *service.UsageLog { + r.mu.Lock() + defer r.mu.Unlock() + return r.logs[0] +} + func (r *cachedBillingUsageRepo) totals() []float64 { r.mu.Lock() defer r.mu.Unlock() @@ -192,6 +199,8 @@ type cachedBillingEnv struct { usage *cachedBillingUsageRepo apiKey *service.APIKey group *service.Group + // accounts 是调度快照桩直接返回的账号对象，测试可改其凭据。 + accounts map[int64]*service.Account } func newCachedBillingEnv(t *testing.T, balance float64, accountIDs ...int64) *cachedBillingEnv { @@ -199,7 +208,10 @@ func newCachedBillingEnv(t *testing.T, balance float64, accountIDs ...int64) *ca dir := t.TempDir() require.NoError(t, os.WriteFile(filepath.Join(dir, "model_pricing.json"), []byte(fmt.Sprintf(`{"gemini-2.5-pro":{ "litellm_provider":"gemini","mode":"chat","input_cost_per_token":%g,"output_cost_per_token":0.00001, - "cache_storage_cost_per_token_per_hour":%g}}`, cachedBillingInputPrice, cachedBillingStoragePrice)), 0o644)) + "cache_storage_cost_per_token_per_hour":%g},"gemini-alias":{ + "litellm_provider":"gemini","mode":"chat","input_cost_per_token":%g,"output_cost_per_token":0.00001, + "cache_storage_cost_per_token_per_hour":%g}}`, + cachedBillingInputPrice, cachedBillingStoragePrice, cachedBillingInputPrice, cachedBillingAliasStoragePrice)), 0o644)) cfg := &config.Config{} cfg.Pricing.DataDir = dir cfg.Pricing.UpdateIntervalHours = 100000 @@ -237,7 +249,25 @@ func newCachedBillingEnv(t *testing.T, balance float64, accountIDs ...int64) *ca } apiKey := &service.APIKey{ID: 7, UserID: 3, GroupID: &groupID, Group: group, Status: service.StatusActive, User: &service.User{ID: 3, Concurrency: 10, Balance: balance}} - return &cachedBillingEnv{h: h, repo: repo, upstream: upstream, usage: usage, apiKey: apiKey, group: group} + return &cachedBillingEnv{h: h, repo: repo, upstream: upstream, usage: usage, apiKey: apiKey, group: group, accounts: accounts} +} + +// TestGeminiExplicitCacheExclusions_FiltersUnsupportedAccountsFromSnapshot 候选过滤直接按快照里的凭据判定： +// 第三方地址且未声明上游显式缓存的 Key 账号在选号前就被排除，官方地址与声明了上游显式缓存的账号保留。 +func TestGeminiExplicitCacheExclusions_FiltersUnsupportedAccountsFromSnapshot(t *testing.T) { + env := newCachedBillingEnv(t, 500, 900, 901, 902) + env.accounts[901].Credentials["base_url"] = "https://third.example.com" + env.accounts[902].Credentials["base_url"] = "https://gw.example.com" + env.accounts[902].Credentials[service.GeminiExplicitCacheUpstreamCredentialKey] = true + + groupID := env.group.ID + excluded, err := env.h.gatewayService.GeminiExplicitCacheExclusions(context.Background(), &groupID, 0) + require.NoError(t, err) + require.Equal(t, map[int64]struct{}{901: {}}, excluded) + + excluded, err = env.h.gatewayService.GeminiExplicitCacheExclusions(context.Background(), &groupID, 902) + require.NoError(t, err) + require.Equal(t, map[int64]struct{}{900: {}, 901: {}}, excluded, "绑定缓存的请求只保留持有账号") } func (e *cachedBillingEnv) serve(method, target, body string, params gin.Params, call func(*gin.Context)) *httptest.ResponseRecorder { @@ -279,6 +309,27 @@ func (e *cachedBillingEnv) patch() *httptest.ResponseRecorder { gin.Params{{Key: "cacheID", Value: "abc"}}, e.h.GeminiCachedContentsPatch) } +// TestGeminiCachedContentsPatch_UsesChannelUsageRecordedAtCreate 延长有效期沿用创建时的渠道计费口径： +// 计费模型来源为请求模型时按别名定价，用量记录带渠道 ID。 +func TestGeminiCachedContentsPatch_UsesChannelUsageRecordedAtCreate(t *testing.T) { + env := newCachedBillingEnv(t, 500, 900) + env.seed(time.Now().Add(time.Hour)) + env.repo.records["abc"].ChannelUsage = service.ChannelUsageFields{ + ChannelID: 5, OriginalModel: "gemini-alias", ChannelMappedModel: "gemini-2.5-pro", + BillingModelSource: service.BillingModelSourceRequested, + } + + rec := env.patch() + require.Equal(t, http.StatusOK, rec.Code, rec.Body.String()) + totals := env.waitForUsage(t, 1) + require.Len(t, totals, 1) + require.InDelta(t, cachedBillingTokens*cachedBillingAliasStoragePrice*23, totals[0], 0.01, "按创建时的计费模型（请求模型别名）计存储费") + usageLog := env.usage.first() + require.NotNil(t, usageLog.ChannelID) + require.Equal(t, int64(5), *usageLog.ChannelID) + require.True(t, strings.HasPrefix(usageLog.RequestID, "gcache:patch:abc:"), usageLog.RequestID) +} + func cachedBillingCreateBody(expire time.Time) string { return fmt.Sprintf(`{"name":"cachedContents/up1","usageMetadata":{"totalTokenCount":%d},"expireTime":%q}`, cachedBillingTokens, expire.UTC().Format(time.RFC3339Nano)) @@ -296,6 +347,7 @@ func TestGeminiCachedContentsCreate_ChargesInputAndStorage(t *testing.T) { require.Len(t, env.repo.records, 1) for _, record := range env.repo.records { require.False(t, record.ExpireTime.After(time.Now().Add(24*time.Hour)), "本地到期时间不超过请求的有效期") + require.Equal(t, "gcache:create:"+record.PublicID, env.usage.first().RequestID, "费用行按缓存 ID 可检索") } } diff --git a/backend/internal/handler/gemini_cached_contents_handler.go b/backend/internal/handler/gemini_cached_contents_handler.go index 00404329d5..ac99eaf1d5 100644 --- a/backend/internal/handler/gemini_cached_contents_handler.go +++ b/backend/internal/handler/gemini_cached_contents_handler.go @@ -17,7 +17,6 @@ import ( "github.com/Wei-Shaw/sub2api/internal/server/middleware" "github.com/Wei-Shaw/sub2api/internal/service" "github.com/gin-gonic/gin" - "github.com/google/uuid" "go.uber.org/zap" ) @@ -246,7 +245,7 @@ func (h *GatewayHandler) GeminiCachedContentsCreate(c *gin.Context) { excluded[account.ID] = struct{}{} continue } - writeGeminiCachedContentUpstreamError(c, upErr) + h.writeGeminiCachedContentUpstreamError(c, upErr) return } // 上游已返回成功但响应不可用：不换号重建，能识别出资源名时删除上游缓存。 @@ -284,7 +283,7 @@ func (h *GatewayHandler) GeminiCachedContentsCreate(c *gin.Context) { reqLog.Info("gemini.cached_content.storage_charge_rejected", zap.Float64("amount", amount), zap.Error(err)) h.cleanupGeminiCachedContentUpstream(reqLog, account, upstream.Name) h.submitGeminiCachedContentUsage(c, apiKey, account, subscription, pricingAt, - "create:"+publicID, modelName, mappedModel, int(tokens), 0, usageFields, body) + service.GeminiCachedContentCreateUsageRequestID(publicID), modelName, mappedModel, int(tokens), 0, usageFields, body) writeGeminiCachedContentBillingError(c, err) return } @@ -303,6 +302,7 @@ func (h *GatewayHandler) GeminiCachedContentsCreate(c *gin.Context) { DisplayName: req.DisplayName, TotalTokenCount: tokens, ExpireTime: expireTime, + ChannelUsage: usageFields, } persistCtx, cancelPersist := geminiCachedContentPersistContext(c) err = svc.Save(persistCtx, record) @@ -315,13 +315,13 @@ func (h *GatewayHandler) GeminiCachedContentsCreate(c *gin.Context) { } h.submitGeminiCachedContentUsage(c, apiKey, account, subscription, pricingAt, - "create:"+publicID, modelName, mappedModel, int(tokens), storageTokenHours, usageFields, body) + service.GeminiCachedContentCreateUsageRequestID(publicID), modelName, mappedModel, int(tokens), storageTokenHours, usageFields, body) c.JSON(http.StatusOK, service.GeminiCachedContentView(record)) return } if lastUpstreamErr != nil { - writeGeminiCachedContentUpstreamError(c, lastUpstreamErr) + h.writeGeminiCachedContentUpstreamError(c, lastUpstreamErr) return } if ttlRejectedLimit > 0 { @@ -436,6 +436,7 @@ func (h *GatewayHandler) GeminiCachedContentsPatch(c *gin.Context) { CacheStorageTokenHours: service.GeminiCachedContentStorageTokenHours(record.TotalTokenCount, extension), }, APIKey: apiKey, User: apiKey.User, Subscription: subscription, PricingAt: pricingAt, + ChannelUsageFields: record.ChannelUsage, } if !h.gatewayService.CacheStoragePriced(c.Request.Context(), usage) { googlePlatformError(c, http.StatusBadRequest, fmt.Sprintf("Explicit context caching is not available for model %s: storage pricing is not configured", record.Model)) @@ -463,7 +464,7 @@ func (h *GatewayHandler) GeminiCachedContentsPatch(c *gin.Context) { reqLog.Warn("gemini.cached_content.forget_failed", zap.Error(forgetErr)) } } - writeGeminiCachedContentUpstreamError(c, upErr) + h.writeGeminiCachedContentUpstreamError(c, upErr) return } reqLog.Warn("gemini.cached_content.upstream_patch_error", zap.Int64("account_id", account.ID), zap.Error(err)) @@ -486,8 +487,8 @@ func (h *GatewayHandler) GeminiCachedContentsPatch(c *gin.Context) { // 按库中实际推进的时长收费：并发延长各自只为自己推进的部分付费。 if extended := expireTime.Sub(previousExpire); extended > 0 { h.submitGeminiCachedContentUsage(c, apiKey, account, subscription, pricingAt, - "patch:"+record.PublicID+":"+uuid.NewString(), record.Model, record.UpstreamModelID(), 0, - service.GeminiCachedContentStorageTokenHours(record.TotalTokenCount, extended), service.ChannelUsageFields{}, body) + service.GeminiCachedContentPatchUsageRequestID(record.PublicID), record.Model, record.UpstreamModelID(), 0, + service.GeminiCachedContentStorageTokenHours(record.TotalTokenCount, extended), record.ChannelUsage, body) } c.JSON(http.StatusOK, service.GeminiCachedContentView(record)) } @@ -520,7 +521,7 @@ func (h *GatewayHandler) GeminiCachedContentsDelete(c *gin.Context) { return } if !upErr.NotFound() { - writeGeminiCachedContentUpstreamError(c, upErr) + h.writeGeminiCachedContentUpstreamError(c, upErr) return } } @@ -628,14 +629,15 @@ func (h *GatewayHandler) acquireGeminiCachedContentAccountSlot(c *gin.Context, r return wrapReleaseOnDone(c.Request.Context(), release), true } -// submitGeminiCachedContentUsage 记录缓存创建 / 延长的计费：inputTokens 按输入价计，storageTokenHours 按计费模型的存储单价计。 +// submitGeminiCachedContentUsage 记录缓存创建 / 延长的计费：inputTokens 按输入价计，storageTokenHours 按计费模型的存储单价计； +// requestID 是带固定前缀的费用行标识（见 GeminiCachedContentUsageRequestIDPrefix），落库时不被请求 id 覆盖。 func (h *GatewayHandler) submitGeminiCachedContentUsage( c *gin.Context, apiKey *service.APIKey, account *service.Account, subscription *service.UserSubscription, pricingAt time.Time, - requestSuffix string, + requestID string, model string, upstreamModel string, inputTokens int, @@ -650,7 +652,7 @@ func (h *GatewayHandler) submitGeminiCachedContentUsage( upstreamModel = "" } result := &service.ForwardResult{ - RequestID: "gemini-cached-content:" + requestSuffix, + RequestID: requestID, Model: model, UpstreamModel: upstreamModel, Usage: service.ClaudeUsage{InputTokens: inputTokens}, @@ -729,9 +731,10 @@ func writeGeminiCachedContentBillingError(c *gin.Context, err error) { googleErrorWithType(c, status, code, "", message) } -// writeGeminiCachedContentUpstreamError 透传上游错误：缓存不存在统一返回标准 403（各上游形态不同且可能含上游资源 ID）； -// 账号凭据或权限类错误不是客户端问题，统一改为 502。 -func writeGeminiCachedContentUpstreamError(c *gin.Context, upErr *service.GeminiCachedContentUpstreamError) { +// writeGeminiCachedContentUpstreamError 回写上游错误：缓存不存在统一返回标准 403（各上游形态不同且可能含上游资源 ID）； +// 账号凭据或权限类错误不是客户端问题，统一改为 502；限流 / 不可用按普通 Gemini 路径处理（透传规则、通用映射）， +// 只有上游网关账号的这类响应反映持有缓存的上游账号状态，原样回写。 +func (h *GatewayHandler) writeGeminiCachedContentUpstreamError(c *gin.Context, upErr *service.GeminiCachedContentUpstreamError) { status := upErr.StatusCode message := strings.TrimSpace(extractGoogleErrorMessage(upErr.Body)) service.SetOpsUpstreamError(c, status, message, "") @@ -740,6 +743,10 @@ func writeGeminiCachedContentUpstreamError(c *gin.Context, upErr *service.Gemini c.Data(http.StatusForbidden, "application/json", []byte(service.GeminiCachedContentNotFoundResponse)) return } + if (status == http.StatusTooManyRequests || status >= http.StatusInternalServerError) && !upErr.UpstreamGateway { + h.handleGeminiFailoverExhausted(c, &service.UpstreamFailoverError{StatusCode: status, ResponseBody: upErr.Body}) + return + } if status == http.StatusUnauthorized || status == http.StatusForbidden { googleError(c, http.StatusBadGateway, "Upstream account rejected the cached content request") return @@ -794,11 +801,11 @@ func mergeAccountIDSets(a, b map[int64]struct{}) map[int64]struct{} { return out } -// handleGeminiCachedContentFailoverExhausted 处理绑定缓存的请求在持有账号上重试耗尽：没有其他账号可换， -// 上游的限流 / 不可用响应原样返回，客户端据此判断缓存所在账号的状态；账号凭据类错误仍按网关映射。 +// handleGeminiCachedContentFailoverExhausted 处理绑定缓存的请求在持有账号上重试耗尽：没有其他账号可换。 +// 上游网关账号的限流 / 不可用响应反映持有缓存的上游账号状态，原样返回供客户端判断； +// 其余（含官方账号）仍按普通 Gemini 路径处理：透传规则、通用映射。 func (h *GatewayHandler) handleGeminiCachedContentFailoverExhausted(c *gin.Context, failoverErr *service.UpstreamFailoverError) { - if failoverErr == nil || len(failoverErr.ResponseBody) == 0 || !json.Valid(failoverErr.ResponseBody) || - (failoverErr.StatusCode != http.StatusTooManyRequests && failoverErr.StatusCode < http.StatusInternalServerError) { + if failoverErr == nil || !failoverErr.BoundUpstreamPassthrough || len(failoverErr.ResponseBody) == 0 || !json.Valid(failoverErr.ResponseBody) { h.handleGeminiFailoverExhausted(c, failoverErr) return } diff --git a/backend/internal/handler/gemini_cached_contents_handler_test.go b/backend/internal/handler/gemini_cached_contents_handler_test.go index 8edbcc86e9..26bc5e7183 100644 --- a/backend/internal/handler/gemini_cached_contents_handler_test.go +++ b/backend/internal/handler/gemini_cached_contents_handler_test.go @@ -204,17 +204,45 @@ func TestHandleGeminiCachedContentFailoverExhausted(t *testing.T) { unavailable := `{"error":{"code":503,"message":"The account holding this cached content is temporarily unavailable","status":"UNAVAILABLE"}}` c, rec := newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", apiKey) - h.handleGeminiCachedContentFailoverExhausted(c, &service.UpstreamFailoverError{StatusCode: http.StatusServiceUnavailable, ResponseBody: []byte(unavailable)}) - require.Equal(t, http.StatusServiceUnavailable, rec.Code, "绑定缓存的请求无账号可换，上游 503 原样返回") + h.handleGeminiCachedContentFailoverExhausted(c, &service.UpstreamFailoverError{StatusCode: http.StatusServiceUnavailable, ResponseBody: []byte(unavailable), BoundUpstreamPassthrough: true}) + require.Equal(t, http.StatusServiceUnavailable, rec.Code, "绑定缓存的请求无账号可换，上游网关的 503 原样返回") require.JSONEq(t, unavailable, rec.Body.String()) c, rec = newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", apiKey) - h.handleGeminiCachedContentFailoverExhausted(c, &service.UpstreamFailoverError{StatusCode: http.StatusTooManyRequests, ResponseBody: []byte(`{"error":{"code":429,"message":"Too many pending requests, please retry later","status":"RESOURCE_EXHAUSTED"}}`)}) + h.handleGeminiCachedContentFailoverExhausted(c, &service.UpstreamFailoverError{StatusCode: http.StatusTooManyRequests, ResponseBody: []byte(`{"error":{"code":429,"message":"Too many pending requests, please retry later","status":"RESOURCE_EXHAUSTED"}}`), BoundUpstreamPassthrough: true}) require.Equal(t, http.StatusTooManyRequests, rec.Code) require.Contains(t, rec.Body.String(), "Too many pending requests") + officialQuota := `{"error":{"code":429,"message":"You exceeded your current quota","status":"RESOURCE_EXHAUSTED","details":[{"@type":"type.googleapis.com/google.rpc.QuotaFailure","violations":[{"quotaId":"GenerateRequestsPerDayPerProjectPerModel-FreeTier"}]}]}}` + c, rec = newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", apiKey) + h.handleGeminiCachedContentFailoverExhausted(c, &service.UpstreamFailoverError{StatusCode: http.StatusTooManyRequests, ResponseBody: []byte(officialQuota)}) + require.Equal(t, http.StatusTooManyRequests, rec.Code, "官方账号的限流按普通 Gemini 路径映射") + require.NotContains(t, rec.Body.String(), "FreeTier", "不向客户端暴露上游账号的配额明细") + require.Contains(t, rec.Body.String(), "Upstream rate limit exceeded") + c, rec = newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", apiKey) h.handleGeminiCachedContentFailoverExhausted(c, &service.UpstreamFailoverError{StatusCode: http.StatusUnauthorized, ResponseBody: []byte(`{"error":{"code":401,"message":"bad key"}}`)}) require.Equal(t, http.StatusBadGateway, rec.Code, "账号凭据类错误仍按网关映射，不透传给客户端") require.NotContains(t, rec.Body.String(), "bad key") } + +func TestWriteGeminiCachedContentUpstreamError_RateLimitMappedUnlessUpstreamGateway(t *testing.T) { + h := &GatewayHandler{} + apiKey := geminiCachedContentTestAPIKey(service.PlatformGemini) + quota := `{"error":{"code":429,"message":"You exceeded your current quota","status":"RESOURCE_EXHAUSTED","details":[{"quotaId":"GenerateRequestsPerDayPerProjectPerModel-FreeTier"}]}}` + + c, rec := newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/cachedContents", apiKey) + h.writeGeminiCachedContentUpstreamError(c, &service.GeminiCachedContentUpstreamError{StatusCode: http.StatusTooManyRequests, Body: []byte(quota)}) + require.Equal(t, http.StatusTooManyRequests, rec.Code) + require.NotContains(t, rec.Body.String(), "FreeTier", "官方账号的限流不透传配额明细") + + c, rec = newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/cachedContents", apiKey) + h.writeGeminiCachedContentUpstreamError(c, &service.GeminiCachedContentUpstreamError{StatusCode: http.StatusTooManyRequests, Body: []byte(quota), UpstreamGateway: true}) + require.Equal(t, http.StatusTooManyRequests, rec.Code) + require.JSONEq(t, quota, rec.Body.String(), "上游网关账号的限流反映持有缓存的上游账号状态，原样回写") + + c, rec = newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/cachedContents", apiKey) + h.writeGeminiCachedContentUpstreamError(c, &service.GeminiCachedContentUpstreamError{StatusCode: http.StatusBadRequest, Body: []byte(`{"error":{"code":400,"message":"Cached content is too small","status":"INVALID_ARGUMENT"}}`)}) + require.Equal(t, http.StatusBadRequest, rec.Code) + require.Contains(t, rec.Body.String(), "too small", "请求校验类错误仍原样回写") +} diff --git a/backend/internal/repository/gemini_cached_content_repo.go b/backend/internal/repository/gemini_cached_content_repo.go index d63929c3e6..ddb05bbdcf 100644 --- a/backend/internal/repository/gemini_cached_content_repo.go +++ b/backend/internal/repository/gemini_cached_content_repo.go @@ -3,6 +3,7 @@ package repository import ( "context" "database/sql" + "encoding/json" "time" "github.com/Wei-Shaw/sub2api/internal/service" @@ -23,16 +24,20 @@ func NewGeminiCachedContentRepository(db *sql.DB) service.GeminiCachedContentRep } const geminiCachedContentColumns = `id, public_id, user_id, api_key_id, group_id, account_id, upstream_name, model, - upstream_model, display_name, total_token_count, expire_time, created_at, updated_at` + upstream_model, display_name, total_token_count, expire_time, channel_usage, created_at, updated_at` func (r *geminiCachedContentRepository) Create(ctx context.Context, record *service.GeminiCachedContent) error { + channelUsage, err := json.Marshal(record.ChannelUsage) + if err != nil { + return err + } row := r.sql.QueryRowContext(ctx, ` INSERT INTO gemini_cached_contents (public_id, user_id, api_key_id, group_id, account_id, upstream_name, model, - upstream_model, display_name, total_token_count, expire_time) -VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11) + upstream_model, display_name, total_token_count, expire_time, channel_usage) +VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12) RETURNING id, created_at, updated_at`, record.PublicID, record.UserID, record.APIKeyID, record.GroupID, record.AccountID, record.UpstreamName, - record.Model, record.UpstreamModel, record.DisplayName, record.TotalTokenCount, record.ExpireTime) + record.Model, record.UpstreamModel, record.DisplayName, record.TotalTokenCount, record.ExpireTime, channelUsage) if err := row.Scan(&record.ID, &record.CreatedAt, &record.UpdatedAt); err != nil { return translatePersistenceError(err, nil, service.ErrGeminiCachedContentExists) } @@ -102,6 +107,7 @@ type geminiCachedContentScanner interface { func scanGeminiCachedContent(row geminiCachedContentScanner) (*service.GeminiCachedContent, error) { var record service.GeminiCachedContent + var channelUsage []byte if err := row.Scan( &record.ID, &record.PublicID, @@ -115,10 +121,16 @@ func scanGeminiCachedContent(row geminiCachedContentScanner) (*service.GeminiCac &record.DisplayName, &record.TotalTokenCount, &record.ExpireTime, + &channelUsage, &record.CreatedAt, &record.UpdatedAt, ); err != nil { return nil, err } + if len(channelUsage) > 0 { + if err := json.Unmarshal(channelUsage, &record.ChannelUsage); err != nil { + return nil, err + } + } return &record, nil } diff --git a/backend/internal/repository/gemini_cached_content_repo_integration_test.go b/backend/internal/repository/gemini_cached_content_repo_integration_test.go index 6f6f3cfe9d..fc537a45f3 100644 --- a/backend/internal/repository/gemini_cached_content_repo_integration_test.go +++ b/backend/internal/repository/gemini_cached_content_repo_integration_test.go @@ -24,6 +24,7 @@ func TestGeminiCachedContentRepository_Lifecycle(t *testing.T) { PublicID: publicID + suffix, UserID: 1, APIKeyID: keyID, GroupID: groupID, AccountID: 99, UpstreamName: "cachedContents/up" + publicID, Model: "gemini-3.8-flash", UpstreamModel: "models/gemini-3.8-flash", DisplayName: "label", TotalTokenCount: 2730, ExpireTime: expire, + ChannelUsage: service.ChannelUsageFields{ChannelID: 7, OriginalModel: "gemini-pro", ChannelMappedModel: "gemini-3.8-flash", BillingModelSource: "requested"}, } } first := newRecord("a", apiKeyID, 42, now.Add(time.Hour)) @@ -47,6 +48,7 @@ func TestGeminiCachedContentRepository_Lifecycle(t *testing.T) { require.Equal(t, "cachedContents/upa", got.UpstreamName) require.Equal(t, int64(2730), got.TotalTokenCount) require.True(t, got.ExpireTime.Equal(first.ExpireTime)) + require.Equal(t, first.ChannelUsage, got.ChannelUsage, "创建时的渠道计费口径随记录持久化") _, err = repo.GetForOwner(ctx, apiKeyID+1, first.PublicID) require.ErrorIs(t, err, service.ErrGeminiCachedContentNotFound) diff --git a/backend/internal/repository/scheduler_cache.go b/backend/internal/repository/scheduler_cache.go index e3715f678e..827a750577 100644 --- a/backend/internal/repository/scheduler_cache.go +++ b/backend/internal/repository/scheduler_cache.go @@ -956,7 +956,11 @@ func filterSchedulerCredentials(credentials map[string]any) map[string]any { } // Candidate-list admission evaluates the account override before hydrating // the full account. Dropping it silently falls back to the platform threshold. - keys := []string{"model_mapping", "compact_model_mapping", "api_key", "project_id", "oauth_type", "plan_type", "account_scheduling_threshold"} + keys := []string{"model_mapping", "compact_model_mapping", "api_key", "project_id", "oauth_type", "plan_type", "account_scheduling_threshold", + // Gemini 显式缓存的候选过滤（GeminiExplicitCacheExclusions）读本投影判定账号是否支持显式缓存： + // 官方地址、Vertex 认证方式与上游网关的显式缓存声明。缺这几个键时只能按类型粗筛， + // 不支持的账号要逐个选中后才被剔除。 + "base_url", "auth_mode", service.GeminiExplicitCacheUpstreamCredentialKey, service.GeminiExplicitCacheUpstreamMaxTTLCredentialKey} filtered := make(map[string]any) for _, key := range keys { if value, ok := credentials[key]; ok && value != nil { diff --git a/backend/internal/repository/scheduler_cache_test.go b/backend/internal/repository/scheduler_cache_test.go index d12e77f74c..0542842fdf 100644 --- a/backend/internal/repository/scheduler_cache_test.go +++ b/backend/internal/repository/scheduler_cache_test.go @@ -20,6 +20,23 @@ func TestFilterSchedulerCredentialsKeepsSubscriptionPlanType(t *testing.T) { require.NotContains(t, filtered, "refresh_token") } +func TestFilterSchedulerCredentialsKeepsGeminiExplicitCacheKeys(t *testing.T) { + filtered := filterSchedulerCredentials(map[string]any{ + "api_key": "k", + "base_url": "https://gw.example.com", + "auth_mode": "apikey", + "service_account_json": "secret", + service.GeminiExplicitCacheUpstreamCredentialKey: true, + service.GeminiExplicitCacheUpstreamMaxTTLCredentialKey: 3600, + }) + + require.Equal(t, "https://gw.example.com", filtered["base_url"]) + require.Equal(t, "apikey", filtered["auth_mode"]) + require.Equal(t, true, filtered[service.GeminiExplicitCacheUpstreamCredentialKey]) + require.Equal(t, 3600, filtered[service.GeminiExplicitCacheUpstreamMaxTTLCredentialKey]) + require.NotContains(t, filtered, "service_account_json") +} + func TestSchedulerMetadataAccountKeepsOpenAISubscriptionIdentity(t *testing.T) { account := service.Account{ ID: 24, diff --git a/backend/internal/service/gateway_service.go b/backend/internal/service/gateway_service.go index bdb90cc193..48c74e7d08 100644 --- a/backend/internal/service/gateway_service.go +++ b/backend/internal/service/gateway_service.go @@ -699,6 +699,7 @@ type UpstreamFailoverError struct { SameAccountRetryMax int // 可选的错误级同账号重试上限，低于 handler 默认预算时优先采用 RequestScopedTransient bool // 故障因素与账号无关（如上游按客户端身份/模型容量降载）：可同账号重试，但不得据此对账号做临时封禁 SafeToFailoverAfterWrite bool // 仅写出 SSE 注释等非语义字节时，仍可在同一客户端流中切换账号 + BoundUpstreamPassthrough bool // 绑定显式缓存的请求在上游网关账号上收到的 429 / 5xx：反映持有缓存的上游账号状态，原样回写不走错误映射 Stage GatewayFailureStage Scope GatewayFailureScope Reason GatewayFailureReason diff --git a/backend/internal/service/gateway_usage_billing.go b/backend/internal/service/gateway_usage_billing.go index 4c164f2ed0..84c604fd04 100644 --- a/backend/internal/service/gateway_usage_billing.go +++ b/backend/internal/service/gateway_usage_billing.go @@ -241,7 +241,8 @@ func isForcedUsageBillingRequestID(requestID string) bool { return strings.HasPrefix(id, "web_search:") || strings.HasPrefix(id, "grok-video:") || strings.HasPrefix(id, "grok_audio:") || - strings.HasPrefix(id, "grok_realtime:") + strings.HasPrefix(id, "grok_realtime:") || + strings.HasPrefix(id, GeminiCachedContentUsageRequestIDPrefix) } // StableGrokAudioBillingRequestID is the durable usage_logs / dedup key for one diff --git a/backend/internal/service/gateway_usage_billing_request_id_test.go b/backend/internal/service/gateway_usage_billing_request_id_test.go index 78a64cb338..fa13918b51 100644 --- a/backend/internal/service/gateway_usage_billing_request_id_test.go +++ b/backend/internal/service/gateway_usage_billing_request_id_test.go @@ -31,9 +31,27 @@ func TestIsForcedUsageBillingRequestID(t *testing.T) { require.True(t, isForcedUsageBillingRequestID("grok-video:task-1")) require.True(t, isForcedUsageBillingRequestID("grok_audio:up-1")) require.True(t, isForcedUsageBillingRequestID("grok_realtime:sess-1")) + require.True(t, isForcedUsageBillingRequestID(GeminiCachedContentCreateUsageRequestID("abc"))) require.False(t, isForcedUsageBillingRequestID("resp_abc")) } +func TestGeminiCachedContentUsageRequestIDs(t *testing.T) { + t.Parallel() + publicID, err := NewGeminiCachedContentPublicID() + require.NoError(t, err) + create := GeminiCachedContentCreateUsageRequestID(publicID) + patch := GeminiCachedContentPatchUsageRequestID(publicID) + require.Equal(t, "gcache:create:"+publicID, create) + require.True(t, strings.HasPrefix(patch, "gcache:patch:"+publicID+":")) + require.NotEqual(t, patch, GeminiCachedContentPatchUsageRequestID(publicID), "每次延长各自唯一") + require.LessOrEqual(t, len(create), 64, "usage_logs.request_id 列长 64") + require.LessOrEqual(t, len(patch), 64, "usage_logs.request_id 列长 64") + + ctx := context.WithValue(context.Background(), ctxkey.ClientRequestID, "client-shared-id") + require.Equal(t, create, resolveUsageBillingRequestID(ctx, create), "缓存费用行不被客户端请求 id 覆盖") + require.Equal(t, patch, resolveUsageBillingRequestID(ctx, patch)) +} + func TestStableGrokAudioBillingRequestID(t *testing.T) { t.Parallel() require.Equal(t, "grok_audio:up-1", StableGrokAudioBillingRequestID("up-1")) diff --git a/backend/internal/service/gemini_cached_content.go b/backend/internal/service/gemini_cached_content.go index 0fe24249c3..92157dffbb 100644 --- a/backend/internal/service/gemini_cached_content.go +++ b/backend/internal/service/gemini_cached_content.go @@ -74,8 +74,10 @@ type GeminiCachedContent struct { DisplayName string TotalTokenCount int64 ExpireTime time.Time - CreatedAt time.Time - UpdatedAt time.Time + // ChannelUsage 是创建时的渠道计费口径（渠道 ID、请求模型、计费模型来源），延长有效期按同一口径计费。 + ChannelUsage ChannelUsageFields + CreatedAt time.Time + UpdatedAt time.Time } // PublicName 返回客户端可见的缓存资源名。 @@ -127,9 +129,11 @@ func IsGeminiExplicitCacheAccount(account *Account) bool { return isOfficialGeminiAPIBaseURL(account.GetCredential("base_url")) || account.GeminiExplicitCacheUpstreamEnabled() } +// GeminiExplicitCacheUpstreamCredentialKey / GeminiExplicitCacheUpstreamMaxTTLCredentialKey 是 API Key 账号凭据里 +// 声明上游实现显式缓存接口及其有效期上限的键；调度快照投影须保留它们，候选过滤才能直接判定账号是否支持显式缓存。 const ( - geminiExplicitCacheUpstreamKey = "explicit_cache_upstream" - geminiExplicitCacheUpstreamMaxTTLKey = "explicit_cache_upstream_max_ttl_seconds" + GeminiExplicitCacheUpstreamCredentialKey = "explicit_cache_upstream" + GeminiExplicitCacheUpstreamMaxTTLCredentialKey = "explicit_cache_upstream_max_ttl_seconds" ) // GeminiExplicitCacheUpstreamEnabled 判断 API Key 账号是否声明其上游实现了显式缓存接口， @@ -138,7 +142,7 @@ func (a *Account) GeminiExplicitCacheUpstreamEnabled() bool { if a == nil || a.Type != AccountTypeAPIKey || a.Credentials == nil { return false } - enabled, _ := a.Credentials[geminiExplicitCacheUpstreamKey].(bool) + enabled, _ := a.Credentials[GeminiExplicitCacheUpstreamCredentialKey].(bool) return enabled } @@ -148,7 +152,7 @@ func (a *Account) GeminiExplicitCacheUpstreamMaxTTL() time.Duration { return 0 } var seconds float64 - switch v := a.Credentials[geminiExplicitCacheUpstreamMaxTTLKey].(type) { + switch v := a.Credentials[GeminiExplicitCacheUpstreamMaxTTLCredentialKey].(type) { case float64: seconds = v case int: @@ -198,6 +202,25 @@ func NewGeminiCachedContentPublicID() (string, error) { return strings.ToLower(geminiCachedContentPublicIDEncoding.EncodeToString(b[:])), nil } +// GeminiCachedContentUsageRequestIDPrefix 是显式缓存创建 / 延长费用在 usage_logs.request_id 中的固定前缀： +// 这类费用行是独立的扣费事件，不随客户端或网关请求 id 变化，按缓存 ID 即可检索一个缓存的全部费用行。 +// 两种格式都不超过 request_id 列的 64 字符。 +const GeminiCachedContentUsageRequestIDPrefix = "gcache:" + +// GeminiCachedContentCreateUsageRequestID 返回创建费用的 request_id（"gcache:create:{id}"）；一个缓存只创建一次。 +func GeminiCachedContentCreateUsageRequestID(publicID string) string { + return GeminiCachedContentUsageRequestIDPrefix + "create:" + publicID +} + +// GeminiCachedContentPatchUsageRequestID 返回一次延长费用的 request_id（"gcache:patch:{id}:{8 位随机}"）；每次延长各自唯一。 +func GeminiCachedContentPatchUsageRequestID(publicID string) string { + var b [5]byte + if _, err := rand.Read(b[:]); err != nil { + return GeminiCachedContentUsageRequestIDPrefix + "patch:" + publicID + ":" + strconv.FormatInt(time.Now().UnixNano()%(1<<40), 32) + } + return GeminiCachedContentUsageRequestIDPrefix + "patch:" + publicID + ":" + strings.ToLower(geminiCachedContentPublicIDEncoding.EncodeToString(b[:])) +} + // ParseGeminiCachedContentPublicID 从 "cachedContents/{id}" 或裸 id 中解析缓存 ID。 func ParseGeminiCachedContentPublicID(name string) (string, bool) { id := strings.TrimPrefix(strings.TrimSpace(name), GeminiCachedContentNamePrefix) @@ -525,6 +548,8 @@ type GeminiCachedContentUpstream struct { type GeminiCachedContentUpstreamError struct { StatusCode int Body []byte + // UpstreamGateway 表示响应来自声明实现显式缓存的上游网关账号：限流 / 不可用反映的是持有缓存的上游账号状态。 + UpstreamGateway bool } func (e *GeminiCachedContentUpstreamError) Error() string { @@ -827,7 +852,7 @@ func (s *GeminiCachedContentService) doUpstream(ctx context.Context, account *Ac return nil, err } if !success { - return nil, &GeminiCachedContentUpstreamError{StatusCode: resp.StatusCode, Body: respBody} + return nil, &GeminiCachedContentUpstreamError{StatusCode: resp.StatusCode, Body: respBody, UpstreamGateway: account.GeminiExplicitCacheUpstreamEnabled()} } return respBody, nil } @@ -936,8 +961,9 @@ func GeminiCachedContentView(record *GeminiCachedContent) map[string]any { } // GeminiExplicitCacheExclusions 返回显式缓存请求在 gemini 分组内需要排除的可调度账号。 -// boundAccountID > 0 时只保留该账号；否则排除平台或类型不支持显式缓存的账号 -// （调度快照不含 base_url，官方地址判定在选中后对完整账号进行）。 +// boundAccountID > 0 时只保留该账号；否则排除不支持显式缓存的账号。判定所需的凭据键 +// （api_key、base_url、auth_mode、上游显式缓存声明）都在调度快照投影内，候选过滤一次算全； +// 选中后仍对完整账号复核。 func (s *GatewayService) GeminiExplicitCacheExclusions(ctx context.Context, groupID *int64, boundAccountID int64) (map[int64]struct{}, error) { accounts, _, err := s.listSchedulableAccounts(ctx, groupID, PlatformGemini, false) if err != nil { @@ -946,7 +972,7 @@ func (s *GatewayService) GeminiExplicitCacheExclusions(ctx context.Context, grou excluded := make(map[int64]struct{}, len(accounts)) for i := range accounts { account := &accounts[i] - keep := IsGeminiExplicitCacheAccountType(account) + keep := IsGeminiExplicitCacheAccount(account) if boundAccountID > 0 { keep = account.ID == boundAccountID } diff --git a/backend/internal/service/gemini_cached_content_hardening_test.go b/backend/internal/service/gemini_cached_content_hardening_test.go index cd9c6fa57a..b1b893187a 100644 --- a/backend/internal/service/gemini_cached_content_hardening_test.go +++ b/backend/internal/service/gemini_cached_content_hardening_test.go @@ -285,6 +285,7 @@ func TestGeminiForwardNative_BoundCacheUpstreamGatewayErrorsSkipAccountPenalty(t require.True(t, errors.As(err, &failoverErr), "交由绑定请求的耗尽处理透传") require.Equal(t, status, failoverErr.StatusCode) require.JSONEq(t, body, string(failoverErr.ResponseBody)) + require.True(t, failoverErr.BoundUpstreamPassthrough, "标记为上游网关账号的响应，handler 原样回写") require.False(t, failoverErr.RetryableOnSameAccount) require.Equal(t, 1, httpStub.calls, "不在同一账号上退避重试") require.Zero(t, repo.rateLimitedCalls, "不限流整个网关账号") diff --git a/backend/internal/service/gemini_cached_content_test.go b/backend/internal/service/gemini_cached_content_test.go index 0aec766105..94e1027e23 100644 --- a/backend/internal/service/gemini_cached_content_test.go +++ b/backend/internal/service/gemini_cached_content_test.go @@ -421,7 +421,7 @@ func TestGeminiCachedContentStorageCostInAccountStats(t *testing.T) { gid := groupID err := svc.RecordUsage(context.Background(), &RecordUsageInput{ Result: &ForwardResult{ - RequestID: "gemini-cached-content:create:" + tc.name, + RequestID: GeminiCachedContentCreateUsageRequestID(tc.name), Model: model, Usage: ClaudeUsage{InputTokens: tc.tokens}, CacheStorageTokenHours: storage / 0.0000005, diff --git a/backend/internal/service/gemini_messages_compat_service.go b/backend/internal/service/gemini_messages_compat_service.go index a6991ea70c..cc420e2ddd 100644 --- a/backend/internal/service/gemini_messages_compat_service.go +++ b/backend/internal/service/gemini_messages_compat_service.go @@ -1470,7 +1470,7 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin. } // 上游网关的限流 / 不可用只针对其持有缓存的账号：不改本账号状态，交由绑定请求的耗尽处理透传。 if isGeminiBoundCacheUpstreamPassthrough(ctx, account, resp.StatusCode) { - return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: respBody} + return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: respBody, BoundUpstreamPassthrough: true} } // Best-effort fallback for OAuth tokens missing AI Studio scopes when calling countTokens. // This avoids Gemini SDKs failing hard during preflight token counting. diff --git a/backend/migrations/244_gemini_cached_contents_channel_usage.sql b/backend/migrations/244_gemini_cached_contents_channel_usage.sql new file mode 100644 index 0000000000..c72f097a73 --- /dev/null +++ b/backend/migrations/244_gemini_cached_contents_channel_usage.sql @@ -0,0 +1,4 @@ +-- Gemini 显式缓存创建时的渠道计费口径（渠道 ID / 请求模型 / 计费模型来源）， +-- 延长有效期按同一口径计费。 +ALTER TABLE gemini_cached_contents + ADD COLUMN IF NOT EXISTS channel_usage JSONB NOT NULL DEFAULT '{}'::jsonb; From f0ee427f0ad2f7eab6ab8f095ba05041af3a48f8 Mon Sep 17 00:00:00 2001 From: feeeei Date: Sun, 11 Oct 2026 11:22:43 +0800 Subject: [PATCH 11/15] =?UTF-8?q?feat(gemini):=20=E6=98=BE=E5=BC=8F?= =?UTF-8?q?=E7=BC=93=E5=AD=98=E5=AD=98=E5=82=A8=E4=BB=B7=E5=86=85=E7=BD=AE?= =?UTF-8?q?=E5=AE=98=E6=96=B9=E4=BB=B7=E5=B9=B6=E6=8C=89=20Pro=20/=20Flash?= =?UTF-8?q?=20=E7=B3=BB=E5=88=97=E5=85=9C=E5=BA=95?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 新增 Gemini 显式缓存存储单价表（Standard 档，USD / token / 小时）： Pro 全系 $4.50、Flash 与 Flash-Lite 全系 $1.00；3.6 / 3.7 / 3.8 Flash 为优惠价 $0.50（有效期至 2026-12-31，2027-01-01 起恢复 $1.00）， AI Studio 与 Vertex 统一按优惠价计 - 型号表未命中的 Gemini 模型按系列兜底：Pro 取 $4.50，Flash / Flash-Lite 取标准价 $1.00；Flash-Lite 变体不落到同版本 Flash 的优惠价 - 取值优先级：价格表（override_file 等）配置值 > 内置型号价 > 内置系列价； 两者都没有的模型仍拒绝创建缓存 --- backend/internal/service/billing_service.go | 25 +++--- .../service/gemini_cache_storage_pricing.go | 72 +++++++++++++++++ .../gemini_cache_storage_pricing_test.go | 77 +++++++++++++++++++ backend/internal/service/pricing_service.go | 3 +- deploy/config.example.yaml | 10 ++- 5 files changed, 169 insertions(+), 18 deletions(-) create mode 100644 backend/internal/service/gemini_cache_storage_pricing.go create mode 100644 backend/internal/service/gemini_cache_storage_pricing_test.go diff --git a/backend/internal/service/billing_service.go b/backend/internal/service/billing_service.go index 6dfb9a0399..84ee77960d 100644 --- a/backend/internal/service/billing_service.go +++ b/backend/internal/service/billing_service.go @@ -1325,22 +1325,21 @@ func (s *BillingService) GetModelPricing(model string) (*ModelPricing, error) { return s.getModelPricingAt(model, timezone.Now()) } +// CacheStoragePricePerTokenHour 返回模型的显式上下文缓存存储单价（每 token 每小时，USD）。 +// 优先级：价格表（override_file 等）按确定性识别配置的值 > 内置型号价 > 内置 Pro/Flash 系列兜底价； +// 均未命中返回 0。 +func (s *BillingService) CacheStoragePricePerTokenHour(model string) float64 { + if s != nil && s.pricingService != nil { + if pricing := s.pricingService.GetIdentifiedModelPricing(model); pricing != nil && pricing.CacheStorageCostPerTokenPerHour > 0 { + return pricing.CacheStorageCostPerTokenPerHour + } + } + return builtinGeminiCacheStoragePrice(model) +} + // getModelPricingAt 是 GetModelPricing 的带计费时点内部变体：pricingAt 显式 // 驱动 DeepSeek pro→Flash 切换判定（切换点前 Pro 价、之后 Flash 价），使 // 展示/估算路径可与历史补账同刻复算，测试也能用固定时点钉住断言。 -// CacheStoragePricePerTokenHour 返回模型的显式上下文缓存存储单价（每 token 每小时，USD）。 -// 只做确定性识别，不按系列兜底；价格表未配置时返回 0。 -func (s *BillingService) CacheStoragePricePerTokenHour(model string) float64 { - if s == nil || s.pricingService == nil { - return 0 - } - pricing := s.pricingService.GetIdentifiedModelPricing(model) - if pricing == nil || pricing.CacheStorageCostPerTokenPerHour <= 0 { - return 0 - } - return pricing.CacheStorageCostPerTokenPerHour -} - func (s *BillingService) getModelPricingAt(model string, pricingAt time.Time) (*ModelPricing, error) { // 标准化模型名称（转小写） model = strings.ToLower(model) diff --git a/backend/internal/service/gemini_cache_storage_pricing.go b/backend/internal/service/gemini_cache_storage_pricing.go new file mode 100644 index 0000000000..c08ed2f683 --- /dev/null +++ b/backend/internal/service/gemini_cache_storage_pricing.go @@ -0,0 +1,72 @@ +package service + +import ( + "slices" + "strings" +) + +// Gemini 显式上下文缓存存储的内置官方单价（USD / token / 小时，Standard 档）。 +// 存储价与上下文长度无关（≤200K 与 >200K 同价）；Priority 档的更高存储价不在此列。 +// 来源：https://ai.google.dev/gemini-api/docs/pricing +// 与 https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing +const ( + geminiProCacheStoragePrice = 4.5e-6 // $4.50 / 1M tokens / 小时 + geminiFlashCacheStoragePrice = 1e-6 // $1.00 / 1M tokens / 小时，Flash-Lite 同价 + + // 3.6/3.7/3.8 Flash 的优惠价 $0.50 / 1M tokens / 小时，有效期至 2026-12-31， + // 2027-01-01 起恢复为 geminiFlashCacheStoragePrice。AI Studio 价目页直接标注该优惠价； + // Vertex 价目页存储价标 $1.00，优惠以「按该模型净消费返还 50% 抵扣金」发放， + // 官方未明确存储费是否在返还范围内，两类账号统一按优惠价计。 + geminiFlashPromoCacheStoragePrice = 0.5e-6 +) + +// geminiCacheStorageModelPrices 是已发布型号的存储单价。模型名等于前缀，或以「前缀-」开头 +// （-preview、-customtools、-001 等变体）即命中；更具体的前缀排在前面。 +var geminiCacheStorageModelPrices = []struct { + prefix string + price float64 +}{ + {"gemini-3.8-flash-cyber", geminiFlashCacheStoragePrice}, + {"gemini-3.8-flash", geminiFlashPromoCacheStoragePrice}, + {"gemini-3.7-flash", geminiFlashPromoCacheStoragePrice}, + {"gemini-3.6-flash", geminiFlashPromoCacheStoragePrice}, + {"gemini-3.5-flash-lite", geminiFlashCacheStoragePrice}, + {"gemini-3.5-flash", geminiFlashCacheStoragePrice}, + {"gemini-3.1-flash-lite", geminiFlashCacheStoragePrice}, + {"gemini-3.1-pro", geminiProCacheStoragePrice}, + {"gemini-3-flash", geminiFlashCacheStoragePrice}, + {"gemini-3-pro", geminiProCacheStoragePrice}, + {"gemini-2.5-flash-lite", geminiFlashCacheStoragePrice}, + {"gemini-2.5-flash", geminiFlashCacheStoragePrice}, + {"gemini-2.5-pro", geminiProCacheStoragePrice}, + {"gemini-2.0-flash-lite", geminiFlashCacheStoragePrice}, + {"gemini-2.0-flash", geminiFlashCacheStoragePrice}, +} + +// builtinGeminiCacheStoragePrice 返回内置的显式缓存存储单价：先按已发布型号匹配，未命中时按系列兜底， +// Pro 系列取 geminiProCacheStoragePrice，Flash 与 Flash-Lite 系列取 geminiFlashCacheStoragePrice（标准价， +// 不套用优惠价）。其他 Gemini 模型与非 Gemini 模型返回 0。 +func builtinGeminiCacheStoragePrice(model string) float64 { + name := lastSegment(strings.ToLower(strings.TrimSpace(model))) + if !strings.HasPrefix(name, "gemini-") { + return 0 + } + parts := strings.Split(name, "-") + lite := slices.Contains(parts, "lite") + for _, entry := range geminiCacheStorageModelPrices { + // Flash-Lite 变体只匹配 Flash-Lite 条目：gemini-3.8-flash-lite 不能落到 3.8 Flash 的优惠价。 + if lite != strings.HasSuffix(entry.prefix, "-lite") { + continue + } + if name == entry.prefix || strings.HasPrefix(name, entry.prefix+"-") { + return entry.price + } + } + switch { + case slices.Contains(parts, "pro"): + return geminiProCacheStoragePrice + case slices.Contains(parts, "flash"): + return geminiFlashCacheStoragePrice + } + return 0 +} diff --git a/backend/internal/service/gemini_cache_storage_pricing_test.go b/backend/internal/service/gemini_cache_storage_pricing_test.go new file mode 100644 index 0000000000..846f7b56ec --- /dev/null +++ b/backend/internal/service/gemini_cache_storage_pricing_test.go @@ -0,0 +1,77 @@ +//go:build unit + +package service + +import ( + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/stretchr/testify/require" +) + +func TestBuiltinGeminiCacheStoragePrice(t *testing.T) { + for _, tc := range []struct { + model string + want float64 + }{ + // 已发布型号 + {"gemini-3.8-flash", 0.5e-6}, + {"gemini-3.7-flash", 0.5e-6}, + {"gemini-3.6-flash", 0.5e-6}, + {"gemini-3.8-flash-cyber", 1e-6}, + {"gemini-3.5-flash", 1e-6}, + {"gemini-3.5-flash-lite", 1e-6}, + {"gemini-3.1-flash-lite", 1e-6}, + {"gemini-3.1-pro-preview", 4.5e-6}, + {"gemini-3.1-pro-preview-customtools", 4.5e-6}, + {"gemini-3-flash-preview", 1e-6}, + {"gemini-3-pro-preview", 4.5e-6}, + {"gemini-2.5-pro", 4.5e-6}, + {"gemini-2.5-flash", 1e-6}, + {"gemini-2.5-flash-lite", 1e-6}, + {"gemini-2.0-flash-001", 1e-6}, + {"gemini-2.0-flash-lite-001", 1e-6}, + // 资源名与大小写 + {"models/gemini-3.8-flash", 0.5e-6}, + {"projects/p/locations/global/publishers/google/models/gemini-2.5-pro", 4.5e-6}, + {" Gemini-3.8-Flash ", 0.5e-6}, + // 系列兜底：未发布型号与别名按标准价 + {"gemini-3.9-flash", 1e-6}, + {"gemini-3.9-flash-lite", 1e-6}, + {"gemini-3.8-flash-lite", 1e-6}, + {"gemini-3.9-pro-preview", 4.5e-6}, + {"gemini-flash-latest", 1e-6}, + {"gemini-flash-lite-latest", 1e-6}, + {"gemini-pro-latest", 4.5e-6}, + // 无存储价 + {"gemini-embedding-001", 0}, + {"gemini-unknown-model", 0}, + {"claude-sonnet-5-5", 0}, + {"", 0}, + } { + require.InDelta(t, tc.want, builtinGeminiCacheStoragePrice(tc.model), 1e-15, tc.model) + } +} + +func TestCacheStoragePricePerTokenHourPriority(t *testing.T) { + pricing := &PricingService{} + data, err := pricing.parsePricingData([]byte(`{ + "gemini-3.8-flash":{"litellm_provider":"gemini","input_cost_per_token":0.00000075,"output_cost_per_token":0.00000375, + "cache_storage_cost_per_token_per_hour":0.000002}, + "gemini-3.9-flash":{"litellm_provider":"gemini","input_cost_per_token":0.0000015,"output_cost_per_token":0.0000075, + "cache_storage_cost_per_token_per_hour":0.0000007}, + "gemini-2.5-pro":{"litellm_provider":"gemini","input_cost_per_token":0.00000125,"output_cost_per_token":0.00001}}`)) + require.NoError(t, err) + pricing.pricingData = data + billing := NewBillingService(&config.Config{}, pricing) + + require.InDelta(t, 2e-6, billing.CacheStoragePricePerTokenHour("gemini-3.8-flash"), 1e-15, "价格表配置值优先于内置型号价") + require.InDelta(t, 0.7e-6, billing.CacheStoragePricePerTokenHour("gemini-3.9-flash"), 1e-15, "价格表配置值优先于内置系列兜底价") + require.InDelta(t, 4.5e-6, billing.CacheStoragePricePerTokenHour("gemini-2.5-pro"), 1e-15, "价格表条目缺存储价时取内置型号价") + require.InDelta(t, 0.5e-6, billing.CacheStoragePricePerTokenHour("gemini-3.7-flash"), 1e-15, "价格表无条目时取内置型号价") + require.InDelta(t, 4.5e-6, billing.CacheStoragePricePerTokenHour("gemini-3.9-pro"), 1e-15, "内置型号未命中时取系列兜底价") + require.Zero(t, billing.CacheStoragePricePerTokenHour("customer-alias")) + + noCatalog := NewBillingService(&config.Config{}, nil) + require.InDelta(t, 0.5e-6, noCatalog.CacheStoragePricePerTokenHour("gemini-3.8-flash"), 1e-15, "无价格服务时取内置价") +} diff --git a/backend/internal/service/pricing_service.go b/backend/internal/service/pricing_service.go index 45b530fd6d..679ee78226 100644 --- a/backend/internal/service/pricing_service.go +++ b/backend/internal/service/pricing_service.go @@ -223,7 +223,8 @@ type LiteLLMModelPricing struct { OutputCostPerImageToken float64 `json:"output_cost_per_image_token"` // 图片输出 token 价格 InputCostPerImageToken float64 `json:"input_cost_per_image_token"` // 图片输入 token 价格（如 gpt-image-2 图片编辑） CacheReadInputImageTokenCost float64 `json:"cache_read_input_image_token_cost"` - // CacheStorageCostPerTokenPerHour 显式上下文缓存存储单价（每 token 每小时，USD），仅由定价覆盖文件提供。 + // CacheStorageCostPerTokenPerHour 显式上下文缓存存储单价（每 token 每小时，USD），一般由定价覆盖文件提供， + // 配置后优先于内置价（见 BillingService.CacheStoragePricePerTokenHour）。 CacheStorageCostPerTokenPerHour float64 `json:"cache_storage_cost_per_token_per_hour"` // TokenPricingAbsent 表示源数据中 input/output token 价格均缺失（仅有图片价）。 diff --git a/deploy/config.example.yaml b/deploy/config.example.yaml index 174d5a4437..6b8b049838 100644 --- a/deploy/config.example.yaml +++ b/deploy/config.example.yaml @@ -1105,11 +1105,13 @@ pricing: # 覆盖补丁文件（可选，优先级最高）：条目按字段浅合并覆盖目录/回退数据， # 字段值为 null 表示删除该字段。模型名须与目录键完全一致。 # 修改后在下一次哈希检查时按内容哈希自动重载，无需重启。 - # Gemini explicit context caching requires a storage price per model, e.g. + # Gemini explicit context caching storage price (USD per token per hour) comes + # from the field cache_storage_cost_per_token_per_hour here, e.g. # {"gemini-3.8-flash": {"cache_storage_cost_per_token_per_hour": 0.0000005}}; - # models without it reject POST /v1beta/cachedContents. - # Gemini 显式缓存需按模型配置存储单价（每 token 每小时，USD）字段 - # cache_storage_cost_per_token_per_hour；未配置的模型拒绝创建缓存。 + # without it, built-in official prices apply (released models first, then the + # Pro / Flash series rate). Models with no price at all reject POST /v1beta/cachedContents. + # Gemini 显式缓存存储单价（每 token 每小时，USD）优先取这里的 cache_storage_cost_per_token_per_hour 字段， + # 未配置时使用内置官方价（先按已发布型号，再按 Pro / Flash 系列兜底）；两者都没有的模型拒绝创建缓存。 # 目录中不存在的模型条目还须同时提供 input_cost_per_token / output_cost_per_token。 # override_file: "./data/model_pricing_overrides.json" # Update interval in hours From 4426593e9d7dc38a009f28da038ac2906a57c17e Mon Sep 17 00:00:00 2001 From: feeeei Date: Sun, 11 Oct 2026 11:22:48 +0800 Subject: [PATCH 12/15] =?UTF-8?q?fix(gemini):=203.6=20Flash=20=E5=85=9C?= =?UTF-8?q?=E5=BA=95=E4=BB=B7=E6=94=B9=E4=B8=BA=E5=AE=98=E6=96=B9=E4=BC=98?= =?UTF-8?q?=E6=83=A0=E4=BB=B7?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 内置兜底价与本地价格表快照的 gemini-3.6-flash 改为 $0.75 输入 / $3.75 输出 / $0.075 缓存读取（Priority 档 $1.35 / $6.75 / $0.135），有效期至 2026-12-31， 2027-01-01 起恢复 $1.50 / $7.50 / $0.15，与 3.7 / 3.8 Flash 及远程价格目录一致 --- .../service/account_stats_pricing_test.go | 2 +- backend/internal/service/billing_service.go | 15 ++++++++------- backend/internal/service/pricing_service_test.go | 14 +++++++------- .../model_prices_and_context_window.json | 12 ++++++------ 4 files changed, 22 insertions(+), 21 deletions(-) diff --git a/backend/internal/service/account_stats_pricing_test.go b/backend/internal/service/account_stats_pricing_test.go index 26eb645935..0f5d236ca1 100644 --- a/backend/internal/service/account_stats_pricing_test.go +++ b/backend/internal/service/account_stats_pricing_test.go @@ -933,7 +933,7 @@ func TestResolveAccountStatsCost_Gemini36FlashTierUsesFallbackPricing(t *testing UsageTokens{InputTokens: 1_000_000, OutputTokens: 1_000_000, CacheReadTokens: 1_000_000}, 1, 0, "", time.Time{}, true, ) require.NotNil(t, result) - require.InDelta(t, 9.15, *result, 1e-12) + require.InDelta(t, 4.575, *result, 1e-12) } func TestResolveAccountStatsCost_AllMiss_ReturnsNil(t *testing.T) { diff --git a/backend/internal/service/billing_service.go b/backend/internal/service/billing_service.go index 84ee77960d..8597b48e79 100644 --- a/backend/internal/service/billing_service.go +++ b/backend/internal/service/billing_service.go @@ -474,14 +474,15 @@ func (s *BillingService) initFallbackPricing() { SupportsCacheBreakdown: false, } - // Gemini 3.6 Flash (Google AI pricing: $1.50 input / $7.50 output / - // $0.15 cached input per MTok). Antigravity's -high/-low/-medium/-tiered - // aliases are matched below so unavailable remote pricing never records - // token-bearing requests at $0. + // Gemini 3.6 Flash (Google AI pricing: $0.75 input / $3.75 output / + // $0.075 cached input per MTok, promotional through 2026-12-31; official + // rates double to $1.50/$7.50/$0.15 from 2027-01-01). Antigravity's + // -high/-low/-medium/-tiered aliases are matched below so unavailable + // remote pricing never records token-bearing requests at $0. s.fallbackPrices["gemini-3.6-flash"] = &ModelPricing{ - InputPricePerToken: 1.5e-6, - OutputPricePerToken: 7.5e-6, - CacheReadPricePerToken: 0.15e-6, + InputPricePerToken: 0.75e-6, + OutputPricePerToken: 3.75e-6, + CacheReadPricePerToken: 0.075e-6, SupportsCacheBreakdown: false, } diff --git a/backend/internal/service/pricing_service_test.go b/backend/internal/service/pricing_service_test.go index 0dce3e435c..3d9f0cb09e 100644 --- a/backend/internal/service/pricing_service_test.go +++ b/backend/internal/service/pricing_service_test.go @@ -613,10 +613,10 @@ func TestBillingService_Gemini36FlashThinkingTierFallbacksAreBillable(t *testing t.Run(model, func(t *testing.T) { cost, err := svc.CalculateCost(model, tokens, 1) require.NoError(t, err) - require.InDelta(t, 1.5, cost.InputCost, 1e-12) - require.InDelta(t, 7.5, cost.OutputCost, 1e-12) - require.InDelta(t, 0.15, cost.CacheReadCost, 1e-12) - require.InDelta(t, 9.15, cost.TotalCost, 1e-12) + require.InDelta(t, 0.75, cost.InputCost, 1e-12) + require.InDelta(t, 3.75, cost.OutputCost, 1e-12) + require.InDelta(t, 0.075, cost.CacheReadCost, 1e-12) + require.InDelta(t, 4.575, cost.TotalCost, 1e-12) }) } } @@ -725,9 +725,9 @@ func TestDefaultPricingIncludesGemini36FlashRates(t *testing.T) { t.Run(model, func(t *testing.T) { pricing, err := billingSvc.GetModelPricing(model) require.NoError(t, err) - require.InDelta(t, 1.5e-6, pricing.InputPricePerToken, 1e-12) - require.InDelta(t, 7.5e-6, pricing.OutputPricePerToken, 1e-12) - require.InDelta(t, 0.15e-6, pricing.CacheReadPricePerToken, 1e-12) + require.InDelta(t, 0.75e-6, pricing.InputPricePerToken, 1e-12) + require.InDelta(t, 3.75e-6, pricing.OutputPricePerToken, 1e-12) + require.InDelta(t, 0.075e-6, pricing.CacheReadPricePerToken, 1e-12) }) } } diff --git a/backend/resources/model-pricing/model_prices_and_context_window.json b/backend/resources/model-pricing/model_prices_and_context_window.json index 99db4dbbdd..07d02a9b6c 100644 --- a/backend/resources/model-pricing/model_prices_and_context_window.json +++ b/backend/resources/model-pricing/model_prices_and_context_window.json @@ -2616,17 +2616,17 @@ "web_search_billing_unit": "per_query" }, "gemini-3.6-flash": { - "cache_read_input_token_cost": 1.5e-07, - "cache_read_input_token_cost_priority": 2.7e-07, - "input_cost_per_token": 1.5e-06, - "input_cost_per_token_priority": 2.7e-06, + "cache_read_input_token_cost": 7.5e-08, + "cache_read_input_token_cost_priority": 1.35e-07, + "input_cost_per_token": 7.5e-07, + "input_cost_per_token_priority": 1.35e-06, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 1048576, "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", - "output_cost_per_token": 7.5e-06, - "output_cost_per_token_priority": 1.35e-05, + "output_cost_per_token": 3.75e-06, + "output_cost_per_token_priority": 6.75e-06, "source": "https://ai.google.dev/gemini-api/docs/latest-model", "supports_prompt_caching": true, "supports_reasoning": true, From daebb7a80c720763a01e5cef8799c030aed0b058 Mon Sep 17 00:00:00 2001 From: feeeei Date: Sun, 11 Oct 2026 21:57:19 +0800 Subject: [PATCH 13/15] =?UTF-8?q?fix(gemini):=20=E6=A8=A1=E5=9E=8B?= =?UTF-8?q?=E5=88=97=E8=A1=A8=E4=B8=8E=E6=A8=A1=E5=9E=8B=E8=AF=A6=E6=83=85?= =?UTF-8?q?=E6=8E=A5=E5=8F=A3=E6=94=AF=E6=8C=81=20Agent=20Platform=20(Vert?= =?UTF-8?q?ex)=20=E8=B4=A6=E5=8F=B7?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Gemini 分组对外提供 AI Studio 兼容协议，分组内账号不一定是 AI Studio。 分组只剩 Vertex 服务账号时，GET /v1beta/models 与 /v1beta/models/{model} 选中服务账号后报 502 unsupported account type: service_account。 - SelectAccountForAIStudioEndpoints 改名 SelectAccountForModelsEndpoints， Vertex 账号排在所有 AI Studio 兼容账号之后参与选号 - ForwardAIStudioGET 拆为 ForwardListModels / ForwardGetModel，按账号类型分流： AI Studio 兼容账号原样转发；Vertex 服务账号读取 Google 发布方模型列表并 转换为 AI Studio 格式，模型详情在该模型所用区域的列表中查找；Vertex API Key 账号返回内置列表 - Vertex 发布方模型分页读取抽成 listVertexGeminiPublisherModels，与后台 同步上游模型共用 - 模型接口响应体上限改用 gateway.models_list_read_max_bytes，超出时报错 --- .../internal/handler/gemini_v1beta_handler.go | 8 +- .../handler/gemini_v1beta_handler_test.go | 4 +- .../handler/gemini_vertex_models_test.go | 83 +++++ .../service/gemini_messages_compat_service.go | 160 ---------- .../service/gemini_models_endpoints.go | 302 ++++++++++++++++++ .../service/gemini_models_endpoints_test.go | 208 ++++++++++++ backend/internal/service/upstream_models.go | 31 +- 7 files changed, 620 insertions(+), 176 deletions(-) create mode 100644 backend/internal/handler/gemini_vertex_models_test.go create mode 100644 backend/internal/service/gemini_models_endpoints.go create mode 100644 backend/internal/service/gemini_models_endpoints_test.go diff --git a/backend/internal/handler/gemini_v1beta_handler.go b/backend/internal/handler/gemini_v1beta_handler.go index 63e944cb5e..8627e0d139 100644 --- a/backend/internal/handler/gemini_v1beta_handler.go +++ b/backend/internal/handler/gemini_v1beta_handler.go @@ -74,7 +74,7 @@ func (h *GatewayHandler) GeminiV1BetaListModels(c *gin.Context) { return } - account, err := h.geminiCompatService.SelectAccountForAIStudioEndpoints(c.Request.Context(), apiKey.GroupID) + account, err := h.geminiCompatService.SelectAccountForModelsEndpoints(c.Request.Context(), apiKey.GroupID) if err != nil { if len(agModels) > 0 { c.JSON(http.StatusOK, gemini.ModelsListResponse{Models: filterGeminiModels(agModels)}) @@ -85,7 +85,7 @@ func (h *GatewayHandler) GeminiV1BetaListModels(c *gin.Context) { return } - res, err := h.geminiCompatService.ForwardAIStudioGET(c.Request.Context(), account, "/v1beta/models") + res, err := h.geminiCompatService.ForwardListModels(c.Request.Context(), account) if err != nil { googleError(c, http.StatusBadGateway, err.Error()) return @@ -251,7 +251,7 @@ func (h *GatewayHandler) GeminiV1BetaGetModel(c *gin.Context) { return } - account, err := h.geminiCompatService.SelectAccountForAIStudioEndpoints(c.Request.Context(), apiKey.GroupID) + account, err := h.geminiCompatService.SelectAccountForModelsEndpoints(c.Request.Context(), apiKey.GroupID) if err != nil { // 没有 gemini 账户，检查是否有 antigravity 账户可用 hasAntigravity, _ := h.geminiCompatService.HasAntigravityAccounts(c.Request.Context(), apiKey.GroupID) @@ -265,7 +265,7 @@ func (h *GatewayHandler) GeminiV1BetaGetModel(c *gin.Context) { return } - res, err := h.geminiCompatService.ForwardAIStudioGET(c.Request.Context(), account, "/v1beta/models/"+modelName) + res, err := h.geminiCompatService.ForwardGetModel(c.Request.Context(), account, modelName) if err != nil { googleError(c, http.StatusBadGateway, err.Error()) return diff --git a/backend/internal/handler/gemini_v1beta_handler_test.go b/backend/internal/handler/gemini_v1beta_handler_test.go index 8604b07b88..44d66d5c97 100644 --- a/backend/internal/handler/gemini_v1beta_handler_test.go +++ b/backend/internal/handler/gemini_v1beta_handler_test.go @@ -109,7 +109,7 @@ func TestGeminiV1BetaHandler_ListModelsAntigravityFallback(t *testing.T) { expectedBehavior string }{ { - name: "有Gemini账户-调用ForwardAIStudioGET", + name: "有Gemini账户-调用ForwardListModels", hasGeminiAccount: true, hasAntigravity: false, expectedBehavior: "forward_to_upstream", @@ -155,7 +155,7 @@ func TestGeminiV1BetaHandler_GetModelAntigravityFallback(t *testing.T) { expectedBehavior string }{ { - name: "有Gemini账户-调用ForwardAIStudioGET", + name: "有Gemini账户-调用ForwardGetModel", hasGeminiAccount: true, hasAntigravity: false, expectedBehavior: "forward_to_upstream", diff --git a/backend/internal/handler/gemini_vertex_models_test.go b/backend/internal/handler/gemini_vertex_models_test.go new file mode 100644 index 0000000000..eaadf6e03f --- /dev/null +++ b/backend/internal/handler/gemini_vertex_models_test.go @@ -0,0 +1,83 @@ +package handler + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/server/middleware" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +type geminiVertexModelsTokenCache struct{} + +func (geminiVertexModelsTokenCache) GetAccessToken(context.Context, string) (string, error) { + return "ya29.vertex-token", nil +} + +func (geminiVertexModelsTokenCache) SetAccessToken(context.Context, string, string, time.Duration) error { + return nil +} + +func (geminiVertexModelsTokenCache) DeleteAccessToken(context.Context, string) error { + return nil +} + +func (geminiVertexModelsTokenCache) AcquireRefreshLock(context.Context, string, time.Duration) (bool, error) { + return false, nil +} + +func (geminiVertexModelsTokenCache) ReleaseRefreshLock(context.Context, string) error { + return nil +} + +// 分组里只有 Agent Platform (Vertex) 服务账号时，模型列表与模型详情接口返回 Vertex 上的 Gemini 模型。 +func TestGeminiV1BetaModelsEndpoints_VertexServiceAccountOnlyGroup(t *testing.T) { + groupID := int64(8) + repo := &geminiAllowlistAccountRepoStub{gatewayModelsAccountRepoStub: gatewayModelsAccountRepoStub{byGroup: map[int64][]service.Account{groupID: {{ + ID: 18, + Platform: service.PlatformGemini, + Type: service.AccountTypeServiceAccount, + Credentials: map[string]any{ + "location": "global", + "service_account_json": map[string]any{ + "type": "service_account", + "project_id": "proj", + "private_key": "-----BEGIN PRIVATE KEY-----\nabc\n-----END PRIVATE KEY-----\n", + "client_email": "svc@proj.iam.gserviceaccount.com", + }, + }, + }}}}} + upstream := &geminiMixedModelsUpstream{status: http.StatusOK, body: `{"publisherModels":[{"name":"publishers/google/models/gemini-3.8-flash"}]}`} + tokenProvider := service.NewGeminiTokenProvider(nil, geminiVertexModelsTokenCache{}, nil) + h := &GatewayHandler{geminiCompatService: service.NewGeminiMessagesCompatService(repo, nil, nil, nil, tokenProvider, nil, upstream, nil, &config.Config{})} + apiKey := &service.APIKey{GroupID: &groupID, Group: &service.Group{ID: groupID, Platform: service.PlatformGemini}} + + t.Run("list", func(t *testing.T) { + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodGet, "/v1beta/models", nil) + c.Set(string(middleware.ContextKeyAPIKey), apiKey) + + h.GeminiV1BetaListModels(c) + require.Equal(t, http.StatusOK, rec.Code, rec.Body.String()) + require.JSONEq(t, `{"models":[{"name":"models/gemini-3.8-flash","supportedGenerationMethods":["generateContent","streamGenerateContent"]}]}`, rec.Body.String()) + }) + + t.Run("get", func(t *testing.T) { + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodGet, "/v1beta/models/gemini-3.8-flash", nil) + c.Params = gin.Params{{Key: "model", Value: "gemini-3.8-flash"}} + c.Set(string(middleware.ContextKeyAPIKey), apiKey) + + h.GeminiV1BetaGetModel(c) + require.Equal(t, http.StatusOK, rec.Code, rec.Body.String()) + require.JSONEq(t, `{"name":"models/gemini-3.8-flash","supportedGenerationMethods":["generateContent","streamGenerateContent"]}`, rec.Body.String()) + }) +} diff --git a/backend/internal/service/gemini_messages_compat_service.go b/backend/internal/service/gemini_messages_compat_service.go index cc420e2ddd..a54c975cb3 100644 --- a/backend/internal/service/gemini_messages_compat_service.go +++ b/backend/internal/service/gemini_messages_compat_service.go @@ -529,94 +529,6 @@ func (s *GeminiMessagesCompatService) HasAntigravityAccounts(ctx context.Context return len(accounts) > 0, nil } -// SelectAccountForAIStudioEndpoints selects an account that is likely to succeed against -// generativelanguage.googleapis.com (e.g. GET /v1beta/models). -// -// Preference order: -// 1) API key accounts (AI Studio) -// 2) OAuth accounts without project_id (AI Studio OAuth) -// 3) OAuth accounts explicitly marked as ai_studio -// 4) Any remaining Gemini accounts (fallback) -func (s *GeminiMessagesCompatService) SelectAccountForAIStudioEndpoints(ctx context.Context, groupID *int64) (*Account, error) { - accounts, err := s.listSchedulableAccountsOnce(ctx, groupID, PlatformGemini, true) - if err != nil { - return nil, fmt.Errorf("query accounts failed: %w", err) - } - if len(accounts) == 0 { - return nil, errors.New("no available Gemini accounts") - } - - rank := func(a *Account) int { - if a == nil { - return 999 - } - switch a.Type { - case AccountTypeAPIKey: - if strings.TrimSpace(a.GetCredential("api_key")) != "" { - return 0 - } - return 9 - case AccountTypeOAuth: - if strings.TrimSpace(a.GetCredential("project_id")) == "" { - return 1 - } - if strings.TrimSpace(a.GetCredential("oauth_type")) == "ai_studio" { - return 2 - } - // Code Assist OAuth tokens often lack AI Studio scopes for models listing. - return 3 - case AccountTypeServiceAccount: - // Vertex service accounts use aiplatform.googleapis.com, not the AI Studio - // endpoint (generativelanguage.googleapis.com), so they cannot serve these requests. - return 999 - default: - return 10 - } - } - - var selected *Account - for i := range accounts { - acc := &accounts[i] - if selected == nil { - selected = acc - continue - } - - r1, r2 := rank(acc), rank(selected) - if r1 < r2 { - selected = acc - continue - } - if r1 > r2 { - continue - } - - if acc.Priority < selected.Priority { - selected = acc - } else if acc.Priority == selected.Priority { - switch { - case acc.LastUsedAt == nil && selected.LastUsedAt != nil: - selected = acc - case acc.LastUsedAt != nil && selected.LastUsedAt == nil: - // keep selected - case acc.LastUsedAt == nil && selected.LastUsedAt == nil: - if acc.Type == AccountTypeOAuth && selected.Type != AccountTypeOAuth { - selected = acc - } - default: - if acc.LastUsedAt.Before(*selected.LastUsedAt) { - selected = acc - } - } - } - } - - if selected == nil { - return nil, errors.New("no available Gemini accounts") - } - return s.hydrateSelectedAccount(ctx, selected) -} - func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Context, account *Account, body []byte) (*ForwardResult, error) { beginUpstreamResponseModelObservation(c) beginGeminiImageOutputObservation(c) @@ -2916,78 +2828,6 @@ func (s *GeminiMessagesCompatService) handleNativeStreamingResponse(c *gin.Conte return &geminiNativeStreamResult{usage: usage, firstTokenMs: firstTokenMs}, nil } -// ForwardAIStudioGET forwards a GET request to AI Studio (generativelanguage.googleapis.com) for -// endpoints like /v1beta/models and /v1beta/models/{model}. -// -// This is used to support Gemini SDKs that call models listing endpoints before generation. -func (s *GeminiMessagesCompatService) ForwardAIStudioGET(ctx context.Context, account *Account, path string) (*UpstreamHTTPResult, error) { - if account == nil { - return nil, errors.New("account is nil") - } - // path 会被直接拼到上游 base URL 后面，因此按路径护栏逐片段校验， - // 见 upstream_path_guard.go。 - sanitizedPath, ok := sanitizedUpstreamPathSuffix(path) - if !ok || sanitizedPath == "" { - return nil, errors.New("invalid path") - } - path = sanitizedPath - - baseURL := account.GetGeminiBaseURL(geminicli.AIStudioBaseURL) - normalizedBaseURL, err := s.validateUpstreamBaseURL(baseURL) - if err != nil { - return nil, err - } - fullURL := strings.TrimRight(normalizedBaseURL, "/") + path - - var proxyURL string - if account.ProxyID != nil && account.Proxy != nil { - proxyURL = account.Proxy.URL() - } - - req, err := http.NewRequestWithContext(ctx, http.MethodGet, fullURL, nil) - if err != nil { - return nil, err - } - - switch account.Type { - case AccountTypeAPIKey: - apiKey := strings.TrimSpace(account.GetCredential("api_key")) - if apiKey == "" { - return nil, errors.New("gemini api_key not configured") - } - req.Header.Set("x-goog-api-key", apiKey) - case AccountTypeOAuth: - if s.tokenProvider == nil { - return nil, errors.New("gemini token provider not configured") - } - accessToken, err := s.tokenProvider.GetAccessToken(ctx, account) - if err != nil { - return nil, err - } - req.Header.Set("Authorization", "Bearer "+accessToken) - default: - return nil, fmt.Errorf("unsupported account type: %s", account.Type) - } - - resp, err := s.httpUpstream.Do(req, proxyURL, account.ID, account.Concurrency) - if err != nil { - return nil, err - } - defer func() { _ = resp.Body.Close() }() - - body, _ := io.ReadAll(io.LimitReader(resp.Body, 8<<20)) - wwwAuthenticate := resp.Header.Get("Www-Authenticate") - filteredHeaders := responseheaders.FilterHeaders(resp.Header, s.responseHeaderFilter) - if wwwAuthenticate != "" { - filteredHeaders.Set("Www-Authenticate", wwwAuthenticate) - } - return &UpstreamHTTPResult{ - StatusCode: resp.StatusCode, - Headers: filteredHeaders, - Body: body, - }, nil -} - // unwrapGeminiResponse 解包 Gemini OAuth 响应中的 response 字段 // 使用 gjson 零拷贝提取，避免完整 Unmarshal+Marshal func unwrapGeminiResponse(raw []byte) ([]byte, error) { diff --git a/backend/internal/service/gemini_models_endpoints.go b/backend/internal/service/gemini_models_endpoints.go new file mode 100644 index 0000000000..02f3b985a1 --- /dev/null +++ b/backend/internal/service/gemini_models_endpoints.go @@ -0,0 +1,302 @@ +package service + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "slices" + "strings" + + "github.com/Wei-Shaw/sub2api/internal/pkg/gemini" + "github.com/Wei-Shaw/sub2api/internal/pkg/geminicli" + "github.com/Wei-Shaw/sub2api/internal/pkg/googleapi" + "github.com/Wei-Shaw/sub2api/internal/util/responseheaders" +) + +// Gemini 分组对外提供 AI Studio 兼容协议（GET /v1beta/models、/v1beta/models/{model}）， +// 分组内的账号可以是 AI Studio 兼容上游（API Key / OAuth），也可以是 Agent Platform (Vertex)。 +// 本文件把不同账号类型的模型接口统一成 AI Studio 格式的响应。 + +// errGeminiModelsUpstreamStatus 在 Vertex 分页读取遇到非 2xx 时中止分页，原始响应由调用方另行保存。 +var errGeminiModelsUpstreamStatus = errors.New("upstream model list returned non-2xx status") + +// SelectAccountForModelsEndpoints 为模型列表与模型详情接口选择账号。 +// +// 优先级： +// 1. 配置了 api_key 的 API Key 账号 +// 2. 不带 project_id 的 OAuth 账号（AI Studio OAuth） +// 3. 显式标记为 ai_studio 的 OAuth 账号 +// 4. Code Assist OAuth 账号（令牌常缺 AI Studio 权限范围，多半退回内置列表） +// 5. 其余类型的账号 +// 6. Vertex 服务账号（服务账号 JSON 认证，读取 Google 发布方模型列表） +// 7. Vertex API Key 账号（列表接口不接受 API Key，只能返回内置列表） +// +// Vertex 账号排在所有 AI Studio 兼容账号之后，混合分组始终由 AI Studio 兼容账号应答。 +// 同级账号按 priority、最近使用时间挑选。 +func (s *GeminiMessagesCompatService) SelectAccountForModelsEndpoints(ctx context.Context, groupID *int64) (*Account, error) { + accounts, err := s.listSchedulableAccountsOnce(ctx, groupID, PlatformGemini, true) + if err != nil { + return nil, fmt.Errorf("query accounts failed: %w", err) + } + if len(accounts) == 0 { + return nil, errors.New("no available Gemini accounts") + } + + var selected *Account + for i := range accounts { + acc := &accounts[i] + if selected == nil { + selected = acc + continue + } + + r1, r2 := modelsEndpointAccountRank(acc), modelsEndpointAccountRank(selected) + if r1 < r2 { + selected = acc + continue + } + if r1 > r2 { + continue + } + + if acc.Priority < selected.Priority { + selected = acc + } else if acc.Priority == selected.Priority { + switch { + case acc.LastUsedAt == nil && selected.LastUsedAt != nil: + selected = acc + case acc.LastUsedAt != nil && selected.LastUsedAt == nil: + // keep selected + case acc.LastUsedAt == nil && selected.LastUsedAt == nil: + if acc.Type == AccountTypeOAuth && selected.Type != AccountTypeOAuth { + selected = acc + } + default: + if acc.LastUsedAt.Before(*selected.LastUsedAt) { + selected = acc + } + } + } + } + + if selected == nil { + return nil, errors.New("no available Gemini accounts") + } + return s.hydrateSelectedAccount(ctx, selected) +} + +// modelsEndpointAccountRank 返回账号在模型接口选号中的优先级，数值越小越优先，顺序见 SelectAccountForModelsEndpoints。 +func modelsEndpointAccountRank(a *Account) int { + if a == nil { + return 999 + } + switch a.Type { + case AccountTypeAPIKey: + if strings.TrimSpace(a.GetCredential("api_key")) != "" { + return 0 + } + return 9 + case AccountTypeOAuth: + if strings.TrimSpace(a.GetCredential("project_id")) == "" { + return 1 + } + if strings.TrimSpace(a.GetCredential("oauth_type")) == "ai_studio" { + return 2 + } + return 3 + case AccountTypeServiceAccount: + if a.IsVertexAPIKey() { + return 12 + } + return 11 + default: + return 10 + } +} + +// ForwardListModels 读取账号上游的模型列表，返回 AI Studio 格式的响应（GET /v1beta/models）。 +// AI Studio 兼容账号原样返回上游响应；Vertex 服务账号读取 Google 发布方模型列表后转换格式， +// 上游非 2xx 时原样返回该响应；Vertex API Key 账号返回内置列表。 +func (s *GeminiMessagesCompatService) ForwardListModels(ctx context.Context, account *Account) (*UpstreamHTTPResult, error) { + if account == nil { + return nil, errors.New("account is nil") + } + if !account.IsVertexServiceAccount() { + return s.forwardGeminiModelsGET(ctx, account, "/v1beta/models") + } + if account.IsVertexAPIKey() { + return geminiModelsJSONResult(http.StatusOK, gemini.FallbackModelsList()) + } + + modelIDs, failure, err := s.listVertexGeminiModels(ctx, account, account.VertexLocation("")) + if err != nil { + return nil, err + } + if failure != nil { + return failure, nil + } + models := make([]gemini.Model, 0, len(modelIDs)) + for _, modelID := range modelIDs { + models = append(models, gemini.FallbackModel(modelID)) + } + return geminiModelsJSONResult(http.StatusOK, gemini.ModelsListResponse{Models: models}) +} + +// ForwardGetModel 读取账号上游的单个模型信息，返回 AI Studio 格式的响应（GET /v1beta/models/{model}）。 +// Vertex 服务账号在该模型生成请求所用区域的模型列表中查找，找不到时返回 404；Vertex API Key 账号返回内置信息。 +func (s *GeminiMessagesCompatService) ForwardGetModel(ctx context.Context, account *Account, model string) (*UpstreamHTTPResult, error) { + if account == nil { + return nil, errors.New("account is nil") + } + if !account.IsVertexServiceAccount() { + return s.forwardGeminiModelsGET(ctx, account, "/v1beta/models/"+model) + } + if account.IsVertexAPIKey() { + return geminiModelsJSONResult(http.StatusOK, gemini.FallbackModel(model)) + } + + modelIDs, failure, err := s.listVertexGeminiModels(ctx, account, account.VertexLocation(model)) + if err != nil { + return nil, err + } + if failure != nil { + return failure, nil + } + if slices.Contains(modelIDs, model) { + return geminiModelsJSONResult(http.StatusOK, gemini.FallbackModel(model)) + } + return geminiModelsJSONResult(http.StatusNotFound, map[string]any{ + "error": map[string]any{ + "code": http.StatusNotFound, + "message": fmt.Sprintf("models/%s is not found", model), + "status": googleapi.HTTPStatusToGoogleStatus(http.StatusNotFound), + }, + }) +} + +// forwardGeminiModelsGET 向 AI Studio 兼容上游转发模型接口的 GET 请求。 +// path 会被直接拼到上游 base URL 后面，因此按路径护栏逐片段校验，见 upstream_path_guard.go。 +func (s *GeminiMessagesCompatService) forwardGeminiModelsGET(ctx context.Context, account *Account, path string) (*UpstreamHTTPResult, error) { + sanitizedPath, ok := sanitizedUpstreamPathSuffix(path) + if !ok || sanitizedPath == "" { + return nil, errors.New("invalid path") + } + + baseURL := account.GetGeminiBaseURL(geminicli.AIStudioBaseURL) + normalizedBaseURL, err := s.validateUpstreamBaseURL(baseURL) + if err != nil { + return nil, err + } + fullURL := strings.TrimRight(normalizedBaseURL, "/") + sanitizedPath + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, fullURL, nil) + if err != nil { + return nil, err + } + + switch account.Type { + case AccountTypeAPIKey: + apiKey := strings.TrimSpace(account.GetCredential("api_key")) + if apiKey == "" { + return nil, errors.New("gemini api_key not configured") + } + req.Header.Set("x-goog-api-key", apiKey) + case AccountTypeOAuth: + if s.tokenProvider == nil { + return nil, errors.New("gemini token provider not configured") + } + accessToken, err := s.tokenProvider.GetAccessToken(ctx, account) + if err != nil { + return nil, err + } + req.Header.Set("Authorization", "Bearer "+accessToken) + default: + return nil, fmt.Errorf("unsupported account type: %s", account.Type) + } + + return s.doGeminiModelsGET(req, account) +} + +// listVertexGeminiModels 读取 Vertex 服务账号在指定区域可用的 Gemini 模型。 +// 上游返回非 2xx 时 failure 为该页的原始响应，modelIDs 与 err 均为空。 +func (s *GeminiMessagesCompatService) listVertexGeminiModels(ctx context.Context, account *Account, location string) (modelIDs []string, failure *UpstreamHTTPResult, err error) { + if s.tokenProvider == nil { + return nil, nil, errors.New("gemini token provider not configured") + } + accessToken, err := s.tokenProvider.GetAccessToken(ctx, account) + if err != nil { + return nil, nil, err + } + + modelIDs, err = listVertexGeminiPublisherModels(location, func(_ int, listURL string) ([]byte, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, listURL, nil) + if err != nil { + return nil, err + } + req.Header.Set("Accept", "application/json") + req.Header.Set("Authorization", "Bearer "+accessToken) + + res, err := s.doGeminiModelsGET(req, account) + if err != nil { + return nil, err + } + if res.StatusCode < http.StatusOK || res.StatusCode >= http.StatusMultipleChoices { + failure = res + return nil, errGeminiModelsUpstreamStatus + } + return res.Body, nil + }) + if failure != nil { + return nil, failure, nil + } + if err != nil { + return nil, nil, err + } + return modelIDs, nil, nil +} + +// doGeminiModelsGET 发送模型接口请求并读取响应。响应体上限取 gateway.models_list_read_max_bytes，超出时报错；响应头按配置过滤。 +func (s *GeminiMessagesCompatService) doGeminiModelsGET(req *http.Request, account *Account) (*UpstreamHTTPResult, error) { + var proxyURL string + if account.ProxyID != nil && account.Proxy != nil { + proxyURL = account.Proxy.URL() + } + + resp, err := s.httpUpstream.Do(req, proxyURL, account.ID, account.Concurrency) + if err != nil { + return nil, err + } + defer func() { _ = resp.Body.Close() }() + + bodyLimit := resolveModelsListReadLimit(s.cfg) + body, _ := io.ReadAll(io.LimitReader(resp.Body, bodyLimit+1)) + if int64(len(body)) > bodyLimit { + return nil, fmt.Errorf("upstream model list response exceeds %d bytes", bodyLimit) + } + wwwAuthenticate := resp.Header.Get("Www-Authenticate") + filteredHeaders := responseheaders.FilterHeaders(resp.Header, s.responseHeaderFilter) + if wwwAuthenticate != "" { + filteredHeaders.Set("Www-Authenticate", wwwAuthenticate) + } + return &UpstreamHTTPResult{ + StatusCode: resp.StatusCode, + Headers: filteredHeaders, + Body: body, + }, nil +} + +// geminiModelsJSONResult 把网关生成的模型接口响应包装成 UpstreamHTTPResult，交给 handler 统一写出。 +func geminiModelsJSONResult(status int, payload any) (*UpstreamHTTPResult, error) { + body, err := json.Marshal(payload) + if err != nil { + return nil, err + } + return &UpstreamHTTPResult{ + StatusCode: status, + Headers: http.Header{"Content-Type": []string{"application/json; charset=utf-8"}}, + Body: body, + }, nil +} diff --git a/backend/internal/service/gemini_models_endpoints_test.go b/backend/internal/service/gemini_models_endpoints_test.go new file mode 100644 index 0000000000..e33b85d809 --- /dev/null +++ b/backend/internal/service/gemini_models_endpoints_test.go @@ -0,0 +1,208 @@ +//go:build unit + +package service + +import ( + "context" + "encoding/json" + "net/http" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/gemini" + "github.com/stretchr/testify/require" +) + +func newGeminiModelsEndpointsTestService(accounts []Account, upstream HTTPUpstream) *GeminiMessagesCompatService { + repo := &mockAccountRepoForGemini{listByGroupFunc: func(context.Context, int64, []string) ([]Account, error) { + return accounts, nil + }} + return &GeminiMessagesCompatService{ + accountRepo: repo, + httpUpstream: upstream, + cfg: &config.Config{}, + tokenProvider: NewGeminiTokenProvider(nil, &vertexModelsTokenCacheStub{token: "ya29.vertex-token"}, nil), + } +} + +func decodeGeminiModelNames(t *testing.T, body []byte) []string { + t.Helper() + var got gemini.ModelsListResponse + require.NoError(t, json.Unmarshal(body, &got)) + names := make([]string, 0, len(got.Models)) + for _, model := range got.Models { + require.Contains(t, model.SupportedGenerationMethods, "generateContent") + names = append(names, model.Name) + } + return names +} + +func TestSelectAccountForModelsEndpoints_VertexServiceAccountOnlyGroup(t *testing.T) { + groupID := int64(8) + vertex := *vertexModelSyncTestAccount(nil) + vertex.ID = 18 + svc := newGeminiModelsEndpointsTestService([]Account{vertex}, nil) + + account, err := svc.SelectAccountForModelsEndpoints(context.Background(), &groupID) + require.NoError(t, err) + require.Equal(t, int64(18), account.ID) +} + +func TestModelsEndpointAccountRank_Order(t *testing.T) { + ordered := []*Account{ + {Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "k"}}, + {Platform: PlatformGemini, Type: AccountTypeOAuth, Credentials: map[string]any{}}, + {Platform: PlatformGemini, Type: AccountTypeOAuth, Credentials: map[string]any{"project_id": "p", "oauth_type": "ai_studio"}}, + {Platform: PlatformGemini, Type: AccountTypeOAuth, Credentials: map[string]any{"project_id": "p", "oauth_type": "code_assist"}}, + {Platform: PlatformGemini, Type: AccountTypeUpstream, Credentials: map[string]any{"api_key": "k"}}, + vertexModelSyncTestAccount(nil), + {Platform: PlatformGemini, Type: AccountTypeServiceAccount, Credentials: map[string]any{"auth_mode": "apikey", "api_key": "AQ.k"}}, + } + for i := 1; i < len(ordered); i++ { + require.Less(t, modelsEndpointAccountRank(ordered[i-1]), modelsEndpointAccountRank(ordered[i]), "position %d", i) + } +} + +func TestSelectAccountForModelsEndpoints_PrefersAPIKeyOverVertex(t *testing.T) { + groupID := int64(8) + vertex := *vertexModelSyncTestAccount(nil) + vertex.ID = 18 + apiKey := Account{ID: 19, Platform: PlatformGemini, Type: AccountTypeAPIKey, Priority: 9, Credentials: map[string]any{"api_key": "k"}} + svc := newGeminiModelsEndpointsTestService([]Account{vertex, apiKey}, nil) + + account, err := svc.SelectAccountForModelsEndpoints(context.Background(), &groupID) + require.NoError(t, err) + require.Equal(t, int64(19), account.ID) +} + +func TestForwardListModels_VertexServiceAccountReturnsAIStudioFormat(t *testing.T) { + upstream := &httpUpstreamRecorder{responses: []*http.Response{ + vertexModelsJSONResponse(http.StatusOK, `{"publisherModels":[ + {"name":"publishers/google/models/gemini-3.8-flash"}, + {"name":"publishers/google/models/gemma4"} + ],"nextPageToken":"page-2"}`), + vertexModelsJSONResponse(http.StatusOK, `{"publisherModels":[{"name":"publishers/google/models/gemini-2.5-flash"}]}`), + }} + svc := newGeminiModelsEndpointsTestService(nil, upstream) + + res, err := svc.ForwardListModels(context.Background(), vertexModelSyncTestAccount(map[string]any{"location": "europe-west4"})) + require.NoError(t, err) + require.Equal(t, http.StatusOK, res.StatusCode) + require.Equal(t, []string{"models/gemini-2.5-flash", "models/gemini-3.8-flash"}, decodeGeminiModelNames(t, res.Body)) + + require.Len(t, upstream.requests, 2) + require.Equal(t, "https://europe-west4-aiplatform.googleapis.com/v1beta1/publishers/google/models?pageSize=300", upstream.requests[0].URL.String()) + require.Equal(t, "https://europe-west4-aiplatform.googleapis.com/v1beta1/publishers/google/models?pageSize=300&pageToken=page-2", upstream.requests[1].URL.String()) + for _, req := range upstream.requests { + require.Equal(t, "Bearer ya29.vertex-token", req.Header.Get("Authorization")) + } +} + +func TestForwardListModels_VertexUpstreamErrorPassesThrough(t *testing.T) { + body := `{"error":{"code":403,"status":"PERMISSION_DENIED"}}` + upstream := &httpUpstreamRecorder{resp: vertexModelsJSONResponse(http.StatusForbidden, body)} + svc := newGeminiModelsEndpointsTestService(nil, upstream) + + res, err := svc.ForwardListModels(context.Background(), vertexModelSyncTestAccount(map[string]any{"location": "global"})) + require.NoError(t, err) + require.Equal(t, http.StatusForbidden, res.StatusCode) + require.JSONEq(t, body, string(res.Body)) +} + +func TestForwardListModels_VertexAPIKeyReturnsBuiltinList(t *testing.T) { + upstream := &httpUpstreamRecorder{} + svc := newGeminiModelsEndpointsTestService(nil, upstream) + account := &Account{ID: 174, Platform: PlatformGemini, Type: AccountTypeServiceAccount, Credentials: map[string]any{ + "auth_mode": "apikey", + "api_key": "AQ.test-key", + }} + + res, err := svc.ForwardListModels(context.Background(), account) + require.NoError(t, err) + require.Equal(t, http.StatusOK, res.StatusCode) + expected, err := json.Marshal(gemini.FallbackModelsList()) + require.NoError(t, err) + require.JSONEq(t, string(expected), string(res.Body)) + require.Empty(t, upstream.requests) +} + +func TestForwardListModels_APIKeyForwardsToBaseURL(t *testing.T) { + upstream := &httpUpstreamRecorder{resp: vertexModelsJSONResponse(http.StatusOK, `{"models":[{"name":"models/gemini-native"}]}`)} + svc := newGeminiModelsEndpointsTestService(nil, upstream) + account := &Account{ID: 261, Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: map[string]any{ + "api_key": "sk-pool", + "base_url": "https://relay.example.com", + }} + + res, err := svc.ForwardListModels(context.Background(), account) + require.NoError(t, err) + require.Equal(t, http.StatusOK, res.StatusCode) + require.JSONEq(t, `{"models":[{"name":"models/gemini-native"}]}`, string(res.Body)) + require.Equal(t, "https://relay.example.com/v1beta/models", upstream.lastReq.URL.String()) + require.Equal(t, "sk-pool", upstream.lastReq.Header.Get("x-goog-api-key")) +} + +func TestForwardGetModel_VertexServiceAccountLooksUpModelList(t *testing.T) { + listBody := `{"publisherModels":[{"name":"publishers/google/models/gemini-3.8-flash"}]}` + + t.Run("listed model", func(t *testing.T) { + upstream := &httpUpstreamRecorder{resp: vertexModelsJSONResponse(http.StatusOK, listBody)} + svc := newGeminiModelsEndpointsTestService(nil, upstream) + + res, err := svc.ForwardGetModel(context.Background(), vertexModelSyncTestAccount(nil), "gemini-3.8-flash") + require.NoError(t, err) + require.Equal(t, http.StatusOK, res.StatusCode) + var model gemini.Model + require.NoError(t, json.Unmarshal(res.Body, &model)) + require.Equal(t, "models/gemini-3.8-flash", model.Name) + }) + + t.Run("unlisted model", func(t *testing.T) { + upstream := &httpUpstreamRecorder{resp: vertexModelsJSONResponse(http.StatusOK, listBody)} + svc := newGeminiModelsEndpointsTestService(nil, upstream) + + res, err := svc.ForwardGetModel(context.Background(), vertexModelSyncTestAccount(nil), "gemini-unknown") + require.NoError(t, err) + require.Equal(t, http.StatusNotFound, res.StatusCode) + require.JSONEq(t, `{"error":{"code":404,"message":"models/gemini-unknown is not found","status":"NOT_FOUND"}}`, string(res.Body)) + }) +} + +func TestForwardGetModel_VertexUsesModelLocation(t *testing.T) { + upstream := &httpUpstreamRecorder{resp: vertexModelsJSONResponse(http.StatusOK, `{"publisherModels":[{"name":"publishers/google/models/gemini-3.8-flash"}]}`)} + svc := newGeminiModelsEndpointsTestService(nil, upstream) + account := vertexModelSyncTestAccount(map[string]any{ + "location": "europe-west4", + "vertex_model_locations": map[string]any{"gemini-3.8-flash": "global"}, + }) + + res, err := svc.ForwardGetModel(context.Background(), account, "gemini-3.8-flash") + require.NoError(t, err) + require.Equal(t, http.StatusOK, res.StatusCode) + require.Equal(t, "https://aiplatform.googleapis.com/v1beta1/publishers/google/models?pageSize=300", upstream.lastReq.URL.String()) +} + +func TestForwardListModels_VertexMidPaginationErrorPassesThrough(t *testing.T) { + body := `{"error":{"code":503,"status":"UNAVAILABLE"}}` + upstream := &httpUpstreamRecorder{responses: []*http.Response{ + vertexModelsJSONResponse(http.StatusOK, `{"publisherModels":[{"name":"publishers/google/models/gemini-3.8-flash"}],"nextPageToken":"page-2"}`), + vertexModelsJSONResponse(http.StatusServiceUnavailable, body), + }} + svc := newGeminiModelsEndpointsTestService(nil, upstream) + + res, err := svc.ForwardListModels(context.Background(), vertexModelSyncTestAccount(nil)) + require.NoError(t, err) + require.Equal(t, http.StatusServiceUnavailable, res.StatusCode) + require.JSONEq(t, body, string(res.Body)) + require.Len(t, upstream.requests, 2) +} + +func TestForwardListModels_ResponseOverReadLimitFails(t *testing.T) { + upstream := &httpUpstreamRecorder{resp: vertexModelsJSONResponse(http.StatusOK, `{"models":[{"name":"models/gemini-native"}]}`)} + svc := newGeminiModelsEndpointsTestService(nil, upstream) + svc.cfg.Gateway.ModelsListReadMaxBytes = 16 + account := &Account{ID: 261, Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "sk-pool"}} + + _, err := svc.ForwardListModels(context.Background(), account) + require.ErrorContains(t, err, "exceeds 16 bytes") +} diff --git a/backend/internal/service/upstream_models.go b/backend/internal/service/upstream_models.go index c7a392c550..447f2fdb98 100644 --- a/backend/internal/service/upstream_models.go +++ b/backend/internal/service/upstream_models.go @@ -804,9 +804,7 @@ func (s *AccountTestService) readUpstreamModelsResponse(req *http.Request, accou return body, nil } -// fetchVertexGeminiUpstreamModels 按账号默认区域分页读取 Vertex 的 Google 发布方模型列表，只保留 gemini- 前缀的模型。 -// 区域端点还会列出 Gemma、Veo、AutoML 等非 Gemini 的模型花园条目；global 端点只含 Gemini。 -// 列表接口不接受 API Key。 +// fetchVertexGeminiUpstreamModels 按账号默认区域读取 Vertex 上可用的 Gemini 模型。列表接口不接受 API Key。 func (s *AccountTestService) fetchVertexGeminiUpstreamModels(ctx context.Context, account *Account) ([]string, error) { if account.IsVertexAPIKey() { return nil, newUpstreamModelSyncUnsupportedError("Vertex API key accounts do not support upstream model listing", nil) @@ -830,13 +828,7 @@ func (s *AccountTestService) fetchVertexGeminiUpstreamModels(ctx context.Context return nil, newUpstreamModelSyncConfigError("No Vertex access token is available", nil) } - models := make([]string, 0) - pageToken := "" - for page := 0; page < vertexPublisherModelsMaxPages; page++ { - listURL, err := buildVertexPublisherModelsURL(location, pageToken) - if err != nil { - return nil, newUpstreamModelSyncConfigError("Invalid Vertex location", err) - } + return listVertexGeminiPublisherModels(location, func(page int, listURL string) ([]byte, error) { req, err := http.NewRequestWithContext(ctx, http.MethodGet, listURL, nil) if err != nil { return nil, newUpstreamModelSyncConfigError("Invalid Vertex model list URL", err) @@ -853,6 +845,25 @@ func (s *AccountTestService) fetchVertexGeminiUpstreamModels(ctx context.Context } return nil, err } + return body, nil + }) +} + +// listVertexGeminiPublisherModels 按区域分页读取 Google 发布方模型列表，只保留 gemini- 前缀的模型（去重排序）。 +// 区域端点还会列出 Gemma、Veo、AutoML 等非 Gemini 的模型花园条目；global 端点只含 Gemini。 +// fetchPage 负责发送请求并处理非 2xx 响应，返回该页的响应体；它返回的错误原样透出。 +func listVertexGeminiPublisherModels(location string, fetchPage func(page int, listURL string) ([]byte, error)) ([]string, error) { + models := make([]string, 0) + pageToken := "" + for page := 0; page < vertexPublisherModelsMaxPages; page++ { + listURL, err := buildVertexPublisherModelsURL(location, pageToken) + if err != nil { + return nil, newUpstreamModelSyncConfigError("Invalid Vertex location", err) + } + body, err := fetchPage(page, listURL) + if err != nil { + return nil, err + } var response struct { PublisherModels []struct { Name string `json:"name"` From 26b31a83e4190a3b38bce112fc597defbfb9515b Mon Sep 17 00:00:00 2001 From: feeeei Date: Mon, 12 Oct 2026 17:14:56 +0800 Subject: [PATCH 14/15] =?UTF-8?q?fix(gemini):=20=E6=98=BE=E5=BC=8F?= =?UTF-8?q?=E7=BC=93=E5=AD=98=E4=BF=AE=E6=94=B9=E4=B8=B2=E8=A1=8C=E5=8C=96?= =?UTF-8?q?=E3=80=81=E5=A4=B1=E8=B4=A5=E5=9B=9E=E6=BB=9A=EF=BC=8C=E7=A1=AE?= =?UTF-8?q?=E5=AE=9A=E9=87=91=E9=A2=9D=E8=B4=B9=E7=94=A8=E6=8C=89=20Key=20?= =?UTF-8?q?/=20=E8=AE=A2=E9=98=85=E9=A2=84=E7=95=99?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 同一缓存的并发 PATCH 可能以与上游提交不同的顺序写回本地到期时间； 确定金额的存储费只预留用户余额，并发续期可越过 API Key 总额度与订阅额度； 上游已修改而本地写回失败时没有可恢复的状态。 - 新增迁移 245：缓存记录增加修改意图（mutation_token / mutation_lease_until）。 PATCH 在调用上游前登记意图并取得 2 分钟租约，条件为没有进行中的修改且到期时间 未变；上游修改、本地写回与计费均由持有者完成，并发的修改返回 409 ABORTED - 上游结果未知、本地写回失败或进程退出时修改意图保留且不计费；租约过期后由下一次 修改或后台每分钟的恢复流程把上游有效期回滚到修改前的到期时间（已到期则删除上游 缓存与本地记录）。上游明确拒绝或请求未发出时直接清除意图 - 修改 / 删除上游缓存的单次调用超时改为 30 秒；请求发出前的失败包装为 ErrGeminiCachedContentRequestNotSent - CheckChargeAffordable 与 ReserveInflightCharge 合并为 ReserveFixedCharge：在途预留 新增订阅、API Key 两个维度（Redis billing:inflight_scope:*），先登记本笔预留再读取 库中权威用量（Key quota_used 与限速窗口、订阅日 / 周 / 月用量）叠加同维度在途合计 校验；预留随计费任务交接，落库后释放 - GetRateLimitData 一并读取 quota_used --- backend/cmd/server/wire.go | 7 + backend/cmd/server/wire_gen.go | 11 +- backend/cmd/server/wire_gen_test.go | 1 + .../gemini_cached_contents_billing_test.go | 409 ++++++++++++++++-- .../handler/gemini_cached_contents_handler.go | 58 ++- .../gemini_cached_contents_handler_test.go | 20 +- backend/internal/repository/api_key_repo.go | 4 +- .../api_key_repo_integration_test.go | 4 + .../repository/billing_inflight_cache.go | 64 +++ .../repository/billing_inflight_cache_test.go | 69 +++ .../repository/gemini_cached_content_repo.go | 93 +++- ...ni_cached_content_repo_integration_test.go | 50 ++- backend/internal/service/api_key_service.go | 2 + .../internal/service/billing_cache_service.go | 114 +++-- .../service/billing_inflight_reservation.go | 68 ++- .../internal/service/gemini_cached_content.go | 264 +++++++++-- .../gemini_cached_content_hardening_test.go | 151 ++++++- .../service/gemini_cached_content_test.go | 23 +- backend/internal/service/wire.go | 9 +- ..._gemini_cached_contents_mutation_lease.sql | 10 + 20 files changed, 1259 insertions(+), 172 deletions(-) create mode 100644 backend/migrations/245_gemini_cached_contents_mutation_lease.sql diff --git a/backend/cmd/server/wire.go b/backend/cmd/server/wire.go index 06674251af..6834ad78df 100644 --- a/backend/cmd/server/wire.go +++ b/backend/cmd/server/wire.go @@ -104,6 +104,7 @@ func provideCleanup( subscriptionExpiry *service.SubscriptionExpiryService, usageCleanup *service.UsageCleanupService, idempotencyCleanup *service.IdempotencyCleanupService, + geminiCachedContent *service.GeminiCachedContentService, batchImageCleanup *service.BatchImageCleanupService, batchImageWorker *service.BatchImageWorkerRuntime, pricing *service.PricingService, @@ -244,6 +245,12 @@ func provideCleanup( } return nil }}, + {"GeminiCachedContentService", func() error { + if geminiCachedContent != nil { + geminiCachedContent.Stop() + } + return nil + }}, {"BatchImageCleanupService", func() error { if batchImageCleanup != nil { batchImageCleanup.Stop() diff --git a/backend/cmd/server/wire_gen.go b/backend/cmd/server/wire_gen.go index 2c2f0cfbc9..53cef32052 100644 --- a/backend/cmd/server/wire_gen.go +++ b/backend/cmd/server/wire_gen.go @@ -297,7 +297,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { legacyEngine := securityaudit.NewLegacyModerationAdapter(contentModerationService) coordinator := securityaudit.NewCoordinator(legacyEngine, promptService) geminiCachedContentRepository := repository.NewGeminiCachedContentRepository(db) - geminiCachedContentService := service.NewGeminiCachedContentService(geminiCachedContentRepository, geminiMessagesCompatService, configConfig) + geminiCachedContentService := service.ProvideGeminiCachedContentService(geminiCachedContentRepository, geminiMessagesCompatService, configConfig) gatewayHandler := handler.ProvideGatewayHandler(gatewayService, openAIGatewayService, geminiMessagesCompatService, antigravityGatewayService, userService, concurrencyService, billingCacheService, usageService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, userMessageQueueService, configConfig, settingService, coordinator, geminiCachedContentService) openAIGatewayHandler := handler.ProvideOpenAIGatewayHandler(openAIGatewayService, pluginManager, concurrencyService, billingCacheService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, opsService, grokQuotaService, configConfig, coordinator, compositeRouteResolver) handlerSettingHandler := handler.ProvideSettingHandler(settingService, buildInfo, notificationEmailService) @@ -355,7 +355,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { channelMonitorRunner := service.ProvideChannelMonitorRunner(channelMonitorService, settingService, channelMonitorQuotaFetcher) channelMonitorV2Aggregator := service.ProvideChannelMonitorV2Aggregator(channelMonitorV2Repository, db, settingService) userPlatformQuotaUsageFlusher := service.ProvideUserPlatformQuotaUsageFlusher(configConfig, billingCache, serviceUserPlatformQuotaRepository, timingWheelService) - v := provideCleanup(client, redisClient, opsMetricsCollector, opsAggregationService, opsAlertEvaluatorService, opsCleanupService, opsScheduledReportService, opsSystemLogSink, opsService, opsIngressRejectAggregator, apiKeyService, authCacheInvalidationWorker, schedulerSnapshotService, tokenRefreshService, accountExpiryService, cnProviderBalanceCheckService, openAICodexVersionSyncService, claudeCodeVersionSyncService, proxyExpiryService, subscriptionExpiryService, usageCleanupService, idempotencyCleanupService, batchImageCleanupService, batchImageWorkerRuntime, pricingService, emailQueueService, billingCacheService, usageRecordWorkerPool, subscriptionService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, grokOAuthService, openAIGatewayService, scheduledTestRunnerService, backupService, paymentOrderExpiryService, channelMonitorRunner, channelMonitorV2Aggregator, userPlatformQuotaUsageFlusher, upstreamBillingProbeService, ollamaCloudUsageService, openCodeGoUsageService, auditLogService, openAIQuotaAutoResetService, promptService, pluginManager) + v := provideCleanup(client, redisClient, opsMetricsCollector, opsAggregationService, opsAlertEvaluatorService, opsCleanupService, opsScheduledReportService, opsSystemLogSink, opsService, opsIngressRejectAggregator, apiKeyService, authCacheInvalidationWorker, schedulerSnapshotService, tokenRefreshService, accountExpiryService, cnProviderBalanceCheckService, openAICodexVersionSyncService, claudeCodeVersionSyncService, proxyExpiryService, subscriptionExpiryService, usageCleanupService, idempotencyCleanupService, geminiCachedContentService, batchImageCleanupService, batchImageWorkerRuntime, pricingService, emailQueueService, billingCacheService, usageRecordWorkerPool, subscriptionService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, grokOAuthService, openAIGatewayService, scheduledTestRunnerService, backupService, paymentOrderExpiryService, channelMonitorRunner, channelMonitorV2Aggregator, userPlatformQuotaUsageFlusher, upstreamBillingProbeService, ollamaCloudUsageService, openCodeGoUsageService, auditLogService, openAIQuotaAutoResetService, promptService, pluginManager) application := &Application{ Server: httpServer, PromptAudit: promptService, @@ -415,6 +415,7 @@ func provideCleanup( subscriptionExpiry *service.SubscriptionExpiryService, usageCleanup *service.UsageCleanupService, idempotencyCleanup *service.IdempotencyCleanupService, + geminiCachedContent *service.GeminiCachedContentService, batchImageCleanup *service.BatchImageCleanupService, batchImageWorker *service.BatchImageWorkerRuntime, pricing *service.PricingService, @@ -554,6 +555,12 @@ func provideCleanup( } return nil }}, + {"GeminiCachedContentService", func() error { + if geminiCachedContent != nil { + geminiCachedContent.Stop() + } + return nil + }}, {"BatchImageCleanupService", func() error { if batchImageCleanup != nil { batchImageCleanup.Stop() diff --git a/backend/cmd/server/wire_gen_test.go b/backend/cmd/server/wire_gen_test.go index d14ca19db6..bb56c6735f 100644 --- a/backend/cmd/server/wire_gen_test.go +++ b/backend/cmd/server/wire_gen_test.go @@ -74,6 +74,7 @@ func TestProvideCleanup_WithMinimalDependencies_NoPanic(t *testing.T) { subscriptionExpirySvc, &service.UsageCleanupService{}, idempotencyCleanupSvc, + service.NewGeminiCachedContentService(nil, nil, cfg), &service.BatchImageCleanupService{}, nil, // batchImageWorker pricingSvc, diff --git a/backend/internal/handler/gemini_cached_contents_billing_test.go b/backend/internal/handler/gemini_cached_contents_billing_test.go index a1fd414519..fdf0faea93 100644 --- a/backend/internal/handler/gemini_cached_contents_billing_test.go +++ b/backend/internal/handler/gemini_cached_contents_billing_test.go @@ -4,6 +4,7 @@ package handler import ( "context" + "encoding/json" "errors" "fmt" "io" @@ -11,6 +12,7 @@ import ( "net/http/httptest" "os" "path/filepath" + "strconv" "strings" "sync" "testing" @@ -32,13 +34,20 @@ const ( ) type cachedBillingRepo struct { - mu sync.Mutex - records map[string]*service.GeminiCachedContent - failUpdate bool + mu sync.Mutex + records map[string]*service.GeminiCachedContent + // mutations 按记录 id 保存修改意图。 + mutations map[int64]cachedBillingMutation + failCommit bool nextID int64 deleted []int64 } +type cachedBillingMutation struct { + token string + leaseUntil time.Time +} + func (r *cachedBillingRepo) Create(_ context.Context, record *service.GeminiCachedContent) error { r.mu.Lock() defer r.mu.Unlock() @@ -64,20 +73,83 @@ func (r *cachedBillingRepo) ListForOwner(context.Context, int64, int64, time.Tim return nil, nil } -func (r *cachedBillingRepo) UpdateExpireTime(_ context.Context, id int64, expireTime time.Time) (time.Time, error) { - r.mu.Lock() - defer r.mu.Unlock() - if r.failUpdate { - return time.Time{}, errors.New("db unavailable") +// liveLocked 返回未删除的记录；调用方须持有 mu。 +func (r *cachedBillingRepo) liveLocked(id int64) *service.GeminiCachedContent { + for _, deleted := range r.deleted { + if deleted == id { + return nil + } } for _, record := range r.records { if record.ID == id { - previous := record.ExpireTime - record.ExpireTime = expireTime - return previous, nil + return record } } - return time.Time{}, service.ErrGeminiCachedContentNotFound + return nil +} + +func (r *cachedBillingRepo) BeginMutation(_ context.Context, id int64, baseline time.Time, token string, lease time.Duration) (bool, error) { + r.mu.Lock() + defer r.mu.Unlock() + record := r.liveLocked(id) + if record == nil { + return false, service.ErrGeminiCachedContentNotFound + } + if _, pending := r.mutations[id]; pending || !record.ExpireTime.Equal(baseline) { + return false, nil + } + r.mutations[id] = cachedBillingMutation{token: token, leaseUntil: time.Now().Add(lease)} + return true, nil +} + +func (r *cachedBillingRepo) TakeOverStaleMutation(_ context.Context, id int64, token string, lease time.Duration) (*service.GeminiCachedContent, error) { + r.mu.Lock() + defer r.mu.Unlock() + record := r.liveLocked(id) + mutation, pending := r.mutations[id] + if record == nil || !pending || mutation.leaseUntil.After(time.Now()) { + return nil, nil + } + r.mutations[id] = cachedBillingMutation{token: token, leaseUntil: time.Now().Add(lease)} + cp := *record + return &cp, nil +} + +func (r *cachedBillingRepo) CommitMutation(_ context.Context, id int64, token string, expireTime time.Time) error { + r.mu.Lock() + defer r.mu.Unlock() + if r.failCommit { + return errors.New("db unavailable") + } + record := r.liveLocked(id) + mutation, pending := r.mutations[id] + if record == nil || !pending || mutation.token != token { + return service.ErrGeminiCachedContentMutationConflict + } + record.ExpireTime = expireTime + delete(r.mutations, id) + return nil +} + +func (r *cachedBillingRepo) ReleaseMutation(_ context.Context, id int64, token string) error { + r.mu.Lock() + defer r.mu.Unlock() + if mutation, pending := r.mutations[id]; pending && mutation.token == token { + delete(r.mutations, id) + } + return nil +} + +func (r *cachedBillingRepo) ListStaleMutations(_ context.Context, limit int) ([]int64, error) { + r.mu.Lock() + defer r.mu.Unlock() + var ids []int64 + for id, mutation := range r.mutations { + if len(ids) < limit && mutation.leaseUntil.Before(time.Now()) && r.liveLocked(id) != nil { + ids = append(ids, id) + } + } + return ids, nil } func (r *cachedBillingRepo) SoftDelete(_ context.Context, id int64) error { @@ -87,6 +159,23 @@ func (r *cachedBillingRepo) SoftDelete(_ context.Context, id int64) error { return nil } +// pendingMutation 返回记录是否有未清除的修改意图。 +func (r *cachedBillingRepo) pendingMutation(id int64) bool { + r.mu.Lock() + defer r.mu.Unlock() + _, pending := r.mutations[id] + return pending +} + +// expireLease 让记录上的修改意图租约立即过期。 +func (r *cachedBillingRepo) expireLease(id int64) { + r.mu.Lock() + defer r.mu.Unlock() + mutation := r.mutations[id] + mutation.leaseUntil = time.Now().Add(-time.Second) + r.mutations[id] = mutation +} + type cachedBillingAccountRepo struct { service.AccountRepository accounts map[int64]*service.Account @@ -106,9 +195,12 @@ type cachedBillingUpstream struct { mu sync.Mutex createBody string calls []string + // patchTTLs 按顺序记录 PATCH 请求体中的 ttl。 + patchTTLs []string beforePost func(n int) - // failures 按请求方法返回上游错误响应。 - failures map[string]cachedBillingUpstreamFailure + // failures 按请求方法返回上游错误响应；transportErrors 按请求方法返回传输层错误（上游结果未知）。 + failures map[string]cachedBillingUpstreamFailure + transportErrors map[string]error } type cachedBillingUpstreamFailure struct { @@ -117,11 +209,24 @@ type cachedBillingUpstreamFailure struct { } func (u *cachedBillingUpstream) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) { + var patch struct { + TTL string `json:"ttl"` + } + if req.Method == http.MethodPatch { + _ = json.NewDecoder(req.Body).Decode(&patch) + } u.mu.Lock() u.calls = append(u.calls, req.Method+" "+req.URL.Path) n := len(u.calls) + if req.Method == http.MethodPatch { + u.patchTTLs = append(u.patchTTLs, patch.TTL) + } failure, failed := u.failures[req.Method] + transportErr := u.transportErrors[req.Method] u.mu.Unlock() + if transportErr != nil { + return nil, transportErr + } if failed { return &http.Response{StatusCode: failure.status, Header: http.Header{"Content-Type": []string{"application/json"}}, Body: io.NopCloser(strings.NewReader(failure.body))}, nil @@ -134,7 +239,8 @@ func (u *cachedBillingUpstream) Do(req *http.Request, _ string, _ int64, _ int) if u.beforePost != nil { u.beforePost(n) } - body = fmt.Sprintf(`{"name":"cachedContents/up1","expireTime":%q}`, time.Now().Add(24*time.Hour).UTC().Format(time.RFC3339Nano)) + seconds, _ := strconv.ParseInt(strings.TrimSuffix(patch.TTL, "s"), 10, 64) + body = fmt.Sprintf(`{"name":"cachedContents/up1","expireTime":%q}`, time.Now().Add(time.Duration(seconds)*time.Second).UTC().Format(time.RFC3339Nano)) } return &http.Response{StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"application/json"}}, Body: io.NopCloser(strings.NewReader(body))}, nil @@ -240,7 +346,7 @@ func newCachedBillingEnv(t *testing.T, balance float64, accountIDs ...int64) *ca billing, nil, &service.BillingCacheService{}, nil, nil, &service.DeferredService{}, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil) upstream := &cachedBillingUpstream{} compat := service.NewGeminiMessagesCompatService(&cachedBillingAccountRepo{accounts: accounts}, nil, nil, nil, nil, nil, upstream, nil, cfg) - repo := &cachedBillingRepo{records: map[string]*service.GeminiCachedContent{}} + repo := &cachedBillingRepo{records: map[string]*service.GeminiCachedContent{}, mutations: map[int64]cachedBillingMutation{}} h := &GatewayHandler{ gatewayService: gw, billingCacheService: bcs, @@ -377,37 +483,258 @@ func TestGeminiCachedContentsCreate_UnusableUpstreamResponseCleansUpWithoutRetry require.Empty(t, env.usage.totals()) } -func TestGeminiCachedContentsPatch_ConcurrentExtensionsChargeOnlyActualProgress(t *testing.T) { - env := newCachedBillingEnv(t, 1000, 900) - original := time.Now().Add(time.Hour) - env.seed(original) - env.upstream.beforePost = func(n int) { - if n == 1 { - // 第一个 PATCH 的上游调用返回前，第二个 PATCH 完整执行（读到的仍是旧到期时间）。 - require.Equal(t, http.StatusOK, env.patch().Code) - } - } - require.Equal(t, http.StatusOK, env.patch().Code) - - totals := env.waitForUsage(t, 1) - final := env.repo.records["abc"].ExpireTime - fair := cachedBillingTokens * cachedBillingStoragePrice * final.Sub(original).Hours() - sum := 0.0 - for _, v := range totals { - sum += v - } - require.InDelta(t, fair, sum, cachedBillingTokens*cachedBillingStoragePrice*(time.Minute.Hours()), "两次延长合计只收实际推进的时长") +func (e *cachedBillingEnv) patchTTL(id, ttl string) *httptest.ResponseRecorder { + return e.serve(http.MethodPatch, "/v1beta/cachedContents/"+id+"?updateMask=ttl", fmt.Sprintf(`{"ttl":%q}`, ttl), + gin.Params{{Key: "cacheID", Value: id}}, e.h.GeminiCachedContentsPatch) } -func TestGeminiCachedContentsPatch_PersistFailureDoesNotCharge(t *testing.T) { +// requireTTLNear 断言发往上游的 ttl（"Ns"）与 want 相差不超过 5 秒。 +func requireTTLNear(t *testing.T, ttl string, want time.Duration) { + t.Helper() + seconds, err := strconv.ParseInt(strings.TrimSuffix(ttl, "s"), 10, 64) + require.NoError(t, err, ttl) + require.InDelta(t, want.Seconds(), float64(seconds), 5, ttl) +} + +// 同一缓存的修改串行执行：第一个 PATCH 的上游调用返回前，第二个 PATCH 直接得到 409 且不调用上游， +// 本地到期时间与上游最后生效的修改一致，只按第一个 PATCH 实际推进的时长收费。 +func TestGeminiCachedContentsPatch_ConcurrentUpdateConflictsAndKeepsUpstreamInSync(t *testing.T) { env := newCachedBillingEnv(t, 1000, 900) env.seed(time.Now().Add(time.Hour)) - env.repo.failUpdate = true + env.upstream.beforePost = func(n int) { + if n == 1 { + rec := env.patchTTL("abc", "86400s") + require.Equal(t, http.StatusConflict, rec.Code, rec.Body.String()) + require.Contains(t, rec.Body.String(), `"ABORTED"`) + } + } + + rec := env.patchTTL("abc", "7200s") + require.Equal(t, http.StatusOK, rec.Code, rec.Body.String()) + require.Equal(t, 1, env.upstream.methodCalls(http.MethodPatch), "冲突的修改不调用上游") + require.WithinDuration(t, time.Now().Add(2*time.Hour), env.repo.records["abc"].ExpireTime, 5*time.Second, "本地到期时间与上游一致") + require.False(t, env.repo.pendingMutation(100)) + totals := env.waitForUsage(t, 1) + require.Len(t, totals, 1) + require.InDelta(t, cachedBillingTokens*cachedBillingStoragePrice, totals[0], cachedBillingTokens*cachedBillingStoragePrice*time.Minute.Hours(), "只收推进的 1 小时") +} + +// 上游已修改但本地写回失败：不计费并保留修改意图；租约内的修改返回 409， +// 租约过期后后台恢复把上游有效期回滚到修改前的到期时间并清除意图。 +func TestGeminiCachedContentsPatch_CommitFailureRollsBackUpstreamWithoutCharge(t *testing.T) { + env := newCachedBillingEnv(t, 1000, 900) + baseline := time.Now().Add(time.Hour) + env.seed(baseline) + env.repo.failCommit = true + + require.Equal(t, http.StatusInternalServerError, env.patch().Code) + require.Equal(t, 1, env.upstream.methodCalls(http.MethodPatch)) + require.True(t, env.repo.pendingMutation(100), "上游可能已修改：保留修改意图") + require.Empty(t, env.usage.totals(), "未提交的修改不收费") + + env.repo.failCommit = false + require.Equal(t, http.StatusConflict, env.patch().Code, "租约内的修改被拒绝") + require.Equal(t, 1, env.upstream.methodCalls(http.MethodPatch)) + + env.repo.expireLease(100) + recovered, err := env.h.geminiCachedContentService.RecoverStaleMutations(context.Background()) + require.NoError(t, err) + require.Equal(t, 1, recovered) + require.Equal(t, 2, env.upstream.methodCalls(http.MethodPatch)) + requireTTLNear(t, env.upstream.patchTTLs[1], time.Until(baseline)) + require.False(t, env.repo.pendingMutation(100)) + require.True(t, env.repo.records["abc"].ExpireTime.Equal(baseline), "本地到期时间保持修改前的值") + require.Empty(t, env.usage.totals()) +} + +// 租约过期的遗留修改意图在下一次修改前先回滚，随后的修改按修改前的到期时间计费。 +func TestGeminiCachedContentsPatch_StaleMutationRolledBackBeforeNextUpdate(t *testing.T) { + env := newCachedBillingEnv(t, 1000, 900) + baseline := time.Now().Add(time.Hour) + env.seed(baseline) + env.repo.failCommit = true + require.Equal(t, http.StatusInternalServerError, env.patch().Code) + env.repo.failCommit = false + env.repo.expireLease(100) rec := env.patch() - require.Equal(t, http.StatusInternalServerError, rec.Code) - require.Equal(t, 1, env.upstream.methodCalls(http.MethodPatch)) - require.Empty(t, env.usage.totals(), "落库失败不收费") + require.Equal(t, http.StatusOK, rec.Code, rec.Body.String()) + require.Equal(t, 3, env.upstream.methodCalls(http.MethodPatch), "失败的修改、回滚、本次修改") + requireTTLNear(t, env.upstream.patchTTLs[1], time.Until(baseline)) + require.Equal(t, "86400s", env.upstream.patchTTLs[2]) + totals := env.waitForUsage(t, 1) + require.Len(t, totals, 1) + require.InDelta(t, cachedBillingTokens*cachedBillingStoragePrice*23, totals[0], 0.01) +} + +// 上游结果未知（传输层错误）时保留修改意图等待回滚；上游明确拒绝时清除修改意图。两种情况都不收费。 +func TestGeminiCachedContentsPatch_MutationKeptOnlyWhenUpstreamOutcomeUnknown(t *testing.T) { + t.Run("transport error", func(t *testing.T) { + env := newCachedBillingEnv(t, 1000, 900) + env.seed(time.Now().Add(time.Hour)) + env.upstream.transportErrors = map[string]error{http.MethodPatch: errors.New("connection reset")} + + require.Equal(t, http.StatusBadGateway, env.patch().Code) + require.True(t, env.repo.pendingMutation(100)) + require.Empty(t, env.usage.totals()) + }) + t.Run("upstream rejected", func(t *testing.T) { + env := newCachedBillingEnv(t, 1000, 900) + env.seed(time.Now().Add(time.Hour)) + env.upstream.failures = map[string]cachedBillingUpstreamFailure{http.MethodPatch: {status: http.StatusBadRequest, + body: `{"error":{"code":400,"message":"invalid ttl","status":"INVALID_ARGUMENT"}}`}} + + require.Equal(t, http.StatusBadRequest, env.patch().Code) + require.False(t, env.repo.pendingMutation(100)) + require.Empty(t, env.usage.totals()) + }) +} + +// 修改前的到期时间已过：恢复时删除上游缓存与本地记录。 +func TestGeminiCachedContentMutationRecovery_DeletesUpstreamWhenBaselineExpired(t *testing.T) { + env := newCachedBillingEnv(t, 1000, 900) + env.seed(time.Now().Add(-time.Minute)) + env.repo.mutations[100] = cachedBillingMutation{token: "stale", leaseUntil: time.Now().Add(-time.Second)} + + recovered, err := env.h.geminiCachedContentService.RecoverStaleMutations(context.Background()) + require.NoError(t, err) + require.Equal(t, 1, recovered) + require.Equal(t, 1, env.upstream.methodCalls(http.MethodDelete)) + require.Zero(t, env.upstream.methodCalls(http.MethodPatch)) + require.Equal(t, []int64{100}, env.repo.deleted) +} + +// cachedBillingReservationCache 是启用在途预留时的计费缓存桩：余额与订阅用量取固定值，各维度预留按 key 记账。 +type cachedBillingReservationCache struct { + service.BillingCache + mu sync.Mutex + balance float64 + sub *service.SubscriptionCacheData + reserved map[string]map[string]float64 +} + +func (c *cachedBillingReservationCache) GetUserBalance(context.Context, int64) (float64, error) { + return c.balance, nil +} + +func (c *cachedBillingReservationCache) GetSubscriptionCache(context.Context, int64, int64) (*service.SubscriptionCacheData, error) { + return c.sub, nil +} + +func (c *cachedBillingReservationCache) reserveLocked(key, requestID string, amount float64) float64 { + if c.reserved == nil { + c.reserved = map[string]map[string]float64{} + } + if c.reserved[key] == nil { + c.reserved[key] = map[string]float64{} + } + others := 0.0 + for _, v := range c.reserved[key] { + others += v + } + c.reserved[key][requestID] = amount + return others +} + +func (c *cachedBillingReservationCache) ReserveInflightBalance(_ context.Context, userID int64, requestID string, amount, balance float64, _ time.Duration) (bool, float64, error) { + c.mu.Lock() + defer c.mu.Unlock() + key := fmt.Sprintf("user:%d", userID) + others := c.reserveLocked(key, requestID, amount) + if len(c.reserved[key]) > 1 && balance-others < amount { + delete(c.reserved[key], requestID) + return false, others, nil + } + return true, others, nil +} + +func (c *cachedBillingReservationCache) ReleaseInflightBalance(_ context.Context, userID int64, requestID string) error { + c.mu.Lock() + defer c.mu.Unlock() + delete(c.reserved[fmt.Sprintf("user:%d", userID)], requestID) + return nil +} + +func (c *cachedBillingReservationCache) ReserveInflightScope(_ context.Context, scope, requestID string, amount float64, _ time.Duration) (float64, error) { + c.mu.Lock() + defer c.mu.Unlock() + return c.reserveLocked(scope, requestID, amount), nil +} + +func (c *cachedBillingReservationCache) ReleaseInflightScope(_ context.Context, scope, requestID string) error { + c.mu.Lock() + defer c.mu.Unlock() + delete(c.reserved[scope], requestID) + return nil +} + +func enableCachedBillingReservations(t *testing.T, env *cachedBillingEnv, cache *cachedBillingReservationCache) { + t.Helper() + cfg := &config.Config{} + cfg.Billing.InflightReservation.Enabled = true + bcs := service.NewBillingCacheService(cache, nil, nil, nil, nil, nil, cfg, nil) + t.Cleanup(bcs.Stop) + env.h.billingCacheService = bcs +} + +// 确定金额的续期费按 API Key 与余额维度预留：第一笔续期（约 $103.5）尚未入账时， +// 第二个缓存的续期超出 Key 剩余额度或余额，在调用上游前被拒绝。 +func TestGeminiCachedContentsPatch_UnsettledChargeHoldsKeyQuotaAndBalance(t *testing.T) { + for _, tc := range []struct { + name string + balance, quota float64 + }{ + {name: "api key quota", balance: 1000, quota: 110}, + {name: "balance", balance: 110}, + } { + t.Run(tc.name, func(t *testing.T) { + env := newCachedBillingEnv(t, tc.balance, 900) + env.seed(time.Now().Add(time.Hour)) + second := *env.repo.records["abc"] + second.ID, second.PublicID, second.UpstreamName = 101, "def", "cachedContents/up2" + env.repo.records["def"] = &second + env.apiKey.Quota = tc.quota + enableCachedBillingReservations(t, env, &cachedBillingReservationCache{balance: tc.balance}) + env.upstream.beforePost = func(n int) { + if n == 1 { + rec := env.patchTTL("def", "86400s") + require.NotEqual(t, http.StatusOK, rec.Code, rec.Body.String()) + } + } + + require.Equal(t, http.StatusOK, env.patch().Code) + require.Equal(t, 1, env.upstream.methodCalls(http.MethodPatch), "第二笔续期没有到达上游") + }) + } +} + +// 订阅模式的确定金额费用按订阅维度预留：第一笔未入账时第二笔超出日剩余额度被拒绝，释放后放行。 +func TestReserveGeminiCachedContentCharge_UnsettledChargeHoldsSubscriptionHeadroom(t *testing.T) { + env := newCachedBillingEnv(t, 1000, 900) + limit := 110.0 + env.group.SubscriptionType = service.SubscriptionTypeSubscription + env.group.DailyLimitUSD = &limit + enableCachedBillingReservations(t, env, &cachedBillingReservationCache{balance: 1000, + sub: &service.SubscriptionCacheData{Status: service.SubscriptionStatusActive, ExpiresAt: time.Now().Add(time.Hour)}}) + sub := &service.UserSubscription{ID: 1, UserID: env.apiKey.UserID, GroupID: env.group.ID, Status: service.SubscriptionStatusActive, + ExpiresAt: time.Now().Add(time.Hour)} + reserve := func() (func(), error) { + var done func() + var err error + env.serve(http.MethodPatch, "/", `{}`, nil, func(c *gin.Context) { + done, err = env.h.reserveGeminiCachedContentCharge(c, env.apiKey, sub, 103.5) + }) + return done, err + } + + firstDone, err := reserve() + require.NoError(t, err) + _, err = reserve() + require.ErrorIs(t, err, service.ErrDailyLimitExceeded) + firstDone() + secondDone, err := reserve() + require.NoError(t, err) + secondDone() } func TestGeminiCachedContentsPatch_UnaffordableExtensionRejectedBeforeUpstream(t *testing.T) { diff --git a/backend/internal/handler/gemini_cached_contents_handler.go b/backend/internal/handler/gemini_cached_contents_handler.go index ac99eaf1d5..2071c16129 100644 --- a/backend/internal/handler/gemini_cached_contents_handler.go +++ b/backend/internal/handler/gemini_cached_contents_handler.go @@ -395,6 +395,8 @@ func (h *GatewayHandler) GeminiCachedContentsGet(c *gin.Context) { // GeminiCachedContentsPatch PATCH /v1beta/cachedContents/:cacheID // 只允许修改有效期；延长部分在本次请求内按增量时长收取存储费，缩短不退费。 +// 同一缓存的修改串行执行：上游修改、本地写回与计费都在修改意图的租约内完成，并发的修改返回 409。 +// 上游结果未知或本地写回失败的修改不计费，租约过期后由恢复流程把上游有效期回滚到修改前的到期时间。 func (h *GatewayHandler) GeminiCachedContentsPatch(c *gin.Context) { apiKey, _, reqLog, ok := h.geminiCachedContentsPreamble(c, "handler.gemini_v1beta.cached_contents.patch") if !ok { @@ -455,10 +457,33 @@ func (h *GatewayHandler) GeminiCachedContentsPatch(c *gin.Context) { } setOpsSelectedAccount(c, account.ID, account.Platform) + mutation, err := svc.BeginMutation(c.Request.Context(), record) + if err != nil { + switch { + case errors.Is(err, service.ErrGeminiCachedContentNotFound): + googleError(c, http.StatusForbidden, service.GeminiCachedContentNotFoundMessage) + case errors.Is(err, service.ErrGeminiCachedContentMutationConflict): + reqLog.Info("gemini.cached_content.mutation_conflict", zap.Int64("account_id", account.ID), zap.Error(err)) + writeGeminiCachedContentMutationConflict(c) + default: + reqLog.Error("gemini.cached_content.begin_mutation_failed", zap.Int64("account_id", account.ID), zap.Error(err)) + googleError(c, http.StatusInternalServerError, "Failed to update cached content") + } + return + } upstream, err := svc.PatchUpstream(c.Request.Context(), account, record, ttl) if err != nil { var upErr *service.GeminiCachedContentUpstreamError - if errors.As(err, &upErr) { + upstreamRejected := errors.As(err, &upErr) + if upstreamRejected || errors.Is(err, service.ErrGeminiCachedContentRequestNotSent) { + // 上游有效期未变：清除修改意图。 + persistCtx, cancelPersist := geminiCachedContentPersistContext(c) + if releaseErr := svc.ReleaseMutation(persistCtx, mutation); releaseErr != nil { + reqLog.Warn("gemini.cached_content.release_mutation_failed", zap.Error(releaseErr)) + } + cancelPersist() + } + if upstreamRejected { if upErr.NotFound() { if forgetErr := svc.Forget(c.Request.Context(), record); forgetErr != nil { reqLog.Warn("gemini.cached_content.forget_failed", zap.Error(forgetErr)) @@ -473,19 +498,19 @@ func (h *GatewayHandler) GeminiCachedContentsPatch(c *gin.Context) { } expireTime := service.ClampGeminiCachedContentExpire(upstream.ExpireTime, svc.Now().Add(ttl)) persistCtx, cancelPersist := geminiCachedContentPersistContext(c) - previousExpire, err := svc.UpdateExpireTime(persistCtx, record, expireTime) + err = svc.CommitMutation(persistCtx, mutation, expireTime) cancelPersist() if err != nil { - reqLog.Error("gemini.cached_content.update_expire_failed", zap.Int64("account_id", account.ID), zap.Error(err)) - if errors.Is(err, service.ErrGeminiCachedContentNotFound) { - googleError(c, http.StatusForbidden, service.GeminiCachedContentNotFoundMessage) + reqLog.Error("gemini.cached_content.commit_mutation_failed", zap.Int64("account_id", account.ID), zap.Error(err)) + if errors.Is(err, service.ErrGeminiCachedContentMutationConflict) { + writeGeminiCachedContentMutationConflict(c) return } googleError(c, http.StatusInternalServerError, "Failed to persist cached content expiration") return } - // 按库中实际推进的时长收费：并发延长各自只为自己推进的部分付费。 - if extended := expireTime.Sub(previousExpire); extended > 0 { + // 提交前库中的到期时间一直是修改前的值：按本次实际推进的时长收费。 + if extended := expireTime.Sub(mutation.Baseline()); extended > 0 { h.submitGeminiCachedContentUsage(c, apiKey, account, subscription, pricingAt, service.GeminiCachedContentPatchUsageRequestID(record.PublicID), record.Model, record.UpstreamModelID(), 0, service.GeminiCachedContentStorageTokenHours(record.TotalTokenCount, extended), record.ChannelUsage, body) @@ -706,13 +731,11 @@ func (h *GatewayHandler) cleanupGeminiCachedContentUpstream(reqLog *zap.Logger, } // reserveGeminiCachedContentCharge 校验余额 / 订阅额度 / API Key 额度足以支付确定金额 amount（倍率后）， -// 并按该金额登记在途预留（不受预留上限封顶）。返回的 done 须在 handler 结束时调用。 +// 并按该金额在用户、订阅、API Key 维度登记在途预留（不受预留上限封顶，见 ReserveFixedCharge）。 +// 返回的 done 须在 handler 结束时调用；预留随计费任务交接，落库后释放。 func (h *GatewayHandler) reserveGeminiCachedContentCharge(c *gin.Context, apiKey *service.APIKey, subscription *service.UserSubscription, amount float64) (func(), error) { ctx := c.Request.Context() - if err := h.billingCacheService.CheckChargeAffordable(ctx, apiKey.User, apiKey, apiKey.Group, subscription, amount); err != nil { - return inflightNoop, err - } - res, err := h.billingCacheService.ReserveInflightCharge(ctx, apiKey.User, apiKey.Group, subscription, amount) + res, err := h.billingCacheService.ReserveFixedCharge(ctx, apiKey.User, apiKey, apiKey.Group, subscription, amount) if err != nil { return inflightNoop, err } @@ -723,6 +746,17 @@ func (h *GatewayHandler) reserveGeminiCachedContentCharge(c *gin.Context, apiKey return res.HandlerDone, nil } +// writeGeminiCachedContentMutationConflict 回写同一缓存的修改冲突（Google 并发冲突语义 ABORTED），客户端可稍后重试。 +func writeGeminiCachedContentMutationConflict(c *gin.Context) { + c.JSON(http.StatusConflict, gin.H{ + "error": gin.H{ + "code": http.StatusConflict, + "message": "Another update to this cached content is in progress, please retry later", + "status": "ABORTED", + }, + }) +} + func writeGeminiCachedContentBillingError(c *gin.Context, err error) { status, code, message, retryAfter := billingErrorDetails(err) if retryAfter > 0 { diff --git a/backend/internal/handler/gemini_cached_contents_handler_test.go b/backend/internal/handler/gemini_cached_contents_handler_test.go index 26bc5e7183..f9a1d2b493 100644 --- a/backend/internal/handler/gemini_cached_contents_handler_test.go +++ b/backend/internal/handler/gemini_cached_contents_handler_test.go @@ -52,8 +52,24 @@ func (r *geminiCachedContentHandlerRepoStub) ListForOwner(_ context.Context, api return out, nil } -func (r *geminiCachedContentHandlerRepoStub) UpdateExpireTime(context.Context, int64, time.Time) (time.Time, error) { - return time.Time{}, nil +func (r *geminiCachedContentHandlerRepoStub) BeginMutation(context.Context, int64, time.Time, string, time.Duration) (bool, error) { + return true, nil +} + +func (r *geminiCachedContentHandlerRepoStub) TakeOverStaleMutation(context.Context, int64, string, time.Duration) (*service.GeminiCachedContent, error) { + return nil, nil +} + +func (r *geminiCachedContentHandlerRepoStub) CommitMutation(context.Context, int64, string, time.Time) error { + return nil +} + +func (r *geminiCachedContentHandlerRepoStub) ReleaseMutation(context.Context, int64, string) error { + return nil +} + +func (r *geminiCachedContentHandlerRepoStub) ListStaleMutations(context.Context, int) ([]int64, error) { + return nil, nil } func (r *geminiCachedContentHandlerRepoStub) SoftDelete(context.Context, int64) error { diff --git a/backend/internal/repository/api_key_repo.go b/backend/internal/repository/api_key_repo.go index dd831283b7..1fa3049ca9 100644 --- a/backend/internal/repository/api_key_repo.go +++ b/backend/internal/repository/api_key_repo.go @@ -857,7 +857,7 @@ func (r *apiKeyRepository) ResetRateLimitWindows(ctx context.Context, id int64) // GetRateLimitData returns the current rate limit usage and window start times for an API key. func (r *apiKeyRepository) GetRateLimitData(ctx context.Context, id int64) (result *service.APIKeyRateLimitData, err error) { rows, err := r.sql.QueryContext(ctx, ` - SELECT usage_5h, usage_1d, usage_7d, window_5h_start, window_1d_start, window_7d_start + SELECT quota_used, usage_5h, usage_1d, usage_7d, window_5h_start, window_1d_start, window_7d_start FROM api_keys WHERE id = $1 AND deleted_at IS NULL`, id) @@ -873,7 +873,7 @@ func (r *apiKeyRepository) GetRateLimitData(ctx context.Context, id int64) (resu return nil, service.ErrAPIKeyNotFound } data := &service.APIKeyRateLimitData{} - if err := rows.Scan(&data.Usage5h, &data.Usage1d, &data.Usage7d, &data.Window5hStart, &data.Window1dStart, &data.Window7dStart); err != nil { + if err := rows.Scan(&data.QuotaUsed, &data.Usage5h, &data.Usage1d, &data.Usage7d, &data.Window5hStart, &data.Window1dStart, &data.Window7dStart); err != nil { return nil, err } return data, rows.Err() diff --git a/backend/internal/repository/api_key_repo_integration_test.go b/backend/internal/repository/api_key_repo_integration_test.go index ada0a201b3..9d41f269b6 100644 --- a/backend/internal/repository/api_key_repo_integration_test.go +++ b/backend/internal/repository/api_key_repo_integration_test.go @@ -464,6 +464,10 @@ func (s *APIKeyRepoSuite) TestIncrementQuotaUsed_Basic() { newQuota, err = s.repo.IncrementQuotaUsed(s.ctx, key.ID, 2.5) s.Require().NoError(err, "IncrementQuotaUsed second") s.Require().Equal(4.0, newQuota, "第二次递增后应为 4.0") + + data, err := s.repo.GetRateLimitData(s.ctx, key.ID) + s.Require().NoError(err, "GetRateLimitData") + s.Require().Equal(4.0, data.QuotaUsed, "用量读取带回总额度已用量") } func (s *APIKeyRepoSuite) TestIncrementQuotaUsed_NotFound() { diff --git a/backend/internal/repository/billing_inflight_cache.go b/backend/internal/repository/billing_inflight_cache.go index 898cf3b52a..8f83717ce6 100644 --- a/backend/internal/repository/billing_inflight_cache.go +++ b/backend/internal/repository/billing_inflight_cache.go @@ -25,6 +25,20 @@ func billingInflightKeys(userID int64) (string, string) { fmt.Sprintf("%s{%d}", billingInflightAmtKeyPrefix, userID) } +// 订阅 / API Key 维度在途预留（确定金额费用）： +// - billing:inflight_scope:{scope} ZSET member=requestID score=过期时间(ms) +// - billing:inflight_scope_amt:{scope} HASH field=requestID value=预留金额(USD) +// +// 过期清理与释放语义与余额维度相同，不续期。 +const ( + billingInflightScopeKeyPrefix = "billing:inflight_scope:" + billingInflightScopeAmtKeyPrefix = "billing:inflight_scope_amt:" +) + +func billingInflightScopeKeys(scope string) (string, string) { + return billingInflightScopeKeyPrefix + "{" + scope + "}", billingInflightScopeAmtKeyPrefix + "{" + scope + "}" +} + var ( // KEYS: [1]=zset [2]=hash // ARGV: [1]=now_ms [2]=expire_at_ms [3]=member [4]=amount [5]=balance [6]=key_ttl_ms @@ -85,6 +99,28 @@ var ( redis.call('HDEL', KEYS[2], ARGV[1]) return 1 `) + + // KEYS: [1]=zset [2]=hash + // ARGV: [1]=now_ms [2]=expire_at_ms [3]=member [4]=amount [5]=key_ttl_ms + // 先清理已过期成员，再登记本次预留；返回登记前其他在途预留的合计（string），是否放行由调用方按权威用量判定。 + reserveInflightScopeScript = redis.NewScript(` + local now = tonumber(ARGV[1]) + local expired = redis.call('ZRANGEBYSCORE', KEYS[1], '-inf', now) + if #expired > 0 then + redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', now) + redis.call('HDEL', KEYS[2], unpack(expired)) + end + local vals = redis.call('HVALS', KEYS[2]) + local sum = 0 + for i = 1, #vals do + sum = sum + (tonumber(vals[i]) or 0) + end + redis.call('ZADD', KEYS[1], tonumber(ARGV[2]), ARGV[3]) + redis.call('HSET', KEYS[2], ARGV[3], ARGV[4]) + redis.call('PEXPIRE', KEYS[1], ARGV[5]) + redis.call('PEXPIRE', KEYS[2], ARGV[5]) + return tostring(sum) + `) ) // ReserveInflightBalance 实现 service.InflightBalanceReservationCache。 @@ -138,7 +174,35 @@ func (c *billingCache) RenewInflightBalance(ctx context.Context, userID int64, r return n == 1, nil } +// ReserveInflightScope 实现 service.InflightScopeReservationCache。 +func (c *billingCache) ReserveInflightScope(ctx context.Context, scope, requestID string, amount float64, ttl time.Duration) (float64, error) { + zkey, hkey := billingInflightScopeKeys(scope) + now := time.Now().UnixMilli() + ttlMs := ttl.Milliseconds() + if ttlMs <= 0 { + ttlMs = 1 + } + res, err := reserveInflightScopeScript.Run(ctx, c.rdb, []string{zkey, hkey}, + now, + now+ttlMs, + requestID, + strconv.FormatFloat(amount, 'f', -1, 64), + ttlMs, + ).Text() + if err != nil { + return 0, err + } + return strconv.ParseFloat(res, 64) +} + +// ReleaseInflightScope 实现 service.InflightScopeReservationCache。 +func (c *billingCache) ReleaseInflightScope(ctx context.Context, scope, requestID string) error { + zkey, hkey := billingInflightScopeKeys(scope) + return releaseInflightBalanceScript.Run(ctx, c.rdb, []string{zkey, hkey}, requestID).Err() +} + var ( _ service.InflightBalanceReservationCache = (*billingCache)(nil) _ service.InflightBalanceReservationRenewer = (*billingCache)(nil) + _ service.InflightScopeReservationCache = (*billingCache)(nil) ) diff --git a/backend/internal/repository/billing_inflight_cache_test.go b/backend/internal/repository/billing_inflight_cache_test.go index ddab7c4de0..2fd642fe38 100644 --- a/backend/internal/repository/billing_inflight_cache_test.go +++ b/backend/internal/repository/billing_inflight_cache_test.go @@ -285,3 +285,72 @@ func TestInflightReservation_RenewDoesNotResurrectExpiredMember(t *testing.T) { require.NoError(t, err) require.False(t, exists) } + +func TestInflightScopeReservation_ReturnsOthersAndReleases(t *testing.T) { + mr, cache, _ := newInflightTestEnv(t, true, 60) + ctx := context.Background() + + others, err := cache.ReserveInflightScope(ctx, "apikey:7", "r1", 3, time.Minute) + require.NoError(t, err) + require.Zero(t, others) + others, err = cache.ReserveInflightScope(ctx, "apikey:7", "r2", 2, time.Minute) + require.NoError(t, err) + require.InDelta(t, 3, others, 1e-9, "返回登记前其他在途预留的合计") + others, err = cache.ReserveInflightScope(ctx, "sub:7", "r3", 5, time.Minute) + require.NoError(t, err) + require.Zero(t, others, "不同维度互不影响") + + require.NoError(t, cache.ReleaseInflightScope(ctx, "apikey:7", "r1")) + require.NoError(t, cache.ReleaseInflightScope(ctx, "apikey:7", "r1")) + others, err = cache.ReserveInflightScope(ctx, "apikey:7", "r4", 1, time.Minute) + require.NoError(t, err) + require.InDelta(t, 2, others, 1e-9) + + mr.FastForward(2 * time.Minute) + others, err = cache.ReserveInflightScope(ctx, "apikey:7", "r5", 1, time.Minute) + require.NoError(t, err) + require.Zero(t, others, "过期预留不再计入") +} + +func TestReserveFixedCharge_ConcurrentChargesAdmitOnlyWhatAPIKeyQuotaCovers(t *testing.T) { + _, cache, svc := newInflightTestEnv(t, true, 60) + ctx := context.Background() + user := &service.User{ID: 42} + apiKey := &service.APIKey{ID: 9, UserID: user.ID, Quota: 1.0} + require.NoError(t, cache.SetUserBalance(ctx, user.ID, 100)) + + const n = 20 + var admitted atomic.Int32 + reservations := make(chan *service.InflightReservation, n) + var wg sync.WaitGroup + start := make(chan struct{}) + for i := 0; i < n; i++ { + wg.Add(1) + go func() { + defer wg.Done() + <-start + res, err := svc.ReserveFixedCharge(ctx, user, apiKey, nil, nil, 0.3) + if err != nil { + if !errors.Is(err, service.ErrAPIKeyQuotaExhausted) { + t.Errorf("unexpected error: %v", err) + } + return + } + admitted.Add(1) + reservations <- res + }() + } + close(start) + wg.Wait() + close(reservations) + + require.Equal(t, int32(3), admitted.Load(), "Key 额度 1.0、每笔 0.3：同时在途的不超过 3 笔") + for res := range reservations { + res.Release() + } + zkey, _ := billingInflightScopeKeys("apikey:9") + remaining, err := cache.rdb.ZCard(ctx, zkey).Result() + require.NoError(t, err) + require.Zero(t, remaining) + require.Equal(t, int64(0), inflightCount(t, cache, user.ID)) +} diff --git a/backend/internal/repository/gemini_cached_content_repo.go b/backend/internal/repository/gemini_cached_content_repo.go index ddb05bbdcf..42775448d7 100644 --- a/backend/internal/repository/gemini_cached_content_repo.go +++ b/backend/internal/repository/gemini_cached_content_repo.go @@ -4,6 +4,7 @@ import ( "context" "database/sql" "encoding/json" + "errors" "time" "github.com/Wei-Shaw/sub2api/internal/service" @@ -78,20 +79,88 @@ func (r *geminiCachedContentRepository) ListForOwner(ctx context.Context, apiKey return out, rows.Err() } -func (r *geminiCachedContentRepository) UpdateExpireTime(ctx context.Context, id int64, expireTime time.Time) (time.Time, error) { - var previous time.Time - err := r.sql.QueryRowContext(ctx, `WITH prev AS ( - SELECT id, expire_time FROM gemini_cached_contents WHERE id = $1 AND deleted_at IS NULL FOR UPDATE -) -UPDATE gemini_cached_contents AS g - SET expire_time = $2, updated_at = NOW() - FROM prev - WHERE g.id = prev.id - RETURNING prev.expire_time`, id, expireTime).Scan(&previous) +func (r *geminiCachedContentRepository) BeginMutation(ctx context.Context, id int64, baseline time.Time, token string, lease time.Duration) (bool, error) { + result, err := r.sql.ExecContext(ctx, `UPDATE gemini_cached_contents + SET mutation_token = $2, mutation_lease_until = NOW() + make_interval(secs => $3::double precision) + WHERE id = $1 AND deleted_at IS NULL AND mutation_token IS NULL AND expire_time = $4`, + id, token, lease.Seconds(), baseline) if err != nil { - return time.Time{}, translatePersistenceError(err, service.ErrGeminiCachedContentNotFound, nil) + return false, err } - return previous, nil + affected, err := result.RowsAffected() + if err != nil { + return false, err + } + if affected > 0 { + return true, nil + } + var exists bool + if err := r.sql.QueryRowContext(ctx, `SELECT EXISTS ( + SELECT 1 FROM gemini_cached_contents WHERE id = $1 AND deleted_at IS NULL)`, id).Scan(&exists); err != nil { + return false, err + } + if !exists { + return false, service.ErrGeminiCachedContentNotFound + } + return false, nil +} + +func (r *geminiCachedContentRepository) TakeOverStaleMutation(ctx context.Context, id int64, token string, lease time.Duration) (*service.GeminiCachedContent, error) { + record, err := scanGeminiCachedContent(r.sql.QueryRowContext(ctx, `UPDATE gemini_cached_contents + SET mutation_token = $2, mutation_lease_until = NOW() + make_interval(secs => $3::double precision) + WHERE id = $1 AND deleted_at IS NULL AND mutation_token IS NOT NULL AND mutation_lease_until < NOW() + RETURNING `+geminiCachedContentColumns, id, token, lease.Seconds())) + if errors.Is(err, sql.ErrNoRows) { + return nil, nil + } + if err != nil { + return nil, err + } + return record, nil +} + +func (r *geminiCachedContentRepository) CommitMutation(ctx context.Context, id int64, token string, expireTime time.Time) error { + result, err := r.sql.ExecContext(ctx, `UPDATE gemini_cached_contents + SET expire_time = $3, updated_at = NOW(), mutation_token = NULL, mutation_lease_until = NULL + WHERE id = $1 AND deleted_at IS NULL AND mutation_token = $2`, id, token, expireTime) + if err != nil { + return err + } + affected, err := result.RowsAffected() + if err != nil { + return err + } + if affected == 0 { + return service.ErrGeminiCachedContentMutationConflict + } + return nil +} + +func (r *geminiCachedContentRepository) ReleaseMutation(ctx context.Context, id int64, token string) error { + _, err := r.sql.ExecContext(ctx, `UPDATE gemini_cached_contents + SET mutation_token = NULL, mutation_lease_until = NULL + WHERE id = $1 AND mutation_token = $2`, id, token) + return err +} + +func (r *geminiCachedContentRepository) ListStaleMutations(ctx context.Context, limit int) ([]int64, error) { + rows, err := r.sql.QueryContext(ctx, `SELECT id FROM gemini_cached_contents + WHERE mutation_token IS NOT NULL AND deleted_at IS NULL AND mutation_lease_until < NOW() + ORDER BY mutation_lease_until + LIMIT $1`, limit) + if err != nil { + return nil, err + } + defer func() { _ = rows.Close() }() + var ids []int64 + for rows.Next() { + var id int64 + if err := rows.Scan(&id); err != nil { + return nil, err + } + ids = append(ids, id) + } + return ids, rows.Err() } func (r *geminiCachedContentRepository) SoftDelete(ctx context.Context, id int64) error { diff --git a/backend/internal/repository/gemini_cached_content_repo_integration_test.go b/backend/internal/repository/gemini_cached_content_repo_integration_test.go index fc537a45f3..89399552c2 100644 --- a/backend/internal/repository/gemini_cached_content_repo_integration_test.go +++ b/backend/internal/repository/gemini_cached_content_repo_integration_test.go @@ -61,22 +61,58 @@ func TestGeminiCachedContentRepository_Lifecycle(t *testing.T) { require.Len(t, page, 1) require.Equal(t, first.PublicID, page[0].PublicID) - newExpire := now.Add(3 * time.Hour) - previous, err := repo.UpdateExpireTime(ctx, first.ID, newExpire) + lease := 2 * time.Minute + acquired, err := repo.BeginMutation(ctx, first.ID, first.ExpireTime, "t1", lease) require.NoError(t, err) - require.True(t, previous.Equal(first.ExpireTime), "返回更新前的到期时间") + require.True(t, acquired) + acquired, err = repo.BeginMutation(ctx, first.ID, first.ExpireTime, "t2", lease) + require.NoError(t, err) + require.False(t, acquired, "有进行中的修改时不能再登记") + taken, err := repo.TakeOverStaleMutation(ctx, first.ID, "t2", lease) + require.NoError(t, err) + require.Nil(t, taken, "租约未过期不能接管") + require.ErrorIs(t, repo.CommitMutation(ctx, first.ID, "t2", now), service.ErrGeminiCachedContentMutationConflict) + + newExpire := now.Add(3 * time.Hour) + require.NoError(t, repo.CommitMutation(ctx, first.ID, "t1", newExpire)) got, err = repo.GetForOwner(ctx, apiKeyID, first.PublicID) require.NoError(t, err) require.True(t, got.ExpireTime.Equal(newExpire)) - previous, err = repo.UpdateExpireTime(ctx, first.ID, newExpire.Add(time.Hour)) + require.ErrorIs(t, repo.CommitMutation(ctx, first.ID, "t1", now), service.ErrGeminiCachedContentMutationConflict, "提交后修改意图已清除") + + acquired, err = repo.BeginMutation(ctx, first.ID, first.ExpireTime, "t3", lease) require.NoError(t, err) - require.True(t, previous.Equal(newExpire), "连续更新各自拿到自己的前值") - newExpire = newExpire.Add(time.Hour) + require.False(t, acquired, "到期时间已被其他修改改变") + acquired, err = repo.BeginMutation(ctx, first.ID, newExpire, "t3", lease) + require.NoError(t, err) + require.True(t, acquired) + stale, err := repo.ListStaleMutations(ctx, 10) + require.NoError(t, err) + require.NotContains(t, stale, first.ID) + + _, err = tx.ExecContext(ctx, `UPDATE gemini_cached_contents SET mutation_lease_until = NOW() - INTERVAL '1 second' WHERE id = $1`, first.ID) + require.NoError(t, err) + stale, err = repo.ListStaleMutations(ctx, 10) + require.NoError(t, err) + require.Contains(t, stale, first.ID) + taken, err = repo.TakeOverStaleMutation(ctx, first.ID, "t4", lease) + require.NoError(t, err) + require.NotNil(t, taken) + require.True(t, taken.ExpireTime.Equal(newExpire), "接管返回修改前的到期时间") + require.ErrorIs(t, repo.CommitMutation(ctx, first.ID, "t3", now), service.ErrGeminiCachedContentMutationConflict, "被接管后原持有者不能提交") + require.NoError(t, repo.ReleaseMutation(ctx, first.ID, "t3")) + acquired, err = repo.BeginMutation(ctx, first.ID, newExpire, "t5", lease) + require.NoError(t, err) + require.False(t, acquired, "原持有者释放不影响接管者") + require.NoError(t, repo.ReleaseMutation(ctx, first.ID, "t4")) + stale, err = repo.ListStaleMutations(ctx, 10) + require.NoError(t, err) + require.NotContains(t, stale, first.ID) require.NoError(t, repo.SoftDelete(ctx, first.ID)) _, err = repo.GetForOwner(ctx, apiKeyID, first.PublicID) require.ErrorIs(t, err, service.ErrGeminiCachedContentNotFound) - _, err = repo.UpdateExpireTime(ctx, first.ID, newExpire) + _, err = repo.BeginMutation(ctx, first.ID, newExpire, "t6", lease) require.ErrorIs(t, err, service.ErrGeminiCachedContentNotFound) list, err = repo.ListForOwner(ctx, apiKeyID, 42, now, 0, 10) require.NoError(t, err) diff --git a/backend/internal/service/api_key_service.go b/backend/internal/service/api_key_service.go index b00b03d43f..e20a2697eb 100644 --- a/backend/internal/service/api_key_service.go +++ b/backend/internal/service/api_key_service.go @@ -130,6 +130,8 @@ type apiKeyAllByUserIDLister interface { // APIKeyRateLimitData holds rate limit usage and window state for an API key. type APIKeyRateLimitData struct { + // QuotaUsed 是总额度已用量（quota_used）。 + QuotaUsed float64 Usage5h float64 Usage1d float64 Usage7d float64 diff --git a/backend/internal/service/billing_cache_service.go b/backend/internal/service/billing_cache_service.go index 70f75a7ae8..ea22efe2eb 100644 --- a/backend/internal/service/billing_cache_service.go +++ b/backend/internal/service/billing_cache_service.go @@ -790,48 +790,106 @@ func (s *BillingCacheService) CheckBillingEligibility(ctx context.Context, user return nil } -// CheckChargeAffordable 校验扣费前即可确定金额的费用（倍率后 USD，如显式缓存存储费）不会透支： -// 余额模式要求缓存余额覆盖 amount，订阅模式要求日 / 周 / 月剩余额度覆盖 amount， -// API Key 总额度与限速窗口同样叠加 amount 检查。须在 CheckBillingEligibility 通过后调用；简易模式不校验。 -func (s *BillingCacheService) CheckChargeAffordable(ctx context.Context, user *User, apiKey *APIKey, group *Group, subscription *UserSubscription, amount float64) error { +// ReserveFixedCharge 为扣费前即可确定金额的费用（倍率后 USD，如显式缓存存储费）做准入并登记在途预留。 +// +// 余额模式要求缓存余额覆盖 amount 及用户的其他在途预留；订阅模式要求日 / 周 / 月剩余额度覆盖 amount； +// API Key 总额度与限速窗口同样叠加 amount。订阅与 API Key 的用量取库中权威值，并叠加同一订阅 / Key 上其他 +// 在途预留：先登记本次预留再读用量，并发的两笔费用中后读用量的一方必然看到另一方（预留或已落库的扣费）。 +// 返回的句柄可能为 nil（未登记任何预留）；调用方须在 handler 结束时调用 HandlerDone，预留在计费任务落库后释放。 +// 须在 CheckBillingEligibility 通过后调用；简易模式不校验。 +func (s *BillingCacheService) ReserveFixedCharge(ctx context.Context, user *User, apiKey *APIKey, group *Group, subscription *UserSubscription, amount float64) (*InflightReservation, error) { if user == nil || amount <= 0 || math.IsNaN(amount) || math.IsInf(amount, 0) { - return nil + return nil, nil } if s.cfg != nil && s.cfg.RunMode == config.RunModeSimple { - return nil + return nil, nil } - if group != nil && group.IsSubscriptionType() && subscription != nil { - subData, err := s.GetSubscriptionStatus(ctx, user.ID, group.ID) - if err != nil { - return ErrBillingServiceUnavailable.WithCause(err) - } - if group.HasDailyLimit() && subData.DailyUsage+amount > *group.DailyLimitUSD { - return ErrDailyLimitExceeded - } - if group.HasWeeklyLimit() && subData.WeeklyUsage+amount > *group.WeeklyLimitUSD { - return ErrWeeklyLimitExceeded - } - if group.HasMonthlyLimit() && subData.MonthlyUsage+amount > *group.MonthlyLimitUSD { - return ErrMonthlyLimitExceeded - } - } else { + subscriptionBill := group != nil && group.IsSubscriptionType() && subscription != nil + if !subscriptionBill { balance, err := s.GetUserBalance(ctx, user.ID) if err != nil { - return ErrBillingServiceUnavailable.WithCause(err) + return nil, ErrBillingServiceUnavailable.WithCause(err) } if balance < amount { - return ErrInsufficientBalance + return nil, ErrInsufficientBalance } } - if apiKey != nil { - if apiKey.Quota > 0 && apiKey.QuotaUsed+amount > apiKey.Quota { + res, err := s.reserveInflight(ctx, user, group, subscription, amount, false, false) + if err != nil { + return nil, err + } + var subscriptionPending, apiKeyPending float64 + if subscriptionBill && (group.HasDailyLimit() || group.HasWeeklyLimit() || group.HasMonthlyLimit()) { + res, subscriptionPending = s.reserveInflightScope(ctx, res, user.ID, inflightSubscriptionScope(subscription.ID), amount) + } + if apiKey != nil && (apiKey.Quota > 0 || apiKey.HasRateLimits()) { + res, apiKeyPending = s.reserveInflightScope(ctx, res, user.ID, inflightAPIKeyScope(apiKey.ID), amount) + } + if subscriptionBill { + err = s.checkFixedChargeSubscription(ctx, user.ID, group, amount+subscriptionPending) + } + if err == nil && apiKey != nil { + err = s.checkFixedChargeAPIKey(ctx, apiKey, amount+apiKeyPending) + } + if err != nil { + res.Release() + return nil, err + } + return res, nil +} + +// checkFixedChargeSubscription 按库中订阅用量校验日 / 周 / 月剩余额度覆盖 pending。 +func (s *BillingCacheService) checkFixedChargeSubscription(ctx context.Context, userID int64, group *Group, pending float64) error { + read := s.GetSubscriptionStatus + if s.subRepo != nil { + read = s.getSubscriptionFromDB + } + subData, err := read(ctx, userID, group.ID) + if err != nil { + return ErrBillingServiceUnavailable.WithCause(err) + } + if group.HasDailyLimit() && subData.DailyUsage+pending > *group.DailyLimitUSD { + return ErrDailyLimitExceeded + } + if group.HasWeeklyLimit() && subData.WeeklyUsage+pending > *group.WeeklyLimitUSD { + return ErrWeeklyLimitExceeded + } + if group.HasMonthlyLimit() && subData.MonthlyUsage+pending > *group.MonthlyLimitUSD { + return ErrMonthlyLimitExceeded + } + return nil +} + +// checkFixedChargeAPIKey 按库中 Key 用量校验总额度与限速窗口覆盖 pending；未接入仓储时使用请求上下文与限速缓存中的用量。 +func (s *BillingCacheService) checkFixedChargeAPIKey(ctx context.Context, apiKey *APIKey, pending float64) error { + if apiKey.Quota <= 0 && !apiKey.HasRateLimits() { + return nil + } + if s.apiKeyRateLimitLoader == nil { + if apiKey.Quota > 0 && apiKey.QuotaUsed+pending > apiKey.Quota { return ErrAPIKeyQuotaExhausted } if apiKey.HasRateLimits() { - if err := s.checkAPIKeyRateLimitsWithPending(ctx, apiKey, amount); err != nil { - return err - } + return s.checkAPIKeyRateLimitsWithPending(ctx, apiKey, pending) } + return nil + } + data, err := s.apiKeyRateLimitLoader.GetRateLimitData(ctx, apiKey.ID) + if err != nil { + return ErrBillingServiceUnavailable.WithCause(err) + } + exceeds := func(usage, limit float64) bool { + return limit > 0 && usage+pending > limit + } + switch { + case exceeds(data.QuotaUsed, apiKey.Quota): + return ErrAPIKeyQuotaExhausted + case exceeds(data.EffectiveUsage5h(), apiKey.RateLimit5h): + return ErrAPIKeyRateLimit5hExceeded + case exceeds(data.EffectiveUsage1d(), apiKey.RateLimit1d): + return ErrAPIKeyRateLimit1dExceeded + case exceeds(data.EffectiveUsage7d(), apiKey.RateLimit7d): + return ErrAPIKeyRateLimit7dExceeded } return nil } diff --git a/backend/internal/service/billing_inflight_reservation.go b/backend/internal/service/billing_inflight_reservation.go index 714143d554..e63c9db339 100644 --- a/backend/internal/service/billing_inflight_reservation.go +++ b/backend/internal/service/billing_inflight_reservation.go @@ -3,6 +3,7 @@ package service import ( "context" "math" + "strconv" "strings" "sync" "sync/atomic" @@ -25,6 +26,16 @@ type InflightBalanceReservationCache interface { ReleaseInflightBalance(ctx context.Context, userID int64, requestID string) error } +// InflightScopeReservationCache 订阅 / API Key 维度在途预留的缓存能力（可选），用于扣费前即可确定金额的费用； +// 未实现时这两个维度不登记预留（fail-open），准入仍按库中用量校验。 +type InflightScopeReservationCache interface { + // ReserveInflightScope 原子地清理 scope 下的过期预留并登记 requestID 的预留（ttl 后自动失效）， + // 返回登记前该 scope 下其他在途预留的合计。 + ReserveInflightScope(ctx context.Context, scope, requestID string, amount float64, ttl time.Duration) (float64, error) + // ReleaseInflightScope 释放 scope 下 requestID 的预留（幂等）。 + ReleaseInflightScope(ctx context.Context, scope, requestID string) error +} + // InflightBalanceReservationRenewer 可选：续期仍存活的预留（流式长请求期间防止 TTL 过期）。 type InflightBalanceReservationRenewer interface { // RenewInflightBalance 若 requestID 仍存在则把其过期时间推迟到 now+ttl；返回是否仍存在。 @@ -54,11 +65,15 @@ const ( // 不再续期，所以最长持有时间受 TTL 约束（计费任务卡死时预留自动过期）。 // 所有方法对 nil 接收者安全。 type InflightReservation struct { + // cache 为 nil 表示未登记余额维度的预留。 cache InflightBalanceReservationCache userID int64 requestID string amount float64 ttl time.Duration + // scopes 是同一金额在订阅 / API Key 维度登记的预留（与余额维度共用 requestID），不续期。 + scopeCache InflightScopeReservationCache + scopes []string refs atomic.Int64 releaseOnce sync.Once @@ -116,8 +131,15 @@ func (r *InflightReservation) Release() { r.refs.Store(0) relCtx, relCancel := context.WithTimeout(context.Background(), inflightReservationReleaseTimeout) defer relCancel() - if err := r.cache.ReleaseInflightBalance(relCtx, r.userID, r.requestID); err != nil { - logger.LegacyPrintf("service.billing_cache", "Warning: inflight reservation release failed for user %d (expires by ttl): %v", r.userID, err) + if r.cache != nil { + if err := r.cache.ReleaseInflightBalance(relCtx, r.userID, r.requestID); err != nil { + logger.LegacyPrintf("service.billing_cache", "Warning: inflight reservation release failed for user %d (expires by ttl): %v", r.userID, err) + } + } + for _, scope := range r.scopes { + if err := r.scopeCache.ReleaseInflightScope(relCtx, scope, r.requestID); err != nil { + logger.LegacyPrintf("service.billing_cache", "Warning: inflight reservation release failed for %s (expires by ttl): %v", scope, err) + } } }) } @@ -235,10 +257,44 @@ func (s *BillingCacheService) ReserveInflight(ctx context.Context, user *User, g return s.reserveInflight(ctx, user, group, subscription, estimate, true, true) } -// ReserveInflightCharge 为扣费前即可确定金额的费用登记在途预留，金额不受 max_reservation_usd 封顶， -// 其余语义与 ReserveInflight 相同。 -func (s *BillingCacheService) ReserveInflightCharge(ctx context.Context, user *User, group *Group, subscription *UserSubscription, amount float64) (*InflightReservation, error) { - return s.reserveInflight(ctx, user, group, subscription, amount, true, false) +// inflightSubscriptionScope / inflightAPIKeyScope 是订阅 / API Key 维度在途预留的 scope 名。 +func inflightSubscriptionScope(subscriptionID int64) string { + return "sub:" + strconv.FormatInt(subscriptionID, 10) +} + +func inflightAPIKeyScope(apiKeyID int64) string { + return "apikey:" + strconv.FormatInt(apiKeyID, 10) +} + +// reserveInflightScope 在 scope 维度为 amount 登记在途预留并挂到 res 上（res 为 nil 时新建句柄）， +// 返回句柄与登记前该 scope 下其他在途预留的合计。开关关闭 / 缓存不支持 / Redis 执行失败时不登记（fail-open），合计按 0。 +func (s *BillingCacheService) reserveInflightScope(ctx context.Context, res *InflightReservation, userID int64, scope string, amount float64) (*InflightReservation, float64) { + cfg, ok := s.inflightReservationConfig() + if !ok { + return res, 0 + } + sc, ok := s.cache.(InflightScopeReservationCache) + if !ok || sc == nil { + return res, 0 + } + if res == nil { + ttl := defaultInflightReservationTTL + if cfg.TTLSeconds > 0 { + ttl = time.Duration(cfg.TTLSeconds) * time.Second + } + res = &InflightReservation{userID: userID, requestID: uuid.NewString(), amount: amount, ttl: ttl} + res.refs.Store(1) + } + reserveCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), inflightReservationReserveTimeout) + pending, err := sc.ReserveInflightScope(reserveCtx, scope, res.requestID, amount, res.ttl) + cancel() + if err != nil { + logger.LegacyPrintf("service.billing_cache", "Warning: inflight reservation failed for %s (fail-open): %v", scope, err) + return res, 0 + } + res.scopeCache = sc + res.scopes = append(res.scopes, scope) + return res, pending } func (s *BillingCacheService) reserveInflight(ctx context.Context, user *User, group *Group, subscription *UserSubscription, estimate float64, renew, capped bool) (*InflightReservation, error) { diff --git a/backend/internal/service/gemini_cached_content.go b/backend/internal/service/gemini_cached_content.go index 92157dffbb..36b3f79a71 100644 --- a/backend/internal/service/gemini_cached_content.go +++ b/backend/internal/service/gemini_cached_content.go @@ -16,12 +16,14 @@ import ( "slices" "strconv" "strings" + "sync" "time" "unicode/utf8" "github.com/Wei-Shaw/sub2api/internal/config" infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" "github.com/Wei-Shaw/sub2api/internal/pkg/geminicli" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" "github.com/tidwall/gjson" "github.com/tidwall/sjson" ) @@ -42,6 +44,17 @@ const ( geminiCachedContentUpstreamTimeout = 5 * time.Minute geminiCachedContentOfficialHost = "generativelanguage.googleapis.com" + // geminiCachedContentMutationUpstreamTimeout 是修改 / 删除上游缓存的单次调用超时。 + geminiCachedContentMutationUpstreamTimeout = 30 * time.Second + // geminiCachedContentMutationLease 是修改意图的租约时长，须覆盖持有者在租约内的全部操作 + // （一次上游修改与一次落库，或恢复流程的一次上游回滚与一次落库），租约过期后持有者不会再改动上游。 + geminiCachedContentMutationLease = 2 * time.Minute + // geminiCachedContentMutationRecoveryInterval / geminiCachedContentMutationRecoveryBatch 约束后台回滚恢复的节奏； + // geminiCachedContentMutationRecoveryTimeout 是恢复单条记录的落库超时，与一次上游回滚合计不超过租约。 + geminiCachedContentMutationRecoveryInterval = time.Minute + geminiCachedContentMutationRecoveryBatch = 50 + geminiCachedContentMutationRecoveryTimeout = time.Minute + // GeminiCachedContentListDefaultPageSize / GeminiCachedContentListMaxPageSize 约束 list 分页。 GeminiCachedContentListDefaultPageSize = 50 GeminiCachedContentListMaxPageSize = 1000 @@ -50,6 +63,10 @@ const ( var ( ErrGeminiCachedContentNotFound = infraerrors.New(http.StatusNotFound, "GEMINI_CACHED_CONTENT_NOT_FOUND", "gemini cached content not found") ErrGeminiCachedContentExists = infraerrors.New(http.StatusConflict, "GEMINI_CACHED_CONTENT_EXISTS", "gemini cached content already exists") + // ErrGeminiCachedContentMutationConflict 表示同一缓存有进行中的修改、到期时间已被其他修改改变，或修改意图已不属于调用方。 + ErrGeminiCachedContentMutationConflict = infraerrors.New(http.StatusConflict, "GEMINI_CACHED_CONTENT_MUTATION_CONFLICT", "another update to this cached content is in progress") + // ErrGeminiCachedContentRequestNotSent 表示上游缓存请求在发出前失败（地址、凭据等），上游状态未被改变。 + ErrGeminiCachedContentRequestNotSent = errors.New("gemini cached content request was not sent") geminiCachedContentPublicIDPattern = regexp.MustCompile(`^[a-z0-9]{1,64}$`) geminiCachedContentAIStudioName = regexp.MustCompile(`^cachedContents/[A-Za-z0-9_-]+$`) @@ -100,8 +117,18 @@ type GeminiCachedContentRepository interface { GetForOwner(ctx context.Context, apiKeyID int64, publicID string) (*GeminiCachedContent, error) // ListForOwner 按 id 倒序返回分组内 expire_time 晚于 activeAt 的记录；beforeID>0 时只返回 id 更小的记录。 ListForOwner(ctx context.Context, apiKeyID, groupID int64, activeAt time.Time, beforeID int64, limit int) ([]*GeminiCachedContent, error) - // UpdateExpireTime 原子地写入新的到期时间并返回写入前的值。 - UpdateExpireTime(ctx context.Context, id int64, expireTime time.Time) (time.Time, error) + // BeginMutation 在记录没有修改意图且 expire_time 仍为 baseline 时登记修改意图（token 与自数据库当前时间起的租约）， + // 返回是否登记成功；记录不存在或已删除时返回 ErrGeminiCachedContentNotFound。 + BeginMutation(ctx context.Context, id int64, baseline time.Time, token string, lease time.Duration) (bool, error) + // TakeOverStaleMutation 把记录上租约已过期的修改意图换成 token 与新租约并返回记录；没有可接管的意图时返回 nil。 + TakeOverStaleMutation(ctx context.Context, id int64, token string, lease time.Duration) (*GeminiCachedContent, error) + // CommitMutation 由 token 的持有者写入新的到期时间并清除修改意图；token 已失效或记录已删除时返回 + // ErrGeminiCachedContentMutationConflict。 + CommitMutation(ctx context.Context, id int64, token string, expireTime time.Time) error + // ReleaseMutation 由 token 的持有者清除修改意图，不改变到期时间。 + ReleaseMutation(ctx context.Context, id int64, token string) error + // ListStaleMutations 返回修改意图租约已过期的未删除记录 id。 + ListStaleMutations(ctx context.Context, limit int) ([]int64, error) SoftDelete(ctx context.Context, id int64) error } @@ -594,10 +621,48 @@ type GeminiCachedContentService struct { compat *GeminiMessagesCompatService cfg *config.Config now func() time.Time + + startOnce sync.Once + stopOnce sync.Once + stopCh chan struct{} } func NewGeminiCachedContentService(repo GeminiCachedContentRepository, compat *GeminiMessagesCompatService, cfg *config.Config) *GeminiCachedContentService { - return &GeminiCachedContentService{repo: repo, compat: compat, cfg: cfg, now: time.Now} + return &GeminiCachedContentService{repo: repo, compat: compat, cfg: cfg, now: time.Now, stopCh: make(chan struct{})} +} + +// Start 启动后台回滚恢复：定期处理租约已过期仍未清除的修改意图。 +func (s *GeminiCachedContentService) Start() { + if s == nil || s.repo == nil { + return + } + s.startOnce.Do(func() { go s.runMutationRecovery() }) +} + +// Stop 停止后台回滚恢复。 +func (s *GeminiCachedContentService) Stop() { + if s == nil { + return + } + s.stopOnce.Do(func() { close(s.stopCh) }) +} + +func (s *GeminiCachedContentService) runMutationRecovery() { + ticker := time.NewTicker(geminiCachedContentMutationRecoveryInterval) + defer ticker.Stop() + for { + select { + case <-s.stopCh: + return + case <-ticker.C: + recovered, err := s.RecoverStaleMutations(context.Background()) + if err != nil { + logger.LegacyPrintf("service.gemini_cached_content", "Warning: gemini cached content mutation recovery failed (recovered=%d): %v", recovered, err) + } else if recovered > 0 { + logger.LegacyPrintf("service.gemini_cached_content", "gemini cached content mutation recovery rolled back %d stale mutations", recovered) + } + } + } } // Now 返回服务时钟。 @@ -662,15 +727,135 @@ func (s *GeminiCachedContentService) List(ctx context.Context, apiKeyID, groupID return s.repo.ListForOwner(ctx, apiKeyID, groupID, s.now(), beforeID, limit) } -// UpdateExpireTime 更新到期时间并返回更新前库中的到期时间（并发更新按提交顺序各自拿到自己的前值）。 -func (s *GeminiCachedContentService) UpdateExpireTime(ctx context.Context, record *GeminiCachedContent, expireTime time.Time) (time.Time, error) { - previous, err := s.repo.UpdateExpireTime(ctx, record.ID, expireTime) - if err != nil { - return time.Time{}, err +// GeminiCachedContentMutation 是一次已登记修改意图、尚未提交的有效期修改。持有期间同一缓存的其他修改被拒绝； +// 未提交的修改（上游结果未知、落库失败、进程退出）在租约过期后由恢复流程把上游有效期回滚到 Baseline，不计费。 +type GeminiCachedContentMutation struct { + record *GeminiCachedContent + baseline time.Time + token string +} + +// Baseline 返回修改前的到期时间：提交前它始终是库中的到期时间。 +func (m *GeminiCachedContentMutation) Baseline() time.Time { + return m.baseline +} + +// BeginMutation 登记修改 record 有效期的意图并取得租约。同一缓存有进行中的修改，或库中到期时间已不是 +// record.ExpireTime 时返回 ErrGeminiCachedContentMutationConflict；租约已过期的遗留意图先回滚恢复再登记。 +func (s *GeminiCachedContentService) BeginMutation(ctx context.Context, record *GeminiCachedContent) (*GeminiCachedContentMutation, error) { + for attempt := 0; ; attempt++ { + token, err := NewGeminiCachedContentPublicID() + if err != nil { + return nil, err + } + acquired, err := s.repo.BeginMutation(ctx, record.ID, record.ExpireTime, token, geminiCachedContentMutationLease) + if err != nil { + return nil, err + } + if acquired { + return &GeminiCachedContentMutation{record: record, baseline: record.ExpireTime, token: token}, nil + } + if attempt > 0 { + return nil, ErrGeminiCachedContentMutationConflict + } + recovered, err := s.recoverStaleMutation(ctx, record.ID) + if err != nil { + return nil, errors.Join(ErrGeminiCachedContentMutationConflict, err) + } + if !recovered { + return nil, ErrGeminiCachedContentMutationConflict + } } - record.ExpireTime = expireTime - record.UpdatedAt = s.now() - return previous, nil +} + +// CommitMutation 写入新的到期时间并清除修改意图。修改意图已不属于 m（缓存已删除或租约已被接管）时返回 +// ErrGeminiCachedContentMutationConflict；其他失败时意图保留，由恢复流程回滚上游。 +func (s *GeminiCachedContentService) CommitMutation(ctx context.Context, m *GeminiCachedContentMutation, expireTime time.Time) error { + if err := s.repo.CommitMutation(ctx, m.record.ID, m.token, expireTime); err != nil { + return err + } + m.record.ExpireTime = expireTime + m.record.UpdatedAt = s.now() + return nil +} + +// ReleaseMutation 放弃修改并清除修改意图，用于上游明确未改变有效期的情形。 +func (s *GeminiCachedContentService) ReleaseMutation(ctx context.Context, m *GeminiCachedContentMutation) error { + return s.repo.ReleaseMutation(ctx, m.record.ID, m.token) +} + +// RecoverStaleMutations 回滚恢复一批租约已过期的修改意图，返回完成恢复的数量。 +func (s *GeminiCachedContentService) RecoverStaleMutations(ctx context.Context) (int, error) { + listCtx, cancel := context.WithTimeout(ctx, geminiCachedContentMutationRecoveryTimeout) + ids, err := s.repo.ListStaleMutations(listCtx, geminiCachedContentMutationRecoveryBatch) + cancel() + if err != nil { + return 0, err + } + recovered := 0 + var errs []error + for _, id := range ids { + select { + case <-s.stopCh: + return recovered, errors.Join(errs...) + default: + } + itemCtx, cancel := context.WithTimeout(ctx, geminiCachedContentMutationRecoveryTimeout) + ok, err := s.recoverStaleMutation(itemCtx, id) + cancel() + if err != nil { + errs = append(errs, fmt.Errorf("cached content %d: %w", id, err)) + continue + } + if ok { + recovered++ + } + } + return recovered, errors.Join(errs...) +} + +// recoverStaleMutation 接管 id 上租约已过期的修改意图，把上游有效期回滚到库中的到期时间后清除意图； +// 库中到期时间已到时删除上游缓存与本地记录。没有可接管的意图时返回 false；回滚失败时意图保留，租约过期后重试。 +func (s *GeminiCachedContentService) recoverStaleMutation(ctx context.Context, id int64) (bool, error) { + token, err := NewGeminiCachedContentPublicID() + if err != nil { + return false, err + } + record, err := s.repo.TakeOverStaleMutation(ctx, id, token, geminiCachedContentMutationLease) + if err != nil || record == nil { + return false, err + } + forget, err := s.rollbackUpstream(ctx, record) + if err != nil { + return false, err + } + if forget { + return true, s.repo.SoftDelete(ctx, record.ID) + } + return true, s.repo.ReleaseMutation(ctx, record.ID, token) +} + +// rollbackUpstream 把上游有效期改回 record.ExpireTime，剩余不足 1 秒时删除上游缓存； +// 返回 true 表示上游缓存已不存在，本地记录应移除。持有账号已不可用时不改动上游。 +func (s *GeminiCachedContentService) rollbackUpstream(ctx context.Context, record *GeminiCachedContent) (bool, error) { + account, err := s.BoundAccount(ctx, record) + if err != nil || account == nil { + return false, err + } + remaining := record.ExpireTime.Sub(s.now()) + if remaining < time.Second { + err = s.DeleteUpstream(ctx, account, record) + } else { + _, err = s.PatchUpstream(ctx, account, record, remaining) + } + var upErr *GeminiCachedContentUpstreamError + if errors.As(err, &upErr) && upErr.NotFound() { + return true, nil + } + if err != nil { + return false, err + } + return remaining < time.Second, nil } func (s *GeminiCachedContentService) Forget(ctx context.Context, record *GeminiCachedContent) error { @@ -743,7 +928,7 @@ func (s *GeminiCachedContentService) CreateUpstream(ctx context.Context, account if err != nil { return nil, err } - respBody, err := s.doUpstream(ctx, account, http.MethodPost, collectionURL, body) + respBody, err := s.doUpstream(ctx, account, http.MethodPost, collectionURL, body, geminiCachedContentUpstreamTimeout) if err != nil { return nil, err } @@ -759,17 +944,13 @@ func (s *GeminiCachedContentService) CreateUpstream(ctx context.Context, account // PatchUpstream 把上游缓存的有效期改为自现在起 ttl。 func (s *GeminiCachedContentService) PatchUpstream(ctx context.Context, account *Account, record *GeminiCachedContent, ttl time.Duration) (*GeminiCachedContentUpstream, error) { - base, err := s.apiBase(account) - if err != nil { - return nil, err - } - resourceURL, err := geminiCachedContentResourceURL(account, base, record.UpstreamName) + resourceURL, err := s.upstreamResourceURL(account, record) if err != nil { return nil, err } seconds := int64(math.Ceil(ttl.Seconds())) body, _ := json.Marshal(map[string]any{"ttl": fmt.Sprintf("%ds", seconds)}) - respBody, err := s.doUpstream(ctx, account, http.MethodPatch, resourceURL+"?updateMask=ttl", body) + respBody, err := s.doUpstream(ctx, account, http.MethodPatch, resourceURL+"?updateMask=ttl", body, geminiCachedContentMutationUpstreamTimeout) if err != nil { return nil, err } @@ -784,25 +965,38 @@ func (s *GeminiCachedContentService) PatchUpstream(ctx context.Context, account // DeleteUpstream 删除上游缓存。 func (s *GeminiCachedContentService) DeleteUpstream(ctx context.Context, account *Account, record *GeminiCachedContent) error { - base, err := s.apiBase(account) + resourceURL, err := s.upstreamResourceURL(account, record) if err != nil { return err } - resourceURL, err := geminiCachedContentResourceURL(account, base, record.UpstreamName) - if err != nil { - return err - } - _, err = s.doUpstream(ctx, account, http.MethodDelete, resourceURL, nil) + _, err = s.doUpstream(ctx, account, http.MethodDelete, resourceURL, nil, geminiCachedContentMutationUpstreamTimeout) return err } -// doUpstream 调用上游缓存接口并返回 2xx 响应体。上游调用与客户端连接脱钩：客户端断开不会中途取消 -// 已发出的创建 / 修改，调用方据结果完成落库、计费或清理。 -func (s *GeminiCachedContentService) doUpstream(ctx context.Context, account *Account, method, target string, body []byte) ([]byte, error) { - if s.compat == nil || s.compat.httpUpstream == nil { - return nil, errors.New("gemini upstream client not configured") +// upstreamResourceURL 返回 record 对应上游缓存资源的地址；失败时请求不会发出。 +func (s *GeminiCachedContentService) upstreamResourceURL(account *Account, record *GeminiCachedContent) (string, error) { + base, err := s.apiBase(account) + if err != nil { + return "", geminiCachedContentNotSent(err) } - ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), geminiCachedContentUpstreamTimeout) + resourceURL, err := geminiCachedContentResourceURL(account, base, record.UpstreamName) + if err != nil { + return "", geminiCachedContentNotSent(err) + } + return resourceURL, nil +} + +func geminiCachedContentNotSent(err error) error { + return fmt.Errorf("%w: %w", ErrGeminiCachedContentRequestNotSent, err) +} + +// doUpstream 调用上游缓存接口并返回 2xx 响应体，单次调用不超过 timeout。上游调用与客户端连接脱钩：客户端断开不会中途取消 +// 已发出的创建 / 修改，调用方据结果完成落库、计费或清理。请求发出前的失败包装为 ErrGeminiCachedContentRequestNotSent。 +func (s *GeminiCachedContentService) doUpstream(ctx context.Context, account *Account, method, target string, body []byte, timeout time.Duration) ([]byte, error) { + if s.compat == nil || s.compat.httpUpstream == nil { + return nil, geminiCachedContentNotSent(errors.New("gemini upstream client not configured")) + } + ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), timeout) defer cancel() var reader io.Reader if body != nil { @@ -810,7 +1004,7 @@ func (s *GeminiCachedContentService) doUpstream(ctx context.Context, account *Ac } req, err := http.NewRequestWithContext(ctx, method, target, reader) if err != nil { - return nil, err + return nil, geminiCachedContentNotSent(err) } if body != nil { req.Header.Set("Content-Type", "application/json") @@ -819,20 +1013,20 @@ func (s *GeminiCachedContentService) doUpstream(ctx context.Context, account *Ac case AccountTypeAPIKey: apiKey := strings.TrimSpace(account.GetCredential("api_key")) if apiKey == "" { - return nil, errors.New("gemini api_key not configured") + return nil, geminiCachedContentNotSent(errors.New("gemini api_key not configured")) } req.Header.Set("x-goog-api-key", apiKey) case AccountTypeServiceAccount: if s.compat.tokenProvider == nil { - return nil, errors.New("gemini token provider not configured") + return nil, geminiCachedContentNotSent(errors.New("gemini token provider not configured")) } accessToken, err := s.compat.tokenProvider.GetAccessToken(ctx, account) if err != nil { - return nil, err + return nil, geminiCachedContentNotSent(err) } req.Header.Set("Authorization", "Bearer "+accessToken) default: - return nil, fmt.Errorf("account type %s does not support explicit context caching", account.Type) + return nil, geminiCachedContentNotSent(fmt.Errorf("account type %s does not support explicit context caching", account.Type)) } proxyURL := "" if account.ProxyID != nil && account.Proxy != nil { diff --git a/backend/internal/service/gemini_cached_content_hardening_test.go b/backend/internal/service/gemini_cached_content_hardening_test.go index b1b893187a..f93b9424a0 100644 --- a/backend/internal/service/gemini_cached_content_hardening_test.go +++ b/backend/internal/service/gemini_cached_content_hardening_test.go @@ -295,75 +295,186 @@ func TestGeminiForwardNative_BoundCacheUpstreamGatewayErrorsSkipAccountPenalty(t } } -type chargeAffordableCacheStub struct { +type fixedChargeCacheStub struct { billingCacheWorkerStub balance float64 sub *SubscriptionCacheData rate *APIKeyRateLimitCacheData + // reserved 按维度（user:{id} / scope 名）记录在途预留。 + reserved map[string]map[string]float64 } -func (s *chargeAffordableCacheStub) GetUserBalance(context.Context, int64) (float64, error) { +func (s *fixedChargeCacheStub) GetUserBalance(context.Context, int64) (float64, error) { return s.balance, nil } -func (s *chargeAffordableCacheStub) GetSubscriptionCache(context.Context, int64, int64) (*SubscriptionCacheData, error) { +func (s *fixedChargeCacheStub) GetSubscriptionCache(context.Context, int64, int64) (*SubscriptionCacheData, error) { if s.sub == nil { return nil, errors.New("cache miss") } return s.sub, nil } -func (s *chargeAffordableCacheStub) GetAPIKeyRateLimit(context.Context, int64) (*APIKeyRateLimitCacheData, error) { +func (s *fixedChargeCacheStub) GetAPIKeyRateLimit(context.Context, int64) (*APIKeyRateLimitCacheData, error) { if s.rate == nil { return nil, errors.New("cache miss") } return s.rate, nil } -func TestCheckChargeAffordable(t *testing.T) { +func (s *fixedChargeCacheStub) reserve(key, requestID string, amount float64) float64 { + if s.reserved == nil { + s.reserved = map[string]map[string]float64{} + } + if s.reserved[key] == nil { + s.reserved[key] = map[string]float64{} + } + sum := 0.0 + for _, v := range s.reserved[key] { + sum += v + } + s.reserved[key][requestID] = amount + return sum +} + +func (s *fixedChargeCacheStub) ReserveInflightBalance(_ context.Context, userID int64, requestID string, amount, balance float64, _ time.Duration) (bool, float64, error) { + key := fmt.Sprintf("user:%d", userID) + others := s.reserve(key, requestID, amount) + if len(s.reserved[key]) > 1 && balance-others < amount { + delete(s.reserved[key], requestID) + return false, others, nil + } + return true, others, nil +} + +func (s *fixedChargeCacheStub) ReleaseInflightBalance(_ context.Context, userID int64, requestID string) error { + delete(s.reserved[fmt.Sprintf("user:%d", userID)], requestID) + return nil +} + +func (s *fixedChargeCacheStub) ReserveInflightScope(_ context.Context, scope, requestID string, amount float64, _ time.Duration) (float64, error) { + return s.reserve(scope, requestID, amount), nil +} + +func (s *fixedChargeCacheStub) ReleaseInflightScope(_ context.Context, scope, requestID string) error { + delete(s.reserved[scope], requestID) + return nil +} + +type fixedChargeAPIKeyRepoStub struct { + APIKeyRepository + data *APIKeyRateLimitData +} + +func (r *fixedChargeAPIKeyRepoStub) GetRateLimitData(context.Context, int64) (*APIKeyRateLimitData, error) { + return r.data, nil +} + +type fixedChargeSubscriptionRepoStub struct { + UserSubscriptionRepository + sub *UserSubscription +} + +func (r *fixedChargeSubscriptionRepoStub) GetActiveByUserIDAndGroupID(context.Context, int64, int64) (*UserSubscription, error) { + return r.sub, nil +} + +func TestReserveFixedCharge(t *testing.T) { ctx := context.Background() user := &User{ID: 1} - newSvc := func(cache *chargeAffordableCacheStub) *BillingCacheService { + newSvc := func(cache *fixedChargeCacheStub) *BillingCacheService { svc := NewBillingCacheService(cache, nil, nil, nil, nil, nil, &config.Config{}, nil) t.Cleanup(svc.Stop) return svc } + reservingSvc := func(cache *fixedChargeCacheStub, subRepo UserSubscriptionRepository, apiKeyRepo APIKeyRepository) *BillingCacheService { + cfg := &config.Config{} + cfg.Billing.InflightReservation.Enabled = true + svc := NewBillingCacheService(cache, nil, subRepo, apiKeyRepo, nil, nil, cfg, nil) + t.Cleanup(svc.Stop) + return svc + } + reserve := func(svc *BillingCacheService, apiKey *APIKey, group *Group, sub *UserSubscription, amount float64) error { + res, err := svc.ReserveFixedCharge(ctx, user, apiKey, group, sub, amount) + res.Release() + return err + } t.Run("balance must cover the amount", func(t *testing.T) { - svc := newSvc(&chargeAffordableCacheStub{balance: 10}) - require.NoError(t, svc.CheckChargeAffordable(ctx, user, nil, nil, nil, 10)) - require.ErrorIs(t, svc.CheckChargeAffordable(ctx, user, nil, nil, nil, 10.01), ErrInsufficientBalance) - require.NoError(t, svc.CheckChargeAffordable(ctx, user, nil, nil, nil, 0)) + svc := newSvc(&fixedChargeCacheStub{balance: 10}) + require.NoError(t, reserve(svc, nil, nil, nil, 10)) + require.ErrorIs(t, reserve(svc, nil, nil, nil, 10.01), ErrInsufficientBalance) + require.NoError(t, reserve(svc, nil, nil, nil, 0)) }) t.Run("subscription windows must cover the amount", func(t *testing.T) { daily := 5.0 group := &Group{ID: 2, SubscriptionType: SubscriptionTypeSubscription, DailyLimitUSD: &daily} - svc := newSvc(&chargeAffordableCacheStub{sub: &SubscriptionCacheData{ + svc := newSvc(&fixedChargeCacheStub{sub: &SubscriptionCacheData{ Status: SubscriptionStatusActive, ExpiresAt: time.Now().Add(time.Hour), DailyUsage: 3, }}) sub := &UserSubscription{ID: 1} - require.NoError(t, svc.CheckChargeAffordable(ctx, user, nil, group, sub, 2)) - require.ErrorIs(t, svc.CheckChargeAffordable(ctx, user, nil, group, sub, 2.5), ErrDailyLimitExceeded) + require.NoError(t, reserve(svc, nil, group, sub, 2)) + require.ErrorIs(t, reserve(svc, nil, group, sub, 2.5), ErrDailyLimitExceeded) }) t.Run("api key quota and rate windows include the amount", func(t *testing.T) { - svc := newSvc(&chargeAffordableCacheStub{balance: 100, rate: &APIKeyRateLimitCacheData{ + svc := newSvc(&fixedChargeCacheStub{balance: 100, rate: &APIKeyRateLimitCacheData{ Usage1d: 8, Window1d: time.Now().Unix(), }}) quotaKey := &APIKey{ID: 5, Quota: 10, QuotaUsed: 9} - require.NoError(t, svc.CheckChargeAffordable(ctx, user, quotaKey, nil, nil, 1)) - require.ErrorIs(t, svc.CheckChargeAffordable(ctx, user, quotaKey, nil, nil, 1.5), ErrAPIKeyQuotaExhausted) + require.NoError(t, reserve(svc, quotaKey, nil, nil, 1)) + require.ErrorIs(t, reserve(svc, quotaKey, nil, nil, 1.5), ErrAPIKeyQuotaExhausted) windowKey := &APIKey{ID: 6, RateLimit1d: 10} - require.NoError(t, svc.CheckChargeAffordable(ctx, user, windowKey, nil, nil, 2)) - require.ErrorIs(t, svc.CheckChargeAffordable(ctx, user, windowKey, nil, nil, 2.5), ErrAPIKeyRateLimit1dExceeded) + require.NoError(t, reserve(svc, windowKey, nil, nil, 2)) + require.ErrorIs(t, reserve(svc, windowKey, nil, nil, 2.5), ErrAPIKeyRateLimit1dExceeded) + }) + + t.Run("usage is read from the database instead of request snapshots", func(t *testing.T) { + now := time.Now() + apiKeyRepo := &fixedChargeAPIKeyRepoStub{data: &APIKeyRateLimitData{QuotaUsed: 9.5, Usage1d: 8, Window1dStart: &now}} + svc := reservingSvc(&fixedChargeCacheStub{balance: 100}, nil, apiKeyRepo) + require.ErrorIs(t, reserve(svc, &APIKey{ID: 5, Quota: 10}, nil, nil, 1), ErrAPIKeyQuotaExhausted, "请求上下文中的 QuotaUsed 可能滞后") + require.ErrorIs(t, reserve(svc, &APIKey{ID: 6, RateLimit1d: 10}, nil, nil, 2.5), ErrAPIKeyRateLimit1dExceeded) + + daily := 5.0 + group := &Group{ID: 2, SubscriptionType: SubscriptionTypeSubscription, DailyLimitUSD: &daily} + subRepo := &fixedChargeSubscriptionRepoStub{sub: &UserSubscription{Status: SubscriptionStatusActive, DailyUsageUSD: 4}} + svc = reservingSvc(&fixedChargeCacheStub{sub: &SubscriptionCacheData{Status: SubscriptionStatusActive}}, subRepo, nil) + require.ErrorIs(t, reserve(svc, nil, group, &UserSubscription{ID: 1}, 2), ErrDailyLimitExceeded, "订阅缓存用量可能滞后") + }) + + t.Run("unsettled charges hold subscription and api key headroom", func(t *testing.T) { + daily := 110.0 + group := &Group{ID: 2, SubscriptionType: SubscriptionTypeSubscription, DailyLimitUSD: &daily} + sub := &UserSubscription{ID: 1} + cache := &fixedChargeCacheStub{balance: 1000, sub: &SubscriptionCacheData{Status: SubscriptionStatusActive}} + svc := reservingSvc(cache, nil, nil) + first, err := svc.ReserveFixedCharge(ctx, user, nil, group, sub, 103.5) + require.NoError(t, err) + _, err = svc.ReserveFixedCharge(ctx, user, nil, group, sub, 103.5) + require.ErrorIs(t, err, ErrDailyLimitExceeded) + first.Release() + second, err := svc.ReserveFixedCharge(ctx, user, nil, group, sub, 103.5) + require.NoError(t, err) + second.Release() + + apiKey := &APIKey{ID: 5, Quota: 110} + first, err = svc.ReserveFixedCharge(ctx, user, apiKey, nil, nil, 103.5) + require.NoError(t, err) + _, err = svc.ReserveFixedCharge(ctx, user, apiKey, nil, nil, 103.5) + require.ErrorIs(t, err, ErrAPIKeyQuotaExhausted) + first.Release() + for key, reservations := range cache.reserved { + require.Empty(t, reservations, "拒绝与释放都不留下预留：%s", key) + } }) t.Run("simple mode skips the check", func(t *testing.T) { cfg := &config.Config{RunMode: config.RunModeSimple} - svc := NewBillingCacheService(&chargeAffordableCacheStub{}, nil, nil, nil, nil, nil, cfg, nil) + svc := NewBillingCacheService(&fixedChargeCacheStub{}, nil, nil, nil, nil, nil, cfg, nil) t.Cleanup(svc.Stop) - require.NoError(t, svc.CheckChargeAffordable(ctx, user, nil, nil, nil, 1000)) + require.NoError(t, reserve(svc, nil, nil, nil, 1000)) }) } diff --git a/backend/internal/service/gemini_cached_content_test.go b/backend/internal/service/gemini_cached_content_test.go index 94e1027e23..e3afbdd5e6 100644 --- a/backend/internal/service/gemini_cached_content_test.go +++ b/backend/internal/service/gemini_cached_content_test.go @@ -53,15 +53,30 @@ func (r *geminiCachedContentRepoStub) ListForOwner(_ context.Context, apiKeyID, return out, nil } -func (r *geminiCachedContentRepoStub) UpdateExpireTime(_ context.Context, id int64, expireTime time.Time) (time.Time, error) { +func (r *geminiCachedContentRepoStub) BeginMutation(context.Context, int64, time.Time, string, time.Duration) (bool, error) { + return true, nil +} + +func (r *geminiCachedContentRepoStub) TakeOverStaleMutation(context.Context, int64, string, time.Duration) (*GeminiCachedContent, error) { + return nil, nil +} + +func (r *geminiCachedContentRepoStub) CommitMutation(_ context.Context, id int64, _ string, expireTime time.Time) error { for _, record := range r.records { if record.ID == id { - previous := record.ExpireTime record.ExpireTime = expireTime - return previous, nil + return nil } } - return time.Time{}, ErrGeminiCachedContentNotFound + return ErrGeminiCachedContentMutationConflict +} + +func (r *geminiCachedContentRepoStub) ReleaseMutation(context.Context, int64, string) error { + return nil +} + +func (r *geminiCachedContentRepoStub) ListStaleMutations(context.Context, int) ([]int64, error) { + return nil, nil } func (r *geminiCachedContentRepoStub) SoftDelete(_ context.Context, id int64) error { diff --git a/backend/internal/service/wire.go b/backend/internal/service/wire.go index 56fe96f52d..688d729b31 100644 --- a/backend/internal/service/wire.go +++ b/backend/internal/service/wire.go @@ -640,6 +640,13 @@ func ProvideIdempotencyCleanupService(repo IdempotencyRepository, cfg *config.Co return svc } +// ProvideGeminiCachedContentService creates GeminiCachedContentService and starts its mutation recovery loop. +func ProvideGeminiCachedContentService(repo GeminiCachedContentRepository, compat *GeminiMessagesCompatService, cfg *config.Config) *GeminiCachedContentService { + svc := NewGeminiCachedContentService(repo, compat, cfg) + svc.Start() + return svc +} + // ProvideScheduledTestService creates ScheduledTestService. func ProvideScheduledTestService( planRepo ScheduledTestPlanRepository, @@ -885,7 +892,7 @@ var ProviderSet = wire.NewSet( ProvideGrokOAuthService, wire.Bind(new(GrokOAuthTokenService), new(*GrokOAuthService)), NewGeminiOAuthService, - NewGeminiCachedContentService, + ProvideGeminiCachedContentService, NewGeminiQuotaService, NewCompositeTokenCacheInvalidator, wire.Bind(new(TokenCacheInvalidator), new(*CompositeTokenCacheInvalidator)), diff --git a/backend/migrations/245_gemini_cached_contents_mutation_lease.sql b/backend/migrations/245_gemini_cached_contents_mutation_lease.sql new file mode 100644 index 0000000000..e0970efebf --- /dev/null +++ b/backend/migrations/245_gemini_cached_contents_mutation_lease.sql @@ -0,0 +1,10 @@ +-- Gemini 显式缓存修改有效期的进行中操作（修改意图）：mutation_token 标识持有者， +-- mutation_lease_until 之前同一缓存的其他修改被拒绝；expire_time 在提交前保持修改前的值， +-- 租约过期仍未清除的意图由恢复流程把上游有效期回滚到 expire_time。 +ALTER TABLE gemini_cached_contents + ADD COLUMN IF NOT EXISTS mutation_token VARCHAR(64) NULL, + ADD COLUMN IF NOT EXISTS mutation_lease_until TIMESTAMPTZ NULL; + +CREATE INDEX IF NOT EXISTS idx_gemini_cached_contents_pending_mutation + ON gemini_cached_contents (mutation_lease_until) + WHERE mutation_token IS NOT NULL AND deleted_at IS NULL; From 43a3c197cb7663582ad1b6b21c761b84776193c9 Mon Sep 17 00:00:00 2001 From: feeeei Date: Tue, 13 Oct 2026 00:54:27 +0800 Subject: [PATCH 15/15] =?UTF-8?q?fix(gemini):=20=E6=98=BE=E5=BC=8F?= =?UTF-8?q?=E7=BC=93=E5=AD=98=E6=81=A2=E5=A4=8D=E6=B5=81=E7=A8=8B=E5=8C=BA?= =?UTF-8?q?=E5=88=86=E8=B4=A6=E5=8F=B7=E5=B7=B2=E5=88=A0=E9=99=A4=E4=B8=8E?= =?UTF-8?q?=E6=9A=82=E4=B8=8D=E6=94=AF=E6=8C=81=E7=BC=93=E5=AD=98=EF=BC=8C?= =?UTF-8?q?=E4=B8=8A=E6=B8=B8=E5=BF=85=E5=B7=B2=E8=BF=87=E6=9C=9F=E5=90=8E?= =?UTF-8?q?=E4=BD=9C=E5=BA=9F=E6=84=8F=E5=9B=BE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 后台回滚在持有账号返回 nil 时把修改意图当作已回滚清除：账号临时关闭显式缓存能力后 意图丢失，重新开启也无法补偿。 - boundAccount 区分两种情形：账号已删除返回 nil，账号存在但不支持显式缓存返回 errGeminiCachedContentAccountUnsupported；BoundAccount 对外行为不变 - 账号暂不支持缓存时不改动上游、保留意图，租约过期后重试（期间同一缓存的修改返回 409）； 账号已删除时没有凭据可回滚上游，软删本地记录 - 修改前到期时间加上单次最长有效期与 1 小时容忍都已过去时上游缓存必已自行过期： 账号不支持缓存或上游回滚持续失败的意图不再重试，直接作废本地记录 - 测试夹具 GetForOwner 过滤已软删记录 --- .../gemini_cached_contents_billing_test.go | 165 +++++++++++++++++- .../internal/service/gemini_cached_content.go | 58 +++++- 2 files changed, 217 insertions(+), 6 deletions(-) diff --git a/backend/internal/handler/gemini_cached_contents_billing_test.go b/backend/internal/handler/gemini_cached_contents_billing_test.go index fdf0faea93..c933421e9b 100644 --- a/backend/internal/handler/gemini_cached_contents_billing_test.go +++ b/backend/internal/handler/gemini_cached_contents_billing_test.go @@ -62,7 +62,7 @@ func (r *cachedBillingRepo) GetForOwner(_ context.Context, apiKeyID int64, publi r.mu.Lock() defer r.mu.Unlock() record, ok := r.records[publicID] - if !ok || record.APIKeyID != apiKeyID { + if !ok || record.APIKeyID != apiKeyID || r.liveLocked(record.ID) == nil { return nil, service.ErrGeminiCachedContentNotFound } cp := *record @@ -790,3 +790,166 @@ func TestGeminiCachedContentsDelete_UpstreamNotFoundStillForgets(t *testing.T) { }) } } + +// 持有账号暂时不支持显式缓存（能力被关闭）时，恢复流程不能把修改意图当作已回滚而清除； +// 能力恢复后再回滚上游并清除意图。 +func TestGeminiCachedContentMutationRecovery_KeepsIntentWhileAccountCannotCache(t *testing.T) { + env := newCachedBillingEnv(t, 1000, 900) + env.accounts[900].Credentials["base_url"] = "https://gw.example.com" + env.accounts[900].Credentials[service.GeminiExplicitCacheUpstreamCredentialKey] = true + baseline := time.Now().Add(time.Hour) + env.seed(baseline) + env.repo.failCommit = true + require.Equal(t, http.StatusInternalServerError, env.patch().Code) + require.True(t, env.repo.pendingMutation(100)) + require.Equal(t, 1, env.upstream.methodCalls(http.MethodPatch)) + env.repo.failCommit = false + env.repo.expireLease(100) + + env.accounts[900].Credentials[service.GeminiExplicitCacheUpstreamCredentialKey] = false + recovered, err := env.h.geminiCachedContentService.RecoverStaleMutations(context.Background()) + require.Error(t, err, "账号暂不支持缓存时恢复失败") + require.Zero(t, recovered) + require.True(t, env.repo.pendingMutation(100), "未回滚上游前保留修改意图") + require.Equal(t, 1, env.upstream.methodCalls(http.MethodPatch), "不改动上游") + require.Empty(t, env.repo.deleted) + + env.accounts[900].Credentials[service.GeminiExplicitCacheUpstreamCredentialKey] = true + env.repo.expireLease(100) + recovered, err = env.h.geminiCachedContentService.RecoverStaleMutations(context.Background()) + require.NoError(t, err) + require.Equal(t, 1, recovered) + require.Equal(t, 2, env.upstream.methodCalls(http.MethodPatch)) + requireTTLNear(t, env.upstream.patchTTLs[1], time.Until(baseline)) + require.False(t, env.repo.pendingMutation(100)) + require.True(t, env.repo.records["abc"].ExpireTime.Equal(baseline)) + require.Empty(t, env.usage.totals()) +} + +// 持有账号长期不支持显式缓存，且修改前到期时间加上单次最长有效期都已过去：上游缓存必已自行过期， +// 恢复流程不再保留意图，直接移除本地记录。 +func TestGeminiCachedContentMutationRecovery_AbandonsIntentOnceUpstreamMustHaveExpired(t *testing.T) { + env := newCachedBillingEnv(t, 1000, 900) + env.accounts[900].Credentials["base_url"] = "https://gw.example.com" + env.accounts[900].Credentials[service.GeminiExplicitCacheUpstreamCredentialKey] = true + env.seed(time.Now().Add(time.Hour)) + env.repo.failCommit = true + require.Equal(t, http.StatusInternalServerError, env.patch().Code) + env.repo.failCommit = false + env.repo.expireLease(100) + env.accounts[900].Credentials[service.GeminiExplicitCacheUpstreamCredentialKey] = false + + // 修改前到期时间 + 24h 上限 + 1h 容忍仍未过去：保留意图。 + env.repo.records["abc"].ExpireTime = time.Now().Add(-20 * time.Hour) + recovered, err := env.h.geminiCachedContentService.RecoverStaleMutations(context.Background()) + require.Error(t, err) + require.Zero(t, recovered) + require.True(t, env.repo.pendingMutation(100)) + require.Empty(t, env.repo.deleted) + + env.repo.expireLease(100) + env.repo.records["abc"].ExpireTime = time.Now().Add(-26 * time.Hour) + recovered, err = env.h.geminiCachedContentService.RecoverStaleMutations(context.Background()) + require.NoError(t, err) + require.Equal(t, 1, recovered) + require.Equal(t, 1, env.upstream.methodCalls(http.MethodPatch), "不改动上游") + require.Zero(t, env.upstream.methodCalls(http.MethodDelete)) + require.Equal(t, []int64{100}, env.repo.deleted) + require.Empty(t, env.usage.totals()) + + recovered, err = env.h.geminiCachedContentService.RecoverStaleMutations(context.Background()) + require.NoError(t, err) + require.Zero(t, recovered, "作废的记录不再被选中") +} + +// 作废判定的边界：默认上限 24h 加 1h 容忍；配置上限低于 24h 时仍按 24h 兜底。 +func TestGeminiCachedContentMutationRecovery_AbandonBoundary(t *testing.T) { + disableCache := func(env *cachedBillingEnv) { + env.accounts[900].Credentials["base_url"] = "https://gw.example.com" + env.accounts[900].Credentials[service.GeminiExplicitCacheUpstreamCredentialKey] = false + } + t.Run("within grace keeps intent", func(t *testing.T) { + env := newCachedBillingEnv(t, 1000, 900) + env.seed(time.Now().Add(-25*time.Hour + time.Minute)) + env.repo.mutations[100] = cachedBillingMutation{token: "stale", leaseUntil: time.Now().Add(-time.Second)} + disableCache(env) + + recovered, err := env.h.geminiCachedContentService.RecoverStaleMutations(context.Background()) + require.Error(t, err) + require.Zero(t, recovered) + require.True(t, env.repo.pendingMutation(100)) + require.Empty(t, env.repo.deleted) + }) + t.Run("configured max ttl below default still waits a day", func(t *testing.T) { + env := newCachedBillingEnv(t, 1000, 900) + env.seed(time.Now().Add(-12 * time.Hour)) + env.repo.mutations[100] = cachedBillingMutation{token: "stale", leaseUntil: time.Now().Add(-time.Second)} + disableCache(env) + cfg := &config.Config{} + cfg.Gateway.GeminiCachedContentMaxTTLSeconds = 3600 + compat := service.NewGeminiMessagesCompatService(&cachedBillingAccountRepo{accounts: env.accounts}, nil, nil, nil, nil, nil, env.upstream, nil, cfg) + svc := service.NewGeminiCachedContentService(env.repo, compat, cfg) + + recovered, err := svc.RecoverStaleMutations(context.Background()) + require.Error(t, err) + require.Zero(t, recovered) + require.True(t, env.repo.pendingMutation(100)) + + env.repo.expireLease(100) + env.repo.records["abc"].ExpireTime = time.Now().Add(-26 * time.Hour) + recovered, err = svc.RecoverStaleMutations(context.Background()) + require.NoError(t, err) + require.Equal(t, 1, recovered) + require.Equal(t, []int64{100}, env.repo.deleted) + }) +} + +// 持有账号可用但上游回滚持续失败（如 Key 已失效）：上游缓存必已自行过期后不再重试，移除本地记录。 +func TestGeminiCachedContentMutationRecovery_AbandonsIntentWhenRollbackKeepsFailing(t *testing.T) { + env := newCachedBillingEnv(t, 1000, 900) + env.seed(time.Now().Add(-time.Minute)) + env.repo.mutations[100] = cachedBillingMutation{token: "stale", leaseUntil: time.Now().Add(-time.Second)} + env.upstream.failures = map[string]cachedBillingUpstreamFailure{http.MethodDelete: {status: http.StatusUnauthorized, + body: `{"error":{"code":401,"message":"API key not valid","status":"UNAUTHENTICATED"}}`}} + + recovered, err := env.h.geminiCachedContentService.RecoverStaleMutations(context.Background()) + require.Error(t, err, "上游缓存可能仍存活：保留意图等待重试") + require.Zero(t, recovered) + require.Equal(t, 1, env.upstream.methodCalls(http.MethodDelete)) + require.True(t, env.repo.pendingMutation(100)) + require.Empty(t, env.repo.deleted) + + env.repo.expireLease(100) + env.repo.records["abc"].ExpireTime = time.Now().Add(-26 * time.Hour) + recovered, err = env.h.geminiCachedContentService.RecoverStaleMutations(context.Background()) + require.NoError(t, err) + require.Equal(t, 1, recovered) + require.Equal(t, 2, env.upstream.methodCalls(http.MethodDelete), "账号可用时仍先尝试回滚") + require.Equal(t, []int64{100}, env.repo.deleted) + require.Empty(t, env.usage.totals()) +} + +// 持有账号已删除时没有凭据可回滚上游：恢复流程移除本地记录并清除意图，不再重试。 +func TestGeminiCachedContentMutationRecovery_ForgetsRecordWhenAccountDeleted(t *testing.T) { + env := newCachedBillingEnv(t, 1000, 900) + env.seed(time.Now().Add(time.Hour)) + env.repo.failCommit = true + require.Equal(t, http.StatusInternalServerError, env.patch().Code) + require.True(t, env.repo.pendingMutation(100)) + env.repo.failCommit = false + env.repo.expireLease(100) + + delete(env.accounts, 900) + recovered, err := env.h.geminiCachedContentService.RecoverStaleMutations(context.Background()) + require.NoError(t, err) + require.Equal(t, 1, recovered) + require.Equal(t, 1, env.upstream.methodCalls(http.MethodPatch), "无凭据，不改动上游") + require.Zero(t, env.upstream.methodCalls(http.MethodDelete)) + require.Equal(t, []int64{100}, env.repo.deleted, "本地记录作废") + require.Empty(t, env.usage.totals()) + require.Equal(t, http.StatusForbidden, env.serve(http.MethodGet, "/v1beta/cachedContents/abc", "", gin.Params{{Key: "cacheID", Value: "abc"}}, env.h.GeminiCachedContentsGet).Code, "作废的记录对客户端不可见") + + recovered, err = env.h.geminiCachedContentService.RecoverStaleMutations(context.Background()) + require.NoError(t, err) + require.Zero(t, recovered, "不再重复处理") +} diff --git a/backend/internal/service/gemini_cached_content.go b/backend/internal/service/gemini_cached_content.go index 36b3f79a71..bc1a3353dc 100644 --- a/backend/internal/service/gemini_cached_content.go +++ b/backend/internal/service/gemini_cached_content.go @@ -54,6 +54,8 @@ const ( geminiCachedContentMutationRecoveryInterval = time.Minute geminiCachedContentMutationRecoveryBatch = 50 geminiCachedContentMutationRecoveryTimeout = time.Minute + // geminiCachedContentMutationAbandonGrace 是判定"上游缓存必已自行过期"时额外容忍的网关与上游时钟偏差。 + geminiCachedContentMutationAbandonGrace = time.Hour // GeminiCachedContentListDefaultPageSize / GeminiCachedContentListMaxPageSize 约束 list 分页。 GeminiCachedContentListDefaultPageSize = 50 @@ -67,6 +69,8 @@ var ( ErrGeminiCachedContentMutationConflict = infraerrors.New(http.StatusConflict, "GEMINI_CACHED_CONTENT_MUTATION_CONFLICT", "another update to this cached content is in progress") // ErrGeminiCachedContentRequestNotSent 表示上游缓存请求在发出前失败（地址、凭据等），上游状态未被改变。 ErrGeminiCachedContentRequestNotSent = errors.New("gemini cached content request was not sent") + // errGeminiCachedContentAccountUnsupported 表示持有缓存的账号存在，但当前不支持显式缓存（能力被关闭或认证方式不支持）。 + errGeminiCachedContentAccountUnsupported = errors.New("gemini cached content account does not support explicit caching") geminiCachedContentPublicIDPattern = regexp.MustCompile(`^[a-z0-9]{1,64}$`) geminiCachedContentAIStudioName = regexp.MustCompile(`^cachedContents/[A-Za-z0-9_-]+$`) @@ -815,7 +819,8 @@ func (s *GeminiCachedContentService) RecoverStaleMutations(ctx context.Context) } // recoverStaleMutation 接管 id 上租约已过期的修改意图，把上游有效期回滚到库中的到期时间后清除意图； -// 库中到期时间已到时删除上游缓存与本地记录。没有可接管的意图时返回 false；回滚失败时意图保留，租约过期后重试。 +// 库中到期时间已到时删除上游缓存与本地记录。没有可接管的意图时返回 false；回滚失败时意图保留， +// 本次接管取得的新租约（geminiCachedContentMutationLease）过期后重试，期间同一缓存的修改返回 409。 func (s *GeminiCachedContentService) recoverStaleMutation(ctx context.Context, id int64) (bool, error) { token, err := NewGeminiCachedContentPublicID() if err != nil { @@ -836,12 +841,25 @@ func (s *GeminiCachedContentService) recoverStaleMutation(ctx context.Context, i } // rollbackUpstream 把上游有效期改回 record.ExpireTime，剩余不足 1 秒时删除上游缓存； -// 返回 true 表示上游缓存已不存在，本地记录应移除。持有账号已不可用时不改动上游。 +// 返回 true 表示上游缓存已不存在、或已没有凭据可触达（持有账号已删除），本地记录应移除。 +// 持有账号存在但当前不支持显式缓存时不改动上游，修改意图保留到账号恢复后重试；账号可用时总是先向上游回滚。 +// 两种情况下只要 mutationUpstreamMustHaveExpired 成立（上游缓存必已自行过期），本地记录直接作废，不再重试。 func (s *GeminiCachedContentService) rollbackUpstream(ctx context.Context, record *GeminiCachedContent) (bool, error) { - account, err := s.BoundAccount(ctx, record) - if err != nil || account == nil { + account, err := s.boundAccount(ctx, record) + if errors.Is(err, errGeminiCachedContentAccountUnsupported) { + if s.mutationUpstreamMustHaveExpired(record) { + logger.LegacyPrintf("service.gemini_cached_content", "Warning: gemini cached content %d (upstream %s) has a pending mutation but account %d stopped supporting explicit caching and the upstream cache has expired since; dropping the local record", record.ID, record.UpstreamName, record.AccountID) + return true, nil + } + return false, fmt.Errorf("%w (account %d)", err, record.AccountID) + } + if err != nil { return false, err } + if account == nil { + logger.LegacyPrintf("service.gemini_cached_content", "Warning: gemini cached content %d (upstream %s) has a pending mutation but account %d no longer exists; dropping the local record without touching upstream", record.ID, record.UpstreamName, record.AccountID) + return true, nil + } remaining := record.ExpireTime.Sub(s.now()) if remaining < time.Second { err = s.DeleteUpstream(ctx, account, record) @@ -853,17 +871,44 @@ func (s *GeminiCachedContentService) rollbackUpstream(ctx context.Context, recor return true, nil } if err != nil { + if s.mutationUpstreamMustHaveExpired(record) { + logger.LegacyPrintf("service.gemini_cached_content", "Warning: gemini cached content %d (upstream %s) rollback keeps failing on account %d but the upstream cache has expired since; dropping the local record: %v", record.ID, record.UpstreamName, record.AccountID, err) + return true, nil + } return false, err } return remaining < time.Second, nil } +// mutationUpstreamMustHaveExpired 判断 record 上未完成修改所涉及的上游缓存是否必已自行过期： +// 修改发生时记录尚未到期（Resolve 拒绝已过期记录），且单次延长不超过 MaxTTL， +// 所以上游有效期最晚为 record.ExpireTime + MaxTTL；再过 geminiCachedContentMutationAbandonGrace 后判定为已过期。 +// MaxTTL 取判定时刻的配置并以默认上限兜底：配置被调低后判定可能早于上游真实过期，此时只影响运营方的上游存储费， +// 不涉及用户计费。 +func (s *GeminiCachedContentService) mutationUpstreamMustHaveExpired(record *GeminiCachedContent) bool { + maxTTL := s.MaxTTL() + if maxTTL < geminiCachedContentDefaultMaxTTL { + maxTTL = geminiCachedContentDefaultMaxTTL + } + return s.now().After(record.ExpireTime.Add(maxTTL + geminiCachedContentMutationAbandonGrace)) +} + func (s *GeminiCachedContentService) Forget(ctx context.Context, record *GeminiCachedContent) error { return s.repo.SoftDelete(ctx, record.ID) } // BoundAccount 返回持有缓存的账号；账号已删除或不再支持显式缓存时返回 nil。 func (s *GeminiCachedContentService) BoundAccount(ctx context.Context, record *GeminiCachedContent) (*Account, error) { + account, err := s.boundAccount(ctx, record) + if errors.Is(err, errGeminiCachedContentAccountUnsupported) { + return nil, nil + } + return account, err +} + +// boundAccount 返回持有缓存的账号：账号已删除返回 nil；账号存在但当前不支持显式缓存返回 +// errGeminiCachedContentAccountUnsupported，调用方据此区分"永久不可用"与"暂时不可用"。 +func (s *GeminiCachedContentService) boundAccount(ctx context.Context, record *GeminiCachedContent) (*Account, error) { if s.compat == nil || s.compat.accountRepo == nil { return nil, errors.New("account repository not configured") } @@ -874,9 +919,12 @@ func (s *GeminiCachedContentService) BoundAccount(ctx context.Context, record *G } return nil, err } - if account == nil || !IsGeminiExplicitCacheAccount(account) { + if account == nil { return nil, nil } + if !IsGeminiExplicitCacheAccount(account) { + return nil, errGeminiCachedContentAccountUnsupported + } return account, nil }