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