fix(gemini): 显式缓存错误透传收窄到上游网关账号,延长沿用创建时计费口径

- 绑定缓存的请求与缓存管理接口收到 429 / 5xx 时,只有上游网关账号的响应原样回写
  (UpstreamFailoverError.BoundUpstreamPassthrough / GeminiCachedContentUpstreamError.UpstreamGateway),
  官方账号仍走错误透传规则与通用映射,不向客户端暴露配额明细
- 记录新增 channel_usage(迁移 244)保存创建时的渠道计费口径,PATCH 延长按同一口径预检、报价与记账
- 创建 / 延长费用行的 request_id 改为固定前缀 gcache:{create|patch}:{id},不被客户端或网关请求 id 覆盖,
  按缓存 ID 即可检索一个缓存的全部费用行
- 调度快照凭据投影补齐 base_url / auth_mode / 上游显式缓存声明,显式缓存候选过滤一次算全,
  不再逐个选号后剔除
This commit is contained in:
feeeei
2026-10-11 01:37:04 +08:00
parent c9e4b267b6
commit 05d7946af8
15 changed files with 218 additions and 45 deletions
@@ -25,9 +25,10 @@ import (
)
const (
cachedBillingInputPrice = 0.00000125
cachedBillingStoragePrice = 0.0000045
cachedBillingTokens = 1_000_000
cachedBillingInputPrice = 0.00000125
cachedBillingStoragePrice = 0.0000045
cachedBillingAliasStoragePrice = 0.000009
cachedBillingTokens = 1_000_000
)
type cachedBillingRepo struct {
@@ -175,6 +176,12 @@ func (r *cachedBillingUsageRepo) Create(_ context.Context, log *service.UsageLog
return true, nil
}
func (r *cachedBillingUsageRepo) first() *service.UsageLog {
r.mu.Lock()
defer r.mu.Unlock()
return r.logs[0]
}
func (r *cachedBillingUsageRepo) totals() []float64 {
r.mu.Lock()
defer r.mu.Unlock()
@@ -192,6 +199,8 @@ type cachedBillingEnv struct {
usage *cachedBillingUsageRepo
apiKey *service.APIKey
group *service.Group
// accounts 是调度快照桩直接返回的账号对象,测试可改其凭据。
accounts map[int64]*service.Account
}
func newCachedBillingEnv(t *testing.T, balance float64, accountIDs ...int64) *cachedBillingEnv {
@@ -199,7 +208,10 @@ func newCachedBillingEnv(t *testing.T, balance float64, accountIDs ...int64) *ca
dir := t.TempDir()
require.NoError(t, os.WriteFile(filepath.Join(dir, "model_pricing.json"), []byte(fmt.Sprintf(`{"gemini-2.5-pro":{
"litellm_provider":"gemini","mode":"chat","input_cost_per_token":%g,"output_cost_per_token":0.00001,
"cache_storage_cost_per_token_per_hour":%g}}`, cachedBillingInputPrice, cachedBillingStoragePrice)), 0o644))
"cache_storage_cost_per_token_per_hour":%g},"gemini-alias":{
"litellm_provider":"gemini","mode":"chat","input_cost_per_token":%g,"output_cost_per_token":0.00001,
"cache_storage_cost_per_token_per_hour":%g}}`,
cachedBillingInputPrice, cachedBillingStoragePrice, cachedBillingInputPrice, cachedBillingAliasStoragePrice)), 0o644))
cfg := &config.Config{}
cfg.Pricing.DataDir = dir
cfg.Pricing.UpdateIntervalHours = 100000
@@ -237,7 +249,25 @@ func newCachedBillingEnv(t *testing.T, balance float64, accountIDs ...int64) *ca
}
apiKey := &service.APIKey{ID: 7, UserID: 3, GroupID: &groupID, Group: group, Status: service.StatusActive,
User: &service.User{ID: 3, Concurrency: 10, Balance: balance}}
return &cachedBillingEnv{h: h, repo: repo, upstream: upstream, usage: usage, apiKey: apiKey, group: group}
return &cachedBillingEnv{h: h, repo: repo, upstream: upstream, usage: usage, apiKey: apiKey, group: group, accounts: accounts}
}
// TestGeminiExplicitCacheExclusions_FiltersUnsupportedAccountsFromSnapshot 候选过滤直接按快照里的凭据判定:
// 第三方地址且未声明上游显式缓存的 Key 账号在选号前就被排除,官方地址与声明了上游显式缓存的账号保留。
func TestGeminiExplicitCacheExclusions_FiltersUnsupportedAccountsFromSnapshot(t *testing.T) {
env := newCachedBillingEnv(t, 500, 900, 901, 902)
env.accounts[901].Credentials["base_url"] = "https://third.example.com"
env.accounts[902].Credentials["base_url"] = "https://gw.example.com"
env.accounts[902].Credentials[service.GeminiExplicitCacheUpstreamCredentialKey] = true
groupID := env.group.ID
excluded, err := env.h.gatewayService.GeminiExplicitCacheExclusions(context.Background(), &groupID, 0)
require.NoError(t, err)
require.Equal(t, map[int64]struct{}{901: {}}, excluded)
excluded, err = env.h.gatewayService.GeminiExplicitCacheExclusions(context.Background(), &groupID, 902)
require.NoError(t, err)
require.Equal(t, map[int64]struct{}{900: {}, 901: {}}, excluded, "绑定缓存的请求只保留持有账号")
}
func (e *cachedBillingEnv) serve(method, target, body string, params gin.Params, call func(*gin.Context)) *httptest.ResponseRecorder {
@@ -279,6 +309,27 @@ func (e *cachedBillingEnv) patch() *httptest.ResponseRecorder {
gin.Params{{Key: "cacheID", Value: "abc"}}, e.h.GeminiCachedContentsPatch)
}
// TestGeminiCachedContentsPatch_UsesChannelUsageRecordedAtCreate 延长有效期沿用创建时的渠道计费口径:
// 计费模型来源为请求模型时按别名定价,用量记录带渠道 ID。
func TestGeminiCachedContentsPatch_UsesChannelUsageRecordedAtCreate(t *testing.T) {
env := newCachedBillingEnv(t, 500, 900)
env.seed(time.Now().Add(time.Hour))
env.repo.records["abc"].ChannelUsage = service.ChannelUsageFields{
ChannelID: 5, OriginalModel: "gemini-alias", ChannelMappedModel: "gemini-2.5-pro",
BillingModelSource: service.BillingModelSourceRequested,
}
rec := env.patch()
require.Equal(t, http.StatusOK, rec.Code, rec.Body.String())
totals := env.waitForUsage(t, 1)
require.Len(t, totals, 1)
require.InDelta(t, cachedBillingTokens*cachedBillingAliasStoragePrice*23, totals[0], 0.01, "按创建时的计费模型(请求模型别名)计存储费")
usageLog := env.usage.first()
require.NotNil(t, usageLog.ChannelID)
require.Equal(t, int64(5), *usageLog.ChannelID)
require.True(t, strings.HasPrefix(usageLog.RequestID, "gcache:patch:abc:"), usageLog.RequestID)
}
func cachedBillingCreateBody(expire time.Time) string {
return fmt.Sprintf(`{"name":"cachedContents/up1","usageMetadata":{"totalTokenCount":%d},"expireTime":%q}`,
cachedBillingTokens, expire.UTC().Format(time.RFC3339Nano))
@@ -296,6 +347,7 @@ func TestGeminiCachedContentsCreate_ChargesInputAndStorage(t *testing.T) {
require.Len(t, env.repo.records, 1)
for _, record := range env.repo.records {
require.False(t, record.ExpireTime.After(time.Now().Add(24*time.Hour)), "本地到期时间不超过请求的有效期")
require.Equal(t, "gcache:create:"+record.PublicID, env.usage.first().RequestID, "费用行按缓存 ID 可检索")
}
}
@@ -17,7 +17,6 @@ import (
"github.com/Wei-Shaw/sub2api/internal/server/middleware"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/google/uuid"
"go.uber.org/zap"
)
@@ -246,7 +245,7 @@ func (h *GatewayHandler) GeminiCachedContentsCreate(c *gin.Context) {
excluded[account.ID] = struct{}{}
continue
}
writeGeminiCachedContentUpstreamError(c, upErr)
h.writeGeminiCachedContentUpstreamError(c, upErr)
return
}
// 上游已返回成功但响应不可用:不换号重建,能识别出资源名时删除上游缓存。
@@ -284,7 +283,7 @@ func (h *GatewayHandler) GeminiCachedContentsCreate(c *gin.Context) {
reqLog.Info("gemini.cached_content.storage_charge_rejected", zap.Float64("amount", amount), zap.Error(err))
h.cleanupGeminiCachedContentUpstream(reqLog, account, upstream.Name)
h.submitGeminiCachedContentUsage(c, apiKey, account, subscription, pricingAt,
"create:"+publicID, modelName, mappedModel, int(tokens), 0, usageFields, body)
service.GeminiCachedContentCreateUsageRequestID(publicID), modelName, mappedModel, int(tokens), 0, usageFields, body)
writeGeminiCachedContentBillingError(c, err)
return
}
@@ -303,6 +302,7 @@ func (h *GatewayHandler) GeminiCachedContentsCreate(c *gin.Context) {
DisplayName: req.DisplayName,
TotalTokenCount: tokens,
ExpireTime: expireTime,
ChannelUsage: usageFields,
}
persistCtx, cancelPersist := geminiCachedContentPersistContext(c)
err = svc.Save(persistCtx, record)
@@ -315,13 +315,13 @@ func (h *GatewayHandler) GeminiCachedContentsCreate(c *gin.Context) {
}
h.submitGeminiCachedContentUsage(c, apiKey, account, subscription, pricingAt,
"create:"+publicID, modelName, mappedModel, int(tokens), storageTokenHours, usageFields, body)
service.GeminiCachedContentCreateUsageRequestID(publicID), modelName, mappedModel, int(tokens), storageTokenHours, usageFields, body)
c.JSON(http.StatusOK, service.GeminiCachedContentView(record))
return
}
if lastUpstreamErr != nil {
writeGeminiCachedContentUpstreamError(c, lastUpstreamErr)
h.writeGeminiCachedContentUpstreamError(c, lastUpstreamErr)
return
}
if ttlRejectedLimit > 0 {
@@ -436,6 +436,7 @@ func (h *GatewayHandler) GeminiCachedContentsPatch(c *gin.Context) {
CacheStorageTokenHours: service.GeminiCachedContentStorageTokenHours(record.TotalTokenCount, extension),
},
APIKey: apiKey, User: apiKey.User, Subscription: subscription, PricingAt: pricingAt,
ChannelUsageFields: record.ChannelUsage,
}
if !h.gatewayService.CacheStoragePriced(c.Request.Context(), usage) {
googlePlatformError(c, http.StatusBadRequest, fmt.Sprintf("Explicit context caching is not available for model %s: storage pricing is not configured", record.Model))
@@ -463,7 +464,7 @@ func (h *GatewayHandler) GeminiCachedContentsPatch(c *gin.Context) {
reqLog.Warn("gemini.cached_content.forget_failed", zap.Error(forgetErr))
}
}
writeGeminiCachedContentUpstreamError(c, upErr)
h.writeGeminiCachedContentUpstreamError(c, upErr)
return
}
reqLog.Warn("gemini.cached_content.upstream_patch_error", zap.Int64("account_id", account.ID), zap.Error(err))
@@ -486,8 +487,8 @@ func (h *GatewayHandler) GeminiCachedContentsPatch(c *gin.Context) {
// 按库中实际推进的时长收费:并发延长各自只为自己推进的部分付费。
if extended := expireTime.Sub(previousExpire); extended > 0 {
h.submitGeminiCachedContentUsage(c, apiKey, account, subscription, pricingAt,
"patch:"+record.PublicID+":"+uuid.NewString(), record.Model, record.UpstreamModelID(), 0,
service.GeminiCachedContentStorageTokenHours(record.TotalTokenCount, extended), service.ChannelUsageFields{}, body)
service.GeminiCachedContentPatchUsageRequestID(record.PublicID), record.Model, record.UpstreamModelID(), 0,
service.GeminiCachedContentStorageTokenHours(record.TotalTokenCount, extended), record.ChannelUsage, body)
}
c.JSON(http.StatusOK, service.GeminiCachedContentView(record))
}
@@ -520,7 +521,7 @@ func (h *GatewayHandler) GeminiCachedContentsDelete(c *gin.Context) {
return
}
if !upErr.NotFound() {
writeGeminiCachedContentUpstreamError(c, upErr)
h.writeGeminiCachedContentUpstreamError(c, upErr)
return
}
}
@@ -628,14 +629,15 @@ func (h *GatewayHandler) acquireGeminiCachedContentAccountSlot(c *gin.Context, r
return wrapReleaseOnDone(c.Request.Context(), release), true
}
// submitGeminiCachedContentUsage 记录缓存创建 / 延长的计费:inputTokens 按输入价计,storageTokenHours 按计费模型的存储单价计。
// submitGeminiCachedContentUsage 记录缓存创建 / 延长的计费:inputTokens 按输入价计,storageTokenHours 按计费模型的存储单价计;
// requestID 是带固定前缀的费用行标识(见 GeminiCachedContentUsageRequestIDPrefix),落库时不被请求 id 覆盖。
func (h *GatewayHandler) submitGeminiCachedContentUsage(
c *gin.Context,
apiKey *service.APIKey,
account *service.Account,
subscription *service.UserSubscription,
pricingAt time.Time,
requestSuffix string,
requestID string,
model string,
upstreamModel string,
inputTokens int,
@@ -650,7 +652,7 @@ func (h *GatewayHandler) submitGeminiCachedContentUsage(
upstreamModel = ""
}
result := &service.ForwardResult{
RequestID: "gemini-cached-content:" + requestSuffix,
RequestID: requestID,
Model: model,
UpstreamModel: upstreamModel,
Usage: service.ClaudeUsage{InputTokens: inputTokens},
@@ -729,9 +731,10 @@ func writeGeminiCachedContentBillingError(c *gin.Context, err error) {
googleErrorWithType(c, status, code, "", message)
}
// writeGeminiCachedContentUpstreamError 透传上游错误:缓存不存在统一返回标准 403(各上游形态不同且可能含上游资源 ID);
// 账号凭据或权限类错误不是客户端问题,统一改为 502。
func writeGeminiCachedContentUpstreamError(c *gin.Context, upErr *service.GeminiCachedContentUpstreamError) {
// writeGeminiCachedContentUpstreamError 回写上游错误:缓存不存在统一返回标准 403(各上游形态不同且可能含上游资源 ID);
// 账号凭据或权限类错误不是客户端问题,统一改为 502;限流 / 不可用按普通 Gemini 路径处理(透传规则、通用映射),
// 只有上游网关账号的这类响应反映持有缓存的上游账号状态,原样回写。
func (h *GatewayHandler) writeGeminiCachedContentUpstreamError(c *gin.Context, upErr *service.GeminiCachedContentUpstreamError) {
status := upErr.StatusCode
message := strings.TrimSpace(extractGoogleErrorMessage(upErr.Body))
service.SetOpsUpstreamError(c, status, message, "")
@@ -740,6 +743,10 @@ func writeGeminiCachedContentUpstreamError(c *gin.Context, upErr *service.Gemini
c.Data(http.StatusForbidden, "application/json", []byte(service.GeminiCachedContentNotFoundResponse))
return
}
if (status == http.StatusTooManyRequests || status >= http.StatusInternalServerError) && !upErr.UpstreamGateway {
h.handleGeminiFailoverExhausted(c, &service.UpstreamFailoverError{StatusCode: status, ResponseBody: upErr.Body})
return
}
if status == http.StatusUnauthorized || status == http.StatusForbidden {
googleError(c, http.StatusBadGateway, "Upstream account rejected the cached content request")
return
@@ -794,11 +801,11 @@ func mergeAccountIDSets(a, b map[int64]struct{}) map[int64]struct{} {
return out
}
// handleGeminiCachedContentFailoverExhausted 处理绑定缓存的请求在持有账号上重试耗尽:没有其他账号可换,
// 上游的限流 / 不可用响应原样返回,客户端据此判断缓存所在账号的状态;账号凭据类错误仍按网关映射。
// handleGeminiCachedContentFailoverExhausted 处理绑定缓存的请求在持有账号上重试耗尽:没有其他账号可换。
// 上游网关账号的限流 / 不可用响应反映持有缓存的上游账号状态,原样返回供客户端判断;
// 其余(含官方账号)仍按普通 Gemini 路径处理:透传规则、通用映射。
func (h *GatewayHandler) handleGeminiCachedContentFailoverExhausted(c *gin.Context, failoverErr *service.UpstreamFailoverError) {
if failoverErr == nil || len(failoverErr.ResponseBody) == 0 || !json.Valid(failoverErr.ResponseBody) ||
(failoverErr.StatusCode != http.StatusTooManyRequests && failoverErr.StatusCode < http.StatusInternalServerError) {
if failoverErr == nil || !failoverErr.BoundUpstreamPassthrough || len(failoverErr.ResponseBody) == 0 || !json.Valid(failoverErr.ResponseBody) {
h.handleGeminiFailoverExhausted(c, failoverErr)
return
}
@@ -204,17 +204,45 @@ func TestHandleGeminiCachedContentFailoverExhausted(t *testing.T) {
unavailable := `{"error":{"code":503,"message":"The account holding this cached content is temporarily unavailable","status":"UNAVAILABLE"}}`
c, rec := newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", apiKey)
h.handleGeminiCachedContentFailoverExhausted(c, &service.UpstreamFailoverError{StatusCode: http.StatusServiceUnavailable, ResponseBody: []byte(unavailable)})
require.Equal(t, http.StatusServiceUnavailable, rec.Code, "绑定缓存的请求无账号可换,上游 503 原样返回")
h.handleGeminiCachedContentFailoverExhausted(c, &service.UpstreamFailoverError{StatusCode: http.StatusServiceUnavailable, ResponseBody: []byte(unavailable), BoundUpstreamPassthrough: true})
require.Equal(t, http.StatusServiceUnavailable, rec.Code, "绑定缓存的请求无账号可换,上游网关的 503 原样返回")
require.JSONEq(t, unavailable, rec.Body.String())
c, rec = newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", apiKey)
h.handleGeminiCachedContentFailoverExhausted(c, &service.UpstreamFailoverError{StatusCode: http.StatusTooManyRequests, ResponseBody: []byte(`{"error":{"code":429,"message":"Too many pending requests, please retry later","status":"RESOURCE_EXHAUSTED"}}`)})
h.handleGeminiCachedContentFailoverExhausted(c, &service.UpstreamFailoverError{StatusCode: http.StatusTooManyRequests, ResponseBody: []byte(`{"error":{"code":429,"message":"Too many pending requests, please retry later","status":"RESOURCE_EXHAUSTED"}}`), BoundUpstreamPassthrough: true})
require.Equal(t, http.StatusTooManyRequests, rec.Code)
require.Contains(t, rec.Body.String(), "Too many pending requests")
officialQuota := `{"error":{"code":429,"message":"You exceeded your current quota","status":"RESOURCE_EXHAUSTED","details":[{"@type":"type.googleapis.com/google.rpc.QuotaFailure","violations":[{"quotaId":"GenerateRequestsPerDayPerProjectPerModel-FreeTier"}]}]}}`
c, rec = newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", apiKey)
h.handleGeminiCachedContentFailoverExhausted(c, &service.UpstreamFailoverError{StatusCode: http.StatusTooManyRequests, ResponseBody: []byte(officialQuota)})
require.Equal(t, http.StatusTooManyRequests, rec.Code, "官方账号的限流按普通 Gemini 路径映射")
require.NotContains(t, rec.Body.String(), "FreeTier", "不向客户端暴露上游账号的配额明细")
require.Contains(t, rec.Body.String(), "Upstream rate limit exceeded")
c, rec = newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", apiKey)
h.handleGeminiCachedContentFailoverExhausted(c, &service.UpstreamFailoverError{StatusCode: http.StatusUnauthorized, ResponseBody: []byte(`{"error":{"code":401,"message":"bad key"}}`)})
require.Equal(t, http.StatusBadGateway, rec.Code, "账号凭据类错误仍按网关映射,不透传给客户端")
require.NotContains(t, rec.Body.String(), "bad key")
}
func TestWriteGeminiCachedContentUpstreamError_RateLimitMappedUnlessUpstreamGateway(t *testing.T) {
h := &GatewayHandler{}
apiKey := geminiCachedContentTestAPIKey(service.PlatformGemini)
quota := `{"error":{"code":429,"message":"You exceeded your current quota","status":"RESOURCE_EXHAUSTED","details":[{"quotaId":"GenerateRequestsPerDayPerProjectPerModel-FreeTier"}]}}`
c, rec := newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/cachedContents", apiKey)
h.writeGeminiCachedContentUpstreamError(c, &service.GeminiCachedContentUpstreamError{StatusCode: http.StatusTooManyRequests, Body: []byte(quota)})
require.Equal(t, http.StatusTooManyRequests, rec.Code)
require.NotContains(t, rec.Body.String(), "FreeTier", "官方账号的限流不透传配额明细")
c, rec = newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/cachedContents", apiKey)
h.writeGeminiCachedContentUpstreamError(c, &service.GeminiCachedContentUpstreamError{StatusCode: http.StatusTooManyRequests, Body: []byte(quota), UpstreamGateway: true})
require.Equal(t, http.StatusTooManyRequests, rec.Code)
require.JSONEq(t, quota, rec.Body.String(), "上游网关账号的限流反映持有缓存的上游账号状态,原样回写")
c, rec = newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/cachedContents", apiKey)
h.writeGeminiCachedContentUpstreamError(c, &service.GeminiCachedContentUpstreamError{StatusCode: http.StatusBadRequest, Body: []byte(`{"error":{"code":400,"message":"Cached content is too small","status":"INVALID_ARGUMENT"}}`)})
require.Equal(t, http.StatusBadRequest, rec.Code)
require.Contains(t, rec.Body.String(), "too small", "请求校验类错误仍原样回写")
}
@@ -3,6 +3,7 @@ package repository
import (
"context"
"database/sql"
"encoding/json"
"time"
"github.com/Wei-Shaw/sub2api/internal/service"
@@ -23,16 +24,20 @@ func NewGeminiCachedContentRepository(db *sql.DB) service.GeminiCachedContentRep
}
const geminiCachedContentColumns = `id, public_id, user_id, api_key_id, group_id, account_id, upstream_name, model,
upstream_model, display_name, total_token_count, expire_time, created_at, updated_at`
upstream_model, display_name, total_token_count, expire_time, channel_usage, created_at, updated_at`
func (r *geminiCachedContentRepository) Create(ctx context.Context, record *service.GeminiCachedContent) error {
channelUsage, err := json.Marshal(record.ChannelUsage)
if err != nil {
return err
}
row := r.sql.QueryRowContext(ctx, `
INSERT INTO gemini_cached_contents (public_id, user_id, api_key_id, group_id, account_id, upstream_name, model,
upstream_model, display_name, total_token_count, expire_time)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)
upstream_model, display_name, total_token_count, expire_time, channel_usage)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)
RETURNING id, created_at, updated_at`,
record.PublicID, record.UserID, record.APIKeyID, record.GroupID, record.AccountID, record.UpstreamName,
record.Model, record.UpstreamModel, record.DisplayName, record.TotalTokenCount, record.ExpireTime)
record.Model, record.UpstreamModel, record.DisplayName, record.TotalTokenCount, record.ExpireTime, channelUsage)
if err := row.Scan(&record.ID, &record.CreatedAt, &record.UpdatedAt); err != nil {
return translatePersistenceError(err, nil, service.ErrGeminiCachedContentExists)
}
@@ -102,6 +107,7 @@ type geminiCachedContentScanner interface {
func scanGeminiCachedContent(row geminiCachedContentScanner) (*service.GeminiCachedContent, error) {
var record service.GeminiCachedContent
var channelUsage []byte
if err := row.Scan(
&record.ID,
&record.PublicID,
@@ -115,10 +121,16 @@ func scanGeminiCachedContent(row geminiCachedContentScanner) (*service.GeminiCac
&record.DisplayName,
&record.TotalTokenCount,
&record.ExpireTime,
&channelUsage,
&record.CreatedAt,
&record.UpdatedAt,
); err != nil {
return nil, err
}
if len(channelUsage) > 0 {
if err := json.Unmarshal(channelUsage, &record.ChannelUsage); err != nil {
return nil, err
}
}
return &record, nil
}
@@ -24,6 +24,7 @@ func TestGeminiCachedContentRepository_Lifecycle(t *testing.T) {
PublicID: publicID + suffix, UserID: 1, APIKeyID: keyID, GroupID: groupID, AccountID: 99,
UpstreamName: "cachedContents/up" + publicID, Model: "gemini-3.8-flash", UpstreamModel: "models/gemini-3.8-flash",
DisplayName: "label", TotalTokenCount: 2730, ExpireTime: expire,
ChannelUsage: service.ChannelUsageFields{ChannelID: 7, OriginalModel: "gemini-pro", ChannelMappedModel: "gemini-3.8-flash", BillingModelSource: "requested"},
}
}
first := newRecord("a", apiKeyID, 42, now.Add(time.Hour))
@@ -47,6 +48,7 @@ func TestGeminiCachedContentRepository_Lifecycle(t *testing.T) {
require.Equal(t, "cachedContents/upa", got.UpstreamName)
require.Equal(t, int64(2730), got.TotalTokenCount)
require.True(t, got.ExpireTime.Equal(first.ExpireTime))
require.Equal(t, first.ChannelUsage, got.ChannelUsage, "创建时的渠道计费口径随记录持久化")
_, err = repo.GetForOwner(ctx, apiKeyID+1, first.PublicID)
require.ErrorIs(t, err, service.ErrGeminiCachedContentNotFound)
@@ -956,7 +956,11 @@ func filterSchedulerCredentials(credentials map[string]any) map[string]any {
}
// Candidate-list admission evaluates the account override before hydrating
// the full account. Dropping it silently falls back to the platform threshold.
keys := []string{"model_mapping", "compact_model_mapping", "api_key", "project_id", "oauth_type", "plan_type", "account_scheduling_threshold"}
keys := []string{"model_mapping", "compact_model_mapping", "api_key", "project_id", "oauth_type", "plan_type", "account_scheduling_threshold",
// Gemini 显式缓存的候选过滤(GeminiExplicitCacheExclusions)读本投影判定账号是否支持显式缓存:
// 官方地址、Vertex 认证方式与上游网关的显式缓存声明。缺这几个键时只能按类型粗筛,
// 不支持的账号要逐个选中后才被剔除。
"base_url", "auth_mode", service.GeminiExplicitCacheUpstreamCredentialKey, service.GeminiExplicitCacheUpstreamMaxTTLCredentialKey}
filtered := make(map[string]any)
for _, key := range keys {
if value, ok := credentials[key]; ok && value != nil {
@@ -20,6 +20,23 @@ func TestFilterSchedulerCredentialsKeepsSubscriptionPlanType(t *testing.T) {
require.NotContains(t, filtered, "refresh_token")
}
func TestFilterSchedulerCredentialsKeepsGeminiExplicitCacheKeys(t *testing.T) {
filtered := filterSchedulerCredentials(map[string]any{
"api_key": "k",
"base_url": "https://gw.example.com",
"auth_mode": "apikey",
"service_account_json": "secret",
service.GeminiExplicitCacheUpstreamCredentialKey: true,
service.GeminiExplicitCacheUpstreamMaxTTLCredentialKey: 3600,
})
require.Equal(t, "https://gw.example.com", filtered["base_url"])
require.Equal(t, "apikey", filtered["auth_mode"])
require.Equal(t, true, filtered[service.GeminiExplicitCacheUpstreamCredentialKey])
require.Equal(t, 3600, filtered[service.GeminiExplicitCacheUpstreamMaxTTLCredentialKey])
require.NotContains(t, filtered, "service_account_json")
}
func TestSchedulerMetadataAccountKeepsOpenAISubscriptionIdentity(t *testing.T) {
account := service.Account{
ID: 24,
@@ -699,6 +699,7 @@ type UpstreamFailoverError struct {
SameAccountRetryMax int // 可选的错误级同账号重试上限,低于 handler 默认预算时优先采用
RequestScopedTransient bool // 故障因素与账号无关(如上游按客户端身份/模型容量降载):可同账号重试,但不得据此对账号做临时封禁
SafeToFailoverAfterWrite bool // 仅写出 SSE 注释等非语义字节时,仍可在同一客户端流中切换账号
BoundUpstreamPassthrough bool // 绑定显式缓存的请求在上游网关账号上收到的 429 / 5xx:反映持有缓存的上游账号状态,原样回写不走错误映射
Stage GatewayFailureStage
Scope GatewayFailureScope
Reason GatewayFailureReason
@@ -241,7 +241,8 @@ func isForcedUsageBillingRequestID(requestID string) bool {
return strings.HasPrefix(id, "web_search:") ||
strings.HasPrefix(id, "grok-video:") ||
strings.HasPrefix(id, "grok_audio:") ||
strings.HasPrefix(id, "grok_realtime:")
strings.HasPrefix(id, "grok_realtime:") ||
strings.HasPrefix(id, GeminiCachedContentUsageRequestIDPrefix)
}
// StableGrokAudioBillingRequestID is the durable usage_logs / dedup key for one
@@ -31,9 +31,27 @@ func TestIsForcedUsageBillingRequestID(t *testing.T) {
require.True(t, isForcedUsageBillingRequestID("grok-video:task-1"))
require.True(t, isForcedUsageBillingRequestID("grok_audio:up-1"))
require.True(t, isForcedUsageBillingRequestID("grok_realtime:sess-1"))
require.True(t, isForcedUsageBillingRequestID(GeminiCachedContentCreateUsageRequestID("abc")))
require.False(t, isForcedUsageBillingRequestID("resp_abc"))
}
func TestGeminiCachedContentUsageRequestIDs(t *testing.T) {
t.Parallel()
publicID, err := NewGeminiCachedContentPublicID()
require.NoError(t, err)
create := GeminiCachedContentCreateUsageRequestID(publicID)
patch := GeminiCachedContentPatchUsageRequestID(publicID)
require.Equal(t, "gcache:create:"+publicID, create)
require.True(t, strings.HasPrefix(patch, "gcache:patch:"+publicID+":"))
require.NotEqual(t, patch, GeminiCachedContentPatchUsageRequestID(publicID), "每次延长各自唯一")
require.LessOrEqual(t, len(create), 64, "usage_logs.request_id 列长 64")
require.LessOrEqual(t, len(patch), 64, "usage_logs.request_id 列长 64")
ctx := context.WithValue(context.Background(), ctxkey.ClientRequestID, "client-shared-id")
require.Equal(t, create, resolveUsageBillingRequestID(ctx, create), "缓存费用行不被客户端请求 id 覆盖")
require.Equal(t, patch, resolveUsageBillingRequestID(ctx, patch))
}
func TestStableGrokAudioBillingRequestID(t *testing.T) {
t.Parallel()
require.Equal(t, "grok_audio:up-1", StableGrokAudioBillingRequestID("up-1"))
@@ -74,8 +74,10 @@ type GeminiCachedContent struct {
DisplayName string
TotalTokenCount int64
ExpireTime time.Time
CreatedAt time.Time
UpdatedAt time.Time
// ChannelUsage 是创建时的渠道计费口径(渠道 ID、请求模型、计费模型来源),延长有效期按同一口径计费。
ChannelUsage ChannelUsageFields
CreatedAt time.Time
UpdatedAt time.Time
}
// PublicName 返回客户端可见的缓存资源名。
@@ -127,9 +129,11 @@ func IsGeminiExplicitCacheAccount(account *Account) bool {
return isOfficialGeminiAPIBaseURL(account.GetCredential("base_url")) || account.GeminiExplicitCacheUpstreamEnabled()
}
// GeminiExplicitCacheUpstreamCredentialKey / GeminiExplicitCacheUpstreamMaxTTLCredentialKey 是 API Key 账号凭据里
// 声明上游实现显式缓存接口及其有效期上限的键;调度快照投影须保留它们,候选过滤才能直接判定账号是否支持显式缓存。
const (
geminiExplicitCacheUpstreamKey = "explicit_cache_upstream"
geminiExplicitCacheUpstreamMaxTTLKey = "explicit_cache_upstream_max_ttl_seconds"
GeminiExplicitCacheUpstreamCredentialKey = "explicit_cache_upstream"
GeminiExplicitCacheUpstreamMaxTTLCredentialKey = "explicit_cache_upstream_max_ttl_seconds"
)
// GeminiExplicitCacheUpstreamEnabled 判断 API Key 账号是否声明其上游实现了显式缓存接口,
@@ -138,7 +142,7 @@ func (a *Account) GeminiExplicitCacheUpstreamEnabled() bool {
if a == nil || a.Type != AccountTypeAPIKey || a.Credentials == nil {
return false
}
enabled, _ := a.Credentials[geminiExplicitCacheUpstreamKey].(bool)
enabled, _ := a.Credentials[GeminiExplicitCacheUpstreamCredentialKey].(bool)
return enabled
}
@@ -148,7 +152,7 @@ func (a *Account) GeminiExplicitCacheUpstreamMaxTTL() time.Duration {
return 0
}
var seconds float64
switch v := a.Credentials[geminiExplicitCacheUpstreamMaxTTLKey].(type) {
switch v := a.Credentials[GeminiExplicitCacheUpstreamMaxTTLCredentialKey].(type) {
case float64:
seconds = v
case int:
@@ -198,6 +202,25 @@ func NewGeminiCachedContentPublicID() (string, error) {
return strings.ToLower(geminiCachedContentPublicIDEncoding.EncodeToString(b[:])), nil
}
// GeminiCachedContentUsageRequestIDPrefix 是显式缓存创建 / 延长费用在 usage_logs.request_id 中的固定前缀:
// 这类费用行是独立的扣费事件,不随客户端或网关请求 id 变化,按缓存 ID 即可检索一个缓存的全部费用行。
// 两种格式都不超过 request_id 列的 64 字符。
const GeminiCachedContentUsageRequestIDPrefix = "gcache:"
// GeminiCachedContentCreateUsageRequestID 返回创建费用的 request_id("gcache:create:{id}");一个缓存只创建一次。
func GeminiCachedContentCreateUsageRequestID(publicID string) string {
return GeminiCachedContentUsageRequestIDPrefix + "create:" + publicID
}
// GeminiCachedContentPatchUsageRequestID 返回一次延长费用的 request_id("gcache:patch:{id}:{8 位随机}");每次延长各自唯一。
func GeminiCachedContentPatchUsageRequestID(publicID string) string {
var b [5]byte
if _, err := rand.Read(b[:]); err != nil {
return GeminiCachedContentUsageRequestIDPrefix + "patch:" + publicID + ":" + strconv.FormatInt(time.Now().UnixNano()%(1<<40), 32)
}
return GeminiCachedContentUsageRequestIDPrefix + "patch:" + publicID + ":" + strings.ToLower(geminiCachedContentPublicIDEncoding.EncodeToString(b[:]))
}
// ParseGeminiCachedContentPublicID 从 "cachedContents/{id}" 或裸 id 中解析缓存 ID。
func ParseGeminiCachedContentPublicID(name string) (string, bool) {
id := strings.TrimPrefix(strings.TrimSpace(name), GeminiCachedContentNamePrefix)
@@ -525,6 +548,8 @@ type GeminiCachedContentUpstream struct {
type GeminiCachedContentUpstreamError struct {
StatusCode int
Body []byte
// UpstreamGateway 表示响应来自声明实现显式缓存的上游网关账号:限流 / 不可用反映的是持有缓存的上游账号状态。
UpstreamGateway bool
}
func (e *GeminiCachedContentUpstreamError) Error() string {
@@ -827,7 +852,7 @@ func (s *GeminiCachedContentService) doUpstream(ctx context.Context, account *Ac
return nil, err
}
if !success {
return nil, &GeminiCachedContentUpstreamError{StatusCode: resp.StatusCode, Body: respBody}
return nil, &GeminiCachedContentUpstreamError{StatusCode: resp.StatusCode, Body: respBody, UpstreamGateway: account.GeminiExplicitCacheUpstreamEnabled()}
}
return respBody, nil
}
@@ -936,8 +961,9 @@ func GeminiCachedContentView(record *GeminiCachedContent) map[string]any {
}
// GeminiExplicitCacheExclusions 返回显式缓存请求在 gemini 分组内需要排除的可调度账号。
// boundAccountID > 0 时只保留该账号;否则排除平台或类型不支持显式缓存的账号
// (调度快照不含 base_url,官方地址判定在选中后对完整账号进行)。
// boundAccountID > 0 时只保留该账号;否则排除不支持显式缓存的账号。判定所需的凭据键
// (api_key、base_url、auth_mode、上游显式缓存声明)都在调度快照投影内,候选过滤一次算全;
// 选中后仍对完整账号复核。
func (s *GatewayService) GeminiExplicitCacheExclusions(ctx context.Context, groupID *int64, boundAccountID int64) (map[int64]struct{}, error) {
accounts, _, err := s.listSchedulableAccounts(ctx, groupID, PlatformGemini, false)
if err != nil {
@@ -946,7 +972,7 @@ func (s *GatewayService) GeminiExplicitCacheExclusions(ctx context.Context, grou
excluded := make(map[int64]struct{}, len(accounts))
for i := range accounts {
account := &accounts[i]
keep := IsGeminiExplicitCacheAccountType(account)
keep := IsGeminiExplicitCacheAccount(account)
if boundAccountID > 0 {
keep = account.ID == boundAccountID
}
@@ -285,6 +285,7 @@ func TestGeminiForwardNative_BoundCacheUpstreamGatewayErrorsSkipAccountPenalty(t
require.True(t, errors.As(err, &failoverErr), "交由绑定请求的耗尽处理透传")
require.Equal(t, status, failoverErr.StatusCode)
require.JSONEq(t, body, string(failoverErr.ResponseBody))
require.True(t, failoverErr.BoundUpstreamPassthrough, "标记为上游网关账号的响应,handler 原样回写")
require.False(t, failoverErr.RetryableOnSameAccount)
require.Equal(t, 1, httpStub.calls, "不在同一账号上退避重试")
require.Zero(t, repo.rateLimitedCalls, "不限流整个网关账号")
@@ -421,7 +421,7 @@ func TestGeminiCachedContentStorageCostInAccountStats(t *testing.T) {
gid := groupID
err := svc.RecordUsage(context.Background(), &RecordUsageInput{
Result: &ForwardResult{
RequestID: "gemini-cached-content:create:" + tc.name,
RequestID: GeminiCachedContentCreateUsageRequestID(tc.name),
Model: model,
Usage: ClaudeUsage{InputTokens: tc.tokens},
CacheStorageTokenHours: storage / 0.0000005,
@@ -1470,7 +1470,7 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
}
// 上游网关的限流 / 不可用只针对其持有缓存的账号:不改本账号状态,交由绑定请求的耗尽处理透传。
if isGeminiBoundCacheUpstreamPassthrough(ctx, account, resp.StatusCode) {
return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: respBody}
return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: respBody, BoundUpstreamPassthrough: true}
}
// Best-effort fallback for OAuth tokens missing AI Studio scopes when calling countTokens.
// This avoids Gemini SDKs failing hard during preflight token counting.
@@ -0,0 +1,4 @@
-- Gemini 显式缓存创建时的渠道计费口径(渠道 ID / 请求模型 / 计费模型来源),
-- 延长有效期按同一口径计费。
ALTER TABLE gemini_cached_contents
ADD COLUMN IF NOT EXISTS channel_usage JSONB NOT NULL DEFAULT '{}'::jsonb;