diff --git a/backend/cmd/server/wire.go b/backend/cmd/server/wire.go index 06674251af..6834ad78df 100644 --- a/backend/cmd/server/wire.go +++ b/backend/cmd/server/wire.go @@ -104,6 +104,7 @@ func provideCleanup( subscriptionExpiry *service.SubscriptionExpiryService, usageCleanup *service.UsageCleanupService, idempotencyCleanup *service.IdempotencyCleanupService, + geminiCachedContent *service.GeminiCachedContentService, batchImageCleanup *service.BatchImageCleanupService, batchImageWorker *service.BatchImageWorkerRuntime, pricing *service.PricingService, @@ -244,6 +245,12 @@ func provideCleanup( } return nil }}, + {"GeminiCachedContentService", func() error { + if geminiCachedContent != nil { + geminiCachedContent.Stop() + } + return nil + }}, {"BatchImageCleanupService", func() error { if batchImageCleanup != nil { batchImageCleanup.Stop() diff --git a/backend/cmd/server/wire_gen.go b/backend/cmd/server/wire_gen.go index 1eda60da76..53cef32052 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.ProvideGeminiCachedContentService(geminiCachedContentRepository, geminiMessagesCompatService, configConfig) + gatewayHandler := handler.ProvideGatewayHandler(gatewayService, openAIGatewayService, geminiMessagesCompatService, antigravityGatewayService, userService, concurrencyService, billingCacheService, usageService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, userMessageQueueService, configConfig, settingService, coordinator, geminiCachedContentService) openAIGatewayHandler := handler.ProvideOpenAIGatewayHandler(openAIGatewayService, pluginManager, concurrencyService, billingCacheService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, opsService, grokQuotaService, configConfig, coordinator, compositeRouteResolver) handlerSettingHandler := handler.ProvideSettingHandler(settingService, buildInfo, notificationEmailService) totpHandler := handler.NewTotpHandler(totpService) @@ -353,7 +355,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { channelMonitorRunner := service.ProvideChannelMonitorRunner(channelMonitorService, settingService, channelMonitorQuotaFetcher) channelMonitorV2Aggregator := service.ProvideChannelMonitorV2Aggregator(channelMonitorV2Repository, db, settingService) userPlatformQuotaUsageFlusher := service.ProvideUserPlatformQuotaUsageFlusher(configConfig, billingCache, serviceUserPlatformQuotaRepository, timingWheelService) - v := provideCleanup(client, redisClient, opsMetricsCollector, opsAggregationService, opsAlertEvaluatorService, opsCleanupService, opsScheduledReportService, opsSystemLogSink, opsService, opsIngressRejectAggregator, apiKeyService, authCacheInvalidationWorker, schedulerSnapshotService, tokenRefreshService, accountExpiryService, cnProviderBalanceCheckService, openAICodexVersionSyncService, claudeCodeVersionSyncService, proxyExpiryService, subscriptionExpiryService, usageCleanupService, idempotencyCleanupService, batchImageCleanupService, batchImageWorkerRuntime, pricingService, emailQueueService, billingCacheService, usageRecordWorkerPool, subscriptionService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, grokOAuthService, openAIGatewayService, scheduledTestRunnerService, backupService, paymentOrderExpiryService, channelMonitorRunner, channelMonitorV2Aggregator, userPlatformQuotaUsageFlusher, upstreamBillingProbeService, ollamaCloudUsageService, openCodeGoUsageService, auditLogService, openAIQuotaAutoResetService, promptService, pluginManager) + v := provideCleanup(client, redisClient, opsMetricsCollector, opsAggregationService, opsAlertEvaluatorService, opsCleanupService, opsScheduledReportService, opsSystemLogSink, opsService, opsIngressRejectAggregator, apiKeyService, authCacheInvalidationWorker, schedulerSnapshotService, tokenRefreshService, accountExpiryService, cnProviderBalanceCheckService, openAICodexVersionSyncService, claudeCodeVersionSyncService, proxyExpiryService, subscriptionExpiryService, usageCleanupService, idempotencyCleanupService, geminiCachedContentService, batchImageCleanupService, batchImageWorkerRuntime, pricingService, emailQueueService, billingCacheService, usageRecordWorkerPool, subscriptionService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, grokOAuthService, openAIGatewayService, scheduledTestRunnerService, backupService, paymentOrderExpiryService, channelMonitorRunner, channelMonitorV2Aggregator, userPlatformQuotaUsageFlusher, upstreamBillingProbeService, ollamaCloudUsageService, openCodeGoUsageService, auditLogService, openAIQuotaAutoResetService, promptService, pluginManager) application := &Application{ Server: httpServer, PromptAudit: promptService, @@ -413,6 +415,7 @@ func provideCleanup( subscriptionExpiry *service.SubscriptionExpiryService, usageCleanup *service.UsageCleanupService, idempotencyCleanup *service.IdempotencyCleanupService, + geminiCachedContent *service.GeminiCachedContentService, batchImageCleanup *service.BatchImageCleanupService, batchImageWorker *service.BatchImageWorkerRuntime, pricing *service.PricingService, @@ -552,6 +555,12 @@ func provideCleanup( } return nil }}, + {"GeminiCachedContentService", func() error { + if geminiCachedContent != nil { + geminiCachedContent.Stop() + } + return nil + }}, {"BatchImageCleanupService", func() error { if batchImageCleanup != nil { batchImageCleanup.Stop() diff --git a/backend/cmd/server/wire_gen_test.go b/backend/cmd/server/wire_gen_test.go index d14ca19db6..bb56c6735f 100644 --- a/backend/cmd/server/wire_gen_test.go +++ b/backend/cmd/server/wire_gen_test.go @@ -74,6 +74,7 @@ func TestProvideCleanup_WithMinimalDependencies_NoPanic(t *testing.T) { subscriptionExpirySvc, &service.UsageCleanupService{}, idempotencyCleanupSvc, + service.NewGeminiCachedContentService(nil, nil, cfg), &service.BatchImageCleanupService{}, nil, // batchImageWorker pricingSvc, diff --git a/backend/internal/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..c933421e9b --- /dev/null +++ b/backend/internal/handler/gemini_cached_contents_billing_test.go @@ -0,0 +1,955 @@ +//go:build unit + +package handler + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strconv" + "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 + cachedBillingAliasStoragePrice = 0.000009 + cachedBillingTokens = 1_000_000 +) + +type cachedBillingRepo struct { + mu sync.Mutex + records map[string]*service.GeminiCachedContent + // mutations 按记录 id 保存修改意图。 + mutations map[int64]cachedBillingMutation + failCommit bool + nextID int64 + deleted []int64 +} + +type cachedBillingMutation struct { + token string + leaseUntil time.Time +} + +func (r *cachedBillingRepo) Create(_ context.Context, record *service.GeminiCachedContent) error { + r.mu.Lock() + defer r.mu.Unlock() + 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 || r.liveLocked(record.ID) == nil { + 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 +} + +// liveLocked 返回未删除的记录；调用方须持有 mu。 +func (r *cachedBillingRepo) liveLocked(id int64) *service.GeminiCachedContent { + for _, deleted := range r.deleted { + if deleted == id { + return nil + } + } + for _, record := range r.records { + if record.ID == id { + return record + } + } + return nil +} + +func (r *cachedBillingRepo) BeginMutation(_ context.Context, id int64, baseline time.Time, token string, lease time.Duration) (bool, error) { + r.mu.Lock() + defer r.mu.Unlock() + record := r.liveLocked(id) + if record == nil { + return false, service.ErrGeminiCachedContentNotFound + } + if _, pending := r.mutations[id]; pending || !record.ExpireTime.Equal(baseline) { + return false, nil + } + r.mutations[id] = cachedBillingMutation{token: token, leaseUntil: time.Now().Add(lease)} + return true, nil +} + +func (r *cachedBillingRepo) TakeOverStaleMutation(_ context.Context, id int64, token string, lease time.Duration) (*service.GeminiCachedContent, error) { + r.mu.Lock() + defer r.mu.Unlock() + record := r.liveLocked(id) + mutation, pending := r.mutations[id] + if record == nil || !pending || mutation.leaseUntil.After(time.Now()) { + return nil, nil + } + r.mutations[id] = cachedBillingMutation{token: token, leaseUntil: time.Now().Add(lease)} + cp := *record + return &cp, nil +} + +func (r *cachedBillingRepo) CommitMutation(_ context.Context, id int64, token string, expireTime time.Time) error { + r.mu.Lock() + defer r.mu.Unlock() + if r.failCommit { + return errors.New("db unavailable") + } + record := r.liveLocked(id) + mutation, pending := r.mutations[id] + if record == nil || !pending || mutation.token != token { + return service.ErrGeminiCachedContentMutationConflict + } + record.ExpireTime = expireTime + delete(r.mutations, id) + return nil +} + +func (r *cachedBillingRepo) ReleaseMutation(_ context.Context, id int64, token string) error { + r.mu.Lock() + defer r.mu.Unlock() + if mutation, pending := r.mutations[id]; pending && mutation.token == token { + delete(r.mutations, id) + } + return nil +} + +func (r *cachedBillingRepo) ListStaleMutations(_ context.Context, limit int) ([]int64, error) { + r.mu.Lock() + defer r.mu.Unlock() + var ids []int64 + for id, mutation := range r.mutations { + if len(ids) < limit && mutation.leaseUntil.Before(time.Now()) && r.liveLocked(id) != nil { + ids = append(ids, id) + } + } + return ids, nil +} + +func (r *cachedBillingRepo) SoftDelete(_ context.Context, id int64) error { + r.mu.Lock() + defer r.mu.Unlock() + r.deleted = append(r.deleted, id) + return nil +} + +// pendingMutation 返回记录是否有未清除的修改意图。 +func (r *cachedBillingRepo) pendingMutation(id int64) bool { + r.mu.Lock() + defer r.mu.Unlock() + _, pending := r.mutations[id] + return pending +} + +// expireLease 让记录上的修改意图租约立即过期。 +func (r *cachedBillingRepo) expireLease(id int64) { + r.mu.Lock() + defer r.mu.Unlock() + mutation := r.mutations[id] + mutation.leaseUntil = time.Now().Add(-time.Second) + r.mutations[id] = mutation +} + +type cachedBillingAccountRepo struct { + service.AccountRepository + accounts map[int64]*service.Account +} + +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 + // patchTTLs 按顺序记录 PATCH 请求体中的 ttl。 + patchTTLs []string + beforePost func(n int) + // failures 按请求方法返回上游错误响应；transportErrors 按请求方法返回传输层错误（上游结果未知）。 + failures map[string]cachedBillingUpstreamFailure + transportErrors map[string]error +} + +type cachedBillingUpstreamFailure struct { + status int + body string +} + +func (u *cachedBillingUpstream) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) { + var patch struct { + TTL string `json:"ttl"` + } + if req.Method == http.MethodPatch { + _ = json.NewDecoder(req.Body).Decode(&patch) + } + u.mu.Lock() + u.calls = append(u.calls, req.Method+" "+req.URL.Path) + n := len(u.calls) + if req.Method == http.MethodPatch { + u.patchTTLs = append(u.patchTTLs, patch.TTL) + } + failure, failed := u.failures[req.Method] + transportErr := u.transportErrors[req.Method] + u.mu.Unlock() + if transportErr != nil { + return nil, transportErr + } + if failed { + return &http.Response{StatusCode: failure.status, Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(failure.body))}, nil + } + body := "{}" + switch req.Method { + case http.MethodPost: + body = u.createBody + case http.MethodPatch: + if u.beforePost != nil { + u.beforePost(n) + } + seconds, _ := strconv.ParseInt(strings.TrimSuffix(patch.TTL, "s"), 10, 64) + body = fmt.Sprintf(`{"name":"cachedContents/up1","expireTime":%q}`, time.Now().Add(time.Duration(seconds)*time.Second).UTC().Format(time.RFC3339Nano)) + } + return &http.Response{StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(body))}, nil +} + +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) first() *service.UsageLog { + r.mu.Lock() + defer r.mu.Unlock() + return r.logs[0] +} + +func (r *cachedBillingUsageRepo) totals() []float64 { + r.mu.Lock() + defer r.mu.Unlock() + 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 + // accounts 是调度快照桩直接返回的账号对象，测试可改其凭据。 + accounts map[int64]*service.Account +} + +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},"gemini-alias":{ + "litellm_provider":"gemini","mode":"chat","input_cost_per_token":%g,"output_cost_per_token":0.00001, + "cache_storage_cost_per_token_per_hour":%g}}`, + cachedBillingInputPrice, cachedBillingStoragePrice, cachedBillingInputPrice, cachedBillingAliasStoragePrice)), 0o644)) + cfg := &config.Config{} + cfg.Pricing.DataDir = dir + cfg.Pricing.UpdateIntervalHours = 100000 + 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{}, mutations: map[int64]cachedBillingMutation{}} + 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, accounts: accounts} +} + +// TestGeminiExplicitCacheExclusions_FiltersUnsupportedAccountsFromSnapshot 候选过滤直接按快照里的凭据判定： +// 第三方地址且未声明上游显式缓存的 Key 账号在选号前就被排除，官方地址与声明了上游显式缓存的账号保留。 +func TestGeminiExplicitCacheExclusions_FiltersUnsupportedAccountsFromSnapshot(t *testing.T) { + env := newCachedBillingEnv(t, 500, 900, 901, 902) + env.accounts[901].Credentials["base_url"] = "https://third.example.com" + env.accounts[902].Credentials["base_url"] = "https://gw.example.com" + env.accounts[902].Credentials[service.GeminiExplicitCacheUpstreamCredentialKey] = true + + groupID := env.group.ID + excluded, err := env.h.gatewayService.GeminiExplicitCacheExclusions(context.Background(), &groupID, 0) + require.NoError(t, err) + require.Equal(t, map[int64]struct{}{901: {}}, excluded) + + excluded, err = env.h.gatewayService.GeminiExplicitCacheExclusions(context.Background(), &groupID, 902) + require.NoError(t, err) + require.Equal(t, map[int64]struct{}{900: {}, 901: {}}, excluded, "绑定缓存的请求只保留持有账号") +} + +func (e *cachedBillingEnv) serve(method, target, body string, params gin.Params, call func(*gin.Context)) *httptest.ResponseRecorder { + 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) +} + +// TestGeminiCachedContentsPatch_UsesChannelUsageRecordedAtCreate 延长有效期沿用创建时的渠道计费口径： +// 计费模型来源为请求模型时按别名定价，用量记录带渠道 ID。 +func TestGeminiCachedContentsPatch_UsesChannelUsageRecordedAtCreate(t *testing.T) { + env := newCachedBillingEnv(t, 500, 900) + env.seed(time.Now().Add(time.Hour)) + env.repo.records["abc"].ChannelUsage = service.ChannelUsageFields{ + ChannelID: 5, OriginalModel: "gemini-alias", ChannelMappedModel: "gemini-2.5-pro", + BillingModelSource: service.BillingModelSourceRequested, + } + + rec := env.patch() + require.Equal(t, http.StatusOK, rec.Code, rec.Body.String()) + totals := env.waitForUsage(t, 1) + require.Len(t, totals, 1) + require.InDelta(t, cachedBillingTokens*cachedBillingAliasStoragePrice*23, totals[0], 0.01, "按创建时的计费模型（请求模型别名）计存储费") + usageLog := env.usage.first() + require.NotNil(t, usageLog.ChannelID) + require.Equal(t, int64(5), *usageLog.ChannelID) + require.True(t, strings.HasPrefix(usageLog.RequestID, "gcache:patch:abc:"), usageLog.RequestID) +} + +func cachedBillingCreateBody(expire time.Time) string { + return fmt.Sprintf(`{"name":"cachedContents/up1","usageMetadata":{"totalTokenCount":%d},"expireTime":%q}`, + cachedBillingTokens, expire.UTC().Format(time.RFC3339Nano)) +} + +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)), "本地到期时间不超过请求的有效期") + require.Equal(t, "gcache:create:"+record.PublicID, env.usage.first().RequestID, "费用行按缓存 ID 可检索") + } +} + +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 (e *cachedBillingEnv) patchTTL(id, ttl string) *httptest.ResponseRecorder { + return e.serve(http.MethodPatch, "/v1beta/cachedContents/"+id+"?updateMask=ttl", fmt.Sprintf(`{"ttl":%q}`, ttl), + gin.Params{{Key: "cacheID", Value: id}}, e.h.GeminiCachedContentsPatch) +} + +// requireTTLNear 断言发往上游的 ttl（"Ns"）与 want 相差不超过 5 秒。 +func requireTTLNear(t *testing.T, ttl string, want time.Duration) { + t.Helper() + seconds, err := strconv.ParseInt(strings.TrimSuffix(ttl, "s"), 10, 64) + require.NoError(t, err, ttl) + require.InDelta(t, want.Seconds(), float64(seconds), 5, ttl) +} + +// 同一缓存的修改串行执行：第一个 PATCH 的上游调用返回前，第二个 PATCH 直接得到 409 且不调用上游， +// 本地到期时间与上游最后生效的修改一致，只按第一个 PATCH 实际推进的时长收费。 +func TestGeminiCachedContentsPatch_ConcurrentUpdateConflictsAndKeepsUpstreamInSync(t *testing.T) { + env := newCachedBillingEnv(t, 1000, 900) + env.seed(time.Now().Add(time.Hour)) + env.upstream.beforePost = func(n int) { + if n == 1 { + rec := env.patchTTL("abc", "86400s") + require.Equal(t, http.StatusConflict, rec.Code, rec.Body.String()) + require.Contains(t, rec.Body.String(), `"ABORTED"`) + } + } + + rec := env.patchTTL("abc", "7200s") + require.Equal(t, http.StatusOK, rec.Code, rec.Body.String()) + require.Equal(t, 1, env.upstream.methodCalls(http.MethodPatch), "冲突的修改不调用上游") + require.WithinDuration(t, time.Now().Add(2*time.Hour), env.repo.records["abc"].ExpireTime, 5*time.Second, "本地到期时间与上游一致") + require.False(t, env.repo.pendingMutation(100)) + totals := env.waitForUsage(t, 1) + require.Len(t, totals, 1) + require.InDelta(t, cachedBillingTokens*cachedBillingStoragePrice, totals[0], cachedBillingTokens*cachedBillingStoragePrice*time.Minute.Hours(), "只收推进的 1 小时") +} + +// 上游已修改但本地写回失败：不计费并保留修改意图；租约内的修改返回 409， +// 租约过期后后台恢复把上游有效期回滚到修改前的到期时间并清除意图。 +func TestGeminiCachedContentsPatch_CommitFailureRollsBackUpstreamWithoutCharge(t *testing.T) { + env := newCachedBillingEnv(t, 1000, 900) + baseline := time.Now().Add(time.Hour) + env.seed(baseline) + env.repo.failCommit = true + + require.Equal(t, http.StatusInternalServerError, env.patch().Code) + require.Equal(t, 1, env.upstream.methodCalls(http.MethodPatch)) + require.True(t, env.repo.pendingMutation(100), "上游可能已修改：保留修改意图") + require.Empty(t, env.usage.totals(), "未提交的修改不收费") + + env.repo.failCommit = false + require.Equal(t, http.StatusConflict, env.patch().Code, "租约内的修改被拒绝") + require.Equal(t, 1, env.upstream.methodCalls(http.MethodPatch)) + + env.repo.expireLease(100) + recovered, err := env.h.geminiCachedContentService.RecoverStaleMutations(context.Background()) + require.NoError(t, err) + require.Equal(t, 1, recovered) + require.Equal(t, 2, env.upstream.methodCalls(http.MethodPatch)) + requireTTLNear(t, env.upstream.patchTTLs[1], time.Until(baseline)) + require.False(t, env.repo.pendingMutation(100)) + require.True(t, env.repo.records["abc"].ExpireTime.Equal(baseline), "本地到期时间保持修改前的值") + require.Empty(t, env.usage.totals()) +} + +// 租约过期的遗留修改意图在下一次修改前先回滚，随后的修改按修改前的到期时间计费。 +func TestGeminiCachedContentsPatch_StaleMutationRolledBackBeforeNextUpdate(t *testing.T) { + env := newCachedBillingEnv(t, 1000, 900) + baseline := time.Now().Add(time.Hour) + env.seed(baseline) + env.repo.failCommit = true + require.Equal(t, http.StatusInternalServerError, env.patch().Code) + env.repo.failCommit = false + env.repo.expireLease(100) + + rec := env.patch() + require.Equal(t, http.StatusOK, rec.Code, rec.Body.String()) + require.Equal(t, 3, env.upstream.methodCalls(http.MethodPatch), "失败的修改、回滚、本次修改") + requireTTLNear(t, env.upstream.patchTTLs[1], time.Until(baseline)) + require.Equal(t, "86400s", env.upstream.patchTTLs[2]) + totals := env.waitForUsage(t, 1) + require.Len(t, totals, 1) + require.InDelta(t, cachedBillingTokens*cachedBillingStoragePrice*23, totals[0], 0.01) +} + +// 上游结果未知（传输层错误）时保留修改意图等待回滚；上游明确拒绝时清除修改意图。两种情况都不收费。 +func TestGeminiCachedContentsPatch_MutationKeptOnlyWhenUpstreamOutcomeUnknown(t *testing.T) { + t.Run("transport error", func(t *testing.T) { + env := newCachedBillingEnv(t, 1000, 900) + env.seed(time.Now().Add(time.Hour)) + env.upstream.transportErrors = map[string]error{http.MethodPatch: errors.New("connection reset")} + + require.Equal(t, http.StatusBadGateway, env.patch().Code) + require.True(t, env.repo.pendingMutation(100)) + require.Empty(t, env.usage.totals()) + }) + t.Run("upstream rejected", func(t *testing.T) { + env := newCachedBillingEnv(t, 1000, 900) + env.seed(time.Now().Add(time.Hour)) + env.upstream.failures = map[string]cachedBillingUpstreamFailure{http.MethodPatch: {status: http.StatusBadRequest, + body: `{"error":{"code":400,"message":"invalid ttl","status":"INVALID_ARGUMENT"}}`}} + + require.Equal(t, http.StatusBadRequest, env.patch().Code) + require.False(t, env.repo.pendingMutation(100)) + require.Empty(t, env.usage.totals()) + }) +} + +// 修改前的到期时间已过：恢复时删除上游缓存与本地记录。 +func TestGeminiCachedContentMutationRecovery_DeletesUpstreamWhenBaselineExpired(t *testing.T) { + env := newCachedBillingEnv(t, 1000, 900) + env.seed(time.Now().Add(-time.Minute)) + env.repo.mutations[100] = cachedBillingMutation{token: "stale", leaseUntil: time.Now().Add(-time.Second)} + + recovered, err := env.h.geminiCachedContentService.RecoverStaleMutations(context.Background()) + require.NoError(t, err) + require.Equal(t, 1, recovered) + require.Equal(t, 1, env.upstream.methodCalls(http.MethodDelete)) + require.Zero(t, env.upstream.methodCalls(http.MethodPatch)) + require.Equal(t, []int64{100}, env.repo.deleted) +} + +// cachedBillingReservationCache 是启用在途预留时的计费缓存桩：余额与订阅用量取固定值，各维度预留按 key 记账。 +type cachedBillingReservationCache struct { + service.BillingCache + mu sync.Mutex + balance float64 + sub *service.SubscriptionCacheData + reserved map[string]map[string]float64 +} + +func (c *cachedBillingReservationCache) GetUserBalance(context.Context, int64) (float64, error) { + return c.balance, nil +} + +func (c *cachedBillingReservationCache) GetSubscriptionCache(context.Context, int64, int64) (*service.SubscriptionCacheData, error) { + return c.sub, nil +} + +func (c *cachedBillingReservationCache) reserveLocked(key, requestID string, amount float64) float64 { + if c.reserved == nil { + c.reserved = map[string]map[string]float64{} + } + if c.reserved[key] == nil { + c.reserved[key] = map[string]float64{} + } + others := 0.0 + for _, v := range c.reserved[key] { + others += v + } + c.reserved[key][requestID] = amount + return others +} + +func (c *cachedBillingReservationCache) ReserveInflightBalance(_ context.Context, userID int64, requestID string, amount, balance float64, _ time.Duration) (bool, float64, error) { + c.mu.Lock() + defer c.mu.Unlock() + key := fmt.Sprintf("user:%d", userID) + others := c.reserveLocked(key, requestID, amount) + if len(c.reserved[key]) > 1 && balance-others < amount { + delete(c.reserved[key], requestID) + return false, others, nil + } + return true, others, nil +} + +func (c *cachedBillingReservationCache) ReleaseInflightBalance(_ context.Context, userID int64, requestID string) error { + c.mu.Lock() + defer c.mu.Unlock() + delete(c.reserved[fmt.Sprintf("user:%d", userID)], requestID) + return nil +} + +func (c *cachedBillingReservationCache) ReserveInflightScope(_ context.Context, scope, requestID string, amount float64, _ time.Duration) (float64, error) { + c.mu.Lock() + defer c.mu.Unlock() + return c.reserveLocked(scope, requestID, amount), nil +} + +func (c *cachedBillingReservationCache) ReleaseInflightScope(_ context.Context, scope, requestID string) error { + c.mu.Lock() + defer c.mu.Unlock() + delete(c.reserved[scope], requestID) + return nil +} + +func enableCachedBillingReservations(t *testing.T, env *cachedBillingEnv, cache *cachedBillingReservationCache) { + t.Helper() + cfg := &config.Config{} + cfg.Billing.InflightReservation.Enabled = true + bcs := service.NewBillingCacheService(cache, nil, nil, nil, nil, nil, cfg, nil) + t.Cleanup(bcs.Stop) + env.h.billingCacheService = bcs +} + +// 确定金额的续期费按 API Key 与余额维度预留：第一笔续期（约 $103.5）尚未入账时， +// 第二个缓存的续期超出 Key 剩余额度或余额，在调用上游前被拒绝。 +func TestGeminiCachedContentsPatch_UnsettledChargeHoldsKeyQuotaAndBalance(t *testing.T) { + for _, tc := range []struct { + name string + balance, quota float64 + }{ + {name: "api key quota", balance: 1000, quota: 110}, + {name: "balance", balance: 110}, + } { + t.Run(tc.name, func(t *testing.T) { + env := newCachedBillingEnv(t, tc.balance, 900) + env.seed(time.Now().Add(time.Hour)) + second := *env.repo.records["abc"] + second.ID, second.PublicID, second.UpstreamName = 101, "def", "cachedContents/up2" + env.repo.records["def"] = &second + env.apiKey.Quota = tc.quota + enableCachedBillingReservations(t, env, &cachedBillingReservationCache{balance: tc.balance}) + env.upstream.beforePost = func(n int) { + if n == 1 { + rec := env.patchTTL("def", "86400s") + require.NotEqual(t, http.StatusOK, rec.Code, rec.Body.String()) + } + } + + require.Equal(t, http.StatusOK, env.patch().Code) + require.Equal(t, 1, env.upstream.methodCalls(http.MethodPatch), "第二笔续期没有到达上游") + }) + } +} + +// 订阅模式的确定金额费用按订阅维度预留：第一笔未入账时第二笔超出日剩余额度被拒绝，释放后放行。 +func TestReserveGeminiCachedContentCharge_UnsettledChargeHoldsSubscriptionHeadroom(t *testing.T) { + env := newCachedBillingEnv(t, 1000, 900) + limit := 110.0 + env.group.SubscriptionType = service.SubscriptionTypeSubscription + env.group.DailyLimitUSD = &limit + enableCachedBillingReservations(t, env, &cachedBillingReservationCache{balance: 1000, + sub: &service.SubscriptionCacheData{Status: service.SubscriptionStatusActive, ExpiresAt: time.Now().Add(time.Hour)}}) + sub := &service.UserSubscription{ID: 1, UserID: env.apiKey.UserID, GroupID: env.group.ID, Status: service.SubscriptionStatusActive, + ExpiresAt: time.Now().Add(time.Hour)} + reserve := func() (func(), error) { + var done func() + var err error + env.serve(http.MethodPatch, "/", `{}`, nil, func(c *gin.Context) { + done, err = env.h.reserveGeminiCachedContentCharge(c, env.apiKey, sub, 103.5) + }) + return done, err + } + + firstDone, err := reserve() + require.NoError(t, err) + _, err = reserve() + require.ErrorIs(t, err, service.ErrDailyLimitExceeded) + firstDone() + secondDone, err := reserve() + require.NoError(t, err) + secondDone() +} + +func TestGeminiCachedContentsPatch_UnaffordableExtensionRejectedBeforeUpstream(t *testing.T) { + 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) + }) + } +} + +// 持有账号暂时不支持显式缓存（能力被关闭）时，恢复流程不能把修改意图当作已回滚而清除； +// 能力恢复后再回滚上游并清除意图。 +func TestGeminiCachedContentMutationRecovery_KeepsIntentWhileAccountCannotCache(t *testing.T) { + env := newCachedBillingEnv(t, 1000, 900) + env.accounts[900].Credentials["base_url"] = "https://gw.example.com" + env.accounts[900].Credentials[service.GeminiExplicitCacheUpstreamCredentialKey] = true + baseline := time.Now().Add(time.Hour) + env.seed(baseline) + env.repo.failCommit = true + require.Equal(t, http.StatusInternalServerError, env.patch().Code) + require.True(t, env.repo.pendingMutation(100)) + require.Equal(t, 1, env.upstream.methodCalls(http.MethodPatch)) + env.repo.failCommit = false + env.repo.expireLease(100) + + env.accounts[900].Credentials[service.GeminiExplicitCacheUpstreamCredentialKey] = false + recovered, err := env.h.geminiCachedContentService.RecoverStaleMutations(context.Background()) + require.Error(t, err, "账号暂不支持缓存时恢复失败") + require.Zero(t, recovered) + require.True(t, env.repo.pendingMutation(100), "未回滚上游前保留修改意图") + require.Equal(t, 1, env.upstream.methodCalls(http.MethodPatch), "不改动上游") + require.Empty(t, env.repo.deleted) + + env.accounts[900].Credentials[service.GeminiExplicitCacheUpstreamCredentialKey] = true + env.repo.expireLease(100) + recovered, err = env.h.geminiCachedContentService.RecoverStaleMutations(context.Background()) + require.NoError(t, err) + require.Equal(t, 1, recovered) + require.Equal(t, 2, env.upstream.methodCalls(http.MethodPatch)) + requireTTLNear(t, env.upstream.patchTTLs[1], time.Until(baseline)) + require.False(t, env.repo.pendingMutation(100)) + require.True(t, env.repo.records["abc"].ExpireTime.Equal(baseline)) + require.Empty(t, env.usage.totals()) +} + +// 持有账号长期不支持显式缓存，且修改前到期时间加上单次最长有效期都已过去：上游缓存必已自行过期， +// 恢复流程不再保留意图，直接移除本地记录。 +func TestGeminiCachedContentMutationRecovery_AbandonsIntentOnceUpstreamMustHaveExpired(t *testing.T) { + env := newCachedBillingEnv(t, 1000, 900) + env.accounts[900].Credentials["base_url"] = "https://gw.example.com" + env.accounts[900].Credentials[service.GeminiExplicitCacheUpstreamCredentialKey] = true + env.seed(time.Now().Add(time.Hour)) + env.repo.failCommit = true + require.Equal(t, http.StatusInternalServerError, env.patch().Code) + env.repo.failCommit = false + env.repo.expireLease(100) + env.accounts[900].Credentials[service.GeminiExplicitCacheUpstreamCredentialKey] = false + + // 修改前到期时间 + 24h 上限 + 1h 容忍仍未过去：保留意图。 + env.repo.records["abc"].ExpireTime = time.Now().Add(-20 * time.Hour) + recovered, err := env.h.geminiCachedContentService.RecoverStaleMutations(context.Background()) + require.Error(t, err) + require.Zero(t, recovered) + require.True(t, env.repo.pendingMutation(100)) + require.Empty(t, env.repo.deleted) + + env.repo.expireLease(100) + env.repo.records["abc"].ExpireTime = time.Now().Add(-26 * time.Hour) + recovered, err = env.h.geminiCachedContentService.RecoverStaleMutations(context.Background()) + require.NoError(t, err) + require.Equal(t, 1, recovered) + require.Equal(t, 1, env.upstream.methodCalls(http.MethodPatch), "不改动上游") + require.Zero(t, env.upstream.methodCalls(http.MethodDelete)) + require.Equal(t, []int64{100}, env.repo.deleted) + require.Empty(t, env.usage.totals()) + + recovered, err = env.h.geminiCachedContentService.RecoverStaleMutations(context.Background()) + require.NoError(t, err) + require.Zero(t, recovered, "作废的记录不再被选中") +} + +// 作废判定的边界：默认上限 24h 加 1h 容忍；配置上限低于 24h 时仍按 24h 兜底。 +func TestGeminiCachedContentMutationRecovery_AbandonBoundary(t *testing.T) { + disableCache := func(env *cachedBillingEnv) { + env.accounts[900].Credentials["base_url"] = "https://gw.example.com" + env.accounts[900].Credentials[service.GeminiExplicitCacheUpstreamCredentialKey] = false + } + t.Run("within grace keeps intent", func(t *testing.T) { + env := newCachedBillingEnv(t, 1000, 900) + env.seed(time.Now().Add(-25*time.Hour + time.Minute)) + env.repo.mutations[100] = cachedBillingMutation{token: "stale", leaseUntil: time.Now().Add(-time.Second)} + disableCache(env) + + recovered, err := env.h.geminiCachedContentService.RecoverStaleMutations(context.Background()) + require.Error(t, err) + require.Zero(t, recovered) + require.True(t, env.repo.pendingMutation(100)) + require.Empty(t, env.repo.deleted) + }) + t.Run("configured max ttl below default still waits a day", func(t *testing.T) { + env := newCachedBillingEnv(t, 1000, 900) + env.seed(time.Now().Add(-12 * time.Hour)) + env.repo.mutations[100] = cachedBillingMutation{token: "stale", leaseUntil: time.Now().Add(-time.Second)} + disableCache(env) + cfg := &config.Config{} + cfg.Gateway.GeminiCachedContentMaxTTLSeconds = 3600 + compat := service.NewGeminiMessagesCompatService(&cachedBillingAccountRepo{accounts: env.accounts}, nil, nil, nil, nil, nil, env.upstream, nil, cfg) + svc := service.NewGeminiCachedContentService(env.repo, compat, cfg) + + recovered, err := svc.RecoverStaleMutations(context.Background()) + require.Error(t, err) + require.Zero(t, recovered) + require.True(t, env.repo.pendingMutation(100)) + + env.repo.expireLease(100) + env.repo.records["abc"].ExpireTime = time.Now().Add(-26 * time.Hour) + recovered, err = svc.RecoverStaleMutations(context.Background()) + require.NoError(t, err) + require.Equal(t, 1, recovered) + require.Equal(t, []int64{100}, env.repo.deleted) + }) +} + +// 持有账号可用但上游回滚持续失败（如 Key 已失效）：上游缓存必已自行过期后不再重试，移除本地记录。 +func TestGeminiCachedContentMutationRecovery_AbandonsIntentWhenRollbackKeepsFailing(t *testing.T) { + env := newCachedBillingEnv(t, 1000, 900) + env.seed(time.Now().Add(-time.Minute)) + env.repo.mutations[100] = cachedBillingMutation{token: "stale", leaseUntil: time.Now().Add(-time.Second)} + env.upstream.failures = map[string]cachedBillingUpstreamFailure{http.MethodDelete: {status: http.StatusUnauthorized, + body: `{"error":{"code":401,"message":"API key not valid","status":"UNAUTHENTICATED"}}`}} + + recovered, err := env.h.geminiCachedContentService.RecoverStaleMutations(context.Background()) + require.Error(t, err, "上游缓存可能仍存活：保留意图等待重试") + require.Zero(t, recovered) + require.Equal(t, 1, env.upstream.methodCalls(http.MethodDelete)) + require.True(t, env.repo.pendingMutation(100)) + require.Empty(t, env.repo.deleted) + + env.repo.expireLease(100) + env.repo.records["abc"].ExpireTime = time.Now().Add(-26 * time.Hour) + recovered, err = env.h.geminiCachedContentService.RecoverStaleMutations(context.Background()) + require.NoError(t, err) + require.Equal(t, 1, recovered) + require.Equal(t, 2, env.upstream.methodCalls(http.MethodDelete), "账号可用时仍先尝试回滚") + require.Equal(t, []int64{100}, env.repo.deleted) + require.Empty(t, env.usage.totals()) +} + +// 持有账号已删除时没有凭据可回滚上游：恢复流程移除本地记录并清除意图，不再重试。 +func TestGeminiCachedContentMutationRecovery_ForgetsRecordWhenAccountDeleted(t *testing.T) { + env := newCachedBillingEnv(t, 1000, 900) + env.seed(time.Now().Add(time.Hour)) + env.repo.failCommit = true + require.Equal(t, http.StatusInternalServerError, env.patch().Code) + require.True(t, env.repo.pendingMutation(100)) + env.repo.failCommit = false + env.repo.expireLease(100) + + delete(env.accounts, 900) + recovered, err := env.h.geminiCachedContentService.RecoverStaleMutations(context.Background()) + require.NoError(t, err) + require.Equal(t, 1, recovered) + require.Equal(t, 1, env.upstream.methodCalls(http.MethodPatch), "无凭据，不改动上游") + require.Zero(t, env.upstream.methodCalls(http.MethodDelete)) + require.Equal(t, []int64{100}, env.repo.deleted, "本地记录作废") + require.Empty(t, env.usage.totals()) + require.Equal(t, http.StatusForbidden, env.serve(http.MethodGet, "/v1beta/cachedContents/abc", "", gin.Params{{Key: "cacheID", Value: "abc"}}, env.h.GeminiCachedContentsGet).Code, "作废的记录对客户端不可见") + + recovered, err = env.h.geminiCachedContentService.RecoverStaleMutations(context.Background()) + require.NoError(t, err) + require.Zero(t, recovered, "不再重复处理") +} diff --git a/backend/internal/handler/gemini_cached_contents_handler.go b/backend/internal/handler/gemini_cached_contents_handler.go new file mode 100644 index 0000000000..2071c16129 --- /dev/null +++ b/backend/internal/handler/gemini_cached_contents_handler.go @@ -0,0 +1,848 @@ +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" + "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 { + googlePlatformError(c, http.StatusNotFound, "Explicit context caching is not enabled") + return nil, middleware.AuthSubject{}, nil, false + } + if apiKey.Group == nil || apiKey.GroupID == nil || apiKey.Group.Platform == service.PlatformComposite { + 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, code, message, retryAfter := billingErrorDetails(err) + if retryAfter > 0 { + c.Header("Retry-After", strconv.Itoa(retryAfter)) + } + googleErrorWithType(c, status, code, "", 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, code, message, retryAfter := billingErrorDetails(err) + if retryAfter > 0 { + c.Header("Retry-After", strconv.Itoa(retryAfter)) + } + googleErrorWithType(c, status, code, "", 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() + googlePlatformError(c, http.StatusBadRequest, fmt.Sprintf("Explicit context caching is not available for model %s: storage pricing is not configured", modelName)) + return + } + upstreamModel, err := service.GeminiCachedContentUpstreamModelResource(account, mappedModel) + 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 + } + h.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, + service.GeminiCachedContentCreateUsageRequestID(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, + ChannelUsage: usageFields, + } + 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, + service.GeminiCachedContentCreateUsageRequestID(publicID), modelName, mappedModel, int(tokens), storageTokenHours, usageFields, body) + c.JSON(http.StatusOK, service.GeminiCachedContentView(record)) + return + } + + if lastUpstreamErr != nil { + h.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 +// 只允许修改有效期；延长部分在本次请求内按增量时长收取存储费，缩短不退费。 +// 同一缓存的修改串行执行：上游修改、本地写回与计费都在修改意图的租约内完成，并发的修改返回 409。 +// 上游结果未知或本地写回失败的修改不计费，租约过期后由恢复流程把上游有效期回滚到修改前的到期时间。 +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, + ChannelUsageFields: record.ChannelUsage, + } + if !h.gatewayService.CacheStoragePriced(c.Request.Context(), usage) { + googlePlatformError(c, http.StatusBadRequest, fmt.Sprintf("Explicit context caching is not available for model %s: storage pricing is not configured", record.Model)) + 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) + mutation, err := svc.BeginMutation(c.Request.Context(), record) + if err != nil { + switch { + case errors.Is(err, service.ErrGeminiCachedContentNotFound): + googleError(c, http.StatusForbidden, service.GeminiCachedContentNotFoundMessage) + case errors.Is(err, service.ErrGeminiCachedContentMutationConflict): + reqLog.Info("gemini.cached_content.mutation_conflict", zap.Int64("account_id", account.ID), zap.Error(err)) + writeGeminiCachedContentMutationConflict(c) + default: + reqLog.Error("gemini.cached_content.begin_mutation_failed", zap.Int64("account_id", account.ID), zap.Error(err)) + googleError(c, http.StatusInternalServerError, "Failed to update cached content") + } + return + } + upstream, err := svc.PatchUpstream(c.Request.Context(), account, record, ttl) + if err != nil { + var upErr *service.GeminiCachedContentUpstreamError + upstreamRejected := errors.As(err, &upErr) + if upstreamRejected || errors.Is(err, service.ErrGeminiCachedContentRequestNotSent) { + // 上游有效期未变：清除修改意图。 + persistCtx, cancelPersist := geminiCachedContentPersistContext(c) + if releaseErr := svc.ReleaseMutation(persistCtx, mutation); releaseErr != nil { + reqLog.Warn("gemini.cached_content.release_mutation_failed", zap.Error(releaseErr)) + } + cancelPersist() + } + if upstreamRejected { + if upErr.NotFound() { + if forgetErr := svc.Forget(c.Request.Context(), record); forgetErr != nil { + reqLog.Warn("gemini.cached_content.forget_failed", zap.Error(forgetErr)) + } + } + h.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) + err = svc.CommitMutation(persistCtx, mutation, expireTime) + cancelPersist() + if err != nil { + reqLog.Error("gemini.cached_content.commit_mutation_failed", zap.Int64("account_id", account.ID), zap.Error(err)) + if errors.Is(err, service.ErrGeminiCachedContentMutationConflict) { + writeGeminiCachedContentMutationConflict(c) + return + } + googleError(c, http.StatusInternalServerError, "Failed to persist cached content expiration") + return + } + // 提交前库中的到期时间一直是修改前的值：按本次实际推进的时长收费。 + if extended := expireTime.Sub(mutation.Baseline()); extended > 0 { + h.submitGeminiCachedContentUsage(c, apiKey, account, subscription, pricingAt, + service.GeminiCachedContentPatchUsageRequestID(record.PublicID), record.Model, record.UpstreamModelID(), 0, + service.GeminiCachedContentStorageTokenHours(record.TotalTokenCount, extended), record.ChannelUsage, body) + } + 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() { + h.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 按计费模型的存储单价计； +// requestID 是带固定前缀的费用行标识（见 GeminiCachedContentUsageRequestIDPrefix），落库时不被请求 id 覆盖。 +func (h *GatewayHandler) submitGeminiCachedContentUsage( + c *gin.Context, + apiKey *service.APIKey, + account *service.Account, + subscription *service.UserSubscription, + pricingAt time.Time, + requestID 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: requestID, + 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（倍率后）， +// 并按该金额在用户、订阅、API Key 维度登记在途预留（不受预留上限封顶，见 ReserveFixedCharge）。 +// 返回的 done 须在 handler 结束时调用；预留随计费任务交接，落库后释放。 +func (h *GatewayHandler) reserveGeminiCachedContentCharge(c *gin.Context, apiKey *service.APIKey, subscription *service.UserSubscription, amount float64) (func(), error) { + ctx := c.Request.Context() + res, err := h.billingCacheService.ReserveFixedCharge(ctx, apiKey.User, apiKey, 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 +} + +// writeGeminiCachedContentMutationConflict 回写同一缓存的修改冲突（Google 并发冲突语义 ABORTED），客户端可稍后重试。 +func writeGeminiCachedContentMutationConflict(c *gin.Context) { + c.JSON(http.StatusConflict, gin.H{ + "error": gin.H{ + "code": http.StatusConflict, + "message": "Another update to this cached content is in progress, please retry later", + "status": "ABORTED", + }, + }) +} + +func writeGeminiCachedContentBillingError(c *gin.Context, err error) { + status, code, message, retryAfter := billingErrorDetails(err) + if retryAfter > 0 { + c.Header("Retry-After", strconv.Itoa(retryAfter)) + } + googleErrorWithType(c, status, code, "", message) +} + +// writeGeminiCachedContentUpstreamError 回写上游错误：缓存不存在统一返回标准 403（各上游形态不同且可能含上游资源 ID）； +// 账号凭据或权限类错误不是客户端问题，统一改为 502；限流 / 不可用按普通 Gemini 路径处理（透传规则、通用映射）， +// 只有上游网关账号的这类响应反映持有缓存的上游账号状态，原样回写。 +func (h *GatewayHandler) writeGeminiCachedContentUpstreamError(c *gin.Context, upErr *service.GeminiCachedContentUpstreamError) { + status := upErr.StatusCode + message := strings.TrimSpace(extractGoogleErrorMessage(upErr.Body)) + service.SetOpsUpstreamError(c, status, message, "") + if upErr.NotFound() { + service.MarkOpsRequestScopedError(c, "permission_error") + c.Data(http.StatusForbidden, "application/json", []byte(service.GeminiCachedContentNotFoundResponse)) + return + } + if (status == http.StatusTooManyRequests || status >= http.StatusInternalServerError) && !upErr.UpstreamGateway { + h.handleGeminiFailoverExhausted(c, &service.UpstreamFailoverError{StatusCode: status, ResponseBody: upErr.Body}) + return + } + if status == http.StatusUnauthorized || status == http.StatusForbidden { + googleError(c, http.StatusBadGateway, "Upstream account rejected the cached content request") + return + } + 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 处理绑定缓存的请求在持有账号上重试耗尽：没有其他账号可换。 +// 上游网关账号的限流 / 不可用响应反映持有缓存的上游账号状态，原样返回供客户端判断； +// 其余（含官方账号）仍按普通 Gemini 路径处理：透传规则、通用映射。 +func (h *GatewayHandler) handleGeminiCachedContentFailoverExhausted(c *gin.Context, failoverErr *service.UpstreamFailoverError) { + if failoverErr == nil || !failoverErr.BoundUpstreamPassthrough || len(failoverErr.ResponseBody) == 0 || !json.Valid(failoverErr.ResponseBody) { + 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..f9a1d2b493 --- /dev/null +++ b/backend/internal/handler/gemini_cached_contents_handler_test.go @@ -0,0 +1,264 @@ +//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) BeginMutation(context.Context, int64, time.Time, string, time.Duration) (bool, error) { + return true, nil +} + +func (r *geminiCachedContentHandlerRepoStub) TakeOverStaleMutation(context.Context, int64, string, time.Duration) (*service.GeminiCachedContent, error) { + return nil, nil +} + +func (r *geminiCachedContentHandlerRepoStub) CommitMutation(context.Context, int64, string, time.Time) error { + return nil +} + +func (r *geminiCachedContentHandlerRepoStub) ReleaseMutation(context.Context, int64, string) error { + return nil +} + +func (r *geminiCachedContentHandlerRepoStub) ListStaleMutations(context.Context, int) ([]int64, error) { + return nil, nil +} + +func (r *geminiCachedContentHandlerRepoStub) SoftDelete(context.Context, int64) error { + 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), BoundUpstreamPassthrough: true}) + require.Equal(t, http.StatusServiceUnavailable, rec.Code, "绑定缓存的请求无账号可换，上游网关的 503 原样返回") + require.JSONEq(t, unavailable, rec.Body.String()) + + c, rec = newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", apiKey) + h.handleGeminiCachedContentFailoverExhausted(c, &service.UpstreamFailoverError{StatusCode: http.StatusTooManyRequests, ResponseBody: []byte(`{"error":{"code":429,"message":"Too many pending requests, please retry later","status":"RESOURCE_EXHAUSTED"}}`), BoundUpstreamPassthrough: true}) + require.Equal(t, http.StatusTooManyRequests, rec.Code) + require.Contains(t, rec.Body.String(), "Too many pending requests") + + officialQuota := `{"error":{"code":429,"message":"You exceeded your current quota","status":"RESOURCE_EXHAUSTED","details":[{"@type":"type.googleapis.com/google.rpc.QuotaFailure","violations":[{"quotaId":"GenerateRequestsPerDayPerProjectPerModel-FreeTier"}]}]}}` + c, rec = newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", apiKey) + h.handleGeminiCachedContentFailoverExhausted(c, &service.UpstreamFailoverError{StatusCode: http.StatusTooManyRequests, ResponseBody: []byte(officialQuota)}) + require.Equal(t, http.StatusTooManyRequests, rec.Code, "官方账号的限流按普通 Gemini 路径映射") + require.NotContains(t, rec.Body.String(), "FreeTier", "不向客户端暴露上游账号的配额明细") + require.Contains(t, rec.Body.String(), "Upstream rate limit exceeded") + + c, rec = newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", apiKey) + h.handleGeminiCachedContentFailoverExhausted(c, &service.UpstreamFailoverError{StatusCode: http.StatusUnauthorized, ResponseBody: []byte(`{"error":{"code":401,"message":"bad key"}}`)}) + require.Equal(t, http.StatusBadGateway, rec.Code, "账号凭据类错误仍按网关映射，不透传给客户端") + require.NotContains(t, rec.Body.String(), "bad key") +} + +func TestWriteGeminiCachedContentUpstreamError_RateLimitMappedUnlessUpstreamGateway(t *testing.T) { + h := &GatewayHandler{} + apiKey := geminiCachedContentTestAPIKey(service.PlatformGemini) + quota := `{"error":{"code":429,"message":"You exceeded your current quota","status":"RESOURCE_EXHAUSTED","details":[{"quotaId":"GenerateRequestsPerDayPerProjectPerModel-FreeTier"}]}}` + + c, rec := newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/cachedContents", apiKey) + h.writeGeminiCachedContentUpstreamError(c, &service.GeminiCachedContentUpstreamError{StatusCode: http.StatusTooManyRequests, Body: []byte(quota)}) + require.Equal(t, http.StatusTooManyRequests, rec.Code) + require.NotContains(t, rec.Body.String(), "FreeTier", "官方账号的限流不透传配额明细") + + c, rec = newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/cachedContents", apiKey) + h.writeGeminiCachedContentUpstreamError(c, &service.GeminiCachedContentUpstreamError{StatusCode: http.StatusTooManyRequests, Body: []byte(quota), UpstreamGateway: true}) + require.Equal(t, http.StatusTooManyRequests, rec.Code) + require.JSONEq(t, quota, rec.Body.String(), "上游网关账号的限流反映持有缓存的上游账号状态，原样回写") + + c, rec = newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/cachedContents", apiKey) + h.writeGeminiCachedContentUpstreamError(c, &service.GeminiCachedContentUpstreamError{StatusCode: http.StatusBadRequest, Body: []byte(`{"error":{"code":400,"message":"Cached content is too small","status":"INVALID_ARGUMENT"}}`)}) + require.Equal(t, http.StatusBadRequest, rec.Code) + require.Contains(t, rec.Body.String(), "too small", "请求校验类错误仍原样回写") +} diff --git a/backend/internal/handler/gemini_v1beta_handler.go b/backend/internal/handler/gemini_v1beta_handler.go index 4e56c37232..8627e0d139 100644 --- a/backend/internal/handler/gemini_v1beta_handler.go +++ b/backend/internal/handler/gemini_v1beta_handler.go @@ -74,7 +74,7 @@ func (h *GatewayHandler) GeminiV1BetaListModels(c *gin.Context) { return } - account, err := h.geminiCompatService.SelectAccountForAIStudioEndpoints(c.Request.Context(), apiKey.GroupID) + account, err := h.geminiCompatService.SelectAccountForModelsEndpoints(c.Request.Context(), apiKey.GroupID) if err != nil { if len(agModels) > 0 { c.JSON(http.StatusOK, gemini.ModelsListResponse{Models: filterGeminiModels(agModels)}) @@ -85,7 +85,7 @@ func (h *GatewayHandler) GeminiV1BetaListModels(c *gin.Context) { return } - res, err := h.geminiCompatService.ForwardAIStudioGET(c.Request.Context(), account, "/v1beta/models") + res, err := h.geminiCompatService.ForwardListModels(c.Request.Context(), account) if err != nil { googleError(c, http.StatusBadGateway, err.Error()) return @@ -251,7 +251,7 @@ func (h *GatewayHandler) GeminiV1BetaGetModel(c *gin.Context) { return } - account, err := h.geminiCompatService.SelectAccountForAIStudioEndpoints(c.Request.Context(), apiKey.GroupID) + account, err := h.geminiCompatService.SelectAccountForModelsEndpoints(c.Request.Context(), apiKey.GroupID) if err != nil { // 没有 gemini 账户，检查是否有 antigravity 账户可用 hasAntigravity, _ := h.geminiCompatService.HasAntigravityAccounts(c.Request.Context(), apiKey.GroupID) @@ -265,7 +265,7 @@ func (h *GatewayHandler) GeminiV1BetaGetModel(c *gin.Context) { return } - res, err := h.geminiCompatService.ForwardAIStudioGET(c.Request.Context(), account, "/v1beta/models/"+modelName) + res, err := h.geminiCompatService.ForwardGetModel(c.Request.Context(), account, modelName) if err != nil { googleError(c, http.StatusBadGateway, err.Error()) return @@ -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) @@ -382,11 +388,11 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) { // 2) billing eligibility check (after wait) if err := h.billingCacheService.CheckBillingEligibility(c.Request.Context(), apiKey.User, apiKey, apiKey.Group, subscription, service.QuotaPlatform(c.Request.Context(), apiKey)); err != nil { reqLog.Info("gemini.billing_eligibility_check_failed", zap.Error(err)) - status, _, message, retryAfter := billingErrorDetails(err) + status, code, message, retryAfter := billingErrorDetails(err) if retryAfter > 0 { c.Header("Retry-After", strconv.Itoa(retryAfter)) } - googleError(c, status, message) + googleErrorWithType(c, status, code, "", message) return } @@ -394,19 +400,22 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) { inflightRelease, err := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(modelName, body)) if err != nil { reqLog.Info("gemini.inflight_reservation_rejected", zap.Error(err)) - status, _, message, retryAfter := billingErrorDetails(err) + status, code, message, retryAfter := billingErrorDetails(err) if retryAfter > 0 { c.Header("Retry-After", strconv.Itoa(retryAfter)) } - googleError(c, status, message) + googleErrorWithType(c, status, code, "", message) return } defer inflightRelease() // 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: @@ -809,7 +861,8 @@ func (h *GatewayHandler) handleGeminiFailoverExhausted(c *gin.Context, failoverE c.Set(service.OpsSkipPassthroughKey, true) } - googleError(c, respCode, msg) + // 按透传规则改写的是上游错误，不登记本地错误类型。 + googleErrorWithType(c, respCode, "", "", msg) return } } @@ -844,7 +897,21 @@ type pathParseError struct{ msg string } func (e *pathParseError) Error() string { return e.msg } +// googleError 以 Google 错误格式回写网关产生的错误，并按状态码向 ops 登记错误类型： +// 4xx 为客户端请求问题，其余状态码不登记（按平台错误统计）。 func googleError(c *gin.Context, status int, message string) { + googleErrorWithType(c, status, googleClientErrorType(status), "", message) +} + +// googlePlatformError 以 Google 错误格式回写由平台侧原因（如管理员未完成配置）导致的错误，ops 按平台错误统计。 +func googlePlatformError(c *gin.Context, status int, message string) { + googleErrorWithType(c, status, "", "", message) +} + +// googleErrorWithType 以 Google 错误格式回写错误，并向 ops 登记错误类型与错误码（Google 错误体没有 type 字段）， +// 与 Claude / OpenAI 格式错误同口径分类；errType 为空时不登记。 +func googleErrorWithType(c *gin.Context, status int, errType, code, message string) { + service.SetOpsLocalErrorType(c, errType, code) c.JSON(status, gin.H{ "error": gin.H{ "code": status, @@ -854,11 +921,27 @@ func googleError(c *gin.Context, status int, message string) { }) } +// googleClientErrorType 返回客户端请求类 4xx 状态码对应的 ops 错误类型；其余状态码返回空串。 +func googleClientErrorType(status int) string { + switch status { + case http.StatusBadRequest, http.StatusRequestEntityTooLarge: + return "invalid_request_error" + case http.StatusUnauthorized: + return "authentication_error" + case http.StatusForbidden: + return "permission_error" + case http.StatusNotFound: + return "not_found_error" + default: + return "" + } +} + // googleConcurrencyError 以 Google 错误格式回写并发槽获取失败，状态码与文案 // 沿用 concurrencyErrorResponse 的统一映射（客户端断开为 499）。 func googleConcurrencyError(c *gin.Context, err error, slotType string) { - status, _, _, message := concurrencyErrorResponse(err, slotType) - googleError(c, status, message) + status, errType, code, message := concurrencyErrorResponse(err, slotType) + googleErrorWithType(c, status, errType, code, message) } func writeUpstreamResponse(c *gin.Context, res *service.UpstreamHTTPResult) { diff --git a/backend/internal/handler/gemini_v1beta_handler_test.go b/backend/internal/handler/gemini_v1beta_handler_test.go index 8604b07b88..44d66d5c97 100644 --- a/backend/internal/handler/gemini_v1beta_handler_test.go +++ b/backend/internal/handler/gemini_v1beta_handler_test.go @@ -109,7 +109,7 @@ func TestGeminiV1BetaHandler_ListModelsAntigravityFallback(t *testing.T) { expectedBehavior string }{ { - name: "有Gemini账户-调用ForwardAIStudioGET", + name: "有Gemini账户-调用ForwardListModels", hasGeminiAccount: true, hasAntigravity: false, expectedBehavior: "forward_to_upstream", @@ -155,7 +155,7 @@ func TestGeminiV1BetaHandler_GetModelAntigravityFallback(t *testing.T) { expectedBehavior string }{ { - name: "有Gemini账户-调用ForwardAIStudioGET", + name: "有Gemini账户-调用ForwardGetModel", hasGeminiAccount: true, hasAntigravity: false, expectedBehavior: "forward_to_upstream", diff --git a/backend/internal/handler/gemini_vertex_models_test.go b/backend/internal/handler/gemini_vertex_models_test.go new file mode 100644 index 0000000000..eaadf6e03f --- /dev/null +++ b/backend/internal/handler/gemini_vertex_models_test.go @@ -0,0 +1,83 @@ +package handler + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/server/middleware" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +type geminiVertexModelsTokenCache struct{} + +func (geminiVertexModelsTokenCache) GetAccessToken(context.Context, string) (string, error) { + return "ya29.vertex-token", nil +} + +func (geminiVertexModelsTokenCache) SetAccessToken(context.Context, string, string, time.Duration) error { + return nil +} + +func (geminiVertexModelsTokenCache) DeleteAccessToken(context.Context, string) error { + return nil +} + +func (geminiVertexModelsTokenCache) AcquireRefreshLock(context.Context, string, time.Duration) (bool, error) { + return false, nil +} + +func (geminiVertexModelsTokenCache) ReleaseRefreshLock(context.Context, string) error { + return nil +} + +// 分组里只有 Agent Platform (Vertex) 服务账号时，模型列表与模型详情接口返回 Vertex 上的 Gemini 模型。 +func TestGeminiV1BetaModelsEndpoints_VertexServiceAccountOnlyGroup(t *testing.T) { + groupID := int64(8) + repo := &geminiAllowlistAccountRepoStub{gatewayModelsAccountRepoStub: gatewayModelsAccountRepoStub{byGroup: map[int64][]service.Account{groupID: {{ + ID: 18, + Platform: service.PlatformGemini, + Type: service.AccountTypeServiceAccount, + Credentials: map[string]any{ + "location": "global", + "service_account_json": map[string]any{ + "type": "service_account", + "project_id": "proj", + "private_key": "-----BEGIN PRIVATE KEY-----\nabc\n-----END PRIVATE KEY-----\n", + "client_email": "svc@proj.iam.gserviceaccount.com", + }, + }, + }}}}} + upstream := &geminiMixedModelsUpstream{status: http.StatusOK, body: `{"publisherModels":[{"name":"publishers/google/models/gemini-3.8-flash"}]}`} + tokenProvider := service.NewGeminiTokenProvider(nil, geminiVertexModelsTokenCache{}, nil) + h := &GatewayHandler{geminiCompatService: service.NewGeminiMessagesCompatService(repo, nil, nil, nil, tokenProvider, nil, upstream, nil, &config.Config{})} + apiKey := &service.APIKey{GroupID: &groupID, Group: &service.Group{ID: groupID, Platform: service.PlatformGemini}} + + t.Run("list", func(t *testing.T) { + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodGet, "/v1beta/models", nil) + c.Set(string(middleware.ContextKeyAPIKey), apiKey) + + h.GeminiV1BetaListModels(c) + require.Equal(t, http.StatusOK, rec.Code, rec.Body.String()) + require.JSONEq(t, `{"models":[{"name":"models/gemini-3.8-flash","supportedGenerationMethods":["generateContent","streamGenerateContent"]}]}`, rec.Body.String()) + }) + + t.Run("get", func(t *testing.T) { + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodGet, "/v1beta/models/gemini-3.8-flash", nil) + c.Params = gin.Params{{Key: "model", Value: "gemini-3.8-flash"}} + c.Set(string(middleware.ContextKeyAPIKey), apiKey) + + h.GeminiV1BetaGetModel(c) + require.Equal(t, http.StatusOK, rec.Code, rec.Body.String()) + require.JSONEq(t, `{"name":"models/gemini-3.8-flash","supportedGenerationMethods":["generateContent","streamGenerateContent"]}`, rec.Body.String()) + }) +} diff --git a/backend/internal/handler/ops_error_logger.go b/backend/internal/handler/ops_error_logger.go index a4e30814ca..679cf4e8a2 100644 --- a/backend/internal/handler/ops_error_logger.go +++ b/backend/internal/handler/ops_error_logger.go @@ -1183,6 +1183,14 @@ func OpsErrorLoggerMiddleware(ops *service.OpsService) gin.HandlerFunc { } } + // Google 格式错误体没有 type 字段：按网关回写时登记的错误类型补齐。 + if errType, code := service.OpsLocalErrorType(c); errType != "" && !parsed.StreamFailure && + (parsed.ErrorType == "" || parsed.ErrorType == "api_error") { + parsed.ErrorType = errType + if code != "" { + parsed.Code = code + } + } normalizedType := normalizeOpsErrorType(parsed.ErrorType, parsed.Code) phase, isBusinessLimited, errorOwner, errorSource := classifyOpsErrorLog(c, normalizedType, parsed.Message, parsed.Code, status) @@ -2214,6 +2222,11 @@ func classifyOpsSeverity(errType string, status int) string { func classifyOpsErrorLog(c *gin.Context, errType, message, code string, status int) (phase string, isBusinessLimited bool, errorOwner string, errorSource string) { phase = classifyOpsPhase(errType, message, code) + if service.IsOpsRequestScopedError(c) { + // 失败由请求内容决定：只按错误类型分类，上游错误上下文不参与归因。 + isBusinessLimited = classifyOpsIsBusinessLimited(errType, phase, code, status, message) + return phase, isBusinessLimited, classifyOpsErrorOwner(phase, message), classifyOpsErrorSource(phase, message) + } routingCapacityLimited := isOpsRoutingCapacityLimited(c) clientBusinessLimited := service.HasOpsClientBusinessLimited(c) localModelConfiguration := clientBusinessLimited && service.OpsClientBusinessLimitedReason(c) == service.OpsClientBusinessLimitedReasonLocalModelConfiguration diff --git a/backend/internal/handler/ops_error_logger_google_test.go b/backend/internal/handler/ops_error_logger_google_test.go new file mode 100644 index 0000000000..392421dffd --- /dev/null +++ b/backend/internal/handler/ops_error_logger_google_test.go @@ -0,0 +1,135 @@ +package handler + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +// recordGoogleFormatOpsError 让 handle 以 Google 错误格式回写错误，返回 ops 中间件记录的错误日志。 +func recordGoogleFormatOpsError(t *testing.T, handle gin.HandlerFunc) *service.OpsInsertErrorLogInput { + t.Helper() + setupOpsErrorLogTestQueue(t, 1) + gin.SetMode(gin.TestMode) + ops := service.NewOpsService(nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil) + router := gin.New() + router.Use(OpsErrorLoggerMiddleware(ops)) + router.POST("/v1beta/cachedContents", handle) + + router.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodPost, "/v1beta/cachedContents", nil)) + + require.Equal(t, int64(1), OpsErrorLogQueueLength()) + return (<-opsErrorLogQueue).entry +} + +func TestOpsErrorLoggerMiddleware_GoogleClientErrorsClassifiedAsClient(t *testing.T) { + for _, tc := range []struct { + status int + errType string + phase string + }{ + {status: http.StatusBadRequest, errType: "invalid_request_error", phase: "request"}, + {status: http.StatusRequestEntityTooLarge, errType: "invalid_request_error", phase: "request"}, + {status: http.StatusUnauthorized, errType: "authentication_error", phase: "auth"}, + {status: http.StatusForbidden, errType: "permission_error", phase: "request"}, + {status: http.StatusNotFound, errType: "not_found_error", phase: "request"}, + } { + t.Run(http.StatusText(tc.status), func(t *testing.T) { + entry := recordGoogleFormatOpsError(t, func(c *gin.Context) { + googleError(c, tc.status, "client request problem") + }) + require.Equal(t, tc.status, entry.StatusCode) + require.Equal(t, tc.errType, entry.ErrorType) + require.Equal(t, tc.phase, entry.ErrorPhase) + require.Equal(t, "client", entry.ErrorOwner) + require.Equal(t, "client_request", entry.ErrorSource) + require.Equal(t, "P3", entry.Severity) + }) + } +} + +func TestOpsErrorLoggerMiddleware_GooglePlatformErrorStaysPlatform(t *testing.T) { + entry := recordGoogleFormatOpsError(t, func(c *gin.Context) { + googlePlatformError(c, http.StatusBadRequest, "Explicit context caching is not available for model m: storage pricing is not configured") + }) + require.Equal(t, "api_error", entry.ErrorType) + require.Equal(t, "internal", entry.ErrorPhase) + require.Equal(t, "platform", entry.ErrorOwner) + require.Equal(t, "gateway", entry.ErrorSource) +} + +func TestOpsErrorLoggerMiddleware_GoogleServerErrorsNotTypedByStatus(t *testing.T) { + entry := recordGoogleFormatOpsError(t, func(c *gin.Context) { + googleError(c, http.StatusInternalServerError, "Failed to persist cached content") + }) + require.Equal(t, "api_error", entry.ErrorType) + require.Equal(t, "platform", entry.ErrorOwner) +} + +func TestOpsErrorLoggerMiddleware_GoogleBillingErrorMatchesClaudeFormat(t *testing.T) { + entry := recordGoogleFormatOpsError(t, func(c *gin.Context) { + status, code, message, _ := billingErrorDetails(service.ErrInsufficientBalance) + googleErrorWithType(c, status, code, "", message) + }) + require.Equal(t, "billing_error", entry.ErrorType) + require.Equal(t, "request", entry.ErrorPhase) + require.Equal(t, "client", entry.ErrorOwner) + require.True(t, entry.IsBusinessLimited) +} + +func TestOpsErrorLoggerMiddleware_GoogleConcurrencyErrorClassifiedAsRequest(t *testing.T) { + entry := recordGoogleFormatOpsError(t, func(c *gin.Context) { + googleConcurrencyError(c, &ConcurrencyError{SlotType: "user"}, "user") + }) + require.Equal(t, http.StatusTooManyRequests, entry.StatusCode) + require.Equal(t, "rate_limit_error", entry.ErrorType) + require.Equal(t, "request", entry.ErrorPhase) + require.Equal(t, "client", entry.ErrorOwner) +} + +func TestOpsErrorLoggerMiddleware_GoogleErrorAfterUpstreamFailureKeepsUpstreamAttribution(t *testing.T) { + entry := recordGoogleFormatOpsError(t, func(c *gin.Context) { + service.SetOpsUpstreamError(c, http.StatusBadRequest, "upstream rejected", "") + googleError(c, http.StatusBadRequest, "upstream rejected") + }) + require.Equal(t, "upstream", entry.ErrorPhase) + require.Equal(t, "provider", entry.ErrorOwner) + require.Equal(t, "upstream_http", entry.ErrorSource) +} + +func TestOpsErrorLoggerMiddleware_RequestScopedUpstreamErrorClassifiedAsClient(t *testing.T) { + entry := recordGoogleFormatOpsError(t, func(c *gin.Context) { + service.SetOpsUpstreamError(c, http.StatusBadRequest, "Invalid resource state for cache content 8079767076522164224.", "") + c.Set(service.OpsUpstreamErrorsKey, []*service.OpsUpstreamErrorEvent{{ + Platform: service.PlatformGemini, AccountID: 173, UpstreamStatusCode: http.StatusBadRequest, Kind: "http_error", + Message: "Invalid resource state for cache content 8079767076522164224.", + }}) + service.MarkOpsRequestScopedError(c, "permission_error") + c.Data(http.StatusForbidden, "application/json", []byte(service.GeminiCachedContentNotFoundResponse)) + }) + require.Equal(t, http.StatusForbidden, entry.StatusCode) + require.Equal(t, "permission_error", entry.ErrorType) + require.Equal(t, "request", entry.ErrorPhase) + require.Equal(t, "client", entry.ErrorOwner) + require.Equal(t, "client_request", entry.ErrorSource) + require.False(t, entry.IsBusinessLimited) + require.NotNil(t, entry.UpstreamStatusCode, "上游原始错误保留为明细") + require.Equal(t, http.StatusBadRequest, *entry.UpstreamStatusCode) +} + +func TestSetOpsLocalErrorTypeClearsRequestScopedMark(t *testing.T) { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + service.MarkOpsRequestScopedError(c, "permission_error") + require.True(t, service.IsOpsRequestScopedError(c)) + + service.SetOpsLocalErrorType(c, "", "") + require.False(t, service.IsOpsRequestScopedError(c)) + errType, code := service.OpsLocalErrorType(c) + require.Empty(t, errType) + require.Empty(t, code) +} diff --git a/backend/internal/handler/security_audit_errors.go b/backend/internal/handler/security_audit_errors.go index b6f42c7111..7c55e9bfde 100644 --- a/backend/internal/handler/security_audit_errors.go +++ b/backend/internal/handler/security_audit_errors.go @@ -100,9 +100,14 @@ func googleSecurityAuditError(c *gin.Context, decision *securityaudit.Decision) return } if decision.Legacy != nil && decision.Legacy.Blocked { - googleError(c, securityAuditStatus(decision), securityAuditMessage(decision)) + googleErrorWithType(c, securityAuditStatus(decision), securityAuditErrorCode(decision), "", securityAuditMessage(decision)) return } + errType := "api_error" + if decision.Kind == securityaudit.DecisionBlock { + errType = "permission_error" + } + service.SetOpsLocalErrorType(c, errType, securityAuditErrorCode(decision)) status := securityAuditStatus(decision) googleStatus := googleapi.HTTPStatusToGoogleStatus(status) if status == http.StatusServiceUnavailable { diff --git a/backend/internal/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/api_key_repo.go b/backend/internal/repository/api_key_repo.go index 99796615fe..cc47f5e623 100644 --- a/backend/internal/repository/api_key_repo.go +++ b/backend/internal/repository/api_key_repo.go @@ -857,7 +857,7 @@ func (r *apiKeyRepository) ResetRateLimitWindows(ctx context.Context, id int64) // GetRateLimitData returns the current rate limit usage and window start times for an API key. func (r *apiKeyRepository) GetRateLimitData(ctx context.Context, id int64) (result *service.APIKeyRateLimitData, err error) { rows, err := r.sql.QueryContext(ctx, ` - SELECT usage_5h, usage_1d, usage_7d, window_5h_start, window_1d_start, window_7d_start + SELECT quota_used, usage_5h, usage_1d, usage_7d, window_5h_start, window_1d_start, window_7d_start FROM api_keys WHERE id = $1 AND deleted_at IS NULL`, id) @@ -873,7 +873,7 @@ func (r *apiKeyRepository) GetRateLimitData(ctx context.Context, id int64) (resu return nil, service.ErrAPIKeyNotFound } data := &service.APIKeyRateLimitData{} - if err := rows.Scan(&data.Usage5h, &data.Usage1d, &data.Usage7d, &data.Window5hStart, &data.Window1dStart, &data.Window7dStart); err != nil { + if err := rows.Scan(&data.QuotaUsed, &data.Usage5h, &data.Usage1d, &data.Usage7d, &data.Window5hStart, &data.Window1dStart, &data.Window7dStart); err != nil { return nil, err } return data, rows.Err() diff --git a/backend/internal/repository/api_key_repo_integration_test.go b/backend/internal/repository/api_key_repo_integration_test.go index ada0a201b3..9d41f269b6 100644 --- a/backend/internal/repository/api_key_repo_integration_test.go +++ b/backend/internal/repository/api_key_repo_integration_test.go @@ -464,6 +464,10 @@ func (s *APIKeyRepoSuite) TestIncrementQuotaUsed_Basic() { newQuota, err = s.repo.IncrementQuotaUsed(s.ctx, key.ID, 2.5) s.Require().NoError(err, "IncrementQuotaUsed second") s.Require().Equal(4.0, newQuota, "第二次递增后应为 4.0") + + data, err := s.repo.GetRateLimitData(s.ctx, key.ID) + s.Require().NoError(err, "GetRateLimitData") + s.Require().Equal(4.0, data.QuotaUsed, "用量读取带回总额度已用量") } func (s *APIKeyRepoSuite) TestIncrementQuotaUsed_NotFound() { diff --git a/backend/internal/repository/billing_inflight_cache.go b/backend/internal/repository/billing_inflight_cache.go index 898cf3b52a..8f83717ce6 100644 --- a/backend/internal/repository/billing_inflight_cache.go +++ b/backend/internal/repository/billing_inflight_cache.go @@ -25,6 +25,20 @@ func billingInflightKeys(userID int64) (string, string) { fmt.Sprintf("%s{%d}", billingInflightAmtKeyPrefix, userID) } +// 订阅 / API Key 维度在途预留（确定金额费用）： +// - billing:inflight_scope:{scope} ZSET member=requestID score=过期时间(ms) +// - billing:inflight_scope_amt:{scope} HASH field=requestID value=预留金额(USD) +// +// 过期清理与释放语义与余额维度相同，不续期。 +const ( + billingInflightScopeKeyPrefix = "billing:inflight_scope:" + billingInflightScopeAmtKeyPrefix = "billing:inflight_scope_amt:" +) + +func billingInflightScopeKeys(scope string) (string, string) { + return billingInflightScopeKeyPrefix + "{" + scope + "}", billingInflightScopeAmtKeyPrefix + "{" + scope + "}" +} + var ( // KEYS: [1]=zset [2]=hash // ARGV: [1]=now_ms [2]=expire_at_ms [3]=member [4]=amount [5]=balance [6]=key_ttl_ms @@ -85,6 +99,28 @@ var ( redis.call('HDEL', KEYS[2], ARGV[1]) return 1 `) + + // KEYS: [1]=zset [2]=hash + // ARGV: [1]=now_ms [2]=expire_at_ms [3]=member [4]=amount [5]=key_ttl_ms + // 先清理已过期成员，再登记本次预留；返回登记前其他在途预留的合计（string），是否放行由调用方按权威用量判定。 + reserveInflightScopeScript = redis.NewScript(` + local now = tonumber(ARGV[1]) + local expired = redis.call('ZRANGEBYSCORE', KEYS[1], '-inf', now) + if #expired > 0 then + redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', now) + redis.call('HDEL', KEYS[2], unpack(expired)) + end + local vals = redis.call('HVALS', KEYS[2]) + local sum = 0 + for i = 1, #vals do + sum = sum + (tonumber(vals[i]) or 0) + end + redis.call('ZADD', KEYS[1], tonumber(ARGV[2]), ARGV[3]) + redis.call('HSET', KEYS[2], ARGV[3], ARGV[4]) + redis.call('PEXPIRE', KEYS[1], ARGV[5]) + redis.call('PEXPIRE', KEYS[2], ARGV[5]) + return tostring(sum) + `) ) // ReserveInflightBalance 实现 service.InflightBalanceReservationCache。 @@ -138,7 +174,35 @@ func (c *billingCache) RenewInflightBalance(ctx context.Context, userID int64, r return n == 1, nil } +// ReserveInflightScope 实现 service.InflightScopeReservationCache。 +func (c *billingCache) ReserveInflightScope(ctx context.Context, scope, requestID string, amount float64, ttl time.Duration) (float64, error) { + zkey, hkey := billingInflightScopeKeys(scope) + now := time.Now().UnixMilli() + ttlMs := ttl.Milliseconds() + if ttlMs <= 0 { + ttlMs = 1 + } + res, err := reserveInflightScopeScript.Run(ctx, c.rdb, []string{zkey, hkey}, + now, + now+ttlMs, + requestID, + strconv.FormatFloat(amount, 'f', -1, 64), + ttlMs, + ).Text() + if err != nil { + return 0, err + } + return strconv.ParseFloat(res, 64) +} + +// ReleaseInflightScope 实现 service.InflightScopeReservationCache。 +func (c *billingCache) ReleaseInflightScope(ctx context.Context, scope, requestID string) error { + zkey, hkey := billingInflightScopeKeys(scope) + return releaseInflightBalanceScript.Run(ctx, c.rdb, []string{zkey, hkey}, requestID).Err() +} + var ( _ service.InflightBalanceReservationCache = (*billingCache)(nil) _ service.InflightBalanceReservationRenewer = (*billingCache)(nil) + _ service.InflightScopeReservationCache = (*billingCache)(nil) ) diff --git a/backend/internal/repository/billing_inflight_cache_test.go b/backend/internal/repository/billing_inflight_cache_test.go index ddab7c4de0..2fd642fe38 100644 --- a/backend/internal/repository/billing_inflight_cache_test.go +++ b/backend/internal/repository/billing_inflight_cache_test.go @@ -285,3 +285,72 @@ func TestInflightReservation_RenewDoesNotResurrectExpiredMember(t *testing.T) { require.NoError(t, err) require.False(t, exists) } + +func TestInflightScopeReservation_ReturnsOthersAndReleases(t *testing.T) { + mr, cache, _ := newInflightTestEnv(t, true, 60) + ctx := context.Background() + + others, err := cache.ReserveInflightScope(ctx, "apikey:7", "r1", 3, time.Minute) + require.NoError(t, err) + require.Zero(t, others) + others, err = cache.ReserveInflightScope(ctx, "apikey:7", "r2", 2, time.Minute) + require.NoError(t, err) + require.InDelta(t, 3, others, 1e-9, "返回登记前其他在途预留的合计") + others, err = cache.ReserveInflightScope(ctx, "sub:7", "r3", 5, time.Minute) + require.NoError(t, err) + require.Zero(t, others, "不同维度互不影响") + + require.NoError(t, cache.ReleaseInflightScope(ctx, "apikey:7", "r1")) + require.NoError(t, cache.ReleaseInflightScope(ctx, "apikey:7", "r1")) + others, err = cache.ReserveInflightScope(ctx, "apikey:7", "r4", 1, time.Minute) + require.NoError(t, err) + require.InDelta(t, 2, others, 1e-9) + + mr.FastForward(2 * time.Minute) + others, err = cache.ReserveInflightScope(ctx, "apikey:7", "r5", 1, time.Minute) + require.NoError(t, err) + require.Zero(t, others, "过期预留不再计入") +} + +func TestReserveFixedCharge_ConcurrentChargesAdmitOnlyWhatAPIKeyQuotaCovers(t *testing.T) { + _, cache, svc := newInflightTestEnv(t, true, 60) + ctx := context.Background() + user := &service.User{ID: 42} + apiKey := &service.APIKey{ID: 9, UserID: user.ID, Quota: 1.0} + require.NoError(t, cache.SetUserBalance(ctx, user.ID, 100)) + + const n = 20 + var admitted atomic.Int32 + reservations := make(chan *service.InflightReservation, n) + var wg sync.WaitGroup + start := make(chan struct{}) + for i := 0; i < n; i++ { + wg.Add(1) + go func() { + defer wg.Done() + <-start + res, err := svc.ReserveFixedCharge(ctx, user, apiKey, nil, nil, 0.3) + if err != nil { + if !errors.Is(err, service.ErrAPIKeyQuotaExhausted) { + t.Errorf("unexpected error: %v", err) + } + return + } + admitted.Add(1) + reservations <- res + }() + } + close(start) + wg.Wait() + close(reservations) + + require.Equal(t, int32(3), admitted.Load(), "Key 额度 1.0、每笔 0.3：同时在途的不超过 3 笔") + for res := range reservations { + res.Release() + } + zkey, _ := billingInflightScopeKeys("apikey:9") + remaining, err := cache.rdb.ZCard(ctx, zkey).Result() + require.NoError(t, err) + require.Zero(t, remaining) + require.Equal(t, int64(0), inflightCount(t, cache, user.ID)) +} diff --git a/backend/internal/repository/gemini_cached_content_repo.go b/backend/internal/repository/gemini_cached_content_repo.go new file mode 100644 index 0000000000..42775448d7 --- /dev/null +++ b/backend/internal/repository/gemini_cached_content_repo.go @@ -0,0 +1,205 @@ +package repository + +import ( + "context" + "database/sql" + "encoding/json" + "errors" + "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, channel_usage, created_at, updated_at` + +func (r *geminiCachedContentRepository) Create(ctx context.Context, record *service.GeminiCachedContent) error { + channelUsage, err := json.Marshal(record.ChannelUsage) + if err != nil { + return err + } + row := r.sql.QueryRowContext(ctx, ` +INSERT INTO gemini_cached_contents (public_id, user_id, api_key_id, group_id, account_id, upstream_name, model, + upstream_model, display_name, total_token_count, expire_time, channel_usage) +VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12) +RETURNING id, created_at, updated_at`, + record.PublicID, record.UserID, record.APIKeyID, record.GroupID, record.AccountID, record.UpstreamName, + record.Model, record.UpstreamModel, record.DisplayName, record.TotalTokenCount, record.ExpireTime, channelUsage) + 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) BeginMutation(ctx context.Context, id int64, baseline time.Time, token string, lease time.Duration) (bool, error) { + result, err := r.sql.ExecContext(ctx, `UPDATE gemini_cached_contents + SET mutation_token = $2, mutation_lease_until = NOW() + make_interval(secs => $3::double precision) + WHERE id = $1 AND deleted_at IS NULL AND mutation_token IS NULL AND expire_time = $4`, + id, token, lease.Seconds(), baseline) + if err != nil { + return false, err + } + affected, err := result.RowsAffected() + if err != nil { + return false, err + } + if affected > 0 { + return true, nil + } + var exists bool + if err := r.sql.QueryRowContext(ctx, `SELECT EXISTS ( + SELECT 1 FROM gemini_cached_contents WHERE id = $1 AND deleted_at IS NULL)`, id).Scan(&exists); err != nil { + return false, err + } + if !exists { + return false, service.ErrGeminiCachedContentNotFound + } + return false, nil +} + +func (r *geminiCachedContentRepository) TakeOverStaleMutation(ctx context.Context, id int64, token string, lease time.Duration) (*service.GeminiCachedContent, error) { + record, err := scanGeminiCachedContent(r.sql.QueryRowContext(ctx, `UPDATE gemini_cached_contents + SET mutation_token = $2, mutation_lease_until = NOW() + make_interval(secs => $3::double precision) + WHERE id = $1 AND deleted_at IS NULL AND mutation_token IS NOT NULL AND mutation_lease_until < NOW() + RETURNING `+geminiCachedContentColumns, id, token, lease.Seconds())) + if errors.Is(err, sql.ErrNoRows) { + return nil, nil + } + if err != nil { + return nil, err + } + return record, nil +} + +func (r *geminiCachedContentRepository) CommitMutation(ctx context.Context, id int64, token string, expireTime time.Time) error { + result, err := r.sql.ExecContext(ctx, `UPDATE gemini_cached_contents + SET expire_time = $3, updated_at = NOW(), mutation_token = NULL, mutation_lease_until = NULL + WHERE id = $1 AND deleted_at IS NULL AND mutation_token = $2`, id, token, expireTime) + if err != nil { + return err + } + affected, err := result.RowsAffected() + if err != nil { + return err + } + if affected == 0 { + return service.ErrGeminiCachedContentMutationConflict + } + return nil +} + +func (r *geminiCachedContentRepository) ReleaseMutation(ctx context.Context, id int64, token string) error { + _, err := r.sql.ExecContext(ctx, `UPDATE gemini_cached_contents + SET mutation_token = NULL, mutation_lease_until = NULL + WHERE id = $1 AND mutation_token = $2`, id, token) + return err +} + +func (r *geminiCachedContentRepository) ListStaleMutations(ctx context.Context, limit int) ([]int64, error) { + rows, err := r.sql.QueryContext(ctx, `SELECT id FROM gemini_cached_contents + WHERE mutation_token IS NOT NULL AND deleted_at IS NULL AND mutation_lease_until < NOW() + ORDER BY mutation_lease_until + LIMIT $1`, limit) + if err != nil { + return nil, err + } + defer func() { _ = rows.Close() }() + var ids []int64 + for rows.Next() { + var id int64 + if err := rows.Scan(&id); err != nil { + return nil, err + } + ids = append(ids, id) + } + return ids, rows.Err() +} + +func (r *geminiCachedContentRepository) SoftDelete(ctx context.Context, id int64) error { + _, 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 + var channelUsage []byte + 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, + &channelUsage, + &record.CreatedAt, + &record.UpdatedAt, + ); err != nil { + return nil, err + } + if len(channelUsage) > 0 { + if err := json.Unmarshal(channelUsage, &record.ChannelUsage); err != nil { + return nil, err + } + } + return &record, nil +} diff --git a/backend/internal/repository/gemini_cached_content_repo_integration_test.go b/backend/internal/repository/gemini_cached_content_repo_integration_test.go new file mode 100644 index 0000000000..89399552c2 --- /dev/null +++ b/backend/internal/repository/gemini_cached_content_repo_integration_test.go @@ -0,0 +1,120 @@ +//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, + ChannelUsage: service.ChannelUsageFields{ChannelID: 7, OriginalModel: "gemini-pro", ChannelMappedModel: "gemini-3.8-flash", BillingModelSource: "requested"}, + } + } + 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)) + require.Equal(t, first.ChannelUsage, got.ChannelUsage, "创建时的渠道计费口径随记录持久化") + _, 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) + + lease := 2 * time.Minute + acquired, err := repo.BeginMutation(ctx, first.ID, first.ExpireTime, "t1", lease) + require.NoError(t, err) + require.True(t, acquired) + acquired, err = repo.BeginMutation(ctx, first.ID, first.ExpireTime, "t2", lease) + require.NoError(t, err) + require.False(t, acquired, "有进行中的修改时不能再登记") + taken, err := repo.TakeOverStaleMutation(ctx, first.ID, "t2", lease) + require.NoError(t, err) + require.Nil(t, taken, "租约未过期不能接管") + require.ErrorIs(t, repo.CommitMutation(ctx, first.ID, "t2", now), service.ErrGeminiCachedContentMutationConflict) + + newExpire := now.Add(3 * time.Hour) + require.NoError(t, repo.CommitMutation(ctx, first.ID, "t1", newExpire)) + got, err = repo.GetForOwner(ctx, apiKeyID, first.PublicID) + require.NoError(t, err) + require.True(t, got.ExpireTime.Equal(newExpire)) + require.ErrorIs(t, repo.CommitMutation(ctx, first.ID, "t1", now), service.ErrGeminiCachedContentMutationConflict, "提交后修改意图已清除") + + acquired, err = repo.BeginMutation(ctx, first.ID, first.ExpireTime, "t3", lease) + require.NoError(t, err) + require.False(t, acquired, "到期时间已被其他修改改变") + acquired, err = repo.BeginMutation(ctx, first.ID, newExpire, "t3", lease) + require.NoError(t, err) + require.True(t, acquired) + stale, err := repo.ListStaleMutations(ctx, 10) + require.NoError(t, err) + require.NotContains(t, stale, first.ID) + + _, err = tx.ExecContext(ctx, `UPDATE gemini_cached_contents SET mutation_lease_until = NOW() - INTERVAL '1 second' WHERE id = $1`, first.ID) + require.NoError(t, err) + stale, err = repo.ListStaleMutations(ctx, 10) + require.NoError(t, err) + require.Contains(t, stale, first.ID) + taken, err = repo.TakeOverStaleMutation(ctx, first.ID, "t4", lease) + require.NoError(t, err) + require.NotNil(t, taken) + require.True(t, taken.ExpireTime.Equal(newExpire), "接管返回修改前的到期时间") + require.ErrorIs(t, repo.CommitMutation(ctx, first.ID, "t3", now), service.ErrGeminiCachedContentMutationConflict, "被接管后原持有者不能提交") + require.NoError(t, repo.ReleaseMutation(ctx, first.ID, "t3")) + acquired, err = repo.BeginMutation(ctx, first.ID, newExpire, "t5", lease) + require.NoError(t, err) + require.False(t, acquired, "原持有者释放不影响接管者") + require.NoError(t, repo.ReleaseMutation(ctx, first.ID, "t4")) + stale, err = repo.ListStaleMutations(ctx, 10) + require.NoError(t, err) + require.NotContains(t, stale, first.ID) + + require.NoError(t, repo.SoftDelete(ctx, first.ID)) + _, err = repo.GetForOwner(ctx, apiKeyID, first.PublicID) + require.ErrorIs(t, err, service.ErrGeminiCachedContentNotFound) + _, err = repo.BeginMutation(ctx, first.ID, newExpire, "t6", lease) + require.ErrorIs(t, err, service.ErrGeminiCachedContentNotFound) + list, err = repo.ListForOwner(ctx, apiKeyID, 42, now, 0, 10) + require.NoError(t, err) + require.Len(t, list, 1) +} diff --git a/backend/internal/repository/scheduler_cache.go b/backend/internal/repository/scheduler_cache.go index e3715f678e..827a750577 100644 --- a/backend/internal/repository/scheduler_cache.go +++ b/backend/internal/repository/scheduler_cache.go @@ -956,7 +956,11 @@ func filterSchedulerCredentials(credentials map[string]any) map[string]any { } // Candidate-list admission evaluates the account override before hydrating // the full account. Dropping it silently falls back to the platform threshold. - keys := []string{"model_mapping", "compact_model_mapping", "api_key", "project_id", "oauth_type", "plan_type", "account_scheduling_threshold"} + keys := []string{"model_mapping", "compact_model_mapping", "api_key", "project_id", "oauth_type", "plan_type", "account_scheduling_threshold", + // Gemini 显式缓存的候选过滤（GeminiExplicitCacheExclusions）读本投影判定账号是否支持显式缓存： + // 官方地址、Vertex 认证方式与上游网关的显式缓存声明。缺这几个键时只能按类型粗筛， + // 不支持的账号要逐个选中后才被剔除。 + "base_url", "auth_mode", service.GeminiExplicitCacheUpstreamCredentialKey, service.GeminiExplicitCacheUpstreamMaxTTLCredentialKey} filtered := make(map[string]any) for _, key := range keys { if value, ok := credentials[key]; ok && value != nil { diff --git a/backend/internal/repository/scheduler_cache_test.go b/backend/internal/repository/scheduler_cache_test.go index d12e77f74c..0542842fdf 100644 --- a/backend/internal/repository/scheduler_cache_test.go +++ b/backend/internal/repository/scheduler_cache_test.go @@ -20,6 +20,23 @@ func TestFilterSchedulerCredentialsKeepsSubscriptionPlanType(t *testing.T) { require.NotContains(t, filtered, "refresh_token") } +func TestFilterSchedulerCredentialsKeepsGeminiExplicitCacheKeys(t *testing.T) { + filtered := filterSchedulerCredentials(map[string]any{ + "api_key": "k", + "base_url": "https://gw.example.com", + "auth_mode": "apikey", + "service_account_json": "secret", + service.GeminiExplicitCacheUpstreamCredentialKey: true, + service.GeminiExplicitCacheUpstreamMaxTTLCredentialKey: 3600, + }) + + require.Equal(t, "https://gw.example.com", filtered["base_url"]) + require.Equal(t, "apikey", filtered["auth_mode"]) + require.Equal(t, true, filtered[service.GeminiExplicitCacheUpstreamCredentialKey]) + require.Equal(t, 3600, filtered[service.GeminiExplicitCacheUpstreamMaxTTLCredentialKey]) + require.NotContains(t, filtered, "service_account_json") +} + func TestSchedulerMetadataAccountKeepsOpenAISubscriptionIdentity(t *testing.T) { account := service.Account{ ID: 24, diff --git a/backend/internal/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/account.go b/backend/internal/service/account.go index f5ce7c8de0..d394e551a2 100644 --- a/backend/internal/service/account.go +++ b/backend/internal/service/account.go @@ -339,6 +339,13 @@ func (a *Account) IsGeminiGoogleOne() bool { return a.Platform == PlatformGemini && a.Type == AccountTypeOAuth && a.GeminiOAuthType() == "google_one" } +// IsGeminiThirdPartyAPIKey reports whether a Gemini API Key account points at a +// non-Google base URL. Such upstreams carry no AI Studio tier or daily quota. +func (a *Account) IsGeminiThirdPartyAPIKey() bool { + return a != nil && a.Platform == PlatformGemini && a.Type == AccountTypeAPIKey && + !isOfficialGeminiAPIBaseURL(a.GetCredential("base_url")) +} + func (a *Account) CanGetUsage() bool { return a.Type == AccountTypeOAuth } diff --git a/backend/internal/service/account_stats_pricing_test.go b/backend/internal/service/account_stats_pricing_test.go index 26eb645935..0f5d236ca1 100644 --- a/backend/internal/service/account_stats_pricing_test.go +++ b/backend/internal/service/account_stats_pricing_test.go @@ -933,7 +933,7 @@ func TestResolveAccountStatsCost_Gemini36FlashTierUsesFallbackPricing(t *testing UsageTokens{InputTokens: 1_000_000, OutputTokens: 1_000_000, CacheReadTokens: 1_000_000}, 1, 0, "", time.Time{}, true, ) require.NotNil(t, result) - require.InDelta(t, 9.15, *result, 1e-12) + require.InDelta(t, 4.575, *result, 1e-12) } func TestResolveAccountStatsCost_AllMiss_ReturnsNil(t *testing.T) { diff --git a/backend/internal/service/account_test_service.go b/backend/internal/service/account_test_service.go index 5ae0859a8f..26060efc78 100644 --- a/backend/internal/service/account_test_service.go +++ b/backend/internal/service/account_test_service.go @@ -2585,14 +2585,7 @@ func (s *AccountTestService) buildGeminiOAuthRequest(ctx context.Context, accoun } func (s *AccountTestService) buildGeminiServiceAccountRequest(ctx context.Context, account *Account, modelID string, payload []byte) (*http.Request, error) { - if s.geminiTokenProvider == nil { - return nil, fmt.Errorf("gemini token provider not configured") - } - accessToken, err := s.geminiTokenProvider.GetAccessToken(ctx, account) - if err != nil { - return nil, fmt.Errorf("failed to get service account access token: %w", err) - } - fullURL, err := buildVertexGeminiURL(account.VertexProjectID(), account.VertexLocation(modelID), modelID, "streamGenerateContent", true) + fullURL, err := buildVertexGeminiAccountURL(account, modelID, "streamGenerateContent", true) if err != nil { return nil, err } @@ -2601,7 +2594,9 @@ func (s *AccountTestService) buildGeminiServiceAccountRequest(ctx context.Contex return nil, err } req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", "Bearer "+accessToken) + if err := setVertexGeminiAuth(ctx, req, account, s.geminiTokenProvider); err != nil { + return nil, fmt.Errorf("failed to authenticate vertex request: %w", err) + } return req, nil } diff --git a/backend/internal/service/api_key_service.go b/backend/internal/service/api_key_service.go index b00b03d43f..e20a2697eb 100644 --- a/backend/internal/service/api_key_service.go +++ b/backend/internal/service/api_key_service.go @@ -130,6 +130,8 @@ type apiKeyAllByUserIDLister interface { // APIKeyRateLimitData holds rate limit usage and window state for an API key. type APIKeyRateLimitData struct { + // QuotaUsed 是总额度已用量（quota_used）。 + QuotaUsed float64 Usage5h float64 Usage1d float64 Usage7d float64 diff --git a/backend/internal/service/billing_cache_service.go b/backend/internal/service/billing_cache_service.go index 691d2558f7..ea22efe2eb 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,110 @@ func (s *BillingCacheService) CheckBillingEligibility(ctx context.Context, user return nil } +// ReserveFixedCharge 为扣费前即可确定金额的费用（倍率后 USD，如显式缓存存储费）做准入并登记在途预留。 +// +// 余额模式要求缓存余额覆盖 amount 及用户的其他在途预留；订阅模式要求日 / 周 / 月剩余额度覆盖 amount； +// API Key 总额度与限速窗口同样叠加 amount。订阅与 API Key 的用量取库中权威值，并叠加同一订阅 / Key 上其他 +// 在途预留：先登记本次预留再读用量，并发的两笔费用中后读用量的一方必然看到另一方（预留或已落库的扣费）。 +// 返回的句柄可能为 nil（未登记任何预留）；调用方须在 handler 结束时调用 HandlerDone，预留在计费任务落库后释放。 +// 须在 CheckBillingEligibility 通过后调用；简易模式不校验。 +func (s *BillingCacheService) ReserveFixedCharge(ctx context.Context, user *User, apiKey *APIKey, group *Group, subscription *UserSubscription, amount float64) (*InflightReservation, error) { + if user == nil || amount <= 0 || math.IsNaN(amount) || math.IsInf(amount, 0) { + return nil, nil + } + if s.cfg != nil && s.cfg.RunMode == config.RunModeSimple { + return nil, nil + } + subscriptionBill := group != nil && group.IsSubscriptionType() && subscription != nil + if !subscriptionBill { + balance, err := s.GetUserBalance(ctx, user.ID) + if err != nil { + return nil, ErrBillingServiceUnavailable.WithCause(err) + } + if balance < amount { + return nil, ErrInsufficientBalance + } + } + res, err := s.reserveInflight(ctx, user, group, subscription, amount, false, false) + if err != nil { + return nil, err + } + var subscriptionPending, apiKeyPending float64 + if subscriptionBill && (group.HasDailyLimit() || group.HasWeeklyLimit() || group.HasMonthlyLimit()) { + res, subscriptionPending = s.reserveInflightScope(ctx, res, user.ID, inflightSubscriptionScope(subscription.ID), amount) + } + if apiKey != nil && (apiKey.Quota > 0 || apiKey.HasRateLimits()) { + res, apiKeyPending = s.reserveInflightScope(ctx, res, user.ID, inflightAPIKeyScope(apiKey.ID), amount) + } + if subscriptionBill { + err = s.checkFixedChargeSubscription(ctx, user.ID, group, amount+subscriptionPending) + } + if err == nil && apiKey != nil { + err = s.checkFixedChargeAPIKey(ctx, apiKey, amount+apiKeyPending) + } + if err != nil { + res.Release() + return nil, err + } + return res, nil +} + +// checkFixedChargeSubscription 按库中订阅用量校验日 / 周 / 月剩余额度覆盖 pending。 +func (s *BillingCacheService) checkFixedChargeSubscription(ctx context.Context, userID int64, group *Group, pending float64) error { + read := s.GetSubscriptionStatus + if s.subRepo != nil { + read = s.getSubscriptionFromDB + } + subData, err := read(ctx, userID, group.ID) + if err != nil { + return ErrBillingServiceUnavailable.WithCause(err) + } + if group.HasDailyLimit() && subData.DailyUsage+pending > *group.DailyLimitUSD { + return ErrDailyLimitExceeded + } + if group.HasWeeklyLimit() && subData.WeeklyUsage+pending > *group.WeeklyLimitUSD { + return ErrWeeklyLimitExceeded + } + if group.HasMonthlyLimit() && subData.MonthlyUsage+pending > *group.MonthlyLimitUSD { + return ErrMonthlyLimitExceeded + } + return nil +} + +// checkFixedChargeAPIKey 按库中 Key 用量校验总额度与限速窗口覆盖 pending；未接入仓储时使用请求上下文与限速缓存中的用量。 +func (s *BillingCacheService) checkFixedChargeAPIKey(ctx context.Context, apiKey *APIKey, pending float64) error { + if apiKey.Quota <= 0 && !apiKey.HasRateLimits() { + return nil + } + if s.apiKeyRateLimitLoader == nil { + if apiKey.Quota > 0 && apiKey.QuotaUsed+pending > apiKey.Quota { + return ErrAPIKeyQuotaExhausted + } + if apiKey.HasRateLimits() { + return s.checkAPIKeyRateLimitsWithPending(ctx, apiKey, pending) + } + return nil + } + data, err := s.apiKeyRateLimitLoader.GetRateLimitData(ctx, apiKey.ID) + if err != nil { + return ErrBillingServiceUnavailable.WithCause(err) + } + exceeds := func(usage, limit float64) bool { + return limit > 0 && usage+pending > limit + } + switch { + case exceeds(data.QuotaUsed, apiKey.Quota): + return ErrAPIKeyQuotaExhausted + case exceeds(data.EffectiveUsage5h(), apiKey.RateLimit5h): + return ErrAPIKeyRateLimit5hExceeded + case exceeds(data.EffectiveUsage1d(), apiKey.RateLimit1d): + return ErrAPIKeyRateLimit1dExceeded + case exceeds(data.EffectiveUsage7d(), apiKey.RateLimit7d): + return ErrAPIKeyRateLimit7dExceeded + } + return nil +} + // 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..e63c9db339 100644 --- a/backend/internal/service/billing_inflight_reservation.go +++ b/backend/internal/service/billing_inflight_reservation.go @@ -3,6 +3,7 @@ package service import ( "context" "math" + "strconv" "strings" "sync" "sync/atomic" @@ -25,6 +26,16 @@ type InflightBalanceReservationCache interface { ReleaseInflightBalance(ctx context.Context, userID int64, requestID string) error } +// InflightScopeReservationCache 订阅 / API Key 维度在途预留的缓存能力（可选），用于扣费前即可确定金额的费用； +// 未实现时这两个维度不登记预留（fail-open），准入仍按库中用量校验。 +type InflightScopeReservationCache interface { + // ReserveInflightScope 原子地清理 scope 下的过期预留并登记 requestID 的预留（ttl 后自动失效）， + // 返回登记前该 scope 下其他在途预留的合计。 + ReserveInflightScope(ctx context.Context, scope, requestID string, amount float64, ttl time.Duration) (float64, error) + // ReleaseInflightScope 释放 scope 下 requestID 的预留（幂等）。 + ReleaseInflightScope(ctx context.Context, scope, requestID string) error +} + // InflightBalanceReservationRenewer 可选：续期仍存活的预留（流式长请求期间防止 TTL 过期）。 type InflightBalanceReservationRenewer interface { // RenewInflightBalance 若 requestID 仍存在则把其过期时间推迟到 now+ttl；返回是否仍存在。 @@ -54,11 +65,15 @@ const ( // 不再续期，所以最长持有时间受 TTL 约束（计费任务卡死时预留自动过期）。 // 所有方法对 nil 接收者安全。 type InflightReservation struct { + // cache 为 nil 表示未登记余额维度的预留。 cache InflightBalanceReservationCache userID int64 requestID string amount float64 ttl time.Duration + // scopes 是同一金额在订阅 / API Key 维度登记的预留（与余额维度共用 requestID），不续期。 + scopeCache InflightScopeReservationCache + scopes []string refs atomic.Int64 releaseOnce sync.Once @@ -116,8 +131,15 @@ func (r *InflightReservation) Release() { r.refs.Store(0) relCtx, relCancel := context.WithTimeout(context.Background(), inflightReservationReleaseTimeout) defer relCancel() - if err := r.cache.ReleaseInflightBalance(relCtx, r.userID, r.requestID); err != nil { - logger.LegacyPrintf("service.billing_cache", "Warning: inflight reservation release failed for user %d (expires by ttl): %v", r.userID, err) + if r.cache != nil { + if err := r.cache.ReleaseInflightBalance(relCtx, r.userID, r.requestID); err != nil { + logger.LegacyPrintf("service.billing_cache", "Warning: inflight reservation release failed for user %d (expires by ttl): %v", r.userID, err) + } + } + for _, scope := range r.scopes { + if err := r.scopeCache.ReleaseInflightScope(relCtx, scope, r.requestID); err != nil { + logger.LegacyPrintf("service.billing_cache", "Warning: inflight reservation release failed for %s (expires by ttl): %v", scope, err) + } } }) } @@ -213,7 +235,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 +254,50 @@ 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) { +// inflightSubscriptionScope / inflightAPIKeyScope 是订阅 / API Key 维度在途预留的 scope 名。 +func inflightSubscriptionScope(subscriptionID int64) string { + return "sub:" + strconv.FormatInt(subscriptionID, 10) +} + +func inflightAPIKeyScope(apiKeyID int64) string { + return "apikey:" + strconv.FormatInt(apiKeyID, 10) +} + +// reserveInflightScope 在 scope 维度为 amount 登记在途预留并挂到 res 上（res 为 nil 时新建句柄）， +// 返回句柄与登记前该 scope 下其他在途预留的合计。开关关闭 / 缓存不支持 / Redis 执行失败时不登记（fail-open），合计按 0。 +func (s *BillingCacheService) reserveInflightScope(ctx context.Context, res *InflightReservation, userID int64, scope string, amount float64) (*InflightReservation, float64) { + cfg, ok := s.inflightReservationConfig() + if !ok { + return res, 0 + } + sc, ok := s.cache.(InflightScopeReservationCache) + if !ok || sc == nil { + return res, 0 + } + if res == nil { + ttl := defaultInflightReservationTTL + if cfg.TTLSeconds > 0 { + ttl = time.Duration(cfg.TTLSeconds) * time.Second + } + res = &InflightReservation{userID: userID, requestID: uuid.NewString(), amount: amount, ttl: ttl} + res.refs.Store(1) + } + reserveCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), inflightReservationReserveTimeout) + pending, err := sc.ReserveInflightScope(reserveCtx, scope, res.requestID, amount, res.ttl) + cancel() + if err != nil { + logger.LegacyPrintf("service.billing_cache", "Warning: inflight reservation failed for %s (fail-open): %v", scope, err) + return res, 0 + } + res.scopeCache = sc + res.scopes = append(res.scopes, scope) + return res, pending +} + +func (s *BillingCacheService) reserveInflight(ctx context.Context, user *User, group *Group, subscription *UserSubscription, estimate float64, renew, capped bool) (*InflightReservation, error) { cfg, ok := s.inflightReservationConfig() if !ok || user == nil { return nil, nil @@ -243,7 +305,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 967f2ba778..a1fe74ebee 100644 --- a/backend/internal/service/billing_service.go +++ b/backend/internal/service/billing_service.go @@ -488,14 +488,15 @@ func (s *BillingService) initFallbackPricing() { SupportsCacheBreakdown: false, } - // Gemini 3.6 Flash (Google AI pricing: $1.50 input / $7.50 output / - // $0.15 cached input per MTok). Antigravity's -high/-low/-medium/-tiered - // aliases are matched below so unavailable remote pricing never records - // token-bearing requests at $0. + // Gemini 3.6 Flash (Google AI pricing: $0.75 input / $3.75 output / + // $0.075 cached input per MTok, promotional through 2026-12-31; official + // rates double to $1.50/$7.50/$0.15 from 2027-01-01). Antigravity's + // -high/-low/-medium/-tiered aliases are matched below so unavailable + // remote pricing never records token-bearing requests at $0. s.fallbackPrices["gemini-3.6-flash"] = &ModelPricing{ - InputPricePerToken: 1.5e-6, - OutputPricePerToken: 7.5e-6, - CacheReadPricePerToken: 0.15e-6, + InputPricePerToken: 0.75e-6, + OutputPricePerToken: 3.75e-6, + CacheReadPricePerToken: 0.075e-6, SupportsCacheBreakdown: false, } @@ -1342,6 +1343,18 @@ func (s *BillingService) GetModelPricing(model string) (*ModelPricing, error) { return s.getModelPricingAt(model, timezone.Now()) } +// CacheStoragePricePerTokenHour 返回模型的显式上下文缓存存储单价（每 token 每小时，USD）。 +// 优先级：价格表（override_file 等）按确定性识别配置的值 > 内置型号价 > 内置 Pro/Flash 系列兜底价； +// 均未命中返回 0。 +func (s *BillingService) CacheStoragePricePerTokenHour(model string) float64 { + if s != nil && s.pricingService != nil { + if pricing := s.pricingService.GetIdentifiedModelPricing(model); pricing != nil && pricing.CacheStorageCostPerTokenPerHour > 0 { + return pricing.CacheStorageCostPerTokenPerHour + } + } + return builtinGeminiCacheStoragePrice(model) +} + // getModelPricingAt 是 GetModelPricing 的带计费时点内部变体：pricingAt 显式 // 驱动 DeepSeek pro→Flash 切换判定（切换点前 Pro 价、之后 Flash 价），使 // 展示/估算路径可与历史补账同刻复算，测试也能用固定时点钉住断言。 diff --git a/backend/internal/service/channel_monitor_checker.go b/backend/internal/service/channel_monitor_checker.go index ddb5e78317..c8b4fdf069 100644 --- a/backend/internal/service/channel_monitor_checker.go +++ b/backend/internal/service/channel_monitor_checker.go @@ -636,6 +636,8 @@ var monitorAPIKeyPatterns = []struct { {regexp.MustCompile(`xai-[A-Za-z0-9_-]{6,}`), "xai-***REDACTED***"}, // Gemini / Google API Key：固定前缀 + 35 位 {regexp.MustCompile(`AIza[A-Za-z0-9_-]{35}`), "AIza***REDACTED***"}, + // Gemini auth key：AQ.Ab 前缀 + 48 位，长度放宽以覆盖格式变体 + {regexp.MustCompile(`AQ\.Ab[A-Za-z0-9_-]{40,}`), "AQ.***REDACTED***"}, // JWT 三段式（Bearer 后常出现）：eyJxxx.eyJxxx.signature {regexp.MustCompile(`eyJ[A-Za-z0-9_-]{8,}\.eyJ[A-Za-z0-9_-]{8,}\.[A-Za-z0-9_-]{8,}`), "eyJ***REDACTED.JWT***"}, } @@ -643,7 +645,7 @@ var monitorAPIKeyPatterns = []struct { // sanitizeErrorMessage 擦除错误/响应文本中可能泄露的 API key。 // 处理两类来源： // 1. URL query 中的 ?key= / ?api_key= 等（Go *url.Error 会回填完整 URL） -// 2. 上游 HTTP body 文本里直接出现的 sk-* / xai-* / AIza* / JWT 等密钥碎片 +// 2. 上游 HTTP body 文本里直接出现的 sk-* / xai-* / AIza* / AQ.* / JWT 等密钥碎片 // // 注意：与 gemini_messages_compat_service.go 的 sanitizeUpstreamErrorMessage 关注点类似但参数集更广， // 监控模块独立维护，避免互相耦合。 diff --git a/backend/internal/service/channel_monitor_checker_body_test.go b/backend/internal/service/channel_monitor_checker_body_test.go index 922eb86732..5c347cae89 100644 --- a/backend/internal/service/channel_monitor_checker_body_test.go +++ b/backend/internal/service/channel_monitor_checker_body_test.go @@ -279,6 +279,17 @@ func TestRunCheckForModel_Grok_RedactsXAIKeyFromUpstreamBody(t *testing.T) { } } +func TestSanitizeErrorMessage_RedactsGeminiAuthKey(t *testing.T) { + authKey := "AQ." + "Ab" + strings.Repeat("Xy9_-z", 8) + got := sanitizeErrorMessage(`{"error":{"message":"API key not valid: ` + authKey + `"}}`) + if strings.Contains(got, authKey) { + t.Fatalf("Gemini auth key leaked: %q", got) + } + if !strings.Contains(got, "AQ.***REDACTED***") { + t.Fatalf("expected redaction marker, got %q", got) + } +} + func TestRunCheckForModel_OpenAIResponses_DefaultRequest(t *testing.T) { h := &openAICaptureHandler{} endpoint := setupFakeOpenAI(t, h) diff --git a/backend/internal/service/gateway_service.go b/backend/internal/service/gateway_service.go index c22097ede1..48c74e7d08 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 @@ -696,6 +699,7 @@ type UpstreamFailoverError struct { SameAccountRetryMax int // 可选的错误级同账号重试上限，低于 handler 默认预算时优先采用 RequestScopedTransient bool // 故障因素与账号无关（如上游按客户端身份/模型容量降载）：可同账号重试，但不得据此对账号做临时封禁 SafeToFailoverAfterWrite bool // 仅写出 SSE 注释等非语义字节时，仍可在同一客户端流中切换账号 + BoundUpstreamPassthrough bool // 绑定显式缓存的请求在上游网关账号上收到的 429 / 5xx：反映持有缓存的上游账号状态，原样回写不走错误映射 Stage GatewayFailureStage Scope GatewayFailureScope Reason GatewayFailureReason diff --git a/backend/internal/service/gateway_usage_billing.go b/backend/internal/service/gateway_usage_billing.go index 3a612d17d0..84c604fd04 100644 --- a/backend/internal/service/gateway_usage_billing.go +++ b/backend/internal/service/gateway_usage_billing.go @@ -241,7 +241,8 @@ func isForcedUsageBillingRequestID(requestID string) bool { return strings.HasPrefix(id, "web_search:") || strings.HasPrefix(id, "grok-video:") || strings.HasPrefix(id, "grok_audio:") || - strings.HasPrefix(id, "grok_realtime:") + strings.HasPrefix(id, "grok_realtime:") || + strings.HasPrefix(id, GeminiCachedContentUsageRequestIDPrefix) } // StableGrokAudioBillingRequestID is the durable usage_logs / dedup key for one @@ -787,42 +788,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 +808,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 +834,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 +848,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 +897,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 +1040,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 +1061,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/gateway_usage_billing_request_id_test.go b/backend/internal/service/gateway_usage_billing_request_id_test.go index 78a64cb338..fa13918b51 100644 --- a/backend/internal/service/gateway_usage_billing_request_id_test.go +++ b/backend/internal/service/gateway_usage_billing_request_id_test.go @@ -31,9 +31,27 @@ func TestIsForcedUsageBillingRequestID(t *testing.T) { require.True(t, isForcedUsageBillingRequestID("grok-video:task-1")) require.True(t, isForcedUsageBillingRequestID("grok_audio:up-1")) require.True(t, isForcedUsageBillingRequestID("grok_realtime:sess-1")) + require.True(t, isForcedUsageBillingRequestID(GeminiCachedContentCreateUsageRequestID("abc"))) require.False(t, isForcedUsageBillingRequestID("resp_abc")) } +func TestGeminiCachedContentUsageRequestIDs(t *testing.T) { + t.Parallel() + publicID, err := NewGeminiCachedContentPublicID() + require.NoError(t, err) + create := GeminiCachedContentCreateUsageRequestID(publicID) + patch := GeminiCachedContentPatchUsageRequestID(publicID) + require.Equal(t, "gcache:create:"+publicID, create) + require.True(t, strings.HasPrefix(patch, "gcache:patch:"+publicID+":")) + require.NotEqual(t, patch, GeminiCachedContentPatchUsageRequestID(publicID), "每次延长各自唯一") + require.LessOrEqual(t, len(create), 64, "usage_logs.request_id 列长 64") + require.LessOrEqual(t, len(patch), 64, "usage_logs.request_id 列长 64") + + ctx := context.WithValue(context.Background(), ctxkey.ClientRequestID, "client-shared-id") + require.Equal(t, create, resolveUsageBillingRequestID(ctx, create), "缓存费用行不被客户端请求 id 覆盖") + require.Equal(t, patch, resolveUsageBillingRequestID(ctx, patch)) +} + func TestStableGrokAudioBillingRequestID(t *testing.T) { t.Parallel() require.Equal(t, "grok_audio:up-1", StableGrokAudioBillingRequestID("up-1")) diff --git a/backend/internal/service/gemini_cache_storage_pricing.go b/backend/internal/service/gemini_cache_storage_pricing.go new file mode 100644 index 0000000000..c08ed2f683 --- /dev/null +++ b/backend/internal/service/gemini_cache_storage_pricing.go @@ -0,0 +1,72 @@ +package service + +import ( + "slices" + "strings" +) + +// Gemini 显式上下文缓存存储的内置官方单价（USD / token / 小时，Standard 档）。 +// 存储价与上下文长度无关（≤200K 与 >200K 同价）；Priority 档的更高存储价不在此列。 +// 来源：https://ai.google.dev/gemini-api/docs/pricing +// 与 https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing +const ( + geminiProCacheStoragePrice = 4.5e-6 // $4.50 / 1M tokens / 小时 + geminiFlashCacheStoragePrice = 1e-6 // $1.00 / 1M tokens / 小时，Flash-Lite 同价 + + // 3.6/3.7/3.8 Flash 的优惠价 $0.50 / 1M tokens / 小时，有效期至 2026-12-31， + // 2027-01-01 起恢复为 geminiFlashCacheStoragePrice。AI Studio 价目页直接标注该优惠价； + // Vertex 价目页存储价标 $1.00，优惠以「按该模型净消费返还 50% 抵扣金」发放， + // 官方未明确存储费是否在返还范围内，两类账号统一按优惠价计。 + geminiFlashPromoCacheStoragePrice = 0.5e-6 +) + +// geminiCacheStorageModelPrices 是已发布型号的存储单价。模型名等于前缀，或以「前缀-」开头 +// （-preview、-customtools、-001 等变体）即命中；更具体的前缀排在前面。 +var geminiCacheStorageModelPrices = []struct { + prefix string + price float64 +}{ + {"gemini-3.8-flash-cyber", geminiFlashCacheStoragePrice}, + {"gemini-3.8-flash", geminiFlashPromoCacheStoragePrice}, + {"gemini-3.7-flash", geminiFlashPromoCacheStoragePrice}, + {"gemini-3.6-flash", geminiFlashPromoCacheStoragePrice}, + {"gemini-3.5-flash-lite", geminiFlashCacheStoragePrice}, + {"gemini-3.5-flash", geminiFlashCacheStoragePrice}, + {"gemini-3.1-flash-lite", geminiFlashCacheStoragePrice}, + {"gemini-3.1-pro", geminiProCacheStoragePrice}, + {"gemini-3-flash", geminiFlashCacheStoragePrice}, + {"gemini-3-pro", geminiProCacheStoragePrice}, + {"gemini-2.5-flash-lite", geminiFlashCacheStoragePrice}, + {"gemini-2.5-flash", geminiFlashCacheStoragePrice}, + {"gemini-2.5-pro", geminiProCacheStoragePrice}, + {"gemini-2.0-flash-lite", geminiFlashCacheStoragePrice}, + {"gemini-2.0-flash", geminiFlashCacheStoragePrice}, +} + +// builtinGeminiCacheStoragePrice 返回内置的显式缓存存储单价：先按已发布型号匹配，未命中时按系列兜底， +// Pro 系列取 geminiProCacheStoragePrice，Flash 与 Flash-Lite 系列取 geminiFlashCacheStoragePrice（标准价， +// 不套用优惠价）。其他 Gemini 模型与非 Gemini 模型返回 0。 +func builtinGeminiCacheStoragePrice(model string) float64 { + name := lastSegment(strings.ToLower(strings.TrimSpace(model))) + if !strings.HasPrefix(name, "gemini-") { + return 0 + } + parts := strings.Split(name, "-") + lite := slices.Contains(parts, "lite") + for _, entry := range geminiCacheStorageModelPrices { + // Flash-Lite 变体只匹配 Flash-Lite 条目：gemini-3.8-flash-lite 不能落到 3.8 Flash 的优惠价。 + if lite != strings.HasSuffix(entry.prefix, "-lite") { + continue + } + if name == entry.prefix || strings.HasPrefix(name, entry.prefix+"-") { + return entry.price + } + } + switch { + case slices.Contains(parts, "pro"): + return geminiProCacheStoragePrice + case slices.Contains(parts, "flash"): + return geminiFlashCacheStoragePrice + } + return 0 +} diff --git a/backend/internal/service/gemini_cache_storage_pricing_test.go b/backend/internal/service/gemini_cache_storage_pricing_test.go new file mode 100644 index 0000000000..846f7b56ec --- /dev/null +++ b/backend/internal/service/gemini_cache_storage_pricing_test.go @@ -0,0 +1,77 @@ +//go:build unit + +package service + +import ( + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/stretchr/testify/require" +) + +func TestBuiltinGeminiCacheStoragePrice(t *testing.T) { + for _, tc := range []struct { + model string + want float64 + }{ + // 已发布型号 + {"gemini-3.8-flash", 0.5e-6}, + {"gemini-3.7-flash", 0.5e-6}, + {"gemini-3.6-flash", 0.5e-6}, + {"gemini-3.8-flash-cyber", 1e-6}, + {"gemini-3.5-flash", 1e-6}, + {"gemini-3.5-flash-lite", 1e-6}, + {"gemini-3.1-flash-lite", 1e-6}, + {"gemini-3.1-pro-preview", 4.5e-6}, + {"gemini-3.1-pro-preview-customtools", 4.5e-6}, + {"gemini-3-flash-preview", 1e-6}, + {"gemini-3-pro-preview", 4.5e-6}, + {"gemini-2.5-pro", 4.5e-6}, + {"gemini-2.5-flash", 1e-6}, + {"gemini-2.5-flash-lite", 1e-6}, + {"gemini-2.0-flash-001", 1e-6}, + {"gemini-2.0-flash-lite-001", 1e-6}, + // 资源名与大小写 + {"models/gemini-3.8-flash", 0.5e-6}, + {"projects/p/locations/global/publishers/google/models/gemini-2.5-pro", 4.5e-6}, + {" Gemini-3.8-Flash ", 0.5e-6}, + // 系列兜底：未发布型号与别名按标准价 + {"gemini-3.9-flash", 1e-6}, + {"gemini-3.9-flash-lite", 1e-6}, + {"gemini-3.8-flash-lite", 1e-6}, + {"gemini-3.9-pro-preview", 4.5e-6}, + {"gemini-flash-latest", 1e-6}, + {"gemini-flash-lite-latest", 1e-6}, + {"gemini-pro-latest", 4.5e-6}, + // 无存储价 + {"gemini-embedding-001", 0}, + {"gemini-unknown-model", 0}, + {"claude-sonnet-5-5", 0}, + {"", 0}, + } { + require.InDelta(t, tc.want, builtinGeminiCacheStoragePrice(tc.model), 1e-15, tc.model) + } +} + +func TestCacheStoragePricePerTokenHourPriority(t *testing.T) { + pricing := &PricingService{} + data, err := pricing.parsePricingData([]byte(`{ + "gemini-3.8-flash":{"litellm_provider":"gemini","input_cost_per_token":0.00000075,"output_cost_per_token":0.00000375, + "cache_storage_cost_per_token_per_hour":0.000002}, + "gemini-3.9-flash":{"litellm_provider":"gemini","input_cost_per_token":0.0000015,"output_cost_per_token":0.0000075, + "cache_storage_cost_per_token_per_hour":0.0000007}, + "gemini-2.5-pro":{"litellm_provider":"gemini","input_cost_per_token":0.00000125,"output_cost_per_token":0.00001}}`)) + require.NoError(t, err) + pricing.pricingData = data + billing := NewBillingService(&config.Config{}, pricing) + + require.InDelta(t, 2e-6, billing.CacheStoragePricePerTokenHour("gemini-3.8-flash"), 1e-15, "价格表配置值优先于内置型号价") + require.InDelta(t, 0.7e-6, billing.CacheStoragePricePerTokenHour("gemini-3.9-flash"), 1e-15, "价格表配置值优先于内置系列兜底价") + require.InDelta(t, 4.5e-6, billing.CacheStoragePricePerTokenHour("gemini-2.5-pro"), 1e-15, "价格表条目缺存储价时取内置型号价") + require.InDelta(t, 0.5e-6, billing.CacheStoragePricePerTokenHour("gemini-3.7-flash"), 1e-15, "价格表无条目时取内置型号价") + require.InDelta(t, 4.5e-6, billing.CacheStoragePricePerTokenHour("gemini-3.9-pro"), 1e-15, "内置型号未命中时取系列兜底价") + require.Zero(t, billing.CacheStoragePricePerTokenHour("customer-alias")) + + noCatalog := NewBillingService(&config.Config{}, nil) + require.InDelta(t, 0.5e-6, noCatalog.CacheStoragePricePerTokenHour("gemini-3.8-flash"), 1e-15, "无价格服务时取内置价") +} diff --git a/backend/internal/service/gemini_cached_content.go b/backend/internal/service/gemini_cached_content.go new file mode 100644 index 0000000000..bc1a3353dc --- /dev/null +++ b/backend/internal/service/gemini_cached_content.go @@ -0,0 +1,1226 @@ +package service + +import ( + "bytes" + "context" + "crypto/rand" + "encoding/base32" + "encoding/json" + "errors" + "fmt" + "io" + "math" + "net/http" + "net/url" + "regexp" + "slices" + "strconv" + "strings" + "sync" + "time" + "unicode/utf8" + + "github.com/Wei-Shaw/sub2api/internal/config" + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + "github.com/Wei-Shaw/sub2api/internal/pkg/geminicli" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" +) + +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" + + // geminiCachedContentMutationUpstreamTimeout 是修改 / 删除上游缓存的单次调用超时。 + geminiCachedContentMutationUpstreamTimeout = 30 * time.Second + // geminiCachedContentMutationLease 是修改意图的租约时长，须覆盖持有者在租约内的全部操作 + // （一次上游修改与一次落库，或恢复流程的一次上游回滚与一次落库），租约过期后持有者不会再改动上游。 + geminiCachedContentMutationLease = 2 * time.Minute + // geminiCachedContentMutationRecoveryInterval / geminiCachedContentMutationRecoveryBatch 约束后台回滚恢复的节奏； + // geminiCachedContentMutationRecoveryTimeout 是恢复单条记录的落库超时，与一次上游回滚合计不超过租约。 + geminiCachedContentMutationRecoveryInterval = time.Minute + geminiCachedContentMutationRecoveryBatch = 50 + geminiCachedContentMutationRecoveryTimeout = time.Minute + // geminiCachedContentMutationAbandonGrace 是判定"上游缓存必已自行过期"时额外容忍的网关与上游时钟偏差。 + geminiCachedContentMutationAbandonGrace = time.Hour + + // 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") + // ErrGeminiCachedContentMutationConflict 表示同一缓存有进行中的修改、到期时间已被其他修改改变，或修改意图已不属于调用方。 + ErrGeminiCachedContentMutationConflict = infraerrors.New(http.StatusConflict, "GEMINI_CACHED_CONTENT_MUTATION_CONFLICT", "another update to this cached content is in progress") + // ErrGeminiCachedContentRequestNotSent 表示上游缓存请求在发出前失败（地址、凭据等），上游状态未被改变。 + ErrGeminiCachedContentRequestNotSent = errors.New("gemini cached content request was not sent") + // errGeminiCachedContentAccountUnsupported 表示持有缓存的账号存在，但当前不支持显式缓存（能力被关闭或认证方式不支持）。 + errGeminiCachedContentAccountUnsupported = errors.New("gemini cached content account does not support explicit caching") + + geminiCachedContentPublicIDPattern = regexp.MustCompile(`^[a-z0-9]{1,64}$`) + geminiCachedContentAIStudioName = regexp.MustCompile(`^cachedContents/[A-Za-z0-9_-]+$`) + 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 + // ChannelUsage 是创建时的渠道计费口径（渠道 ID、请求模型、计费模型来源），延长有效期按同一口径计费。 + ChannelUsage ChannelUsageFields + 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) + // BeginMutation 在记录没有修改意图且 expire_time 仍为 baseline 时登记修改意图（token 与自数据库当前时间起的租约）， + // 返回是否登记成功；记录不存在或已删除时返回 ErrGeminiCachedContentNotFound。 + BeginMutation(ctx context.Context, id int64, baseline time.Time, token string, lease time.Duration) (bool, error) + // TakeOverStaleMutation 把记录上租约已过期的修改意图换成 token 与新租约并返回记录；没有可接管的意图时返回 nil。 + TakeOverStaleMutation(ctx context.Context, id int64, token string, lease time.Duration) (*GeminiCachedContent, error) + // CommitMutation 由 token 的持有者写入新的到期时间并清除修改意图；token 已失效或记录已删除时返回 + // ErrGeminiCachedContentMutationConflict。 + CommitMutation(ctx context.Context, id int64, token string, expireTime time.Time) error + // ReleaseMutation 由 token 的持有者清除修改意图，不改变到期时间。 + ReleaseMutation(ctx context.Context, id int64, token string) error + // ListStaleMutations 返回修改意图租约已过期的未删除记录 id。 + ListStaleMutations(ctx context.Context, limit int) ([]int64, error) + SoftDelete(ctx context.Context, id int64) error +} + +// 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 账号，或 Service Account JSON 认证的 Vertex 账号。 +// API Key 认证的 Vertex 账号不参与显式缓存：Vertex 缓存接口只接受 OAuth access token，不接受 API Key。 +func IsGeminiExplicitCacheAccount(account *Account) bool { + if !IsGeminiExplicitCacheAccountType(account) { + return false + } + if account.Type == AccountTypeServiceAccount { + return !account.IsVertexAPIKey() + } + if strings.TrimSpace(account.GetCredential("api_key")) == "" { + return false + } + return isOfficialGeminiAPIBaseURL(account.GetCredential("base_url")) || account.GeminiExplicitCacheUpstreamEnabled() +} + +// GeminiExplicitCacheUpstreamCredentialKey / GeminiExplicitCacheUpstreamMaxTTLCredentialKey 是 API Key 账号凭据里 +// 声明上游实现显式缓存接口及其有效期上限的键；调度快照投影须保留它们，候选过滤才能直接判定账号是否支持显式缓存。 +const ( + GeminiExplicitCacheUpstreamCredentialKey = "explicit_cache_upstream" + GeminiExplicitCacheUpstreamMaxTTLCredentialKey = "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[GeminiExplicitCacheUpstreamCredentialKey].(bool) + return enabled +} + +// GeminiExplicitCacheUpstreamMaxTTL 返回上游网关允许的缓存最长有效期；0 表示未配置。 +func (a *Account) GeminiExplicitCacheUpstreamMaxTTL() time.Duration { + if !a.GeminiExplicitCacheUpstreamEnabled() { + return 0 + } + var seconds float64 + switch v := a.Credentials[GeminiExplicitCacheUpstreamMaxTTLCredentialKey].(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 +} + +// GeminiCachedContentUsageRequestIDPrefix 是显式缓存创建 / 延长费用在 usage_logs.request_id 中的固定前缀： +// 这类费用行是独立的扣费事件，不随客户端或网关请求 id 变化，按缓存 ID 即可检索一个缓存的全部费用行。 +// 两种格式都不超过 request_id 列的 64 字符。 +const GeminiCachedContentUsageRequestIDPrefix = "gcache:" + +// GeminiCachedContentCreateUsageRequestID 返回创建费用的 request_id（"gcache:create:{id}"）；一个缓存只创建一次。 +func GeminiCachedContentCreateUsageRequestID(publicID string) string { + return GeminiCachedContentUsageRequestIDPrefix + "create:" + publicID +} + +// GeminiCachedContentPatchUsageRequestID 返回一次延长费用的 request_id（"gcache:patch:{id}:{8 位随机}"）；每次延长各自唯一。 +func GeminiCachedContentPatchUsageRequestID(publicID string) string { + var b [5]byte + if _, err := rand.Read(b[:]); err != nil { + return GeminiCachedContentUsageRequestIDPrefix + "patch:" + publicID + ":" + strconv.FormatInt(time.Now().UnixNano()%(1<<40), 32) + } + return GeminiCachedContentUsageRequestIDPrefix + "patch:" + publicID + ":" + strings.ToLower(geminiCachedContentPublicIDEncoding.EncodeToString(b[:])) +} + +// ParseGeminiCachedContentPublicID 从 "cachedContents/{id}" 或裸 id 中解析缓存 ID。 +func ParseGeminiCachedContentPublicID(name string) (string, bool) { + id := strings.TrimPrefix(strings.TrimSpace(name), GeminiCachedContentNamePrefix) + 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 + // UpstreamGateway 表示响应来自声明实现显式缓存的上游网关账号：限流 / 不可用反映的是持有缓存的上游账号状态。 + UpstreamGateway bool +} + +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 + + startOnce sync.Once + stopOnce sync.Once + stopCh chan struct{} +} + +func NewGeminiCachedContentService(repo GeminiCachedContentRepository, compat *GeminiMessagesCompatService, cfg *config.Config) *GeminiCachedContentService { + return &GeminiCachedContentService{repo: repo, compat: compat, cfg: cfg, now: time.Now, stopCh: make(chan struct{})} +} + +// Start 启动后台回滚恢复：定期处理租约已过期仍未清除的修改意图。 +func (s *GeminiCachedContentService) Start() { + if s == nil || s.repo == nil { + return + } + s.startOnce.Do(func() { go s.runMutationRecovery() }) +} + +// Stop 停止后台回滚恢复。 +func (s *GeminiCachedContentService) Stop() { + if s == nil { + return + } + s.stopOnce.Do(func() { close(s.stopCh) }) +} + +func (s *GeminiCachedContentService) runMutationRecovery() { + ticker := time.NewTicker(geminiCachedContentMutationRecoveryInterval) + defer ticker.Stop() + for { + select { + case <-s.stopCh: + return + case <-ticker.C: + recovered, err := s.RecoverStaleMutations(context.Background()) + if err != nil { + logger.LegacyPrintf("service.gemini_cached_content", "Warning: gemini cached content mutation recovery failed (recovered=%d): %v", recovered, err) + } else if recovered > 0 { + logger.LegacyPrintf("service.gemini_cached_content", "gemini cached content mutation recovery rolled back %d stale mutations", recovered) + } + } + } +} + +// Now 返回服务时钟。 +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) +} + +// GeminiCachedContentMutation 是一次已登记修改意图、尚未提交的有效期修改。持有期间同一缓存的其他修改被拒绝； +// 未提交的修改（上游结果未知、落库失败、进程退出）在租约过期后由恢复流程把上游有效期回滚到 Baseline，不计费。 +type GeminiCachedContentMutation struct { + record *GeminiCachedContent + baseline time.Time + token string +} + +// Baseline 返回修改前的到期时间：提交前它始终是库中的到期时间。 +func (m *GeminiCachedContentMutation) Baseline() time.Time { + return m.baseline +} + +// BeginMutation 登记修改 record 有效期的意图并取得租约。同一缓存有进行中的修改，或库中到期时间已不是 +// record.ExpireTime 时返回 ErrGeminiCachedContentMutationConflict；租约已过期的遗留意图先回滚恢复再登记。 +func (s *GeminiCachedContentService) BeginMutation(ctx context.Context, record *GeminiCachedContent) (*GeminiCachedContentMutation, error) { + for attempt := 0; ; attempt++ { + token, err := NewGeminiCachedContentPublicID() + if err != nil { + return nil, err + } + acquired, err := s.repo.BeginMutation(ctx, record.ID, record.ExpireTime, token, geminiCachedContentMutationLease) + if err != nil { + return nil, err + } + if acquired { + return &GeminiCachedContentMutation{record: record, baseline: record.ExpireTime, token: token}, nil + } + if attempt > 0 { + return nil, ErrGeminiCachedContentMutationConflict + } + recovered, err := s.recoverStaleMutation(ctx, record.ID) + if err != nil { + return nil, errors.Join(ErrGeminiCachedContentMutationConflict, err) + } + if !recovered { + return nil, ErrGeminiCachedContentMutationConflict + } + } +} + +// CommitMutation 写入新的到期时间并清除修改意图。修改意图已不属于 m（缓存已删除或租约已被接管）时返回 +// ErrGeminiCachedContentMutationConflict；其他失败时意图保留，由恢复流程回滚上游。 +func (s *GeminiCachedContentService) CommitMutation(ctx context.Context, m *GeminiCachedContentMutation, expireTime time.Time) error { + if err := s.repo.CommitMutation(ctx, m.record.ID, m.token, expireTime); err != nil { + return err + } + m.record.ExpireTime = expireTime + m.record.UpdatedAt = s.now() + return nil +} + +// ReleaseMutation 放弃修改并清除修改意图，用于上游明确未改变有效期的情形。 +func (s *GeminiCachedContentService) ReleaseMutation(ctx context.Context, m *GeminiCachedContentMutation) error { + return s.repo.ReleaseMutation(ctx, m.record.ID, m.token) +} + +// RecoverStaleMutations 回滚恢复一批租约已过期的修改意图，返回完成恢复的数量。 +func (s *GeminiCachedContentService) RecoverStaleMutations(ctx context.Context) (int, error) { + listCtx, cancel := context.WithTimeout(ctx, geminiCachedContentMutationRecoveryTimeout) + ids, err := s.repo.ListStaleMutations(listCtx, geminiCachedContentMutationRecoveryBatch) + cancel() + if err != nil { + return 0, err + } + recovered := 0 + var errs []error + for _, id := range ids { + select { + case <-s.stopCh: + return recovered, errors.Join(errs...) + default: + } + itemCtx, cancel := context.WithTimeout(ctx, geminiCachedContentMutationRecoveryTimeout) + ok, err := s.recoverStaleMutation(itemCtx, id) + cancel() + if err != nil { + errs = append(errs, fmt.Errorf("cached content %d: %w", id, err)) + continue + } + if ok { + recovered++ + } + } + return recovered, errors.Join(errs...) +} + +// recoverStaleMutation 接管 id 上租约已过期的修改意图，把上游有效期回滚到库中的到期时间后清除意图； +// 库中到期时间已到时删除上游缓存与本地记录。没有可接管的意图时返回 false；回滚失败时意图保留， +// 本次接管取得的新租约（geminiCachedContentMutationLease）过期后重试，期间同一缓存的修改返回 409。 +func (s *GeminiCachedContentService) recoverStaleMutation(ctx context.Context, id int64) (bool, error) { + token, err := NewGeminiCachedContentPublicID() + if err != nil { + return false, err + } + record, err := s.repo.TakeOverStaleMutation(ctx, id, token, geminiCachedContentMutationLease) + if err != nil || record == nil { + return false, err + } + forget, err := s.rollbackUpstream(ctx, record) + if err != nil { + return false, err + } + if forget { + return true, s.repo.SoftDelete(ctx, record.ID) + } + return true, s.repo.ReleaseMutation(ctx, record.ID, token) +} + +// rollbackUpstream 把上游有效期改回 record.ExpireTime，剩余不足 1 秒时删除上游缓存； +// 返回 true 表示上游缓存已不存在、或已没有凭据可触达（持有账号已删除），本地记录应移除。 +// 持有账号存在但当前不支持显式缓存时不改动上游，修改意图保留到账号恢复后重试；账号可用时总是先向上游回滚。 +// 两种情况下只要 mutationUpstreamMustHaveExpired 成立（上游缓存必已自行过期），本地记录直接作废，不再重试。 +func (s *GeminiCachedContentService) rollbackUpstream(ctx context.Context, record *GeminiCachedContent) (bool, error) { + account, err := s.boundAccount(ctx, record) + if errors.Is(err, errGeminiCachedContentAccountUnsupported) { + if s.mutationUpstreamMustHaveExpired(record) { + logger.LegacyPrintf("service.gemini_cached_content", "Warning: gemini cached content %d (upstream %s) has a pending mutation but account %d stopped supporting explicit caching and the upstream cache has expired since; dropping the local record", record.ID, record.UpstreamName, record.AccountID) + return true, nil + } + return false, fmt.Errorf("%w (account %d)", err, record.AccountID) + } + if err != nil { + return false, err + } + if account == nil { + logger.LegacyPrintf("service.gemini_cached_content", "Warning: gemini cached content %d (upstream %s) has a pending mutation but account %d no longer exists; dropping the local record without touching upstream", record.ID, record.UpstreamName, record.AccountID) + return true, nil + } + remaining := record.ExpireTime.Sub(s.now()) + if remaining < time.Second { + err = s.DeleteUpstream(ctx, account, record) + } else { + _, err = s.PatchUpstream(ctx, account, record, remaining) + } + var upErr *GeminiCachedContentUpstreamError + if errors.As(err, &upErr) && upErr.NotFound() { + return true, nil + } + if err != nil { + if s.mutationUpstreamMustHaveExpired(record) { + logger.LegacyPrintf("service.gemini_cached_content", "Warning: gemini cached content %d (upstream %s) rollback keeps failing on account %d but the upstream cache has expired since; dropping the local record: %v", record.ID, record.UpstreamName, record.AccountID, err) + return true, nil + } + return false, err + } + return remaining < time.Second, nil +} + +// mutationUpstreamMustHaveExpired 判断 record 上未完成修改所涉及的上游缓存是否必已自行过期： +// 修改发生时记录尚未到期（Resolve 拒绝已过期记录），且单次延长不超过 MaxTTL， +// 所以上游有效期最晚为 record.ExpireTime + MaxTTL；再过 geminiCachedContentMutationAbandonGrace 后判定为已过期。 +// MaxTTL 取判定时刻的配置并以默认上限兜底：配置被调低后判定可能早于上游真实过期，此时只影响运营方的上游存储费， +// 不涉及用户计费。 +func (s *GeminiCachedContentService) mutationUpstreamMustHaveExpired(record *GeminiCachedContent) bool { + maxTTL := s.MaxTTL() + if maxTTL < geminiCachedContentDefaultMaxTTL { + maxTTL = geminiCachedContentDefaultMaxTTL + } + return s.now().After(record.ExpireTime.Add(maxTTL + geminiCachedContentMutationAbandonGrace)) +} + +func (s *GeminiCachedContentService) Forget(ctx context.Context, record *GeminiCachedContent) error { + return s.repo.SoftDelete(ctx, record.ID) +} + +// BoundAccount 返回持有缓存的账号；账号已删除或不再支持显式缓存时返回 nil。 +func (s *GeminiCachedContentService) BoundAccount(ctx context.Context, record *GeminiCachedContent) (*Account, error) { + account, err := s.boundAccount(ctx, record) + if errors.Is(err, errGeminiCachedContentAccountUnsupported) { + return nil, nil + } + return account, err +} + +// boundAccount 返回持有缓存的账号：账号已删除返回 nil；账号存在但当前不支持显式缓存返回 +// errGeminiCachedContentAccountUnsupported，调用方据此区分"永久不可用"与"暂时不可用"。 +func (s *GeminiCachedContentService) boundAccount(ctx context.Context, record *GeminiCachedContent) (*Account, error) { + if s.compat == nil || s.compat.accountRepo == nil { + return nil, errors.New("account repository not configured") + } + 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 { + return nil, nil + } + if !IsGeminiExplicitCacheAccount(account) { + return nil, errGeminiCachedContentAccountUnsupported + } + 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, geminiCachedContentUpstreamTimeout) + 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) { + resourceURL, err := s.upstreamResourceURL(account, record) + if err != nil { + return nil, err + } + seconds := int64(math.Ceil(ttl.Seconds())) + body, _ := json.Marshal(map[string]any{"ttl": fmt.Sprintf("%ds", seconds)}) + respBody, err := s.doUpstream(ctx, account, http.MethodPatch, resourceURL+"?updateMask=ttl", body, geminiCachedContentMutationUpstreamTimeout) + 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 { + resourceURL, err := s.upstreamResourceURL(account, record) + if err != nil { + return err + } + _, err = s.doUpstream(ctx, account, http.MethodDelete, resourceURL, nil, geminiCachedContentMutationUpstreamTimeout) + return err +} + +// upstreamResourceURL 返回 record 对应上游缓存资源的地址；失败时请求不会发出。 +func (s *GeminiCachedContentService) upstreamResourceURL(account *Account, record *GeminiCachedContent) (string, error) { + base, err := s.apiBase(account) + if err != nil { + return "", geminiCachedContentNotSent(err) + } + resourceURL, err := geminiCachedContentResourceURL(account, base, record.UpstreamName) + if err != nil { + return "", geminiCachedContentNotSent(err) + } + return resourceURL, nil +} + +func geminiCachedContentNotSent(err error) error { + return fmt.Errorf("%w: %w", ErrGeminiCachedContentRequestNotSent, err) +} + +// doUpstream 调用上游缓存接口并返回 2xx 响应体，单次调用不超过 timeout。上游调用与客户端连接脱钩：客户端断开不会中途取消 +// 已发出的创建 / 修改，调用方据结果完成落库、计费或清理。请求发出前的失败包装为 ErrGeminiCachedContentRequestNotSent。 +func (s *GeminiCachedContentService) doUpstream(ctx context.Context, account *Account, method, target string, body []byte, timeout time.Duration) ([]byte, error) { + if s.compat == nil || s.compat.httpUpstream == nil { + return nil, geminiCachedContentNotSent(errors.New("gemini upstream client not configured")) + } + ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), timeout) + defer cancel() + var reader io.Reader + if body != nil { + reader = bytes.NewReader(body) + } + req, err := http.NewRequestWithContext(ctx, method, target, reader) + if err != nil { + return nil, geminiCachedContentNotSent(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, geminiCachedContentNotSent(errors.New("gemini api_key not configured")) + } + req.Header.Set("x-goog-api-key", apiKey) + case AccountTypeServiceAccount: + if s.compat.tokenProvider == nil { + return nil, geminiCachedContentNotSent(errors.New("gemini token provider not configured")) + } + accessToken, err := s.compat.tokenProvider.GetAccessToken(ctx, account) + if err != nil { + return nil, geminiCachedContentNotSent(err) + } + req.Header.Set("Authorization", "Bearer "+accessToken) + default: + return nil, geminiCachedContentNotSent(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, UpstreamGateway: account.GeminiExplicitCacheUpstreamEnabled()} + } + 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 时只保留该账号；否则排除不支持显式缓存的账号。判定所需的凭据键 +// （api_key、base_url、auth_mode、上游显式缓存声明）都在调度快照投影内，候选过滤一次算全； +// 选中后仍对完整账号复核。 +func (s *GatewayService) GeminiExplicitCacheExclusions(ctx context.Context, groupID *int64, boundAccountID int64) (map[int64]struct{}, error) { + accounts, _, err := s.listSchedulableAccounts(ctx, groupID, PlatformGemini, false) + if err != nil { + return nil, err + } + excluded := make(map[int64]struct{}, len(accounts)) + for i := range accounts { + account := &accounts[i] + keep := IsGeminiExplicitCacheAccount(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..f93b9424a0 --- /dev/null +++ b/backend/internal/service/gemini_cached_content_hardening_test.go @@ -0,0 +1,480 @@ +//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.True(t, failoverErr.BoundUpstreamPassthrough, "标记为上游网关账号的响应，handler 原样回写") + 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 fixedChargeCacheStub struct { + billingCacheWorkerStub + balance float64 + sub *SubscriptionCacheData + rate *APIKeyRateLimitCacheData + // reserved 按维度（user:{id} / scope 名）记录在途预留。 + reserved map[string]map[string]float64 +} + +func (s *fixedChargeCacheStub) GetUserBalance(context.Context, int64) (float64, error) { + return s.balance, nil +} + +func (s *fixedChargeCacheStub) GetSubscriptionCache(context.Context, int64, int64) (*SubscriptionCacheData, error) { + if s.sub == nil { + return nil, errors.New("cache miss") + } + return s.sub, nil +} + +func (s *fixedChargeCacheStub) GetAPIKeyRateLimit(context.Context, int64) (*APIKeyRateLimitCacheData, error) { + if s.rate == nil { + return nil, errors.New("cache miss") + } + return s.rate, nil +} + +func (s *fixedChargeCacheStub) reserve(key, requestID string, amount float64) float64 { + if s.reserved == nil { + s.reserved = map[string]map[string]float64{} + } + if s.reserved[key] == nil { + s.reserved[key] = map[string]float64{} + } + sum := 0.0 + for _, v := range s.reserved[key] { + sum += v + } + s.reserved[key][requestID] = amount + return sum +} + +func (s *fixedChargeCacheStub) ReserveInflightBalance(_ context.Context, userID int64, requestID string, amount, balance float64, _ time.Duration) (bool, float64, error) { + key := fmt.Sprintf("user:%d", userID) + others := s.reserve(key, requestID, amount) + if len(s.reserved[key]) > 1 && balance-others < amount { + delete(s.reserved[key], requestID) + return false, others, nil + } + return true, others, nil +} + +func (s *fixedChargeCacheStub) ReleaseInflightBalance(_ context.Context, userID int64, requestID string) error { + delete(s.reserved[fmt.Sprintf("user:%d", userID)], requestID) + return nil +} + +func (s *fixedChargeCacheStub) ReserveInflightScope(_ context.Context, scope, requestID string, amount float64, _ time.Duration) (float64, error) { + return s.reserve(scope, requestID, amount), nil +} + +func (s *fixedChargeCacheStub) ReleaseInflightScope(_ context.Context, scope, requestID string) error { + delete(s.reserved[scope], requestID) + return nil +} + +type fixedChargeAPIKeyRepoStub struct { + APIKeyRepository + data *APIKeyRateLimitData +} + +func (r *fixedChargeAPIKeyRepoStub) GetRateLimitData(context.Context, int64) (*APIKeyRateLimitData, error) { + return r.data, nil +} + +type fixedChargeSubscriptionRepoStub struct { + UserSubscriptionRepository + sub *UserSubscription +} + +func (r *fixedChargeSubscriptionRepoStub) GetActiveByUserIDAndGroupID(context.Context, int64, int64) (*UserSubscription, error) { + return r.sub, nil +} + +func TestReserveFixedCharge(t *testing.T) { + ctx := context.Background() + user := &User{ID: 1} + newSvc := func(cache *fixedChargeCacheStub) *BillingCacheService { + svc := NewBillingCacheService(cache, nil, nil, nil, nil, nil, &config.Config{}, nil) + t.Cleanup(svc.Stop) + return svc + } + reservingSvc := func(cache *fixedChargeCacheStub, subRepo UserSubscriptionRepository, apiKeyRepo APIKeyRepository) *BillingCacheService { + cfg := &config.Config{} + cfg.Billing.InflightReservation.Enabled = true + svc := NewBillingCacheService(cache, nil, subRepo, apiKeyRepo, nil, nil, cfg, nil) + t.Cleanup(svc.Stop) + return svc + } + reserve := func(svc *BillingCacheService, apiKey *APIKey, group *Group, sub *UserSubscription, amount float64) error { + res, err := svc.ReserveFixedCharge(ctx, user, apiKey, group, sub, amount) + res.Release() + return err + } + + t.Run("balance must cover the amount", func(t *testing.T) { + svc := newSvc(&fixedChargeCacheStub{balance: 10}) + require.NoError(t, reserve(svc, nil, nil, nil, 10)) + require.ErrorIs(t, reserve(svc, nil, nil, nil, 10.01), ErrInsufficientBalance) + require.NoError(t, reserve(svc, nil, nil, nil, 0)) + }) + + t.Run("subscription windows must cover the amount", func(t *testing.T) { + daily := 5.0 + group := &Group{ID: 2, SubscriptionType: SubscriptionTypeSubscription, DailyLimitUSD: &daily} + svc := newSvc(&fixedChargeCacheStub{sub: &SubscriptionCacheData{ + Status: SubscriptionStatusActive, ExpiresAt: time.Now().Add(time.Hour), DailyUsage: 3, + }}) + sub := &UserSubscription{ID: 1} + require.NoError(t, reserve(svc, nil, group, sub, 2)) + require.ErrorIs(t, reserve(svc, nil, group, sub, 2.5), ErrDailyLimitExceeded) + }) + + t.Run("api key quota and rate windows include the amount", func(t *testing.T) { + svc := newSvc(&fixedChargeCacheStub{balance: 100, rate: &APIKeyRateLimitCacheData{ + Usage1d: 8, Window1d: time.Now().Unix(), + }}) + quotaKey := &APIKey{ID: 5, Quota: 10, QuotaUsed: 9} + require.NoError(t, reserve(svc, quotaKey, nil, nil, 1)) + require.ErrorIs(t, reserve(svc, quotaKey, nil, nil, 1.5), ErrAPIKeyQuotaExhausted) + + windowKey := &APIKey{ID: 6, RateLimit1d: 10} + require.NoError(t, reserve(svc, windowKey, nil, nil, 2)) + require.ErrorIs(t, reserve(svc, windowKey, nil, nil, 2.5), ErrAPIKeyRateLimit1dExceeded) + }) + + t.Run("usage is read from the database instead of request snapshots", func(t *testing.T) { + now := time.Now() + apiKeyRepo := &fixedChargeAPIKeyRepoStub{data: &APIKeyRateLimitData{QuotaUsed: 9.5, Usage1d: 8, Window1dStart: &now}} + svc := reservingSvc(&fixedChargeCacheStub{balance: 100}, nil, apiKeyRepo) + require.ErrorIs(t, reserve(svc, &APIKey{ID: 5, Quota: 10}, nil, nil, 1), ErrAPIKeyQuotaExhausted, "请求上下文中的 QuotaUsed 可能滞后") + require.ErrorIs(t, reserve(svc, &APIKey{ID: 6, RateLimit1d: 10}, nil, nil, 2.5), ErrAPIKeyRateLimit1dExceeded) + + daily := 5.0 + group := &Group{ID: 2, SubscriptionType: SubscriptionTypeSubscription, DailyLimitUSD: &daily} + subRepo := &fixedChargeSubscriptionRepoStub{sub: &UserSubscription{Status: SubscriptionStatusActive, DailyUsageUSD: 4}} + svc = reservingSvc(&fixedChargeCacheStub{sub: &SubscriptionCacheData{Status: SubscriptionStatusActive}}, subRepo, nil) + require.ErrorIs(t, reserve(svc, nil, group, &UserSubscription{ID: 1}, 2), ErrDailyLimitExceeded, "订阅缓存用量可能滞后") + }) + + t.Run("unsettled charges hold subscription and api key headroom", func(t *testing.T) { + daily := 110.0 + group := &Group{ID: 2, SubscriptionType: SubscriptionTypeSubscription, DailyLimitUSD: &daily} + sub := &UserSubscription{ID: 1} + cache := &fixedChargeCacheStub{balance: 1000, sub: &SubscriptionCacheData{Status: SubscriptionStatusActive}} + svc := reservingSvc(cache, nil, nil) + first, err := svc.ReserveFixedCharge(ctx, user, nil, group, sub, 103.5) + require.NoError(t, err) + _, err = svc.ReserveFixedCharge(ctx, user, nil, group, sub, 103.5) + require.ErrorIs(t, err, ErrDailyLimitExceeded) + first.Release() + second, err := svc.ReserveFixedCharge(ctx, user, nil, group, sub, 103.5) + require.NoError(t, err) + second.Release() + + apiKey := &APIKey{ID: 5, Quota: 110} + first, err = svc.ReserveFixedCharge(ctx, user, apiKey, nil, nil, 103.5) + require.NoError(t, err) + _, err = svc.ReserveFixedCharge(ctx, user, apiKey, nil, nil, 103.5) + require.ErrorIs(t, err, ErrAPIKeyQuotaExhausted) + first.Release() + for key, reservations := range cache.reserved { + require.Empty(t, reservations, "拒绝与释放都不留下预留：%s", key) + } + }) + + t.Run("simple mode skips the check", func(t *testing.T) { + cfg := &config.Config{RunMode: config.RunModeSimple} + svc := NewBillingCacheService(&fixedChargeCacheStub{}, nil, nil, nil, nil, nil, cfg, nil) + t.Cleanup(svc.Stop) + require.NoError(t, reserve(svc, 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 new file mode 100644 index 0000000000..fc4daa1181 --- /dev/null +++ b/backend/internal/service/gemini_cached_content_not_found_test.go @@ -0,0 +1,186 @@ +//go:build unit + +package service + +import ( + "context" + "errors" + "io" + "net/http" + "strings" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +// 上游对不存在、已过期、已删除或不属于当前项目的显式缓存一律返回这个 403。 +const geminiCachedContentNotFoundBody = `{"error":{"code":403,"message":"CachedContent not found (or permission denied)","status":"PERMISSION_DENIED"}}` + +func newGeminiCachedContentTestService(status int, body string) (*GeminiMessagesCompatService, *geminiCompatHTTPUpstreamStub, *errorPolicyRepoStub) { + repo := &errorPolicyRepoStub{} + httpStub := &geminiCompatHTTPUpstreamStub{ + response: &http.Response{ + StatusCode: status, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(body)), + }, + } + svc := &GeminiMessagesCompatService{ + httpUpstream: httpStub, + cfg: &config.Config{}, + rateLimitService: NewRateLimitService(repo, nil, &config.Config{}, nil, nil), + } + return svc, httpStub, repo +} + +func geminiDefaultPolicyAPIKeyAccount() *Account { + return &Account{ + ID: 710, + Platform: PlatformGemini, + Type: AccountTypeAPIKey, + Credentials: map[string]any{ + "api_key": "test-key", + }, + } +} + +const geminiCachedContentRequestBody = `{"cachedContent":"cachedContents/doesnotexist","contents":[{"role":"user","parts":[{"text":"hi"}]}]}` + +func TestGeminiForwardNative_CachedContentNotFoundPassthroughWithoutAccountPenalty(t *testing.T) { + gin.SetMode(gin.TestMode) + for _, tc := range []struct { + name string + action string + stream bool + }{ + {name: "generateContent", action: "generateContent"}, + {name: "streamGenerateContent", action: "streamGenerateContent", stream: true}, + {name: "countTokens", action: "countTokens"}, + } { + t.Run(tc.name, func(t *testing.T) { + svc, httpStub, repo := newGeminiCachedContentTestService(http.StatusForbidden, geminiCachedContentNotFoundBody) + c, rec := newGeminiNativeTestContext(t) + + result, err := svc.ForwardNative(context.Background(), c, geminiDefaultPolicyAPIKeyAccount(), + "gemini-3.8-flash", tc.action, tc.stream, []byte(geminiCachedContentRequestBody)) + + require.Nil(t, result) + require.Error(t, err) + var failoverErr *UpstreamFailoverError + require.False(t, errors.As(err, &failoverErr), "缓存不存在是请求问题，不应换号") + require.Equal(t, 0, repo.setErrCalls, "不应把账号标记为错误") + require.Equal(t, 0, repo.tempCalls, "不应临时停调度") + require.Equal(t, 1, httpStub.calls, "不应重试") + require.Equal(t, http.StatusForbidden, rec.Code) + require.Equal(t, GeminiCachedContentNotFoundResponse, rec.Body.String(), "应返回标准缓存不存在响应") + require.True(t, IsOpsRequestScopedError(c), "缓存不存在按客户端请求错误统计") + }) + } +} + +func TestGeminiForwardNative_CachedContentNotFoundBypassesTempUnschedulableRules(t *testing.T) { + gin.SetMode(gin.TestMode) + svc, _, repo := newGeminiCachedContentTestService(http.StatusForbidden, geminiCachedContentNotFoundBody) + c, rec := newGeminiNativeTestContext(t) + account := geminiDefaultPolicyAPIKeyAccount() + account.Credentials["temp_unschedulable_enabled"] = true + account.Credentials["temp_unschedulable_rules"] = []any{ + map[string]any{ + "error_code": float64(403), + "keywords": []any{"permission denied"}, + "duration_minutes": float64(10), + "description": "403 permission denied", + }, + } + + _, err := svc.ForwardNative(context.Background(), c, account, + "gemini-3.8-flash", "generateContent", false, []byte(geminiCachedContentRequestBody)) + + require.Error(t, err) + var failoverErr *UpstreamFailoverError + require.False(t, errors.As(err, &failoverErr)) + require.Equal(t, 0, repo.tempCalls, "临时不可调度规则不应命中") + require.Equal(t, 0, repo.setErrCalls) + require.Equal(t, http.StatusForbidden, rec.Code) +} + +func TestGeminiForwardNative_OtherForbiddenStillPenalizesAccount(t *testing.T) { + gin.SetMode(gin.TestMode) + body := `{"error":{"code":403,"message":"Your project has been denied access. Please contact support.","status":"PERMISSION_DENIED"}}` + svc, _, repo := newGeminiCachedContentTestService(http.StatusForbidden, body) + c, _ := newGeminiNativeTestContext(t) + + _, err := svc.ForwardNative(context.Background(), c, geminiDefaultPolicyAPIKeyAccount(), + "gemini-3.8-flash", "generateContent", false, []byte(geminiCachedContentRequestBody)) + + var failoverErr *UpstreamFailoverError + require.True(t, errors.As(err, &failoverErr), "其他 403 仍按账号问题换号") + require.Equal(t, 1, repo.setErrCalls, "其他 403 仍标记账号错误") +} + +// 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") + require.True(t, IsOpsRequestScopedError(c), "缓存不存在按客户端请求错误统计") + errType, _ := OpsLocalErrorType(c) + require.Equal(t, "permission_error", errType) + }) + } +} 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..e3afbdd5e6 --- /dev/null +++ b/backend/internal/service/gemini_cached_content_test.go @@ -0,0 +1,522 @@ +//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) BeginMutation(context.Context, int64, time.Time, string, time.Duration) (bool, error) { + return true, nil +} + +func (r *geminiCachedContentRepoStub) TakeOverStaleMutation(context.Context, int64, string, time.Duration) (*GeminiCachedContent, error) { + return nil, nil +} + +func (r *geminiCachedContentRepoStub) CommitMutation(_ context.Context, id int64, _ string, expireTime time.Time) error { + for _, record := range r.records { + if record.ID == id { + record.ExpireTime = expireTime + return nil + } + } + return ErrGeminiCachedContentMutationConflict +} + +func (r *geminiCachedContentRepoStub) ReleaseMutation(context.Context, int64, string) error { + return nil +} + +func (r *geminiCachedContentRepoStub) ListStaleMutations(context.Context, int) ([]int64, error) { + return nil, nil +} + +func (r *geminiCachedContentRepoStub) SoftDelete(_ context.Context, id int64) error { + 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}, + {"vertex api key", &Account{Platform: PlatformGemini, Type: AccountTypeServiceAccount, Credentials: map[string]any{"auth_mode": "apikey", "api_key": "k", "project_id": "my-project"}}, false}, + {"oauth", &Account{Platform: PlatformGemini, Type: AccountTypeOAuth}, false}, + {"antigravity api key", &Account{Platform: PlatformAntigravity, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "k"}}, false}, + } + 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: GeminiCachedContentCreateUsageRequestID(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_chat_completions_compat_service.go b/backend/internal/service/gemini_chat_completions_compat_service.go index a98c94cf93..1ad2b487e8 100644 --- a/backend/internal/service/gemini_chat_completions_compat_service.go +++ b/backend/internal/service/gemini_chat_completions_compat_service.go @@ -399,19 +399,11 @@ func (s *GeminiMessagesCompatService) buildGeminiChatCompletionsUpstreamRequestF case AccountTypeServiceAccount: return func(ctx context.Context) (*http.Request, string, error) { - if s.tokenProvider == nil { - return nil, "", errors.New("gemini token provider not configured") - } - accessToken, err := s.tokenProvider.GetAccessToken(ctx, account) - if err != nil { - return nil, "", err - } - action := "generateContent" if clientStream { action = "streamGenerateContent" } - fullURL, err := buildVertexGeminiURL(account.VertexProjectID(), account.VertexLocation(mappedModel), mappedModel, action, clientStream) + fullURL, err := buildVertexGeminiAccountURL(account, mappedModel, action, clientStream) if err != nil { return nil, "", err } @@ -422,7 +414,9 @@ func (s *GeminiMessagesCompatService) buildGeminiChatCompletionsUpstreamRequestF return nil, "", err } upstreamReq.Header.Set("Content-Type", "application/json") - upstreamReq.Header.Set("Authorization", "Bearer "+accessToken) + if err := setVertexGeminiAuth(ctx, upstreamReq, account, s.tokenProvider); err != nil { + return nil, "", err + } return upstreamReq, "x-request-id", nil }, "x-request-id" diff --git a/backend/internal/service/gemini_error_policy_test.go b/backend/internal/service/gemini_error_policy_test.go index c868cf6417..2f119ff5e4 100644 --- a/backend/internal/service/gemini_error_policy_test.go +++ b/backend/internal/service/gemini_error_policy_test.go @@ -677,6 +677,72 @@ func TestHandleGeminiUpstreamError_APIKeyFallbackStillPSTMidnight(t *testing.T) require.WithinDuration(t, expected, repo.lastRateLimitReset, 5*time.Second) } +// 第三方上游 API Key 的 429 不套用 AI Studio 每日配额语义：只认显式重试时长， +// 否则走全局 429 兜底冷却（未配置时默认 5 秒）。 +func TestHandleGeminiUpstreamError_ThirdPartyAPIKey429(t *testing.T) { + const relayBaseURL = "https://relay.example.com" + tests := []struct { + name string + baseURL string + body string + expectWait time.Duration + pstNight bool + }{ + { + name: "no_reset_time_uses_global_fallback", + baseURL: relayBaseURL, + body: `{"error":{"code":429,"message":"You have exhausted your capacity on this model. Your quota will reset after 6h53m10s."}}`, + expectWait: defaultRateLimit429CooldownSeconds * time.Second, + }, + { + name: "daily_quota_message_not_treated_as_pst_midnight", + baseURL: relayBaseURL, + body: `{"error":{"code":429,"message":"Quota exceeded for metric: generate_content_free_tier_requests, limit: 50 per day"}}`, + expectWait: defaultRateLimit429CooldownSeconds * time.Second, + }, + { + name: "explicit_quota_reset_delay_respected", + baseURL: relayBaseURL, + body: `{"error":{"code":429,"message":"rate limit","details":[{"metadata":{"quotaResetDelay":"42s"}}]}}`, + expectWait: 42 * time.Second, + }, + { + name: "official_base_url_keeps_pst_midnight", + baseURL: "https://generativelanguage.googleapis.com/", + body: `{"error":{"code":429,"message":"rate limit"}}`, + pstNight: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + repo := &rateLimit429AccountRepoStub{} + svc := &GeminiMessagesCompatService{ + accountRepo: repo, + rateLimitService: NewRateLimitService(repo, nil, &config.Config{}, nil, nil), + } + account := &Account{ + ID: 700, + Platform: PlatformGemini, + Type: AccountTypeAPIKey, + Credentials: map[string]any{"base_url": tt.baseURL}, + } + + before := time.Now() + svc.handleGeminiUpstreamError(context.Background(), account, http.StatusTooManyRequests, http.Header{}, []byte(tt.body)) + + require.Equal(t, 1, repo.rateLimitCalls) + require.Equal(t, int64(700), repo.lastRateLimitID) + if tt.pstNight { + expected := time.Unix(*nextGeminiDailyResetUnix(), 0) + require.WithinDuration(t, expected, repo.lastRateLimitReset, 5*time.Second) + return + } + require.WithinDuration(t, before.Add(tt.expectWait), repo.lastRateLimitReset, 2*time.Second) + }) + } +} + type geminiErrorPolicyRepo struct { mockAccountRepoForGemini setErrorCalls int diff --git a/backend/internal/service/gemini_messages_compat_service.go b/backend/internal/service/gemini_messages_compat_service.go index 278cbfa6fa..a54c975cb3 100644 --- a/backend/internal/service/gemini_messages_compat_service.go +++ b/backend/internal/service/gemini_messages_compat_service.go @@ -529,94 +529,6 @@ func (s *GeminiMessagesCompatService) HasAntigravityAccounts(ctx context.Context return len(accounts) > 0, nil } -// SelectAccountForAIStudioEndpoints selects an account that is likely to succeed against -// generativelanguage.googleapis.com (e.g. GET /v1beta/models). -// -// Preference order: -// 1) API key accounts (AI Studio) -// 2) OAuth accounts without project_id (AI Studio OAuth) -// 3) OAuth accounts explicitly marked as ai_studio -// 4) Any remaining Gemini accounts (fallback) -func (s *GeminiMessagesCompatService) SelectAccountForAIStudioEndpoints(ctx context.Context, groupID *int64) (*Account, error) { - accounts, err := s.listSchedulableAccountsOnce(ctx, groupID, PlatformGemini, true) - if err != nil { - return nil, fmt.Errorf("query accounts failed: %w", err) - } - if len(accounts) == 0 { - return nil, errors.New("no available Gemini accounts") - } - - rank := func(a *Account) int { - if a == nil { - return 999 - } - switch a.Type { - case AccountTypeAPIKey: - if strings.TrimSpace(a.GetCredential("api_key")) != "" { - return 0 - } - return 9 - case AccountTypeOAuth: - if strings.TrimSpace(a.GetCredential("project_id")) == "" { - return 1 - } - if strings.TrimSpace(a.GetCredential("oauth_type")) == "ai_studio" { - return 2 - } - // Code Assist OAuth tokens often lack AI Studio scopes for models listing. - return 3 - case AccountTypeServiceAccount: - // Vertex service accounts use aiplatform.googleapis.com, not the AI Studio - // endpoint (generativelanguage.googleapis.com), so they cannot serve these requests. - return 999 - default: - return 10 - } - } - - var selected *Account - for i := range accounts { - acc := &accounts[i] - if selected == nil { - selected = acc - continue - } - - r1, r2 := rank(acc), rank(selected) - if r1 < r2 { - selected = acc - continue - } - if r1 > r2 { - continue - } - - if acc.Priority < selected.Priority { - selected = acc - } else if acc.Priority == selected.Priority { - switch { - case acc.LastUsedAt == nil && selected.LastUsedAt != nil: - selected = acc - case acc.LastUsedAt != nil && selected.LastUsedAt == nil: - // keep selected - case acc.LastUsedAt == nil && selected.LastUsedAt == nil: - if acc.Type == AccountTypeOAuth && selected.Type != AccountTypeOAuth { - selected = acc - } - default: - if acc.LastUsedAt.Before(*selected.LastUsedAt) { - selected = acc - } - } - } - } - - if selected == nil { - return nil, errors.New("no available Gemini accounts") - } - return s.hydrateSelectedAccount(ctx, selected) -} - func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Context, account *Account, body []byte) (*ForwardResult, error) { beginUpstreamResponseModelObservation(c) beginGeminiImageOutputObservation(c) @@ -770,19 +682,11 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex case AccountTypeServiceAccount: buildReq = func(ctx context.Context) (*http.Request, string, error) { - if s.tokenProvider == nil { - return nil, "", errors.New("gemini token provider not configured") - } - accessToken, err := s.tokenProvider.GetAccessToken(ctx, account) - if err != nil { - return nil, "", err - } - action := "generateContent" if req.Stream { action = "streamGenerateContent" } - fullURL, err := buildVertexGeminiURL(account.VertexProjectID(), account.VertexLocation(mappedModel), mappedModel, action, req.Stream) + fullURL, err := buildVertexGeminiAccountURL(account, mappedModel, action, req.Stream) if err != nil { return nil, "", err } @@ -793,7 +697,9 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex return nil, "", err } upstreamReq.Header.Set("Content-Type", "application/json") - upstreamReq.Header.Set("Authorization", "Bearer "+accessToken) + if err := setVertexGeminiAuth(ctx, upstreamReq, account, s.tokenProvider); err != nil { + return nil, "", err + } return upstreamReq, "x-request-id", nil } requestIDHeader = "x-request-id" @@ -1311,15 +1217,7 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin. case AccountTypeServiceAccount: buildReq = func(ctx context.Context) (*http.Request, string, error) { - if s.tokenProvider == nil { - return nil, "", errors.New("gemini token provider not configured") - } - accessToken, err := s.tokenProvider.GetAccessToken(ctx, account) - if err != nil { - return nil, "", err - } - - fullURL, err := buildVertexGeminiURL(account.VertexProjectID(), account.VertexLocation(mappedModel), mappedModel, upstreamAction, useUpstreamStream) + fullURL, err := buildVertexGeminiAccountURL(account, mappedModel, upstreamAction, useUpstreamStream) if err != nil { return nil, "", err } @@ -1329,7 +1227,9 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin. return nil, "", err } upstreamReq.Header.Set("Content-Type", "application/json") - upstreamReq.Header.Set("Authorization", "Bearer "+accessToken) + if err := setVertexGeminiAuth(ctx, upstreamReq, account, s.tokenProvider); err != nil { + return nil, "", err + } return upstreamReq, "x-request-id", nil } requestIDHeader = "x-request-id" @@ -1374,6 +1274,15 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin. return nil, transportErr } + // 引用缓存不存在、上游网关持有缓存的账号不可用先于错误策略识别：错误策略会写临时不可调度， + // 重试与换号也无法恢复。 + if passthrough, rebuilt := s.cachedContentPassthroughInLoop(ctx, account, resp); passthrough { + resp = rebuilt + break + } else { + resp = rebuilt + } + // 错误策略优先：匹配则跳过重试直接处理。 if matched, rebuilt := s.checkErrorPolicyInLoop(ctx, account, resp, mappedModel); matched { resp = rebuilt @@ -1467,6 +1376,14 @@ 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.writeGeminiCachedContentNotFound(c, account, resp, respBody, requestID, isOAuth) + } + // 上游网关的限流 / 不可用只针对其持有缓存的账号：不改本账号状态，交由绑定请求的耗尽处理透传。 + if isGeminiBoundCacheUpstreamPassthrough(ctx, account, resp.StatusCode) { + return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: respBody, BoundUpstreamPassthrough: true} + } // Best-effort fallback for OAuth tokens missing AI Studio scopes when calling countTokens. // This avoids Gemini SDKs failing hard during preflight token counting. // Checked before error policy so it always works regardless of custom error codes. @@ -1676,6 +1593,27 @@ func (s *GeminiMessagesCompatService) checkErrorPolicyInLoop( return policy != ErrorPolicyNone, rebuilt } +// 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) + _ = resp.Body.Close() + rebuilt = &http.Response{ + StatusCode: resp.StatusCode, + Header: resp.Header.Clone(), + Body: io.NopCloser(bytes.NewReader(body)), + } + return isGeminiCachedContentNotFound(resp.StatusCode, body), rebuilt +} + func (s *GeminiMessagesCompatService) shouldRetryGeminiUpstreamError(account *Account, statusCode int) bool { switch statusCode { case 429, 500, 502, 503, 504, 529: @@ -1781,25 +1719,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 == "" { @@ -1813,6 +1733,41 @@ func (s *GeminiMessagesCompatService) writeGeminiNativeUpstreamError(c *gin.Cont return fmt.Errorf("gemini upstream error: %d message=%s", resp.StatusCode, upstreamMsg) } +// writeGeminiCachedContentNotFound 记录上游原始错误，并向客户端统一返回缓存不存在的 403： +// 各上游的响应形态不同，且 Vertex 的文案含上游缓存资源 ID。失败由请求引用的缓存决定，ops 按客户端请求错误统计。 +func (s *GeminiMessagesCompatService) writeGeminiCachedContentNotFound(c *gin.Context, account *Account, resp *http.Response, respBody []byte, requestID string, isOAuth bool) error { + respBody = unwrapIfNeeded(isOAuth, respBody) + upstreamMsg := s.recordGeminiNativeUpstreamError(c, account, resp.StatusCode, respBody, requestID) + MarkOpsRequestScopedError(c, "permission_error") + MarkResponseCommitted(c) + c.Data(http.StatusForbidden, "application/json", []byte(GeminiCachedContentNotFoundResponse)) + return fmt.Errorf("gemini upstream error: %d message=%s", resp.StatusCode, upstreamMsg) +} + +// 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 { @@ -2635,6 +2590,25 @@ func isGeminiInsufficientScope(headers http.Header, body []byte) bool { return strings.Contains(lower, "insufficient authentication scopes") || strings.Contains(lower, "access_token_scope_insufficient") } +// 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 { + 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 + } +} + func estimateGeminiCountTokens(reqBody []byte) int { total := 0 @@ -2854,78 +2828,6 @@ func (s *GeminiMessagesCompatService) handleNativeStreamingResponse(c *gin.Conte return &geminiNativeStreamResult{usage: usage, firstTokenMs: firstTokenMs}, nil } -// ForwardAIStudioGET forwards a GET request to AI Studio (generativelanguage.googleapis.com) for -// endpoints like /v1beta/models and /v1beta/models/{model}. -// -// This is used to support Gemini SDKs that call models listing endpoints before generation. -func (s *GeminiMessagesCompatService) ForwardAIStudioGET(ctx context.Context, account *Account, path string) (*UpstreamHTTPResult, error) { - if account == nil { - return nil, errors.New("account is nil") - } - // path 会被直接拼到上游 base URL 后面，因此按路径护栏逐片段校验， - // 见 upstream_path_guard.go。 - sanitizedPath, ok := sanitizedUpstreamPathSuffix(path) - if !ok || sanitizedPath == "" { - return nil, errors.New("invalid path") - } - path = sanitizedPath - - baseURL := account.GetGeminiBaseURL(geminicli.AIStudioBaseURL) - normalizedBaseURL, err := s.validateUpstreamBaseURL(baseURL) - if err != nil { - return nil, err - } - fullURL := strings.TrimRight(normalizedBaseURL, "/") + path - - var proxyURL string - if account.ProxyID != nil && account.Proxy != nil { - proxyURL = account.Proxy.URL() - } - - req, err := http.NewRequestWithContext(ctx, http.MethodGet, fullURL, nil) - if err != nil { - return nil, err - } - - switch account.Type { - case AccountTypeAPIKey: - apiKey := strings.TrimSpace(account.GetCredential("api_key")) - if apiKey == "" { - return nil, errors.New("gemini api_key not configured") - } - req.Header.Set("x-goog-api-key", apiKey) - case AccountTypeOAuth: - if s.tokenProvider == nil { - return nil, errors.New("gemini token provider not configured") - } - accessToken, err := s.tokenProvider.GetAccessToken(ctx, account) - if err != nil { - return nil, err - } - req.Header.Set("Authorization", "Bearer "+accessToken) - default: - return nil, fmt.Errorf("unsupported account type: %s", account.Type) - } - - resp, err := s.httpUpstream.Do(req, proxyURL, account.ID, account.Concurrency) - if err != nil { - return nil, err - } - defer func() { _ = resp.Body.Close() }() - - body, _ := io.ReadAll(io.LimitReader(resp.Body, 8<<20)) - wwwAuthenticate := resp.Header.Get("Www-Authenticate") - filteredHeaders := responseheaders.FilterHeaders(resp.Header, s.responseHeaderFilter) - if wwwAuthenticate != "" { - filteredHeaders.Set("Www-Authenticate", wwwAuthenticate) - } - return &UpstreamHTTPResult{ - StatusCode: resp.StatusCode, - Headers: filteredHeaders, - Body: body, - }, nil -} - // unwrapGeminiResponse 解包 Gemini OAuth 响应中的 response 字段 // 使用 gjson 零拷贝提取，避免完整 Unmarshal+Marshal func unwrapGeminiResponse(raw []byte) ([]byte, error) { @@ -3097,6 +2999,10 @@ func (s *GeminiMessagesCompatService) handleGeminiUpstreamError(ctx context.Cont if account.IsPoolMode() && !account.IsCustomErrorCodesEnabled() { return } + if account.IsGeminiThirdPartyAPIKey() { + s.handleGeminiThirdParty429(ctx, account, body) + return + } oauthType := account.GeminiOAuthType() tierID := account.GeminiTierID() @@ -3145,16 +3051,36 @@ func (s *GeminiMessagesCompatService) handleGeminiUpstreamError(ctx context.Cont account.ID, resetTime, oauthType, tierID) } +// handleGeminiThirdParty429 处理第三方上游的 429：不套用 AI Studio 每日配额语义， +// 只认响应里显式给出的重试时间，否则走全局 429 兜底冷却。 +func (s *GeminiMessagesCompatService) handleGeminiThirdParty429(ctx context.Context, account *Account, body []byte) { + if resetAt := parseGeminiRetryDelayResetTime(body); resetAt != nil { + resetTime := time.Unix(*resetAt, 0) + _ = s.accountRepo.SetRateLimited(ctx, account.ID, resetTime) + logger.LegacyPrintf("service.gemini_messages_compat", "[Gemini 429] Account %d (third-party API Key) rate limited until %v", account.ID, resetTime) + return + } + if s.rateLimitService != nil { + s.rateLimitService.apply429FallbackRateLimit(ctx, account, "gemini_third_party_no_reset_time") + return + } + _ = s.accountRepo.SetRateLimited(ctx, account.ID, time.Now().Add(defaultRateLimit429CooldownSeconds*time.Second)) +} + // ParseGeminiRateLimitResetTime 解析 Gemini 格式的 429 响应，返回重置时间的 Unix 时间戳 func ParseGeminiRateLimitResetTime(body []byte) *int64 { - // 第一阶段：gjson 结构化提取 errMsg := gjson.GetBytes(body, "error.message").String() if looksLikeGeminiDailyQuota(errMsg) { if ts := nextGeminiDailyResetUnix(); ts != nil { return ts } } + return parseGeminiRetryDelayResetTime(body) +} +// parseGeminiRetryDelayResetTime 只解析响应里显式给出的重试时长，返回重置时间的 Unix 时间戳。 +func parseGeminiRetryDelayResetTime(body []byte) *int64 { + // 第一阶段：gjson 结构化提取 // 遍历 error.details 查找 quotaResetDelay（AI Studio）或 RetryInfo.retryDelay（Vertex） var found *int64 gjson.GetBytes(body, "error.details").ForEach(func(_, detail gjson.Result) bool { diff --git a/backend/internal/service/gemini_models_endpoints.go b/backend/internal/service/gemini_models_endpoints.go new file mode 100644 index 0000000000..02f3b985a1 --- /dev/null +++ b/backend/internal/service/gemini_models_endpoints.go @@ -0,0 +1,302 @@ +package service + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "slices" + "strings" + + "github.com/Wei-Shaw/sub2api/internal/pkg/gemini" + "github.com/Wei-Shaw/sub2api/internal/pkg/geminicli" + "github.com/Wei-Shaw/sub2api/internal/pkg/googleapi" + "github.com/Wei-Shaw/sub2api/internal/util/responseheaders" +) + +// Gemini 分组对外提供 AI Studio 兼容协议（GET /v1beta/models、/v1beta/models/{model}）， +// 分组内的账号可以是 AI Studio 兼容上游（API Key / OAuth），也可以是 Agent Platform (Vertex)。 +// 本文件把不同账号类型的模型接口统一成 AI Studio 格式的响应。 + +// errGeminiModelsUpstreamStatus 在 Vertex 分页读取遇到非 2xx 时中止分页，原始响应由调用方另行保存。 +var errGeminiModelsUpstreamStatus = errors.New("upstream model list returned non-2xx status") + +// SelectAccountForModelsEndpoints 为模型列表与模型详情接口选择账号。 +// +// 优先级： +// 1. 配置了 api_key 的 API Key 账号 +// 2. 不带 project_id 的 OAuth 账号（AI Studio OAuth） +// 3. 显式标记为 ai_studio 的 OAuth 账号 +// 4. Code Assist OAuth 账号（令牌常缺 AI Studio 权限范围，多半退回内置列表） +// 5. 其余类型的账号 +// 6. Vertex 服务账号（服务账号 JSON 认证，读取 Google 发布方模型列表） +// 7. Vertex API Key 账号（列表接口不接受 API Key，只能返回内置列表） +// +// Vertex 账号排在所有 AI Studio 兼容账号之后，混合分组始终由 AI Studio 兼容账号应答。 +// 同级账号按 priority、最近使用时间挑选。 +func (s *GeminiMessagesCompatService) SelectAccountForModelsEndpoints(ctx context.Context, groupID *int64) (*Account, error) { + accounts, err := s.listSchedulableAccountsOnce(ctx, groupID, PlatformGemini, true) + if err != nil { + return nil, fmt.Errorf("query accounts failed: %w", err) + } + if len(accounts) == 0 { + return nil, errors.New("no available Gemini accounts") + } + + var selected *Account + for i := range accounts { + acc := &accounts[i] + if selected == nil { + selected = acc + continue + } + + r1, r2 := modelsEndpointAccountRank(acc), modelsEndpointAccountRank(selected) + if r1 < r2 { + selected = acc + continue + } + if r1 > r2 { + continue + } + + if acc.Priority < selected.Priority { + selected = acc + } else if acc.Priority == selected.Priority { + switch { + case acc.LastUsedAt == nil && selected.LastUsedAt != nil: + selected = acc + case acc.LastUsedAt != nil && selected.LastUsedAt == nil: + // keep selected + case acc.LastUsedAt == nil && selected.LastUsedAt == nil: + if acc.Type == AccountTypeOAuth && selected.Type != AccountTypeOAuth { + selected = acc + } + default: + if acc.LastUsedAt.Before(*selected.LastUsedAt) { + selected = acc + } + } + } + } + + if selected == nil { + return nil, errors.New("no available Gemini accounts") + } + return s.hydrateSelectedAccount(ctx, selected) +} + +// modelsEndpointAccountRank 返回账号在模型接口选号中的优先级，数值越小越优先，顺序见 SelectAccountForModelsEndpoints。 +func modelsEndpointAccountRank(a *Account) int { + if a == nil { + return 999 + } + switch a.Type { + case AccountTypeAPIKey: + if strings.TrimSpace(a.GetCredential("api_key")) != "" { + return 0 + } + return 9 + case AccountTypeOAuth: + if strings.TrimSpace(a.GetCredential("project_id")) == "" { + return 1 + } + if strings.TrimSpace(a.GetCredential("oauth_type")) == "ai_studio" { + return 2 + } + return 3 + case AccountTypeServiceAccount: + if a.IsVertexAPIKey() { + return 12 + } + return 11 + default: + return 10 + } +} + +// ForwardListModels 读取账号上游的模型列表，返回 AI Studio 格式的响应（GET /v1beta/models）。 +// AI Studio 兼容账号原样返回上游响应；Vertex 服务账号读取 Google 发布方模型列表后转换格式， +// 上游非 2xx 时原样返回该响应；Vertex API Key 账号返回内置列表。 +func (s *GeminiMessagesCompatService) ForwardListModels(ctx context.Context, account *Account) (*UpstreamHTTPResult, error) { + if account == nil { + return nil, errors.New("account is nil") + } + if !account.IsVertexServiceAccount() { + return s.forwardGeminiModelsGET(ctx, account, "/v1beta/models") + } + if account.IsVertexAPIKey() { + return geminiModelsJSONResult(http.StatusOK, gemini.FallbackModelsList()) + } + + modelIDs, failure, err := s.listVertexGeminiModels(ctx, account, account.VertexLocation("")) + if err != nil { + return nil, err + } + if failure != nil { + return failure, nil + } + models := make([]gemini.Model, 0, len(modelIDs)) + for _, modelID := range modelIDs { + models = append(models, gemini.FallbackModel(modelID)) + } + return geminiModelsJSONResult(http.StatusOK, gemini.ModelsListResponse{Models: models}) +} + +// ForwardGetModel 读取账号上游的单个模型信息，返回 AI Studio 格式的响应（GET /v1beta/models/{model}）。 +// Vertex 服务账号在该模型生成请求所用区域的模型列表中查找，找不到时返回 404；Vertex API Key 账号返回内置信息。 +func (s *GeminiMessagesCompatService) ForwardGetModel(ctx context.Context, account *Account, model string) (*UpstreamHTTPResult, error) { + if account == nil { + return nil, errors.New("account is nil") + } + if !account.IsVertexServiceAccount() { + return s.forwardGeminiModelsGET(ctx, account, "/v1beta/models/"+model) + } + if account.IsVertexAPIKey() { + return geminiModelsJSONResult(http.StatusOK, gemini.FallbackModel(model)) + } + + modelIDs, failure, err := s.listVertexGeminiModels(ctx, account, account.VertexLocation(model)) + if err != nil { + return nil, err + } + if failure != nil { + return failure, nil + } + if slices.Contains(modelIDs, model) { + return geminiModelsJSONResult(http.StatusOK, gemini.FallbackModel(model)) + } + return geminiModelsJSONResult(http.StatusNotFound, map[string]any{ + "error": map[string]any{ + "code": http.StatusNotFound, + "message": fmt.Sprintf("models/%s is not found", model), + "status": googleapi.HTTPStatusToGoogleStatus(http.StatusNotFound), + }, + }) +} + +// forwardGeminiModelsGET 向 AI Studio 兼容上游转发模型接口的 GET 请求。 +// path 会被直接拼到上游 base URL 后面，因此按路径护栏逐片段校验，见 upstream_path_guard.go。 +func (s *GeminiMessagesCompatService) forwardGeminiModelsGET(ctx context.Context, account *Account, path string) (*UpstreamHTTPResult, error) { + sanitizedPath, ok := sanitizedUpstreamPathSuffix(path) + if !ok || sanitizedPath == "" { + return nil, errors.New("invalid path") + } + + baseURL := account.GetGeminiBaseURL(geminicli.AIStudioBaseURL) + normalizedBaseURL, err := s.validateUpstreamBaseURL(baseURL) + if err != nil { + return nil, err + } + fullURL := strings.TrimRight(normalizedBaseURL, "/") + sanitizedPath + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, fullURL, nil) + if err != nil { + return nil, err + } + + switch account.Type { + case AccountTypeAPIKey: + apiKey := strings.TrimSpace(account.GetCredential("api_key")) + if apiKey == "" { + return nil, errors.New("gemini api_key not configured") + } + req.Header.Set("x-goog-api-key", apiKey) + case AccountTypeOAuth: + if s.tokenProvider == nil { + return nil, errors.New("gemini token provider not configured") + } + accessToken, err := s.tokenProvider.GetAccessToken(ctx, account) + if err != nil { + return nil, err + } + req.Header.Set("Authorization", "Bearer "+accessToken) + default: + return nil, fmt.Errorf("unsupported account type: %s", account.Type) + } + + return s.doGeminiModelsGET(req, account) +} + +// listVertexGeminiModels 读取 Vertex 服务账号在指定区域可用的 Gemini 模型。 +// 上游返回非 2xx 时 failure 为该页的原始响应，modelIDs 与 err 均为空。 +func (s *GeminiMessagesCompatService) listVertexGeminiModels(ctx context.Context, account *Account, location string) (modelIDs []string, failure *UpstreamHTTPResult, err error) { + if s.tokenProvider == nil { + return nil, nil, errors.New("gemini token provider not configured") + } + accessToken, err := s.tokenProvider.GetAccessToken(ctx, account) + if err != nil { + return nil, nil, err + } + + modelIDs, err = listVertexGeminiPublisherModels(location, func(_ int, listURL string) ([]byte, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, listURL, nil) + if err != nil { + return nil, err + } + req.Header.Set("Accept", "application/json") + req.Header.Set("Authorization", "Bearer "+accessToken) + + res, err := s.doGeminiModelsGET(req, account) + if err != nil { + return nil, err + } + if res.StatusCode < http.StatusOK || res.StatusCode >= http.StatusMultipleChoices { + failure = res + return nil, errGeminiModelsUpstreamStatus + } + return res.Body, nil + }) + if failure != nil { + return nil, failure, nil + } + if err != nil { + return nil, nil, err + } + return modelIDs, nil, nil +} + +// doGeminiModelsGET 发送模型接口请求并读取响应。响应体上限取 gateway.models_list_read_max_bytes，超出时报错；响应头按配置过滤。 +func (s *GeminiMessagesCompatService) doGeminiModelsGET(req *http.Request, account *Account) (*UpstreamHTTPResult, error) { + var proxyURL string + if account.ProxyID != nil && account.Proxy != nil { + proxyURL = account.Proxy.URL() + } + + resp, err := s.httpUpstream.Do(req, proxyURL, account.ID, account.Concurrency) + if err != nil { + return nil, err + } + defer func() { _ = resp.Body.Close() }() + + bodyLimit := resolveModelsListReadLimit(s.cfg) + body, _ := io.ReadAll(io.LimitReader(resp.Body, bodyLimit+1)) + if int64(len(body)) > bodyLimit { + return nil, fmt.Errorf("upstream model list response exceeds %d bytes", bodyLimit) + } + wwwAuthenticate := resp.Header.Get("Www-Authenticate") + filteredHeaders := responseheaders.FilterHeaders(resp.Header, s.responseHeaderFilter) + if wwwAuthenticate != "" { + filteredHeaders.Set("Www-Authenticate", wwwAuthenticate) + } + return &UpstreamHTTPResult{ + StatusCode: resp.StatusCode, + Headers: filteredHeaders, + Body: body, + }, nil +} + +// geminiModelsJSONResult 把网关生成的模型接口响应包装成 UpstreamHTTPResult，交给 handler 统一写出。 +func geminiModelsJSONResult(status int, payload any) (*UpstreamHTTPResult, error) { + body, err := json.Marshal(payload) + if err != nil { + return nil, err + } + return &UpstreamHTTPResult{ + StatusCode: status, + Headers: http.Header{"Content-Type": []string{"application/json; charset=utf-8"}}, + Body: body, + }, nil +} diff --git a/backend/internal/service/gemini_models_endpoints_test.go b/backend/internal/service/gemini_models_endpoints_test.go new file mode 100644 index 0000000000..e33b85d809 --- /dev/null +++ b/backend/internal/service/gemini_models_endpoints_test.go @@ -0,0 +1,208 @@ +//go:build unit + +package service + +import ( + "context" + "encoding/json" + "net/http" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/gemini" + "github.com/stretchr/testify/require" +) + +func newGeminiModelsEndpointsTestService(accounts []Account, upstream HTTPUpstream) *GeminiMessagesCompatService { + repo := &mockAccountRepoForGemini{listByGroupFunc: func(context.Context, int64, []string) ([]Account, error) { + return accounts, nil + }} + return &GeminiMessagesCompatService{ + accountRepo: repo, + httpUpstream: upstream, + cfg: &config.Config{}, + tokenProvider: NewGeminiTokenProvider(nil, &vertexModelsTokenCacheStub{token: "ya29.vertex-token"}, nil), + } +} + +func decodeGeminiModelNames(t *testing.T, body []byte) []string { + t.Helper() + var got gemini.ModelsListResponse + require.NoError(t, json.Unmarshal(body, &got)) + names := make([]string, 0, len(got.Models)) + for _, model := range got.Models { + require.Contains(t, model.SupportedGenerationMethods, "generateContent") + names = append(names, model.Name) + } + return names +} + +func TestSelectAccountForModelsEndpoints_VertexServiceAccountOnlyGroup(t *testing.T) { + groupID := int64(8) + vertex := *vertexModelSyncTestAccount(nil) + vertex.ID = 18 + svc := newGeminiModelsEndpointsTestService([]Account{vertex}, nil) + + account, err := svc.SelectAccountForModelsEndpoints(context.Background(), &groupID) + require.NoError(t, err) + require.Equal(t, int64(18), account.ID) +} + +func TestModelsEndpointAccountRank_Order(t *testing.T) { + ordered := []*Account{ + {Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "k"}}, + {Platform: PlatformGemini, Type: AccountTypeOAuth, Credentials: map[string]any{}}, + {Platform: PlatformGemini, Type: AccountTypeOAuth, Credentials: map[string]any{"project_id": "p", "oauth_type": "ai_studio"}}, + {Platform: PlatformGemini, Type: AccountTypeOAuth, Credentials: map[string]any{"project_id": "p", "oauth_type": "code_assist"}}, + {Platform: PlatformGemini, Type: AccountTypeUpstream, Credentials: map[string]any{"api_key": "k"}}, + vertexModelSyncTestAccount(nil), + {Platform: PlatformGemini, Type: AccountTypeServiceAccount, Credentials: map[string]any{"auth_mode": "apikey", "api_key": "AQ.k"}}, + } + for i := 1; i < len(ordered); i++ { + require.Less(t, modelsEndpointAccountRank(ordered[i-1]), modelsEndpointAccountRank(ordered[i]), "position %d", i) + } +} + +func TestSelectAccountForModelsEndpoints_PrefersAPIKeyOverVertex(t *testing.T) { + groupID := int64(8) + vertex := *vertexModelSyncTestAccount(nil) + vertex.ID = 18 + apiKey := Account{ID: 19, Platform: PlatformGemini, Type: AccountTypeAPIKey, Priority: 9, Credentials: map[string]any{"api_key": "k"}} + svc := newGeminiModelsEndpointsTestService([]Account{vertex, apiKey}, nil) + + account, err := svc.SelectAccountForModelsEndpoints(context.Background(), &groupID) + require.NoError(t, err) + require.Equal(t, int64(19), account.ID) +} + +func TestForwardListModels_VertexServiceAccountReturnsAIStudioFormat(t *testing.T) { + upstream := &httpUpstreamRecorder{responses: []*http.Response{ + vertexModelsJSONResponse(http.StatusOK, `{"publisherModels":[ + {"name":"publishers/google/models/gemini-3.8-flash"}, + {"name":"publishers/google/models/gemma4"} + ],"nextPageToken":"page-2"}`), + vertexModelsJSONResponse(http.StatusOK, `{"publisherModels":[{"name":"publishers/google/models/gemini-2.5-flash"}]}`), + }} + svc := newGeminiModelsEndpointsTestService(nil, upstream) + + res, err := svc.ForwardListModels(context.Background(), vertexModelSyncTestAccount(map[string]any{"location": "europe-west4"})) + require.NoError(t, err) + require.Equal(t, http.StatusOK, res.StatusCode) + require.Equal(t, []string{"models/gemini-2.5-flash", "models/gemini-3.8-flash"}, decodeGeminiModelNames(t, res.Body)) + + require.Len(t, upstream.requests, 2) + require.Equal(t, "https://europe-west4-aiplatform.googleapis.com/v1beta1/publishers/google/models?pageSize=300", upstream.requests[0].URL.String()) + require.Equal(t, "https://europe-west4-aiplatform.googleapis.com/v1beta1/publishers/google/models?pageSize=300&pageToken=page-2", upstream.requests[1].URL.String()) + for _, req := range upstream.requests { + require.Equal(t, "Bearer ya29.vertex-token", req.Header.Get("Authorization")) + } +} + +func TestForwardListModels_VertexUpstreamErrorPassesThrough(t *testing.T) { + body := `{"error":{"code":403,"status":"PERMISSION_DENIED"}}` + upstream := &httpUpstreamRecorder{resp: vertexModelsJSONResponse(http.StatusForbidden, body)} + svc := newGeminiModelsEndpointsTestService(nil, upstream) + + res, err := svc.ForwardListModels(context.Background(), vertexModelSyncTestAccount(map[string]any{"location": "global"})) + require.NoError(t, err) + require.Equal(t, http.StatusForbidden, res.StatusCode) + require.JSONEq(t, body, string(res.Body)) +} + +func TestForwardListModels_VertexAPIKeyReturnsBuiltinList(t *testing.T) { + upstream := &httpUpstreamRecorder{} + svc := newGeminiModelsEndpointsTestService(nil, upstream) + account := &Account{ID: 174, Platform: PlatformGemini, Type: AccountTypeServiceAccount, Credentials: map[string]any{ + "auth_mode": "apikey", + "api_key": "AQ.test-key", + }} + + res, err := svc.ForwardListModels(context.Background(), account) + require.NoError(t, err) + require.Equal(t, http.StatusOK, res.StatusCode) + expected, err := json.Marshal(gemini.FallbackModelsList()) + require.NoError(t, err) + require.JSONEq(t, string(expected), string(res.Body)) + require.Empty(t, upstream.requests) +} + +func TestForwardListModels_APIKeyForwardsToBaseURL(t *testing.T) { + upstream := &httpUpstreamRecorder{resp: vertexModelsJSONResponse(http.StatusOK, `{"models":[{"name":"models/gemini-native"}]}`)} + svc := newGeminiModelsEndpointsTestService(nil, upstream) + account := &Account{ID: 261, Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: map[string]any{ + "api_key": "sk-pool", + "base_url": "https://relay.example.com", + }} + + res, err := svc.ForwardListModels(context.Background(), account) + require.NoError(t, err) + require.Equal(t, http.StatusOK, res.StatusCode) + require.JSONEq(t, `{"models":[{"name":"models/gemini-native"}]}`, string(res.Body)) + require.Equal(t, "https://relay.example.com/v1beta/models", upstream.lastReq.URL.String()) + require.Equal(t, "sk-pool", upstream.lastReq.Header.Get("x-goog-api-key")) +} + +func TestForwardGetModel_VertexServiceAccountLooksUpModelList(t *testing.T) { + listBody := `{"publisherModels":[{"name":"publishers/google/models/gemini-3.8-flash"}]}` + + t.Run("listed model", func(t *testing.T) { + upstream := &httpUpstreamRecorder{resp: vertexModelsJSONResponse(http.StatusOK, listBody)} + svc := newGeminiModelsEndpointsTestService(nil, upstream) + + res, err := svc.ForwardGetModel(context.Background(), vertexModelSyncTestAccount(nil), "gemini-3.8-flash") + require.NoError(t, err) + require.Equal(t, http.StatusOK, res.StatusCode) + var model gemini.Model + require.NoError(t, json.Unmarshal(res.Body, &model)) + require.Equal(t, "models/gemini-3.8-flash", model.Name) + }) + + t.Run("unlisted model", func(t *testing.T) { + upstream := &httpUpstreamRecorder{resp: vertexModelsJSONResponse(http.StatusOK, listBody)} + svc := newGeminiModelsEndpointsTestService(nil, upstream) + + res, err := svc.ForwardGetModel(context.Background(), vertexModelSyncTestAccount(nil), "gemini-unknown") + require.NoError(t, err) + require.Equal(t, http.StatusNotFound, res.StatusCode) + require.JSONEq(t, `{"error":{"code":404,"message":"models/gemini-unknown is not found","status":"NOT_FOUND"}}`, string(res.Body)) + }) +} + +func TestForwardGetModel_VertexUsesModelLocation(t *testing.T) { + upstream := &httpUpstreamRecorder{resp: vertexModelsJSONResponse(http.StatusOK, `{"publisherModels":[{"name":"publishers/google/models/gemini-3.8-flash"}]}`)} + svc := newGeminiModelsEndpointsTestService(nil, upstream) + account := vertexModelSyncTestAccount(map[string]any{ + "location": "europe-west4", + "vertex_model_locations": map[string]any{"gemini-3.8-flash": "global"}, + }) + + res, err := svc.ForwardGetModel(context.Background(), account, "gemini-3.8-flash") + require.NoError(t, err) + require.Equal(t, http.StatusOK, res.StatusCode) + require.Equal(t, "https://aiplatform.googleapis.com/v1beta1/publishers/google/models?pageSize=300", upstream.lastReq.URL.String()) +} + +func TestForwardListModels_VertexMidPaginationErrorPassesThrough(t *testing.T) { + body := `{"error":{"code":503,"status":"UNAVAILABLE"}}` + upstream := &httpUpstreamRecorder{responses: []*http.Response{ + vertexModelsJSONResponse(http.StatusOK, `{"publisherModels":[{"name":"publishers/google/models/gemini-3.8-flash"}],"nextPageToken":"page-2"}`), + vertexModelsJSONResponse(http.StatusServiceUnavailable, body), + }} + svc := newGeminiModelsEndpointsTestService(nil, upstream) + + res, err := svc.ForwardListModels(context.Background(), vertexModelSyncTestAccount(nil)) + require.NoError(t, err) + require.Equal(t, http.StatusServiceUnavailable, res.StatusCode) + require.JSONEq(t, body, string(res.Body)) + require.Len(t, upstream.requests, 2) +} + +func TestForwardListModels_ResponseOverReadLimitFails(t *testing.T) { + upstream := &httpUpstreamRecorder{resp: vertexModelsJSONResponse(http.StatusOK, `{"models":[{"name":"models/gemini-native"}]}`)} + svc := newGeminiModelsEndpointsTestService(nil, upstream) + svc.cfg.Gateway.ModelsListReadMaxBytes = 16 + account := &Account{ID: 261, Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "sk-pool"}} + + _, err := svc.ForwardListModels(context.Background(), account) + require.ErrorContains(t, err, "exceeds 16 bytes") +} diff --git a/backend/internal/service/gemini_quota.go b/backend/internal/service/gemini_quota.go index 3a70232c87..be927754cf 100644 --- a/backend/internal/service/gemini_quota.go +++ b/backend/internal/service/gemini_quota.go @@ -373,6 +373,10 @@ func geminiQuotaTierKeyForAccount(account *Account) string { if account == nil || account.Platform != PlatformGemini { return "" } + // 本地模拟配额只适用于按 Google 账号等级限额的通道；第三方上游与 Vertex 按量付费不套用。 + if account.IsGeminiThirdPartyAPIKey() || account.IsVertexServiceAccount() { + return "" + } // Note: GeminiOAuthType() already defaults legacy (project_id present) to code_assist. oauthType := strings.ToLower(strings.TrimSpace(account.GeminiOAuthType())) diff --git a/backend/internal/service/gemini_quota_test.go b/backend/internal/service/gemini_quota_test.go new file mode 100644 index 0000000000..36e60f291e --- /dev/null +++ b/backend/internal/service/gemini_quota_test.go @@ -0,0 +1,80 @@ +//go:build unit + +package service + +import ( + "context" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/stretchr/testify/require" +) + +func TestAccountIsGeminiThirdPartyAPIKey(t *testing.T) { + tests := []struct { + name string + account *Account + expected bool + }{ + {"nil_account", nil, false}, + {"empty_base_url_is_official", geminiQuotaAPIKeyAccount(""), false}, + {"official_base_url", geminiQuotaAPIKeyAccount("https://generativelanguage.googleapis.com"), false}, + {"official_base_url_trailing_slash", geminiQuotaAPIKeyAccount("https://generativelanguage.googleapis.com/"), false}, + {"official_host_over_http", geminiQuotaAPIKeyAccount("http://generativelanguage.googleapis.com"), true}, + {"official_host_with_path_prefix", geminiQuotaAPIKeyAccount("https://generativelanguage.googleapis.com/proxy"), true}, + {"relay_base_url", geminiQuotaAPIKeyAccount("https://relay.example.com"), true}, + {"oauth_account", &Account{Platform: PlatformGemini, Type: AccountTypeOAuth, Credentials: map[string]any{"base_url": "https://relay.example.com"}}, false}, + {"vertex_service_account", &Account{Platform: PlatformGemini, Type: AccountTypeServiceAccount}, false}, + {"non_gemini_platform", &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Credentials: map[string]any{"base_url": "https://relay.example.com"}}, false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + require.Equal(t, tt.expected, tt.account.IsGeminiThirdPartyAPIKey()) + }) + } +} + +func TestGeminiQuotaService_QuotaForAccountSkipsThirdPartyAndVertex(t *testing.T) { + svc := NewGeminiQuotaService(&config.Config{}, nil) + ctx := context.Background() + + quota, ok := svc.QuotaForAccount(ctx, geminiQuotaAPIKeyAccount("")) + require.True(t, ok) + require.Equal(t, int64(50), quota.ProRPD) + + paid := geminiQuotaAPIKeyAccount("https://generativelanguage.googleapis.com") + paid.Credentials["tier_id"] = GeminiTierAIStudioPaid + quota, ok = svc.QuotaForAccount(ctx, paid) + require.True(t, ok) + require.Equal(t, int64(-1), quota.ProRPD) + + relay := geminiQuotaAPIKeyAccount("https://relay.example.com") + relay.Credentials["tier_id"] = GeminiTierAIStudioFree + _, ok = svc.QuotaForAccount(ctx, relay) + require.False(t, ok, "第三方上游不套用 AI Studio 模拟配额") + + vertex := &Account{ + Platform: PlatformGemini, + Type: AccountTypeServiceAccount, + Credentials: map[string]any{"project_id": "my-project", "tier_id": "vertex"}, + } + _, ok = svc.QuotaForAccount(ctx, vertex) + require.False(t, ok, "Vertex 按量付费不套用模拟配额") + + googleOne := &Account{ + Platform: PlatformGemini, + Type: AccountTypeOAuth, + Credentials: map[string]any{"oauth_type": "google_one", "tier_id": GeminiTierGoogleOneFree}, + } + quota, ok = svc.QuotaForAccount(ctx, googleOne) + require.True(t, ok) + require.Positive(t, quota.SharedRPD) +} + +func geminiQuotaAPIKeyAccount(baseURL string) *Account { + credentials := map[string]any{} + if baseURL != "" { + credentials["base_url"] = baseURL + } + return &Account{Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: credentials} +} diff --git a/backend/internal/service/gemini_vertex_api_key_forward_test.go b/backend/internal/service/gemini_vertex_api_key_forward_test.go new file mode 100644 index 0000000000..bf8be6c1c0 --- /dev/null +++ b/backend/internal/service/gemini_vertex_api_key_forward_test.go @@ -0,0 +1,101 @@ +//go:build unit + +package service + +import ( + "bytes" + "context" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +const geminiVertexAPIKeyTestResponse = `{"candidates":[{"content":{"parts":[{"text":"hello"}],"role":"model"},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":5,"totalTokenCount":15}}` + +func newGeminiVertexAPIKeyTestService() (*GeminiMessagesCompatService, *geminiCompatHTTPUpstreamStub) { + httpStub := &geminiCompatHTTPUpstreamStub{ + response: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}, "X-Request-Id": []string{"vertex-req-1"}}, + Body: io.NopCloser(strings.NewReader(geminiVertexAPIKeyTestResponse)), + }, + } + // tokenProvider 留空：API Key 认证不依赖 access token。 + return &GeminiMessagesCompatService{httpUpstream: httpStub, cfg: &config.Config{}}, httpStub +} + +func geminiVertexAPIKeyTestAccount(projectID string) *Account { + credentials := map[string]any{ + "auth_mode": "apikey", + "api_key": "vertex-test-key", + "location": "us-central1", + "tier_id": "vertex", + } + if projectID != "" { + credentials["project_id"] = projectID + } + return &Account{ + ID: 811, + Platform: PlatformGemini, + Type: AccountTypeServiceAccount, + Credentials: credentials, + Concurrency: 1, + } +} + +func requireVertexAPIKeyAuth(t *testing.T, req *http.Request) { + t.Helper() + require.NotNil(t, req) + require.Equal(t, "vertex-test-key", req.Header.Get("x-goog-api-key")) + require.Empty(t, req.Header.Get("Authorization")) +} + +func TestGeminiForwardNative_VertexAPIKeyWithoutProjectUsesExpressEndpoint(t *testing.T) { + gin.SetMode(gin.TestMode) + svc, httpStub := newGeminiVertexAPIKeyTestService() + c, _ := newGeminiNativeTestContext(t) + + result, err := svc.ForwardNative(context.Background(), c, geminiVertexAPIKeyTestAccount(""), + "gemini-3-flash-preview", "generateContent", false, []byte(`{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`)) + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, 1, httpStub.calls) + require.Equal(t, "https://aiplatform.googleapis.com/v1/publishers/google/models/gemini-3-flash-preview:generateContent", httpStub.lastReq.URL.String()) + requireVertexAPIKeyAuth(t, httpStub.lastReq) +} + +func TestGeminiMessagesCompatServiceForward_VertexAPIKeyWithProjectUsesProjectEndpoint(t *testing.T) { + gin.SetMode(gin.TestMode) + svc, httpStub := newGeminiVertexAPIKeyTestService() + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) + body := []byte(`{"model":"gemini-3-flash-preview","max_tokens":16,"messages":[{"role":"user","content":"hello"}]}`) + + result, err := svc.Forward(context.Background(), c, geminiVertexAPIKeyTestAccount("my-project"), body) + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, "https://us-central1-aiplatform.googleapis.com/v1/projects/my-project/locations/us-central1/publishers/google/models/gemini-3-flash-preview:generateContent", httpStub.lastReq.URL.String()) + requireVertexAPIKeyAuth(t, httpStub.lastReq) +} + +func TestGeminiForwardAsChatCompletions_VertexAPIKeyUsesAPIKeyHeader(t *testing.T) { + gin.SetMode(gin.TestMode) + svc, httpStub := newGeminiVertexAPIKeyTestService() + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + body := []byte(`{"model":"gemini-3-flash-preview","messages":[{"role":"user","content":"hi"}]}`) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body)) + + result, err := svc.ForwardAsChatCompletions(context.Background(), c, geminiVertexAPIKeyTestAccount(""), body) + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, http.StatusOK, rec.Code) + require.Equal(t, "https://aiplatform.googleapis.com/v1/publishers/google/models/gemini-3-flash-preview:generateContent", httpStub.lastReq.URL.String()) + requireVertexAPIKeyAuth(t, httpStub.lastReq) +} diff --git a/backend/internal/service/ops_upstream_context.go b/backend/internal/service/ops_upstream_context.go index 4d09a43ae6..a8d0625fd8 100644 --- a/backend/internal/service/ops_upstream_context.go +++ b/backend/internal/service/ops_upstream_context.go @@ -52,6 +52,14 @@ const ( // ensureForwardErrorResponse 检查此 key，为 true 时跳过兜底写入，避免在已完成的 JSON 后追加 SSE。 ResponseCommittedKey = "response_committed" + // OpsLocalErrorTypeKey / OpsLocalErrorCodeKey 记录网关以 Google 错误格式回写的错误类型与错误码： + // Google 错误体没有 type 字段，ops_error_logger 据此按与 Claude / OpenAI 格式错误相同的口径分类。 + OpsLocalErrorTypeKey = "ops_local_error_type" + OpsLocalErrorCodeKey = "ops_local_error_code" + // OpsRequestScopedErrorKey 标记失败由请求内容决定（如引用的显式缓存不存在）：只按错误类型分类， + // 上游错误上下文保留为明细但不参与归因。 + OpsRequestScopedErrorKey = "ops_request_scoped_error" + OpsClientBusinessLimitedKey = "ops_client_business_limited" OpsClientBusinessLimitedReasonKey = "ops_client_business_limited_reason" OpsClientBusinessLimitedReasonIPRestriction = "api_key_ip_restriction" @@ -134,6 +142,42 @@ func OpsClientBusinessLimitedReason(c *gin.Context) string { return strings.TrimSpace(reason) } +// SetOpsLocalErrorType 登记响应错误体缺少的错误类型与错误码，并清除请求级失败标记；errType 为空表示不登记。 +func SetOpsLocalErrorType(c *gin.Context, errType, code string) { + if c == nil { + return + } + c.Set(OpsLocalErrorTypeKey, strings.TrimSpace(errType)) + c.Set(OpsLocalErrorCodeKey, strings.TrimSpace(code)) + c.Set(OpsRequestScopedErrorKey, false) +} + +// OpsLocalErrorType 返回 SetOpsLocalErrorType 登记的错误类型与错误码；未登记时返回空串。 +func OpsLocalErrorType(c *gin.Context) (errType, code string) { + if c == nil { + return "", "" + } + errType = c.GetString(OpsLocalErrorTypeKey) + if errType == "" { + return "", "" + } + return errType, c.GetString(OpsLocalErrorCodeKey) +} + +// MarkOpsRequestScopedError 登记错误类型，并标记失败由请求内容决定：上游错误上下文不参与 ops 归因。 +func MarkOpsRequestScopedError(c *gin.Context, errType string) { + if c == nil { + return + } + SetOpsLocalErrorType(c, errType, "") + c.Set(OpsRequestScopedErrorKey, true) +} + +// IsOpsRequestScopedError 判断 MarkOpsRequestScopedError 是否标记了本次请求的失败。 +func IsOpsRequestScopedError(c *gin.Context) bool { + return c != nil && c.GetBool(OpsRequestScopedErrorKey) +} + // OpsStreamError 描述承载在 2xx 响应上的带内错误：网关在响应状态已固化为 200 之后 // 就地以 SSE error 帧返回的错误（并发限流回退、Wait 后二次计费校验失败、流开始后才无 // 可用账号等），以及上游 2xx 正文或事件里携带的错误结果（NonStream 标记非流式正文）。 diff --git a/backend/internal/service/pricing_service.go b/backend/internal/service/pricing_service.go index ac29fd0f3a..408de5d8b5 100644 --- a/backend/internal/service/pricing_service.go +++ b/backend/internal/service/pricing_service.go @@ -231,6 +231,9 @@ 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），一般由定价覆盖文件提供， + // 配置后优先于内置价（见 BillingService.CacheStoragePricePerTokenHour）。 + CacheStorageCostPerTokenPerHour float64 `json:"cache_storage_cost_per_token_per_hour"` // TokenPricingAbsent 表示源数据中 input/output token 价格均缺失（仅有图片价）。 // 此类条目只可用于图片计费，token 计费必须回退到 fallback 或 fail-closed， @@ -266,6 +269,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 动态价格服务 @@ -738,6 +742,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/pricing_service_test.go b/backend/internal/service/pricing_service_test.go index 0dce3e435c..3d9f0cb09e 100644 --- a/backend/internal/service/pricing_service_test.go +++ b/backend/internal/service/pricing_service_test.go @@ -613,10 +613,10 @@ func TestBillingService_Gemini36FlashThinkingTierFallbacksAreBillable(t *testing t.Run(model, func(t *testing.T) { cost, err := svc.CalculateCost(model, tokens, 1) require.NoError(t, err) - require.InDelta(t, 1.5, cost.InputCost, 1e-12) - require.InDelta(t, 7.5, cost.OutputCost, 1e-12) - require.InDelta(t, 0.15, cost.CacheReadCost, 1e-12) - require.InDelta(t, 9.15, cost.TotalCost, 1e-12) + require.InDelta(t, 0.75, cost.InputCost, 1e-12) + require.InDelta(t, 3.75, cost.OutputCost, 1e-12) + require.InDelta(t, 0.075, cost.CacheReadCost, 1e-12) + require.InDelta(t, 4.575, cost.TotalCost, 1e-12) }) } } @@ -725,9 +725,9 @@ func TestDefaultPricingIncludesGemini36FlashRates(t *testing.T) { t.Run(model, func(t *testing.T) { pricing, err := billingSvc.GetModelPricing(model) require.NoError(t, err) - require.InDelta(t, 1.5e-6, pricing.InputPricePerToken, 1e-12) - require.InDelta(t, 7.5e-6, pricing.OutputPricePerToken, 1e-12) - require.InDelta(t, 0.15e-6, pricing.CacheReadPricePerToken, 1e-12) + require.InDelta(t, 0.75e-6, pricing.InputPricePerToken, 1e-12) + require.InDelta(t, 3.75e-6, pricing.OutputPricePerToken, 1e-12) + require.InDelta(t, 0.075e-6, pricing.CacheReadPricePerToken, 1e-12) }) } } diff --git a/backend/internal/service/upstream_models.go b/backend/internal/service/upstream_models.go index 3c5eb2fcd2..447f2fdb98 100644 --- a/backend/internal/service/upstream_models.go +++ b/backend/internal/service/upstream_models.go @@ -10,6 +10,7 @@ import ( "net/http" "net/url" "sort" + "strconv" "strings" "time" @@ -25,6 +26,10 @@ const ( UpstreamModelMetadataExtraKey = "upstream_model_metadata" UpstreamModelMetadataIncompleteCode = "upstream_model_metadata_incomplete" UpstreamModelMetadataPartialCode = "upstream_model_metadata_partial" + + vertexPublisherModelsPageSize = 300 + vertexPublisherModelsMaxPages = 20 + vertexGooglePublisherModelPrefix = "publishers/google/models/" ) type UpstreamModelMetadata struct { @@ -741,34 +746,19 @@ func (s *AccountTestService) fetchUpstreamModelList(ctx context.Context, account return nil, nil, newUpstreamModelSyncConfigError("Upstream HTTP client is not configured", nil) } + if account.IsGemini() && account.IsVertexServiceAccount() { + models, err := s.fetchVertexGeminiUpstreamModels(ctx, account) + return models, nil, err + } + req, err := s.buildUpstreamModelsRequest(ctx, account) if err != nil { return nil, nil, err } - proxyURL := upstreamModelsProxyURL(account) - resp, err := s.doUpstreamModelsRequest(req, proxyURL, account) + body, err := s.readUpstreamModelsResponse(req, account) if err != nil { - return nil, nil, newUpstreamModelSyncUpstreamError("Failed to request upstream model list", err) - } - defer func() { _ = resp.Body.Close() }() - - bodyLimit := resolveModelsListReadLimit(s.cfg) - body, err := io.ReadAll(io.LimitReader(resp.Body, bodyLimit+1)) - if err != nil { - return nil, nil, newUpstreamModelSyncUpstreamError("Failed to read upstream model list", err) - } - if int64(len(body)) > bodyLimit { - return nil, nil, newUpstreamModelSyncUpstreamError("Upstream model list response is too large", fmt.Errorf("response exceeds %d bytes", bodyLimit)) - } - - if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices { - return nil, nil, &UpstreamModelSyncError{ - Kind: UpstreamModelSyncErrorUpstream, - Message: fmt.Sprintf("Upstream model list request failed with HTTP %d", resp.StatusCode), - StatusCode: resp.StatusCode, - Err: fmt.Errorf("upstream model list returned HTTP %d", resp.StatusCode), - } + return nil, nil, err } extractModels := extractUpstreamModelIDs @@ -786,6 +776,128 @@ func (s *AccountTestService) fetchUpstreamModelList(ctx context.Context, account return models, body, nil } +// readUpstreamModelsResponse 发送模型列表请求并读取受限大小的响应体；非 2xx 返回带状态码的同步错误。 +func (s *AccountTestService) readUpstreamModelsResponse(req *http.Request, account *Account) ([]byte, error) { + resp, err := s.doUpstreamModelsRequest(req, upstreamModelsProxyURL(account), account) + if err != nil { + return nil, newUpstreamModelSyncUpstreamError("Failed to request upstream model list", err) + } + defer func() { _ = resp.Body.Close() }() + + bodyLimit := resolveModelsListReadLimit(s.cfg) + body, err := io.ReadAll(io.LimitReader(resp.Body, bodyLimit+1)) + if err != nil { + return nil, newUpstreamModelSyncUpstreamError("Failed to read upstream model list", err) + } + if int64(len(body)) > bodyLimit { + return nil, newUpstreamModelSyncUpstreamError("Upstream model list response is too large", fmt.Errorf("response exceeds %d bytes", bodyLimit)) + } + + if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices { + return nil, &UpstreamModelSyncError{ + Kind: UpstreamModelSyncErrorUpstream, + Message: fmt.Sprintf("Upstream model list request failed with HTTP %d", resp.StatusCode), + StatusCode: resp.StatusCode, + Err: fmt.Errorf("upstream model list returned HTTP %d", resp.StatusCode), + } + } + return body, nil +} + +// fetchVertexGeminiUpstreamModels 按账号默认区域读取 Vertex 上可用的 Gemini 模型。列表接口不接受 API Key。 +func (s *AccountTestService) fetchVertexGeminiUpstreamModels(ctx context.Context, account *Account) ([]string, error) { + if account.IsVertexAPIKey() { + return nil, newUpstreamModelSyncUnsupportedError("Vertex API key accounts do not support upstream model listing", nil) + } + if _, err := parseVertexServiceAccountKey(account); err != nil { + return nil, newUpstreamModelSyncConfigError("Invalid Vertex service account credentials", err) + } + location := account.VertexLocation("") + if _, err := buildVertexPublisherModelsURL(location, ""); err != nil { + return nil, newUpstreamModelSyncConfigError("Invalid Vertex location", err) + } + if s.geminiTokenProvider == nil { + return nil, newUpstreamModelSyncConfigError("Gemini token provider is not configured", nil) + } + accessToken, err := s.geminiTokenProvider.GetAccessToken(ctx, account) + if err != nil { + return nil, newUpstreamModelSyncUpstreamError("Failed to get Vertex access token", err) + } + accessToken = strings.TrimSpace(accessToken) + if accessToken == "" { + return nil, newUpstreamModelSyncConfigError("No Vertex access token is available", nil) + } + + return listVertexGeminiPublisherModels(location, func(page int, listURL string) ([]byte, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, listURL, nil) + if err != nil { + return nil, newUpstreamModelSyncConfigError("Invalid Vertex model list URL", err) + } + req.Header.Set("Accept", "application/json") + req.Header.Set("Authorization", "Bearer "+accessToken) + + body, err := s.readUpstreamModelsResponse(req, account) + if err != nil { + // 只有首页的 404/405 表示列表端点不可用；翻页途中的失败不触发配置模型回退。 + var syncErr *UpstreamModelSyncError + if page > 0 && errors.As(err, &syncErr) { + syncErr.StatusCode = 0 + } + return nil, err + } + return body, nil + }) +} + +// listVertexGeminiPublisherModels 按区域分页读取 Google 发布方模型列表，只保留 gemini- 前缀的模型（去重排序）。 +// 区域端点还会列出 Gemma、Veo、AutoML 等非 Gemini 的模型花园条目；global 端点只含 Gemini。 +// fetchPage 负责发送请求并处理非 2xx 响应，返回该页的响应体；它返回的错误原样透出。 +func listVertexGeminiPublisherModels(location string, fetchPage func(page int, listURL string) ([]byte, error)) ([]string, error) { + models := make([]string, 0) + pageToken := "" + for page := 0; page < vertexPublisherModelsMaxPages; page++ { + listURL, err := buildVertexPublisherModelsURL(location, pageToken) + if err != nil { + return nil, newUpstreamModelSyncConfigError("Invalid Vertex location", err) + } + body, err := fetchPage(page, listURL) + if err != nil { + return nil, err + } + var response struct { + PublisherModels []struct { + Name string `json:"name"` + } `json:"publisherModels"` + NextPageToken string `json:"nextPageToken"` + } + if err := json.Unmarshal(body, &response); err != nil { + return nil, newUpstreamModelSyncUpstreamError("Upstream model list response was not valid JSON", err) + } + for _, entry := range response.PublisherModels { + name := strings.TrimSpace(entry.Name) + if !strings.HasPrefix(name, vertexGooglePublisherModelPrefix) { + continue + } + modelID := strings.TrimPrefix(name, vertexGooglePublisherModelPrefix) + if strings.HasPrefix(modelID, "gemini-") && len(modelID) > len("gemini-") { + models = append(models, modelID) + } + } + + pageToken = strings.TrimSpace(response.NextPageToken) + if pageToken == "" { + if len(models) == 0 { + return nil, newUpstreamModelSyncUpstreamError("Upstream returned no supported models", nil) + } + return dedupeAndSortModelIDs(models), nil + } + } + return nil, newUpstreamModelSyncUpstreamError( + "Upstream model list pagination did not finish", + fmt.Errorf("more than %d pages", vertexPublisherModelsMaxPages), + ) +} + func (s *AccountTestService) buildUpstreamModelsRequest(ctx context.Context, account *Account) (*http.Request, error) { switch { case account.Platform == PlatformAntigravity: @@ -1218,6 +1330,26 @@ func buildOpenAIModelsURL(base string) string { return buildOpenAIEndpointURL(base, "/v1/models") } +// buildVertexPublisherModelsURL 构造 Google 发布方模型列表地址。该接口只有 v1beta1 版本，pageSize 上限 300。 +func buildVertexPublisherModelsURL(location, pageToken string) (string, error) { + location = strings.TrimSpace(location) + if location == "" { + location = vertexDefaultLocation + } + if !vertexLocationPattern.MatchString(location) { + return "", fmt.Errorf("invalid vertex location: %s", location) + } + host := location + "-aiplatform.googleapis.com" + if location == "global" { + host = vertexGlobalHost + } + query := url.Values{"pageSize": []string{strconv.Itoa(vertexPublisherModelsPageSize)}} + if pageToken != "" { + query.Set("pageToken", pageToken) + } + return fmt.Sprintf("https://%s/v1beta1/publishers/google/models?%s", host, query.Encode()), nil +} + func buildGeminiModelsURL(base string) string { normalized := strings.TrimRight(strings.TrimSpace(base), "/") if strings.HasSuffix(normalized, "/v1beta/models") { diff --git a/backend/internal/service/upstream_models_vertex_test.go b/backend/internal/service/upstream_models_vertex_test.go new file mode 100644 index 0000000000..bd4678614e --- /dev/null +++ b/backend/internal/service/upstream_models_vertex_test.go @@ -0,0 +1,276 @@ +package service + +import ( + "context" + "errors" + "io" + "net/http" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +type vertexModelsTokenCacheStub struct { + token string + calls int +} + +func (c *vertexModelsTokenCacheStub) GetAccessToken(context.Context, string) (string, error) { + c.calls++ + if c.token == "" { + return "", errors.New("missing token") + } + return c.token, nil +} + +func (c *vertexModelsTokenCacheStub) SetAccessToken(context.Context, string, string, time.Duration) error { + return nil +} + +func (c *vertexModelsTokenCacheStub) DeleteAccessToken(context.Context, string) error { + return nil +} + +func (c *vertexModelsTokenCacheStub) AcquireRefreshLock(context.Context, string, time.Duration) (bool, error) { + return false, nil +} + +func (c *vertexModelsTokenCacheStub) ReleaseRefreshLock(context.Context, string) error { + return nil +} + +func vertexModelSyncTestAccount(credentials map[string]any) *Account { + merged := map[string]any{ + "service_account_json": map[string]any{ + "type": "service_account", + "project_id": "proj", + "private_key": "-----BEGIN PRIVATE KEY-----\nabc\n-----END PRIVATE KEY-----\n", + "client_email": "svc@proj.iam.gserviceaccount.com", + }, + } + for key, value := range credentials { + merged[key] = value + } + return &Account{ + ID: 173, + Platform: PlatformGemini, + Type: AccountTypeServiceAccount, + Credentials: merged, + } +} + +func vertexModelSyncTestService(upstream *httpUpstreamRecorder) *AccountTestService { + return &AccountTestService{ + httpUpstream: upstream, + cfg: upstreamModelSyncTestConfig(), + geminiTokenProvider: NewGeminiTokenProvider(nil, &vertexModelsTokenCacheStub{token: "ya29.vertex-token"}, nil), + } +} + +func vertexModelsJSONResponse(status int, body string) *http.Response { + return &http.Response{ + StatusCode: status, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(body)), + } +} + +func TestBuildVertexPublisherModelsURL(t *testing.T) { + t.Parallel() + + u, err := buildVertexPublisherModelsURL("global", "") + require.NoError(t, err) + require.Equal(t, "https://aiplatform.googleapis.com/v1beta1/publishers/google/models?pageSize=300", u) + + u, err = buildVertexPublisherModelsURL("us-central1", "next/page+token") + require.NoError(t, err) + require.Equal(t, "https://us-central1-aiplatform.googleapis.com/v1beta1/publishers/google/models?pageSize=300&pageToken=next%2Fpage%2Btoken", u) + + u, err = buildVertexPublisherModelsURL("", "") + require.NoError(t, err) + require.Equal(t, "https://us-central1-aiplatform.googleapis.com/v1beta1/publishers/google/models?pageSize=300", u) + + _, err = buildVertexPublisherModelsURL("evil.com/x", "") + require.Error(t, err) +} + +// Scenario: Vertex SA 账号按账号区域分页读取发布方模型列表，只保留 gemini- 前缀模型。 +func TestFetchUpstreamSupportedModelsVertexServiceAccountPaginatesAndKeepsGemini(t *testing.T) { + t.Parallel() + + upstream := &httpUpstreamRecorder{responses: []*http.Response{ + vertexModelsJSONResponse(http.StatusOK, `{"publisherModels":[ + {"name":"publishers/google/models/gemini-3.8-flash"}, + {"name":"publishers/google/models/gemma4"}, + {"name":"publishers/google/models/veo-3.1-generate-001"} + ],"nextPageToken":"page-2"}`), + vertexModelsJSONResponse(http.StatusOK, `{"publisherModels":[ + {"name":" publishers/google/models/gemini-2.5-flash "}, + {"name":"publishers/google/models/gemini-3.8-flash"}, + {"name":"gemini-without-publisher-prefix"}, + {"name":"publishers/anthropic/models/gemini-other-publisher"}, + {"name":"publishers/google/models/gemini-"} + ]}`), + }} + svc := vertexModelSyncTestService(upstream) + + models, err := svc.FetchUpstreamSupportedModels(context.Background(), vertexModelSyncTestAccount(map[string]any{ + "location": "europe-west4", + })) + require.NoError(t, err) + require.Equal(t, []string{"gemini-2.5-flash", "gemini-3.8-flash"}, models) + + require.Len(t, upstream.requests, 2) + require.Equal(t, "https://europe-west4-aiplatform.googleapis.com/v1beta1/publishers/google/models?pageSize=300", upstream.requests[0].URL.String()) + require.Equal(t, "https://europe-west4-aiplatform.googleapis.com/v1beta1/publishers/google/models?pageSize=300&pageToken=page-2", upstream.requests[1].URL.String()) + for _, req := range upstream.requests { + require.Equal(t, http.MethodGet, req.Method) + require.Equal(t, "Bearer ya29.vertex-token", req.Header.Get("Authorization")) + } +} + +func TestFetchUpstreamSupportedModelsVertexUpstreamErrorKeepsStatus(t *testing.T) { + t.Parallel() + + upstream := &httpUpstreamRecorder{resp: vertexModelsJSONResponse(http.StatusForbidden, `{"error":{"code":403,"status":"PERMISSION_DENIED"}}`)} + svc := vertexModelSyncTestService(upstream) + + _, err := svc.FetchUpstreamSupportedModels(context.Background(), vertexModelSyncTestAccount(map[string]any{"location": "global"})) + var syncErr *UpstreamModelSyncError + require.ErrorAs(t, err, &syncErr) + require.Equal(t, UpstreamModelSyncErrorUpstream, syncErr.Kind) + require.Equal(t, http.StatusForbidden, syncErr.StatusCode) + require.Equal(t, "https://aiplatform.googleapis.com/v1beta1/publishers/google/models?pageSize=300", upstream.lastReq.URL.String()) +} + +func TestFetchUpstreamSupportedModelsVertexWithoutGeminiModelsFails(t *testing.T) { + t.Parallel() + + upstream := &httpUpstreamRecorder{resp: vertexModelsJSONResponse(http.StatusOK, `{"publisherModels":[{"name":"publishers/google/models/gemma4"}]}`)} + svc := vertexModelSyncTestService(upstream) + + _, err := svc.FetchUpstreamSupportedModels(context.Background(), vertexModelSyncTestAccount(nil)) + var syncErr *UpstreamModelSyncError + require.ErrorAs(t, err, &syncErr) + require.Equal(t, UpstreamModelSyncErrorUpstream, syncErr.Kind) + require.Equal(t, "Upstream returned no supported models", syncErr.SafeMessage()) +} + +func TestFetchUpstreamSupportedModelsVertexStopsEndlessPagination(t *testing.T) { + t.Parallel() + + responses := make([]*http.Response, 0, vertexPublisherModelsMaxPages) + for i := 0; i < vertexPublisherModelsMaxPages; i++ { + responses = append(responses, vertexModelsJSONResponse(http.StatusOK, `{"publisherModels":[{"name":"publishers/google/models/gemini-3.8-flash"}],"nextPageToken":"again"}`)) + } + upstream := &httpUpstreamRecorder{responses: responses} + svc := vertexModelSyncTestService(upstream) + + _, err := svc.FetchUpstreamSupportedModels(context.Background(), vertexModelSyncTestAccount(nil)) + var syncErr *UpstreamModelSyncError + require.ErrorAs(t, err, &syncErr) + require.Equal(t, "Upstream model list pagination did not finish", syncErr.SafeMessage()) + require.Len(t, upstream.requests, vertexPublisherModelsMaxPages) +} + +func TestFetchUpstreamSupportedModelsVertexAPIKeyUnsupported(t *testing.T) { + t.Parallel() + + upstream := &httpUpstreamRecorder{} + svc := vertexModelSyncTestService(upstream) + account := &Account{ + ID: 174, + Platform: PlatformGemini, + Type: AccountTypeServiceAccount, + Credentials: map[string]any{ + "auth_mode": "apikey", + "api_key": "AQ.test-key", + }, + } + + _, err := svc.FetchUpstreamSupportedModels(context.Background(), account) + var syncErr *UpstreamModelSyncError + require.ErrorAs(t, err, &syncErr) + require.Equal(t, UpstreamModelSyncErrorUnsupported, syncErr.Kind) + require.Empty(t, upstream.requests) +} + +func TestFetchUpstreamSupportedModelsVertexInvalidCredentialsIsConfigError(t *testing.T) { + t.Parallel() + + cache := &vertexModelsTokenCacheStub{token: "ya29.vertex-token"} + upstream := &httpUpstreamRecorder{} + svc := &AccountTestService{ + httpUpstream: upstream, + cfg: upstreamModelSyncTestConfig(), + geminiTokenProvider: NewGeminiTokenProvider(nil, cache, nil), + } + account := &Account{ + ID: 175, + Platform: PlatformGemini, + Type: AccountTypeServiceAccount, + Credentials: map[string]any{"api_key": "not-a-service-account"}, + } + + _, err := svc.FetchUpstreamSupportedModels(context.Background(), account) + var syncErr *UpstreamModelSyncError + require.ErrorAs(t, err, &syncErr) + require.Equal(t, UpstreamModelSyncErrorConfiguration, syncErr.Kind) + require.Equal(t, "Invalid Vertex service account credentials", syncErr.SafeMessage()) + require.Zero(t, cache.calls) + require.Empty(t, upstream.requests) +} + +func TestFetchUpstreamSupportedModelsVertexInvalidLocationSkipsTokenExchange(t *testing.T) { + t.Parallel() + + cache := &vertexModelsTokenCacheStub{token: "ya29.vertex-token"} + upstream := &httpUpstreamRecorder{} + svc := &AccountTestService{ + httpUpstream: upstream, + cfg: upstreamModelSyncTestConfig(), + geminiTokenProvider: NewGeminiTokenProvider(nil, cache, nil), + } + + _, err := svc.FetchUpstreamSupportedModels(context.Background(), vertexModelSyncTestAccount(map[string]any{ + "location": "us-central1.evil.com", + })) + var syncErr *UpstreamModelSyncError + require.ErrorAs(t, err, &syncErr) + require.Equal(t, UpstreamModelSyncErrorConfiguration, syncErr.Kind) + require.Equal(t, "Invalid Vertex location", syncErr.SafeMessage()) + require.Zero(t, cache.calls) + require.Empty(t, upstream.requests) +} + +// Scenario: 翻页途中的 404 不能被当成列表端点不可用而回退到配置模型。 +func TestSyncUpstreamModelCatalogVertexMidPagination404DoesNotFallback(t *testing.T) { + t.Parallel() + + upstream := &httpUpstreamRecorder{responses: []*http.Response{ + vertexModelsJSONResponse(http.StatusOK, `{"publisherModels":[{"name":"publishers/google/models/gemini-a"}],"nextPageToken":"page-2"}`), + vertexModelsJSONResponse(http.StatusNotFound, `{"error":{"code":404}}`), + }} + svc := vertexModelSyncTestService(upstream) + + _, err := svc.SyncUpstreamModelCatalog(context.Background(), vertexModelSyncTestAccount(map[string]any{ + "model_mapping": map[string]any{"x": "gemini-mapped"}, + })) + var syncErr *UpstreamModelSyncError + require.ErrorAs(t, err, &syncErr) + require.Equal(t, UpstreamModelSyncErrorUpstream, syncErr.Kind) + require.Zero(t, syncErr.StatusCode) + require.Len(t, upstream.requests, 2) +} + +func TestFetchUpstreamSupportedModelsVertexFirstPage404KeepsStatus(t *testing.T) { + t.Parallel() + + upstream := &httpUpstreamRecorder{resp: vertexModelsJSONResponse(http.StatusNotFound, `{"error":{"code":404}}`)} + svc := vertexModelSyncTestService(upstream) + + _, err := svc.FetchUpstreamSupportedModels(context.Background(), vertexModelSyncTestAccount(nil)) + require.True(t, upstreamModelListEndpointUnsupported(err)) +} diff --git a/backend/internal/service/vertex_service_account.go b/backend/internal/service/vertex_service_account.go index 7ccbeee43c..2a19a50e21 100644 --- a/backend/internal/service/vertex_service_account.go +++ b/backend/internal/service/vertex_service_account.go @@ -29,6 +29,8 @@ const ( vertexServiceAccountCacheSkew = 5 * time.Minute vertexLockWaitTime = 200 * time.Millisecond vertexAnthropicVersion = "vertex-2023-10-16" + vertexGlobalHost = "aiplatform.googleapis.com" + vertexAuthModeAPIKey = "apikey" ) var ( @@ -58,6 +60,13 @@ func (a *Account) IsVertexServiceAccount() bool { return a != nil && a.Type == AccountTypeServiceAccount } +// IsVertexAPIKey 判断 Gemini 平台的 Vertex 账号是否使用 API Key 认证（credentials.auth_mode=apikey）。 +// Anthropic 平台的 Vertex 账号只支持 Service Account JSON。 +func (a *Account) IsVertexAPIKey() bool { + return a.IsVertexServiceAccount() && a.Platform == PlatformGemini && + strings.TrimSpace(a.GetCredential("auth_mode")) == vertexAuthModeAPIKey +} + func (a *Account) VertexProjectID() string { if a == nil { return "" @@ -283,6 +292,63 @@ func exchangeVertexServiceAccountToken(ctx context.Context, key *vertexServiceAc return parsed.AccessToken, ttl, nil } +// buildVertexGeminiAccountURL 按账号构造 Vertex Gemini 请求地址：未配置 project_id 的 API Key 账号 +// 走 Express mode 全局端点，其余账号走项目级端点。 +func buildVertexGeminiAccountURL(account *Account, model, action string, stream bool) (string, error) { + if account.IsVertexAPIKey() && account.VertexProjectID() == "" { + return buildVertexGeminiExpressURL(model, action, stream) + } + return buildVertexGeminiURL(account.VertexProjectID(), account.VertexLocation(model), model, action, stream) +} + +// buildVertexGeminiExpressURL 构造不带项目与区域的全局端点地址，仅 API Key 认证可用。 +func buildVertexGeminiExpressURL(model, action string, stream bool) (string, error) { + model = strings.TrimSpace(model) + action = strings.TrimSpace(action) + if model == "" { + return "", errors.New("vertex model is required") + } + if err := validateVertexGeminiAction(action); err != nil { + return "", err + } + u := fmt.Sprintf("https://%s/v1/publishers/google/models/%s:%s", vertexGlobalHost, url.PathEscape(model), action) + if stream { + u += "?alt=sse" + } + return u, nil +} + +func validateVertexGeminiAction(action string) error { + switch action { + case "generateContent", "streamGenerateContent", "countTokens": + return nil + default: + return fmt.Errorf("unsupported vertex gemini action: %s", action) + } +} + +// setVertexGeminiAuth 给 Vertex Gemini 请求设置认证头：API Key 账号用 x-goog-api-key， +// Service Account JSON 账号用换取的 Bearer access token。 +func setVertexGeminiAuth(ctx context.Context, req *http.Request, account *Account, tokenProvider *GeminiTokenProvider) error { + if account.IsVertexAPIKey() { + apiKey := strings.TrimSpace(account.GetCredential("api_key")) + if apiKey == "" { + return errors.New("vertex api_key not configured") + } + req.Header.Set("x-goog-api-key", apiKey) + return nil + } + if tokenProvider == nil { + return errors.New("gemini token provider not configured") + } + accessToken, err := tokenProvider.GetAccessToken(ctx, account) + if err != nil { + return err + } + req.Header.Set("Authorization", "Bearer "+accessToken) + return nil +} + func buildVertexGeminiURL(projectID, location, model, action string, stream bool) (string, error) { projectID = strings.TrimSpace(projectID) location = strings.TrimSpace(location) @@ -300,14 +366,12 @@ func buildVertexGeminiURL(projectID, location, model, action string, stream bool if model == "" { return "", errors.New("vertex model is required") } - switch action { - case "generateContent", "streamGenerateContent", "countTokens": - default: - return "", fmt.Errorf("unsupported vertex gemini action: %s", action) + if err := validateVertexGeminiAction(action); err != nil { + return "", err } host := fmt.Sprintf("%s-aiplatform.googleapis.com", location) if location == "global" { - host = "aiplatform.googleapis.com" + host = vertexGlobalHost } u := fmt.Sprintf( "https://%s/v1/projects/%s/locations/%s/publishers/google/models/%s:%s", diff --git a/backend/internal/service/vertex_service_account_test.go b/backend/internal/service/vertex_service_account_test.go index 68c756eaa2..77cb6617c7 100644 --- a/backend/internal/service/vertex_service_account_test.go +++ b/backend/internal/service/vertex_service_account_test.go @@ -30,6 +30,62 @@ func TestBuildVertexGeminiURLUsesGlobalEndpointHost(t *testing.T) { require.Equal(t, "https://aiplatform.googleapis.com/v1/projects/my-project/locations/global/publishers/google/models/gemini-3-flash-preview:streamGenerateContent?alt=sse", got) } +func TestIsVertexAPIKey(t *testing.T) { + apiKeyCreds := map[string]any{"auth_mode": "apikey", "api_key": "k"} + require.True(t, (&Account{Platform: PlatformGemini, Type: AccountTypeServiceAccount, Credentials: apiKeyCreds}).IsVertexAPIKey()) + require.False(t, (&Account{Platform: PlatformGemini, Type: AccountTypeServiceAccount, Credentials: map[string]any{"service_account_json": "{}"}}).IsVertexAPIKey()) + require.False(t, (&Account{Platform: PlatformAnthropic, Type: AccountTypeServiceAccount, Credentials: apiKeyCreds}).IsVertexAPIKey()) + require.False(t, (&Account{Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: apiKeyCreds}).IsVertexAPIKey()) +} + +func TestBuildVertexGeminiAccountURL(t *testing.T) { + cases := []struct { + name string + credentials map[string]any + want string + }{ + { + name: "api key without project uses express global endpoint", + credentials: map[string]any{"auth_mode": "apikey", "api_key": "k", "location": "us-central1"}, + want: "https://aiplatform.googleapis.com/v1/publishers/google/models/gemini-3-flash-preview:streamGenerateContent?alt=sse", + }, + { + name: "api key with project uses project endpoint", + credentials: map[string]any{"auth_mode": "apikey", "api_key": "k", "project_id": "my-project", "location": "us-central1"}, + want: "https://us-central1-aiplatform.googleapis.com/v1/projects/my-project/locations/us-central1/publishers/google/models/gemini-3-flash-preview:streamGenerateContent?alt=sse", + }, + { + name: "service account uses project endpoint", + credentials: map[string]any{"project_id": "sa-project", "location": "global"}, + want: "https://aiplatform.googleapis.com/v1/projects/sa-project/locations/global/publishers/google/models/gemini-3-flash-preview:streamGenerateContent?alt=sse", + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + account := &Account{Platform: PlatformGemini, Type: AccountTypeServiceAccount, Credentials: tc.credentials} + got, err := buildVertexGeminiAccountURL(account, "gemini-3-flash-preview", "streamGenerateContent", true) + require.NoError(t, err) + require.Equal(t, tc.want, got) + }) + } + + account := &Account{Platform: PlatformGemini, Type: AccountTypeServiceAccount, Credentials: map[string]any{"auth_mode": "apikey", "api_key": "k"}} + _, err := buildVertexGeminiAccountURL(account, "gemini-3-flash-preview", "rawPredict", false) + require.ErrorContains(t, err, "unsupported vertex gemini action") +} + +func TestSetVertexGeminiAuthUsesAPIKeyHeader(t *testing.T) { + account := &Account{Platform: PlatformGemini, Type: AccountTypeServiceAccount, Credentials: map[string]any{"auth_mode": "apikey", "api_key": " vertex-key "}} + req := httptest.NewRequest(http.MethodPost, "https://aiplatform.googleapis.com/v1/publishers/google/models/gemini-3-flash-preview:generateContent", nil) + + require.NoError(t, setVertexGeminiAuth(context.Background(), req, account, nil)) + require.Equal(t, "vertex-key", req.Header.Get("x-goog-api-key")) + require.Empty(t, req.Header.Get("Authorization")) + + account.Credentials["api_key"] = "" + require.ErrorContains(t, setVertexGeminiAuth(context.Background(), req, account, nil), "vertex api_key not configured") +} + func TestBuildVertexAnthropicURL(t *testing.T) { got, err := buildVertexAnthropicURL("my-project", "us-east5", "claude-sonnet-4-5@20250929", false) require.NoError(t, err) diff --git a/backend/internal/service/wire.go b/backend/internal/service/wire.go index ec3058dc66..688d729b31 100644 --- a/backend/internal/service/wire.go +++ b/backend/internal/service/wire.go @@ -640,6 +640,13 @@ func ProvideIdempotencyCleanupService(repo IdempotencyRepository, cfg *config.Co return svc } +// ProvideGeminiCachedContentService creates GeminiCachedContentService and starts its mutation recovery loop. +func ProvideGeminiCachedContentService(repo GeminiCachedContentRepository, compat *GeminiMessagesCompatService, cfg *config.Config) *GeminiCachedContentService { + svc := NewGeminiCachedContentService(repo, compat, cfg) + svc.Start() + return svc +} + // ProvideScheduledTestService creates ScheduledTestService. func ProvideScheduledTestService( planRepo ScheduledTestPlanRepository, @@ -885,6 +892,7 @@ var ProviderSet = wire.NewSet( ProvideGrokOAuthService, wire.Bind(new(GrokOAuthTokenService), new(*GrokOAuthService)), NewGeminiOAuthService, + ProvideGeminiCachedContentService, NewGeminiQuotaService, NewCompositeTokenCacheInvalidator, wire.Bind(new(TokenCacheInvalidator), new(*CompositeTokenCacheInvalidator)), diff --git a/backend/internal/util/logredact/redact.go b/backend/internal/util/logredact/redact.go index 9249b761c7..8be6448a18 100644 --- a/backend/internal/util/logredact/redact.go +++ b/backend/internal/util/logredact/redact.go @@ -42,6 +42,8 @@ type textRedactPatterns struct { var ( reGOCSPX = regexp.MustCompile(`GOCSPX-[0-9A-Za-z_-]{24,}`) reAIza = regexp.MustCompile(`AIza[0-9A-Za-z_-]{35}`) + // Gemini auth key：AQ.Ab 前缀 + 48 位，长度放宽以覆盖格式变体。 + reGeminiAuthKey = regexp.MustCompile(`AQ\.Ab[0-9A-Za-z_-]{40,}`) defaultTextRedactPatterns = compileTextRedactPatterns(nil) extraTextPatternCache sync.Map // map[string]*textRedactPatterns @@ -99,6 +101,7 @@ func RedactText(input string, extraKeys ...string) string { out := input out = reGOCSPX.ReplaceAllString(out, "GOCSPX-***") out = reAIza.ReplaceAllString(out, "AIza***") + out = reGeminiAuthKey.ReplaceAllString(out, "AQ.***") out = patterns.reJSONLike.ReplaceAllString(out, `$1***$3`) out = patterns.reQueryLike.ReplaceAllString(out, `$1=***`) out = patterns.rePlain.ReplaceAllString(out, `$1$2***`) diff --git a/backend/internal/util/logredact/redact_test.go b/backend/internal/util/logredact/redact_test.go index 266db69dbd..e615f33e0e 100644 --- a/backend/internal/util/logredact/redact_test.go +++ b/backend/internal/util/logredact/redact_test.go @@ -38,6 +38,28 @@ func TestRedactText_GOCSPX(t *testing.T) { } } +func TestRedactText_GeminiAPIKeys(t *testing.T) { + legacyKey := "AI" + "za" + strings.Repeat("Ab1_-", 7) + authKey := "AQ." + "Ab" + strings.Repeat("Xy9_-z", 8) + in := "upstream rejected " + legacyKey + " and " + authKey + " for model gemini-3.1-pro" + out := RedactText(in) + if strings.Contains(out, legacyKey) || strings.Contains(out, authKey) { + t.Fatalf("expected Gemini keys redacted, got %q", out) + } + for _, want := range []string{"AIza***", "AQ.***", "for model gemini-3.1-pro"} { + if !strings.Contains(out, want) { + t.Fatalf("expected %q in %q", want, out) + } + } +} + +func TestRedactText_GeminiAuthKeyPrefixAloneUnchanged(t *testing.T) { + in := "key format AQ.Ab... is documented" + if out := RedactText(in); out != in { + t.Fatalf("expected short AQ. prefix unchanged, got %q", out) + } +} + func TestRedactText_ExtraKeyCacheUsesNormalizedSortedKey(t *testing.T) { clearExtraTextPatternCache() diff --git a/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/backend/migrations/244_gemini_cached_contents_channel_usage.sql b/backend/migrations/244_gemini_cached_contents_channel_usage.sql new file mode 100644 index 0000000000..c72f097a73 --- /dev/null +++ b/backend/migrations/244_gemini_cached_contents_channel_usage.sql @@ -0,0 +1,4 @@ +-- Gemini 显式缓存创建时的渠道计费口径（渠道 ID / 请求模型 / 计费模型来源）， +-- 延长有效期按同一口径计费。 +ALTER TABLE gemini_cached_contents + ADD COLUMN IF NOT EXISTS channel_usage JSONB NOT NULL DEFAULT '{}'::jsonb; diff --git a/backend/migrations/245_gemini_cached_contents_mutation_lease.sql b/backend/migrations/245_gemini_cached_contents_mutation_lease.sql new file mode 100644 index 0000000000..e0970efebf --- /dev/null +++ b/backend/migrations/245_gemini_cached_contents_mutation_lease.sql @@ -0,0 +1,10 @@ +-- Gemini 显式缓存修改有效期的进行中操作（修改意图）：mutation_token 标识持有者， +-- mutation_lease_until 之前同一缓存的其他修改被拒绝；expire_time 在提交前保持修改前的值， +-- 租约过期仍未清除的意图由恢复流程把上游有效期回滚到 expire_time。 +ALTER TABLE gemini_cached_contents + ADD COLUMN IF NOT EXISTS mutation_token VARCHAR(64) NULL, + ADD COLUMN IF NOT EXISTS mutation_lease_until TIMESTAMPTZ NULL; + +CREATE INDEX IF NOT EXISTS idx_gemini_cached_contents_pending_mutation + ON gemini_cached_contents (mutation_lease_until) + WHERE mutation_token IS NOT NULL AND deleted_at IS NULL; diff --git a/backend/resources/model-pricing/model_prices_and_context_window.json b/backend/resources/model-pricing/model_prices_and_context_window.json index 174366de04..b11a75f688 100644 --- a/backend/resources/model-pricing/model_prices_and_context_window.json +++ b/backend/resources/model-pricing/model_prices_and_context_window.json @@ -2650,17 +2650,17 @@ "web_search_billing_unit": "per_query" }, "gemini-3.6-flash": { - "cache_read_input_token_cost": 1.5e-07, - "cache_read_input_token_cost_priority": 2.7e-07, - "input_cost_per_token": 1.5e-06, - "input_cost_per_token_priority": 2.7e-06, + "cache_read_input_token_cost": 7.5e-08, + "cache_read_input_token_cost_priority": 1.35e-07, + "input_cost_per_token": 7.5e-07, + "input_cost_per_token_priority": 1.35e-06, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 1048576, "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", - "output_cost_per_token": 7.5e-06, - "output_cost_per_token_priority": 1.35e-05, + "output_cost_per_token": 3.75e-06, + "output_cost_per_token_priority": 6.75e-06, "source": "https://ai.google.dev/gemini-api/docs/latest-model", "supports_prompt_caching": true, "supports_reasoning": true, diff --git a/deploy/config.example.yaml b/deploy/config.example.yaml index 7066f17208..6b8b049838 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,14 @@ pricing: # 覆盖补丁文件（可选，优先级最高）：条目按字段浅合并覆盖目录/回退数据， # 字段值为 null 表示删除该字段。模型名须与目录键完全一致。 # 修改后在下一次哈希检查时按内容哈希自动重载，无需重启。 + # Gemini explicit context caching storage price (USD per token per hour) comes + # from the field cache_storage_cost_per_token_per_hour here, e.g. + # {"gemini-3.8-flash": {"cache_storage_cost_per_token_per_hour": 0.0000005}}; + # without it, built-in official prices apply (released models first, then the + # Pro / Flash series rate). Models with no price at all reject POST /v1beta/cachedContents. + # Gemini 显式缓存存储单价（每 token 每小时，USD）优先取这里的 cache_storage_cost_per_token_per_hour 字段， + # 未配置时使用内置官方价（先按已发布型号，再按 Pro / Flash 系列兜底）；两者都没有的模型拒绝创建缓存。 + # 目录中不存在的模型条目还须同时提供 input_cost_per_token / output_cost_per_token。 # override_file: "./data/model_pricing_overrides.json" # Update interval in hours # 更新间隔（小时） diff --git a/frontend/src/components/account/AccountUsageCell.vue b/frontend/src/components/account/AccountUsageCell.vue index fd5b9177bc..6bf5f12af5 100644 --- a/frontend/src/components/account/AccountUsageCell.vue +++ b/frontend/src/components/account/AccountUsageCell.vue @@ -489,6 +489,7 @@ -
+
{{ t('admin.accounts.gemini.rateLimit.unlimited') }}
@@ -594,7 +595,7 @@
- +
void) | null = null let visibilityObserver: IntersectionObserver | null = null +// Gemini 第三方上游 API Key：没有 AI Studio 等级与本地模拟配额 +const isGeminiThirdPartyApiKey = computed(() => { + if (props.account.platform !== 'gemini' || props.account.type !== 'apikey') return false + const creds = props.account.credentials as GeminiCredentials | undefined + return !isOfficialGeminiBaseUrl(creds?.base_url) +}) + +const isGeminiVertex = computed(() => { + return props.account.platform === 'gemini' && props.account.type === 'service_account' +}) + // Show usage windows for OAuth and Setup Token accounts const showUsageWindows = computed(() => { // Gemini: we can always compute local usage windows from DB logs (simulated quotas). - if (props.account.platform === 'gemini') return true + // 第三方上游 API Key 没有模拟配额，按普通 Key 账号展示今日统计与额度进度条。 + if (props.account.platform === 'gemini') return !isGeminiThirdPartyApiKey.value // 多协议 API Key 供应商：apikey 账号也有滚动用量窗口（coding plan / 订阅）或余额， // 由 CNProviderQuotaCell / CNProviderBalanceCell 自行探测与展示。 if (isMultiProtocolApiKeyPlatform(props.account.platform)) { @@ -770,7 +784,7 @@ const shouldFetchUsage = computed(() => { return props.account.type === 'oauth' || props.account.type === 'setup-token' } if (props.account.platform === 'gemini') { - return true + return !isGeminiThirdPartyApiKey.value && !isGeminiVertex.value } if (props.account.platform === 'antigravity') { return props.account.type === 'oauth' @@ -791,8 +805,9 @@ const cnBalanceCellVisible = computed(() => cnBalanceCellVisibleFn(props.account const isBatchManaged = computed(() => typeof props.requestBatchedUsage === 'function') +// Vertex 没有本地模拟配额，用量窗口展示今日统计 const showGeminiTodayStats = computed(() => { - return props.account.platform === 'gemini' && props.account.type === 'service_account' + return isGeminiVertex.value }) const geminiUsageAvailable = computed(() => { @@ -956,6 +971,10 @@ const geminiOAuthType = computed(() => { return (creds?.oauth_type || '').trim() || null }) +const showGeminiQuotaPolicy = computed(() => { + return !isGeminiVertex.value +}) + // Gemini 是否为 Code Assist OAuth const isGeminiCodeAssist = computed(() => { if (props.account.platform !== 'gemini') return false @@ -963,9 +982,10 @@ const isGeminiCodeAssist = computed(() => { return creds?.oauth_type === 'code_assist' || (!creds?.oauth_type && !!creds?.project_id) }) -const geminiChannelShort = computed((): 'ai studio' | 'gcp' | 'google one' | 'client' | null => { +const geminiChannelShort = computed((): 'ai studio' | 'vertex' | 'gcp' | 'google one' | 'client' | null => { if (props.account.platform !== 'gemini') return null + if (isGeminiVertex.value) return 'vertex' // API Key accounts are AI Studio. if (props.account.type === 'apikey') return 'ai studio' @@ -979,6 +999,7 @@ const geminiChannelShort = computed((): 'ai studio' | 'gcp' | 'google one' | 'cl const geminiUserLevel = computed((): string | null => { if (props.account.platform !== 'gemini') return null + if (isGeminiVertex.value) return null const tier = (geminiTier.value || '').toString().trim() const tierLower = tier.toLowerCase() @@ -1036,7 +1057,7 @@ const geminiTierClass = computed(() => { const channel = geminiChannelShort.value const level = geminiUserLevel.value - if (channel === 'client' || channel === 'ai studio') { + if (channel === 'client' || channel === 'ai studio' || channel === 'vertex') { return 'bg-blue-100 text-blue-600 dark:bg-blue-900/40 dark:text-blue-300' } diff --git a/frontend/src/components/account/CreateAccountModal.vue b/frontend/src/components/account/CreateAccountModal.vue index 2ea44ee11c..315feae303 100644 --- a/frontend/src/components/account/CreateAccountModal.vue +++ b/frontend/src/components/account/CreateAccountModal.vue @@ -381,8 +381,8 @@
- Vertex - Service Account + {{ t('admin.accounts.vertexLabel') }} + {{ t('admin.accounts.vertexDesc') }}
@@ -830,17 +830,18 @@
- Vertex + {{ t('admin.accounts.vertexLabel') }} - Service Account + {{ t('admin.accounts.vertexGeminiDesc') }}

{{ t('admin.accounts.gemini.accountType.apiKeyNote') }}

@@ -1068,7 +1069,7 @@
-
+
-
-
-
-
- - {{ vertexClientEmail ? t('admin.accounts.vertexSaJsonLoaded') : t('admin.accounts.vertexSaJsonDrop') }} -
-

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

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

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

-
+ + +
@@ -1502,9 +1580,9 @@
-
+
- @@ -1767,6 +1845,48 @@
+ +
+
+
+ +

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

+
+ +
+
+

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

+
+
+ + +

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

+
+
+
@@ -4005,6 +4125,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' @@ -4132,7 +4253,7 @@ const apiKeyValuePlaceholder = computed(() => { case 'openai': return 'sk-proj-...' case 'gemini': - return 'AIza...' + return 'AQ.Ab...' case 'grok': return 'xai-...' case 'kimi': @@ -4473,6 +4594,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('') @@ -4606,6 +4729,17 @@ const vertexProjectId = ref('') const vertexClientEmail = ref('') const vertexLocation = ref('global') const vertexServiceAccountDragActive = ref(false) +const vertexAuthMode = ref<'service_account' | 'apikey'>('service_account') +const vertexApiKey = ref('') +// API Key 认证只对 Gemini 平台开放，Anthropic 平台的 Vertex 账号始终使用 Service Account JSON。 +const vertexUsesApiKey = computed(() => form.platform === 'gemini' && vertexAuthMode.value === 'apikey') +// 两种认证方式的 Project ID 来源不同（JSON 读取 / 手填），切换时清空凭证字段避免串用。 +watch(vertexAuthMode, () => { + vertexServiceAccountJson.value = '' + vertexProjectId.value = '' + vertexClientEmail.value = '' + vertexApiKey.value = '' +}) const tempUnschedEnabled = ref(false) const tempUnschedRules = ref([]) const getModelMappingKey = createStableObjectKeyResolver('create-model-mapping') @@ -5010,6 +5144,8 @@ watch( vertexProjectId.value = '' vertexClientEmail.value = '' vertexLocation.value = 'global' + vertexAuthMode.value = 'service_account' + vertexApiKey.value = '' // Reset Anthropic/Antigravity-specific settings when switching to other platforms if (newPlatform !== 'anthropic' && newPlatform !== 'antigravity') { interceptWarmupRequests.value = false @@ -5458,6 +5594,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 @@ -5511,6 +5649,8 @@ const resetForm = () => { vertexProjectId.value = '' vertexClientEmail.value = '' vertexLocation.value = 'global' + vertexAuthMode.value = 'service_account' + vertexApiKey.value = '' tempUnschedEnabled.value = false tempUnschedRules.value = [] geminiOAuthType.value = 'code_assist' @@ -5859,6 +5999,28 @@ const handleSubmit = async () => { appStore.showError(t('admin.accounts.pleaseEnterAccountName')) return } + if (vertexUsesApiKey.value) { + if (!vertexApiKey.value.trim()) { + appStore.showError(t('admin.accounts.vertexApiKeyRequired')) + return + } + const projectId = vertexProjectId.value.trim() + if (projectId && !vertexLocation.value.trim()) { + appStore.showError(t('admin.accounts.vertexLocationRequired')) + return + } + const credentials: Record = { + auth_mode: 'apikey', + api_key: vertexApiKey.value.trim(), + location: vertexLocation.value.trim() || 'global', + tier_id: 'vertex' + } + if (projectId) { + credentials.project_id = projectId + } + await createAccountAndFinish(form.platform, 'service_account' as AccountType, credentials) + return + } if (!parseVertexServiceAccountJson()) { return } @@ -5900,7 +6062,7 @@ const handleSubmit = async () => { base_url: apiKeyBaseUrl.value.trim() || defaultBaseUrl, api_key: apiKeyValue.value.trim() } - if (form.platform === 'gemini') { + if (form.platform === 'gemini' && isOfficialGeminiBaseUrl(apiKeyBaseUrl.value)) { credentials.tier_id = geminiTierAIStudio.value } @@ -5963,6 +6125,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..04b1c9f68f 100644 --- a/frontend/src/components/account/EditAccountModal.vue +++ b/frontend/src/components/account/EditAccountModal.vue @@ -239,7 +239,7 @@ account.platform === 'openai' ? 'sk-proj-...' : account.platform === 'gemini' - ? 'AIza...' + ? 'AQ.Ab...' : account.platform === 'antigravity' ? 'sk-...' : account.platform === 'grok' @@ -500,6 +500,48 @@
+ +
+
+
+ +

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

+
+ +
+
+

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

+
+
+ + +

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

+
+
+
@@ -915,8 +957,30 @@
+
+ + +

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

+
-
+
+ + +

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

+
+

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

-
+