mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-14 03:52:26 +08:00
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:
@@ -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;
|
||||
Reference in New Issue
Block a user