From d5a802f84a6ae3405bee2a6e0160675785abf237 Mon Sep 17 00:00:00 2001 From: feeeei Date: Sat, 10 Oct 2026 22:21:28 +0800 Subject: [PATCH] =?UTF-8?q?feat(gemini):=20=E6=94=AF=E6=8C=81=E6=98=BE?= =?UTF-8?q?=E5=BC=8F=E4=B8=8A=E4=B8=8B=E6=96=87=E7=BC=93=E5=AD=98=EF=BC=88?= =?UTF-8?q?AI=20Studio=20/=20Vertex=20/=20=E8=87=AA=E5=BB=BA=E7=BD=91?= =?UTF-8?q?=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 + } +}