mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-15 23:22:59 +08:00
Merge pull request #8006 from feeeei/feat/gemini-explicit-cache
feat(gemini): 支持显式上下文缓存,并完善 Vertex 与第三方 API Key 账号
This commit is contained in:
@@ -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,955 @@
|
||||
//go:build unit
|
||||
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||||
middleware "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
const (
|
||||
cachedBillingInputPrice = 0.00000125
|
||||
cachedBillingStoragePrice = 0.0000045
|
||||
cachedBillingAliasStoragePrice = 0.000009
|
||||
cachedBillingTokens = 1_000_000
|
||||
)
|
||||
|
||||
type cachedBillingRepo struct {
|
||||
mu sync.Mutex
|
||||
records map[string]*service.GeminiCachedContent
|
||||
// mutations 按记录 id 保存修改意图。
|
||||
mutations map[int64]cachedBillingMutation
|
||||
failCommit bool
|
||||
nextID int64
|
||||
deleted []int64
|
||||
}
|
||||
|
||||
type cachedBillingMutation struct {
|
||||
token string
|
||||
leaseUntil time.Time
|
||||
}
|
||||
|
||||
func (r *cachedBillingRepo) Create(_ context.Context, record *service.GeminiCachedContent) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.nextID++
|
||||
record.ID = r.nextID
|
||||
cp := *record
|
||||
r.records[record.PublicID] = &cp
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *cachedBillingRepo) GetForOwner(_ context.Context, apiKeyID int64, publicID string) (*service.GeminiCachedContent, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
record, ok := r.records[publicID]
|
||||
if !ok || record.APIKeyID != apiKeyID || r.liveLocked(record.ID) == nil {
|
||||
return nil, service.ErrGeminiCachedContentNotFound
|
||||
}
|
||||
cp := *record
|
||||
return &cp, nil
|
||||
}
|
||||
|
||||
func (r *cachedBillingRepo) ListForOwner(context.Context, int64, int64, time.Time, int64, int) ([]*service.GeminiCachedContent, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// liveLocked 返回未删除的记录;调用方须持有 mu。
|
||||
func (r *cachedBillingRepo) liveLocked(id int64) *service.GeminiCachedContent {
|
||||
for _, deleted := range r.deleted {
|
||||
if deleted == id {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
for _, record := range r.records {
|
||||
if record.ID == id {
|
||||
return record
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *cachedBillingRepo) BeginMutation(_ context.Context, id int64, baseline time.Time, token string, lease time.Duration) (bool, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
record := r.liveLocked(id)
|
||||
if record == nil {
|
||||
return false, service.ErrGeminiCachedContentNotFound
|
||||
}
|
||||
if _, pending := r.mutations[id]; pending || !record.ExpireTime.Equal(baseline) {
|
||||
return false, nil
|
||||
}
|
||||
r.mutations[id] = cachedBillingMutation{token: token, leaseUntil: time.Now().Add(lease)}
|
||||
return true, nil
|
||||
}
|
||||
|
||||
func (r *cachedBillingRepo) TakeOverStaleMutation(_ context.Context, id int64, token string, lease time.Duration) (*service.GeminiCachedContent, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
record := r.liveLocked(id)
|
||||
mutation, pending := r.mutations[id]
|
||||
if record == nil || !pending || mutation.leaseUntil.After(time.Now()) {
|
||||
return nil, nil
|
||||
}
|
||||
r.mutations[id] = cachedBillingMutation{token: token, leaseUntil: time.Now().Add(lease)}
|
||||
cp := *record
|
||||
return &cp, nil
|
||||
}
|
||||
|
||||
func (r *cachedBillingRepo) CommitMutation(_ context.Context, id int64, token string, expireTime time.Time) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
if r.failCommit {
|
||||
return errors.New("db unavailable")
|
||||
}
|
||||
record := r.liveLocked(id)
|
||||
mutation, pending := r.mutations[id]
|
||||
if record == nil || !pending || mutation.token != token {
|
||||
return service.ErrGeminiCachedContentMutationConflict
|
||||
}
|
||||
record.ExpireTime = expireTime
|
||||
delete(r.mutations, id)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *cachedBillingRepo) ReleaseMutation(_ context.Context, id int64, token string) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
if mutation, pending := r.mutations[id]; pending && mutation.token == token {
|
||||
delete(r.mutations, id)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *cachedBillingRepo) ListStaleMutations(_ context.Context, limit int) ([]int64, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
var ids []int64
|
||||
for id, mutation := range r.mutations {
|
||||
if len(ids) < limit && mutation.leaseUntil.Before(time.Now()) && r.liveLocked(id) != nil {
|
||||
ids = append(ids, id)
|
||||
}
|
||||
}
|
||||
return ids, nil
|
||||
}
|
||||
|
||||
func (r *cachedBillingRepo) SoftDelete(_ context.Context, id int64) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.deleted = append(r.deleted, id)
|
||||
return nil
|
||||
}
|
||||
|
||||
// pendingMutation 返回记录是否有未清除的修改意图。
|
||||
func (r *cachedBillingRepo) pendingMutation(id int64) bool {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
_, pending := r.mutations[id]
|
||||
return pending
|
||||
}
|
||||
|
||||
// expireLease 让记录上的修改意图租约立即过期。
|
||||
func (r *cachedBillingRepo) expireLease(id int64) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
mutation := r.mutations[id]
|
||||
mutation.leaseUntil = time.Now().Add(-time.Second)
|
||||
r.mutations[id] = mutation
|
||||
}
|
||||
|
||||
type cachedBillingAccountRepo struct {
|
||||
service.AccountRepository
|
||||
accounts map[int64]*service.Account
|
||||
}
|
||||
|
||||
func (r *cachedBillingAccountRepo) GetByID(_ context.Context, id int64) (*service.Account, error) {
|
||||
account, ok := r.accounts[id]
|
||||
if !ok {
|
||||
return nil, service.ErrAccountNotFound
|
||||
}
|
||||
cp := *account
|
||||
return &cp, nil
|
||||
}
|
||||
|
||||
type cachedBillingUpstream struct {
|
||||
service.HTTPUpstream
|
||||
mu sync.Mutex
|
||||
createBody string
|
||||
calls []string
|
||||
// patchTTLs 按顺序记录 PATCH 请求体中的 ttl。
|
||||
patchTTLs []string
|
||||
beforePost func(n int)
|
||||
// failures 按请求方法返回上游错误响应;transportErrors 按请求方法返回传输层错误(上游结果未知)。
|
||||
failures map[string]cachedBillingUpstreamFailure
|
||||
transportErrors map[string]error
|
||||
}
|
||||
|
||||
type cachedBillingUpstreamFailure struct {
|
||||
status int
|
||||
body string
|
||||
}
|
||||
|
||||
func (u *cachedBillingUpstream) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) {
|
||||
var patch struct {
|
||||
TTL string `json:"ttl"`
|
||||
}
|
||||
if req.Method == http.MethodPatch {
|
||||
_ = json.NewDecoder(req.Body).Decode(&patch)
|
||||
}
|
||||
u.mu.Lock()
|
||||
u.calls = append(u.calls, req.Method+" "+req.URL.Path)
|
||||
n := len(u.calls)
|
||||
if req.Method == http.MethodPatch {
|
||||
u.patchTTLs = append(u.patchTTLs, patch.TTL)
|
||||
}
|
||||
failure, failed := u.failures[req.Method]
|
||||
transportErr := u.transportErrors[req.Method]
|
||||
u.mu.Unlock()
|
||||
if transportErr != nil {
|
||||
return nil, transportErr
|
||||
}
|
||||
if failed {
|
||||
return &http.Response{StatusCode: failure.status, Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(failure.body))}, nil
|
||||
}
|
||||
body := "{}"
|
||||
switch req.Method {
|
||||
case http.MethodPost:
|
||||
body = u.createBody
|
||||
case http.MethodPatch:
|
||||
if u.beforePost != nil {
|
||||
u.beforePost(n)
|
||||
}
|
||||
seconds, _ := strconv.ParseInt(strings.TrimSuffix(patch.TTL, "s"), 10, 64)
|
||||
body = fmt.Sprintf(`{"name":"cachedContents/up1","expireTime":%q}`, time.Now().Add(time.Duration(seconds)*time.Second).UTC().Format(time.RFC3339Nano))
|
||||
}
|
||||
return &http.Response{StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(body))}, nil
|
||||
}
|
||||
|
||||
func (u *cachedBillingUpstream) methodCalls(method string) int {
|
||||
u.mu.Lock()
|
||||
defer u.mu.Unlock()
|
||||
n := 0
|
||||
for _, call := range u.calls {
|
||||
if strings.HasPrefix(call, method+" ") {
|
||||
n++
|
||||
}
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
type cachedBillingUserRepo struct {
|
||||
service.UserRepository
|
||||
balance float64
|
||||
}
|
||||
|
||||
func (r *cachedBillingUserRepo) GetByID(_ context.Context, id int64) (*service.User, error) {
|
||||
return &service.User{ID: id, Balance: r.balance}, nil
|
||||
}
|
||||
|
||||
func (r *cachedBillingUserRepo) DeductBalance(context.Context, int64, float64) error { return nil }
|
||||
|
||||
type cachedBillingUsageRepo struct {
|
||||
service.UsageLogRepository
|
||||
mu sync.Mutex
|
||||
logs []*service.UsageLog
|
||||
}
|
||||
|
||||
func (r *cachedBillingUsageRepo) Create(_ context.Context, log *service.UsageLog) (bool, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.logs = append(r.logs, log)
|
||||
return true, nil
|
||||
}
|
||||
|
||||
func (r *cachedBillingUsageRepo) first() *service.UsageLog {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
return r.logs[0]
|
||||
}
|
||||
|
||||
func (r *cachedBillingUsageRepo) totals() []float64 {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
out := make([]float64, 0, len(r.logs))
|
||||
for _, log := range r.logs {
|
||||
out = append(out, log.TotalCost)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
type cachedBillingEnv struct {
|
||||
h *GatewayHandler
|
||||
repo *cachedBillingRepo
|
||||
upstream *cachedBillingUpstream
|
||||
usage *cachedBillingUsageRepo
|
||||
apiKey *service.APIKey
|
||||
group *service.Group
|
||||
// accounts 是调度快照桩直接返回的账号对象,测试可改其凭据。
|
||||
accounts map[int64]*service.Account
|
||||
}
|
||||
|
||||
func newCachedBillingEnv(t *testing.T, balance float64, accountIDs ...int64) *cachedBillingEnv {
|
||||
t.Helper()
|
||||
dir := t.TempDir()
|
||||
require.NoError(t, os.WriteFile(filepath.Join(dir, "model_pricing.json"), []byte(fmt.Sprintf(`{"gemini-2.5-pro":{
|
||||
"litellm_provider":"gemini","mode":"chat","input_cost_per_token":%g,"output_cost_per_token":0.00001,
|
||||
"cache_storage_cost_per_token_per_hour":%g},"gemini-alias":{
|
||||
"litellm_provider":"gemini","mode":"chat","input_cost_per_token":%g,"output_cost_per_token":0.00001,
|
||||
"cache_storage_cost_per_token_per_hour":%g}}`,
|
||||
cachedBillingInputPrice, cachedBillingStoragePrice, cachedBillingInputPrice, cachedBillingAliasStoragePrice)), 0o644))
|
||||
cfg := &config.Config{}
|
||||
cfg.Pricing.DataDir = dir
|
||||
cfg.Pricing.UpdateIntervalHours = 100000
|
||||
cfg.Default.RateMultiplier = 1
|
||||
pricing := service.NewPricingService(cfg, nil)
|
||||
require.NoError(t, pricing.Initialize())
|
||||
billing := service.NewBillingService(cfg, pricing)
|
||||
|
||||
groupID := int64(42)
|
||||
group := &service.Group{ID: groupID, Hydrated: true, Platform: service.PlatformGemini, Status: service.StatusActive, RateMultiplier: 1}
|
||||
accounts := map[int64]*service.Account{}
|
||||
var scheduled []*service.Account
|
||||
for _, id := range accountIDs {
|
||||
account := &service.Account{ID: id, Platform: service.PlatformGemini, Type: service.AccountTypeAPIKey,
|
||||
Credentials: map[string]any{"api_key": "k"}, Status: service.StatusActive, Schedulable: true,
|
||||
Concurrency: 5, Priority: 1, AccountGroups: []service.AccountGroup{{AccountID: id, GroupID: groupID}}}
|
||||
accounts[id] = account
|
||||
scheduled = append(scheduled, account)
|
||||
}
|
||||
userRepo := &cachedBillingUserRepo{balance: balance}
|
||||
usage := &cachedBillingUsageRepo{}
|
||||
bcs := service.NewBillingCacheService(nil, userRepo, nil, nil, nil, nil, cfg, nil)
|
||||
t.Cleanup(bcs.Stop)
|
||||
snapshot := service.NewSchedulerSnapshotService(&fakeSchedulerCache{accounts: scheduled}, nil, nil, nil, nil)
|
||||
gw := service.NewGatewayService(nil, &fakeGroupRepo{group: group}, usage, nil, userRepo, nil, nil, nil, cfg, snapshot, nil,
|
||||
billing, nil, &service.BillingCacheService{}, nil, nil, &service.DeferredService{}, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
|
||||
upstream := &cachedBillingUpstream{}
|
||||
compat := service.NewGeminiMessagesCompatService(&cachedBillingAccountRepo{accounts: accounts}, nil, nil, nil, nil, nil, upstream, nil, cfg)
|
||||
repo := &cachedBillingRepo{records: map[string]*service.GeminiCachedContent{}, mutations: map[int64]cachedBillingMutation{}}
|
||||
h := &GatewayHandler{
|
||||
gatewayService: gw,
|
||||
billingCacheService: bcs,
|
||||
concurrencyHelper: NewConcurrencyHelper(service.NewConcurrencyService(&fakeConcurrencyCache{}), SSEPingFormatNone, 0),
|
||||
geminiCachedContentService: service.NewGeminiCachedContentService(repo, compat, cfg),
|
||||
}
|
||||
apiKey := &service.APIKey{ID: 7, UserID: 3, GroupID: &groupID, Group: group, Status: service.StatusActive,
|
||||
User: &service.User{ID: 3, Concurrency: 10, Balance: balance}}
|
||||
return &cachedBillingEnv{h: h, repo: repo, upstream: upstream, usage: usage, apiKey: apiKey, group: group, accounts: accounts}
|
||||
}
|
||||
|
||||
// TestGeminiExplicitCacheExclusions_FiltersUnsupportedAccountsFromSnapshot 候选过滤直接按快照里的凭据判定:
|
||||
// 第三方地址且未声明上游显式缓存的 Key 账号在选号前就被排除,官方地址与声明了上游显式缓存的账号保留。
|
||||
func TestGeminiExplicitCacheExclusions_FiltersUnsupportedAccountsFromSnapshot(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 500, 900, 901, 902)
|
||||
env.accounts[901].Credentials["base_url"] = "https://third.example.com"
|
||||
env.accounts[902].Credentials["base_url"] = "https://gw.example.com"
|
||||
env.accounts[902].Credentials[service.GeminiExplicitCacheUpstreamCredentialKey] = true
|
||||
|
||||
groupID := env.group.ID
|
||||
excluded, err := env.h.gatewayService.GeminiExplicitCacheExclusions(context.Background(), &groupID, 0)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, map[int64]struct{}{901: {}}, excluded)
|
||||
|
||||
excluded, err = env.h.gatewayService.GeminiExplicitCacheExclusions(context.Background(), &groupID, 902)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, map[int64]struct{}{900: {}, 901: {}}, excluded, "绑定缓存的请求只保留持有账号")
|
||||
}
|
||||
|
||||
func (e *cachedBillingEnv) serve(method, target, body string, params gin.Params, call func(*gin.Context)) *httptest.ResponseRecorder {
|
||||
gin.SetMode(gin.TestMode)
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
req := httptest.NewRequest(method, target, strings.NewReader(body))
|
||||
c.Request = req.WithContext(context.WithValue(req.Context(), ctxkey.Group, e.group))
|
||||
c.Params = params
|
||||
c.Set(string(middleware.ContextKeyAPIKey), e.apiKey)
|
||||
c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: e.apiKey.UserID, Concurrency: 10})
|
||||
call(c)
|
||||
return rec
|
||||
}
|
||||
|
||||
// waitForUsage 等待异步计费任务落库(无 worker 池时同步执行,直接返回)。
|
||||
func (e *cachedBillingEnv) waitForUsage(t *testing.T, n int) []float64 {
|
||||
t.Helper()
|
||||
require.Eventually(t, func() bool { return len(e.usage.totals()) >= n }, time.Second, 10*time.Millisecond)
|
||||
return e.usage.totals()
|
||||
}
|
||||
|
||||
func (e *cachedBillingEnv) create() *httptest.ResponseRecorder {
|
||||
return e.serve(http.MethodPost, "/v1beta/cachedContents",
|
||||
`{"model":"models/gemini-2.5-pro","ttl":"86400s","contents":[{"role":"user","parts":[{"text":"doc"}]}]}`,
|
||||
nil, e.h.GeminiCachedContentsCreate)
|
||||
}
|
||||
|
||||
func (e *cachedBillingEnv) seed(expire time.Time) {
|
||||
e.repo.records["abc"] = &service.GeminiCachedContent{
|
||||
ID: 100, PublicID: "abc", UserID: 3, APIKeyID: 7, GroupID: 42, AccountID: 900,
|
||||
UpstreamName: "cachedContents/up1", Model: "gemini-2.5-pro", UpstreamModel: "models/gemini-2.5-pro",
|
||||
TotalTokenCount: cachedBillingTokens, ExpireTime: expire,
|
||||
}
|
||||
}
|
||||
|
||||
func (e *cachedBillingEnv) patch() *httptest.ResponseRecorder {
|
||||
return e.serve(http.MethodPatch, "/v1beta/cachedContents/abc?updateMask=ttl", `{"ttl":"86400s"}`,
|
||||
gin.Params{{Key: "cacheID", Value: "abc"}}, e.h.GeminiCachedContentsPatch)
|
||||
}
|
||||
|
||||
// TestGeminiCachedContentsPatch_UsesChannelUsageRecordedAtCreate 延长有效期沿用创建时的渠道计费口径:
|
||||
// 计费模型来源为请求模型时按别名定价,用量记录带渠道 ID。
|
||||
func TestGeminiCachedContentsPatch_UsesChannelUsageRecordedAtCreate(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 500, 900)
|
||||
env.seed(time.Now().Add(time.Hour))
|
||||
env.repo.records["abc"].ChannelUsage = service.ChannelUsageFields{
|
||||
ChannelID: 5, OriginalModel: "gemini-alias", ChannelMappedModel: "gemini-2.5-pro",
|
||||
BillingModelSource: service.BillingModelSourceRequested,
|
||||
}
|
||||
|
||||
rec := env.patch()
|
||||
require.Equal(t, http.StatusOK, rec.Code, rec.Body.String())
|
||||
totals := env.waitForUsage(t, 1)
|
||||
require.Len(t, totals, 1)
|
||||
require.InDelta(t, cachedBillingTokens*cachedBillingAliasStoragePrice*23, totals[0], 0.01, "按创建时的计费模型(请求模型别名)计存储费")
|
||||
usageLog := env.usage.first()
|
||||
require.NotNil(t, usageLog.ChannelID)
|
||||
require.Equal(t, int64(5), *usageLog.ChannelID)
|
||||
require.True(t, strings.HasPrefix(usageLog.RequestID, "gcache:patch:abc:"), usageLog.RequestID)
|
||||
}
|
||||
|
||||
func cachedBillingCreateBody(expire time.Time) string {
|
||||
return fmt.Sprintf(`{"name":"cachedContents/up1","usageMetadata":{"totalTokenCount":%d},"expireTime":%q}`,
|
||||
cachedBillingTokens, expire.UTC().Format(time.RFC3339Nano))
|
||||
}
|
||||
|
||||
func TestGeminiCachedContentsCreate_ChargesInputAndStorage(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 500, 900)
|
||||
env.upstream.createBody = cachedBillingCreateBody(time.Now().Add(100 * time.Hour))
|
||||
|
||||
rec := env.create()
|
||||
require.Equal(t, http.StatusOK, rec.Code, rec.Body.String())
|
||||
totals := env.waitForUsage(t, 1)
|
||||
require.Len(t, totals, 1)
|
||||
require.InDelta(t, cachedBillingTokens*cachedBillingInputPrice+cachedBillingTokens*cachedBillingStoragePrice*24, totals[0], 1e-9)
|
||||
require.Len(t, env.repo.records, 1)
|
||||
for _, record := range env.repo.records {
|
||||
require.False(t, record.ExpireTime.After(time.Now().Add(24*time.Hour)), "本地到期时间不超过请求的有效期")
|
||||
require.Equal(t, "gcache:create:"+record.PublicID, env.usage.first().RequestID, "费用行按缓存 ID 可检索")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeminiCachedContentsCreate_UnaffordableStorageDeletesUpstreamAndBillsInputOnly(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 1, 900)
|
||||
env.upstream.createBody = cachedBillingCreateBody(time.Now().Add(24 * time.Hour))
|
||||
|
||||
rec := env.create()
|
||||
require.NotEqual(t, http.StatusOK, rec.Code)
|
||||
require.Contains(t, rec.Body.String(), "error")
|
||||
require.Equal(t, 1, env.upstream.methodCalls(http.MethodDelete), "删除已创建的上游缓存,避免产生存储费")
|
||||
require.Empty(t, env.repo.records, "不登记缓存")
|
||||
totals := env.waitForUsage(t, 1)
|
||||
require.Len(t, totals, 1)
|
||||
require.InDelta(t, cachedBillingTokens*cachedBillingInputPrice, totals[0], 1e-9, "只收已发生的缓存 token 输入费")
|
||||
}
|
||||
|
||||
func TestGeminiCachedContentsCreate_UnusableUpstreamResponseCleansUpWithoutRetry(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 500, 900, 901)
|
||||
env.upstream.createBody = `{"name":"cachedContents/up1","expireTime":"2099-01-01T00:00:00Z"}`
|
||||
|
||||
rec := env.create()
|
||||
require.Equal(t, http.StatusBadGateway, rec.Code)
|
||||
require.Equal(t, 1, env.upstream.methodCalls(http.MethodPost), "不换号重建")
|
||||
require.Equal(t, 1, env.upstream.methodCalls(http.MethodDelete), "删除已创建的上游缓存")
|
||||
require.Empty(t, env.repo.records)
|
||||
require.Empty(t, env.usage.totals())
|
||||
}
|
||||
|
||||
func (e *cachedBillingEnv) patchTTL(id, ttl string) *httptest.ResponseRecorder {
|
||||
return e.serve(http.MethodPatch, "/v1beta/cachedContents/"+id+"?updateMask=ttl", fmt.Sprintf(`{"ttl":%q}`, ttl),
|
||||
gin.Params{{Key: "cacheID", Value: id}}, e.h.GeminiCachedContentsPatch)
|
||||
}
|
||||
|
||||
// requireTTLNear 断言发往上游的 ttl("Ns")与 want 相差不超过 5 秒。
|
||||
func requireTTLNear(t *testing.T, ttl string, want time.Duration) {
|
||||
t.Helper()
|
||||
seconds, err := strconv.ParseInt(strings.TrimSuffix(ttl, "s"), 10, 64)
|
||||
require.NoError(t, err, ttl)
|
||||
require.InDelta(t, want.Seconds(), float64(seconds), 5, ttl)
|
||||
}
|
||||
|
||||
// 同一缓存的修改串行执行:第一个 PATCH 的上游调用返回前,第二个 PATCH 直接得到 409 且不调用上游,
|
||||
// 本地到期时间与上游最后生效的修改一致,只按第一个 PATCH 实际推进的时长收费。
|
||||
func TestGeminiCachedContentsPatch_ConcurrentUpdateConflictsAndKeepsUpstreamInSync(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 1000, 900)
|
||||
env.seed(time.Now().Add(time.Hour))
|
||||
env.upstream.beforePost = func(n int) {
|
||||
if n == 1 {
|
||||
rec := env.patchTTL("abc", "86400s")
|
||||
require.Equal(t, http.StatusConflict, rec.Code, rec.Body.String())
|
||||
require.Contains(t, rec.Body.String(), `"ABORTED"`)
|
||||
}
|
||||
}
|
||||
|
||||
rec := env.patchTTL("abc", "7200s")
|
||||
require.Equal(t, http.StatusOK, rec.Code, rec.Body.String())
|
||||
require.Equal(t, 1, env.upstream.methodCalls(http.MethodPatch), "冲突的修改不调用上游")
|
||||
require.WithinDuration(t, time.Now().Add(2*time.Hour), env.repo.records["abc"].ExpireTime, 5*time.Second, "本地到期时间与上游一致")
|
||||
require.False(t, env.repo.pendingMutation(100))
|
||||
totals := env.waitForUsage(t, 1)
|
||||
require.Len(t, totals, 1)
|
||||
require.InDelta(t, cachedBillingTokens*cachedBillingStoragePrice, totals[0], cachedBillingTokens*cachedBillingStoragePrice*time.Minute.Hours(), "只收推进的 1 小时")
|
||||
}
|
||||
|
||||
// 上游已修改但本地写回失败:不计费并保留修改意图;租约内的修改返回 409,
|
||||
// 租约过期后后台恢复把上游有效期回滚到修改前的到期时间并清除意图。
|
||||
func TestGeminiCachedContentsPatch_CommitFailureRollsBackUpstreamWithoutCharge(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 1000, 900)
|
||||
baseline := time.Now().Add(time.Hour)
|
||||
env.seed(baseline)
|
||||
env.repo.failCommit = true
|
||||
|
||||
require.Equal(t, http.StatusInternalServerError, env.patch().Code)
|
||||
require.Equal(t, 1, env.upstream.methodCalls(http.MethodPatch))
|
||||
require.True(t, env.repo.pendingMutation(100), "上游可能已修改:保留修改意图")
|
||||
require.Empty(t, env.usage.totals(), "未提交的修改不收费")
|
||||
|
||||
env.repo.failCommit = false
|
||||
require.Equal(t, http.StatusConflict, env.patch().Code, "租约内的修改被拒绝")
|
||||
require.Equal(t, 1, env.upstream.methodCalls(http.MethodPatch))
|
||||
|
||||
env.repo.expireLease(100)
|
||||
recovered, err := env.h.geminiCachedContentService.RecoverStaleMutations(context.Background())
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, recovered)
|
||||
require.Equal(t, 2, env.upstream.methodCalls(http.MethodPatch))
|
||||
requireTTLNear(t, env.upstream.patchTTLs[1], time.Until(baseline))
|
||||
require.False(t, env.repo.pendingMutation(100))
|
||||
require.True(t, env.repo.records["abc"].ExpireTime.Equal(baseline), "本地到期时间保持修改前的值")
|
||||
require.Empty(t, env.usage.totals())
|
||||
}
|
||||
|
||||
// 租约过期的遗留修改意图在下一次修改前先回滚,随后的修改按修改前的到期时间计费。
|
||||
func TestGeminiCachedContentsPatch_StaleMutationRolledBackBeforeNextUpdate(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 1000, 900)
|
||||
baseline := time.Now().Add(time.Hour)
|
||||
env.seed(baseline)
|
||||
env.repo.failCommit = true
|
||||
require.Equal(t, http.StatusInternalServerError, env.patch().Code)
|
||||
env.repo.failCommit = false
|
||||
env.repo.expireLease(100)
|
||||
|
||||
rec := env.patch()
|
||||
require.Equal(t, http.StatusOK, rec.Code, rec.Body.String())
|
||||
require.Equal(t, 3, env.upstream.methodCalls(http.MethodPatch), "失败的修改、回滚、本次修改")
|
||||
requireTTLNear(t, env.upstream.patchTTLs[1], time.Until(baseline))
|
||||
require.Equal(t, "86400s", env.upstream.patchTTLs[2])
|
||||
totals := env.waitForUsage(t, 1)
|
||||
require.Len(t, totals, 1)
|
||||
require.InDelta(t, cachedBillingTokens*cachedBillingStoragePrice*23, totals[0], 0.01)
|
||||
}
|
||||
|
||||
// 上游结果未知(传输层错误)时保留修改意图等待回滚;上游明确拒绝时清除修改意图。两种情况都不收费。
|
||||
func TestGeminiCachedContentsPatch_MutationKeptOnlyWhenUpstreamOutcomeUnknown(t *testing.T) {
|
||||
t.Run("transport error", func(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 1000, 900)
|
||||
env.seed(time.Now().Add(time.Hour))
|
||||
env.upstream.transportErrors = map[string]error{http.MethodPatch: errors.New("connection reset")}
|
||||
|
||||
require.Equal(t, http.StatusBadGateway, env.patch().Code)
|
||||
require.True(t, env.repo.pendingMutation(100))
|
||||
require.Empty(t, env.usage.totals())
|
||||
})
|
||||
t.Run("upstream rejected", func(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 1000, 900)
|
||||
env.seed(time.Now().Add(time.Hour))
|
||||
env.upstream.failures = map[string]cachedBillingUpstreamFailure{http.MethodPatch: {status: http.StatusBadRequest,
|
||||
body: `{"error":{"code":400,"message":"invalid ttl","status":"INVALID_ARGUMENT"}}`}}
|
||||
|
||||
require.Equal(t, http.StatusBadRequest, env.patch().Code)
|
||||
require.False(t, env.repo.pendingMutation(100))
|
||||
require.Empty(t, env.usage.totals())
|
||||
})
|
||||
}
|
||||
|
||||
// 修改前的到期时间已过:恢复时删除上游缓存与本地记录。
|
||||
func TestGeminiCachedContentMutationRecovery_DeletesUpstreamWhenBaselineExpired(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 1000, 900)
|
||||
env.seed(time.Now().Add(-time.Minute))
|
||||
env.repo.mutations[100] = cachedBillingMutation{token: "stale", leaseUntil: time.Now().Add(-time.Second)}
|
||||
|
||||
recovered, err := env.h.geminiCachedContentService.RecoverStaleMutations(context.Background())
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, recovered)
|
||||
require.Equal(t, 1, env.upstream.methodCalls(http.MethodDelete))
|
||||
require.Zero(t, env.upstream.methodCalls(http.MethodPatch))
|
||||
require.Equal(t, []int64{100}, env.repo.deleted)
|
||||
}
|
||||
|
||||
// cachedBillingReservationCache 是启用在途预留时的计费缓存桩:余额与订阅用量取固定值,各维度预留按 key 记账。
|
||||
type cachedBillingReservationCache struct {
|
||||
service.BillingCache
|
||||
mu sync.Mutex
|
||||
balance float64
|
||||
sub *service.SubscriptionCacheData
|
||||
reserved map[string]map[string]float64
|
||||
}
|
||||
|
||||
func (c *cachedBillingReservationCache) GetUserBalance(context.Context, int64) (float64, error) {
|
||||
return c.balance, nil
|
||||
}
|
||||
|
||||
func (c *cachedBillingReservationCache) GetSubscriptionCache(context.Context, int64, int64) (*service.SubscriptionCacheData, error) {
|
||||
return c.sub, nil
|
||||
}
|
||||
|
||||
func (c *cachedBillingReservationCache) reserveLocked(key, requestID string, amount float64) float64 {
|
||||
if c.reserved == nil {
|
||||
c.reserved = map[string]map[string]float64{}
|
||||
}
|
||||
if c.reserved[key] == nil {
|
||||
c.reserved[key] = map[string]float64{}
|
||||
}
|
||||
others := 0.0
|
||||
for _, v := range c.reserved[key] {
|
||||
others += v
|
||||
}
|
||||
c.reserved[key][requestID] = amount
|
||||
return others
|
||||
}
|
||||
|
||||
func (c *cachedBillingReservationCache) ReserveInflightBalance(_ context.Context, userID int64, requestID string, amount, balance float64, _ time.Duration) (bool, float64, error) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
key := fmt.Sprintf("user:%d", userID)
|
||||
others := c.reserveLocked(key, requestID, amount)
|
||||
if len(c.reserved[key]) > 1 && balance-others < amount {
|
||||
delete(c.reserved[key], requestID)
|
||||
return false, others, nil
|
||||
}
|
||||
return true, others, nil
|
||||
}
|
||||
|
||||
func (c *cachedBillingReservationCache) ReleaseInflightBalance(_ context.Context, userID int64, requestID string) error {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
delete(c.reserved[fmt.Sprintf("user:%d", userID)], requestID)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *cachedBillingReservationCache) ReserveInflightScope(_ context.Context, scope, requestID string, amount float64, _ time.Duration) (float64, error) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
return c.reserveLocked(scope, requestID, amount), nil
|
||||
}
|
||||
|
||||
func (c *cachedBillingReservationCache) ReleaseInflightScope(_ context.Context, scope, requestID string) error {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
delete(c.reserved[scope], requestID)
|
||||
return nil
|
||||
}
|
||||
|
||||
func enableCachedBillingReservations(t *testing.T, env *cachedBillingEnv, cache *cachedBillingReservationCache) {
|
||||
t.Helper()
|
||||
cfg := &config.Config{}
|
||||
cfg.Billing.InflightReservation.Enabled = true
|
||||
bcs := service.NewBillingCacheService(cache, nil, nil, nil, nil, nil, cfg, nil)
|
||||
t.Cleanup(bcs.Stop)
|
||||
env.h.billingCacheService = bcs
|
||||
}
|
||||
|
||||
// 确定金额的续期费按 API Key 与余额维度预留:第一笔续期(约 $103.5)尚未入账时,
|
||||
// 第二个缓存的续期超出 Key 剩余额度或余额,在调用上游前被拒绝。
|
||||
func TestGeminiCachedContentsPatch_UnsettledChargeHoldsKeyQuotaAndBalance(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
balance, quota float64
|
||||
}{
|
||||
{name: "api key quota", balance: 1000, quota: 110},
|
||||
{name: "balance", balance: 110},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, tc.balance, 900)
|
||||
env.seed(time.Now().Add(time.Hour))
|
||||
second := *env.repo.records["abc"]
|
||||
second.ID, second.PublicID, second.UpstreamName = 101, "def", "cachedContents/up2"
|
||||
env.repo.records["def"] = &second
|
||||
env.apiKey.Quota = tc.quota
|
||||
enableCachedBillingReservations(t, env, &cachedBillingReservationCache{balance: tc.balance})
|
||||
env.upstream.beforePost = func(n int) {
|
||||
if n == 1 {
|
||||
rec := env.patchTTL("def", "86400s")
|
||||
require.NotEqual(t, http.StatusOK, rec.Code, rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
require.Equal(t, http.StatusOK, env.patch().Code)
|
||||
require.Equal(t, 1, env.upstream.methodCalls(http.MethodPatch), "第二笔续期没有到达上游")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// 订阅模式的确定金额费用按订阅维度预留:第一笔未入账时第二笔超出日剩余额度被拒绝,释放后放行。
|
||||
func TestReserveGeminiCachedContentCharge_UnsettledChargeHoldsSubscriptionHeadroom(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 1000, 900)
|
||||
limit := 110.0
|
||||
env.group.SubscriptionType = service.SubscriptionTypeSubscription
|
||||
env.group.DailyLimitUSD = &limit
|
||||
enableCachedBillingReservations(t, env, &cachedBillingReservationCache{balance: 1000,
|
||||
sub: &service.SubscriptionCacheData{Status: service.SubscriptionStatusActive, ExpiresAt: time.Now().Add(time.Hour)}})
|
||||
sub := &service.UserSubscription{ID: 1, UserID: env.apiKey.UserID, GroupID: env.group.ID, Status: service.SubscriptionStatusActive,
|
||||
ExpiresAt: time.Now().Add(time.Hour)}
|
||||
reserve := func() (func(), error) {
|
||||
var done func()
|
||||
var err error
|
||||
env.serve(http.MethodPatch, "/", `{}`, nil, func(c *gin.Context) {
|
||||
done, err = env.h.reserveGeminiCachedContentCharge(c, env.apiKey, sub, 103.5)
|
||||
})
|
||||
return done, err
|
||||
}
|
||||
|
||||
firstDone, err := reserve()
|
||||
require.NoError(t, err)
|
||||
_, err = reserve()
|
||||
require.ErrorIs(t, err, service.ErrDailyLimitExceeded)
|
||||
firstDone()
|
||||
secondDone, err := reserve()
|
||||
require.NoError(t, err)
|
||||
secondDone()
|
||||
}
|
||||
|
||||
func TestGeminiCachedContentsPatch_UnaffordableExtensionRejectedBeforeUpstream(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 1, 900)
|
||||
env.seed(time.Now().Add(time.Hour))
|
||||
|
||||
rec := env.patch()
|
||||
require.NotEqual(t, http.StatusOK, rec.Code)
|
||||
require.Zero(t, env.upstream.methodCalls(http.MethodPatch), "余额不足以支付延长部分时不调用上游")
|
||||
require.Empty(t, env.usage.totals())
|
||||
}
|
||||
|
||||
// 上游对缓存不存在的响应:AI Studio 403;Vertex 缓存资源接口 404,文案含上游缓存数字 ID。
|
||||
var cachedBillingUpstreamNotFoundCases = []struct {
|
||||
name string
|
||||
failure cachedBillingUpstreamFailure
|
||||
}{
|
||||
{name: "ai studio 403", failure: cachedBillingUpstreamFailure{status: http.StatusForbidden,
|
||||
body: `{"error":{"code":403,"message":"CachedContent not found (or permission denied)","status":"PERMISSION_DENIED"}}`}},
|
||||
{name: "vertex 404", failure: cachedBillingUpstreamFailure{status: http.StatusNotFound,
|
||||
body: `{"error":{"code":404,"message":"Cached content 8079767076522164224 is not found.","status":"NOT_FOUND"}}`}},
|
||||
}
|
||||
|
||||
func TestGeminiCachedContentsPatch_UpstreamNotFoundForgetsAndReturnsStandardForbidden(t *testing.T) {
|
||||
for _, tc := range cachedBillingUpstreamNotFoundCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 1000, 900)
|
||||
env.seed(time.Now().Add(time.Hour))
|
||||
env.upstream.failures = map[string]cachedBillingUpstreamFailure{http.MethodPatch: tc.failure}
|
||||
|
||||
rec := env.patch()
|
||||
require.Equal(t, http.StatusForbidden, rec.Code)
|
||||
require.Equal(t, service.GeminiCachedContentNotFoundResponse, rec.Body.String())
|
||||
require.NotContains(t, rec.Body.String(), "8079767076522164224", "不应暴露上游缓存 ID")
|
||||
require.Equal(t, []int64{100}, env.repo.deleted, "上游已不存在的缓存应从本地移除")
|
||||
require.Empty(t, env.usage.totals(), "上游未延长不收费")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeminiCachedContentsDelete_UpstreamNotFoundStillForgets(t *testing.T) {
|
||||
for _, tc := range cachedBillingUpstreamNotFoundCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 1000, 900)
|
||||
env.seed(time.Now().Add(time.Hour))
|
||||
env.upstream.failures = map[string]cachedBillingUpstreamFailure{http.MethodDelete: tc.failure}
|
||||
|
||||
rec := env.serve(http.MethodDelete, "/v1beta/cachedContents/abc", "",
|
||||
gin.Params{{Key: "cacheID", Value: "abc"}}, env.h.GeminiCachedContentsDelete)
|
||||
require.Equal(t, http.StatusOK, rec.Code, rec.Body.String())
|
||||
require.JSONEq(t, `{}`, rec.Body.String())
|
||||
require.Equal(t, []int64{100}, env.repo.deleted)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// 持有账号暂时不支持显式缓存(能力被关闭)时,恢复流程不能把修改意图当作已回滚而清除;
|
||||
// 能力恢复后再回滚上游并清除意图。
|
||||
func TestGeminiCachedContentMutationRecovery_KeepsIntentWhileAccountCannotCache(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 1000, 900)
|
||||
env.accounts[900].Credentials["base_url"] = "https://gw.example.com"
|
||||
env.accounts[900].Credentials[service.GeminiExplicitCacheUpstreamCredentialKey] = true
|
||||
baseline := time.Now().Add(time.Hour)
|
||||
env.seed(baseline)
|
||||
env.repo.failCommit = true
|
||||
require.Equal(t, http.StatusInternalServerError, env.patch().Code)
|
||||
require.True(t, env.repo.pendingMutation(100))
|
||||
require.Equal(t, 1, env.upstream.methodCalls(http.MethodPatch))
|
||||
env.repo.failCommit = false
|
||||
env.repo.expireLease(100)
|
||||
|
||||
env.accounts[900].Credentials[service.GeminiExplicitCacheUpstreamCredentialKey] = false
|
||||
recovered, err := env.h.geminiCachedContentService.RecoverStaleMutations(context.Background())
|
||||
require.Error(t, err, "账号暂不支持缓存时恢复失败")
|
||||
require.Zero(t, recovered)
|
||||
require.True(t, env.repo.pendingMutation(100), "未回滚上游前保留修改意图")
|
||||
require.Equal(t, 1, env.upstream.methodCalls(http.MethodPatch), "不改动上游")
|
||||
require.Empty(t, env.repo.deleted)
|
||||
|
||||
env.accounts[900].Credentials[service.GeminiExplicitCacheUpstreamCredentialKey] = true
|
||||
env.repo.expireLease(100)
|
||||
recovered, err = env.h.geminiCachedContentService.RecoverStaleMutations(context.Background())
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, recovered)
|
||||
require.Equal(t, 2, env.upstream.methodCalls(http.MethodPatch))
|
||||
requireTTLNear(t, env.upstream.patchTTLs[1], time.Until(baseline))
|
||||
require.False(t, env.repo.pendingMutation(100))
|
||||
require.True(t, env.repo.records["abc"].ExpireTime.Equal(baseline))
|
||||
require.Empty(t, env.usage.totals())
|
||||
}
|
||||
|
||||
// 持有账号长期不支持显式缓存,且修改前到期时间加上单次最长有效期都已过去:上游缓存必已自行过期,
|
||||
// 恢复流程不再保留意图,直接移除本地记录。
|
||||
func TestGeminiCachedContentMutationRecovery_AbandonsIntentOnceUpstreamMustHaveExpired(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 1000, 900)
|
||||
env.accounts[900].Credentials["base_url"] = "https://gw.example.com"
|
||||
env.accounts[900].Credentials[service.GeminiExplicitCacheUpstreamCredentialKey] = true
|
||||
env.seed(time.Now().Add(time.Hour))
|
||||
env.repo.failCommit = true
|
||||
require.Equal(t, http.StatusInternalServerError, env.patch().Code)
|
||||
env.repo.failCommit = false
|
||||
env.repo.expireLease(100)
|
||||
env.accounts[900].Credentials[service.GeminiExplicitCacheUpstreamCredentialKey] = false
|
||||
|
||||
// 修改前到期时间 + 24h 上限 + 1h 容忍仍未过去:保留意图。
|
||||
env.repo.records["abc"].ExpireTime = time.Now().Add(-20 * time.Hour)
|
||||
recovered, err := env.h.geminiCachedContentService.RecoverStaleMutations(context.Background())
|
||||
require.Error(t, err)
|
||||
require.Zero(t, recovered)
|
||||
require.True(t, env.repo.pendingMutation(100))
|
||||
require.Empty(t, env.repo.deleted)
|
||||
|
||||
env.repo.expireLease(100)
|
||||
env.repo.records["abc"].ExpireTime = time.Now().Add(-26 * time.Hour)
|
||||
recovered, err = env.h.geminiCachedContentService.RecoverStaleMutations(context.Background())
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, recovered)
|
||||
require.Equal(t, 1, env.upstream.methodCalls(http.MethodPatch), "不改动上游")
|
||||
require.Zero(t, env.upstream.methodCalls(http.MethodDelete))
|
||||
require.Equal(t, []int64{100}, env.repo.deleted)
|
||||
require.Empty(t, env.usage.totals())
|
||||
|
||||
recovered, err = env.h.geminiCachedContentService.RecoverStaleMutations(context.Background())
|
||||
require.NoError(t, err)
|
||||
require.Zero(t, recovered, "作废的记录不再被选中")
|
||||
}
|
||||
|
||||
// 作废判定的边界:默认上限 24h 加 1h 容忍;配置上限低于 24h 时仍按 24h 兜底。
|
||||
func TestGeminiCachedContentMutationRecovery_AbandonBoundary(t *testing.T) {
|
||||
disableCache := func(env *cachedBillingEnv) {
|
||||
env.accounts[900].Credentials["base_url"] = "https://gw.example.com"
|
||||
env.accounts[900].Credentials[service.GeminiExplicitCacheUpstreamCredentialKey] = false
|
||||
}
|
||||
t.Run("within grace keeps intent", func(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 1000, 900)
|
||||
env.seed(time.Now().Add(-25*time.Hour + time.Minute))
|
||||
env.repo.mutations[100] = cachedBillingMutation{token: "stale", leaseUntil: time.Now().Add(-time.Second)}
|
||||
disableCache(env)
|
||||
|
||||
recovered, err := env.h.geminiCachedContentService.RecoverStaleMutations(context.Background())
|
||||
require.Error(t, err)
|
||||
require.Zero(t, recovered)
|
||||
require.True(t, env.repo.pendingMutation(100))
|
||||
require.Empty(t, env.repo.deleted)
|
||||
})
|
||||
t.Run("configured max ttl below default still waits a day", func(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 1000, 900)
|
||||
env.seed(time.Now().Add(-12 * time.Hour))
|
||||
env.repo.mutations[100] = cachedBillingMutation{token: "stale", leaseUntil: time.Now().Add(-time.Second)}
|
||||
disableCache(env)
|
||||
cfg := &config.Config{}
|
||||
cfg.Gateway.GeminiCachedContentMaxTTLSeconds = 3600
|
||||
compat := service.NewGeminiMessagesCompatService(&cachedBillingAccountRepo{accounts: env.accounts}, nil, nil, nil, nil, nil, env.upstream, nil, cfg)
|
||||
svc := service.NewGeminiCachedContentService(env.repo, compat, cfg)
|
||||
|
||||
recovered, err := svc.RecoverStaleMutations(context.Background())
|
||||
require.Error(t, err)
|
||||
require.Zero(t, recovered)
|
||||
require.True(t, env.repo.pendingMutation(100))
|
||||
|
||||
env.repo.expireLease(100)
|
||||
env.repo.records["abc"].ExpireTime = time.Now().Add(-26 * time.Hour)
|
||||
recovered, err = svc.RecoverStaleMutations(context.Background())
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, recovered)
|
||||
require.Equal(t, []int64{100}, env.repo.deleted)
|
||||
})
|
||||
}
|
||||
|
||||
// 持有账号可用但上游回滚持续失败(如 Key 已失效):上游缓存必已自行过期后不再重试,移除本地记录。
|
||||
func TestGeminiCachedContentMutationRecovery_AbandonsIntentWhenRollbackKeepsFailing(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 1000, 900)
|
||||
env.seed(time.Now().Add(-time.Minute))
|
||||
env.repo.mutations[100] = cachedBillingMutation{token: "stale", leaseUntil: time.Now().Add(-time.Second)}
|
||||
env.upstream.failures = map[string]cachedBillingUpstreamFailure{http.MethodDelete: {status: http.StatusUnauthorized,
|
||||
body: `{"error":{"code":401,"message":"API key not valid","status":"UNAUTHENTICATED"}}`}}
|
||||
|
||||
recovered, err := env.h.geminiCachedContentService.RecoverStaleMutations(context.Background())
|
||||
require.Error(t, err, "上游缓存可能仍存活:保留意图等待重试")
|
||||
require.Zero(t, recovered)
|
||||
require.Equal(t, 1, env.upstream.methodCalls(http.MethodDelete))
|
||||
require.True(t, env.repo.pendingMutation(100))
|
||||
require.Empty(t, env.repo.deleted)
|
||||
|
||||
env.repo.expireLease(100)
|
||||
env.repo.records["abc"].ExpireTime = time.Now().Add(-26 * time.Hour)
|
||||
recovered, err = env.h.geminiCachedContentService.RecoverStaleMutations(context.Background())
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, recovered)
|
||||
require.Equal(t, 2, env.upstream.methodCalls(http.MethodDelete), "账号可用时仍先尝试回滚")
|
||||
require.Equal(t, []int64{100}, env.repo.deleted)
|
||||
require.Empty(t, env.usage.totals())
|
||||
}
|
||||
|
||||
// 持有账号已删除时没有凭据可回滚上游:恢复流程移除本地记录并清除意图,不再重试。
|
||||
func TestGeminiCachedContentMutationRecovery_ForgetsRecordWhenAccountDeleted(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 1000, 900)
|
||||
env.seed(time.Now().Add(time.Hour))
|
||||
env.repo.failCommit = true
|
||||
require.Equal(t, http.StatusInternalServerError, env.patch().Code)
|
||||
require.True(t, env.repo.pendingMutation(100))
|
||||
env.repo.failCommit = false
|
||||
env.repo.expireLease(100)
|
||||
|
||||
delete(env.accounts, 900)
|
||||
recovered, err := env.h.geminiCachedContentService.RecoverStaleMutations(context.Background())
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, recovered)
|
||||
require.Equal(t, 1, env.upstream.methodCalls(http.MethodPatch), "无凭据,不改动上游")
|
||||
require.Zero(t, env.upstream.methodCalls(http.MethodDelete))
|
||||
require.Equal(t, []int64{100}, env.repo.deleted, "本地记录作废")
|
||||
require.Empty(t, env.usage.totals())
|
||||
require.Equal(t, http.StatusForbidden, env.serve(http.MethodGet, "/v1beta/cachedContents/abc", "", gin.Params{{Key: "cacheID", Value: "abc"}}, env.h.GeminiCachedContentsGet).Code, "作废的记录对客户端不可见")
|
||||
|
||||
recovered, err = env.h.geminiCachedContentService.RecoverStaleMutations(context.Background())
|
||||
require.NoError(t, err)
|
||||
require.Zero(t, recovered, "不再重复处理")
|
||||
}
|
||||
@@ -0,0 +1,848 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
pkghttputil "github.com/Wei-Shaw/sub2api/internal/pkg/httputil"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ip"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||||
"github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// geminiCachedContentCreateMaxAttempts 限制创建缓存时实际发往上游的账号数。
|
||||
const geminiCachedContentCreateMaxAttempts = 3
|
||||
|
||||
// geminiCachedContentPersistTimeout 是上游操作完成后落库的超时;落库与客户端连接脱钩。
|
||||
const geminiCachedContentPersistTimeout = 15 * time.Second
|
||||
|
||||
const geminiCachedContentUnavailableMessage = "The account holding this cached content is temporarily unavailable"
|
||||
|
||||
// geminiCachedContentsPreamble 校验显式缓存管理接口的公共前置条件:gemini 分组、非 composite 分组、功能已装配。
|
||||
func (h *GatewayHandler) geminiCachedContentsPreamble(c *gin.Context, component string) (*service.APIKey, middleware.AuthSubject, *zap.Logger, bool) {
|
||||
apiKey, ok := middleware.GetAPIKeyFromContext(c)
|
||||
if !ok || apiKey == nil {
|
||||
googleError(c, http.StatusUnauthorized, "Invalid API key")
|
||||
return nil, middleware.AuthSubject{}, nil, false
|
||||
}
|
||||
authSubject, ok := middleware.GetAuthSubjectFromContext(c)
|
||||
if !ok {
|
||||
googleError(c, http.StatusInternalServerError, "User context not found")
|
||||
return nil, middleware.AuthSubject{}, nil, false
|
||||
}
|
||||
if h.geminiCachedContentService == nil {
|
||||
googlePlatformError(c, http.StatusNotFound, "Explicit context caching is not enabled")
|
||||
return nil, middleware.AuthSubject{}, nil, false
|
||||
}
|
||||
if apiKey.Group == nil || apiKey.GroupID == nil || apiKey.Group.Platform == service.PlatformComposite {
|
||||
googleError(c, http.StatusBadRequest, "Explicit context caching is not supported for this API key group")
|
||||
return nil, middleware.AuthSubject{}, nil, false
|
||||
}
|
||||
if effectiveAPIKeyPlatform(c, apiKey) != service.PlatformGemini {
|
||||
googleError(c, http.StatusBadRequest, "API key group platform is not gemini")
|
||||
return nil, middleware.AuthSubject{}, nil, false
|
||||
}
|
||||
reqLog := requestLogger(
|
||||
c,
|
||||
component,
|
||||
zap.Int64("user_id", authSubject.UserID),
|
||||
zap.Int64("api_key_id", apiKey.ID),
|
||||
zap.Any("group_id", apiKey.GroupID),
|
||||
)
|
||||
return apiKey, authSubject, reqLog, true
|
||||
}
|
||||
|
||||
func readGeminiCachedContentBody(c *gin.Context) ([]byte, bool) {
|
||||
body, err := pkghttputil.ReadRequestBodyWithPrealloc(c.Request)
|
||||
if err != nil {
|
||||
if maxErr, ok := extractMaxBytesError(err); ok {
|
||||
googleError(c, http.StatusRequestEntityTooLarge, buildBodyTooLargeMessage(maxErr.Limit))
|
||||
return nil, false
|
||||
}
|
||||
googleError(c, http.StatusBadRequest, "Failed to read request body")
|
||||
return nil, false
|
||||
}
|
||||
if len(body) == 0 {
|
||||
googleError(c, http.StatusBadRequest, "Request body is empty")
|
||||
return nil, false
|
||||
}
|
||||
return body, true
|
||||
}
|
||||
|
||||
func (h *GatewayHandler) checkGeminiCachedContentBilling(c *gin.Context, reqLog *zap.Logger, apiKey *service.APIKey, subscription *service.UserSubscription) bool {
|
||||
if err := h.billingCacheService.CheckBillingEligibility(c.Request.Context(), apiKey.User, apiKey, apiKey.Group, subscription, service.QuotaPlatform(c.Request.Context(), apiKey)); err != nil {
|
||||
reqLog.Info("gemini.cached_content.billing_eligibility_check_failed", zap.Error(err))
|
||||
status, code, message, retryAfter := billingErrorDetails(err)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
googleErrorWithType(c, status, code, "", message)
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// GeminiCachedContentsCreate POST /v1beta/cachedContents
|
||||
// 创建请求内一次性收取缓存 token 的输入费与整个有效期的存储费。
|
||||
func (h *GatewayHandler) GeminiCachedContentsCreate(c *gin.Context) {
|
||||
apiKey, authSubject, reqLog, ok := h.geminiCachedContentsPreamble(c, "handler.gemini_v1beta.cached_contents.create")
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
svc := h.geminiCachedContentService
|
||||
|
||||
body, ok := readGeminiCachedContentBody(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
req, err := service.ParseGeminiCachedContentCreateRequest(body, svc.Now(), svc.MaxTTL())
|
||||
if err != nil {
|
||||
googleError(c, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
modelName := req.Model
|
||||
reqLog = reqLog.With(zap.String("model", modelName))
|
||||
|
||||
setOpsRequestContext(c, modelName, false)
|
||||
setOpsEndpointContext(c, "", int16(service.RequestTypeFromLegacy(false, false)))
|
||||
pricingCtx, pricingAt := service.WithGatewayTokenRequestPricing(c.Request.Context())
|
||||
c.Request = c.Request.WithContext(pricingCtx)
|
||||
|
||||
if decision := h.checkSecurityAudit(c, reqLog, apiKey, authSubject, service.ContentModerationProtocolGemini, modelName, body); decision != nil && !decision.AllowNextStage {
|
||||
googleSecurityAuditError(c, decision)
|
||||
return
|
||||
}
|
||||
|
||||
channelMapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(c.Request.Context(), apiKey.GroupID, modelName)
|
||||
reqModel := modelName
|
||||
if channelMapping.Mapped {
|
||||
modelName = channelMapping.MappedModel
|
||||
}
|
||||
|
||||
subscription, _ := middleware.GetSubscriptionFromContext(c)
|
||||
geminiConcurrency := NewConcurrencyHelper(h.concurrencyHelper.concurrencyService, SSEPingFormatNone, 0)
|
||||
streamStarted := false
|
||||
userReleaseFunc, err := geminiConcurrency.AcquireUserSlotWithWait(c, authSubject.UserID, authSubject.Concurrency, false, &streamStarted)
|
||||
if err != nil {
|
||||
reqLog.Warn("gemini.cached_content.user_slot_acquire_failed", zap.Error(err))
|
||||
googleConcurrencyError(c, err, "user")
|
||||
return
|
||||
}
|
||||
userReleaseFunc = wrapReleaseOnDone(c.Request.Context(), userReleaseFunc)
|
||||
if userReleaseFunc != nil {
|
||||
defer userReleaseFunc()
|
||||
}
|
||||
|
||||
if !h.checkGeminiCachedContentBilling(c, reqLog, apiKey, subscription) {
|
||||
return
|
||||
}
|
||||
inflightRelease, err := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(modelName, body))
|
||||
if err != nil {
|
||||
reqLog.Info("gemini.cached_content.inflight_reservation_rejected", zap.Error(err))
|
||||
status, code, message, retryAfter := billingErrorDetails(err)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
googleErrorWithType(c, status, code, "", message)
|
||||
return
|
||||
}
|
||||
defer inflightRelease()
|
||||
|
||||
excluded, err := h.gatewayService.GeminiExplicitCacheExclusions(c.Request.Context(), apiKey.GroupID, 0)
|
||||
if err != nil {
|
||||
reqLog.Warn("gemini.cached_content.list_accounts_failed", zap.Error(err))
|
||||
googleError(c, http.StatusServiceUnavailable, "No available Gemini accounts support explicit context caching")
|
||||
return
|
||||
}
|
||||
|
||||
var lastUpstreamErr *service.GeminiCachedContentUpstreamError
|
||||
var ttlRejectedLimit time.Duration
|
||||
for attempts := 0; attempts < geminiCachedContentCreateMaxAttempts; {
|
||||
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, "", modelName, excluded, "", int64(0))
|
||||
if err != nil {
|
||||
if failoverClientGone(c) {
|
||||
return
|
||||
}
|
||||
reqLog.Info("gemini.cached_content.account_select_exhausted", zap.Error(err))
|
||||
break
|
||||
}
|
||||
account := selection.Account
|
||||
if !service.IsGeminiExplicitCacheAccount(account) {
|
||||
if selection.Acquired && selection.ReleaseFunc != nil {
|
||||
selection.ReleaseFunc()
|
||||
}
|
||||
excluded[account.ID] = struct{}{}
|
||||
continue
|
||||
}
|
||||
ttl, ttlAllowed := req.TTLForAccount(account)
|
||||
if !ttlAllowed {
|
||||
if selection.Acquired && selection.ReleaseFunc != nil {
|
||||
selection.ReleaseFunc()
|
||||
}
|
||||
excluded[account.ID] = struct{}{}
|
||||
ttlRejectedLimit = max(ttlRejectedLimit, account.GeminiExplicitCacheUpstreamMaxTTL())
|
||||
continue
|
||||
}
|
||||
accountReleaseFunc, ok := h.acquireGeminiCachedContentAccountSlot(c, reqLog, geminiConcurrency, selection)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
admissionCtx := service.ContextWithSelectionProfitGate(c.Request.Context(), selection)
|
||||
latest, vetoed, reason := h.gatewayService.GatewayProfitControlVetoLatest(admissionCtx, account)
|
||||
if vetoed {
|
||||
accountReleaseFunc()
|
||||
reqLog.Debug("gemini.cached_content.account_profit_vetoed", zap.Int64("account_id", account.ID), zap.String("reason", reason))
|
||||
excluded[account.ID] = struct{}{}
|
||||
continue
|
||||
}
|
||||
account = latest
|
||||
attempts++
|
||||
setOpsSelectedAccount(c, account.ID, account.Platform)
|
||||
|
||||
mappedModel := account.GetMappedModel(modelName)
|
||||
usageFields := clientRequestedUsageFields(c, channelMapping, reqModel, mappedModel)
|
||||
if !h.gatewayService.CacheStoragePriced(c.Request.Context(), &service.RecordUsageInput{
|
||||
Result: &service.ForwardResult{Model: modelName, UpstreamModel: mappedModel, CacheStorageTokenHours: 1},
|
||||
APIKey: apiKey,
|
||||
ChannelUsageFields: usageFields,
|
||||
}) {
|
||||
accountReleaseFunc()
|
||||
googlePlatformError(c, http.StatusBadRequest, fmt.Sprintf("Explicit context caching is not available for model %s: storage pricing is not configured", modelName))
|
||||
return
|
||||
}
|
||||
upstreamModel, err := service.GeminiCachedContentUpstreamModelResource(account, mappedModel)
|
||||
if err != nil {
|
||||
accountReleaseFunc()
|
||||
reqLog.Warn("gemini.cached_content.upstream_model_invalid", zap.Int64("account_id", account.ID), zap.Error(err))
|
||||
excluded[account.ID] = struct{}{}
|
||||
continue
|
||||
}
|
||||
publicID, err := service.NewGeminiCachedContentPublicID()
|
||||
if err != nil {
|
||||
accountReleaseFunc()
|
||||
googleError(c, http.StatusInternalServerError, "Failed to allocate cached content id")
|
||||
return
|
||||
}
|
||||
|
||||
upstream, err := svc.CreateUpstream(c.Request.Context(), account, req, mappedModel, publicID, ttl)
|
||||
accountReleaseFunc()
|
||||
if err != nil {
|
||||
var upErr *service.GeminiCachedContentUpstreamError
|
||||
if errors.As(err, &upErr) {
|
||||
reqLog.Info("gemini.cached_content.upstream_create_failed", zap.Int64("account_id", account.ID), zap.Int("status", upErr.StatusCode))
|
||||
if upErr.RetryableOnOtherAccount() {
|
||||
lastUpstreamErr = upErr
|
||||
excluded[account.ID] = struct{}{}
|
||||
continue
|
||||
}
|
||||
h.writeGeminiCachedContentUpstreamError(c, upErr)
|
||||
return
|
||||
}
|
||||
// 上游已返回成功但响应不可用:不换号重建,能识别出资源名时删除上游缓存。
|
||||
var malformed *service.GeminiCachedContentMalformedResponseError
|
||||
if errors.As(err, &malformed) {
|
||||
reqLog.Error("gemini.cached_content.upstream_create_malformed", zap.Int64("account_id", account.ID), zap.Error(err))
|
||||
if malformed.Name != "" {
|
||||
h.cleanupGeminiCachedContentUpstream(reqLog, account, malformed.Name)
|
||||
}
|
||||
googleError(c, http.StatusBadGateway, "Upstream returned an invalid cached content response")
|
||||
return
|
||||
}
|
||||
if failoverClientGone(c) {
|
||||
return
|
||||
}
|
||||
reqLog.Warn("gemini.cached_content.upstream_create_error", zap.Int64("account_id", account.ID), zap.Error(err))
|
||||
excluded[account.ID] = struct{}{}
|
||||
continue
|
||||
}
|
||||
|
||||
// 缓存 token 数在上游创建后才确定:按确定金额校验余额 / 额度并预留,不足时删除上游缓存,
|
||||
// 只收取已发生的缓存 token 输入费。
|
||||
tokens := upstream.TotalTokenCount
|
||||
storageTokenHours := service.GeminiCachedContentStorageTokenHours(tokens, ttl)
|
||||
amount := h.gatewayService.QuoteUsageCost(c.Request.Context(), &service.RecordUsageInput{
|
||||
Result: &service.ForwardResult{
|
||||
Model: modelName, UpstreamModel: mappedModel,
|
||||
Usage: service.ClaudeUsage{InputTokens: int(tokens)}, CacheStorageTokenHours: storageTokenHours,
|
||||
},
|
||||
APIKey: apiKey, User: apiKey.User, Subscription: subscription, PricingAt: pricingAt, ChannelUsageFields: usageFields,
|
||||
})
|
||||
inflightRelease()
|
||||
chargeDone, err := h.reserveGeminiCachedContentCharge(c, apiKey, subscription, amount)
|
||||
if err != nil {
|
||||
reqLog.Info("gemini.cached_content.storage_charge_rejected", zap.Float64("amount", amount), zap.Error(err))
|
||||
h.cleanupGeminiCachedContentUpstream(reqLog, account, upstream.Name)
|
||||
h.submitGeminiCachedContentUsage(c, apiKey, account, subscription, pricingAt,
|
||||
service.GeminiCachedContentCreateUsageRequestID(publicID), modelName, mappedModel, int(tokens), 0, usageFields, body)
|
||||
writeGeminiCachedContentBillingError(c, err)
|
||||
return
|
||||
}
|
||||
defer chargeDone()
|
||||
|
||||
expireTime := service.ClampGeminiCachedContentExpire(upstream.ExpireTime, svc.Now().Add(ttl))
|
||||
record := &service.GeminiCachedContent{
|
||||
PublicID: publicID,
|
||||
UserID: authSubject.UserID,
|
||||
APIKeyID: apiKey.ID,
|
||||
GroupID: *apiKey.GroupID,
|
||||
AccountID: account.ID,
|
||||
UpstreamName: upstream.Name,
|
||||
Model: modelName,
|
||||
UpstreamModel: upstreamModel,
|
||||
DisplayName: req.DisplayName,
|
||||
TotalTokenCount: tokens,
|
||||
ExpireTime: expireTime,
|
||||
ChannelUsage: usageFields,
|
||||
}
|
||||
persistCtx, cancelPersist := geminiCachedContentPersistContext(c)
|
||||
err = svc.Save(persistCtx, record)
|
||||
cancelPersist()
|
||||
if err != nil {
|
||||
reqLog.Error("gemini.cached_content.save_failed", zap.Int64("account_id", account.ID), zap.Error(err))
|
||||
h.cleanupGeminiCachedContentUpstream(reqLog, account, record.UpstreamName)
|
||||
googleError(c, http.StatusInternalServerError, "Failed to persist cached content")
|
||||
return
|
||||
}
|
||||
|
||||
h.submitGeminiCachedContentUsage(c, apiKey, account, subscription, pricingAt,
|
||||
service.GeminiCachedContentCreateUsageRequestID(publicID), modelName, mappedModel, int(tokens), storageTokenHours, usageFields, body)
|
||||
c.JSON(http.StatusOK, service.GeminiCachedContentView(record))
|
||||
return
|
||||
}
|
||||
|
||||
if lastUpstreamErr != nil {
|
||||
h.writeGeminiCachedContentUpstreamError(c, lastUpstreamErr)
|
||||
return
|
||||
}
|
||||
if ttlRejectedLimit > 0 {
|
||||
googleError(c, http.StatusBadRequest, fmt.Sprintf("ttl exceeds the maximum allowed %ds by the upstream accounts", int64(ttlRejectedLimit.Seconds())))
|
||||
return
|
||||
}
|
||||
markOpsRoutingCapacityLimited(c)
|
||||
googleError(c, http.StatusServiceUnavailable, "No available Gemini accounts support explicit context caching")
|
||||
}
|
||||
|
||||
// GeminiCachedContentsList GET /v1beta/cachedContents
|
||||
// 只返回当前 API Key 在当前分组下仍有效的缓存,不透传上游 list(同一账号上其他租户的缓存不可见)。
|
||||
func (h *GatewayHandler) GeminiCachedContentsList(c *gin.Context) {
|
||||
apiKey, _, reqLog, ok := h.geminiCachedContentsPreamble(c, "handler.gemini_v1beta.cached_contents.list")
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
pageSize := service.GeminiCachedContentListDefaultPageSize
|
||||
if raw := strings.TrimSpace(c.Query("pageSize")); raw != "" {
|
||||
n, err := strconv.Atoi(raw)
|
||||
if err != nil || n < 0 {
|
||||
googleError(c, http.StatusBadRequest, "pageSize must be a non-negative integer")
|
||||
return
|
||||
}
|
||||
if n > 0 {
|
||||
pageSize = min(n, service.GeminiCachedContentListMaxPageSize)
|
||||
}
|
||||
}
|
||||
beforeID := int64(0)
|
||||
if raw := strings.TrimSpace(c.Query("pageToken")); raw != "" {
|
||||
id, ok := decodeGeminiCachedContentPageToken(raw)
|
||||
if !ok {
|
||||
googleError(c, http.StatusBadRequest, "Invalid pageToken")
|
||||
return
|
||||
}
|
||||
beforeID = id
|
||||
}
|
||||
records, err := h.geminiCachedContentService.List(c.Request.Context(), apiKey.ID, *apiKey.GroupID, beforeID, pageSize+1)
|
||||
if err != nil {
|
||||
reqLog.Error("gemini.cached_content.list_failed", zap.Error(err))
|
||||
googleError(c, http.StatusInternalServerError, "Failed to list cached contents")
|
||||
return
|
||||
}
|
||||
resp := gin.H{}
|
||||
if len(records) > pageSize {
|
||||
records = records[:pageSize]
|
||||
resp["nextPageToken"] = encodeGeminiCachedContentPageToken(records[len(records)-1].ID)
|
||||
}
|
||||
if len(records) > 0 {
|
||||
views := make([]map[string]any, 0, len(records))
|
||||
for _, record := range records {
|
||||
views = append(views, service.GeminiCachedContentView(record))
|
||||
}
|
||||
resp["cachedContents"] = views
|
||||
}
|
||||
c.JSON(http.StatusOK, resp)
|
||||
}
|
||||
|
||||
// GeminiCachedContentsGet GET /v1beta/cachedContents/:cacheID
|
||||
func (h *GatewayHandler) GeminiCachedContentsGet(c *gin.Context) {
|
||||
apiKey, _, reqLog, ok := h.geminiCachedContentsPreamble(c, "handler.gemini_v1beta.cached_contents.get")
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
record, ok := h.resolveGeminiCachedContentOrWrite(c, reqLog, apiKey, c.Param("cacheID"))
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, service.GeminiCachedContentView(record))
|
||||
}
|
||||
|
||||
// GeminiCachedContentsPatch PATCH /v1beta/cachedContents/:cacheID
|
||||
// 只允许修改有效期;延长部分在本次请求内按增量时长收取存储费,缩短不退费。
|
||||
// 同一缓存的修改串行执行:上游修改、本地写回与计费都在修改意图的租约内完成,并发的修改返回 409。
|
||||
// 上游结果未知或本地写回失败的修改不计费,租约过期后由恢复流程把上游有效期回滚到修改前的到期时间。
|
||||
func (h *GatewayHandler) GeminiCachedContentsPatch(c *gin.Context) {
|
||||
apiKey, _, reqLog, ok := h.geminiCachedContentsPreamble(c, "handler.gemini_v1beta.cached_contents.patch")
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
svc := h.geminiCachedContentService
|
||||
body, ok := readGeminiCachedContentBody(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
record, ok := h.resolveGeminiCachedContentOrWrite(c, reqLog, apiKey, c.Param("cacheID"))
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
account, err := svc.BoundAccount(c.Request.Context(), record)
|
||||
if err != nil {
|
||||
reqLog.Error("gemini.cached_content.load_account_failed", zap.Int64("account_id", record.AccountID), zap.Error(err))
|
||||
googleError(c, http.StatusInternalServerError, "Failed to load cached content account")
|
||||
return
|
||||
}
|
||||
if account == nil {
|
||||
googleError(c, http.StatusServiceUnavailable, geminiCachedContentUnavailableMessage)
|
||||
return
|
||||
}
|
||||
ttl, err := service.ParseGeminiCachedContentPatchRequest(body, c.Query("updateMask"), svc.Now(), svc.MaxTTLForAccount(account))
|
||||
if err != nil {
|
||||
googleError(c, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
setOpsRequestContext(c, record.Model, false)
|
||||
pricingCtx, pricingAt := service.WithGatewayTokenRequestPricing(c.Request.Context())
|
||||
c.Request = c.Request.WithContext(pricingCtx)
|
||||
|
||||
subscription, _ := middleware.GetSubscriptionFromContext(c)
|
||||
if extension := svc.Now().Add(ttl).Sub(record.ExpireTime); extension > 0 {
|
||||
usage := &service.RecordUsageInput{
|
||||
Result: &service.ForwardResult{
|
||||
Model: record.Model, UpstreamModel: record.UpstreamModelID(),
|
||||
CacheStorageTokenHours: service.GeminiCachedContentStorageTokenHours(record.TotalTokenCount, extension),
|
||||
},
|
||||
APIKey: apiKey, User: apiKey.User, Subscription: subscription, PricingAt: pricingAt,
|
||||
ChannelUsageFields: record.ChannelUsage,
|
||||
}
|
||||
if !h.gatewayService.CacheStoragePriced(c.Request.Context(), usage) {
|
||||
googlePlatformError(c, http.StatusBadRequest, fmt.Sprintf("Explicit context caching is not available for model %s: storage pricing is not configured", record.Model))
|
||||
return
|
||||
}
|
||||
if !h.checkGeminiCachedContentBilling(c, reqLog, apiKey, subscription) {
|
||||
return
|
||||
}
|
||||
chargeDone, err := h.reserveGeminiCachedContentCharge(c, apiKey, subscription, h.gatewayService.QuoteUsageCost(c.Request.Context(), usage))
|
||||
if err != nil {
|
||||
reqLog.Info("gemini.cached_content.storage_charge_rejected", zap.Error(err))
|
||||
writeGeminiCachedContentBillingError(c, err)
|
||||
return
|
||||
}
|
||||
defer chargeDone()
|
||||
}
|
||||
|
||||
setOpsSelectedAccount(c, account.ID, account.Platform)
|
||||
mutation, err := svc.BeginMutation(c.Request.Context(), record)
|
||||
if err != nil {
|
||||
switch {
|
||||
case errors.Is(err, service.ErrGeminiCachedContentNotFound):
|
||||
googleError(c, http.StatusForbidden, service.GeminiCachedContentNotFoundMessage)
|
||||
case errors.Is(err, service.ErrGeminiCachedContentMutationConflict):
|
||||
reqLog.Info("gemini.cached_content.mutation_conflict", zap.Int64("account_id", account.ID), zap.Error(err))
|
||||
writeGeminiCachedContentMutationConflict(c)
|
||||
default:
|
||||
reqLog.Error("gemini.cached_content.begin_mutation_failed", zap.Int64("account_id", account.ID), zap.Error(err))
|
||||
googleError(c, http.StatusInternalServerError, "Failed to update cached content")
|
||||
}
|
||||
return
|
||||
}
|
||||
upstream, err := svc.PatchUpstream(c.Request.Context(), account, record, ttl)
|
||||
if err != nil {
|
||||
var upErr *service.GeminiCachedContentUpstreamError
|
||||
upstreamRejected := errors.As(err, &upErr)
|
||||
if upstreamRejected || errors.Is(err, service.ErrGeminiCachedContentRequestNotSent) {
|
||||
// 上游有效期未变:清除修改意图。
|
||||
persistCtx, cancelPersist := geminiCachedContentPersistContext(c)
|
||||
if releaseErr := svc.ReleaseMutation(persistCtx, mutation); releaseErr != nil {
|
||||
reqLog.Warn("gemini.cached_content.release_mutation_failed", zap.Error(releaseErr))
|
||||
}
|
||||
cancelPersist()
|
||||
}
|
||||
if upstreamRejected {
|
||||
if upErr.NotFound() {
|
||||
if forgetErr := svc.Forget(c.Request.Context(), record); forgetErr != nil {
|
||||
reqLog.Warn("gemini.cached_content.forget_failed", zap.Error(forgetErr))
|
||||
}
|
||||
}
|
||||
h.writeGeminiCachedContentUpstreamError(c, upErr)
|
||||
return
|
||||
}
|
||||
reqLog.Warn("gemini.cached_content.upstream_patch_error", zap.Int64("account_id", account.ID), zap.Error(err))
|
||||
googleError(c, http.StatusBadGateway, "Upstream request failed")
|
||||
return
|
||||
}
|
||||
expireTime := service.ClampGeminiCachedContentExpire(upstream.ExpireTime, svc.Now().Add(ttl))
|
||||
persistCtx, cancelPersist := geminiCachedContentPersistContext(c)
|
||||
err = svc.CommitMutation(persistCtx, mutation, expireTime)
|
||||
cancelPersist()
|
||||
if err != nil {
|
||||
reqLog.Error("gemini.cached_content.commit_mutation_failed", zap.Int64("account_id", account.ID), zap.Error(err))
|
||||
if errors.Is(err, service.ErrGeminiCachedContentMutationConflict) {
|
||||
writeGeminiCachedContentMutationConflict(c)
|
||||
return
|
||||
}
|
||||
googleError(c, http.StatusInternalServerError, "Failed to persist cached content expiration")
|
||||
return
|
||||
}
|
||||
// 提交前库中的到期时间一直是修改前的值:按本次实际推进的时长收费。
|
||||
if extended := expireTime.Sub(mutation.Baseline()); extended > 0 {
|
||||
h.submitGeminiCachedContentUsage(c, apiKey, account, subscription, pricingAt,
|
||||
service.GeminiCachedContentPatchUsageRequestID(record.PublicID), record.Model, record.UpstreamModelID(), 0,
|
||||
service.GeminiCachedContentStorageTokenHours(record.TotalTokenCount, extended), record.ChannelUsage, body)
|
||||
}
|
||||
c.JSON(http.StatusOK, service.GeminiCachedContentView(record))
|
||||
}
|
||||
|
||||
// GeminiCachedContentsDelete DELETE /v1beta/cachedContents/:cacheID
|
||||
// 删除不退还已收取的存储费。
|
||||
func (h *GatewayHandler) GeminiCachedContentsDelete(c *gin.Context) {
|
||||
apiKey, _, reqLog, ok := h.geminiCachedContentsPreamble(c, "handler.gemini_v1beta.cached_contents.delete")
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
svc := h.geminiCachedContentService
|
||||
record, ok := h.resolveGeminiCachedContentOrWrite(c, reqLog, apiKey, c.Param("cacheID"))
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
account, err := svc.BoundAccount(c.Request.Context(), record)
|
||||
if err != nil {
|
||||
reqLog.Error("gemini.cached_content.load_account_failed", zap.Int64("account_id", record.AccountID), zap.Error(err))
|
||||
googleError(c, http.StatusInternalServerError, "Failed to load cached content account")
|
||||
return
|
||||
}
|
||||
if account != nil {
|
||||
setOpsSelectedAccount(c, account.ID, account.Platform)
|
||||
if err := svc.DeleteUpstream(c.Request.Context(), account, record); err != nil {
|
||||
var upErr *service.GeminiCachedContentUpstreamError
|
||||
if !errors.As(err, &upErr) {
|
||||
reqLog.Warn("gemini.cached_content.upstream_delete_error", zap.Int64("account_id", account.ID), zap.Error(err))
|
||||
googleError(c, http.StatusBadGateway, "Upstream request failed")
|
||||
return
|
||||
}
|
||||
if !upErr.NotFound() {
|
||||
h.writeGeminiCachedContentUpstreamError(c, upErr)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
persistCtx, cancelPersist := geminiCachedContentPersistContext(c)
|
||||
err = svc.Forget(persistCtx, record)
|
||||
cancelPersist()
|
||||
if err != nil {
|
||||
reqLog.Error("gemini.cached_content.forget_failed", zap.Error(err))
|
||||
googleError(c, http.StatusInternalServerError, "Failed to delete cached content")
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{})
|
||||
}
|
||||
|
||||
func (h *GatewayHandler) resolveGeminiCachedContentOrWrite(c *gin.Context, reqLog *zap.Logger, apiKey *service.APIKey, cacheID string) (*service.GeminiCachedContent, bool) {
|
||||
record, err := h.geminiCachedContentService.Resolve(c.Request.Context(), apiKey.ID, apiKey.GroupID, cacheID)
|
||||
if err != nil {
|
||||
if errors.Is(err, service.ErrGeminiCachedContentNotFound) {
|
||||
googleError(c, http.StatusForbidden, service.GeminiCachedContentNotFoundMessage)
|
||||
return nil, false
|
||||
}
|
||||
reqLog.Error("gemini.cached_content.lookup_failed", zap.Error(err))
|
||||
googleError(c, http.StatusInternalServerError, "Failed to load cached content")
|
||||
return nil, false
|
||||
}
|
||||
return record, true
|
||||
}
|
||||
|
||||
// resolveGeminiCachedContentBinding 处理 generateContent / countTokens 请求对显式缓存的引用:
|
||||
// 校验归属与模型后把引用改写为上游资源名,返回绑定记录;请求未引用缓存时返回 nil。
|
||||
// 返回 false 表示已写出错误响应。
|
||||
func (h *GatewayHandler) resolveGeminiCachedContentBinding(c *gin.Context, reqLog *zap.Logger, apiKey *service.APIKey, modelName string, body []byte) (*service.GeminiCachedContent, []byte, bool) {
|
||||
ref, err := service.FindGeminiCachedContentReference(body)
|
||||
if err != nil {
|
||||
googleError(c, http.StatusBadRequest, err.Error())
|
||||
return nil, nil, false
|
||||
}
|
||||
if ref == nil {
|
||||
return nil, body, true
|
||||
}
|
||||
if middleware.HasForcePlatform(c) || h.geminiCachedContentService == nil ||
|
||||
apiKey.Group == nil || apiKey.Group.Platform == service.PlatformComposite {
|
||||
googleError(c, http.StatusBadRequest, "cachedContent is not supported on this route")
|
||||
return nil, nil, false
|
||||
}
|
||||
record, err := h.geminiCachedContentService.Resolve(c.Request.Context(), apiKey.ID, apiKey.GroupID, ref.Name)
|
||||
if err != nil {
|
||||
if errors.Is(err, service.ErrGeminiCachedContentNotFound) {
|
||||
googleError(c, http.StatusForbidden, service.GeminiCachedContentNotFoundMessage)
|
||||
return nil, nil, false
|
||||
}
|
||||
reqLog.Error("gemini.cached_content.lookup_failed", zap.Error(err))
|
||||
googleError(c, http.StatusInternalServerError, "Failed to load cached content")
|
||||
return nil, nil, false
|
||||
}
|
||||
if record.Model != modelName {
|
||||
googleError(c, http.StatusBadRequest, fmt.Sprintf(
|
||||
"Model used by GenerateContent request (models/%s) and CachedContent (models/%s) has to be the same.", modelName, record.Model))
|
||||
return nil, nil, false
|
||||
}
|
||||
rewritten, err := service.RewriteGeminiCachedContentReference(body, ref, record.UpstreamName)
|
||||
if err != nil {
|
||||
googleError(c, http.StatusBadRequest, "Invalid cachedContent reference")
|
||||
return nil, nil, false
|
||||
}
|
||||
return record, rewritten, true
|
||||
}
|
||||
|
||||
// acquireGeminiCachedContentAccountSlot 获取选中账号的并发槽(含等待计划);返回 false 表示已写出错误响应。
|
||||
func (h *GatewayHandler) acquireGeminiCachedContentAccountSlot(c *gin.Context, reqLog *zap.Logger, helper *ConcurrencyHelper, selection *service.AccountSelectionResult) (func(), bool) {
|
||||
if selection.Acquired {
|
||||
release := selection.ReleaseFunc
|
||||
if release == nil {
|
||||
release = func() {}
|
||||
}
|
||||
return wrapReleaseOnDone(c.Request.Context(), release), true
|
||||
}
|
||||
account := selection.Account
|
||||
if selection.WaitPlan == nil {
|
||||
markOpsRoutingCapacityLimited(c)
|
||||
googleError(c, http.StatusServiceUnavailable, "No available Gemini accounts")
|
||||
return nil, false
|
||||
}
|
||||
canWait, err := helper.IncrementAccountWaitCount(c.Request.Context(), account.ID, selection.WaitPlan.MaxWaiting)
|
||||
if err != nil {
|
||||
reqLog.Warn("gemini.cached_content.account_wait_counter_increment_failed", zap.Int64("account_id", account.ID), zap.Error(err))
|
||||
} else if !canWait {
|
||||
googleError(c, http.StatusTooManyRequests, "Too many pending requests, please retry later")
|
||||
return nil, false
|
||||
}
|
||||
streamStarted := false
|
||||
release, acquireErr := helper.AcquireAccountSlotWithWaitTimeout(c, account.ID, selection.WaitPlan.MaxConcurrency, selection.WaitPlan.Timeout, false, &streamStarted)
|
||||
if err == nil {
|
||||
helper.DecrementAccountWaitCount(c.Request.Context(), account.ID)
|
||||
}
|
||||
if acquireErr != nil {
|
||||
reqLog.Warn("gemini.cached_content.account_slot_acquire_failed", zap.Int64("account_id", account.ID), zap.Error(acquireErr))
|
||||
googleConcurrencyError(c, acquireErr, "account")
|
||||
return nil, false
|
||||
}
|
||||
if release == nil {
|
||||
release = func() {}
|
||||
}
|
||||
return wrapReleaseOnDone(c.Request.Context(), release), true
|
||||
}
|
||||
|
||||
// submitGeminiCachedContentUsage 记录缓存创建 / 延长的计费:inputTokens 按输入价计,storageTokenHours 按计费模型的存储单价计;
|
||||
// requestID 是带固定前缀的费用行标识(见 GeminiCachedContentUsageRequestIDPrefix),落库时不被请求 id 覆盖。
|
||||
func (h *GatewayHandler) submitGeminiCachedContentUsage(
|
||||
c *gin.Context,
|
||||
apiKey *service.APIKey,
|
||||
account *service.Account,
|
||||
subscription *service.UserSubscription,
|
||||
pricingAt time.Time,
|
||||
requestID string,
|
||||
model string,
|
||||
upstreamModel string,
|
||||
inputTokens int,
|
||||
storageTokenHours float64,
|
||||
usageFields service.ChannelUsageFields,
|
||||
payload []byte,
|
||||
) {
|
||||
if inputTokens <= 0 && storageTokenHours <= 0 {
|
||||
return
|
||||
}
|
||||
if upstreamModel == model {
|
||||
upstreamModel = ""
|
||||
}
|
||||
result := &service.ForwardResult{
|
||||
RequestID: requestID,
|
||||
Model: model,
|
||||
UpstreamModel: upstreamModel,
|
||||
Usage: service.ClaudeUsage{InputTokens: inputTokens},
|
||||
CacheStorageTokenHours: storageTokenHours,
|
||||
}
|
||||
inboundEndpoint := GetInboundEndpoint(c)
|
||||
upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform)
|
||||
userAgent := c.GetHeader("User-Agent")
|
||||
clientIP := ip.GetClientIP(c)
|
||||
requestPayloadHash := service.HashUsageRequestPayload(payload)
|
||||
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
|
||||
h.submitMandatoryUsageRecordTask(c.Request.Context(), func(ctx context.Context) {
|
||||
if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{
|
||||
Result: result,
|
||||
APIKey: apiKey,
|
||||
User: apiKey.User,
|
||||
Account: account,
|
||||
Subscription: subscription,
|
||||
PricingAt: pricingAt,
|
||||
InboundEndpoint: inboundEndpoint,
|
||||
UpstreamEndpoint: upstreamEndpoint,
|
||||
UserAgent: userAgent,
|
||||
IPAddress: clientIP,
|
||||
RequestPayloadHash: requestPayloadHash,
|
||||
APIKeyService: h.apiKeyService,
|
||||
QuotaPlatform: quotaPlatform,
|
||||
ChannelUsageFields: usageFields,
|
||||
}); err != nil {
|
||||
logger.L().With(
|
||||
zap.String("component", "handler.gemini_v1beta.cached_contents"),
|
||||
zap.Int64("api_key_id", apiKey.ID),
|
||||
zap.Int64("account_id", account.ID),
|
||||
zap.String("request_id", result.RequestID),
|
||||
).Error("gemini.cached_content.record_usage_failed", zap.Error(err))
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// geminiCachedContentPersistContext 返回与客户端连接脱钩的落库 context:上游操作已生效时,客户端断开不应让本地记录缺失。
|
||||
func geminiCachedContentPersistContext(c *gin.Context) (context.Context, context.CancelFunc) {
|
||||
return context.WithTimeout(context.WithoutCancel(c.Request.Context()), geminiCachedContentPersistTimeout)
|
||||
}
|
||||
|
||||
// cleanupGeminiCachedContentUpstream 尽力删除已在上游创建、但网关不会登记的缓存,避免其持续产生存储费。
|
||||
func (h *GatewayHandler) cleanupGeminiCachedContentUpstream(reqLog *zap.Logger, account *service.Account, upstreamName string) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), geminiCachedContentPersistTimeout)
|
||||
defer cancel()
|
||||
if err := h.geminiCachedContentService.DeleteUpstream(ctx, account, &service.GeminiCachedContent{UpstreamName: upstreamName}); err != nil {
|
||||
reqLog.Error("gemini.cached_content.orphan_cleanup_failed", zap.Int64("account_id", account.ID), zap.String("upstream_name", upstreamName), zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
// reserveGeminiCachedContentCharge 校验余额 / 订阅额度 / API Key 额度足以支付确定金额 amount(倍率后),
|
||||
// 并按该金额在用户、订阅、API Key 维度登记在途预留(不受预留上限封顶,见 ReserveFixedCharge)。
|
||||
// 返回的 done 须在 handler 结束时调用;预留随计费任务交接,落库后释放。
|
||||
func (h *GatewayHandler) reserveGeminiCachedContentCharge(c *gin.Context, apiKey *service.APIKey, subscription *service.UserSubscription, amount float64) (func(), error) {
|
||||
ctx := c.Request.Context()
|
||||
res, err := h.billingCacheService.ReserveFixedCharge(ctx, apiKey.User, apiKey, apiKey.Group, subscription, amount)
|
||||
if err != nil {
|
||||
return inflightNoop, err
|
||||
}
|
||||
if res == nil {
|
||||
return inflightNoop, nil
|
||||
}
|
||||
c.Request = c.Request.WithContext(service.WithInflightReservation(ctx, res))
|
||||
return res.HandlerDone, nil
|
||||
}
|
||||
|
||||
// writeGeminiCachedContentMutationConflict 回写同一缓存的修改冲突(Google 并发冲突语义 ABORTED),客户端可稍后重试。
|
||||
func writeGeminiCachedContentMutationConflict(c *gin.Context) {
|
||||
c.JSON(http.StatusConflict, gin.H{
|
||||
"error": gin.H{
|
||||
"code": http.StatusConflict,
|
||||
"message": "Another update to this cached content is in progress, please retry later",
|
||||
"status": "ABORTED",
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
func writeGeminiCachedContentBillingError(c *gin.Context, err error) {
|
||||
status, code, message, retryAfter := billingErrorDetails(err)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
googleErrorWithType(c, status, code, "", message)
|
||||
}
|
||||
|
||||
// writeGeminiCachedContentUpstreamError 回写上游错误:缓存不存在统一返回标准 403(各上游形态不同且可能含上游资源 ID);
|
||||
// 账号凭据或权限类错误不是客户端问题,统一改为 502;限流 / 不可用按普通 Gemini 路径处理(透传规则、通用映射),
|
||||
// 只有上游网关账号的这类响应反映持有缓存的上游账号状态,原样回写。
|
||||
func (h *GatewayHandler) writeGeminiCachedContentUpstreamError(c *gin.Context, upErr *service.GeminiCachedContentUpstreamError) {
|
||||
status := upErr.StatusCode
|
||||
message := strings.TrimSpace(extractGoogleErrorMessage(upErr.Body))
|
||||
service.SetOpsUpstreamError(c, status, message, "")
|
||||
if upErr.NotFound() {
|
||||
service.MarkOpsRequestScopedError(c, "permission_error")
|
||||
c.Data(http.StatusForbidden, "application/json", []byte(service.GeminiCachedContentNotFoundResponse))
|
||||
return
|
||||
}
|
||||
if (status == http.StatusTooManyRequests || status >= http.StatusInternalServerError) && !upErr.UpstreamGateway {
|
||||
h.handleGeminiFailoverExhausted(c, &service.UpstreamFailoverError{StatusCode: status, ResponseBody: upErr.Body})
|
||||
return
|
||||
}
|
||||
if status == http.StatusUnauthorized || status == http.StatusForbidden {
|
||||
googleError(c, http.StatusBadGateway, "Upstream account rejected the cached content request")
|
||||
return
|
||||
}
|
||||
if len(upErr.Body) == 0 || !json.Valid(upErr.Body) {
|
||||
if message == "" {
|
||||
message = http.StatusText(status)
|
||||
}
|
||||
googleError(c, status, message)
|
||||
return
|
||||
}
|
||||
c.Data(status, "application/json", upErr.Body)
|
||||
}
|
||||
|
||||
func extractGoogleErrorMessage(body []byte) string {
|
||||
var envelope struct {
|
||||
Error struct {
|
||||
Message string `json:"message"`
|
||||
} `json:"error"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &envelope); err != nil {
|
||||
return ""
|
||||
}
|
||||
return envelope.Error.Message
|
||||
}
|
||||
|
||||
func encodeGeminiCachedContentPageToken(id int64) string {
|
||||
return base64.RawURLEncoding.EncodeToString([]byte(strconv.FormatInt(id, 10)))
|
||||
}
|
||||
|
||||
func decodeGeminiCachedContentPageToken(raw string) (int64, bool) {
|
||||
decoded, err := base64.RawURLEncoding.DecodeString(raw)
|
||||
if err != nil {
|
||||
return 0, false
|
||||
}
|
||||
id, err := strconv.ParseInt(string(decoded), 10, 64)
|
||||
if err != nil || id <= 0 {
|
||||
return 0, false
|
||||
}
|
||||
return id, true
|
||||
}
|
||||
|
||||
// mergeAccountIDSets 返回两个账号 ID 集合的并集(不修改入参)。
|
||||
func mergeAccountIDSets(a, b map[int64]struct{}) map[int64]struct{} {
|
||||
out := make(map[int64]struct{}, len(a)+len(b))
|
||||
for id := range a {
|
||||
out[id] = struct{}{}
|
||||
}
|
||||
for id := range b {
|
||||
out[id] = struct{}{}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// handleGeminiCachedContentFailoverExhausted 处理绑定缓存的请求在持有账号上重试耗尽:没有其他账号可换。
|
||||
// 上游网关账号的限流 / 不可用响应反映持有缓存的上游账号状态,原样返回供客户端判断;
|
||||
// 其余(含官方账号)仍按普通 Gemini 路径处理:透传规则、通用映射。
|
||||
func (h *GatewayHandler) handleGeminiCachedContentFailoverExhausted(c *gin.Context, failoverErr *service.UpstreamFailoverError) {
|
||||
if failoverErr == nil || !failoverErr.BoundUpstreamPassthrough || len(failoverErr.ResponseBody) == 0 || !json.Valid(failoverErr.ResponseBody) {
|
||||
h.handleGeminiFailoverExhausted(c, failoverErr)
|
||||
return
|
||||
}
|
||||
service.SetOpsUpstreamError(c, failoverErr.StatusCode, service.ExtractUpstreamErrorMessage(failoverErr.ResponseBody), "")
|
||||
c.Data(failoverErr.StatusCode, "application/json", failoverErr.ResponseBody)
|
||||
}
|
||||
@@ -0,0 +1,264 @@
|
||||
//go:build unit
|
||||
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
middleware "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
type geminiCachedContentHandlerRepoStub struct {
|
||||
records []*service.GeminiCachedContent
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentHandlerRepoStub) Create(_ context.Context, record *service.GeminiCachedContent) error {
|
||||
record.ID = int64(len(r.records) + 1)
|
||||
r.records = append(r.records, record)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentHandlerRepoStub) GetForOwner(_ context.Context, apiKeyID int64, publicID string) (*service.GeminiCachedContent, error) {
|
||||
for _, record := range r.records {
|
||||
if record.PublicID == publicID && record.APIKeyID == apiKeyID {
|
||||
return record, nil
|
||||
}
|
||||
}
|
||||
return nil, service.ErrGeminiCachedContentNotFound
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentHandlerRepoStub) ListForOwner(_ context.Context, apiKeyID, groupID int64, activeAt time.Time, beforeID int64, limit int) ([]*service.GeminiCachedContent, error) {
|
||||
var out []*service.GeminiCachedContent
|
||||
for i := len(r.records) - 1; i >= 0; i-- {
|
||||
record := r.records[i]
|
||||
if record.APIKeyID == apiKeyID && record.GroupID == groupID && record.ExpireTime.After(activeAt) && (beforeID <= 0 || record.ID < beforeID) {
|
||||
out = append(out, record)
|
||||
}
|
||||
if len(out) == limit {
|
||||
break
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentHandlerRepoStub) BeginMutation(context.Context, int64, time.Time, string, time.Duration) (bool, error) {
|
||||
return true, nil
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentHandlerRepoStub) TakeOverStaleMutation(context.Context, int64, string, time.Duration) (*service.GeminiCachedContent, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentHandlerRepoStub) CommitMutation(context.Context, int64, string, time.Time) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentHandlerRepoStub) ReleaseMutation(context.Context, int64, string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentHandlerRepoStub) ListStaleMutations(context.Context, int) ([]int64, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentHandlerRepoStub) SoftDelete(context.Context, int64) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func newGeminiCachedContentHandlerForTest(repo *geminiCachedContentHandlerRepoStub) *GatewayHandler {
|
||||
return &GatewayHandler{geminiCachedContentService: service.NewGeminiCachedContentService(repo, nil, &config.Config{})}
|
||||
}
|
||||
|
||||
func geminiCachedContentTestAPIKey(platform string) *service.APIKey {
|
||||
groupID := int64(42)
|
||||
return &service.APIKey{ID: 7, UserID: 3, GroupID: &groupID, Group: &service.Group{ID: groupID, Platform: platform}}
|
||||
}
|
||||
|
||||
func newGeminiCachedContentTestContext(method, target string, apiKey *service.APIKey) (*gin.Context, *httptest.ResponseRecorder) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(method, target, nil)
|
||||
c.Set(string(middleware.ContextKeyAPIKey), apiKey)
|
||||
c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: apiKey.UserID, Concurrency: 10})
|
||||
return c, rec
|
||||
}
|
||||
|
||||
func seedGeminiCachedContent(repo *geminiCachedContentHandlerRepoStub, publicID string, apiKeyID int64, expire time.Time) *service.GeminiCachedContent {
|
||||
record := &service.GeminiCachedContent{
|
||||
PublicID: publicID,
|
||||
UserID: 3,
|
||||
APIKeyID: apiKeyID,
|
||||
GroupID: 42,
|
||||
AccountID: 900,
|
||||
UpstreamName: "cachedContents/upstream" + publicID,
|
||||
Model: "gemini-3.8-flash",
|
||||
UpstreamModel: "models/gemini-3.8-flash",
|
||||
TotalTokenCount: 6273,
|
||||
ExpireTime: expire,
|
||||
CreatedAt: time.Date(2026, 10, 9, 8, 0, 0, 0, time.UTC),
|
||||
UpdatedAt: time.Date(2026, 10, 9, 8, 0, 0, 0, time.UTC),
|
||||
}
|
||||
_ = repo.Create(context.Background(), record)
|
||||
return record
|
||||
}
|
||||
|
||||
func TestResolveGeminiCachedContentBinding(t *testing.T) {
|
||||
repo := &geminiCachedContentHandlerRepoStub{}
|
||||
seedGeminiCachedContent(repo, "abc", 7, time.Now().Add(time.Hour))
|
||||
seedGeminiCachedContent(repo, "foreign", 8, time.Now().Add(time.Hour))
|
||||
h := newGeminiCachedContentHandlerForTest(repo)
|
||||
apiKey := geminiCachedContentTestAPIKey(service.PlatformGemini)
|
||||
|
||||
t.Run("request without cache passes through", func(t *testing.T) {
|
||||
c, _ := newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", apiKey)
|
||||
body := []byte(`{"contents":[]}`)
|
||||
record, out, ok := h.resolveGeminiCachedContentBinding(c, zap.NewNop(), apiKey, "gemini-3.8-flash", body)
|
||||
require.True(t, ok)
|
||||
require.Nil(t, record)
|
||||
require.Equal(t, body, out)
|
||||
})
|
||||
|
||||
t.Run("owned cache rewrites reference and binds account", func(t *testing.T) {
|
||||
c, _ := newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", apiKey)
|
||||
record, out, ok := h.resolveGeminiCachedContentBinding(c, zap.NewNop(), apiKey, "gemini-3.8-flash",
|
||||
[]byte(`{"cachedContent":"cachedContents/abc","contents":[]}`))
|
||||
require.True(t, ok)
|
||||
require.Equal(t, int64(900), record.AccountID)
|
||||
require.JSONEq(t, `{"cachedContent":"cachedContents/upstreamabc","contents":[]}`, string(out))
|
||||
})
|
||||
|
||||
t.Run("cache owned by another key looks missing", func(t *testing.T) {
|
||||
c, rec := newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", apiKey)
|
||||
_, _, ok := h.resolveGeminiCachedContentBinding(c, zap.NewNop(), apiKey, "gemini-3.8-flash",
|
||||
[]byte(`{"cachedContent":"cachedContents/foreign","contents":[]}`))
|
||||
require.False(t, ok)
|
||||
require.Equal(t, http.StatusForbidden, rec.Code)
|
||||
require.Contains(t, rec.Body.String(), service.GeminiCachedContentNotFoundMessage)
|
||||
})
|
||||
|
||||
t.Run("model must match the cache literally", func(t *testing.T) {
|
||||
c, rec := newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-flash-latest:generateContent", apiKey)
|
||||
_, _, ok := h.resolveGeminiCachedContentBinding(c, zap.NewNop(), apiKey, "gemini-flash-latest",
|
||||
[]byte(`{"cachedContent":"cachedContents/abc","contents":[]}`))
|
||||
require.False(t, ok)
|
||||
require.Equal(t, http.StatusBadRequest, rec.Code)
|
||||
require.Contains(t, rec.Body.String(), "has to be the same")
|
||||
})
|
||||
|
||||
t.Run("composite groups are rejected", func(t *testing.T) {
|
||||
composite := geminiCachedContentTestAPIKey(service.PlatformComposite)
|
||||
c, rec := newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", composite)
|
||||
_, _, ok := h.resolveGeminiCachedContentBinding(c, zap.NewNop(), composite, "gemini-3.8-flash",
|
||||
[]byte(`{"cachedContent":"cachedContents/abc","contents":[]}`))
|
||||
require.False(t, ok)
|
||||
require.Equal(t, http.StatusBadRequest, rec.Code)
|
||||
})
|
||||
}
|
||||
|
||||
func TestGeminiCachedContentsListAndGet(t *testing.T) {
|
||||
repo := &geminiCachedContentHandlerRepoStub{}
|
||||
seedGeminiCachedContent(repo, "first", 7, time.Now().Add(time.Hour))
|
||||
seedGeminiCachedContent(repo, "second", 7, time.Now().Add(time.Hour))
|
||||
seedGeminiCachedContent(repo, "expired", 7, time.Now().Add(-time.Minute))
|
||||
seedGeminiCachedContent(repo, "foreign", 8, time.Now().Add(time.Hour))
|
||||
h := newGeminiCachedContentHandlerForTest(repo)
|
||||
apiKey := geminiCachedContentTestAPIKey(service.PlatformGemini)
|
||||
|
||||
c, rec := newGeminiCachedContentTestContext(http.MethodGet, "/v1beta/cachedContents?pageSize=1", apiKey)
|
||||
h.GeminiCachedContentsList(c)
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
var page struct {
|
||||
CachedContents []map[string]any `json:"cachedContents"`
|
||||
NextPageToken string `json:"nextPageToken"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &page))
|
||||
require.Len(t, page.CachedContents, 1)
|
||||
require.Equal(t, "cachedContents/second", page.CachedContents[0]["name"])
|
||||
require.NotEmpty(t, page.NextPageToken)
|
||||
|
||||
c, rec = newGeminiCachedContentTestContext(http.MethodGet, "/v1beta/cachedContents?pageSize=1&pageToken="+page.NextPageToken, apiKey)
|
||||
h.GeminiCachedContentsList(c)
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
page.NextPageToken = ""
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &page))
|
||||
require.Len(t, page.CachedContents, 1)
|
||||
require.Equal(t, "cachedContents/first", page.CachedContents[0]["name"])
|
||||
require.Empty(t, page.NextPageToken, "过期与其他 Key 的缓存不出现在列表里")
|
||||
|
||||
c, rec = newGeminiCachedContentTestContext(http.MethodGet, "/v1beta/cachedContents/first", apiKey)
|
||||
c.Params = gin.Params{{Key: "cacheID", Value: "first"}}
|
||||
h.GeminiCachedContentsGet(c)
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
var view map[string]any
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &view))
|
||||
require.Equal(t, "cachedContents/first", view["name"])
|
||||
require.Equal(t, "models/gemini-3.8-flash", view["model"])
|
||||
require.NotContains(t, rec.Body.String(), "upstreamfirst", "上游资源名不暴露给客户端")
|
||||
|
||||
c, rec = newGeminiCachedContentTestContext(http.MethodGet, "/v1beta/cachedContents/foreign", apiKey)
|
||||
c.Params = gin.Params{{Key: "cacheID", Value: "foreign"}}
|
||||
h.GeminiCachedContentsGet(c)
|
||||
require.Equal(t, http.StatusForbidden, rec.Code)
|
||||
require.True(t, strings.Contains(rec.Body.String(), service.GeminiCachedContentNotFoundMessage))
|
||||
}
|
||||
|
||||
func TestHandleGeminiCachedContentFailoverExhausted(t *testing.T) {
|
||||
h := &GatewayHandler{}
|
||||
apiKey := geminiCachedContentTestAPIKey(service.PlatformGemini)
|
||||
unavailable := `{"error":{"code":503,"message":"The account holding this cached content is temporarily unavailable","status":"UNAVAILABLE"}}`
|
||||
|
||||
c, rec := newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", apiKey)
|
||||
h.handleGeminiCachedContentFailoverExhausted(c, &service.UpstreamFailoverError{StatusCode: http.StatusServiceUnavailable, ResponseBody: []byte(unavailable), BoundUpstreamPassthrough: true})
|
||||
require.Equal(t, http.StatusServiceUnavailable, rec.Code, "绑定缓存的请求无账号可换,上游网关的 503 原样返回")
|
||||
require.JSONEq(t, unavailable, rec.Body.String())
|
||||
|
||||
c, rec = newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", apiKey)
|
||||
h.handleGeminiCachedContentFailoverExhausted(c, &service.UpstreamFailoverError{StatusCode: http.StatusTooManyRequests, ResponseBody: []byte(`{"error":{"code":429,"message":"Too many pending requests, please retry later","status":"RESOURCE_EXHAUSTED"}}`), BoundUpstreamPassthrough: true})
|
||||
require.Equal(t, http.StatusTooManyRequests, rec.Code)
|
||||
require.Contains(t, rec.Body.String(), "Too many pending requests")
|
||||
|
||||
officialQuota := `{"error":{"code":429,"message":"You exceeded your current quota","status":"RESOURCE_EXHAUSTED","details":[{"@type":"type.googleapis.com/google.rpc.QuotaFailure","violations":[{"quotaId":"GenerateRequestsPerDayPerProjectPerModel-FreeTier"}]}]}}`
|
||||
c, rec = newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", apiKey)
|
||||
h.handleGeminiCachedContentFailoverExhausted(c, &service.UpstreamFailoverError{StatusCode: http.StatusTooManyRequests, ResponseBody: []byte(officialQuota)})
|
||||
require.Equal(t, http.StatusTooManyRequests, rec.Code, "官方账号的限流按普通 Gemini 路径映射")
|
||||
require.NotContains(t, rec.Body.String(), "FreeTier", "不向客户端暴露上游账号的配额明细")
|
||||
require.Contains(t, rec.Body.String(), "Upstream rate limit exceeded")
|
||||
|
||||
c, rec = newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", apiKey)
|
||||
h.handleGeminiCachedContentFailoverExhausted(c, &service.UpstreamFailoverError{StatusCode: http.StatusUnauthorized, ResponseBody: []byte(`{"error":{"code":401,"message":"bad key"}}`)})
|
||||
require.Equal(t, http.StatusBadGateway, rec.Code, "账号凭据类错误仍按网关映射,不透传给客户端")
|
||||
require.NotContains(t, rec.Body.String(), "bad key")
|
||||
}
|
||||
|
||||
func TestWriteGeminiCachedContentUpstreamError_RateLimitMappedUnlessUpstreamGateway(t *testing.T) {
|
||||
h := &GatewayHandler{}
|
||||
apiKey := geminiCachedContentTestAPIKey(service.PlatformGemini)
|
||||
quota := `{"error":{"code":429,"message":"You exceeded your current quota","status":"RESOURCE_EXHAUSTED","details":[{"quotaId":"GenerateRequestsPerDayPerProjectPerModel-FreeTier"}]}}`
|
||||
|
||||
c, rec := newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/cachedContents", apiKey)
|
||||
h.writeGeminiCachedContentUpstreamError(c, &service.GeminiCachedContentUpstreamError{StatusCode: http.StatusTooManyRequests, Body: []byte(quota)})
|
||||
require.Equal(t, http.StatusTooManyRequests, rec.Code)
|
||||
require.NotContains(t, rec.Body.String(), "FreeTier", "官方账号的限流不透传配额明细")
|
||||
|
||||
c, rec = newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/cachedContents", apiKey)
|
||||
h.writeGeminiCachedContentUpstreamError(c, &service.GeminiCachedContentUpstreamError{StatusCode: http.StatusTooManyRequests, Body: []byte(quota), UpstreamGateway: true})
|
||||
require.Equal(t, http.StatusTooManyRequests, rec.Code)
|
||||
require.JSONEq(t, quota, rec.Body.String(), "上游网关账号的限流反映持有缓存的上游账号状态,原样回写")
|
||||
|
||||
c, rec = newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/cachedContents", apiKey)
|
||||
h.writeGeminiCachedContentUpstreamError(c, &service.GeminiCachedContentUpstreamError{StatusCode: http.StatusBadRequest, Body: []byte(`{"error":{"code":400,"message":"Cached content is too small","status":"INVALID_ARGUMENT"}}`)})
|
||||
require.Equal(t, http.StatusBadRequest, rec.Code)
|
||||
require.Contains(t, rec.Body.String(), "too small", "请求校验类错误仍原样回写")
|
||||
}
|
||||
@@ -74,7 +74,7 @@ func (h *GatewayHandler) GeminiV1BetaListModels(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
account, err := h.geminiCompatService.SelectAccountForAIStudioEndpoints(c.Request.Context(), apiKey.GroupID)
|
||||
account, err := h.geminiCompatService.SelectAccountForModelsEndpoints(c.Request.Context(), apiKey.GroupID)
|
||||
if err != nil {
|
||||
if len(agModels) > 0 {
|
||||
c.JSON(http.StatusOK, gemini.ModelsListResponse{Models: filterGeminiModels(agModels)})
|
||||
@@ -85,7 +85,7 @@ func (h *GatewayHandler) GeminiV1BetaListModels(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
res, err := h.geminiCompatService.ForwardAIStudioGET(c.Request.Context(), account, "/v1beta/models")
|
||||
res, err := h.geminiCompatService.ForwardListModels(c.Request.Context(), account)
|
||||
if err != nil {
|
||||
googleError(c, http.StatusBadGateway, err.Error())
|
||||
return
|
||||
@@ -251,7 +251,7 @@ func (h *GatewayHandler) GeminiV1BetaGetModel(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
account, err := h.geminiCompatService.SelectAccountForAIStudioEndpoints(c.Request.Context(), apiKey.GroupID)
|
||||
account, err := h.geminiCompatService.SelectAccountForModelsEndpoints(c.Request.Context(), apiKey.GroupID)
|
||||
if err != nil {
|
||||
// 没有 gemini 账户,检查是否有 antigravity 账户可用
|
||||
hasAntigravity, _ := h.geminiCompatService.HasAntigravityAccounts(c.Request.Context(), apiKey.GroupID)
|
||||
@@ -265,7 +265,7 @@ func (h *GatewayHandler) GeminiV1BetaGetModel(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
res, err := h.geminiCompatService.ForwardAIStudioGET(c.Request.Context(), account, "/v1beta/models/"+modelName)
|
||||
res, err := h.geminiCompatService.ForwardGetModel(c.Request.Context(), account, modelName)
|
||||
if err != nil {
|
||||
googleError(c, http.StatusBadGateway, err.Error())
|
||||
return
|
||||
@@ -356,6 +356,12 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
|
||||
modelName = channelMapping.MappedModel
|
||||
}
|
||||
|
||||
// 引用显式缓存的请求只能发往持有该缓存的账号,且不参与会话粘性与换号。
|
||||
cachedContentBinding, body, ok := h.resolveGeminiCachedContentBinding(c, reqLog, apiKey, modelName, body)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
// Get subscription (may be nil)
|
||||
subscription, _ := middleware.GetSubscriptionFromContext(c)
|
||||
|
||||
@@ -382,11 +388,11 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
|
||||
// 2) billing eligibility check (after wait)
|
||||
if err := h.billingCacheService.CheckBillingEligibility(c.Request.Context(), apiKey.User, apiKey, apiKey.Group, subscription, service.QuotaPlatform(c.Request.Context(), apiKey)); err != nil {
|
||||
reqLog.Info("gemini.billing_eligibility_check_failed", zap.Error(err))
|
||||
status, _, message, retryAfter := billingErrorDetails(err)
|
||||
status, code, message, retryAfter := billingErrorDetails(err)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
googleError(c, status, message)
|
||||
googleErrorWithType(c, status, code, "", message)
|
||||
return
|
||||
}
|
||||
|
||||
@@ -394,19 +400,22 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
|
||||
inflightRelease, err := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(modelName, body))
|
||||
if err != nil {
|
||||
reqLog.Info("gemini.inflight_reservation_rejected", zap.Error(err))
|
||||
status, _, message, retryAfter := billingErrorDetails(err)
|
||||
status, code, message, retryAfter := billingErrorDetails(err)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
googleError(c, status, message)
|
||||
googleErrorWithType(c, status, code, "", message)
|
||||
return
|
||||
}
|
||||
defer inflightRelease()
|
||||
|
||||
// 3) select account (sticky session based on request body)
|
||||
// 优先使用 Gemini CLI 的会话标识(privileged-user-id + tmp 目录哈希)
|
||||
sessionHash := extractGeminiCLISessionHash(c, body)
|
||||
if sessionHash == "" {
|
||||
sessionHash := ""
|
||||
if cachedContentBinding == nil {
|
||||
sessionHash = extractGeminiCLISessionHash(c, body)
|
||||
}
|
||||
if sessionHash == "" && cachedContentBinding == nil {
|
||||
// Fallback: 使用通用的会话哈希生成逻辑(适用于其他客户端)
|
||||
parsedReq, _ := service.ParseGatewayRequest(service.NewRequestBodyRef(body), domain.PlatformGemini)
|
||||
if parsedReq != nil {
|
||||
@@ -443,7 +452,7 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
|
||||
var geminiPrefixHash string
|
||||
var geminiSessionUUID string
|
||||
var matchedDigestChain string
|
||||
useDigestFallback := sessionBoundAccountID == 0
|
||||
useDigestFallback := sessionBoundAccountID == 0 && cachedContentBinding == nil
|
||||
|
||||
if useDigestFallback {
|
||||
// 解析 Gemini 请求体
|
||||
@@ -507,7 +516,21 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
|
||||
hasBoundSession := sessionKey != "" && sessionBoundAccountID > 0
|
||||
cleanedForUnknownBinding := false
|
||||
|
||||
fs := NewFailoverState(h.maxAccountSwitchesGemini, hasBoundSession)
|
||||
maxAccountSwitches := h.maxAccountSwitchesGemini
|
||||
var cachedContentExclusions map[int64]struct{}
|
||||
if cachedContentBinding != nil {
|
||||
// 缓存只存在于绑定账号:排除分组内其余账号,失败不换号。
|
||||
maxAccountSwitches = 0
|
||||
sessionBoundAccountID = cachedContentBinding.AccountID
|
||||
c.Request = c.Request.WithContext(service.WithGeminiCachedContentBound(c.Request.Context()))
|
||||
cachedContentExclusions, err = h.gatewayService.GeminiExplicitCacheExclusions(c.Request.Context(), apiKey.GroupID, cachedContentBinding.AccountID)
|
||||
if err != nil {
|
||||
reqLog.Warn("gemini.cached_content.list_accounts_failed", zap.Error(err))
|
||||
googleError(c, http.StatusServiceUnavailable, geminiCachedContentUnavailableMessage)
|
||||
return
|
||||
}
|
||||
}
|
||||
fs := NewFailoverState(maxAccountSwitches, hasBoundSession)
|
||||
|
||||
// 单账号分组提前设置 SingleAccountRetry 标记,让 Service 层首次 503 就不设模型限流标记。
|
||||
// 避免单账号分组收到 503 (MODEL_CAPACITY_EXHAUSTED) 时设 29s 限流,导致后续请求连续快速失败。
|
||||
@@ -517,12 +540,32 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
|
||||
}
|
||||
|
||||
for {
|
||||
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, sessionKey, modelName, fs.FailedAccountIDs, "", int64(0)) // Gemini 不使用会话限制
|
||||
excludedAccountIDs := fs.FailedAccountIDs
|
||||
if cachedContentExclusions != nil {
|
||||
excludedAccountIDs = mergeAccountIDSets(cachedContentExclusions, fs.FailedAccountIDs)
|
||||
}
|
||||
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, sessionKey, modelName, excludedAccountIDs, "", int64(0)) // Gemini 不使用会话限制
|
||||
if err == nil && cachedContentBinding != nil && selection.Account.ID != cachedContentBinding.AccountID {
|
||||
if selection.Acquired && selection.ReleaseFunc != nil {
|
||||
selection.ReleaseFunc()
|
||||
}
|
||||
err = service.ErrNoAvailableAccounts
|
||||
}
|
||||
if err != nil {
|
||||
if failoverClientGone(c) {
|
||||
reqLog.Info("gemini.account_select_aborted_client_disconnected", zap.Error(err))
|
||||
return
|
||||
}
|
||||
if cachedContentBinding != nil {
|
||||
if fs.LastFailoverErr != nil {
|
||||
h.handleGeminiCachedContentFailoverExhausted(c, fs.LastFailoverErr)
|
||||
return
|
||||
}
|
||||
reqLog.Info("gemini.cached_content.bound_account_unavailable", zap.Int64("account_id", cachedContentBinding.AccountID), zap.Error(err))
|
||||
markOpsRoutingCapacityLimited(c)
|
||||
googleError(c, http.StatusServiceUnavailable, geminiCachedContentUnavailableMessage)
|
||||
return
|
||||
}
|
||||
if len(fs.FailedAccountIDs) == 0 {
|
||||
cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, modelName, modelName, service.PlatformGemini)
|
||||
if !cls.ModelNotFound {
|
||||
@@ -631,6 +674,11 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
|
||||
accountReleaseFunc()
|
||||
}
|
||||
reqLog.Debug("gemini.account_slot_profit_vetoed", zap.Int64("account_id", account.ID), zap.String("reason", reason))
|
||||
if cachedContentBinding != nil {
|
||||
markOpsRoutingCapacityLimited(c)
|
||||
googleError(c, http.StatusServiceUnavailable, geminiCachedContentUnavailableMessage)
|
||||
return
|
||||
}
|
||||
if fs.RecordProfitVeto(account.ID) == FailoverExhausted {
|
||||
reqLog.Warn("gemini.profit_veto_attempts_exhausted", zap.Int("profit_veto_count", fs.ProfitVetoCount()))
|
||||
markOpsRoutingCapacityLimited(c)
|
||||
@@ -684,6 +732,10 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
|
||||
case FailoverContinue:
|
||||
continue
|
||||
case FailoverExhausted:
|
||||
if cachedContentBinding != nil {
|
||||
h.handleGeminiCachedContentFailoverExhausted(c, fs.LastFailoverErr)
|
||||
return
|
||||
}
|
||||
h.handleGeminiFailoverExhausted(c, fs.LastFailoverErr)
|
||||
return
|
||||
case FailoverCanceled:
|
||||
@@ -809,7 +861,8 @@ func (h *GatewayHandler) handleGeminiFailoverExhausted(c *gin.Context, failoverE
|
||||
c.Set(service.OpsSkipPassthroughKey, true)
|
||||
}
|
||||
|
||||
googleError(c, respCode, msg)
|
||||
// 按透传规则改写的是上游错误,不登记本地错误类型。
|
||||
googleErrorWithType(c, respCode, "", "", msg)
|
||||
return
|
||||
}
|
||||
}
|
||||
@@ -844,7 +897,21 @@ type pathParseError struct{ msg string }
|
||||
|
||||
func (e *pathParseError) Error() string { return e.msg }
|
||||
|
||||
// googleError 以 Google 错误格式回写网关产生的错误,并按状态码向 ops 登记错误类型:
|
||||
// 4xx 为客户端请求问题,其余状态码不登记(按平台错误统计)。
|
||||
func googleError(c *gin.Context, status int, message string) {
|
||||
googleErrorWithType(c, status, googleClientErrorType(status), "", message)
|
||||
}
|
||||
|
||||
// googlePlatformError 以 Google 错误格式回写由平台侧原因(如管理员未完成配置)导致的错误,ops 按平台错误统计。
|
||||
func googlePlatformError(c *gin.Context, status int, message string) {
|
||||
googleErrorWithType(c, status, "", "", message)
|
||||
}
|
||||
|
||||
// googleErrorWithType 以 Google 错误格式回写错误,并向 ops 登记错误类型与错误码(Google 错误体没有 type 字段),
|
||||
// 与 Claude / OpenAI 格式错误同口径分类;errType 为空时不登记。
|
||||
func googleErrorWithType(c *gin.Context, status int, errType, code, message string) {
|
||||
service.SetOpsLocalErrorType(c, errType, code)
|
||||
c.JSON(status, gin.H{
|
||||
"error": gin.H{
|
||||
"code": status,
|
||||
@@ -854,11 +921,27 @@ func googleError(c *gin.Context, status int, message string) {
|
||||
})
|
||||
}
|
||||
|
||||
// googleClientErrorType 返回客户端请求类 4xx 状态码对应的 ops 错误类型;其余状态码返回空串。
|
||||
func googleClientErrorType(status int) string {
|
||||
switch status {
|
||||
case http.StatusBadRequest, http.StatusRequestEntityTooLarge:
|
||||
return "invalid_request_error"
|
||||
case http.StatusUnauthorized:
|
||||
return "authentication_error"
|
||||
case http.StatusForbidden:
|
||||
return "permission_error"
|
||||
case http.StatusNotFound:
|
||||
return "not_found_error"
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
// googleConcurrencyError 以 Google 错误格式回写并发槽获取失败,状态码与文案
|
||||
// 沿用 concurrencyErrorResponse 的统一映射(客户端断开为 499)。
|
||||
func googleConcurrencyError(c *gin.Context, err error, slotType string) {
|
||||
status, _, _, message := concurrencyErrorResponse(err, slotType)
|
||||
googleError(c, status, message)
|
||||
status, errType, code, message := concurrencyErrorResponse(err, slotType)
|
||||
googleErrorWithType(c, status, errType, code, message)
|
||||
}
|
||||
|
||||
func writeUpstreamResponse(c *gin.Context, res *service.UpstreamHTTPResult) {
|
||||
|
||||
@@ -109,7 +109,7 @@ func TestGeminiV1BetaHandler_ListModelsAntigravityFallback(t *testing.T) {
|
||||
expectedBehavior string
|
||||
}{
|
||||
{
|
||||
name: "有Gemini账户-调用ForwardAIStudioGET",
|
||||
name: "有Gemini账户-调用ForwardListModels",
|
||||
hasGeminiAccount: true,
|
||||
hasAntigravity: false,
|
||||
expectedBehavior: "forward_to_upstream",
|
||||
@@ -155,7 +155,7 @@ func TestGeminiV1BetaHandler_GetModelAntigravityFallback(t *testing.T) {
|
||||
expectedBehavior string
|
||||
}{
|
||||
{
|
||||
name: "有Gemini账户-调用ForwardAIStudioGET",
|
||||
name: "有Gemini账户-调用ForwardGetModel",
|
||||
hasGeminiAccount: true,
|
||||
hasAntigravity: false,
|
||||
expectedBehavior: "forward_to_upstream",
|
||||
|
||||
@@ -0,0 +1,83 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type geminiVertexModelsTokenCache struct{}
|
||||
|
||||
func (geminiVertexModelsTokenCache) GetAccessToken(context.Context, string) (string, error) {
|
||||
return "ya29.vertex-token", nil
|
||||
}
|
||||
|
||||
func (geminiVertexModelsTokenCache) SetAccessToken(context.Context, string, string, time.Duration) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (geminiVertexModelsTokenCache) DeleteAccessToken(context.Context, string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (geminiVertexModelsTokenCache) AcquireRefreshLock(context.Context, string, time.Duration) (bool, error) {
|
||||
return false, nil
|
||||
}
|
||||
|
||||
func (geminiVertexModelsTokenCache) ReleaseRefreshLock(context.Context, string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// 分组里只有 Agent Platform (Vertex) 服务账号时,模型列表与模型详情接口返回 Vertex 上的 Gemini 模型。
|
||||
func TestGeminiV1BetaModelsEndpoints_VertexServiceAccountOnlyGroup(t *testing.T) {
|
||||
groupID := int64(8)
|
||||
repo := &geminiAllowlistAccountRepoStub{gatewayModelsAccountRepoStub: gatewayModelsAccountRepoStub{byGroup: map[int64][]service.Account{groupID: {{
|
||||
ID: 18,
|
||||
Platform: service.PlatformGemini,
|
||||
Type: service.AccountTypeServiceAccount,
|
||||
Credentials: map[string]any{
|
||||
"location": "global",
|
||||
"service_account_json": map[string]any{
|
||||
"type": "service_account",
|
||||
"project_id": "proj",
|
||||
"private_key": "-----BEGIN PRIVATE KEY-----\nabc\n-----END PRIVATE KEY-----\n",
|
||||
"client_email": "svc@proj.iam.gserviceaccount.com",
|
||||
},
|
||||
},
|
||||
}}}}}
|
||||
upstream := &geminiMixedModelsUpstream{status: http.StatusOK, body: `{"publisherModels":[{"name":"publishers/google/models/gemini-3.8-flash"}]}`}
|
||||
tokenProvider := service.NewGeminiTokenProvider(nil, geminiVertexModelsTokenCache{}, nil)
|
||||
h := &GatewayHandler{geminiCompatService: service.NewGeminiMessagesCompatService(repo, nil, nil, nil, tokenProvider, nil, upstream, nil, &config.Config{})}
|
||||
apiKey := &service.APIKey{GroupID: &groupID, Group: &service.Group{ID: groupID, Platform: service.PlatformGemini}}
|
||||
|
||||
t.Run("list", func(t *testing.T) {
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v1beta/models", nil)
|
||||
c.Set(string(middleware.ContextKeyAPIKey), apiKey)
|
||||
|
||||
h.GeminiV1BetaListModels(c)
|
||||
require.Equal(t, http.StatusOK, rec.Code, rec.Body.String())
|
||||
require.JSONEq(t, `{"models":[{"name":"models/gemini-3.8-flash","supportedGenerationMethods":["generateContent","streamGenerateContent"]}]}`, rec.Body.String())
|
||||
})
|
||||
|
||||
t.Run("get", func(t *testing.T) {
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v1beta/models/gemini-3.8-flash", nil)
|
||||
c.Params = gin.Params{{Key: "model", Value: "gemini-3.8-flash"}}
|
||||
c.Set(string(middleware.ContextKeyAPIKey), apiKey)
|
||||
|
||||
h.GeminiV1BetaGetModel(c)
|
||||
require.Equal(t, http.StatusOK, rec.Code, rec.Body.String())
|
||||
require.JSONEq(t, `{"name":"models/gemini-3.8-flash","supportedGenerationMethods":["generateContent","streamGenerateContent"]}`, rec.Body.String())
|
||||
})
|
||||
}
|
||||
@@ -1183,6 +1183,14 @@ func OpsErrorLoggerMiddleware(ops *service.OpsService) gin.HandlerFunc {
|
||||
}
|
||||
}
|
||||
|
||||
// Google 格式错误体没有 type 字段:按网关回写时登记的错误类型补齐。
|
||||
if errType, code := service.OpsLocalErrorType(c); errType != "" && !parsed.StreamFailure &&
|
||||
(parsed.ErrorType == "" || parsed.ErrorType == "api_error") {
|
||||
parsed.ErrorType = errType
|
||||
if code != "" {
|
||||
parsed.Code = code
|
||||
}
|
||||
}
|
||||
normalizedType := normalizeOpsErrorType(parsed.ErrorType, parsed.Code)
|
||||
|
||||
phase, isBusinessLimited, errorOwner, errorSource := classifyOpsErrorLog(c, normalizedType, parsed.Message, parsed.Code, status)
|
||||
@@ -2214,6 +2222,11 @@ func classifyOpsSeverity(errType string, status int) string {
|
||||
|
||||
func classifyOpsErrorLog(c *gin.Context, errType, message, code string, status int) (phase string, isBusinessLimited bool, errorOwner string, errorSource string) {
|
||||
phase = classifyOpsPhase(errType, message, code)
|
||||
if service.IsOpsRequestScopedError(c) {
|
||||
// 失败由请求内容决定:只按错误类型分类,上游错误上下文不参与归因。
|
||||
isBusinessLimited = classifyOpsIsBusinessLimited(errType, phase, code, status, message)
|
||||
return phase, isBusinessLimited, classifyOpsErrorOwner(phase, message), classifyOpsErrorSource(phase, message)
|
||||
}
|
||||
routingCapacityLimited := isOpsRoutingCapacityLimited(c)
|
||||
clientBusinessLimited := service.HasOpsClientBusinessLimited(c)
|
||||
localModelConfiguration := clientBusinessLimited && service.OpsClientBusinessLimitedReason(c) == service.OpsClientBusinessLimitedReasonLocalModelConfiguration
|
||||
|
||||
@@ -0,0 +1,135 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// recordGoogleFormatOpsError 让 handle 以 Google 错误格式回写错误,返回 ops 中间件记录的错误日志。
|
||||
func recordGoogleFormatOpsError(t *testing.T, handle gin.HandlerFunc) *service.OpsInsertErrorLogInput {
|
||||
t.Helper()
|
||||
setupOpsErrorLogTestQueue(t, 1)
|
||||
gin.SetMode(gin.TestMode)
|
||||
ops := service.NewOpsService(nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
|
||||
router := gin.New()
|
||||
router.Use(OpsErrorLoggerMiddleware(ops))
|
||||
router.POST("/v1beta/cachedContents", handle)
|
||||
|
||||
router.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodPost, "/v1beta/cachedContents", nil))
|
||||
|
||||
require.Equal(t, int64(1), OpsErrorLogQueueLength())
|
||||
return (<-opsErrorLogQueue).entry
|
||||
}
|
||||
|
||||
func TestOpsErrorLoggerMiddleware_GoogleClientErrorsClassifiedAsClient(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
status int
|
||||
errType string
|
||||
phase string
|
||||
}{
|
||||
{status: http.StatusBadRequest, errType: "invalid_request_error", phase: "request"},
|
||||
{status: http.StatusRequestEntityTooLarge, errType: "invalid_request_error", phase: "request"},
|
||||
{status: http.StatusUnauthorized, errType: "authentication_error", phase: "auth"},
|
||||
{status: http.StatusForbidden, errType: "permission_error", phase: "request"},
|
||||
{status: http.StatusNotFound, errType: "not_found_error", phase: "request"},
|
||||
} {
|
||||
t.Run(http.StatusText(tc.status), func(t *testing.T) {
|
||||
entry := recordGoogleFormatOpsError(t, func(c *gin.Context) {
|
||||
googleError(c, tc.status, "client request problem")
|
||||
})
|
||||
require.Equal(t, tc.status, entry.StatusCode)
|
||||
require.Equal(t, tc.errType, entry.ErrorType)
|
||||
require.Equal(t, tc.phase, entry.ErrorPhase)
|
||||
require.Equal(t, "client", entry.ErrorOwner)
|
||||
require.Equal(t, "client_request", entry.ErrorSource)
|
||||
require.Equal(t, "P3", entry.Severity)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpsErrorLoggerMiddleware_GooglePlatformErrorStaysPlatform(t *testing.T) {
|
||||
entry := recordGoogleFormatOpsError(t, func(c *gin.Context) {
|
||||
googlePlatformError(c, http.StatusBadRequest, "Explicit context caching is not available for model m: storage pricing is not configured")
|
||||
})
|
||||
require.Equal(t, "api_error", entry.ErrorType)
|
||||
require.Equal(t, "internal", entry.ErrorPhase)
|
||||
require.Equal(t, "platform", entry.ErrorOwner)
|
||||
require.Equal(t, "gateway", entry.ErrorSource)
|
||||
}
|
||||
|
||||
func TestOpsErrorLoggerMiddleware_GoogleServerErrorsNotTypedByStatus(t *testing.T) {
|
||||
entry := recordGoogleFormatOpsError(t, func(c *gin.Context) {
|
||||
googleError(c, http.StatusInternalServerError, "Failed to persist cached content")
|
||||
})
|
||||
require.Equal(t, "api_error", entry.ErrorType)
|
||||
require.Equal(t, "platform", entry.ErrorOwner)
|
||||
}
|
||||
|
||||
func TestOpsErrorLoggerMiddleware_GoogleBillingErrorMatchesClaudeFormat(t *testing.T) {
|
||||
entry := recordGoogleFormatOpsError(t, func(c *gin.Context) {
|
||||
status, code, message, _ := billingErrorDetails(service.ErrInsufficientBalance)
|
||||
googleErrorWithType(c, status, code, "", message)
|
||||
})
|
||||
require.Equal(t, "billing_error", entry.ErrorType)
|
||||
require.Equal(t, "request", entry.ErrorPhase)
|
||||
require.Equal(t, "client", entry.ErrorOwner)
|
||||
require.True(t, entry.IsBusinessLimited)
|
||||
}
|
||||
|
||||
func TestOpsErrorLoggerMiddleware_GoogleConcurrencyErrorClassifiedAsRequest(t *testing.T) {
|
||||
entry := recordGoogleFormatOpsError(t, func(c *gin.Context) {
|
||||
googleConcurrencyError(c, &ConcurrencyError{SlotType: "user"}, "user")
|
||||
})
|
||||
require.Equal(t, http.StatusTooManyRequests, entry.StatusCode)
|
||||
require.Equal(t, "rate_limit_error", entry.ErrorType)
|
||||
require.Equal(t, "request", entry.ErrorPhase)
|
||||
require.Equal(t, "client", entry.ErrorOwner)
|
||||
}
|
||||
|
||||
func TestOpsErrorLoggerMiddleware_GoogleErrorAfterUpstreamFailureKeepsUpstreamAttribution(t *testing.T) {
|
||||
entry := recordGoogleFormatOpsError(t, func(c *gin.Context) {
|
||||
service.SetOpsUpstreamError(c, http.StatusBadRequest, "upstream rejected", "")
|
||||
googleError(c, http.StatusBadRequest, "upstream rejected")
|
||||
})
|
||||
require.Equal(t, "upstream", entry.ErrorPhase)
|
||||
require.Equal(t, "provider", entry.ErrorOwner)
|
||||
require.Equal(t, "upstream_http", entry.ErrorSource)
|
||||
}
|
||||
|
||||
func TestOpsErrorLoggerMiddleware_RequestScopedUpstreamErrorClassifiedAsClient(t *testing.T) {
|
||||
entry := recordGoogleFormatOpsError(t, func(c *gin.Context) {
|
||||
service.SetOpsUpstreamError(c, http.StatusBadRequest, "Invalid resource state for cache content 8079767076522164224.", "")
|
||||
c.Set(service.OpsUpstreamErrorsKey, []*service.OpsUpstreamErrorEvent{{
|
||||
Platform: service.PlatformGemini, AccountID: 173, UpstreamStatusCode: http.StatusBadRequest, Kind: "http_error",
|
||||
Message: "Invalid resource state for cache content 8079767076522164224.",
|
||||
}})
|
||||
service.MarkOpsRequestScopedError(c, "permission_error")
|
||||
c.Data(http.StatusForbidden, "application/json", []byte(service.GeminiCachedContentNotFoundResponse))
|
||||
})
|
||||
require.Equal(t, http.StatusForbidden, entry.StatusCode)
|
||||
require.Equal(t, "permission_error", entry.ErrorType)
|
||||
require.Equal(t, "request", entry.ErrorPhase)
|
||||
require.Equal(t, "client", entry.ErrorOwner)
|
||||
require.Equal(t, "client_request", entry.ErrorSource)
|
||||
require.False(t, entry.IsBusinessLimited)
|
||||
require.NotNil(t, entry.UpstreamStatusCode, "上游原始错误保留为明细")
|
||||
require.Equal(t, http.StatusBadRequest, *entry.UpstreamStatusCode)
|
||||
}
|
||||
|
||||
func TestSetOpsLocalErrorTypeClearsRequestScopedMark(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
service.MarkOpsRequestScopedError(c, "permission_error")
|
||||
require.True(t, service.IsOpsRequestScopedError(c))
|
||||
|
||||
service.SetOpsLocalErrorType(c, "", "")
|
||||
require.False(t, service.IsOpsRequestScopedError(c))
|
||||
errType, code := service.OpsLocalErrorType(c)
|
||||
require.Empty(t, errType)
|
||||
require.Empty(t, code)
|
||||
}
|
||||
@@ -100,9 +100,14 @@ func googleSecurityAuditError(c *gin.Context, decision *securityaudit.Decision)
|
||||
return
|
||||
}
|
||||
if decision.Legacy != nil && decision.Legacy.Blocked {
|
||||
googleError(c, securityAuditStatus(decision), securityAuditMessage(decision))
|
||||
googleErrorWithType(c, securityAuditStatus(decision), securityAuditErrorCode(decision), "", securityAuditMessage(decision))
|
||||
return
|
||||
}
|
||||
errType := "api_error"
|
||||
if decision.Kind == securityaudit.DecisionBlock {
|
||||
errType = "permission_error"
|
||||
}
|
||||
service.SetOpsLocalErrorType(c, errType, securityAuditErrorCode(decision))
|
||||
status := securityAuditStatus(decision)
|
||||
googleStatus := googleapi.HTTPStatusToGoogleStatus(status)
|
||||
if status == http.StatusServiceUnavailable {
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -857,7 +857,7 @@ func (r *apiKeyRepository) ResetRateLimitWindows(ctx context.Context, id int64)
|
||||
// GetRateLimitData returns the current rate limit usage and window start times for an API key.
|
||||
func (r *apiKeyRepository) GetRateLimitData(ctx context.Context, id int64) (result *service.APIKeyRateLimitData, err error) {
|
||||
rows, err := r.sql.QueryContext(ctx, `
|
||||
SELECT usage_5h, usage_1d, usage_7d, window_5h_start, window_1d_start, window_7d_start
|
||||
SELECT quota_used, usage_5h, usage_1d, usage_7d, window_5h_start, window_1d_start, window_7d_start
|
||||
FROM api_keys
|
||||
WHERE id = $1 AND deleted_at IS NULL`,
|
||||
id)
|
||||
@@ -873,7 +873,7 @@ func (r *apiKeyRepository) GetRateLimitData(ctx context.Context, id int64) (resu
|
||||
return nil, service.ErrAPIKeyNotFound
|
||||
}
|
||||
data := &service.APIKeyRateLimitData{}
|
||||
if err := rows.Scan(&data.Usage5h, &data.Usage1d, &data.Usage7d, &data.Window5hStart, &data.Window1dStart, &data.Window7dStart); err != nil {
|
||||
if err := rows.Scan(&data.QuotaUsed, &data.Usage5h, &data.Usage1d, &data.Usage7d, &data.Window5hStart, &data.Window1dStart, &data.Window7dStart); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return data, rows.Err()
|
||||
|
||||
@@ -464,6 +464,10 @@ func (s *APIKeyRepoSuite) TestIncrementQuotaUsed_Basic() {
|
||||
newQuota, err = s.repo.IncrementQuotaUsed(s.ctx, key.ID, 2.5)
|
||||
s.Require().NoError(err, "IncrementQuotaUsed second")
|
||||
s.Require().Equal(4.0, newQuota, "第二次递增后应为 4.0")
|
||||
|
||||
data, err := s.repo.GetRateLimitData(s.ctx, key.ID)
|
||||
s.Require().NoError(err, "GetRateLimitData")
|
||||
s.Require().Equal(4.0, data.QuotaUsed, "用量读取带回总额度已用量")
|
||||
}
|
||||
|
||||
func (s *APIKeyRepoSuite) TestIncrementQuotaUsed_NotFound() {
|
||||
|
||||
@@ -25,6 +25,20 @@ func billingInflightKeys(userID int64) (string, string) {
|
||||
fmt.Sprintf("%s{%d}", billingInflightAmtKeyPrefix, userID)
|
||||
}
|
||||
|
||||
// 订阅 / API Key 维度在途预留(确定金额费用):
|
||||
// - billing:inflight_scope:{scope} ZSET member=requestID score=过期时间(ms)
|
||||
// - billing:inflight_scope_amt:{scope} HASH field=requestID value=预留金额(USD)
|
||||
//
|
||||
// 过期清理与释放语义与余额维度相同,不续期。
|
||||
const (
|
||||
billingInflightScopeKeyPrefix = "billing:inflight_scope:"
|
||||
billingInflightScopeAmtKeyPrefix = "billing:inflight_scope_amt:"
|
||||
)
|
||||
|
||||
func billingInflightScopeKeys(scope string) (string, string) {
|
||||
return billingInflightScopeKeyPrefix + "{" + scope + "}", billingInflightScopeAmtKeyPrefix + "{" + scope + "}"
|
||||
}
|
||||
|
||||
var (
|
||||
// KEYS: [1]=zset [2]=hash
|
||||
// ARGV: [1]=now_ms [2]=expire_at_ms [3]=member [4]=amount [5]=balance [6]=key_ttl_ms
|
||||
@@ -85,6 +99,28 @@ var (
|
||||
redis.call('HDEL', KEYS[2], ARGV[1])
|
||||
return 1
|
||||
`)
|
||||
|
||||
// KEYS: [1]=zset [2]=hash
|
||||
// ARGV: [1]=now_ms [2]=expire_at_ms [3]=member [4]=amount [5]=key_ttl_ms
|
||||
// 先清理已过期成员,再登记本次预留;返回登记前其他在途预留的合计(string),是否放行由调用方按权威用量判定。
|
||||
reserveInflightScopeScript = redis.NewScript(`
|
||||
local now = tonumber(ARGV[1])
|
||||
local expired = redis.call('ZRANGEBYSCORE', KEYS[1], '-inf', now)
|
||||
if #expired > 0 then
|
||||
redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', now)
|
||||
redis.call('HDEL', KEYS[2], unpack(expired))
|
||||
end
|
||||
local vals = redis.call('HVALS', KEYS[2])
|
||||
local sum = 0
|
||||
for i = 1, #vals do
|
||||
sum = sum + (tonumber(vals[i]) or 0)
|
||||
end
|
||||
redis.call('ZADD', KEYS[1], tonumber(ARGV[2]), ARGV[3])
|
||||
redis.call('HSET', KEYS[2], ARGV[3], ARGV[4])
|
||||
redis.call('PEXPIRE', KEYS[1], ARGV[5])
|
||||
redis.call('PEXPIRE', KEYS[2], ARGV[5])
|
||||
return tostring(sum)
|
||||
`)
|
||||
)
|
||||
|
||||
// ReserveInflightBalance 实现 service.InflightBalanceReservationCache。
|
||||
@@ -138,7 +174,35 @@ func (c *billingCache) RenewInflightBalance(ctx context.Context, userID int64, r
|
||||
return n == 1, nil
|
||||
}
|
||||
|
||||
// ReserveInflightScope 实现 service.InflightScopeReservationCache。
|
||||
func (c *billingCache) ReserveInflightScope(ctx context.Context, scope, requestID string, amount float64, ttl time.Duration) (float64, error) {
|
||||
zkey, hkey := billingInflightScopeKeys(scope)
|
||||
now := time.Now().UnixMilli()
|
||||
ttlMs := ttl.Milliseconds()
|
||||
if ttlMs <= 0 {
|
||||
ttlMs = 1
|
||||
}
|
||||
res, err := reserveInflightScopeScript.Run(ctx, c.rdb, []string{zkey, hkey},
|
||||
now,
|
||||
now+ttlMs,
|
||||
requestID,
|
||||
strconv.FormatFloat(amount, 'f', -1, 64),
|
||||
ttlMs,
|
||||
).Text()
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return strconv.ParseFloat(res, 64)
|
||||
}
|
||||
|
||||
// ReleaseInflightScope 实现 service.InflightScopeReservationCache。
|
||||
func (c *billingCache) ReleaseInflightScope(ctx context.Context, scope, requestID string) error {
|
||||
zkey, hkey := billingInflightScopeKeys(scope)
|
||||
return releaseInflightBalanceScript.Run(ctx, c.rdb, []string{zkey, hkey}, requestID).Err()
|
||||
}
|
||||
|
||||
var (
|
||||
_ service.InflightBalanceReservationCache = (*billingCache)(nil)
|
||||
_ service.InflightBalanceReservationRenewer = (*billingCache)(nil)
|
||||
_ service.InflightScopeReservationCache = (*billingCache)(nil)
|
||||
)
|
||||
|
||||
@@ -285,3 +285,72 @@ func TestInflightReservation_RenewDoesNotResurrectExpiredMember(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
require.False(t, exists)
|
||||
}
|
||||
|
||||
func TestInflightScopeReservation_ReturnsOthersAndReleases(t *testing.T) {
|
||||
mr, cache, _ := newInflightTestEnv(t, true, 60)
|
||||
ctx := context.Background()
|
||||
|
||||
others, err := cache.ReserveInflightScope(ctx, "apikey:7", "r1", 3, time.Minute)
|
||||
require.NoError(t, err)
|
||||
require.Zero(t, others)
|
||||
others, err = cache.ReserveInflightScope(ctx, "apikey:7", "r2", 2, time.Minute)
|
||||
require.NoError(t, err)
|
||||
require.InDelta(t, 3, others, 1e-9, "返回登记前其他在途预留的合计")
|
||||
others, err = cache.ReserveInflightScope(ctx, "sub:7", "r3", 5, time.Minute)
|
||||
require.NoError(t, err)
|
||||
require.Zero(t, others, "不同维度互不影响")
|
||||
|
||||
require.NoError(t, cache.ReleaseInflightScope(ctx, "apikey:7", "r1"))
|
||||
require.NoError(t, cache.ReleaseInflightScope(ctx, "apikey:7", "r1"))
|
||||
others, err = cache.ReserveInflightScope(ctx, "apikey:7", "r4", 1, time.Minute)
|
||||
require.NoError(t, err)
|
||||
require.InDelta(t, 2, others, 1e-9)
|
||||
|
||||
mr.FastForward(2 * time.Minute)
|
||||
others, err = cache.ReserveInflightScope(ctx, "apikey:7", "r5", 1, time.Minute)
|
||||
require.NoError(t, err)
|
||||
require.Zero(t, others, "过期预留不再计入")
|
||||
}
|
||||
|
||||
func TestReserveFixedCharge_ConcurrentChargesAdmitOnlyWhatAPIKeyQuotaCovers(t *testing.T) {
|
||||
_, cache, svc := newInflightTestEnv(t, true, 60)
|
||||
ctx := context.Background()
|
||||
user := &service.User{ID: 42}
|
||||
apiKey := &service.APIKey{ID: 9, UserID: user.ID, Quota: 1.0}
|
||||
require.NoError(t, cache.SetUserBalance(ctx, user.ID, 100))
|
||||
|
||||
const n = 20
|
||||
var admitted atomic.Int32
|
||||
reservations := make(chan *service.InflightReservation, n)
|
||||
var wg sync.WaitGroup
|
||||
start := make(chan struct{})
|
||||
for i := 0; i < n; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
<-start
|
||||
res, err := svc.ReserveFixedCharge(ctx, user, apiKey, nil, nil, 0.3)
|
||||
if err != nil {
|
||||
if !errors.Is(err, service.ErrAPIKeyQuotaExhausted) {
|
||||
t.Errorf("unexpected error: %v", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
admitted.Add(1)
|
||||
reservations <- res
|
||||
}()
|
||||
}
|
||||
close(start)
|
||||
wg.Wait()
|
||||
close(reservations)
|
||||
|
||||
require.Equal(t, int32(3), admitted.Load(), "Key 额度 1.0、每笔 0.3:同时在途的不超过 3 笔")
|
||||
for res := range reservations {
|
||||
res.Release()
|
||||
}
|
||||
zkey, _ := billingInflightScopeKeys("apikey:9")
|
||||
remaining, err := cache.rdb.ZCard(ctx, zkey).Result()
|
||||
require.NoError(t, err)
|
||||
require.Zero(t, remaining)
|
||||
require.Equal(t, int64(0), inflightCount(t, cache, user.ID))
|
||||
}
|
||||
|
||||
@@ -0,0 +1,205 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
)
|
||||
|
||||
type geminiCachedContentSQLExecutor interface {
|
||||
ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error)
|
||||
QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error)
|
||||
QueryRowContext(ctx context.Context, query string, args ...any) *sql.Row
|
||||
}
|
||||
|
||||
type geminiCachedContentRepository struct {
|
||||
sql geminiCachedContentSQLExecutor
|
||||
}
|
||||
|
||||
func NewGeminiCachedContentRepository(db *sql.DB) service.GeminiCachedContentRepository {
|
||||
return &geminiCachedContentRepository{sql: db}
|
||||
}
|
||||
|
||||
const geminiCachedContentColumns = `id, public_id, user_id, api_key_id, group_id, account_id, upstream_name, model,
|
||||
upstream_model, display_name, total_token_count, expire_time, channel_usage, created_at, updated_at`
|
||||
|
||||
func (r *geminiCachedContentRepository) Create(ctx context.Context, record *service.GeminiCachedContent) error {
|
||||
channelUsage, err := json.Marshal(record.ChannelUsage)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
row := r.sql.QueryRowContext(ctx, `
|
||||
INSERT INTO gemini_cached_contents (public_id, user_id, api_key_id, group_id, account_id, upstream_name, model,
|
||||
upstream_model, display_name, total_token_count, expire_time, channel_usage)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)
|
||||
RETURNING id, created_at, updated_at`,
|
||||
record.PublicID, record.UserID, record.APIKeyID, record.GroupID, record.AccountID, record.UpstreamName,
|
||||
record.Model, record.UpstreamModel, record.DisplayName, record.TotalTokenCount, record.ExpireTime, channelUsage)
|
||||
if err := row.Scan(&record.ID, &record.CreatedAt, &record.UpdatedAt); err != nil {
|
||||
return translatePersistenceError(err, nil, service.ErrGeminiCachedContentExists)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentRepository) GetForOwner(ctx context.Context, apiKeyID int64, publicID string) (*service.GeminiCachedContent, error) {
|
||||
record, err := scanGeminiCachedContent(r.sql.QueryRowContext(ctx, `SELECT `+geminiCachedContentColumns+`
|
||||
FROM gemini_cached_contents
|
||||
WHERE public_id = $1 AND api_key_id = $2 AND deleted_at IS NULL`, publicID, apiKeyID))
|
||||
if err != nil {
|
||||
return nil, translatePersistenceError(err, service.ErrGeminiCachedContentNotFound, nil)
|
||||
}
|
||||
return record, nil
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentRepository) ListForOwner(ctx context.Context, apiKeyID, groupID int64, activeAt time.Time, beforeID int64, limit int) ([]*service.GeminiCachedContent, error) {
|
||||
if limit <= 0 {
|
||||
limit = service.GeminiCachedContentListDefaultPageSize
|
||||
}
|
||||
rows, err := r.sql.QueryContext(ctx, `SELECT `+geminiCachedContentColumns+`
|
||||
FROM gemini_cached_contents
|
||||
WHERE api_key_id = $1 AND group_id = $2 AND deleted_at IS NULL AND expire_time > $3 AND ($4 <= 0 OR id < $4)
|
||||
ORDER BY id DESC
|
||||
LIMIT $5`, apiKeyID, groupID, activeAt, beforeID, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer func() { _ = rows.Close() }()
|
||||
var out []*service.GeminiCachedContent
|
||||
for rows.Next() {
|
||||
record, err := scanGeminiCachedContent(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, record)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentRepository) BeginMutation(ctx context.Context, id int64, baseline time.Time, token string, lease time.Duration) (bool, error) {
|
||||
result, err := r.sql.ExecContext(ctx, `UPDATE gemini_cached_contents
|
||||
SET mutation_token = $2, mutation_lease_until = NOW() + make_interval(secs => $3::double precision)
|
||||
WHERE id = $1 AND deleted_at IS NULL AND mutation_token IS NULL AND expire_time = $4`,
|
||||
id, token, lease.Seconds(), baseline)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
affected, err := result.RowsAffected()
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
if affected > 0 {
|
||||
return true, nil
|
||||
}
|
||||
var exists bool
|
||||
if err := r.sql.QueryRowContext(ctx, `SELECT EXISTS (
|
||||
SELECT 1 FROM gemini_cached_contents WHERE id = $1 AND deleted_at IS NULL)`, id).Scan(&exists); err != nil {
|
||||
return false, err
|
||||
}
|
||||
if !exists {
|
||||
return false, service.ErrGeminiCachedContentNotFound
|
||||
}
|
||||
return false, nil
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentRepository) TakeOverStaleMutation(ctx context.Context, id int64, token string, lease time.Duration) (*service.GeminiCachedContent, error) {
|
||||
record, err := scanGeminiCachedContent(r.sql.QueryRowContext(ctx, `UPDATE gemini_cached_contents
|
||||
SET mutation_token = $2, mutation_lease_until = NOW() + make_interval(secs => $3::double precision)
|
||||
WHERE id = $1 AND deleted_at IS NULL AND mutation_token IS NOT NULL AND mutation_lease_until < NOW()
|
||||
RETURNING `+geminiCachedContentColumns, id, token, lease.Seconds()))
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return nil, nil
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return record, nil
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentRepository) CommitMutation(ctx context.Context, id int64, token string, expireTime time.Time) error {
|
||||
result, err := r.sql.ExecContext(ctx, `UPDATE gemini_cached_contents
|
||||
SET expire_time = $3, updated_at = NOW(), mutation_token = NULL, mutation_lease_until = NULL
|
||||
WHERE id = $1 AND deleted_at IS NULL AND mutation_token = $2`, id, token, expireTime)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
affected, err := result.RowsAffected()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if affected == 0 {
|
||||
return service.ErrGeminiCachedContentMutationConflict
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentRepository) ReleaseMutation(ctx context.Context, id int64, token string) error {
|
||||
_, err := r.sql.ExecContext(ctx, `UPDATE gemini_cached_contents
|
||||
SET mutation_token = NULL, mutation_lease_until = NULL
|
||||
WHERE id = $1 AND mutation_token = $2`, id, token)
|
||||
return err
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentRepository) ListStaleMutations(ctx context.Context, limit int) ([]int64, error) {
|
||||
rows, err := r.sql.QueryContext(ctx, `SELECT id FROM gemini_cached_contents
|
||||
WHERE mutation_token IS NOT NULL AND deleted_at IS NULL AND mutation_lease_until < NOW()
|
||||
ORDER BY mutation_lease_until
|
||||
LIMIT $1`, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer func() { _ = rows.Close() }()
|
||||
var ids []int64
|
||||
for rows.Next() {
|
||||
var id int64
|
||||
if err := rows.Scan(&id); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
ids = append(ids, id)
|
||||
}
|
||||
return ids, rows.Err()
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentRepository) SoftDelete(ctx context.Context, id int64) error {
|
||||
_, err := r.sql.ExecContext(ctx, `UPDATE gemini_cached_contents
|
||||
SET deleted_at = NOW(), updated_at = NOW()
|
||||
WHERE id = $1 AND deleted_at IS NULL`, id)
|
||||
return err
|
||||
}
|
||||
|
||||
type geminiCachedContentScanner interface {
|
||||
Scan(dest ...any) error
|
||||
}
|
||||
|
||||
func scanGeminiCachedContent(row geminiCachedContentScanner) (*service.GeminiCachedContent, error) {
|
||||
var record service.GeminiCachedContent
|
||||
var channelUsage []byte
|
||||
if err := row.Scan(
|
||||
&record.ID,
|
||||
&record.PublicID,
|
||||
&record.UserID,
|
||||
&record.APIKeyID,
|
||||
&record.GroupID,
|
||||
&record.AccountID,
|
||||
&record.UpstreamName,
|
||||
&record.Model,
|
||||
&record.UpstreamModel,
|
||||
&record.DisplayName,
|
||||
&record.TotalTokenCount,
|
||||
&record.ExpireTime,
|
||||
&channelUsage,
|
||||
&record.CreatedAt,
|
||||
&record.UpdatedAt,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(channelUsage) > 0 {
|
||||
if err := json.Unmarshal(channelUsage, &record.ChannelUsage); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
return &record, nil
|
||||
}
|
||||
@@ -0,0 +1,120 @@
|
||||
//go:build integration
|
||||
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestGeminiCachedContentRepository_Lifecycle(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
tx := testTx(t)
|
||||
repo := &geminiCachedContentRepository{sql: tx}
|
||||
now := time.Now().UTC().Truncate(time.Microsecond)
|
||||
suffix := time.Now().Format("150405.000000")
|
||||
apiKeyID := int64(910001)
|
||||
|
||||
newRecord := func(publicID string, keyID, groupID int64, expire time.Time) *service.GeminiCachedContent {
|
||||
return &service.GeminiCachedContent{
|
||||
PublicID: publicID + suffix, UserID: 1, APIKeyID: keyID, GroupID: groupID, AccountID: 99,
|
||||
UpstreamName: "cachedContents/up" + publicID, Model: "gemini-3.8-flash", UpstreamModel: "models/gemini-3.8-flash",
|
||||
DisplayName: "label", TotalTokenCount: 2730, ExpireTime: expire,
|
||||
ChannelUsage: service.ChannelUsageFields{ChannelID: 7, OriginalModel: "gemini-pro", ChannelMappedModel: "gemini-3.8-flash", BillingModelSource: "requested"},
|
||||
}
|
||||
}
|
||||
first := newRecord("a", apiKeyID, 42, now.Add(time.Hour))
|
||||
require.NoError(t, repo.Create(ctx, first))
|
||||
require.Positive(t, first.ID)
|
||||
require.False(t, first.CreatedAt.IsZero())
|
||||
_, err := tx.ExecContext(ctx, `SAVEPOINT gemini_cached_content_duplicate`)
|
||||
require.NoError(t, err)
|
||||
require.ErrorIs(t, repo.Create(ctx, newRecord("a", apiKeyID, 42, now.Add(time.Hour))), service.ErrGeminiCachedContentExists)
|
||||
_, err = tx.ExecContext(ctx, `ROLLBACK TO SAVEPOINT gemini_cached_content_duplicate`)
|
||||
require.NoError(t, err)
|
||||
|
||||
second := newRecord("b", apiKeyID, 42, now.Add(time.Hour))
|
||||
require.NoError(t, repo.Create(ctx, second))
|
||||
require.NoError(t, repo.Create(ctx, newRecord("expired", apiKeyID, 42, now.Add(-time.Minute))))
|
||||
require.NoError(t, repo.Create(ctx, newRecord("othergroup", apiKeyID, 43, now.Add(time.Hour))))
|
||||
require.NoError(t, repo.Create(ctx, newRecord("otherkey", apiKeyID+1, 42, now.Add(time.Hour))))
|
||||
|
||||
got, err := repo.GetForOwner(ctx, apiKeyID, first.PublicID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "cachedContents/upa", got.UpstreamName)
|
||||
require.Equal(t, int64(2730), got.TotalTokenCount)
|
||||
require.True(t, got.ExpireTime.Equal(first.ExpireTime))
|
||||
require.Equal(t, first.ChannelUsage, got.ChannelUsage, "创建时的渠道计费口径随记录持久化")
|
||||
_, err = repo.GetForOwner(ctx, apiKeyID+1, first.PublicID)
|
||||
require.ErrorIs(t, err, service.ErrGeminiCachedContentNotFound)
|
||||
|
||||
list, err := repo.ListForOwner(ctx, apiKeyID, 42, now, 0, 10)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, list, 2)
|
||||
require.Equal(t, second.PublicID, list[0].PublicID)
|
||||
page, err := repo.ListForOwner(ctx, apiKeyID, 42, now, list[0].ID, 10)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, page, 1)
|
||||
require.Equal(t, first.PublicID, page[0].PublicID)
|
||||
|
||||
lease := 2 * time.Minute
|
||||
acquired, err := repo.BeginMutation(ctx, first.ID, first.ExpireTime, "t1", lease)
|
||||
require.NoError(t, err)
|
||||
require.True(t, acquired)
|
||||
acquired, err = repo.BeginMutation(ctx, first.ID, first.ExpireTime, "t2", lease)
|
||||
require.NoError(t, err)
|
||||
require.False(t, acquired, "有进行中的修改时不能再登记")
|
||||
taken, err := repo.TakeOverStaleMutation(ctx, first.ID, "t2", lease)
|
||||
require.NoError(t, err)
|
||||
require.Nil(t, taken, "租约未过期不能接管")
|
||||
require.ErrorIs(t, repo.CommitMutation(ctx, first.ID, "t2", now), service.ErrGeminiCachedContentMutationConflict)
|
||||
|
||||
newExpire := now.Add(3 * time.Hour)
|
||||
require.NoError(t, repo.CommitMutation(ctx, first.ID, "t1", newExpire))
|
||||
got, err = repo.GetForOwner(ctx, apiKeyID, first.PublicID)
|
||||
require.NoError(t, err)
|
||||
require.True(t, got.ExpireTime.Equal(newExpire))
|
||||
require.ErrorIs(t, repo.CommitMutation(ctx, first.ID, "t1", now), service.ErrGeminiCachedContentMutationConflict, "提交后修改意图已清除")
|
||||
|
||||
acquired, err = repo.BeginMutation(ctx, first.ID, first.ExpireTime, "t3", lease)
|
||||
require.NoError(t, err)
|
||||
require.False(t, acquired, "到期时间已被其他修改改变")
|
||||
acquired, err = repo.BeginMutation(ctx, first.ID, newExpire, "t3", lease)
|
||||
require.NoError(t, err)
|
||||
require.True(t, acquired)
|
||||
stale, err := repo.ListStaleMutations(ctx, 10)
|
||||
require.NoError(t, err)
|
||||
require.NotContains(t, stale, first.ID)
|
||||
|
||||
_, err = tx.ExecContext(ctx, `UPDATE gemini_cached_contents SET mutation_lease_until = NOW() - INTERVAL '1 second' WHERE id = $1`, first.ID)
|
||||
require.NoError(t, err)
|
||||
stale, err = repo.ListStaleMutations(ctx, 10)
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, stale, first.ID)
|
||||
taken, err = repo.TakeOverStaleMutation(ctx, first.ID, "t4", lease)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, taken)
|
||||
require.True(t, taken.ExpireTime.Equal(newExpire), "接管返回修改前的到期时间")
|
||||
require.ErrorIs(t, repo.CommitMutation(ctx, first.ID, "t3", now), service.ErrGeminiCachedContentMutationConflict, "被接管后原持有者不能提交")
|
||||
require.NoError(t, repo.ReleaseMutation(ctx, first.ID, "t3"))
|
||||
acquired, err = repo.BeginMutation(ctx, first.ID, newExpire, "t5", lease)
|
||||
require.NoError(t, err)
|
||||
require.False(t, acquired, "原持有者释放不影响接管者")
|
||||
require.NoError(t, repo.ReleaseMutation(ctx, first.ID, "t4"))
|
||||
stale, err = repo.ListStaleMutations(ctx, 10)
|
||||
require.NoError(t, err)
|
||||
require.NotContains(t, stale, first.ID)
|
||||
|
||||
require.NoError(t, repo.SoftDelete(ctx, first.ID))
|
||||
_, err = repo.GetForOwner(ctx, apiKeyID, first.PublicID)
|
||||
require.ErrorIs(t, err, service.ErrGeminiCachedContentNotFound)
|
||||
_, err = repo.BeginMutation(ctx, first.ID, newExpire, "t6", lease)
|
||||
require.ErrorIs(t, err, service.ErrGeminiCachedContentNotFound)
|
||||
list, err = repo.ListForOwner(ctx, apiKeyID, 42, now, 0, 10)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, list, 1)
|
||||
}
|
||||
@@ -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,
|
||||
|
||||
@@ -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"},
|
||||
|
||||
@@ -339,6 +339,13 @@ func (a *Account) IsGeminiGoogleOne() bool {
|
||||
return a.Platform == PlatformGemini && a.Type == AccountTypeOAuth && a.GeminiOAuthType() == "google_one"
|
||||
}
|
||||
|
||||
// IsGeminiThirdPartyAPIKey reports whether a Gemini API Key account points at a
|
||||
// non-Google base URL. Such upstreams carry no AI Studio tier or daily quota.
|
||||
func (a *Account) IsGeminiThirdPartyAPIKey() bool {
|
||||
return a != nil && a.Platform == PlatformGemini && a.Type == AccountTypeAPIKey &&
|
||||
!isOfficialGeminiAPIBaseURL(a.GetCredential("base_url"))
|
||||
}
|
||||
|
||||
func (a *Account) CanGetUsage() bool {
|
||||
return a.Type == AccountTypeOAuth
|
||||
}
|
||||
|
||||
@@ -933,7 +933,7 @@ func TestResolveAccountStatsCost_Gemini36FlashTierUsesFallbackPricing(t *testing
|
||||
UsageTokens{InputTokens: 1_000_000, OutputTokens: 1_000_000, CacheReadTokens: 1_000_000}, 1, 0, "", time.Time{}, true,
|
||||
)
|
||||
require.NotNil(t, result)
|
||||
require.InDelta(t, 9.15, *result, 1e-12)
|
||||
require.InDelta(t, 4.575, *result, 1e-12)
|
||||
}
|
||||
|
||||
func TestResolveAccountStatsCost_AllMiss_ReturnsNil(t *testing.T) {
|
||||
|
||||
@@ -2585,14 +2585,7 @@ func (s *AccountTestService) buildGeminiOAuthRequest(ctx context.Context, accoun
|
||||
}
|
||||
|
||||
func (s *AccountTestService) buildGeminiServiceAccountRequest(ctx context.Context, account *Account, modelID string, payload []byte) (*http.Request, error) {
|
||||
if s.geminiTokenProvider == nil {
|
||||
return nil, fmt.Errorf("gemini token provider not configured")
|
||||
}
|
||||
accessToken, err := s.geminiTokenProvider.GetAccessToken(ctx, account)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get service account access token: %w", err)
|
||||
}
|
||||
fullURL, err := buildVertexGeminiURL(account.VertexProjectID(), account.VertexLocation(modelID), modelID, "streamGenerateContent", true)
|
||||
fullURL, err := buildVertexGeminiAccountURL(account, modelID, "streamGenerateContent", true)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -2601,7 +2594,9 @@ func (s *AccountTestService) buildGeminiServiceAccountRequest(ctx context.Contex
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Authorization", "Bearer "+accessToken)
|
||||
if err := setVertexGeminiAuth(ctx, req, account, s.geminiTokenProvider); err != nil {
|
||||
return nil, fmt.Errorf("failed to authenticate vertex request: %w", err)
|
||||
}
|
||||
return req, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -130,6 +130,8 @@ type apiKeyAllByUserIDLister interface {
|
||||
|
||||
// APIKeyRateLimitData holds rate limit usage and window state for an API key.
|
||||
type APIKeyRateLimitData struct {
|
||||
// QuotaUsed 是总额度已用量(quota_used)。
|
||||
QuotaUsed float64
|
||||
Usage5h float64
|
||||
Usage1d float64
|
||||
Usage7d float64
|
||||
|
||||
@@ -3,6 +3,7 @@ package service
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"math"
|
||||
"strconv"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
@@ -573,6 +574,11 @@ func (s *BillingCacheService) InvalidateAPIKeyRateLimit(ctx context.Context, key
|
||||
// resets expired windows in-memory and triggers async DB reset,
|
||||
// and returns an error if any window limit is exceeded.
|
||||
func (s *BillingCacheService) checkAPIKeyRateLimits(ctx context.Context, apiKey *APIKey) error {
|
||||
return s.checkAPIKeyRateLimitsWithPending(ctx, apiKey, 0)
|
||||
}
|
||||
|
||||
// checkAPIKeyRateLimitsWithPending 在窗口用量上叠加 pending(即将扣除的金额)后检查限速窗口。
|
||||
func (s *BillingCacheService) checkAPIKeyRateLimitsWithPending(ctx context.Context, apiKey *APIKey, pending float64) error {
|
||||
if s.cache == nil {
|
||||
// No cache: fall back to reading from DB directly
|
||||
if s.apiKeyRateLimitLoader == nil {
|
||||
@@ -583,7 +589,7 @@ func (s *BillingCacheService) checkAPIKeyRateLimits(ctx context.Context, apiKey
|
||||
return nil // Don't block requests on DB errors
|
||||
}
|
||||
return s.evaluateRateLimits(ctx, apiKey, data.Usage5h, data.Usage1d, data.Usage7d,
|
||||
data.Window5hStart, data.Window1dStart, data.Window7dStart)
|
||||
data.Window5hStart, data.Window1dStart, data.Window7dStart, pending)
|
||||
}
|
||||
|
||||
cacheData, err := s.cache.GetAPIKeyRateLimit(ctx, apiKey.ID)
|
||||
@@ -628,11 +634,11 @@ func (s *BillingCacheService) checkAPIKeyRateLimits(ctx context.Context, apiKey
|
||||
t := time.Unix(cacheData.Window7d, 0)
|
||||
w7d = &t
|
||||
}
|
||||
return s.evaluateRateLimits(ctx, apiKey, cacheData.Usage5h, cacheData.Usage1d, cacheData.Usage7d, w5h, w1d, w7d)
|
||||
return s.evaluateRateLimits(ctx, apiKey, cacheData.Usage5h, cacheData.Usage1d, cacheData.Usage7d, w5h, w1d, w7d, pending)
|
||||
}
|
||||
|
||||
// evaluateRateLimits checks usage against limits, triggering async resets for expired windows.
|
||||
func (s *BillingCacheService) evaluateRateLimits(ctx context.Context, apiKey *APIKey, usage5h, usage1d, usage7d float64, w5h, w1d, w7d *time.Time) error {
|
||||
// evaluateRateLimits checks usage (plus pending) against limits, triggering async resets for expired windows.
|
||||
func (s *BillingCacheService) evaluateRateLimits(ctx context.Context, apiKey *APIKey, usage5h, usage1d, usage7d float64, w5h, w1d, w7d *time.Time, pending float64) error {
|
||||
needsReset := false
|
||||
|
||||
// Reset expired windows in-memory for check purposes
|
||||
@@ -675,13 +681,16 @@ func (s *BillingCacheService) evaluateRateLimits(ctx context.Context, apiKey *AP
|
||||
}
|
||||
|
||||
// Check limits
|
||||
if apiKey.RateLimit5h > 0 && usage5h >= apiKey.RateLimit5h {
|
||||
exceeds := func(usage, limit float64) bool {
|
||||
return limit > 0 && (usage >= limit || usage+pending > limit)
|
||||
}
|
||||
if exceeds(usage5h, apiKey.RateLimit5h) {
|
||||
return ErrAPIKeyRateLimit5hExceeded
|
||||
}
|
||||
if apiKey.RateLimit1d > 0 && usage1d >= apiKey.RateLimit1d {
|
||||
if exceeds(usage1d, apiKey.RateLimit1d) {
|
||||
return ErrAPIKeyRateLimit1dExceeded
|
||||
}
|
||||
if apiKey.RateLimit7d > 0 && usage7d >= apiKey.RateLimit7d {
|
||||
if exceeds(usage7d, apiKey.RateLimit7d) {
|
||||
return ErrAPIKeyRateLimit7dExceeded
|
||||
}
|
||||
return nil
|
||||
@@ -781,6 +790,110 @@ func (s *BillingCacheService) CheckBillingEligibility(ctx context.Context, user
|
||||
return nil
|
||||
}
|
||||
|
||||
// ReserveFixedCharge 为扣费前即可确定金额的费用(倍率后 USD,如显式缓存存储费)做准入并登记在途预留。
|
||||
//
|
||||
// 余额模式要求缓存余额覆盖 amount 及用户的其他在途预留;订阅模式要求日 / 周 / 月剩余额度覆盖 amount;
|
||||
// API Key 总额度与限速窗口同样叠加 amount。订阅与 API Key 的用量取库中权威值,并叠加同一订阅 / Key 上其他
|
||||
// 在途预留:先登记本次预留再读用量,并发的两笔费用中后读用量的一方必然看到另一方(预留或已落库的扣费)。
|
||||
// 返回的句柄可能为 nil(未登记任何预留);调用方须在 handler 结束时调用 HandlerDone,预留在计费任务落库后释放。
|
||||
// 须在 CheckBillingEligibility 通过后调用;简易模式不校验。
|
||||
func (s *BillingCacheService) ReserveFixedCharge(ctx context.Context, user *User, apiKey *APIKey, group *Group, subscription *UserSubscription, amount float64) (*InflightReservation, error) {
|
||||
if user == nil || amount <= 0 || math.IsNaN(amount) || math.IsInf(amount, 0) {
|
||||
return nil, nil
|
||||
}
|
||||
if s.cfg != nil && s.cfg.RunMode == config.RunModeSimple {
|
||||
return nil, nil
|
||||
}
|
||||
subscriptionBill := group != nil && group.IsSubscriptionType() && subscription != nil
|
||||
if !subscriptionBill {
|
||||
balance, err := s.GetUserBalance(ctx, user.ID)
|
||||
if err != nil {
|
||||
return nil, ErrBillingServiceUnavailable.WithCause(err)
|
||||
}
|
||||
if balance < amount {
|
||||
return nil, ErrInsufficientBalance
|
||||
}
|
||||
}
|
||||
res, err := s.reserveInflight(ctx, user, group, subscription, amount, false, false)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var subscriptionPending, apiKeyPending float64
|
||||
if subscriptionBill && (group.HasDailyLimit() || group.HasWeeklyLimit() || group.HasMonthlyLimit()) {
|
||||
res, subscriptionPending = s.reserveInflightScope(ctx, res, user.ID, inflightSubscriptionScope(subscription.ID), amount)
|
||||
}
|
||||
if apiKey != nil && (apiKey.Quota > 0 || apiKey.HasRateLimits()) {
|
||||
res, apiKeyPending = s.reserveInflightScope(ctx, res, user.ID, inflightAPIKeyScope(apiKey.ID), amount)
|
||||
}
|
||||
if subscriptionBill {
|
||||
err = s.checkFixedChargeSubscription(ctx, user.ID, group, amount+subscriptionPending)
|
||||
}
|
||||
if err == nil && apiKey != nil {
|
||||
err = s.checkFixedChargeAPIKey(ctx, apiKey, amount+apiKeyPending)
|
||||
}
|
||||
if err != nil {
|
||||
res.Release()
|
||||
return nil, err
|
||||
}
|
||||
return res, nil
|
||||
}
|
||||
|
||||
// checkFixedChargeSubscription 按库中订阅用量校验日 / 周 / 月剩余额度覆盖 pending。
|
||||
func (s *BillingCacheService) checkFixedChargeSubscription(ctx context.Context, userID int64, group *Group, pending float64) error {
|
||||
read := s.GetSubscriptionStatus
|
||||
if s.subRepo != nil {
|
||||
read = s.getSubscriptionFromDB
|
||||
}
|
||||
subData, err := read(ctx, userID, group.ID)
|
||||
if err != nil {
|
||||
return ErrBillingServiceUnavailable.WithCause(err)
|
||||
}
|
||||
if group.HasDailyLimit() && subData.DailyUsage+pending > *group.DailyLimitUSD {
|
||||
return ErrDailyLimitExceeded
|
||||
}
|
||||
if group.HasWeeklyLimit() && subData.WeeklyUsage+pending > *group.WeeklyLimitUSD {
|
||||
return ErrWeeklyLimitExceeded
|
||||
}
|
||||
if group.HasMonthlyLimit() && subData.MonthlyUsage+pending > *group.MonthlyLimitUSD {
|
||||
return ErrMonthlyLimitExceeded
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// checkFixedChargeAPIKey 按库中 Key 用量校验总额度与限速窗口覆盖 pending;未接入仓储时使用请求上下文与限速缓存中的用量。
|
||||
func (s *BillingCacheService) checkFixedChargeAPIKey(ctx context.Context, apiKey *APIKey, pending float64) error {
|
||||
if apiKey.Quota <= 0 && !apiKey.HasRateLimits() {
|
||||
return nil
|
||||
}
|
||||
if s.apiKeyRateLimitLoader == nil {
|
||||
if apiKey.Quota > 0 && apiKey.QuotaUsed+pending > apiKey.Quota {
|
||||
return ErrAPIKeyQuotaExhausted
|
||||
}
|
||||
if apiKey.HasRateLimits() {
|
||||
return s.checkAPIKeyRateLimitsWithPending(ctx, apiKey, pending)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
data, err := s.apiKeyRateLimitLoader.GetRateLimitData(ctx, apiKey.ID)
|
||||
if err != nil {
|
||||
return ErrBillingServiceUnavailable.WithCause(err)
|
||||
}
|
||||
exceeds := func(usage, limit float64) bool {
|
||||
return limit > 0 && usage+pending > limit
|
||||
}
|
||||
switch {
|
||||
case exceeds(data.QuotaUsed, apiKey.Quota):
|
||||
return ErrAPIKeyQuotaExhausted
|
||||
case exceeds(data.EffectiveUsage5h(), apiKey.RateLimit5h):
|
||||
return ErrAPIKeyRateLimit5hExceeded
|
||||
case exceeds(data.EffectiveUsage1d(), apiKey.RateLimit1d):
|
||||
return ErrAPIKeyRateLimit1dExceeded
|
||||
case exceeds(data.EffectiveUsage7d(), apiKey.RateLimit7d):
|
||||
return ErrAPIKeyRateLimit7dExceeded
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// checkSimpleModeAPIKeyRateLimits is deliberately DB-authoritative. Redis
|
||||
// updates are asynchronous and can be dropped or missed after a committed
|
||||
// transaction, so using the cache here could let a limited simple-mode key
|
||||
|
||||
@@ -3,6 +3,7 @@ package service
|
||||
import (
|
||||
"context"
|
||||
"math"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
@@ -25,6 +26,16 @@ type InflightBalanceReservationCache interface {
|
||||
ReleaseInflightBalance(ctx context.Context, userID int64, requestID string) error
|
||||
}
|
||||
|
||||
// InflightScopeReservationCache 订阅 / API Key 维度在途预留的缓存能力(可选),用于扣费前即可确定金额的费用;
|
||||
// 未实现时这两个维度不登记预留(fail-open),准入仍按库中用量校验。
|
||||
type InflightScopeReservationCache interface {
|
||||
// ReserveInflightScope 原子地清理 scope 下的过期预留并登记 requestID 的预留(ttl 后自动失效),
|
||||
// 返回登记前该 scope 下其他在途预留的合计。
|
||||
ReserveInflightScope(ctx context.Context, scope, requestID string, amount float64, ttl time.Duration) (float64, error)
|
||||
// ReleaseInflightScope 释放 scope 下 requestID 的预留(幂等)。
|
||||
ReleaseInflightScope(ctx context.Context, scope, requestID string) error
|
||||
}
|
||||
|
||||
// InflightBalanceReservationRenewer 可选:续期仍存活的预留(流式长请求期间防止 TTL 过期)。
|
||||
type InflightBalanceReservationRenewer interface {
|
||||
// RenewInflightBalance 若 requestID 仍存在则把其过期时间推迟到 now+ttl;返回是否仍存在。
|
||||
@@ -54,11 +65,15 @@ const (
|
||||
// 不再续期,所以最长持有时间受 TTL 约束(计费任务卡死时预留自动过期)。
|
||||
// 所有方法对 nil 接收者安全。
|
||||
type InflightReservation struct {
|
||||
// cache 为 nil 表示未登记余额维度的预留。
|
||||
cache InflightBalanceReservationCache
|
||||
userID int64
|
||||
requestID string
|
||||
amount float64
|
||||
ttl time.Duration
|
||||
// scopes 是同一金额在订阅 / API Key 维度登记的预留(与余额维度共用 requestID),不续期。
|
||||
scopeCache InflightScopeReservationCache
|
||||
scopes []string
|
||||
|
||||
refs atomic.Int64
|
||||
releaseOnce sync.Once
|
||||
@@ -116,8 +131,15 @@ func (r *InflightReservation) Release() {
|
||||
r.refs.Store(0)
|
||||
relCtx, relCancel := context.WithTimeout(context.Background(), inflightReservationReleaseTimeout)
|
||||
defer relCancel()
|
||||
if err := r.cache.ReleaseInflightBalance(relCtx, r.userID, r.requestID); err != nil {
|
||||
logger.LegacyPrintf("service.billing_cache", "Warning: inflight reservation release failed for user %d (expires by ttl): %v", r.userID, err)
|
||||
if r.cache != nil {
|
||||
if err := r.cache.ReleaseInflightBalance(relCtx, r.userID, r.requestID); err != nil {
|
||||
logger.LegacyPrintf("service.billing_cache", "Warning: inflight reservation release failed for user %d (expires by ttl): %v", r.userID, err)
|
||||
}
|
||||
}
|
||||
for _, scope := range r.scopes {
|
||||
if err := r.scopeCache.ReleaseInflightScope(relCtx, scope, r.requestID); err != nil {
|
||||
logger.LegacyPrintf("service.billing_cache", "Warning: inflight reservation release failed for %s (expires by ttl): %v", scope, err)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -213,7 +235,7 @@ func (s *BillingCacheService) InflightReservationFailClosedOnUnpriced() bool {
|
||||
|
||||
// ReserveInflightBalance 简化封装:返回释放函数,不续期、不做计费交接(预留最长存活 TTL)。
|
||||
func (s *BillingCacheService) ReserveInflightBalance(ctx context.Context, user *User, group *Group, subscription *UserSubscription, estimate float64) (func(), error) {
|
||||
r, err := s.reserveInflight(ctx, user, group, subscription, estimate, false)
|
||||
r, err := s.reserveInflight(ctx, user, group, subscription, estimate, false, true)
|
||||
if err != nil {
|
||||
return noopRelease, err
|
||||
}
|
||||
@@ -232,10 +254,50 @@ func (s *BillingCacheService) ReserveInflightBalance(ctx context.Context, user *
|
||||
// 开关关闭 / 简易模式 / 订阅模式 / estimate <= 0 / 缓存不支持 / 余额读取失败 / Redis 执行失败。
|
||||
// 仅当 Redis 明确判定 缓存余额 - 在途合计 < estimate(且已有在途请求)时返回 ErrInsufficientBalance。
|
||||
func (s *BillingCacheService) ReserveInflight(ctx context.Context, user *User, group *Group, subscription *UserSubscription, estimate float64) (*InflightReservation, error) {
|
||||
return s.reserveInflight(ctx, user, group, subscription, estimate, true)
|
||||
return s.reserveInflight(ctx, user, group, subscription, estimate, true, true)
|
||||
}
|
||||
|
||||
func (s *BillingCacheService) reserveInflight(ctx context.Context, user *User, group *Group, subscription *UserSubscription, estimate float64, renew bool) (*InflightReservation, error) {
|
||||
// inflightSubscriptionScope / inflightAPIKeyScope 是订阅 / API Key 维度在途预留的 scope 名。
|
||||
func inflightSubscriptionScope(subscriptionID int64) string {
|
||||
return "sub:" + strconv.FormatInt(subscriptionID, 10)
|
||||
}
|
||||
|
||||
func inflightAPIKeyScope(apiKeyID int64) string {
|
||||
return "apikey:" + strconv.FormatInt(apiKeyID, 10)
|
||||
}
|
||||
|
||||
// reserveInflightScope 在 scope 维度为 amount 登记在途预留并挂到 res 上(res 为 nil 时新建句柄),
|
||||
// 返回句柄与登记前该 scope 下其他在途预留的合计。开关关闭 / 缓存不支持 / Redis 执行失败时不登记(fail-open),合计按 0。
|
||||
func (s *BillingCacheService) reserveInflightScope(ctx context.Context, res *InflightReservation, userID int64, scope string, amount float64) (*InflightReservation, float64) {
|
||||
cfg, ok := s.inflightReservationConfig()
|
||||
if !ok {
|
||||
return res, 0
|
||||
}
|
||||
sc, ok := s.cache.(InflightScopeReservationCache)
|
||||
if !ok || sc == nil {
|
||||
return res, 0
|
||||
}
|
||||
if res == nil {
|
||||
ttl := defaultInflightReservationTTL
|
||||
if cfg.TTLSeconds > 0 {
|
||||
ttl = time.Duration(cfg.TTLSeconds) * time.Second
|
||||
}
|
||||
res = &InflightReservation{userID: userID, requestID: uuid.NewString(), amount: amount, ttl: ttl}
|
||||
res.refs.Store(1)
|
||||
}
|
||||
reserveCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), inflightReservationReserveTimeout)
|
||||
pending, err := sc.ReserveInflightScope(reserveCtx, scope, res.requestID, amount, res.ttl)
|
||||
cancel()
|
||||
if err != nil {
|
||||
logger.LegacyPrintf("service.billing_cache", "Warning: inflight reservation failed for %s (fail-open): %v", scope, err)
|
||||
return res, 0
|
||||
}
|
||||
res.scopeCache = sc
|
||||
res.scopes = append(res.scopes, scope)
|
||||
return res, pending
|
||||
}
|
||||
|
||||
func (s *BillingCacheService) reserveInflight(ctx context.Context, user *User, group *Group, subscription *UserSubscription, estimate float64, renew, capped bool) (*InflightReservation, error) {
|
||||
cfg, ok := s.inflightReservationConfig()
|
||||
if !ok || user == nil {
|
||||
return nil, nil
|
||||
@@ -243,7 +305,7 @@ func (s *BillingCacheService) reserveInflight(ctx context.Context, user *User, g
|
||||
if group != nil && group.IsSubscriptionType() && subscription != nil {
|
||||
return nil, nil
|
||||
}
|
||||
if cfg.MaxReservationUSD > 0 && estimate > cfg.MaxReservationUSD {
|
||||
if capped && cfg.MaxReservationUSD > 0 && estimate > cfg.MaxReservationUSD {
|
||||
estimate = cfg.MaxReservationUSD
|
||||
}
|
||||
if estimate <= 0 || math.IsNaN(estimate) || math.IsInf(estimate, 0) {
|
||||
|
||||
@@ -488,14 +488,15 @@ func (s *BillingService) initFallbackPricing() {
|
||||
SupportsCacheBreakdown: false,
|
||||
}
|
||||
|
||||
// Gemini 3.6 Flash (Google AI pricing: $1.50 input / $7.50 output /
|
||||
// $0.15 cached input per MTok). Antigravity's -high/-low/-medium/-tiered
|
||||
// aliases are matched below so unavailable remote pricing never records
|
||||
// token-bearing requests at $0.
|
||||
// Gemini 3.6 Flash (Google AI pricing: $0.75 input / $3.75 output /
|
||||
// $0.075 cached input per MTok, promotional through 2026-12-31; official
|
||||
// rates double to $1.50/$7.50/$0.15 from 2027-01-01). Antigravity's
|
||||
// -high/-low/-medium/-tiered aliases are matched below so unavailable
|
||||
// remote pricing never records token-bearing requests at $0.
|
||||
s.fallbackPrices["gemini-3.6-flash"] = &ModelPricing{
|
||||
InputPricePerToken: 1.5e-6,
|
||||
OutputPricePerToken: 7.5e-6,
|
||||
CacheReadPricePerToken: 0.15e-6,
|
||||
InputPricePerToken: 0.75e-6,
|
||||
OutputPricePerToken: 3.75e-6,
|
||||
CacheReadPricePerToken: 0.075e-6,
|
||||
SupportsCacheBreakdown: false,
|
||||
}
|
||||
|
||||
@@ -1342,6 +1343,18 @@ func (s *BillingService) GetModelPricing(model string) (*ModelPricing, error) {
|
||||
return s.getModelPricingAt(model, timezone.Now())
|
||||
}
|
||||
|
||||
// CacheStoragePricePerTokenHour 返回模型的显式上下文缓存存储单价(每 token 每小时,USD)。
|
||||
// 优先级:价格表(override_file 等)按确定性识别配置的值 > 内置型号价 > 内置 Pro/Flash 系列兜底价;
|
||||
// 均未命中返回 0。
|
||||
func (s *BillingService) CacheStoragePricePerTokenHour(model string) float64 {
|
||||
if s != nil && s.pricingService != nil {
|
||||
if pricing := s.pricingService.GetIdentifiedModelPricing(model); pricing != nil && pricing.CacheStorageCostPerTokenPerHour > 0 {
|
||||
return pricing.CacheStorageCostPerTokenPerHour
|
||||
}
|
||||
}
|
||||
return builtinGeminiCacheStoragePrice(model)
|
||||
}
|
||||
|
||||
// getModelPricingAt 是 GetModelPricing 的带计费时点内部变体:pricingAt 显式
|
||||
// 驱动 DeepSeek pro→Flash 切换判定(切换点前 Pro 价、之后 Flash 价),使
|
||||
// 展示/估算路径可与历史补账同刻复算,测试也能用固定时点钉住断言。
|
||||
|
||||
@@ -636,6 +636,8 @@ var monitorAPIKeyPatterns = []struct {
|
||||
{regexp.MustCompile(`xai-[A-Za-z0-9_-]{6,}`), "xai-***REDACTED***"},
|
||||
// Gemini / Google API Key:固定前缀 + 35 位
|
||||
{regexp.MustCompile(`AIza[A-Za-z0-9_-]{35}`), "AIza***REDACTED***"},
|
||||
// Gemini auth key:AQ.Ab 前缀 + 48 位,长度放宽以覆盖格式变体
|
||||
{regexp.MustCompile(`AQ\.Ab[A-Za-z0-9_-]{40,}`), "AQ.***REDACTED***"},
|
||||
// JWT 三段式(Bearer 后常出现):eyJxxx.eyJxxx.signature
|
||||
{regexp.MustCompile(`eyJ[A-Za-z0-9_-]{8,}\.eyJ[A-Za-z0-9_-]{8,}\.[A-Za-z0-9_-]{8,}`), "eyJ***REDACTED.JWT***"},
|
||||
}
|
||||
@@ -643,7 +645,7 @@ var monitorAPIKeyPatterns = []struct {
|
||||
// sanitizeErrorMessage 擦除错误/响应文本中可能泄露的 API key。
|
||||
// 处理两类来源:
|
||||
// 1. URL query 中的 ?key= / ?api_key= 等(Go *url.Error 会回填完整 URL)
|
||||
// 2. 上游 HTTP body 文本里直接出现的 sk-* / xai-* / AIza* / JWT 等密钥碎片
|
||||
// 2. 上游 HTTP body 文本里直接出现的 sk-* / xai-* / AIza* / AQ.* / JWT 等密钥碎片
|
||||
//
|
||||
// 注意:与 gemini_messages_compat_service.go 的 sanitizeUpstreamErrorMessage 关注点类似但参数集更广,
|
||||
// 监控模块独立维护,避免互相耦合。
|
||||
|
||||
@@ -279,6 +279,17 @@ func TestRunCheckForModel_Grok_RedactsXAIKeyFromUpstreamBody(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestSanitizeErrorMessage_RedactsGeminiAuthKey(t *testing.T) {
|
||||
authKey := "AQ." + "Ab" + strings.Repeat("Xy9_-z", 8)
|
||||
got := sanitizeErrorMessage(`{"error":{"message":"API key not valid: ` + authKey + `"}}`)
|
||||
if strings.Contains(got, authKey) {
|
||||
t.Fatalf("Gemini auth key leaked: %q", got)
|
||||
}
|
||||
if !strings.Contains(got, "AQ.***REDACTED***") {
|
||||
t.Fatalf("expected redaction marker, got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunCheckForModel_OpenAIResponses_DefaultRequest(t *testing.T) {
|
||||
h := &openAICaptureHandler{}
|
||||
endpoint := setupFakeOpenAI(t, h)
|
||||
|
||||
@@ -649,6 +649,9 @@ type ForwardResult struct {
|
||||
ImageSizeBreakdown map[string]int
|
||||
SearchCount int
|
||||
AudioUsage *AudioUsage
|
||||
// CacheStorageTokenHours 显式上下文缓存的存储量(缓存 token 数 × 存储小时数),按计费模型的
|
||||
// 存储单价折算为费用后叠加在 token 费用之上。
|
||||
CacheStorageTokenHours float64
|
||||
}
|
||||
|
||||
// GatewayFailureStage identifies which request stage failed. The zero value is
|
||||
@@ -696,6 +699,7 @@ type UpstreamFailoverError struct {
|
||||
SameAccountRetryMax int // 可选的错误级同账号重试上限,低于 handler 默认预算时优先采用
|
||||
RequestScopedTransient bool // 故障因素与账号无关(如上游按客户端身份/模型容量降载):可同账号重试,但不得据此对账号做临时封禁
|
||||
SafeToFailoverAfterWrite bool // 仅写出 SSE 注释等非语义字节时,仍可在同一客户端流中切换账号
|
||||
BoundUpstreamPassthrough bool // 绑定显式缓存的请求在上游网关账号上收到的 429 / 5xx:反映持有缓存的上游账号状态,原样回写不走错误映射
|
||||
Stage GatewayFailureStage
|
||||
Scope GatewayFailureScope
|
||||
Reason GatewayFailureReason
|
||||
|
||||
@@ -241,7 +241,8 @@ func isForcedUsageBillingRequestID(requestID string) bool {
|
||||
return strings.HasPrefix(id, "web_search:") ||
|
||||
strings.HasPrefix(id, "grok-video:") ||
|
||||
strings.HasPrefix(id, "grok_audio:") ||
|
||||
strings.HasPrefix(id, "grok_realtime:")
|
||||
strings.HasPrefix(id, "grok_realtime:") ||
|
||||
strings.HasPrefix(id, GeminiCachedContentUsageRequestIDPrefix)
|
||||
}
|
||||
|
||||
// StableGrokAudioBillingRequestID is the durable usage_logs / dedup key for one
|
||||
@@ -787,42 +788,9 @@ func (s *GatewayService) recordUsageCore(ctx context.Context, input *recordUsage
|
||||
cacheTTLOverridden = (result.Usage.CacheCreation5mTokens + result.Usage.CacheCreation1hTokens) > 0
|
||||
}
|
||||
|
||||
// 获取费率倍数(优先级:用户专属 > 分组默认 > 系统默认)
|
||||
multiplier := 1.0
|
||||
if s.cfg != nil {
|
||||
multiplier = s.cfg.Default.RateMultiplier
|
||||
}
|
||||
if apiKey.GroupID != nil && apiKey.Group != nil {
|
||||
groupDefault := apiKey.Group.RateMultiplier
|
||||
multiplier = s.ResolveUserGroupRateMultiplier(ctx, user.ID, *apiKey.GroupID, groupDefault)
|
||||
}
|
||||
// token 倍率叠加高峰因子(token 计费含图片 token,图片按次倍率不受影响)。高峰因子按请求时刻现算,
|
||||
// 不并入上面的 getUserGroupRateMultiplier,以免污染 user:group 倍率缓存。
|
||||
pricingAt := input.PricingAt
|
||||
if pricingAt.IsZero() {
|
||||
pricingAt = timezone.Now()
|
||||
}
|
||||
multiplier, imageMultiplier := computePeakAwareMultipliers(apiKey, multiplier, pricingAt)
|
||||
|
||||
// 确定计费模型
|
||||
concreteBillingModel := forwardResultBillingModel(result.Model, result.UpstreamModel)
|
||||
billingModel := concreteBillingModel
|
||||
if input.BillingModelSource == BillingModelSourceChannelMapped && input.ChannelMappedModel != "" {
|
||||
billingModel = input.ChannelMappedModel
|
||||
}
|
||||
if input.BillingModelSource == BillingModelSourceRequested && input.OriginalModel != "" {
|
||||
billingModel = input.OriginalModel
|
||||
}
|
||||
// composite 分组的公开别名(如 all/claude)会经 OriginalModel/ChannelMappedModel
|
||||
// 进入上面的来源覆盖:任意别名查无价会静默落 $0,含家族词的别名则被价格表的
|
||||
// 家族模糊匹配错计(如 Opus 流量按 Sonnet 兜底价)。除非管理员为别名显式配置了
|
||||
// 渠道定价(OpenRouter 式自定价),composite 请求一律按实际转发的具体模型计费。
|
||||
if apiKey.Group != nil && apiKey.Group.Platform == PlatformComposite {
|
||||
billingModel = s.compositeBillableModel(ctx, apiKey, billingModel, concreteBillingModel)
|
||||
}
|
||||
// 通用兜底(与 OpenAI 路径的 usageBillingModelCandidates 语义对齐):
|
||||
// 选定模型查不到任何价格时回退到实际转发的具体模型。已定价流量不受影响。
|
||||
billingModel = s.billableModelWithFallback(ctx, apiKey, billingModel, result.UpstreamModel, result.Model)
|
||||
multiplier, imageMultiplier, pricingAt := s.recordUsageMultipliers(ctx, apiKey, user, input.PricingAt)
|
||||
billingModel := s.recordUsageBillingModel(ctx, apiKey, result, input.ChannelUsageFields)
|
||||
cacheStorageCost := s.cacheStorageCost(result, billingModel)
|
||||
|
||||
// 确定 RequestedModel(渠道映射前的原始模型)
|
||||
requestedModel := result.Model
|
||||
@@ -840,7 +808,7 @@ func (s *GatewayService) recordUsageCore(ctx context.Context, input *recordUsage
|
||||
input.BillingModelSource,
|
||||
result.UpstreamResponseModel,
|
||||
result.UpstreamResponseModelConflict,
|
||||
result.ImageCount > 0 || result.AudioUsage != nil || result.SearchCount > 0,
|
||||
result.ImageCount > 0 || result.AudioUsage != nil || result.SearchCount > 0 || result.CacheStorageTokenHours > 0,
|
||||
); responseModel != "" && !strings.EqualFold(responseModel, strings.TrimSpace(billingModel)) {
|
||||
if identified, responseChannelPriced := s.hasIdentifiedResponseModelPricing(ctx, responseModel, apiKey); identified {
|
||||
responseCost := s.calculateRecordUsageCost(ctx, result, apiKey, responseModel, multiplier, imageMultiplier, pricingAt)
|
||||
@@ -866,8 +834,9 @@ func (s *GatewayService) recordUsageCore(ctx context.Context, input *recordUsage
|
||||
usageLog := s.buildRecordUsageLog(ctx, input, result, apiKey, user, account, subscription,
|
||||
requestedModel, multiplier, imageMultiplier, accountRateMultiplier, billingType, cacheTTLOverridden, cost)
|
||||
|
||||
// 计算账号统计定价费用(使用最终上游模型匹配自定义规则)
|
||||
if apiKey.GroupID != nil {
|
||||
// 计算账号统计定价费用(使用最终上游模型匹配自定义规则)。纯存储费不是一次模型请求,
|
||||
// 账号统计成本回落 total_cost(官方存储价)。
|
||||
if apiKey.GroupID != nil && !isCacheStorageOnlyResult(result) {
|
||||
applyAccountStatsCost(ctx, usageLog, s.channelService, s.billingService,
|
||||
account.ID, *apiKey.GroupID, result.UpstreamModel, result.Model,
|
||||
// Anthropic's input_tokens excludes cache_read and cache_creation (billed separately);
|
||||
@@ -879,9 +848,10 @@ func (s *GatewayService) recordUsageCore(ctx context.Context, input *recordUsage
|
||||
CacheReadTokens: result.Usage.CacheReadInputTokens,
|
||||
ImageOutputTokens: result.Usage.ImageOutputTokens,
|
||||
},
|
||||
cost.TotalCost, pricingAt,
|
||||
cost.TotalCost-cacheStorageCost, pricingAt,
|
||||
accountStatsLongContextPricingEnabled(nil),
|
||||
)
|
||||
addCacheStorageToAccountStatsCost(usageLog, cacheStorageCost)
|
||||
}
|
||||
|
||||
simpleModeKeyRateLimitOnly := simpleModeKeyRateLimitBillingEnabled(s.cfg, apiKey)
|
||||
@@ -927,6 +897,112 @@ func (s *GatewayService) recordUsageCore(ctx context.Context, input *recordUsage
|
||||
return nil
|
||||
}
|
||||
|
||||
// recordUsageMultipliers 返回 token / 图片售价倍率(用户专属 > 分组默认 > 系统默认,token 倍率叠加高峰因子)
|
||||
// 与计价时刻;pricingAt 为零值时取当前时刻。
|
||||
func (s *GatewayService) recordUsageMultipliers(ctx context.Context, apiKey *APIKey, user *User, pricingAt time.Time) (float64, float64, time.Time) {
|
||||
multiplier := 1.0
|
||||
if s.cfg != nil {
|
||||
multiplier = s.cfg.Default.RateMultiplier
|
||||
}
|
||||
if apiKey.GroupID != nil && apiKey.Group != nil {
|
||||
multiplier = apiKey.Group.RateMultiplier
|
||||
if user != nil {
|
||||
multiplier = s.ResolveUserGroupRateMultiplier(ctx, user.ID, *apiKey.GroupID, multiplier)
|
||||
}
|
||||
}
|
||||
// 高峰因子按请求时刻现算,不并入 getUserGroupRateMultiplier,以免污染 user:group 倍率缓存。
|
||||
if pricingAt.IsZero() {
|
||||
pricingAt = timezone.Now()
|
||||
}
|
||||
multiplier, imageMultiplier := computePeakAwareMultipliers(apiKey, multiplier, pricingAt)
|
||||
return multiplier, imageMultiplier, pricingAt
|
||||
}
|
||||
|
||||
// recordUsageBillingModel 按计费模型来源、composite 规则与无价兜底确定计费模型。
|
||||
func (s *GatewayService) recordUsageBillingModel(ctx context.Context, apiKey *APIKey, result *ForwardResult, fields ChannelUsageFields) string {
|
||||
concreteBillingModel := forwardResultBillingModel(result.Model, result.UpstreamModel)
|
||||
billingModel := concreteBillingModel
|
||||
if fields.BillingModelSource == BillingModelSourceChannelMapped && fields.ChannelMappedModel != "" {
|
||||
billingModel = fields.ChannelMappedModel
|
||||
}
|
||||
if fields.BillingModelSource == BillingModelSourceRequested && fields.OriginalModel != "" {
|
||||
billingModel = fields.OriginalModel
|
||||
}
|
||||
// composite 分组的公开别名(如 all/claude)会经 OriginalModel/ChannelMappedModel
|
||||
// 进入上面的来源覆盖:任意别名查无价会静默落 $0,含家族词的别名则被价格表的
|
||||
// 家族模糊匹配错计(如 Opus 流量按 Sonnet 兜底价)。除非管理员为别名显式配置了
|
||||
// 渠道定价(OpenRouter 式自定价),composite 请求一律按实际转发的具体模型计费。
|
||||
if apiKey.Group != nil && apiKey.Group.Platform == PlatformComposite {
|
||||
billingModel = s.compositeBillableModel(ctx, apiKey, billingModel, concreteBillingModel)
|
||||
}
|
||||
// 通用兜底(与 OpenAI 路径的 usageBillingModelCandidates 语义对齐):
|
||||
// 选定模型查不到任何价格时回退到实际转发的具体模型。已定价流量不受影响。
|
||||
return s.billableModelWithFallback(ctx, apiKey, billingModel, result.UpstreamModel, result.Model)
|
||||
}
|
||||
|
||||
// cacheStoragePrice 返回计费模型的显式缓存存储单价;计费模型未配置时依次回退到实际转发模型与请求模型。
|
||||
func (s *GatewayService) cacheStoragePrice(result *ForwardResult, billingModel string) float64 {
|
||||
if s.billingService == nil {
|
||||
return 0
|
||||
}
|
||||
for _, model := range []string{billingModel, result.UpstreamModel, result.Model} {
|
||||
if strings.TrimSpace(model) == "" {
|
||||
continue
|
||||
}
|
||||
if price := s.billingService.CacheStoragePricePerTokenHour(model); price > 0 {
|
||||
return price
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// cacheStorageCost 返回 result 的显式缓存存储费(USD,倍率前)。
|
||||
func (s *GatewayService) cacheStorageCost(result *ForwardResult, billingModel string) float64 {
|
||||
if result.CacheStorageTokenHours <= 0 {
|
||||
return 0
|
||||
}
|
||||
return result.CacheStorageTokenHours * s.cacheStoragePrice(result, billingModel)
|
||||
}
|
||||
|
||||
// isCacheStorageOnlyResult 判断 result 是否只有显式缓存存储费(延长有效期):不含 token 用量,不算一次模型请求。
|
||||
func isCacheStorageOnlyResult(result *ForwardResult) bool {
|
||||
u := result.Usage
|
||||
return result.CacheStorageTokenHours > 0 && u.InputTokens == 0 && u.OutputTokens == 0 &&
|
||||
u.CacheCreationInputTokens == 0 && u.CacheReadInputTokens == 0 && u.ImageOutputTokens == 0 &&
|
||||
result.ImageCount == 0 && result.AudioUsage == nil && result.SearchCount == 0
|
||||
}
|
||||
|
||||
// CacheStoragePriced 判断 RecordUsage 对 input.Result 采用的计费模型(含回退候选)是否配置了显式缓存存储单价。
|
||||
func (s *GatewayService) CacheStoragePriced(ctx context.Context, input *RecordUsageInput) bool {
|
||||
billingModel := s.recordUsageBillingModel(ctx, input.APIKey, input.Result, input.ChannelUsageFields)
|
||||
return s.cacheStoragePrice(input.Result, billingModel) > 0
|
||||
}
|
||||
|
||||
// QuoteUsageCost 按与 RecordUsage 相同的倍率、计费模型与定价计算 input.Result 的应付金额(倍率后 USD),不记录也不扣费。
|
||||
func (s *GatewayService) QuoteUsageCost(ctx context.Context, input *RecordUsageInput) float64 {
|
||||
user := input.User
|
||||
if user == nil {
|
||||
user = input.APIKey.User
|
||||
}
|
||||
multiplier, imageMultiplier, pricingAt := s.recordUsageMultipliers(ctx, input.APIKey, user, input.PricingAt)
|
||||
billingModel := s.recordUsageBillingModel(ctx, input.APIKey, input.Result, input.ChannelUsageFields)
|
||||
cost := s.calculateRecordUsageCost(ctx, input.Result, input.APIKey, billingModel, multiplier, imageMultiplier, pricingAt)
|
||||
if cost == nil {
|
||||
return 0
|
||||
}
|
||||
return cost.ActualCost
|
||||
}
|
||||
|
||||
// addCacheStorageToAccountStatsCost 把显式缓存存储费按官方价叠加到账号统计成本。
|
||||
// 账号统计定价只覆盖 token 部分;统计成本为空时成本统计回落 total_cost,其中已含存储费。
|
||||
func addCacheStorageToAccountStatsCost(usageLog *UsageLog, storageCost float64) {
|
||||
if usageLog == nil || usageLog.AccountStatsCost == nil || storageCost <= 0 {
|
||||
return
|
||||
}
|
||||
cost := *usageLog.AccountStatsCost + storageCost
|
||||
usageLog.AccountStatsCost = &cost
|
||||
}
|
||||
|
||||
// calculateRecordUsageCost 根据请求类型计算费用。
|
||||
func (s *GatewayService) calculateRecordUsageCost(
|
||||
ctx context.Context,
|
||||
@@ -964,8 +1040,13 @@ func (s *GatewayService) calculateRecordUsageCost(
|
||||
return s.billingService.CalculateAudioCost(result.AudioUsage.Mode, result.AudioUsage.DurationOrUnits, cfg, multiplier)
|
||||
}
|
||||
|
||||
// Token 计费;SearchCount 为叠加 surcharge(不替代 token)。
|
||||
tokenCost := s.calculateTokenCost(ctx, result, apiKey, billingModel, multiplier, pricingAt)
|
||||
// Token 计费;SearchCount 为叠加 surcharge(不替代 token)。纯存储费不是一次模型请求,不计 token / 按次费用。
|
||||
var tokenCost *CostBreakdown
|
||||
if isCacheStorageOnlyResult(result) {
|
||||
tokenCost = &CostBreakdown{}
|
||||
} else {
|
||||
tokenCost = s.calculateTokenCost(ctx, result, apiKey, billingModel, multiplier, pricingAt)
|
||||
}
|
||||
if result.SearchCount > 0 {
|
||||
price := groupSearchPricePer1kFromAPIKey(apiKey)
|
||||
if price != nil && *price == 0 {
|
||||
@@ -980,6 +1061,14 @@ func (s *GatewayService) calculateRecordUsageCost(
|
||||
tokenCost.ActualCost += searchCost.ActualCost
|
||||
}
|
||||
}
|
||||
// 显式缓存存储费同为叠加 surcharge,与 token 费用使用同一计费模型与倍率。
|
||||
if storageCost := s.cacheStorageCost(result, billingModel); storageCost > 0 {
|
||||
if tokenCost == nil {
|
||||
tokenCost = &CostBreakdown{}
|
||||
}
|
||||
tokenCost.TotalCost += storageCost
|
||||
tokenCost.ActualCost += storageCost * multiplier
|
||||
}
|
||||
return tokenCost
|
||||
}
|
||||
|
||||
|
||||
@@ -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"))
|
||||
|
||||
@@ -0,0 +1,72 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"slices"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Gemini 显式上下文缓存存储的内置官方单价(USD / token / 小时,Standard 档)。
|
||||
// 存储价与上下文长度无关(≤200K 与 >200K 同价);Priority 档的更高存储价不在此列。
|
||||
// 来源:https://ai.google.dev/gemini-api/docs/pricing
|
||||
// 与 https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing
|
||||
const (
|
||||
geminiProCacheStoragePrice = 4.5e-6 // $4.50 / 1M tokens / 小时
|
||||
geminiFlashCacheStoragePrice = 1e-6 // $1.00 / 1M tokens / 小时,Flash-Lite 同价
|
||||
|
||||
// 3.6/3.7/3.8 Flash 的优惠价 $0.50 / 1M tokens / 小时,有效期至 2026-12-31,
|
||||
// 2027-01-01 起恢复为 geminiFlashCacheStoragePrice。AI Studio 价目页直接标注该优惠价;
|
||||
// Vertex 价目页存储价标 $1.00,优惠以「按该模型净消费返还 50% 抵扣金」发放,
|
||||
// 官方未明确存储费是否在返还范围内,两类账号统一按优惠价计。
|
||||
geminiFlashPromoCacheStoragePrice = 0.5e-6
|
||||
)
|
||||
|
||||
// geminiCacheStorageModelPrices 是已发布型号的存储单价。模型名等于前缀,或以「前缀-」开头
|
||||
// (-preview、-customtools、-001 等变体)即命中;更具体的前缀排在前面。
|
||||
var geminiCacheStorageModelPrices = []struct {
|
||||
prefix string
|
||||
price float64
|
||||
}{
|
||||
{"gemini-3.8-flash-cyber", geminiFlashCacheStoragePrice},
|
||||
{"gemini-3.8-flash", geminiFlashPromoCacheStoragePrice},
|
||||
{"gemini-3.7-flash", geminiFlashPromoCacheStoragePrice},
|
||||
{"gemini-3.6-flash", geminiFlashPromoCacheStoragePrice},
|
||||
{"gemini-3.5-flash-lite", geminiFlashCacheStoragePrice},
|
||||
{"gemini-3.5-flash", geminiFlashCacheStoragePrice},
|
||||
{"gemini-3.1-flash-lite", geminiFlashCacheStoragePrice},
|
||||
{"gemini-3.1-pro", geminiProCacheStoragePrice},
|
||||
{"gemini-3-flash", geminiFlashCacheStoragePrice},
|
||||
{"gemini-3-pro", geminiProCacheStoragePrice},
|
||||
{"gemini-2.5-flash-lite", geminiFlashCacheStoragePrice},
|
||||
{"gemini-2.5-flash", geminiFlashCacheStoragePrice},
|
||||
{"gemini-2.5-pro", geminiProCacheStoragePrice},
|
||||
{"gemini-2.0-flash-lite", geminiFlashCacheStoragePrice},
|
||||
{"gemini-2.0-flash", geminiFlashCacheStoragePrice},
|
||||
}
|
||||
|
||||
// builtinGeminiCacheStoragePrice 返回内置的显式缓存存储单价:先按已发布型号匹配,未命中时按系列兜底,
|
||||
// Pro 系列取 geminiProCacheStoragePrice,Flash 与 Flash-Lite 系列取 geminiFlashCacheStoragePrice(标准价,
|
||||
// 不套用优惠价)。其他 Gemini 模型与非 Gemini 模型返回 0。
|
||||
func builtinGeminiCacheStoragePrice(model string) float64 {
|
||||
name := lastSegment(strings.ToLower(strings.TrimSpace(model)))
|
||||
if !strings.HasPrefix(name, "gemini-") {
|
||||
return 0
|
||||
}
|
||||
parts := strings.Split(name, "-")
|
||||
lite := slices.Contains(parts, "lite")
|
||||
for _, entry := range geminiCacheStorageModelPrices {
|
||||
// Flash-Lite 变体只匹配 Flash-Lite 条目:gemini-3.8-flash-lite 不能落到 3.8 Flash 的优惠价。
|
||||
if lite != strings.HasSuffix(entry.prefix, "-lite") {
|
||||
continue
|
||||
}
|
||||
if name == entry.prefix || strings.HasPrefix(name, entry.prefix+"-") {
|
||||
return entry.price
|
||||
}
|
||||
}
|
||||
switch {
|
||||
case slices.Contains(parts, "pro"):
|
||||
return geminiProCacheStoragePrice
|
||||
case slices.Contains(parts, "flash"):
|
||||
return geminiFlashCacheStoragePrice
|
||||
}
|
||||
return 0
|
||||
}
|
||||
@@ -0,0 +1,77 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestBuiltinGeminiCacheStoragePrice(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
model string
|
||||
want float64
|
||||
}{
|
||||
// 已发布型号
|
||||
{"gemini-3.8-flash", 0.5e-6},
|
||||
{"gemini-3.7-flash", 0.5e-6},
|
||||
{"gemini-3.6-flash", 0.5e-6},
|
||||
{"gemini-3.8-flash-cyber", 1e-6},
|
||||
{"gemini-3.5-flash", 1e-6},
|
||||
{"gemini-3.5-flash-lite", 1e-6},
|
||||
{"gemini-3.1-flash-lite", 1e-6},
|
||||
{"gemini-3.1-pro-preview", 4.5e-6},
|
||||
{"gemini-3.1-pro-preview-customtools", 4.5e-6},
|
||||
{"gemini-3-flash-preview", 1e-6},
|
||||
{"gemini-3-pro-preview", 4.5e-6},
|
||||
{"gemini-2.5-pro", 4.5e-6},
|
||||
{"gemini-2.5-flash", 1e-6},
|
||||
{"gemini-2.5-flash-lite", 1e-6},
|
||||
{"gemini-2.0-flash-001", 1e-6},
|
||||
{"gemini-2.0-flash-lite-001", 1e-6},
|
||||
// 资源名与大小写
|
||||
{"models/gemini-3.8-flash", 0.5e-6},
|
||||
{"projects/p/locations/global/publishers/google/models/gemini-2.5-pro", 4.5e-6},
|
||||
{" Gemini-3.8-Flash ", 0.5e-6},
|
||||
// 系列兜底:未发布型号与别名按标准价
|
||||
{"gemini-3.9-flash", 1e-6},
|
||||
{"gemini-3.9-flash-lite", 1e-6},
|
||||
{"gemini-3.8-flash-lite", 1e-6},
|
||||
{"gemini-3.9-pro-preview", 4.5e-6},
|
||||
{"gemini-flash-latest", 1e-6},
|
||||
{"gemini-flash-lite-latest", 1e-6},
|
||||
{"gemini-pro-latest", 4.5e-6},
|
||||
// 无存储价
|
||||
{"gemini-embedding-001", 0},
|
||||
{"gemini-unknown-model", 0},
|
||||
{"claude-sonnet-5-5", 0},
|
||||
{"", 0},
|
||||
} {
|
||||
require.InDelta(t, tc.want, builtinGeminiCacheStoragePrice(tc.model), 1e-15, tc.model)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCacheStoragePricePerTokenHourPriority(t *testing.T) {
|
||||
pricing := &PricingService{}
|
||||
data, err := pricing.parsePricingData([]byte(`{
|
||||
"gemini-3.8-flash":{"litellm_provider":"gemini","input_cost_per_token":0.00000075,"output_cost_per_token":0.00000375,
|
||||
"cache_storage_cost_per_token_per_hour":0.000002},
|
||||
"gemini-3.9-flash":{"litellm_provider":"gemini","input_cost_per_token":0.0000015,"output_cost_per_token":0.0000075,
|
||||
"cache_storage_cost_per_token_per_hour":0.0000007},
|
||||
"gemini-2.5-pro":{"litellm_provider":"gemini","input_cost_per_token":0.00000125,"output_cost_per_token":0.00001}}`))
|
||||
require.NoError(t, err)
|
||||
pricing.pricingData = data
|
||||
billing := NewBillingService(&config.Config{}, pricing)
|
||||
|
||||
require.InDelta(t, 2e-6, billing.CacheStoragePricePerTokenHour("gemini-3.8-flash"), 1e-15, "价格表配置值优先于内置型号价")
|
||||
require.InDelta(t, 0.7e-6, billing.CacheStoragePricePerTokenHour("gemini-3.9-flash"), 1e-15, "价格表配置值优先于内置系列兜底价")
|
||||
require.InDelta(t, 4.5e-6, billing.CacheStoragePricePerTokenHour("gemini-2.5-pro"), 1e-15, "价格表条目缺存储价时取内置型号价")
|
||||
require.InDelta(t, 0.5e-6, billing.CacheStoragePricePerTokenHour("gemini-3.7-flash"), 1e-15, "价格表无条目时取内置型号价")
|
||||
require.InDelta(t, 4.5e-6, billing.CacheStoragePricePerTokenHour("gemini-3.9-pro"), 1e-15, "内置型号未命中时取系列兜底价")
|
||||
require.Zero(t, billing.CacheStoragePricePerTokenHour("customer-alias"))
|
||||
|
||||
noCatalog := NewBillingService(&config.Config{}, nil)
|
||||
require.InDelta(t, 0.5e-6, noCatalog.CacheStoragePricePerTokenHour("gemini-3.8-flash"), 1e-15, "无价格服务时取内置价")
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,480 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
const geminiCacheHardeningPricing = `{
|
||||
"gemini-3.8-flash":{"litellm_provider":"gemini","input_cost_per_token":0.00000075,"output_cost_per_token":0.00000375,
|
||||
"cache_storage_cost_per_token_per_hour":0.0000005},
|
||||
"gemini-2.5-pro":{"litellm_provider":"gemini","input_cost_per_token":0.00000125,"output_cost_per_token":0.00001,
|
||||
"cache_storage_cost_per_token_per_hour":0.0000045}}`
|
||||
|
||||
func newGeminiCacheBillingGatewayForTest(t *testing.T) (*GatewayService, *openAIRecordUsageLogRepoStub) {
|
||||
t.Helper()
|
||||
usageRepo := &openAIRecordUsageLogRepoStub{inserted: true}
|
||||
svc := newGatewayRecordUsageServiceForTest(usageRepo, &openAIRecordUsageUserRepoStub{}, &openAIRecordUsageSubRepoStub{})
|
||||
pricing := &PricingService{}
|
||||
data, err := pricing.parsePricingData([]byte(geminiCacheHardeningPricing))
|
||||
require.NoError(t, err)
|
||||
pricing.pricingData = data
|
||||
svc.billingService = NewBillingService(svc.cfg, pricing)
|
||||
return svc, usageRepo
|
||||
}
|
||||
|
||||
func TestGeminiCachedContentStoragePriceFollowsBillingModel(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
svc, usageRepo := newGeminiCacheBillingGatewayForTest(t)
|
||||
groupID := int64(33)
|
||||
apiKey := &APIKey{ID: 821, GroupID: &groupID, User: &User{ID: 621},
|
||||
Group: &Group{ID: groupID, Platform: PlatformGemini, RateMultiplier: 2, Hydrated: true}}
|
||||
tokenHours := GeminiCachedContentStorageTokenHours(1000, 24*time.Hour)
|
||||
|
||||
mapped := &RecordUsageInput{
|
||||
Result: &ForwardResult{RequestID: "cc-mapped", Model: "gemini-2.5-pro", UpstreamModel: "gemini-3.8-flash",
|
||||
Usage: ClaudeUsage{InputTokens: 1000}, CacheStorageTokenHours: tokenHours},
|
||||
APIKey: apiKey, User: apiKey.User, Account: &Account{ID: 713, Platform: PlatformGemini},
|
||||
}
|
||||
require.True(t, svc.CacheStoragePriced(ctx, mapped))
|
||||
want := (1000*0.00000125 + tokenHours*0.0000045) * 2
|
||||
require.InDelta(t, want, svc.QuoteUsageCost(ctx, mapped), 1e-12, "存储与 token 同按计费模型(请求模型)定价")
|
||||
require.NoError(t, svc.RecordUsage(ctx, mapped))
|
||||
require.InDelta(t, want, usageRepo.lastLog.ActualCost, 1e-12, "报价与实际计费同口径")
|
||||
|
||||
alias := &RecordUsageInput{
|
||||
Result: &ForwardResult{Model: "customer-alias", UpstreamModel: "gemini-3.8-flash", CacheStorageTokenHours: tokenHours},
|
||||
APIKey: apiKey,
|
||||
}
|
||||
require.True(t, svc.CacheStoragePriced(ctx, alias))
|
||||
require.InDelta(t, tokenHours*0.0000005*2, svc.QuoteUsageCost(ctx, alias), 1e-12, "计费模型无存储价时回退到实际转发模型")
|
||||
|
||||
unpriced := &RecordUsageInput{
|
||||
Result: &ForwardResult{Model: "customer-alias", UpstreamModel: "unknown-upstream", CacheStorageTokenHours: tokenHours},
|
||||
APIKey: apiKey,
|
||||
}
|
||||
require.False(t, svc.CacheStoragePriced(ctx, unpriced))
|
||||
}
|
||||
|
||||
func TestGeminiCachedContentStorageOnlyIsNotARequest(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
svc, usageRepo := newGeminiCacheBillingGatewayForTest(t)
|
||||
groupID := int64(34)
|
||||
perRequest := 0.05
|
||||
pricing := []ChannelModelPricing{{Models: []string{"gemini-3.8-flash"}, BillingMode: BillingModePerRequest, PerRequestPrice: &perRequest}}
|
||||
svc.channelService = newTestChannelServiceForStats(t, &Channel{ID: 1, Status: StatusActive,
|
||||
AccountStatsPricingRules: []AccountStatsPricingRule{{GroupIDs: []int64{groupID}, Pricing: pricing}}}, groupID, PlatformGemini)
|
||||
svc.resolver = NewModelPricingResolver(svc.channelService, svc.billingService)
|
||||
apiKey := &APIKey{ID: 822, GroupID: &groupID, User: &User{ID: 622},
|
||||
Group: &Group{ID: groupID, Platform: PlatformGemini, RateMultiplier: 1, Hydrated: true, ModelPricing: pricing}}
|
||||
tokenHours := 20000.0
|
||||
storage := tokenHours * 0.0000005
|
||||
|
||||
record := func(tokens int) *UsageLog {
|
||||
require.NoError(t, svc.RecordUsage(ctx, &RecordUsageInput{
|
||||
Result: &ForwardResult{RequestID: fmt.Sprintf("cc-%d", tokens), Model: "gemini-3.8-flash",
|
||||
Usage: ClaudeUsage{InputTokens: tokens}, CacheStorageTokenHours: tokenHours},
|
||||
APIKey: apiKey, User: apiKey.User, Account: &Account{ID: 712, Platform: PlatformGemini},
|
||||
}))
|
||||
return usageRepo.lastLog
|
||||
}
|
||||
|
||||
created := record(10000)
|
||||
require.InDelta(t, perRequest+storage, created.TotalCost, 1e-12, "创建是一次请求:按次价 + 存储费")
|
||||
require.NotNil(t, created.AccountStatsCost)
|
||||
require.InDelta(t, perRequest+storage, *created.AccountStatsCost, 1e-12)
|
||||
|
||||
extended := record(0)
|
||||
require.InDelta(t, storage, extended.TotalCost, 1e-12, "延长有效期不算一次请求,只收存储费")
|
||||
require.Nil(t, extended.AccountStatsCost, "账号统计成本回落 total_cost,不套按次规则")
|
||||
}
|
||||
|
||||
func TestFindGeminiCachedContentReferenceRejectsAmbiguousKeys(t *testing.T) {
|
||||
for name, body := range map[string]string{
|
||||
"duplicate top-level": `{"cachedContent":"cachedContents/a","contents":[],"cachedContent":"cachedContents/b"}`,
|
||||
"snake top-level": `{"cached_content":"cachedContents/a"}`,
|
||||
"snake beside camel": `{"cachedContent":"cachedContents/a","cached_content":"cachedContents/b"}`,
|
||||
"duplicate nested": `{"generateContentRequest":{"cachedContent":"cachedContents/a","cachedContent":"cachedContents/b"}}`,
|
||||
"snake nested": `{"generateContentRequest":{"cached_content":"cachedContents/a"}}`,
|
||||
"snake parent with cache": `{"generate_content_request":{"cachedContent":"cachedContents/a"}}`,
|
||||
"duplicate nested parent": `{"generateContentRequest":{},"generateContentRequest":{"cachedContent":"cachedContents/a"}}`,
|
||||
"conflicting two locations": `{"cachedContent":"cachedContents/a","generateContentRequest":{"cachedContent":"cachedContents/b"}}`,
|
||||
} {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
_, err := FindGeminiCachedContentReference([]byte(body))
|
||||
require.Error(t, err)
|
||||
})
|
||||
}
|
||||
|
||||
for name, body := range map[string]string{
|
||||
"snake parent without cache": `{"generate_content_request":{"contents":[]}}`,
|
||||
"key name inside content": `{"contents":[{"role":"user","cached_content":"x","parts":[{"text":"cached_content"}]}]}`,
|
||||
"not an object": `[{"cachedContent":"cachedContents/a"}]`,
|
||||
} {
|
||||
t.Run("allows "+name, func(t *testing.T) {
|
||||
ref, err := FindGeminiCachedContentReference([]byte(body))
|
||||
require.NoError(t, err)
|
||||
require.Nil(t, ref)
|
||||
})
|
||||
}
|
||||
|
||||
ref, err := FindGeminiCachedContentReference([]byte(`{"cachedContent":"cachedContents/a","generateContentRequest":{"cachedContent":"cachedContents/a"}}`))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []string{"cachedContent", "generateContentRequest.cachedContent"}, ref.Paths)
|
||||
}
|
||||
|
||||
func TestGeminiCachedContentPayloadNormalization(t *testing.T) {
|
||||
now := time.Date(2026, 10, 9, 8, 0, 0, 0, time.UTC)
|
||||
maxTTL := 24 * time.Hour
|
||||
|
||||
req, err := ParseGeminiCachedContentCreateRequest([]byte(`{"model":"gemini-3.8-flash","display_name":"docs","expire_time":"2026-10-09T10:00:00Z"}`), now, maxTTL)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "docs", req.DisplayName)
|
||||
require.Equal(t, 2*time.Hour, req.TTL)
|
||||
for _, key := range []string{"display_name", "expire_time", "displayName", "expireTime"} {
|
||||
require.NotContains(t, req.payload, key, "元数据字段不随内容发往上游")
|
||||
}
|
||||
|
||||
for name, body := range map[string]string{
|
||||
"both spellings": `{"model":"gemini-3.8-flash","displayName":"a","display_name":"b"}`,
|
||||
"trailing data": `{"model":"gemini-3.8-flash"} {"model":"x"}`,
|
||||
} {
|
||||
t.Run("rejects "+name, func(t *testing.T) {
|
||||
_, err := ParseGeminiCachedContentCreateRequest([]byte(body), now, maxTTL)
|
||||
require.Error(t, err)
|
||||
})
|
||||
}
|
||||
|
||||
ttl, err := ParseGeminiCachedContentPatchRequest([]byte(`{"expire_time":"2026-10-09T09:00:00Z"}`), "", now, maxTTL)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, time.Hour, ttl)
|
||||
_, err = ParseGeminiCachedContentPatchRequest([]byte(`{"display_name":"x"}`), "", now, maxTTL)
|
||||
require.ErrorContains(t, err, "immutable")
|
||||
}
|
||||
|
||||
func TestGeminiCachedContentTTLRoundsUpToWholeSeconds(t *testing.T) {
|
||||
now := time.Date(2026, 10, 9, 8, 0, 0, 0, time.UTC)
|
||||
req, err := ParseGeminiCachedContentCreateRequest([]byte(`{"model":"gemini-3.8-flash","ttl":"1e-10s"}`), now, 24*time.Hour)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, time.Second, req.TTL, "极小的显式有效期按 1 秒处理,不回落到默认 1 小时")
|
||||
require.True(t, req.ttlExplicit)
|
||||
|
||||
ttl, err := ParseGeminiCachedContentPatchRequest([]byte(`{"ttl":"1e-10s"}`), "", now, 24*time.Hour)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, time.Second, ttl)
|
||||
|
||||
req, err = ParseGeminiCachedContentCreateRequest([]byte(`{"model":"gemini-3.8-flash","expireTime":"2026-10-09T08:00:10.2Z"}`), now, 24*time.Hour)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 11*time.Second, req.TTL, "expireTime 折算的有效期与发往上游的整秒 ttl 一致")
|
||||
}
|
||||
|
||||
func TestGeminiCachedContentCreateKeepsNumbersVerbatim(t *testing.T) {
|
||||
svc, httpStub := newGeminiCachedContentServiceForTest(t, http.StatusOK,
|
||||
`{"name":"cachedContents/up1","usageMetadata":{"totalTokenCount":2000}}`)
|
||||
req, err := ParseGeminiCachedContentCreateRequest([]byte(`{"model":"models/gemini-3.8-flash","contents":[{"role":"user","parts":[
|
||||
{"functionResponse":{"name":"lookup","response":{"order_id":1790000000000000001,"ratio":0.10}}}]}]}`), time.Now(), 24*time.Hour)
|
||||
require.NoError(t, err)
|
||||
account := &Account{ID: 9, Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "secret"}}
|
||||
_, err = svc.CreateUpstream(context.Background(), account, req, "gemini-3.8-flash", "pub1", time.Hour)
|
||||
require.NoError(t, err)
|
||||
raw, err := io.ReadAll(httpStub.lastReq.Body)
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, string(raw), `"order_id":1790000000000000001`)
|
||||
require.Contains(t, string(raw), `"ratio":0.10`)
|
||||
}
|
||||
|
||||
func TestGeminiCachedContentCreateUpstreamUnusableResponse(t *testing.T) {
|
||||
account := &Account{ID: 9, Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "secret"}}
|
||||
for name, tc := range map[string]struct {
|
||||
body string
|
||||
wantName string
|
||||
}{
|
||||
"missing token count": {`{"name":"cachedContents/up9","expireTime":"2099-01-01T00:00:00Z"}`, "cachedContents/up9"},
|
||||
"zero token count": {`{"name":"cachedContents/up9","usageMetadata":{"totalTokenCount":0}}`, "cachedContents/up9"},
|
||||
"unexpected name": {`{"name":"cachedContents/a/../b","usageMetadata":{"totalTokenCount":5}}`, ""},
|
||||
"not json": {`<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.True(t, failoverErr.BoundUpstreamPassthrough, "标记为上游网关账号的响应,handler 原样回写")
|
||||
require.False(t, failoverErr.RetryableOnSameAccount)
|
||||
require.Equal(t, 1, httpStub.calls, "不在同一账号上退避重试")
|
||||
require.Zero(t, repo.rateLimitedCalls, "不限流整个网关账号")
|
||||
require.Zero(t, repo.tempCalls)
|
||||
require.Zero(t, repo.setErrCalls)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
type fixedChargeCacheStub struct {
|
||||
billingCacheWorkerStub
|
||||
balance float64
|
||||
sub *SubscriptionCacheData
|
||||
rate *APIKeyRateLimitCacheData
|
||||
// reserved 按维度(user:{id} / scope 名)记录在途预留。
|
||||
reserved map[string]map[string]float64
|
||||
}
|
||||
|
||||
func (s *fixedChargeCacheStub) GetUserBalance(context.Context, int64) (float64, error) {
|
||||
return s.balance, nil
|
||||
}
|
||||
|
||||
func (s *fixedChargeCacheStub) GetSubscriptionCache(context.Context, int64, int64) (*SubscriptionCacheData, error) {
|
||||
if s.sub == nil {
|
||||
return nil, errors.New("cache miss")
|
||||
}
|
||||
return s.sub, nil
|
||||
}
|
||||
|
||||
func (s *fixedChargeCacheStub) GetAPIKeyRateLimit(context.Context, int64) (*APIKeyRateLimitCacheData, error) {
|
||||
if s.rate == nil {
|
||||
return nil, errors.New("cache miss")
|
||||
}
|
||||
return s.rate, nil
|
||||
}
|
||||
|
||||
func (s *fixedChargeCacheStub) reserve(key, requestID string, amount float64) float64 {
|
||||
if s.reserved == nil {
|
||||
s.reserved = map[string]map[string]float64{}
|
||||
}
|
||||
if s.reserved[key] == nil {
|
||||
s.reserved[key] = map[string]float64{}
|
||||
}
|
||||
sum := 0.0
|
||||
for _, v := range s.reserved[key] {
|
||||
sum += v
|
||||
}
|
||||
s.reserved[key][requestID] = amount
|
||||
return sum
|
||||
}
|
||||
|
||||
func (s *fixedChargeCacheStub) ReserveInflightBalance(_ context.Context, userID int64, requestID string, amount, balance float64, _ time.Duration) (bool, float64, error) {
|
||||
key := fmt.Sprintf("user:%d", userID)
|
||||
others := s.reserve(key, requestID, amount)
|
||||
if len(s.reserved[key]) > 1 && balance-others < amount {
|
||||
delete(s.reserved[key], requestID)
|
||||
return false, others, nil
|
||||
}
|
||||
return true, others, nil
|
||||
}
|
||||
|
||||
func (s *fixedChargeCacheStub) ReleaseInflightBalance(_ context.Context, userID int64, requestID string) error {
|
||||
delete(s.reserved[fmt.Sprintf("user:%d", userID)], requestID)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *fixedChargeCacheStub) ReserveInflightScope(_ context.Context, scope, requestID string, amount float64, _ time.Duration) (float64, error) {
|
||||
return s.reserve(scope, requestID, amount), nil
|
||||
}
|
||||
|
||||
func (s *fixedChargeCacheStub) ReleaseInflightScope(_ context.Context, scope, requestID string) error {
|
||||
delete(s.reserved[scope], requestID)
|
||||
return nil
|
||||
}
|
||||
|
||||
type fixedChargeAPIKeyRepoStub struct {
|
||||
APIKeyRepository
|
||||
data *APIKeyRateLimitData
|
||||
}
|
||||
|
||||
func (r *fixedChargeAPIKeyRepoStub) GetRateLimitData(context.Context, int64) (*APIKeyRateLimitData, error) {
|
||||
return r.data, nil
|
||||
}
|
||||
|
||||
type fixedChargeSubscriptionRepoStub struct {
|
||||
UserSubscriptionRepository
|
||||
sub *UserSubscription
|
||||
}
|
||||
|
||||
func (r *fixedChargeSubscriptionRepoStub) GetActiveByUserIDAndGroupID(context.Context, int64, int64) (*UserSubscription, error) {
|
||||
return r.sub, nil
|
||||
}
|
||||
|
||||
func TestReserveFixedCharge(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
user := &User{ID: 1}
|
||||
newSvc := func(cache *fixedChargeCacheStub) *BillingCacheService {
|
||||
svc := NewBillingCacheService(cache, nil, nil, nil, nil, nil, &config.Config{}, nil)
|
||||
t.Cleanup(svc.Stop)
|
||||
return svc
|
||||
}
|
||||
reservingSvc := func(cache *fixedChargeCacheStub, subRepo UserSubscriptionRepository, apiKeyRepo APIKeyRepository) *BillingCacheService {
|
||||
cfg := &config.Config{}
|
||||
cfg.Billing.InflightReservation.Enabled = true
|
||||
svc := NewBillingCacheService(cache, nil, subRepo, apiKeyRepo, nil, nil, cfg, nil)
|
||||
t.Cleanup(svc.Stop)
|
||||
return svc
|
||||
}
|
||||
reserve := func(svc *BillingCacheService, apiKey *APIKey, group *Group, sub *UserSubscription, amount float64) error {
|
||||
res, err := svc.ReserveFixedCharge(ctx, user, apiKey, group, sub, amount)
|
||||
res.Release()
|
||||
return err
|
||||
}
|
||||
|
||||
t.Run("balance must cover the amount", func(t *testing.T) {
|
||||
svc := newSvc(&fixedChargeCacheStub{balance: 10})
|
||||
require.NoError(t, reserve(svc, nil, nil, nil, 10))
|
||||
require.ErrorIs(t, reserve(svc, nil, nil, nil, 10.01), ErrInsufficientBalance)
|
||||
require.NoError(t, reserve(svc, nil, nil, nil, 0))
|
||||
})
|
||||
|
||||
t.Run("subscription windows must cover the amount", func(t *testing.T) {
|
||||
daily := 5.0
|
||||
group := &Group{ID: 2, SubscriptionType: SubscriptionTypeSubscription, DailyLimitUSD: &daily}
|
||||
svc := newSvc(&fixedChargeCacheStub{sub: &SubscriptionCacheData{
|
||||
Status: SubscriptionStatusActive, ExpiresAt: time.Now().Add(time.Hour), DailyUsage: 3,
|
||||
}})
|
||||
sub := &UserSubscription{ID: 1}
|
||||
require.NoError(t, reserve(svc, nil, group, sub, 2))
|
||||
require.ErrorIs(t, reserve(svc, nil, group, sub, 2.5), ErrDailyLimitExceeded)
|
||||
})
|
||||
|
||||
t.Run("api key quota and rate windows include the amount", func(t *testing.T) {
|
||||
svc := newSvc(&fixedChargeCacheStub{balance: 100, rate: &APIKeyRateLimitCacheData{
|
||||
Usage1d: 8, Window1d: time.Now().Unix(),
|
||||
}})
|
||||
quotaKey := &APIKey{ID: 5, Quota: 10, QuotaUsed: 9}
|
||||
require.NoError(t, reserve(svc, quotaKey, nil, nil, 1))
|
||||
require.ErrorIs(t, reserve(svc, quotaKey, nil, nil, 1.5), ErrAPIKeyQuotaExhausted)
|
||||
|
||||
windowKey := &APIKey{ID: 6, RateLimit1d: 10}
|
||||
require.NoError(t, reserve(svc, windowKey, nil, nil, 2))
|
||||
require.ErrorIs(t, reserve(svc, windowKey, nil, nil, 2.5), ErrAPIKeyRateLimit1dExceeded)
|
||||
})
|
||||
|
||||
t.Run("usage is read from the database instead of request snapshots", func(t *testing.T) {
|
||||
now := time.Now()
|
||||
apiKeyRepo := &fixedChargeAPIKeyRepoStub{data: &APIKeyRateLimitData{QuotaUsed: 9.5, Usage1d: 8, Window1dStart: &now}}
|
||||
svc := reservingSvc(&fixedChargeCacheStub{balance: 100}, nil, apiKeyRepo)
|
||||
require.ErrorIs(t, reserve(svc, &APIKey{ID: 5, Quota: 10}, nil, nil, 1), ErrAPIKeyQuotaExhausted, "请求上下文中的 QuotaUsed 可能滞后")
|
||||
require.ErrorIs(t, reserve(svc, &APIKey{ID: 6, RateLimit1d: 10}, nil, nil, 2.5), ErrAPIKeyRateLimit1dExceeded)
|
||||
|
||||
daily := 5.0
|
||||
group := &Group{ID: 2, SubscriptionType: SubscriptionTypeSubscription, DailyLimitUSD: &daily}
|
||||
subRepo := &fixedChargeSubscriptionRepoStub{sub: &UserSubscription{Status: SubscriptionStatusActive, DailyUsageUSD: 4}}
|
||||
svc = reservingSvc(&fixedChargeCacheStub{sub: &SubscriptionCacheData{Status: SubscriptionStatusActive}}, subRepo, nil)
|
||||
require.ErrorIs(t, reserve(svc, nil, group, &UserSubscription{ID: 1}, 2), ErrDailyLimitExceeded, "订阅缓存用量可能滞后")
|
||||
})
|
||||
|
||||
t.Run("unsettled charges hold subscription and api key headroom", func(t *testing.T) {
|
||||
daily := 110.0
|
||||
group := &Group{ID: 2, SubscriptionType: SubscriptionTypeSubscription, DailyLimitUSD: &daily}
|
||||
sub := &UserSubscription{ID: 1}
|
||||
cache := &fixedChargeCacheStub{balance: 1000, sub: &SubscriptionCacheData{Status: SubscriptionStatusActive}}
|
||||
svc := reservingSvc(cache, nil, nil)
|
||||
first, err := svc.ReserveFixedCharge(ctx, user, nil, group, sub, 103.5)
|
||||
require.NoError(t, err)
|
||||
_, err = svc.ReserveFixedCharge(ctx, user, nil, group, sub, 103.5)
|
||||
require.ErrorIs(t, err, ErrDailyLimitExceeded)
|
||||
first.Release()
|
||||
second, err := svc.ReserveFixedCharge(ctx, user, nil, group, sub, 103.5)
|
||||
require.NoError(t, err)
|
||||
second.Release()
|
||||
|
||||
apiKey := &APIKey{ID: 5, Quota: 110}
|
||||
first, err = svc.ReserveFixedCharge(ctx, user, apiKey, nil, nil, 103.5)
|
||||
require.NoError(t, err)
|
||||
_, err = svc.ReserveFixedCharge(ctx, user, apiKey, nil, nil, 103.5)
|
||||
require.ErrorIs(t, err, ErrAPIKeyQuotaExhausted)
|
||||
first.Release()
|
||||
for key, reservations := range cache.reserved {
|
||||
require.Empty(t, reservations, "拒绝与释放都不留下预留:%s", key)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("simple mode skips the check", func(t *testing.T) {
|
||||
cfg := &config.Config{RunMode: config.RunModeSimple}
|
||||
svc := NewBillingCacheService(&fixedChargeCacheStub{}, nil, nil, nil, nil, nil, cfg, nil)
|
||||
t.Cleanup(svc.Stop)
|
||||
require.NoError(t, reserve(svc, nil, nil, nil, 1000))
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,186 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// 上游对不存在、已过期、已删除或不属于当前项目的显式缓存一律返回这个 403。
|
||||
const geminiCachedContentNotFoundBody = `{"error":{"code":403,"message":"CachedContent not found (or permission denied)","status":"PERMISSION_DENIED"}}`
|
||||
|
||||
func newGeminiCachedContentTestService(status int, body string) (*GeminiMessagesCompatService, *geminiCompatHTTPUpstreamStub, *errorPolicyRepoStub) {
|
||||
repo := &errorPolicyRepoStub{}
|
||||
httpStub := &geminiCompatHTTPUpstreamStub{
|
||||
response: &http.Response{
|
||||
StatusCode: status,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(body)),
|
||||
},
|
||||
}
|
||||
svc := &GeminiMessagesCompatService{
|
||||
httpUpstream: httpStub,
|
||||
cfg: &config.Config{},
|
||||
rateLimitService: NewRateLimitService(repo, nil, &config.Config{}, nil, nil),
|
||||
}
|
||||
return svc, httpStub, repo
|
||||
}
|
||||
|
||||
func geminiDefaultPolicyAPIKeyAccount() *Account {
|
||||
return &Account{
|
||||
ID: 710,
|
||||
Platform: PlatformGemini,
|
||||
Type: AccountTypeAPIKey,
|
||||
Credentials: map[string]any{
|
||||
"api_key": "test-key",
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
const geminiCachedContentRequestBody = `{"cachedContent":"cachedContents/doesnotexist","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`
|
||||
|
||||
func TestGeminiForwardNative_CachedContentNotFoundPassthroughWithoutAccountPenalty(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
action string
|
||||
stream bool
|
||||
}{
|
||||
{name: "generateContent", action: "generateContent"},
|
||||
{name: "streamGenerateContent", action: "streamGenerateContent", stream: true},
|
||||
{name: "countTokens", action: "countTokens"},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
svc, httpStub, repo := newGeminiCachedContentTestService(http.StatusForbidden, geminiCachedContentNotFoundBody)
|
||||
c, rec := newGeminiNativeTestContext(t)
|
||||
|
||||
result, err := svc.ForwardNative(context.Background(), c, geminiDefaultPolicyAPIKeyAccount(),
|
||||
"gemini-3.8-flash", tc.action, tc.stream, []byte(geminiCachedContentRequestBody))
|
||||
|
||||
require.Nil(t, result)
|
||||
require.Error(t, err)
|
||||
var failoverErr *UpstreamFailoverError
|
||||
require.False(t, errors.As(err, &failoverErr), "缓存不存在是请求问题,不应换号")
|
||||
require.Equal(t, 0, repo.setErrCalls, "不应把账号标记为错误")
|
||||
require.Equal(t, 0, repo.tempCalls, "不应临时停调度")
|
||||
require.Equal(t, 1, httpStub.calls, "不应重试")
|
||||
require.Equal(t, http.StatusForbidden, rec.Code)
|
||||
require.Equal(t, GeminiCachedContentNotFoundResponse, rec.Body.String(), "应返回标准缓存不存在响应")
|
||||
require.True(t, IsOpsRequestScopedError(c), "缓存不存在按客户端请求错误统计")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeminiForwardNative_CachedContentNotFoundBypassesTempUnschedulableRules(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc, _, repo := newGeminiCachedContentTestService(http.StatusForbidden, geminiCachedContentNotFoundBody)
|
||||
c, rec := newGeminiNativeTestContext(t)
|
||||
account := geminiDefaultPolicyAPIKeyAccount()
|
||||
account.Credentials["temp_unschedulable_enabled"] = true
|
||||
account.Credentials["temp_unschedulable_rules"] = []any{
|
||||
map[string]any{
|
||||
"error_code": float64(403),
|
||||
"keywords": []any{"permission denied"},
|
||||
"duration_minutes": float64(10),
|
||||
"description": "403 permission denied",
|
||||
},
|
||||
}
|
||||
|
||||
_, err := svc.ForwardNative(context.Background(), c, account,
|
||||
"gemini-3.8-flash", "generateContent", false, []byte(geminiCachedContentRequestBody))
|
||||
|
||||
require.Error(t, err)
|
||||
var failoverErr *UpstreamFailoverError
|
||||
require.False(t, errors.As(err, &failoverErr))
|
||||
require.Equal(t, 0, repo.tempCalls, "临时不可调度规则不应命中")
|
||||
require.Equal(t, 0, repo.setErrCalls)
|
||||
require.Equal(t, http.StatusForbidden, rec.Code)
|
||||
}
|
||||
|
||||
func TestGeminiForwardNative_OtherForbiddenStillPenalizesAccount(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
body := `{"error":{"code":403,"message":"Your project has been denied access. Please contact support.","status":"PERMISSION_DENIED"}}`
|
||||
svc, _, repo := newGeminiCachedContentTestService(http.StatusForbidden, body)
|
||||
c, _ := newGeminiNativeTestContext(t)
|
||||
|
||||
_, err := svc.ForwardNative(context.Background(), c, geminiDefaultPolicyAPIKeyAccount(),
|
||||
"gemini-3.8-flash", "generateContent", false, []byte(geminiCachedContentRequestBody))
|
||||
|
||||
var failoverErr *UpstreamFailoverError
|
||||
require.True(t, errors.As(err, &failoverErr), "其他 403 仍按账号问题换号")
|
||||
require.Equal(t, 1, repo.setErrCalls, "其他 403 仍标记账号错误")
|
||||
}
|
||||
|
||||
// Vertex 对已删除 / 已过期缓存的响应(2026-10 实测):生成请求 400,缓存资源接口 404,文案含上游缓存数字 ID。
|
||||
const (
|
||||
geminiVertexCachedContentInvalidStateBody = `{"error":{"code":400,"message":"Invalid resource state for cache content 8079767076522164224.","status":"FAILED_PRECONDITION"}}`
|
||||
geminiVertexCachedContentNotFoundBody = `{"error":{"code":404,"message":"Cached content 8079767076522164224 is not found.","status":"NOT_FOUND"}}`
|
||||
)
|
||||
|
||||
func TestIsGeminiCachedContentNotFound(t *testing.T) {
|
||||
require.True(t, isGeminiCachedContentNotFound(http.StatusForbidden, []byte(geminiCachedContentNotFoundBody)))
|
||||
require.True(t, isGeminiCachedContentNotFound(http.StatusForbidden,
|
||||
[]byte(`[{"error":{"code":403,"message":"CachedContent not found (or permission denied)","status":"PERMISSION_DENIED"}}]`)))
|
||||
require.True(t, isGeminiCachedContentNotFound(http.StatusBadRequest, []byte(geminiVertexCachedContentInvalidStateBody)))
|
||||
require.True(t, isGeminiCachedContentNotFound(http.StatusNotFound, []byte(geminiVertexCachedContentNotFoundBody)))
|
||||
|
||||
require.False(t, isGeminiCachedContentNotFound(http.StatusBadRequest, []byte(geminiCachedContentNotFoundBody)))
|
||||
require.False(t, isGeminiCachedContentNotFound(http.StatusForbidden,
|
||||
[]byte(`{"error":{"code":403,"message":"Permission denied: Consumer has been suspended.","status":"PERMISSION_DENIED"}}`)))
|
||||
require.False(t, isGeminiCachedContentNotFound(http.StatusNotFound,
|
||||
[]byte(`{"error":{"code":404,"message":"Publisher Model `+"`projects/p/locations/global/publishers/google/models/gemini-9`"+` was not found.","status":"NOT_FOUND"}}`)))
|
||||
require.False(t, isGeminiCachedContentNotFound(http.StatusBadRequest,
|
||||
[]byte(`{"error":{"code":400,"message":"Tool config, tools and system instruction should not be set in the request when using cached content.","status":"INVALID_ARGUMENT"}}`)))
|
||||
require.False(t, isGeminiCachedContentNotFound(http.StatusInternalServerError, []byte(geminiVertexCachedContentNotFoundBody)))
|
||||
}
|
||||
|
||||
func TestGeminiForwardNative_VertexCachedContentNotFoundReturnsStandardForbidden(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
status int
|
||||
body string
|
||||
action string
|
||||
stream bool
|
||||
}{
|
||||
{name: "generate invalid resource state", status: http.StatusBadRequest, body: geminiVertexCachedContentInvalidStateBody, action: "generateContent"},
|
||||
{name: "stream invalid resource state", status: http.StatusBadRequest, body: geminiVertexCachedContentInvalidStateBody, action: "streamGenerateContent", stream: true},
|
||||
{name: "generate not found", status: http.StatusNotFound, body: geminiVertexCachedContentNotFoundBody, action: "generateContent"},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
svc, httpStub, repo := newGeminiCachedContentTestService(tc.status, tc.body)
|
||||
svc.tokenProvider = NewGeminiTokenProvider(nil, &fakeGeminiTokenCache{token: "ya29.test-token"}, nil)
|
||||
c, rec := newGeminiNativeTestContext(t)
|
||||
account := vertexServiceAccount()
|
||||
account.ID = 811
|
||||
|
||||
result, err := svc.ForwardNative(context.Background(), c, account,
|
||||
"gemini-3.8-flash", tc.action, tc.stream,
|
||||
[]byte(`{"cachedContent":"projects/proj-1/locations/global/cachedContents/8079767076522164224","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`))
|
||||
|
||||
require.Nil(t, result)
|
||||
require.Error(t, err)
|
||||
var failoverErr *UpstreamFailoverError
|
||||
require.False(t, errors.As(err, &failoverErr), "缓存不存在是请求问题,不应换号")
|
||||
require.Equal(t, 0, repo.setErrCalls)
|
||||
require.Equal(t, 0, repo.tempCalls)
|
||||
require.Empty(t, repo.modelRateLimitCalls, "404 不应按模型不存在记模型限流")
|
||||
require.Equal(t, 1, httpStub.calls, "不应重试")
|
||||
require.Equal(t, http.StatusForbidden, rec.Code)
|
||||
require.Equal(t, GeminiCachedContentNotFoundResponse, rec.Body.String())
|
||||
require.NotContains(t, rec.Body.String(), "8079767076522164224", "不应暴露上游缓存 ID")
|
||||
require.True(t, IsOpsRequestScopedError(c), "缓存不存在按客户端请求错误统计")
|
||||
errType, _ := OpsLocalErrorType(c)
|
||||
require.Equal(t, "permission_error", errType)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,522 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type geminiCachedContentRepoStub struct {
|
||||
records map[string]*GeminiCachedContent
|
||||
nextID int64
|
||||
deleted []int64
|
||||
}
|
||||
|
||||
func newGeminiCachedContentRepoStub() *geminiCachedContentRepoStub {
|
||||
return &geminiCachedContentRepoStub{records: map[string]*GeminiCachedContent{}}
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentRepoStub) Create(_ context.Context, record *GeminiCachedContent) error {
|
||||
r.nextID++
|
||||
record.ID = r.nextID
|
||||
r.records[record.PublicID] = record
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentRepoStub) GetForOwner(_ context.Context, apiKeyID int64, publicID string) (*GeminiCachedContent, error) {
|
||||
record, ok := r.records[publicID]
|
||||
if !ok || record.APIKeyID != apiKeyID {
|
||||
return nil, ErrGeminiCachedContentNotFound
|
||||
}
|
||||
return record, nil
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentRepoStub) ListForOwner(_ context.Context, apiKeyID, groupID int64, activeAt time.Time, beforeID int64, limit int) ([]*GeminiCachedContent, error) {
|
||||
var out []*GeminiCachedContent
|
||||
for _, record := range r.records {
|
||||
if record.APIKeyID == apiKeyID && record.GroupID == groupID && record.ExpireTime.After(activeAt) && (beforeID <= 0 || record.ID < beforeID) {
|
||||
out = append(out, record)
|
||||
}
|
||||
}
|
||||
if len(out) > limit {
|
||||
out = out[:limit]
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentRepoStub) BeginMutation(context.Context, int64, time.Time, string, time.Duration) (bool, error) {
|
||||
return true, nil
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentRepoStub) TakeOverStaleMutation(context.Context, int64, string, time.Duration) (*GeminiCachedContent, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentRepoStub) CommitMutation(_ context.Context, id int64, _ string, expireTime time.Time) error {
|
||||
for _, record := range r.records {
|
||||
if record.ID == id {
|
||||
record.ExpireTime = expireTime
|
||||
return nil
|
||||
}
|
||||
}
|
||||
return ErrGeminiCachedContentMutationConflict
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentRepoStub) ReleaseMutation(context.Context, int64, string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentRepoStub) ListStaleMutations(context.Context, int) ([]int64, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentRepoStub) SoftDelete(_ context.Context, id int64) error {
|
||||
r.deleted = append(r.deleted, id)
|
||||
for key, record := range r.records {
|
||||
if record.ID == id {
|
||||
delete(r.records, key)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestParseGeminiCachedContentCreateRequest(t *testing.T) {
|
||||
now := time.Date(2026, 10, 9, 8, 0, 0, 0, time.UTC)
|
||||
maxTTL := 24 * time.Hour
|
||||
|
||||
t.Run("defaults to one hour and strips output-only fields", func(t *testing.T) {
|
||||
req, err := ParseGeminiCachedContentCreateRequest([]byte(`{"model":"models/gemini-3.8-flash","displayName":"docs",
|
||||
"contents":[{"role":"user","parts":[{"text":"x"}]}],"name":"cachedContents/forged","usageMetadata":{"totalTokenCount":1}}`), now, maxTTL)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "gemini-3.8-flash", req.Model)
|
||||
require.Equal(t, time.Hour, req.TTL)
|
||||
require.Equal(t, int64(3600), req.TTLSeconds())
|
||||
require.Equal(t, "docs", req.DisplayName)
|
||||
require.NotContains(t, req.payload, "name")
|
||||
require.NotContains(t, req.payload, "usageMetadata")
|
||||
require.Contains(t, req.payload, "contents")
|
||||
})
|
||||
|
||||
t.Run("default is capped by a shorter max ttl", func(t *testing.T) {
|
||||
req, err := ParseGeminiCachedContentCreateRequest([]byte(`{"model":"gemini-3.8-flash"}`), now, 30*time.Minute)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 30*time.Minute, req.TTL)
|
||||
})
|
||||
|
||||
t.Run("fractional ttl rounds up to whole seconds", func(t *testing.T) {
|
||||
req, err := ParseGeminiCachedContentCreateRequest([]byte(`{"model":"gemini-3.8-flash","ttl":"3.5s"}`), now, maxTTL)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(4), req.TTLSeconds())
|
||||
})
|
||||
|
||||
t.Run("expireTime converts to ttl", func(t *testing.T) {
|
||||
req, err := ParseGeminiCachedContentCreateRequest([]byte(`{"model":"gemini-3.8-flash","expireTime":"2026-10-09T10:00:00Z"}`), now, maxTTL)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 2*time.Hour, req.TTL)
|
||||
})
|
||||
|
||||
for name, body := range map[string]string{
|
||||
"missing model": `{"ttl":"60s"}`,
|
||||
"vertex style model": `{"model":"projects/p/locations/l/publishers/google/models/x"}`,
|
||||
"ttl and expireTime": `{"model":"m","ttl":"60s","expireTime":"2026-10-09T10:00:00Z"}`,
|
||||
"ttl over max": `{"model":"m","ttl":"86401s"}`,
|
||||
"non positive ttl": `{"model":"m","ttl":"0s"}`,
|
||||
"ttl without unit": `{"model":"m","ttl":"60"}`,
|
||||
"expireTime in past": `{"model":"m","expireTime":"2026-10-09T07:00:00Z"}`,
|
||||
"displayName too long": `{"model":"m","displayName":"` + strings.Repeat("字", 129) + `"}`,
|
||||
"not an object": `[]`,
|
||||
} {
|
||||
t.Run("rejects "+name, func(t *testing.T) {
|
||||
_, err := ParseGeminiCachedContentCreateRequest([]byte(body), now, maxTTL)
|
||||
require.Error(t, err)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseGeminiCachedContentPatchRequest(t *testing.T) {
|
||||
now := time.Date(2026, 10, 9, 8, 0, 0, 0, time.UTC)
|
||||
ttl, err := ParseGeminiCachedContentPatchRequest([]byte(`{"ttl":"7200s"}`), "ttl", now, 24*time.Hour)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 2*time.Hour, ttl)
|
||||
|
||||
ttl, err = ParseGeminiCachedContentPatchRequest([]byte(`{"expireTime":"2026-10-09T09:00:00Z"}`), "", now, 24*time.Hour)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, time.Hour, ttl)
|
||||
|
||||
_, err = ParseGeminiCachedContentPatchRequest([]byte(`{"displayName":"x"}`), "displayName", now, 24*time.Hour)
|
||||
require.ErrorContains(t, err, "immutable")
|
||||
_, err = ParseGeminiCachedContentPatchRequest([]byte(`{"ttl":"60s"}`), "ttl,displayName", now, 24*time.Hour)
|
||||
require.ErrorContains(t, err, "immutable")
|
||||
_, err = ParseGeminiCachedContentPatchRequest([]byte(`{}`), "", now, 24*time.Hour)
|
||||
require.Error(t, err)
|
||||
_, err = ParseGeminiCachedContentPatchRequest([]byte(`{"ttl":"90000s"}`), "", now, 24*time.Hour)
|
||||
require.ErrorContains(t, err, "maximum")
|
||||
}
|
||||
|
||||
func TestGeminiCachedContentReferenceRewrite(t *testing.T) {
|
||||
ref, err := FindGeminiCachedContentReference([]byte(`{"contents":[]}`))
|
||||
require.NoError(t, err)
|
||||
require.Nil(t, ref)
|
||||
|
||||
body := []byte(`{"cachedContent":"cachedContents/abc","contents":[]}`)
|
||||
ref, err = FindGeminiCachedContentReference(body)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "cachedContents/abc", ref.Name)
|
||||
out, err := RewriteGeminiCachedContentReference(body, ref, "cachedContents/upstream1")
|
||||
require.NoError(t, err)
|
||||
require.JSONEq(t, `{"cachedContent":"cachedContents/upstream1","contents":[]}`, string(out))
|
||||
|
||||
countBody := []byte(`{"generateContentRequest":{"model":"models/m","cachedContent":"cachedContents/abc","contents":[]}}`)
|
||||
ref, err = FindGeminiCachedContentReference(countBody)
|
||||
require.NoError(t, err)
|
||||
out, err = RewriteGeminiCachedContentReference(countBody, ref, "projects/1/locations/us-central1/cachedContents/9")
|
||||
require.NoError(t, err)
|
||||
require.JSONEq(t, `{"generateContentRequest":{"model":"models/m","cachedContent":"projects/1/locations/us-central1/cachedContents/9","contents":[]}}`, string(out))
|
||||
|
||||
_, err = FindGeminiCachedContentReference([]byte(`{"cachedContent":"cachedContents/a","generateContentRequest":{"cachedContent":"cachedContents/b"}}`))
|
||||
require.Error(t, err)
|
||||
_, err = FindGeminiCachedContentReference([]byte(`{"cachedContent":1}`))
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func TestIsGeminiExplicitCacheAccount(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
account *Account
|
||||
want bool
|
||||
}{
|
||||
{"official api key", &Account{Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "k"}}, true},
|
||||
{"explicit official base url", &Account{Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "k", "base_url": "https://generativelanguage.googleapis.com/"}}, true},
|
||||
{"reseller base url", &Account{Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "k", "base_url": "https://relay.example.com"}}, false},
|
||||
{"upstream gateway with explicit cache switch", &Account{Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "k", "base_url": "https://relay.example.com", "explicit_cache_upstream": true}}, true},
|
||||
{"official host with path prefix", &Account{Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "k", "base_url": "https://generativelanguage.googleapis.com/antigravity"}}, false},
|
||||
{"missing api key", &Account{Platform: PlatformGemini, Type: AccountTypeAPIKey}, false},
|
||||
{"vertex service account", &Account{Platform: PlatformGemini, Type: AccountTypeServiceAccount}, true},
|
||||
{"vertex api key", &Account{Platform: PlatformGemini, Type: AccountTypeServiceAccount, Credentials: map[string]any{"auth_mode": "apikey", "api_key": "k", "project_id": "my-project"}}, false},
|
||||
{"oauth", &Account{Platform: PlatformGemini, Type: AccountTypeOAuth}, false},
|
||||
{"antigravity api key", &Account{Platform: PlatformAntigravity, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "k"}}, false},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
require.Equal(t, tc.want, IsGeminiExplicitCacheAccount(tc.account))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeminiCachedContentVertexURLs(t *testing.T) {
|
||||
account := &Account{Platform: PlatformGemini, Type: AccountTypeServiceAccount, Credentials: map[string]any{
|
||||
"project_id": "proj-1",
|
||||
"location": "us-central1",
|
||||
"vertex_model_locations": map[string]any{"gemini-3.8-flash": "global"},
|
||||
}}
|
||||
model, err := GeminiCachedContentUpstreamModelResource(account, "gemini-3.8-flash")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "projects/proj-1/locations/global/publishers/google/models/gemini-3.8-flash", model)
|
||||
|
||||
collection, err := geminiCachedContentCollectionURL(account, "", "gemini-3.8-flash")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "https://aiplatform.googleapis.com/v1/projects/proj-1/locations/global/cachedContents", collection)
|
||||
|
||||
collection, err = geminiCachedContentCollectionURL(account, "", "gemini-3.1-pro-preview")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "https://us-central1-aiplatform.googleapis.com/v1/projects/proj-1/locations/us-central1/cachedContents", collection)
|
||||
|
||||
resource, err := geminiCachedContentResourceURL(account, "", "projects/123/locations/europe-west4/cachedContents/456")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "https://europe-west4-aiplatform.googleapis.com/v1/projects/123/locations/europe-west4/cachedContents/456", resource)
|
||||
|
||||
_, err = geminiCachedContentResourceURL(account, "", "projects/123/locations/europe-west4/cachedContents/../x")
|
||||
require.Error(t, err)
|
||||
_, err = geminiCachedContentResourceURL(&Account{Type: AccountTypeAPIKey}, "https://generativelanguage.googleapis.com", "cachedContents/a/../b")
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func newGeminiCachedContentServiceForTest(t *testing.T, status int, respBody string) (*GeminiCachedContentService, *geminiCompatHTTPUpstreamStub) {
|
||||
t.Helper()
|
||||
httpStub := &geminiCompatHTTPUpstreamStub{response: &http.Response{
|
||||
StatusCode: status,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(respBody)),
|
||||
}}
|
||||
compat := &GeminiMessagesCompatService{httpUpstream: httpStub, cfg: &config.Config{}}
|
||||
return NewGeminiCachedContentService(newGeminiCachedContentRepoStub(), compat, &config.Config{}), httpStub
|
||||
}
|
||||
|
||||
func TestGeminiCachedContentCreateUpstreamAIStudio(t *testing.T) {
|
||||
svc, httpStub := newGeminiCachedContentServiceForTest(t, http.StatusOK,
|
||||
`{"name":"cachedContents/up123","model":"models/gemini-3.8-flash","usageMetadata":{"totalTokenCount":6273},"expireTime":"2026-10-09T09:50:26.591575799Z"}`)
|
||||
account := &Account{ID: 9, Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "secret"}}
|
||||
req, err := ParseGeminiCachedContentCreateRequest([]byte(`{"model":"models/gemini-3.8-flash","displayName":"user label",
|
||||
"systemInstruction":{"parts":[{"text":"sys"}]},"contents":[{"role":"user","parts":[{"text":"doc"}]}]}`), time.Now(), 24*time.Hour)
|
||||
require.NoError(t, err)
|
||||
|
||||
upstream, err := svc.CreateUpstream(context.Background(), account, req, "gemini-3.8-flash", "pub1", req.TTL)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "cachedContents/up123", upstream.Name)
|
||||
require.Equal(t, int64(6273), upstream.TotalTokenCount)
|
||||
require.Equal(t, 2026, upstream.ExpireTime.Year())
|
||||
|
||||
sent := httpStub.lastReq
|
||||
require.Equal(t, http.MethodPost, sent.Method)
|
||||
require.Equal(t, "https://generativelanguage.googleapis.com/v1beta/cachedContents", sent.URL.String())
|
||||
require.Equal(t, "secret", sent.Header.Get("x-goog-api-key"))
|
||||
raw, err := io.ReadAll(sent.Body)
|
||||
require.NoError(t, err)
|
||||
var payload map[string]any
|
||||
require.NoError(t, json.Unmarshal(raw, &payload))
|
||||
require.Equal(t, "models/gemini-3.8-flash", payload["model"])
|
||||
require.Equal(t, "3600s", payload["ttl"])
|
||||
require.Equal(t, "s2a-pub1", payload["displayName"], "上游 displayName 写网关标记,用户标签只存本地")
|
||||
require.Contains(t, payload, "systemInstruction")
|
||||
require.Contains(t, payload, "contents")
|
||||
}
|
||||
|
||||
func TestGeminiCachedContentUpstreamErrorClassification(t *testing.T) {
|
||||
notFoundBody := `{"error":{"code":403,"message":"CachedContent not found (or permission denied)","status":"PERMISSION_DENIED"}}`
|
||||
svc, _ := newGeminiCachedContentServiceForTest(t, http.StatusForbidden, notFoundBody)
|
||||
account := &Account{ID: 9, Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "secret"}}
|
||||
record := &GeminiCachedContent{UpstreamName: "cachedContents/up123"}
|
||||
|
||||
err := svc.DeleteUpstream(context.Background(), account, record)
|
||||
var upErr *GeminiCachedContentUpstreamError
|
||||
require.ErrorAs(t, err, &upErr)
|
||||
require.True(t, upErr.NotFound())
|
||||
require.True(t, upErr.RetryableOnOtherAccount())
|
||||
|
||||
quota := &GeminiCachedContentUpstreamError{StatusCode: http.StatusTooManyRequests}
|
||||
require.True(t, quota.RetryableOnOtherAccount())
|
||||
badRequest := &GeminiCachedContentUpstreamError{StatusCode: http.StatusBadRequest}
|
||||
require.False(t, badRequest.RetryableOnOtherAccount())
|
||||
require.False(t, badRequest.NotFound())
|
||||
}
|
||||
|
||||
func TestGeminiCachedContentResolve(t *testing.T) {
|
||||
repo := newGeminiCachedContentRepoStub()
|
||||
now := time.Date(2026, 10, 9, 8, 0, 0, 0, time.UTC)
|
||||
svc := NewGeminiCachedContentService(repo, nil, &config.Config{})
|
||||
svc.now = func() time.Time { return now }
|
||||
groupID := int64(5)
|
||||
require.NoError(t, repo.Create(context.Background(), &GeminiCachedContent{
|
||||
PublicID: "abc123", APIKeyID: 1, GroupID: groupID, ExpireTime: now.Add(time.Hour),
|
||||
}))
|
||||
require.NoError(t, repo.Create(context.Background(), &GeminiCachedContent{
|
||||
PublicID: "expired1", APIKeyID: 1, GroupID: groupID, ExpireTime: now.Add(-time.Second),
|
||||
}))
|
||||
|
||||
record, err := svc.Resolve(context.Background(), 1, &groupID, "cachedContents/abc123")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "abc123", record.PublicID)
|
||||
|
||||
otherGroup := int64(6)
|
||||
for name, call := range map[string]func() error{
|
||||
"other api key": func() error {
|
||||
_, err := svc.Resolve(context.Background(), 2, &groupID, "cachedContents/abc123")
|
||||
return err
|
||||
},
|
||||
"other group": func() error {
|
||||
_, err := svc.Resolve(context.Background(), 1, &otherGroup, "cachedContents/abc123")
|
||||
return err
|
||||
},
|
||||
"expired": func() error {
|
||||
_, err := svc.Resolve(context.Background(), 1, &groupID, "cachedContents/expired1")
|
||||
return err
|
||||
},
|
||||
"malformed": func() error {
|
||||
_, err := svc.Resolve(context.Background(), 1, &groupID, "cachedContents/../x")
|
||||
return err
|
||||
},
|
||||
} {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
require.ErrorIs(t, call(), ErrGeminiCachedContentNotFound)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeminiCachedContentPublicID(t *testing.T) {
|
||||
id, err := NewGeminiCachedContentPublicID()
|
||||
require.NoError(t, err)
|
||||
require.Len(t, id, 40)
|
||||
parsed, ok := ParseGeminiCachedContentPublicID(GeminiCachedContentNamePrefix + id)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, id, parsed)
|
||||
}
|
||||
|
||||
func TestGeminiCachedContentStoragePricingAndSurcharge(t *testing.T) {
|
||||
pricing := &PricingService{}
|
||||
data, err := pricing.parsePricingData([]byte(`{"gemini-3.8-flash":{"litellm_provider":"gemini","input_cost_per_token":0.00000075,
|
||||
"output_cost_per_token":0.00000375,"cache_read_input_token_cost":0.000000075,"cache_storage_cost_per_token_per_hour":0.0000005}}`))
|
||||
require.NoError(t, err)
|
||||
pricing.pricingData = data
|
||||
billing := NewBillingService(&config.Config{}, pricing)
|
||||
|
||||
price := billing.CacheStoragePricePerTokenHour("models/gemini-3.8-flash")
|
||||
require.InDelta(t, 0.0000005, price, 1e-15)
|
||||
require.Zero(t, billing.CacheStoragePricePerTokenHour("gemini-unknown-model"))
|
||||
|
||||
tokenHours := GeminiCachedContentStorageTokenHours(10000, 2*time.Hour)
|
||||
require.InDelta(t, 20000, tokenHours, 1e-9)
|
||||
storage := tokenHours * price
|
||||
require.InDelta(t, 0.01, storage, 1e-12)
|
||||
|
||||
gateway := &GatewayService{billingService: billing}
|
||||
groupID := int64(3)
|
||||
apiKey := &APIKey{GroupID: &groupID, Group: &Group{ID: groupID, Hydrated: true}}
|
||||
cost := gateway.calculateRecordUsageCost(context.Background(), &ForwardResult{
|
||||
Model: "gemini-3.8-flash",
|
||||
Usage: ClaudeUsage{InputTokens: 10000},
|
||||
CacheStorageTokenHours: tokenHours,
|
||||
}, apiKey, "gemini-3.8-flash", 2, 2, time.Now())
|
||||
inputCost := 10000 * 0.00000075
|
||||
require.InDelta(t, inputCost+storage, cost.TotalCost, 1e-12)
|
||||
require.InDelta(t, (inputCost+storage)*2, cost.ActualCost, 1e-12)
|
||||
|
||||
storageOnly := gateway.calculateRecordUsageCost(context.Background(), &ForwardResult{
|
||||
Model: "gemini-3.8-flash",
|
||||
CacheStorageTokenHours: tokenHours,
|
||||
}, apiKey, "gemini-3.8-flash", 1.5, 1.5, time.Now())
|
||||
require.InDelta(t, storage, storageOnly.TotalCost, 1e-12)
|
||||
require.InDelta(t, storage*1.5, storageOnly.ActualCost, 1e-12)
|
||||
}
|
||||
|
||||
func TestGeminiCachedContentStorageCostInAccountStats(t *testing.T) {
|
||||
const (
|
||||
model = "gemini-3.8-flash"
|
||||
groupID = int64(31)
|
||||
inputTokens = 10000
|
||||
storage = 0.01
|
||||
inputPrice = 0.00000075
|
||||
rulePrice = 0.0000002
|
||||
)
|
||||
newPricing := func() *PricingService {
|
||||
pricing := &PricingService{}
|
||||
data, err := pricing.parsePricingData([]byte(`{"gemini-3.8-flash":{"litellm_provider":"gemini","input_cost_per_token":0.00000075,
|
||||
"output_cost_per_token":0.00000375,"cache_storage_cost_per_token_per_hour":0.0000005}}`))
|
||||
require.NoError(t, err)
|
||||
pricing.pricingData = data
|
||||
return pricing
|
||||
}
|
||||
officialInput := inputTokens * inputPrice
|
||||
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
channel *Channel
|
||||
tokens int
|
||||
wantStats float64
|
||||
}{
|
||||
{name: "no_channel", tokens: inputTokens, wantStats: officialInput + storage},
|
||||
{name: "model_file_pricing", channel: &Channel{ID: 1, Status: StatusActive}, tokens: inputTokens, wantStats: officialInput + storage},
|
||||
{name: "apply_pricing_to_account_stats", channel: &Channel{ID: 1, Status: StatusActive, ApplyPricingToAccountStats: true},
|
||||
tokens: inputTokens, wantStats: officialInput + storage},
|
||||
{name: "custom_rule", channel: &Channel{ID: 1, Status: StatusActive, AccountStatsPricingRules: []AccountStatsPricingRule{{
|
||||
GroupIDs: []int64{groupID},
|
||||
Pricing: []ChannelModelPricing{{Models: []string{model}, InputPrice: testPtrFloat64(rulePrice)}},
|
||||
}}}, tokens: inputTokens, wantStats: inputTokens*rulePrice + storage},
|
||||
{name: "storage_only_with_custom_rule", channel: &Channel{ID: 1, Status: StatusActive, AccountStatsPricingRules: []AccountStatsPricingRule{{
|
||||
GroupIDs: []int64{groupID},
|
||||
Pricing: []ChannelModelPricing{{Models: []string{model}, InputPrice: testPtrFloat64(rulePrice)}},
|
||||
}}}, wantStats: storage},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
usageRepo := &openAIRecordUsageLogRepoStub{inserted: true}
|
||||
svc := newGatewayRecordUsageServiceForTest(usageRepo, &openAIRecordUsageUserRepoStub{}, &openAIRecordUsageSubRepoStub{})
|
||||
svc.billingService = NewBillingService(svc.cfg, newPricing())
|
||||
if tc.channel != nil {
|
||||
svc.channelService = newTestChannelServiceForStats(t, tc.channel, groupID, PlatformGemini)
|
||||
}
|
||||
groupRate, accountRate := 1.5, 0.6
|
||||
gid := groupID
|
||||
err := svc.RecordUsage(context.Background(), &RecordUsageInput{
|
||||
Result: &ForwardResult{
|
||||
RequestID: GeminiCachedContentCreateUsageRequestID(tc.name),
|
||||
Model: model,
|
||||
Usage: ClaudeUsage{InputTokens: tc.tokens},
|
||||
CacheStorageTokenHours: storage / 0.0000005,
|
||||
},
|
||||
APIKey: &APIKey{ID: 811, GroupID: &gid, Group: &Group{ID: groupID, Platform: PlatformGemini, RateMultiplier: groupRate, Hydrated: true}},
|
||||
User: &User{ID: 611},
|
||||
Account: &Account{ID: 711, Platform: PlatformGemini, RateMultiplier: &accountRate},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
log := usageRepo.lastLog
|
||||
require.NotNil(t, log)
|
||||
|
||||
wantTotal := float64(tc.tokens)*inputPrice + storage
|
||||
require.InDelta(t, wantTotal, log.TotalCost, 1e-12)
|
||||
require.InDelta(t, wantTotal*groupRate, log.ActualCost, 1e-12, "售价 = 官方价 × 分组倍率")
|
||||
|
||||
statsCost := log.TotalCost
|
||||
if log.AccountStatsCost != nil {
|
||||
statsCost = *log.AccountStatsCost
|
||||
}
|
||||
require.InDelta(t, tc.wantStats, statsCost, 1e-12, "账号统计成本须含存储费且只计一次")
|
||||
require.NotNil(t, log.AccountRateMultiplier)
|
||||
require.InDelta(t, tc.wantStats*accountRate, statsCost**log.AccountRateMultiplier, 1e-12, "成本 = 统计成本 × 账号倍率")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeminiExplicitCacheUpstreamTTL(t *testing.T) {
|
||||
now := time.Date(2026, 10, 9, 8, 0, 0, 0, time.UTC)
|
||||
upstream := &Account{Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: map[string]any{
|
||||
"api_key": "k", "base_url": "https://relay.example.com", "explicit_cache_upstream": true,
|
||||
"explicit_cache_upstream_max_ttl_seconds": float64(1800),
|
||||
}}
|
||||
official := &Account{Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: map[string]any{
|
||||
"api_key": "k", "explicit_cache_upstream_max_ttl_seconds": float64(1800),
|
||||
}}
|
||||
require.Equal(t, 30*time.Minute, upstream.GeminiExplicitCacheUpstreamMaxTTL())
|
||||
require.Zero(t, official.GeminiExplicitCacheUpstreamMaxTTL(), "未开启开关时上限字段不生效")
|
||||
|
||||
svc := NewGeminiCachedContentService(newGeminiCachedContentRepoStub(), nil, &config.Config{})
|
||||
require.Equal(t, 30*time.Minute, svc.MaxTTLForAccount(upstream))
|
||||
require.Equal(t, 24*time.Hour, svc.MaxTTLForAccount(official))
|
||||
|
||||
implicit, err := ParseGeminiCachedContentCreateRequest([]byte(`{"model":"gemini-3.8-flash"}`), now, 24*time.Hour)
|
||||
require.NoError(t, err)
|
||||
ttl, ok := implicit.TTLForAccount(upstream)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, 30*time.Minute, ttl, "未显式指定有效期时按上游上限创建")
|
||||
ttl, ok = implicit.TTLForAccount(official)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, time.Hour, ttl)
|
||||
|
||||
explicit, err := ParseGeminiCachedContentCreateRequest([]byte(`{"model":"gemini-3.8-flash","ttl":"3600s"}`), now, 24*time.Hour)
|
||||
require.NoError(t, err)
|
||||
_, ok = explicit.TTLForAccount(upstream)
|
||||
require.False(t, ok, "显式有效期超过上游上限时该账号不可用")
|
||||
ttl, ok = explicit.TTLForAccount(official)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, time.Hour, ttl)
|
||||
}
|
||||
|
||||
func TestGeminiCachedContentCreateUpstreamGateway(t *testing.T) {
|
||||
svc, httpStub := newGeminiCachedContentServiceForTest(t, http.StatusOK,
|
||||
`{"name":"cachedContents/gatewayid0001","model":"models/gemini-3.8-flash","usageMetadata":{"totalTokenCount":2730},"expireTime":"2026-10-09T09:00:00Z"}`)
|
||||
account := &Account{ID: 11, Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: map[string]any{
|
||||
"api_key": "downstream-key", "base_url": "https://upstream.example.com/gw/", "explicit_cache_upstream": true,
|
||||
}}
|
||||
req, err := ParseGeminiCachedContentCreateRequest([]byte(`{"model":"gemini-3.8-flash","contents":[{"role":"user","parts":[{"text":"doc"}]}]}`), time.Now(), 24*time.Hour)
|
||||
require.NoError(t, err)
|
||||
|
||||
upstream, err := svc.CreateUpstream(context.Background(), account, req, "gemini-3.8-flash", "pub2", 30*time.Minute)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "cachedContents/gatewayid0001", upstream.Name)
|
||||
require.Equal(t, "https://upstream.example.com/gw/v1beta/cachedContents", httpStub.lastReq.URL.String())
|
||||
require.Equal(t, "downstream-key", httpStub.lastReq.Header.Get("x-goog-api-key"))
|
||||
raw, err := io.ReadAll(httpStub.lastReq.Body)
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, string(raw), `"ttl":"1800s"`)
|
||||
|
||||
resource, err := geminiCachedContentResourceURL(account, geminiCachedContentAPIBase(account), upstream.Name)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "https://upstream.example.com/gw/v1beta/cachedContents/gatewayid0001", resource)
|
||||
}
|
||||
@@ -399,19 +399,11 @@ func (s *GeminiMessagesCompatService) buildGeminiChatCompletionsUpstreamRequestF
|
||||
|
||||
case AccountTypeServiceAccount:
|
||||
return func(ctx context.Context) (*http.Request, string, error) {
|
||||
if s.tokenProvider == nil {
|
||||
return nil, "", errors.New("gemini token provider not configured")
|
||||
}
|
||||
accessToken, err := s.tokenProvider.GetAccessToken(ctx, account)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
|
||||
action := "generateContent"
|
||||
if clientStream {
|
||||
action = "streamGenerateContent"
|
||||
}
|
||||
fullURL, err := buildVertexGeminiURL(account.VertexProjectID(), account.VertexLocation(mappedModel), mappedModel, action, clientStream)
|
||||
fullURL, err := buildVertexGeminiAccountURL(account, mappedModel, action, clientStream)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
@@ -422,7 +414,9 @@ func (s *GeminiMessagesCompatService) buildGeminiChatCompletionsUpstreamRequestF
|
||||
return nil, "", err
|
||||
}
|
||||
upstreamReq.Header.Set("Content-Type", "application/json")
|
||||
upstreamReq.Header.Set("Authorization", "Bearer "+accessToken)
|
||||
if err := setVertexGeminiAuth(ctx, upstreamReq, account, s.tokenProvider); err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
return upstreamReq, "x-request-id", nil
|
||||
}, "x-request-id"
|
||||
|
||||
|
||||
@@ -677,6 +677,72 @@ func TestHandleGeminiUpstreamError_APIKeyFallbackStillPSTMidnight(t *testing.T)
|
||||
require.WithinDuration(t, expected, repo.lastRateLimitReset, 5*time.Second)
|
||||
}
|
||||
|
||||
// 第三方上游 API Key 的 429 不套用 AI Studio 每日配额语义:只认显式重试时长,
|
||||
// 否则走全局 429 兜底冷却(未配置时默认 5 秒)。
|
||||
func TestHandleGeminiUpstreamError_ThirdPartyAPIKey429(t *testing.T) {
|
||||
const relayBaseURL = "https://relay.example.com"
|
||||
tests := []struct {
|
||||
name string
|
||||
baseURL string
|
||||
body string
|
||||
expectWait time.Duration
|
||||
pstNight bool
|
||||
}{
|
||||
{
|
||||
name: "no_reset_time_uses_global_fallback",
|
||||
baseURL: relayBaseURL,
|
||||
body: `{"error":{"code":429,"message":"You have exhausted your capacity on this model. Your quota will reset after 6h53m10s."}}`,
|
||||
expectWait: defaultRateLimit429CooldownSeconds * time.Second,
|
||||
},
|
||||
{
|
||||
name: "daily_quota_message_not_treated_as_pst_midnight",
|
||||
baseURL: relayBaseURL,
|
||||
body: `{"error":{"code":429,"message":"Quota exceeded for metric: generate_content_free_tier_requests, limit: 50 per day"}}`,
|
||||
expectWait: defaultRateLimit429CooldownSeconds * time.Second,
|
||||
},
|
||||
{
|
||||
name: "explicit_quota_reset_delay_respected",
|
||||
baseURL: relayBaseURL,
|
||||
body: `{"error":{"code":429,"message":"rate limit","details":[{"metadata":{"quotaResetDelay":"42s"}}]}}`,
|
||||
expectWait: 42 * time.Second,
|
||||
},
|
||||
{
|
||||
name: "official_base_url_keeps_pst_midnight",
|
||||
baseURL: "https://generativelanguage.googleapis.com/",
|
||||
body: `{"error":{"code":429,"message":"rate limit"}}`,
|
||||
pstNight: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
repo := &rateLimit429AccountRepoStub{}
|
||||
svc := &GeminiMessagesCompatService{
|
||||
accountRepo: repo,
|
||||
rateLimitService: NewRateLimitService(repo, nil, &config.Config{}, nil, nil),
|
||||
}
|
||||
account := &Account{
|
||||
ID: 700,
|
||||
Platform: PlatformGemini,
|
||||
Type: AccountTypeAPIKey,
|
||||
Credentials: map[string]any{"base_url": tt.baseURL},
|
||||
}
|
||||
|
||||
before := time.Now()
|
||||
svc.handleGeminiUpstreamError(context.Background(), account, http.StatusTooManyRequests, http.Header{}, []byte(tt.body))
|
||||
|
||||
require.Equal(t, 1, repo.rateLimitCalls)
|
||||
require.Equal(t, int64(700), repo.lastRateLimitID)
|
||||
if tt.pstNight {
|
||||
expected := time.Unix(*nextGeminiDailyResetUnix(), 0)
|
||||
require.WithinDuration(t, expected, repo.lastRateLimitReset, 5*time.Second)
|
||||
return
|
||||
}
|
||||
require.WithinDuration(t, before.Add(tt.expectWait), repo.lastRateLimitReset, 2*time.Second)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
type geminiErrorPolicyRepo struct {
|
||||
mockAccountRepoForGemini
|
||||
setErrorCalls int
|
||||
|
||||
@@ -529,94 +529,6 @@ func (s *GeminiMessagesCompatService) HasAntigravityAccounts(ctx context.Context
|
||||
return len(accounts) > 0, nil
|
||||
}
|
||||
|
||||
// SelectAccountForAIStudioEndpoints selects an account that is likely to succeed against
|
||||
// generativelanguage.googleapis.com (e.g. GET /v1beta/models).
|
||||
//
|
||||
// Preference order:
|
||||
// 1) API key accounts (AI Studio)
|
||||
// 2) OAuth accounts without project_id (AI Studio OAuth)
|
||||
// 3) OAuth accounts explicitly marked as ai_studio
|
||||
// 4) Any remaining Gemini accounts (fallback)
|
||||
func (s *GeminiMessagesCompatService) SelectAccountForAIStudioEndpoints(ctx context.Context, groupID *int64) (*Account, error) {
|
||||
accounts, err := s.listSchedulableAccountsOnce(ctx, groupID, PlatformGemini, true)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("query accounts failed: %w", err)
|
||||
}
|
||||
if len(accounts) == 0 {
|
||||
return nil, errors.New("no available Gemini accounts")
|
||||
}
|
||||
|
||||
rank := func(a *Account) int {
|
||||
if a == nil {
|
||||
return 999
|
||||
}
|
||||
switch a.Type {
|
||||
case AccountTypeAPIKey:
|
||||
if strings.TrimSpace(a.GetCredential("api_key")) != "" {
|
||||
return 0
|
||||
}
|
||||
return 9
|
||||
case AccountTypeOAuth:
|
||||
if strings.TrimSpace(a.GetCredential("project_id")) == "" {
|
||||
return 1
|
||||
}
|
||||
if strings.TrimSpace(a.GetCredential("oauth_type")) == "ai_studio" {
|
||||
return 2
|
||||
}
|
||||
// Code Assist OAuth tokens often lack AI Studio scopes for models listing.
|
||||
return 3
|
||||
case AccountTypeServiceAccount:
|
||||
// Vertex service accounts use aiplatform.googleapis.com, not the AI Studio
|
||||
// endpoint (generativelanguage.googleapis.com), so they cannot serve these requests.
|
||||
return 999
|
||||
default:
|
||||
return 10
|
||||
}
|
||||
}
|
||||
|
||||
var selected *Account
|
||||
for i := range accounts {
|
||||
acc := &accounts[i]
|
||||
if selected == nil {
|
||||
selected = acc
|
||||
continue
|
||||
}
|
||||
|
||||
r1, r2 := rank(acc), rank(selected)
|
||||
if r1 < r2 {
|
||||
selected = acc
|
||||
continue
|
||||
}
|
||||
if r1 > r2 {
|
||||
continue
|
||||
}
|
||||
|
||||
if acc.Priority < selected.Priority {
|
||||
selected = acc
|
||||
} else if acc.Priority == selected.Priority {
|
||||
switch {
|
||||
case acc.LastUsedAt == nil && selected.LastUsedAt != nil:
|
||||
selected = acc
|
||||
case acc.LastUsedAt != nil && selected.LastUsedAt == nil:
|
||||
// keep selected
|
||||
case acc.LastUsedAt == nil && selected.LastUsedAt == nil:
|
||||
if acc.Type == AccountTypeOAuth && selected.Type != AccountTypeOAuth {
|
||||
selected = acc
|
||||
}
|
||||
default:
|
||||
if acc.LastUsedAt.Before(*selected.LastUsedAt) {
|
||||
selected = acc
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if selected == nil {
|
||||
return nil, errors.New("no available Gemini accounts")
|
||||
}
|
||||
return s.hydrateSelectedAccount(ctx, selected)
|
||||
}
|
||||
|
||||
func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Context, account *Account, body []byte) (*ForwardResult, error) {
|
||||
beginUpstreamResponseModelObservation(c)
|
||||
beginGeminiImageOutputObservation(c)
|
||||
@@ -770,19 +682,11 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex
|
||||
|
||||
case AccountTypeServiceAccount:
|
||||
buildReq = func(ctx context.Context) (*http.Request, string, error) {
|
||||
if s.tokenProvider == nil {
|
||||
return nil, "", errors.New("gemini token provider not configured")
|
||||
}
|
||||
accessToken, err := s.tokenProvider.GetAccessToken(ctx, account)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
|
||||
action := "generateContent"
|
||||
if req.Stream {
|
||||
action = "streamGenerateContent"
|
||||
}
|
||||
fullURL, err := buildVertexGeminiURL(account.VertexProjectID(), account.VertexLocation(mappedModel), mappedModel, action, req.Stream)
|
||||
fullURL, err := buildVertexGeminiAccountURL(account, mappedModel, action, req.Stream)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
@@ -793,7 +697,9 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex
|
||||
return nil, "", err
|
||||
}
|
||||
upstreamReq.Header.Set("Content-Type", "application/json")
|
||||
upstreamReq.Header.Set("Authorization", "Bearer "+accessToken)
|
||||
if err := setVertexGeminiAuth(ctx, upstreamReq, account, s.tokenProvider); err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
return upstreamReq, "x-request-id", nil
|
||||
}
|
||||
requestIDHeader = "x-request-id"
|
||||
@@ -1311,15 +1217,7 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
|
||||
|
||||
case AccountTypeServiceAccount:
|
||||
buildReq = func(ctx context.Context) (*http.Request, string, error) {
|
||||
if s.tokenProvider == nil {
|
||||
return nil, "", errors.New("gemini token provider not configured")
|
||||
}
|
||||
accessToken, err := s.tokenProvider.GetAccessToken(ctx, account)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
|
||||
fullURL, err := buildVertexGeminiURL(account.VertexProjectID(), account.VertexLocation(mappedModel), mappedModel, upstreamAction, useUpstreamStream)
|
||||
fullURL, err := buildVertexGeminiAccountURL(account, mappedModel, upstreamAction, useUpstreamStream)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
@@ -1329,7 +1227,9 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
|
||||
return nil, "", err
|
||||
}
|
||||
upstreamReq.Header.Set("Content-Type", "application/json")
|
||||
upstreamReq.Header.Set("Authorization", "Bearer "+accessToken)
|
||||
if err := setVertexGeminiAuth(ctx, upstreamReq, account, s.tokenProvider); err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
return upstreamReq, "x-request-id", nil
|
||||
}
|
||||
requestIDHeader = "x-request-id"
|
||||
@@ -1374,6 +1274,15 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
|
||||
return nil, transportErr
|
||||
}
|
||||
|
||||
// 引用缓存不存在、上游网关持有缓存的账号不可用先于错误策略识别:错误策略会写临时不可调度,
|
||||
// 重试与换号也无法恢复。
|
||||
if passthrough, rebuilt := s.cachedContentPassthroughInLoop(ctx, account, resp); passthrough {
|
||||
resp = rebuilt
|
||||
break
|
||||
} else {
|
||||
resp = rebuilt
|
||||
}
|
||||
|
||||
// 错误策略优先:匹配则跳过重试直接处理。
|
||||
if matched, rebuilt := s.checkErrorPolicyInLoop(ctx, account, resp, mappedModel); matched {
|
||||
resp = rebuilt
|
||||
@@ -1467,6 +1376,14 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
|
||||
|
||||
if resp.StatusCode >= 400 {
|
||||
respBody := s.readUpstreamErrorBody(resp)
|
||||
// 引用缓存不存在是请求问题:统一返回 403,不改账号状态、不换号。
|
||||
if isGeminiCachedContentNotFound(resp.StatusCode, respBody) {
|
||||
return nil, s.writeGeminiCachedContentNotFound(c, account, resp, respBody, requestID, isOAuth)
|
||||
}
|
||||
// 上游网关的限流 / 不可用只针对其持有缓存的账号:不改本账号状态,交由绑定请求的耗尽处理透传。
|
||||
if isGeminiBoundCacheUpstreamPassthrough(ctx, account, resp.StatusCode) {
|
||||
return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: respBody, BoundUpstreamPassthrough: true}
|
||||
}
|
||||
// Best-effort fallback for OAuth tokens missing AI Studio scopes when calling countTokens.
|
||||
// This avoids Gemini SDKs failing hard during preflight token counting.
|
||||
// Checked before error policy so it always works regardless of custom error codes.
|
||||
@@ -1676,6 +1593,27 @@ func (s *GeminiMessagesCompatService) checkErrorPolicyInLoop(
|
||||
return policy != ErrorPolicyNone, rebuilt
|
||||
}
|
||||
|
||||
// cachedContentPassthroughInLoop 判断上游响应是否跳过同号重试与错误策略直接处理:引用缓存不存在,
|
||||
// 或绑定缓存的请求在上游网关账号上收到的 429 / 5xx。读取过的响应体以内存副本重建。
|
||||
func (s *GeminiMessagesCompatService) cachedContentPassthroughInLoop(ctx context.Context, account *Account, resp *http.Response) (passthrough bool, rebuilt *http.Response) {
|
||||
if isGeminiBoundCacheUpstreamPassthrough(ctx, account, resp.StatusCode) {
|
||||
return true, resp
|
||||
}
|
||||
switch resp.StatusCode {
|
||||
case http.StatusBadRequest, http.StatusForbidden, http.StatusNotFound:
|
||||
default:
|
||||
return false, resp
|
||||
}
|
||||
body := s.readUpstreamErrorBody(resp)
|
||||
_ = resp.Body.Close()
|
||||
rebuilt = &http.Response{
|
||||
StatusCode: resp.StatusCode,
|
||||
Header: resp.Header.Clone(),
|
||||
Body: io.NopCloser(bytes.NewReader(body)),
|
||||
}
|
||||
return isGeminiCachedContentNotFound(resp.StatusCode, body), rebuilt
|
||||
}
|
||||
|
||||
func (s *GeminiMessagesCompatService) shouldRetryGeminiUpstreamError(account *Account, statusCode int) bool {
|
||||
switch statusCode {
|
||||
case 429, 500, 502, 503, 504, 529:
|
||||
@@ -1781,25 +1719,7 @@ func (s *GeminiMessagesCompatService) writeGeminiCustomCodeSkippedError(c *gin.C
|
||||
// 并记录 ops 错误事件。状态码保真:下游据此区分请求级错误与可重试的链路故障。
|
||||
func (s *GeminiMessagesCompatService) writeGeminiNativeUpstreamError(c *gin.Context, account *Account, resp *http.Response, respBody []byte, requestID string, isOAuth bool) error {
|
||||
respBody = unwrapIfNeeded(isOAuth, respBody)
|
||||
upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody))
|
||||
upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
|
||||
upstreamDetail := s.upstreamErrorDetail(respBody)
|
||||
if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody {
|
||||
logger.LegacyPrintf("service.gemini_messages_compat", "[Gemini] native upstream error %d: %s", resp.StatusCode, truncateForLog(respBody, s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes))
|
||||
}
|
||||
setOpsUpstreamError(c, resp.StatusCode, upstreamMsg, upstreamDetail)
|
||||
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
|
||||
ProxyID: opsUpstreamProxyID(account),
|
||||
ProxyName: opsUpstreamProxyName(account),
|
||||
Platform: account.Platform,
|
||||
AccountID: account.ID,
|
||||
AccountName: account.Name,
|
||||
UpstreamStatusCode: resp.StatusCode,
|
||||
UpstreamRequestID: requestID,
|
||||
Kind: "http_error",
|
||||
Message: upstreamMsg,
|
||||
Detail: upstreamDetail,
|
||||
})
|
||||
upstreamMsg := s.recordGeminiNativeUpstreamError(c, account, resp.StatusCode, respBody, requestID)
|
||||
|
||||
contentType := resp.Header.Get("Content-Type")
|
||||
if contentType == "" {
|
||||
@@ -1813,6 +1733,41 @@ func (s *GeminiMessagesCompatService) writeGeminiNativeUpstreamError(c *gin.Cont
|
||||
return fmt.Errorf("gemini upstream error: %d message=%s", resp.StatusCode, upstreamMsg)
|
||||
}
|
||||
|
||||
// writeGeminiCachedContentNotFound 记录上游原始错误,并向客户端统一返回缓存不存在的 403:
|
||||
// 各上游的响应形态不同,且 Vertex 的文案含上游缓存资源 ID。失败由请求引用的缓存决定,ops 按客户端请求错误统计。
|
||||
func (s *GeminiMessagesCompatService) writeGeminiCachedContentNotFound(c *gin.Context, account *Account, resp *http.Response, respBody []byte, requestID string, isOAuth bool) error {
|
||||
respBody = unwrapIfNeeded(isOAuth, respBody)
|
||||
upstreamMsg := s.recordGeminiNativeUpstreamError(c, account, resp.StatusCode, respBody, requestID)
|
||||
MarkOpsRequestScopedError(c, "permission_error")
|
||||
MarkResponseCommitted(c)
|
||||
c.Data(http.StatusForbidden, "application/json", []byte(GeminiCachedContentNotFoundResponse))
|
||||
return fmt.Errorf("gemini upstream error: %d message=%s", resp.StatusCode, upstreamMsg)
|
||||
}
|
||||
|
||||
// recordGeminiNativeUpstreamError 把上游错误响应登记到 ops 上下文,返回脱敏后的上游错误文案。
|
||||
func (s *GeminiMessagesCompatService) recordGeminiNativeUpstreamError(c *gin.Context, account *Account, statusCode int, respBody []byte, requestID string) string {
|
||||
upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody))
|
||||
upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
|
||||
upstreamDetail := s.upstreamErrorDetail(respBody)
|
||||
if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody {
|
||||
logger.LegacyPrintf("service.gemini_messages_compat", "[Gemini] native upstream error %d: %s", statusCode, truncateForLog(respBody, s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes))
|
||||
}
|
||||
setOpsUpstreamError(c, statusCode, upstreamMsg, upstreamDetail)
|
||||
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
|
||||
ProxyID: opsUpstreamProxyID(account),
|
||||
ProxyName: opsUpstreamProxyName(account),
|
||||
Platform: account.Platform,
|
||||
AccountID: account.ID,
|
||||
AccountName: account.Name,
|
||||
UpstreamStatusCode: statusCode,
|
||||
UpstreamRequestID: requestID,
|
||||
Kind: "http_error",
|
||||
Message: upstreamMsg,
|
||||
Detail: upstreamDetail,
|
||||
})
|
||||
return upstreamMsg
|
||||
}
|
||||
|
||||
func sleepGeminiBackoff(attempt int) {
|
||||
delay := geminiRetryBaseDelay * time.Duration(1<<uint(attempt-1))
|
||||
if delay > geminiRetryMaxDelay {
|
||||
@@ -2635,6 +2590,25 @@ func isGeminiInsufficientScope(headers http.Header, body []byte) bool {
|
||||
return strings.Contains(lower, "insufficient authentication scopes") || strings.Contains(lower, "access_token_scope_insufficient")
|
||||
}
|
||||
|
||||
// isGeminiCachedContentNotFound 识别请求引用的显式缓存不存在(含已过期、已删除、不属于当前项目)时上游的响应,
|
||||
// 这类响应由请求的 cachedContent 决定,与账号凭据状态无关:
|
||||
// - AI Studio 及本系统网关:一律 403 "CachedContent not found (or permission denied)";
|
||||
// - Vertex 生成请求:400 "Invalid resource state for cache content {id}.";
|
||||
// - Vertex 缓存资源 GET / PATCH / DELETE:404 "Cached content {id} is not found."。
|
||||
func isGeminiCachedContentNotFound(statusCode int, body []byte) bool {
|
||||
lower := strings.ToLower(string(body))
|
||||
switch statusCode {
|
||||
case http.StatusForbidden:
|
||||
return strings.Contains(lower, "cachedcontent not found")
|
||||
case http.StatusBadRequest:
|
||||
return strings.Contains(lower, "invalid resource state for cache content")
|
||||
case http.StatusNotFound:
|
||||
return strings.Contains(lower, "cached content") && strings.Contains(lower, "not found")
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func estimateGeminiCountTokens(reqBody []byte) int {
|
||||
total := 0
|
||||
|
||||
@@ -2854,78 +2828,6 @@ func (s *GeminiMessagesCompatService) handleNativeStreamingResponse(c *gin.Conte
|
||||
return &geminiNativeStreamResult{usage: usage, firstTokenMs: firstTokenMs}, nil
|
||||
}
|
||||
|
||||
// ForwardAIStudioGET forwards a GET request to AI Studio (generativelanguage.googleapis.com) for
|
||||
// endpoints like /v1beta/models and /v1beta/models/{model}.
|
||||
//
|
||||
// This is used to support Gemini SDKs that call models listing endpoints before generation.
|
||||
func (s *GeminiMessagesCompatService) ForwardAIStudioGET(ctx context.Context, account *Account, path string) (*UpstreamHTTPResult, error) {
|
||||
if account == nil {
|
||||
return nil, errors.New("account is nil")
|
||||
}
|
||||
// path 会被直接拼到上游 base URL 后面,因此按路径护栏逐片段校验,
|
||||
// 见 upstream_path_guard.go。
|
||||
sanitizedPath, ok := sanitizedUpstreamPathSuffix(path)
|
||||
if !ok || sanitizedPath == "" {
|
||||
return nil, errors.New("invalid path")
|
||||
}
|
||||
path = sanitizedPath
|
||||
|
||||
baseURL := account.GetGeminiBaseURL(geminicli.AIStudioBaseURL)
|
||||
normalizedBaseURL, err := s.validateUpstreamBaseURL(baseURL)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
fullURL := strings.TrimRight(normalizedBaseURL, "/") + path
|
||||
|
||||
var proxyURL string
|
||||
if account.ProxyID != nil && account.Proxy != nil {
|
||||
proxyURL = account.Proxy.URL()
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, fullURL, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
switch account.Type {
|
||||
case AccountTypeAPIKey:
|
||||
apiKey := strings.TrimSpace(account.GetCredential("api_key"))
|
||||
if apiKey == "" {
|
||||
return nil, errors.New("gemini api_key not configured")
|
||||
}
|
||||
req.Header.Set("x-goog-api-key", apiKey)
|
||||
case AccountTypeOAuth:
|
||||
if s.tokenProvider == nil {
|
||||
return nil, errors.New("gemini token provider not configured")
|
||||
}
|
||||
accessToken, err := s.tokenProvider.GetAccessToken(ctx, account)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Authorization", "Bearer "+accessToken)
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported account type: %s", account.Type)
|
||||
}
|
||||
|
||||
resp, err := s.httpUpstream.Do(req, proxyURL, account.ID, account.Concurrency)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
body, _ := io.ReadAll(io.LimitReader(resp.Body, 8<<20))
|
||||
wwwAuthenticate := resp.Header.Get("Www-Authenticate")
|
||||
filteredHeaders := responseheaders.FilterHeaders(resp.Header, s.responseHeaderFilter)
|
||||
if wwwAuthenticate != "" {
|
||||
filteredHeaders.Set("Www-Authenticate", wwwAuthenticate)
|
||||
}
|
||||
return &UpstreamHTTPResult{
|
||||
StatusCode: resp.StatusCode,
|
||||
Headers: filteredHeaders,
|
||||
Body: body,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// unwrapGeminiResponse 解包 Gemini OAuth 响应中的 response 字段
|
||||
// 使用 gjson 零拷贝提取,避免完整 Unmarshal+Marshal
|
||||
func unwrapGeminiResponse(raw []byte) ([]byte, error) {
|
||||
@@ -3097,6 +2999,10 @@ func (s *GeminiMessagesCompatService) handleGeminiUpstreamError(ctx context.Cont
|
||||
if account.IsPoolMode() && !account.IsCustomErrorCodesEnabled() {
|
||||
return
|
||||
}
|
||||
if account.IsGeminiThirdPartyAPIKey() {
|
||||
s.handleGeminiThirdParty429(ctx, account, body)
|
||||
return
|
||||
}
|
||||
|
||||
oauthType := account.GeminiOAuthType()
|
||||
tierID := account.GeminiTierID()
|
||||
@@ -3145,16 +3051,36 @@ func (s *GeminiMessagesCompatService) handleGeminiUpstreamError(ctx context.Cont
|
||||
account.ID, resetTime, oauthType, tierID)
|
||||
}
|
||||
|
||||
// handleGeminiThirdParty429 处理第三方上游的 429:不套用 AI Studio 每日配额语义,
|
||||
// 只认响应里显式给出的重试时间,否则走全局 429 兜底冷却。
|
||||
func (s *GeminiMessagesCompatService) handleGeminiThirdParty429(ctx context.Context, account *Account, body []byte) {
|
||||
if resetAt := parseGeminiRetryDelayResetTime(body); resetAt != nil {
|
||||
resetTime := time.Unix(*resetAt, 0)
|
||||
_ = s.accountRepo.SetRateLimited(ctx, account.ID, resetTime)
|
||||
logger.LegacyPrintf("service.gemini_messages_compat", "[Gemini 429] Account %d (third-party API Key) rate limited until %v", account.ID, resetTime)
|
||||
return
|
||||
}
|
||||
if s.rateLimitService != nil {
|
||||
s.rateLimitService.apply429FallbackRateLimit(ctx, account, "gemini_third_party_no_reset_time")
|
||||
return
|
||||
}
|
||||
_ = s.accountRepo.SetRateLimited(ctx, account.ID, time.Now().Add(defaultRateLimit429CooldownSeconds*time.Second))
|
||||
}
|
||||
|
||||
// ParseGeminiRateLimitResetTime 解析 Gemini 格式的 429 响应,返回重置时间的 Unix 时间戳
|
||||
func ParseGeminiRateLimitResetTime(body []byte) *int64 {
|
||||
// 第一阶段:gjson 结构化提取
|
||||
errMsg := gjson.GetBytes(body, "error.message").String()
|
||||
if looksLikeGeminiDailyQuota(errMsg) {
|
||||
if ts := nextGeminiDailyResetUnix(); ts != nil {
|
||||
return ts
|
||||
}
|
||||
}
|
||||
return parseGeminiRetryDelayResetTime(body)
|
||||
}
|
||||
|
||||
// parseGeminiRetryDelayResetTime 只解析响应里显式给出的重试时长,返回重置时间的 Unix 时间戳。
|
||||
func parseGeminiRetryDelayResetTime(body []byte) *int64 {
|
||||
// 第一阶段:gjson 结构化提取
|
||||
// 遍历 error.details 查找 quotaResetDelay(AI Studio)或 RetryInfo.retryDelay(Vertex)
|
||||
var found *int64
|
||||
gjson.GetBytes(body, "error.details").ForEach(func(_, detail gjson.Result) bool {
|
||||
|
||||
@@ -0,0 +1,302 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"slices"
|
||||
"strings"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/gemini"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/geminicli"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/googleapi"
|
||||
"github.com/Wei-Shaw/sub2api/internal/util/responseheaders"
|
||||
)
|
||||
|
||||
// Gemini 分组对外提供 AI Studio 兼容协议(GET /v1beta/models、/v1beta/models/{model}),
|
||||
// 分组内的账号可以是 AI Studio 兼容上游(API Key / OAuth),也可以是 Agent Platform (Vertex)。
|
||||
// 本文件把不同账号类型的模型接口统一成 AI Studio 格式的响应。
|
||||
|
||||
// errGeminiModelsUpstreamStatus 在 Vertex 分页读取遇到非 2xx 时中止分页,原始响应由调用方另行保存。
|
||||
var errGeminiModelsUpstreamStatus = errors.New("upstream model list returned non-2xx status")
|
||||
|
||||
// SelectAccountForModelsEndpoints 为模型列表与模型详情接口选择账号。
|
||||
//
|
||||
// 优先级:
|
||||
// 1. 配置了 api_key 的 API Key 账号
|
||||
// 2. 不带 project_id 的 OAuth 账号(AI Studio OAuth)
|
||||
// 3. 显式标记为 ai_studio 的 OAuth 账号
|
||||
// 4. Code Assist OAuth 账号(令牌常缺 AI Studio 权限范围,多半退回内置列表)
|
||||
// 5. 其余类型的账号
|
||||
// 6. Vertex 服务账号(服务账号 JSON 认证,读取 Google 发布方模型列表)
|
||||
// 7. Vertex API Key 账号(列表接口不接受 API Key,只能返回内置列表)
|
||||
//
|
||||
// Vertex 账号排在所有 AI Studio 兼容账号之后,混合分组始终由 AI Studio 兼容账号应答。
|
||||
// 同级账号按 priority、最近使用时间挑选。
|
||||
func (s *GeminiMessagesCompatService) SelectAccountForModelsEndpoints(ctx context.Context, groupID *int64) (*Account, error) {
|
||||
accounts, err := s.listSchedulableAccountsOnce(ctx, groupID, PlatformGemini, true)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("query accounts failed: %w", err)
|
||||
}
|
||||
if len(accounts) == 0 {
|
||||
return nil, errors.New("no available Gemini accounts")
|
||||
}
|
||||
|
||||
var selected *Account
|
||||
for i := range accounts {
|
||||
acc := &accounts[i]
|
||||
if selected == nil {
|
||||
selected = acc
|
||||
continue
|
||||
}
|
||||
|
||||
r1, r2 := modelsEndpointAccountRank(acc), modelsEndpointAccountRank(selected)
|
||||
if r1 < r2 {
|
||||
selected = acc
|
||||
continue
|
||||
}
|
||||
if r1 > r2 {
|
||||
continue
|
||||
}
|
||||
|
||||
if acc.Priority < selected.Priority {
|
||||
selected = acc
|
||||
} else if acc.Priority == selected.Priority {
|
||||
switch {
|
||||
case acc.LastUsedAt == nil && selected.LastUsedAt != nil:
|
||||
selected = acc
|
||||
case acc.LastUsedAt != nil && selected.LastUsedAt == nil:
|
||||
// keep selected
|
||||
case acc.LastUsedAt == nil && selected.LastUsedAt == nil:
|
||||
if acc.Type == AccountTypeOAuth && selected.Type != AccountTypeOAuth {
|
||||
selected = acc
|
||||
}
|
||||
default:
|
||||
if acc.LastUsedAt.Before(*selected.LastUsedAt) {
|
||||
selected = acc
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if selected == nil {
|
||||
return nil, errors.New("no available Gemini accounts")
|
||||
}
|
||||
return s.hydrateSelectedAccount(ctx, selected)
|
||||
}
|
||||
|
||||
// modelsEndpointAccountRank 返回账号在模型接口选号中的优先级,数值越小越优先,顺序见 SelectAccountForModelsEndpoints。
|
||||
func modelsEndpointAccountRank(a *Account) int {
|
||||
if a == nil {
|
||||
return 999
|
||||
}
|
||||
switch a.Type {
|
||||
case AccountTypeAPIKey:
|
||||
if strings.TrimSpace(a.GetCredential("api_key")) != "" {
|
||||
return 0
|
||||
}
|
||||
return 9
|
||||
case AccountTypeOAuth:
|
||||
if strings.TrimSpace(a.GetCredential("project_id")) == "" {
|
||||
return 1
|
||||
}
|
||||
if strings.TrimSpace(a.GetCredential("oauth_type")) == "ai_studio" {
|
||||
return 2
|
||||
}
|
||||
return 3
|
||||
case AccountTypeServiceAccount:
|
||||
if a.IsVertexAPIKey() {
|
||||
return 12
|
||||
}
|
||||
return 11
|
||||
default:
|
||||
return 10
|
||||
}
|
||||
}
|
||||
|
||||
// ForwardListModels 读取账号上游的模型列表,返回 AI Studio 格式的响应(GET /v1beta/models)。
|
||||
// AI Studio 兼容账号原样返回上游响应;Vertex 服务账号读取 Google 发布方模型列表后转换格式,
|
||||
// 上游非 2xx 时原样返回该响应;Vertex API Key 账号返回内置列表。
|
||||
func (s *GeminiMessagesCompatService) ForwardListModels(ctx context.Context, account *Account) (*UpstreamHTTPResult, error) {
|
||||
if account == nil {
|
||||
return nil, errors.New("account is nil")
|
||||
}
|
||||
if !account.IsVertexServiceAccount() {
|
||||
return s.forwardGeminiModelsGET(ctx, account, "/v1beta/models")
|
||||
}
|
||||
if account.IsVertexAPIKey() {
|
||||
return geminiModelsJSONResult(http.StatusOK, gemini.FallbackModelsList())
|
||||
}
|
||||
|
||||
modelIDs, failure, err := s.listVertexGeminiModels(ctx, account, account.VertexLocation(""))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if failure != nil {
|
||||
return failure, nil
|
||||
}
|
||||
models := make([]gemini.Model, 0, len(modelIDs))
|
||||
for _, modelID := range modelIDs {
|
||||
models = append(models, gemini.FallbackModel(modelID))
|
||||
}
|
||||
return geminiModelsJSONResult(http.StatusOK, gemini.ModelsListResponse{Models: models})
|
||||
}
|
||||
|
||||
// ForwardGetModel 读取账号上游的单个模型信息,返回 AI Studio 格式的响应(GET /v1beta/models/{model})。
|
||||
// Vertex 服务账号在该模型生成请求所用区域的模型列表中查找,找不到时返回 404;Vertex API Key 账号返回内置信息。
|
||||
func (s *GeminiMessagesCompatService) ForwardGetModel(ctx context.Context, account *Account, model string) (*UpstreamHTTPResult, error) {
|
||||
if account == nil {
|
||||
return nil, errors.New("account is nil")
|
||||
}
|
||||
if !account.IsVertexServiceAccount() {
|
||||
return s.forwardGeminiModelsGET(ctx, account, "/v1beta/models/"+model)
|
||||
}
|
||||
if account.IsVertexAPIKey() {
|
||||
return geminiModelsJSONResult(http.StatusOK, gemini.FallbackModel(model))
|
||||
}
|
||||
|
||||
modelIDs, failure, err := s.listVertexGeminiModels(ctx, account, account.VertexLocation(model))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if failure != nil {
|
||||
return failure, nil
|
||||
}
|
||||
if slices.Contains(modelIDs, model) {
|
||||
return geminiModelsJSONResult(http.StatusOK, gemini.FallbackModel(model))
|
||||
}
|
||||
return geminiModelsJSONResult(http.StatusNotFound, map[string]any{
|
||||
"error": map[string]any{
|
||||
"code": http.StatusNotFound,
|
||||
"message": fmt.Sprintf("models/%s is not found", model),
|
||||
"status": googleapi.HTTPStatusToGoogleStatus(http.StatusNotFound),
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// forwardGeminiModelsGET 向 AI Studio 兼容上游转发模型接口的 GET 请求。
|
||||
// path 会被直接拼到上游 base URL 后面,因此按路径护栏逐片段校验,见 upstream_path_guard.go。
|
||||
func (s *GeminiMessagesCompatService) forwardGeminiModelsGET(ctx context.Context, account *Account, path string) (*UpstreamHTTPResult, error) {
|
||||
sanitizedPath, ok := sanitizedUpstreamPathSuffix(path)
|
||||
if !ok || sanitizedPath == "" {
|
||||
return nil, errors.New("invalid path")
|
||||
}
|
||||
|
||||
baseURL := account.GetGeminiBaseURL(geminicli.AIStudioBaseURL)
|
||||
normalizedBaseURL, err := s.validateUpstreamBaseURL(baseURL)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
fullURL := strings.TrimRight(normalizedBaseURL, "/") + sanitizedPath
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, fullURL, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
switch account.Type {
|
||||
case AccountTypeAPIKey:
|
||||
apiKey := strings.TrimSpace(account.GetCredential("api_key"))
|
||||
if apiKey == "" {
|
||||
return nil, errors.New("gemini api_key not configured")
|
||||
}
|
||||
req.Header.Set("x-goog-api-key", apiKey)
|
||||
case AccountTypeOAuth:
|
||||
if s.tokenProvider == nil {
|
||||
return nil, errors.New("gemini token provider not configured")
|
||||
}
|
||||
accessToken, err := s.tokenProvider.GetAccessToken(ctx, account)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Authorization", "Bearer "+accessToken)
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported account type: %s", account.Type)
|
||||
}
|
||||
|
||||
return s.doGeminiModelsGET(req, account)
|
||||
}
|
||||
|
||||
// listVertexGeminiModels 读取 Vertex 服务账号在指定区域可用的 Gemini 模型。
|
||||
// 上游返回非 2xx 时 failure 为该页的原始响应,modelIDs 与 err 均为空。
|
||||
func (s *GeminiMessagesCompatService) listVertexGeminiModels(ctx context.Context, account *Account, location string) (modelIDs []string, failure *UpstreamHTTPResult, err error) {
|
||||
if s.tokenProvider == nil {
|
||||
return nil, nil, errors.New("gemini token provider not configured")
|
||||
}
|
||||
accessToken, err := s.tokenProvider.GetAccessToken(ctx, account)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
|
||||
modelIDs, err = listVertexGeminiPublisherModels(location, func(_ int, listURL string) ([]byte, error) {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, listURL, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Accept", "application/json")
|
||||
req.Header.Set("Authorization", "Bearer "+accessToken)
|
||||
|
||||
res, err := s.doGeminiModelsGET(req, account)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if res.StatusCode < http.StatusOK || res.StatusCode >= http.StatusMultipleChoices {
|
||||
failure = res
|
||||
return nil, errGeminiModelsUpstreamStatus
|
||||
}
|
||||
return res.Body, nil
|
||||
})
|
||||
if failure != nil {
|
||||
return nil, failure, nil
|
||||
}
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
return modelIDs, nil, nil
|
||||
}
|
||||
|
||||
// doGeminiModelsGET 发送模型接口请求并读取响应。响应体上限取 gateway.models_list_read_max_bytes,超出时报错;响应头按配置过滤。
|
||||
func (s *GeminiMessagesCompatService) doGeminiModelsGET(req *http.Request, account *Account) (*UpstreamHTTPResult, error) {
|
||||
var proxyURL string
|
||||
if account.ProxyID != nil && account.Proxy != nil {
|
||||
proxyURL = account.Proxy.URL()
|
||||
}
|
||||
|
||||
resp, err := s.httpUpstream.Do(req, proxyURL, account.ID, account.Concurrency)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
bodyLimit := resolveModelsListReadLimit(s.cfg)
|
||||
body, _ := io.ReadAll(io.LimitReader(resp.Body, bodyLimit+1))
|
||||
if int64(len(body)) > bodyLimit {
|
||||
return nil, fmt.Errorf("upstream model list response exceeds %d bytes", bodyLimit)
|
||||
}
|
||||
wwwAuthenticate := resp.Header.Get("Www-Authenticate")
|
||||
filteredHeaders := responseheaders.FilterHeaders(resp.Header, s.responseHeaderFilter)
|
||||
if wwwAuthenticate != "" {
|
||||
filteredHeaders.Set("Www-Authenticate", wwwAuthenticate)
|
||||
}
|
||||
return &UpstreamHTTPResult{
|
||||
StatusCode: resp.StatusCode,
|
||||
Headers: filteredHeaders,
|
||||
Body: body,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// geminiModelsJSONResult 把网关生成的模型接口响应包装成 UpstreamHTTPResult,交给 handler 统一写出。
|
||||
func geminiModelsJSONResult(status int, payload any) (*UpstreamHTTPResult, error) {
|
||||
body, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &UpstreamHTTPResult{
|
||||
StatusCode: status,
|
||||
Headers: http.Header{"Content-Type": []string{"application/json; charset=utf-8"}},
|
||||
Body: body,
|
||||
}, nil
|
||||
}
|
||||
@@ -0,0 +1,208 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/gemini"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func newGeminiModelsEndpointsTestService(accounts []Account, upstream HTTPUpstream) *GeminiMessagesCompatService {
|
||||
repo := &mockAccountRepoForGemini{listByGroupFunc: func(context.Context, int64, []string) ([]Account, error) {
|
||||
return accounts, nil
|
||||
}}
|
||||
return &GeminiMessagesCompatService{
|
||||
accountRepo: repo,
|
||||
httpUpstream: upstream,
|
||||
cfg: &config.Config{},
|
||||
tokenProvider: NewGeminiTokenProvider(nil, &vertexModelsTokenCacheStub{token: "ya29.vertex-token"}, nil),
|
||||
}
|
||||
}
|
||||
|
||||
func decodeGeminiModelNames(t *testing.T, body []byte) []string {
|
||||
t.Helper()
|
||||
var got gemini.ModelsListResponse
|
||||
require.NoError(t, json.Unmarshal(body, &got))
|
||||
names := make([]string, 0, len(got.Models))
|
||||
for _, model := range got.Models {
|
||||
require.Contains(t, model.SupportedGenerationMethods, "generateContent")
|
||||
names = append(names, model.Name)
|
||||
}
|
||||
return names
|
||||
}
|
||||
|
||||
func TestSelectAccountForModelsEndpoints_VertexServiceAccountOnlyGroup(t *testing.T) {
|
||||
groupID := int64(8)
|
||||
vertex := *vertexModelSyncTestAccount(nil)
|
||||
vertex.ID = 18
|
||||
svc := newGeminiModelsEndpointsTestService([]Account{vertex}, nil)
|
||||
|
||||
account, err := svc.SelectAccountForModelsEndpoints(context.Background(), &groupID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(18), account.ID)
|
||||
}
|
||||
|
||||
func TestModelsEndpointAccountRank_Order(t *testing.T) {
|
||||
ordered := []*Account{
|
||||
{Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "k"}},
|
||||
{Platform: PlatformGemini, Type: AccountTypeOAuth, Credentials: map[string]any{}},
|
||||
{Platform: PlatformGemini, Type: AccountTypeOAuth, Credentials: map[string]any{"project_id": "p", "oauth_type": "ai_studio"}},
|
||||
{Platform: PlatformGemini, Type: AccountTypeOAuth, Credentials: map[string]any{"project_id": "p", "oauth_type": "code_assist"}},
|
||||
{Platform: PlatformGemini, Type: AccountTypeUpstream, Credentials: map[string]any{"api_key": "k"}},
|
||||
vertexModelSyncTestAccount(nil),
|
||||
{Platform: PlatformGemini, Type: AccountTypeServiceAccount, Credentials: map[string]any{"auth_mode": "apikey", "api_key": "AQ.k"}},
|
||||
}
|
||||
for i := 1; i < len(ordered); i++ {
|
||||
require.Less(t, modelsEndpointAccountRank(ordered[i-1]), modelsEndpointAccountRank(ordered[i]), "position %d", i)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSelectAccountForModelsEndpoints_PrefersAPIKeyOverVertex(t *testing.T) {
|
||||
groupID := int64(8)
|
||||
vertex := *vertexModelSyncTestAccount(nil)
|
||||
vertex.ID = 18
|
||||
apiKey := Account{ID: 19, Platform: PlatformGemini, Type: AccountTypeAPIKey, Priority: 9, Credentials: map[string]any{"api_key": "k"}}
|
||||
svc := newGeminiModelsEndpointsTestService([]Account{vertex, apiKey}, nil)
|
||||
|
||||
account, err := svc.SelectAccountForModelsEndpoints(context.Background(), &groupID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(19), account.ID)
|
||||
}
|
||||
|
||||
func TestForwardListModels_VertexServiceAccountReturnsAIStudioFormat(t *testing.T) {
|
||||
upstream := &httpUpstreamRecorder{responses: []*http.Response{
|
||||
vertexModelsJSONResponse(http.StatusOK, `{"publisherModels":[
|
||||
{"name":"publishers/google/models/gemini-3.8-flash"},
|
||||
{"name":"publishers/google/models/gemma4"}
|
||||
],"nextPageToken":"page-2"}`),
|
||||
vertexModelsJSONResponse(http.StatusOK, `{"publisherModels":[{"name":"publishers/google/models/gemini-2.5-flash"}]}`),
|
||||
}}
|
||||
svc := newGeminiModelsEndpointsTestService(nil, upstream)
|
||||
|
||||
res, err := svc.ForwardListModels(context.Background(), vertexModelSyncTestAccount(map[string]any{"location": "europe-west4"}))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, http.StatusOK, res.StatusCode)
|
||||
require.Equal(t, []string{"models/gemini-2.5-flash", "models/gemini-3.8-flash"}, decodeGeminiModelNames(t, res.Body))
|
||||
|
||||
require.Len(t, upstream.requests, 2)
|
||||
require.Equal(t, "https://europe-west4-aiplatform.googleapis.com/v1beta1/publishers/google/models?pageSize=300", upstream.requests[0].URL.String())
|
||||
require.Equal(t, "https://europe-west4-aiplatform.googleapis.com/v1beta1/publishers/google/models?pageSize=300&pageToken=page-2", upstream.requests[1].URL.String())
|
||||
for _, req := range upstream.requests {
|
||||
require.Equal(t, "Bearer ya29.vertex-token", req.Header.Get("Authorization"))
|
||||
}
|
||||
}
|
||||
|
||||
func TestForwardListModels_VertexUpstreamErrorPassesThrough(t *testing.T) {
|
||||
body := `{"error":{"code":403,"status":"PERMISSION_DENIED"}}`
|
||||
upstream := &httpUpstreamRecorder{resp: vertexModelsJSONResponse(http.StatusForbidden, body)}
|
||||
svc := newGeminiModelsEndpointsTestService(nil, upstream)
|
||||
|
||||
res, err := svc.ForwardListModels(context.Background(), vertexModelSyncTestAccount(map[string]any{"location": "global"}))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, http.StatusForbidden, res.StatusCode)
|
||||
require.JSONEq(t, body, string(res.Body))
|
||||
}
|
||||
|
||||
func TestForwardListModels_VertexAPIKeyReturnsBuiltinList(t *testing.T) {
|
||||
upstream := &httpUpstreamRecorder{}
|
||||
svc := newGeminiModelsEndpointsTestService(nil, upstream)
|
||||
account := &Account{ID: 174, Platform: PlatformGemini, Type: AccountTypeServiceAccount, Credentials: map[string]any{
|
||||
"auth_mode": "apikey",
|
||||
"api_key": "AQ.test-key",
|
||||
}}
|
||||
|
||||
res, err := svc.ForwardListModels(context.Background(), account)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, http.StatusOK, res.StatusCode)
|
||||
expected, err := json.Marshal(gemini.FallbackModelsList())
|
||||
require.NoError(t, err)
|
||||
require.JSONEq(t, string(expected), string(res.Body))
|
||||
require.Empty(t, upstream.requests)
|
||||
}
|
||||
|
||||
func TestForwardListModels_APIKeyForwardsToBaseURL(t *testing.T) {
|
||||
upstream := &httpUpstreamRecorder{resp: vertexModelsJSONResponse(http.StatusOK, `{"models":[{"name":"models/gemini-native"}]}`)}
|
||||
svc := newGeminiModelsEndpointsTestService(nil, upstream)
|
||||
account := &Account{ID: 261, Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: map[string]any{
|
||||
"api_key": "sk-pool",
|
||||
"base_url": "https://relay.example.com",
|
||||
}}
|
||||
|
||||
res, err := svc.ForwardListModels(context.Background(), account)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, http.StatusOK, res.StatusCode)
|
||||
require.JSONEq(t, `{"models":[{"name":"models/gemini-native"}]}`, string(res.Body))
|
||||
require.Equal(t, "https://relay.example.com/v1beta/models", upstream.lastReq.URL.String())
|
||||
require.Equal(t, "sk-pool", upstream.lastReq.Header.Get("x-goog-api-key"))
|
||||
}
|
||||
|
||||
func TestForwardGetModel_VertexServiceAccountLooksUpModelList(t *testing.T) {
|
||||
listBody := `{"publisherModels":[{"name":"publishers/google/models/gemini-3.8-flash"}]}`
|
||||
|
||||
t.Run("listed model", func(t *testing.T) {
|
||||
upstream := &httpUpstreamRecorder{resp: vertexModelsJSONResponse(http.StatusOK, listBody)}
|
||||
svc := newGeminiModelsEndpointsTestService(nil, upstream)
|
||||
|
||||
res, err := svc.ForwardGetModel(context.Background(), vertexModelSyncTestAccount(nil), "gemini-3.8-flash")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, http.StatusOK, res.StatusCode)
|
||||
var model gemini.Model
|
||||
require.NoError(t, json.Unmarshal(res.Body, &model))
|
||||
require.Equal(t, "models/gemini-3.8-flash", model.Name)
|
||||
})
|
||||
|
||||
t.Run("unlisted model", func(t *testing.T) {
|
||||
upstream := &httpUpstreamRecorder{resp: vertexModelsJSONResponse(http.StatusOK, listBody)}
|
||||
svc := newGeminiModelsEndpointsTestService(nil, upstream)
|
||||
|
||||
res, err := svc.ForwardGetModel(context.Background(), vertexModelSyncTestAccount(nil), "gemini-unknown")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, http.StatusNotFound, res.StatusCode)
|
||||
require.JSONEq(t, `{"error":{"code":404,"message":"models/gemini-unknown is not found","status":"NOT_FOUND"}}`, string(res.Body))
|
||||
})
|
||||
}
|
||||
|
||||
func TestForwardGetModel_VertexUsesModelLocation(t *testing.T) {
|
||||
upstream := &httpUpstreamRecorder{resp: vertexModelsJSONResponse(http.StatusOK, `{"publisherModels":[{"name":"publishers/google/models/gemini-3.8-flash"}]}`)}
|
||||
svc := newGeminiModelsEndpointsTestService(nil, upstream)
|
||||
account := vertexModelSyncTestAccount(map[string]any{
|
||||
"location": "europe-west4",
|
||||
"vertex_model_locations": map[string]any{"gemini-3.8-flash": "global"},
|
||||
})
|
||||
|
||||
res, err := svc.ForwardGetModel(context.Background(), account, "gemini-3.8-flash")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, http.StatusOK, res.StatusCode)
|
||||
require.Equal(t, "https://aiplatform.googleapis.com/v1beta1/publishers/google/models?pageSize=300", upstream.lastReq.URL.String())
|
||||
}
|
||||
|
||||
func TestForwardListModels_VertexMidPaginationErrorPassesThrough(t *testing.T) {
|
||||
body := `{"error":{"code":503,"status":"UNAVAILABLE"}}`
|
||||
upstream := &httpUpstreamRecorder{responses: []*http.Response{
|
||||
vertexModelsJSONResponse(http.StatusOK, `{"publisherModels":[{"name":"publishers/google/models/gemini-3.8-flash"}],"nextPageToken":"page-2"}`),
|
||||
vertexModelsJSONResponse(http.StatusServiceUnavailable, body),
|
||||
}}
|
||||
svc := newGeminiModelsEndpointsTestService(nil, upstream)
|
||||
|
||||
res, err := svc.ForwardListModels(context.Background(), vertexModelSyncTestAccount(nil))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, http.StatusServiceUnavailable, res.StatusCode)
|
||||
require.JSONEq(t, body, string(res.Body))
|
||||
require.Len(t, upstream.requests, 2)
|
||||
}
|
||||
|
||||
func TestForwardListModels_ResponseOverReadLimitFails(t *testing.T) {
|
||||
upstream := &httpUpstreamRecorder{resp: vertexModelsJSONResponse(http.StatusOK, `{"models":[{"name":"models/gemini-native"}]}`)}
|
||||
svc := newGeminiModelsEndpointsTestService(nil, upstream)
|
||||
svc.cfg.Gateway.ModelsListReadMaxBytes = 16
|
||||
account := &Account{ID: 261, Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "sk-pool"}}
|
||||
|
||||
_, err := svc.ForwardListModels(context.Background(), account)
|
||||
require.ErrorContains(t, err, "exceeds 16 bytes")
|
||||
}
|
||||
@@ -373,6 +373,10 @@ func geminiQuotaTierKeyForAccount(account *Account) string {
|
||||
if account == nil || account.Platform != PlatformGemini {
|
||||
return ""
|
||||
}
|
||||
// 本地模拟配额只适用于按 Google 账号等级限额的通道;第三方上游与 Vertex 按量付费不套用。
|
||||
if account.IsGeminiThirdPartyAPIKey() || account.IsVertexServiceAccount() {
|
||||
return ""
|
||||
}
|
||||
|
||||
// Note: GeminiOAuthType() already defaults legacy (project_id present) to code_assist.
|
||||
oauthType := strings.ToLower(strings.TrimSpace(account.GeminiOAuthType()))
|
||||
|
||||
@@ -0,0 +1,80 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestAccountIsGeminiThirdPartyAPIKey(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
account *Account
|
||||
expected bool
|
||||
}{
|
||||
{"nil_account", nil, false},
|
||||
{"empty_base_url_is_official", geminiQuotaAPIKeyAccount(""), false},
|
||||
{"official_base_url", geminiQuotaAPIKeyAccount("https://generativelanguage.googleapis.com"), false},
|
||||
{"official_base_url_trailing_slash", geminiQuotaAPIKeyAccount("https://generativelanguage.googleapis.com/"), false},
|
||||
{"official_host_over_http", geminiQuotaAPIKeyAccount("http://generativelanguage.googleapis.com"), true},
|
||||
{"official_host_with_path_prefix", geminiQuotaAPIKeyAccount("https://generativelanguage.googleapis.com/proxy"), true},
|
||||
{"relay_base_url", geminiQuotaAPIKeyAccount("https://relay.example.com"), true},
|
||||
{"oauth_account", &Account{Platform: PlatformGemini, Type: AccountTypeOAuth, Credentials: map[string]any{"base_url": "https://relay.example.com"}}, false},
|
||||
{"vertex_service_account", &Account{Platform: PlatformGemini, Type: AccountTypeServiceAccount}, false},
|
||||
{"non_gemini_platform", &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Credentials: map[string]any{"base_url": "https://relay.example.com"}}, false},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
require.Equal(t, tt.expected, tt.account.IsGeminiThirdPartyAPIKey())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeminiQuotaService_QuotaForAccountSkipsThirdPartyAndVertex(t *testing.T) {
|
||||
svc := NewGeminiQuotaService(&config.Config{}, nil)
|
||||
ctx := context.Background()
|
||||
|
||||
quota, ok := svc.QuotaForAccount(ctx, geminiQuotaAPIKeyAccount(""))
|
||||
require.True(t, ok)
|
||||
require.Equal(t, int64(50), quota.ProRPD)
|
||||
|
||||
paid := geminiQuotaAPIKeyAccount("https://generativelanguage.googleapis.com")
|
||||
paid.Credentials["tier_id"] = GeminiTierAIStudioPaid
|
||||
quota, ok = svc.QuotaForAccount(ctx, paid)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, int64(-1), quota.ProRPD)
|
||||
|
||||
relay := geminiQuotaAPIKeyAccount("https://relay.example.com")
|
||||
relay.Credentials["tier_id"] = GeminiTierAIStudioFree
|
||||
_, ok = svc.QuotaForAccount(ctx, relay)
|
||||
require.False(t, ok, "第三方上游不套用 AI Studio 模拟配额")
|
||||
|
||||
vertex := &Account{
|
||||
Platform: PlatformGemini,
|
||||
Type: AccountTypeServiceAccount,
|
||||
Credentials: map[string]any{"project_id": "my-project", "tier_id": "vertex"},
|
||||
}
|
||||
_, ok = svc.QuotaForAccount(ctx, vertex)
|
||||
require.False(t, ok, "Vertex 按量付费不套用模拟配额")
|
||||
|
||||
googleOne := &Account{
|
||||
Platform: PlatformGemini,
|
||||
Type: AccountTypeOAuth,
|
||||
Credentials: map[string]any{"oauth_type": "google_one", "tier_id": GeminiTierGoogleOneFree},
|
||||
}
|
||||
quota, ok = svc.QuotaForAccount(ctx, googleOne)
|
||||
require.True(t, ok)
|
||||
require.Positive(t, quota.SharedRPD)
|
||||
}
|
||||
|
||||
func geminiQuotaAPIKeyAccount(baseURL string) *Account {
|
||||
credentials := map[string]any{}
|
||||
if baseURL != "" {
|
||||
credentials["base_url"] = baseURL
|
||||
}
|
||||
return &Account{Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: credentials}
|
||||
}
|
||||
@@ -0,0 +1,101 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
const geminiVertexAPIKeyTestResponse = `{"candidates":[{"content":{"parts":[{"text":"hello"}],"role":"model"},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":5,"totalTokenCount":15}}`
|
||||
|
||||
func newGeminiVertexAPIKeyTestService() (*GeminiMessagesCompatService, *geminiCompatHTTPUpstreamStub) {
|
||||
httpStub := &geminiCompatHTTPUpstreamStub{
|
||||
response: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}, "X-Request-Id": []string{"vertex-req-1"}},
|
||||
Body: io.NopCloser(strings.NewReader(geminiVertexAPIKeyTestResponse)),
|
||||
},
|
||||
}
|
||||
// tokenProvider 留空:API Key 认证不依赖 access token。
|
||||
return &GeminiMessagesCompatService{httpUpstream: httpStub, cfg: &config.Config{}}, httpStub
|
||||
}
|
||||
|
||||
func geminiVertexAPIKeyTestAccount(projectID string) *Account {
|
||||
credentials := map[string]any{
|
||||
"auth_mode": "apikey",
|
||||
"api_key": "vertex-test-key",
|
||||
"location": "us-central1",
|
||||
"tier_id": "vertex",
|
||||
}
|
||||
if projectID != "" {
|
||||
credentials["project_id"] = projectID
|
||||
}
|
||||
return &Account{
|
||||
ID: 811,
|
||||
Platform: PlatformGemini,
|
||||
Type: AccountTypeServiceAccount,
|
||||
Credentials: credentials,
|
||||
Concurrency: 1,
|
||||
}
|
||||
}
|
||||
|
||||
func requireVertexAPIKeyAuth(t *testing.T, req *http.Request) {
|
||||
t.Helper()
|
||||
require.NotNil(t, req)
|
||||
require.Equal(t, "vertex-test-key", req.Header.Get("x-goog-api-key"))
|
||||
require.Empty(t, req.Header.Get("Authorization"))
|
||||
}
|
||||
|
||||
func TestGeminiForwardNative_VertexAPIKeyWithoutProjectUsesExpressEndpoint(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc, httpStub := newGeminiVertexAPIKeyTestService()
|
||||
c, _ := newGeminiNativeTestContext(t)
|
||||
|
||||
result, err := svc.ForwardNative(context.Background(), c, geminiVertexAPIKeyTestAccount(""),
|
||||
"gemini-3-flash-preview", "generateContent", false, []byte(`{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`))
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.Equal(t, 1, httpStub.calls)
|
||||
require.Equal(t, "https://aiplatform.googleapis.com/v1/publishers/google/models/gemini-3-flash-preview:generateContent", httpStub.lastReq.URL.String())
|
||||
requireVertexAPIKeyAuth(t, httpStub.lastReq)
|
||||
}
|
||||
|
||||
func TestGeminiMessagesCompatServiceForward_VertexAPIKeyWithProjectUsesProjectEndpoint(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc, httpStub := newGeminiVertexAPIKeyTestService()
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
|
||||
body := []byte(`{"model":"gemini-3-flash-preview","max_tokens":16,"messages":[{"role":"user","content":"hello"}]}`)
|
||||
|
||||
result, err := svc.Forward(context.Background(), c, geminiVertexAPIKeyTestAccount("my-project"), body)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.Equal(t, "https://us-central1-aiplatform.googleapis.com/v1/projects/my-project/locations/us-central1/publishers/google/models/gemini-3-flash-preview:generateContent", httpStub.lastReq.URL.String())
|
||||
requireVertexAPIKeyAuth(t, httpStub.lastReq)
|
||||
}
|
||||
|
||||
func TestGeminiForwardAsChatCompletions_VertexAPIKeyUsesAPIKeyHeader(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc, httpStub := newGeminiVertexAPIKeyTestService()
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
body := []byte(`{"model":"gemini-3-flash-preview","messages":[{"role":"user","content":"hi"}]}`)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
|
||||
|
||||
result, err := svc.ForwardAsChatCompletions(context.Background(), c, geminiVertexAPIKeyTestAccount(""), body)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
require.Equal(t, "https://aiplatform.googleapis.com/v1/publishers/google/models/gemini-3-flash-preview:generateContent", httpStub.lastReq.URL.String())
|
||||
requireVertexAPIKeyAuth(t, httpStub.lastReq)
|
||||
}
|
||||
@@ -52,6 +52,14 @@ const (
|
||||
// ensureForwardErrorResponse 检查此 key,为 true 时跳过兜底写入,避免在已完成的 JSON 后追加 SSE。
|
||||
ResponseCommittedKey = "response_committed"
|
||||
|
||||
// OpsLocalErrorTypeKey / OpsLocalErrorCodeKey 记录网关以 Google 错误格式回写的错误类型与错误码:
|
||||
// Google 错误体没有 type 字段,ops_error_logger 据此按与 Claude / OpenAI 格式错误相同的口径分类。
|
||||
OpsLocalErrorTypeKey = "ops_local_error_type"
|
||||
OpsLocalErrorCodeKey = "ops_local_error_code"
|
||||
// OpsRequestScopedErrorKey 标记失败由请求内容决定(如引用的显式缓存不存在):只按错误类型分类,
|
||||
// 上游错误上下文保留为明细但不参与归因。
|
||||
OpsRequestScopedErrorKey = "ops_request_scoped_error"
|
||||
|
||||
OpsClientBusinessLimitedKey = "ops_client_business_limited"
|
||||
OpsClientBusinessLimitedReasonKey = "ops_client_business_limited_reason"
|
||||
OpsClientBusinessLimitedReasonIPRestriction = "api_key_ip_restriction"
|
||||
@@ -134,6 +142,42 @@ func OpsClientBusinessLimitedReason(c *gin.Context) string {
|
||||
return strings.TrimSpace(reason)
|
||||
}
|
||||
|
||||
// SetOpsLocalErrorType 登记响应错误体缺少的错误类型与错误码,并清除请求级失败标记;errType 为空表示不登记。
|
||||
func SetOpsLocalErrorType(c *gin.Context, errType, code string) {
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
c.Set(OpsLocalErrorTypeKey, strings.TrimSpace(errType))
|
||||
c.Set(OpsLocalErrorCodeKey, strings.TrimSpace(code))
|
||||
c.Set(OpsRequestScopedErrorKey, false)
|
||||
}
|
||||
|
||||
// OpsLocalErrorType 返回 SetOpsLocalErrorType 登记的错误类型与错误码;未登记时返回空串。
|
||||
func OpsLocalErrorType(c *gin.Context) (errType, code string) {
|
||||
if c == nil {
|
||||
return "", ""
|
||||
}
|
||||
errType = c.GetString(OpsLocalErrorTypeKey)
|
||||
if errType == "" {
|
||||
return "", ""
|
||||
}
|
||||
return errType, c.GetString(OpsLocalErrorCodeKey)
|
||||
}
|
||||
|
||||
// MarkOpsRequestScopedError 登记错误类型,并标记失败由请求内容决定:上游错误上下文不参与 ops 归因。
|
||||
func MarkOpsRequestScopedError(c *gin.Context, errType string) {
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
SetOpsLocalErrorType(c, errType, "")
|
||||
c.Set(OpsRequestScopedErrorKey, true)
|
||||
}
|
||||
|
||||
// IsOpsRequestScopedError 判断 MarkOpsRequestScopedError 是否标记了本次请求的失败。
|
||||
func IsOpsRequestScopedError(c *gin.Context) bool {
|
||||
return c != nil && c.GetBool(OpsRequestScopedErrorKey)
|
||||
}
|
||||
|
||||
// OpsStreamError 描述承载在 2xx 响应上的带内错误:网关在响应状态已固化为 200 之后
|
||||
// 就地以 SSE error 帧返回的错误(并发限流回退、Wait 后二次计费校验失败、流开始后才无
|
||||
// 可用账号等),以及上游 2xx 正文或事件里携带的错误结果(NonStream 标记非流式正文)。
|
||||
|
||||
@@ -231,6 +231,9 @@ type LiteLLMModelPricing struct {
|
||||
OutputCostPerImageToken float64 `json:"output_cost_per_image_token"` // 图片输出 token 价格
|
||||
InputCostPerImageToken float64 `json:"input_cost_per_image_token"` // 图片输入 token 价格(如 gpt-image-2 图片编辑)
|
||||
CacheReadInputImageTokenCost float64 `json:"cache_read_input_image_token_cost"`
|
||||
// CacheStorageCostPerTokenPerHour 显式上下文缓存存储单价(每 token 每小时,USD),一般由定价覆盖文件提供,
|
||||
// 配置后优先于内置价(见 BillingService.CacheStoragePricePerTokenHour)。
|
||||
CacheStorageCostPerTokenPerHour float64 `json:"cache_storage_cost_per_token_per_hour"`
|
||||
|
||||
// TokenPricingAbsent 表示源数据中 input/output token 价格均缺失(仅有图片价)。
|
||||
// 此类条目只可用于图片计费,token 计费必须回退到 fallback 或 fail-closed,
|
||||
@@ -266,6 +269,7 @@ type LiteLLMRawEntry struct {
|
||||
OutputCostPerImageToken *float64 `json:"output_cost_per_image_token"`
|
||||
InputCostPerImageToken *float64 `json:"input_cost_per_image_token"`
|
||||
CacheReadInputImageTokenCost *float64 `json:"cache_read_input_image_token_cost"`
|
||||
CacheStorageCostPerTokenPerHour *float64 `json:"cache_storage_cost_per_token_per_hour"`
|
||||
}
|
||||
|
||||
// PricingService 动态价格服务
|
||||
@@ -738,6 +742,9 @@ func (s *PricingService) parsePricingData(body []byte) (map[string]*LiteLLMModel
|
||||
if entry.CacheReadInputImageTokenCost != nil {
|
||||
pricing.CacheReadInputImageTokenCost = *entry.CacheReadInputImageTokenCost
|
||||
}
|
||||
if entry.CacheStorageCostPerTokenPerHour != nil {
|
||||
pricing.CacheStorageCostPerTokenPerHour = *entry.CacheStorageCostPerTokenPerHour
|
||||
}
|
||||
|
||||
hasExplicitLongContext := entry.LongContextInputTokenThreshold != nil ||
|
||||
entry.LongContextInputCostMultiplier != nil ||
|
||||
|
||||
@@ -613,10 +613,10 @@ func TestBillingService_Gemini36FlashThinkingTierFallbacksAreBillable(t *testing
|
||||
t.Run(model, func(t *testing.T) {
|
||||
cost, err := svc.CalculateCost(model, tokens, 1)
|
||||
require.NoError(t, err)
|
||||
require.InDelta(t, 1.5, cost.InputCost, 1e-12)
|
||||
require.InDelta(t, 7.5, cost.OutputCost, 1e-12)
|
||||
require.InDelta(t, 0.15, cost.CacheReadCost, 1e-12)
|
||||
require.InDelta(t, 9.15, cost.TotalCost, 1e-12)
|
||||
require.InDelta(t, 0.75, cost.InputCost, 1e-12)
|
||||
require.InDelta(t, 3.75, cost.OutputCost, 1e-12)
|
||||
require.InDelta(t, 0.075, cost.CacheReadCost, 1e-12)
|
||||
require.InDelta(t, 4.575, cost.TotalCost, 1e-12)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -725,9 +725,9 @@ func TestDefaultPricingIncludesGemini36FlashRates(t *testing.T) {
|
||||
t.Run(model, func(t *testing.T) {
|
||||
pricing, err := billingSvc.GetModelPricing(model)
|
||||
require.NoError(t, err)
|
||||
require.InDelta(t, 1.5e-6, pricing.InputPricePerToken, 1e-12)
|
||||
require.InDelta(t, 7.5e-6, pricing.OutputPricePerToken, 1e-12)
|
||||
require.InDelta(t, 0.15e-6, pricing.CacheReadPricePerToken, 1e-12)
|
||||
require.InDelta(t, 0.75e-6, pricing.InputPricePerToken, 1e-12)
|
||||
require.InDelta(t, 3.75e-6, pricing.OutputPricePerToken, 1e-12)
|
||||
require.InDelta(t, 0.075e-6, pricing.CacheReadPricePerToken, 1e-12)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -10,6 +10,7 @@ import (
|
||||
"net/http"
|
||||
"net/url"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
@@ -25,6 +26,10 @@ const (
|
||||
UpstreamModelMetadataExtraKey = "upstream_model_metadata"
|
||||
UpstreamModelMetadataIncompleteCode = "upstream_model_metadata_incomplete"
|
||||
UpstreamModelMetadataPartialCode = "upstream_model_metadata_partial"
|
||||
|
||||
vertexPublisherModelsPageSize = 300
|
||||
vertexPublisherModelsMaxPages = 20
|
||||
vertexGooglePublisherModelPrefix = "publishers/google/models/"
|
||||
)
|
||||
|
||||
type UpstreamModelMetadata struct {
|
||||
@@ -741,34 +746,19 @@ func (s *AccountTestService) fetchUpstreamModelList(ctx context.Context, account
|
||||
return nil, nil, newUpstreamModelSyncConfigError("Upstream HTTP client is not configured", nil)
|
||||
}
|
||||
|
||||
if account.IsGemini() && account.IsVertexServiceAccount() {
|
||||
models, err := s.fetchVertexGeminiUpstreamModels(ctx, account)
|
||||
return models, nil, err
|
||||
}
|
||||
|
||||
req, err := s.buildUpstreamModelsRequest(ctx, account)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
|
||||
proxyURL := upstreamModelsProxyURL(account)
|
||||
resp, err := s.doUpstreamModelsRequest(req, proxyURL, account)
|
||||
body, err := s.readUpstreamModelsResponse(req, account)
|
||||
if err != nil {
|
||||
return nil, nil, newUpstreamModelSyncUpstreamError("Failed to request upstream model list", err)
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
bodyLimit := resolveModelsListReadLimit(s.cfg)
|
||||
body, err := io.ReadAll(io.LimitReader(resp.Body, bodyLimit+1))
|
||||
if err != nil {
|
||||
return nil, nil, newUpstreamModelSyncUpstreamError("Failed to read upstream model list", err)
|
||||
}
|
||||
if int64(len(body)) > bodyLimit {
|
||||
return nil, nil, newUpstreamModelSyncUpstreamError("Upstream model list response is too large", fmt.Errorf("response exceeds %d bytes", bodyLimit))
|
||||
}
|
||||
|
||||
if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
|
||||
return nil, nil, &UpstreamModelSyncError{
|
||||
Kind: UpstreamModelSyncErrorUpstream,
|
||||
Message: fmt.Sprintf("Upstream model list request failed with HTTP %d", resp.StatusCode),
|
||||
StatusCode: resp.StatusCode,
|
||||
Err: fmt.Errorf("upstream model list returned HTTP %d", resp.StatusCode),
|
||||
}
|
||||
return nil, nil, err
|
||||
}
|
||||
|
||||
extractModels := extractUpstreamModelIDs
|
||||
@@ -786,6 +776,128 @@ func (s *AccountTestService) fetchUpstreamModelList(ctx context.Context, account
|
||||
return models, body, nil
|
||||
}
|
||||
|
||||
// readUpstreamModelsResponse 发送模型列表请求并读取受限大小的响应体;非 2xx 返回带状态码的同步错误。
|
||||
func (s *AccountTestService) readUpstreamModelsResponse(req *http.Request, account *Account) ([]byte, error) {
|
||||
resp, err := s.doUpstreamModelsRequest(req, upstreamModelsProxyURL(account), account)
|
||||
if err != nil {
|
||||
return nil, newUpstreamModelSyncUpstreamError("Failed to request upstream model list", err)
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
bodyLimit := resolveModelsListReadLimit(s.cfg)
|
||||
body, err := io.ReadAll(io.LimitReader(resp.Body, bodyLimit+1))
|
||||
if err != nil {
|
||||
return nil, newUpstreamModelSyncUpstreamError("Failed to read upstream model list", err)
|
||||
}
|
||||
if int64(len(body)) > bodyLimit {
|
||||
return nil, newUpstreamModelSyncUpstreamError("Upstream model list response is too large", fmt.Errorf("response exceeds %d bytes", bodyLimit))
|
||||
}
|
||||
|
||||
if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
|
||||
return nil, &UpstreamModelSyncError{
|
||||
Kind: UpstreamModelSyncErrorUpstream,
|
||||
Message: fmt.Sprintf("Upstream model list request failed with HTTP %d", resp.StatusCode),
|
||||
StatusCode: resp.StatusCode,
|
||||
Err: fmt.Errorf("upstream model list returned HTTP %d", resp.StatusCode),
|
||||
}
|
||||
}
|
||||
return body, nil
|
||||
}
|
||||
|
||||
// fetchVertexGeminiUpstreamModels 按账号默认区域读取 Vertex 上可用的 Gemini 模型。列表接口不接受 API Key。
|
||||
func (s *AccountTestService) fetchVertexGeminiUpstreamModels(ctx context.Context, account *Account) ([]string, error) {
|
||||
if account.IsVertexAPIKey() {
|
||||
return nil, newUpstreamModelSyncUnsupportedError("Vertex API key accounts do not support upstream model listing", nil)
|
||||
}
|
||||
if _, err := parseVertexServiceAccountKey(account); err != nil {
|
||||
return nil, newUpstreamModelSyncConfigError("Invalid Vertex service account credentials", err)
|
||||
}
|
||||
location := account.VertexLocation("")
|
||||
if _, err := buildVertexPublisherModelsURL(location, ""); err != nil {
|
||||
return nil, newUpstreamModelSyncConfigError("Invalid Vertex location", err)
|
||||
}
|
||||
if s.geminiTokenProvider == nil {
|
||||
return nil, newUpstreamModelSyncConfigError("Gemini token provider is not configured", nil)
|
||||
}
|
||||
accessToken, err := s.geminiTokenProvider.GetAccessToken(ctx, account)
|
||||
if err != nil {
|
||||
return nil, newUpstreamModelSyncUpstreamError("Failed to get Vertex access token", err)
|
||||
}
|
||||
accessToken = strings.TrimSpace(accessToken)
|
||||
if accessToken == "" {
|
||||
return nil, newUpstreamModelSyncConfigError("No Vertex access token is available", nil)
|
||||
}
|
||||
|
||||
return listVertexGeminiPublisherModels(location, func(page int, listURL string) ([]byte, error) {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, listURL, nil)
|
||||
if err != nil {
|
||||
return nil, newUpstreamModelSyncConfigError("Invalid Vertex model list URL", err)
|
||||
}
|
||||
req.Header.Set("Accept", "application/json")
|
||||
req.Header.Set("Authorization", "Bearer "+accessToken)
|
||||
|
||||
body, err := s.readUpstreamModelsResponse(req, account)
|
||||
if err != nil {
|
||||
// 只有首页的 404/405 表示列表端点不可用;翻页途中的失败不触发配置模型回退。
|
||||
var syncErr *UpstreamModelSyncError
|
||||
if page > 0 && errors.As(err, &syncErr) {
|
||||
syncErr.StatusCode = 0
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
return body, nil
|
||||
})
|
||||
}
|
||||
|
||||
// listVertexGeminiPublisherModels 按区域分页读取 Google 发布方模型列表,只保留 gemini- 前缀的模型(去重排序)。
|
||||
// 区域端点还会列出 Gemma、Veo、AutoML 等非 Gemini 的模型花园条目;global 端点只含 Gemini。
|
||||
// fetchPage 负责发送请求并处理非 2xx 响应,返回该页的响应体;它返回的错误原样透出。
|
||||
func listVertexGeminiPublisherModels(location string, fetchPage func(page int, listURL string) ([]byte, error)) ([]string, error) {
|
||||
models := make([]string, 0)
|
||||
pageToken := ""
|
||||
for page := 0; page < vertexPublisherModelsMaxPages; page++ {
|
||||
listURL, err := buildVertexPublisherModelsURL(location, pageToken)
|
||||
if err != nil {
|
||||
return nil, newUpstreamModelSyncConfigError("Invalid Vertex location", err)
|
||||
}
|
||||
body, err := fetchPage(page, listURL)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var response struct {
|
||||
PublisherModels []struct {
|
||||
Name string `json:"name"`
|
||||
} `json:"publisherModels"`
|
||||
NextPageToken string `json:"nextPageToken"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &response); err != nil {
|
||||
return nil, newUpstreamModelSyncUpstreamError("Upstream model list response was not valid JSON", err)
|
||||
}
|
||||
for _, entry := range response.PublisherModels {
|
||||
name := strings.TrimSpace(entry.Name)
|
||||
if !strings.HasPrefix(name, vertexGooglePublisherModelPrefix) {
|
||||
continue
|
||||
}
|
||||
modelID := strings.TrimPrefix(name, vertexGooglePublisherModelPrefix)
|
||||
if strings.HasPrefix(modelID, "gemini-") && len(modelID) > len("gemini-") {
|
||||
models = append(models, modelID)
|
||||
}
|
||||
}
|
||||
|
||||
pageToken = strings.TrimSpace(response.NextPageToken)
|
||||
if pageToken == "" {
|
||||
if len(models) == 0 {
|
||||
return nil, newUpstreamModelSyncUpstreamError("Upstream returned no supported models", nil)
|
||||
}
|
||||
return dedupeAndSortModelIDs(models), nil
|
||||
}
|
||||
}
|
||||
return nil, newUpstreamModelSyncUpstreamError(
|
||||
"Upstream model list pagination did not finish",
|
||||
fmt.Errorf("more than %d pages", vertexPublisherModelsMaxPages),
|
||||
)
|
||||
}
|
||||
|
||||
func (s *AccountTestService) buildUpstreamModelsRequest(ctx context.Context, account *Account) (*http.Request, error) {
|
||||
switch {
|
||||
case account.Platform == PlatformAntigravity:
|
||||
@@ -1218,6 +1330,26 @@ func buildOpenAIModelsURL(base string) string {
|
||||
return buildOpenAIEndpointURL(base, "/v1/models")
|
||||
}
|
||||
|
||||
// buildVertexPublisherModelsURL 构造 Google 发布方模型列表地址。该接口只有 v1beta1 版本,pageSize 上限 300。
|
||||
func buildVertexPublisherModelsURL(location, pageToken string) (string, error) {
|
||||
location = strings.TrimSpace(location)
|
||||
if location == "" {
|
||||
location = vertexDefaultLocation
|
||||
}
|
||||
if !vertexLocationPattern.MatchString(location) {
|
||||
return "", fmt.Errorf("invalid vertex location: %s", location)
|
||||
}
|
||||
host := location + "-aiplatform.googleapis.com"
|
||||
if location == "global" {
|
||||
host = vertexGlobalHost
|
||||
}
|
||||
query := url.Values{"pageSize": []string{strconv.Itoa(vertexPublisherModelsPageSize)}}
|
||||
if pageToken != "" {
|
||||
query.Set("pageToken", pageToken)
|
||||
}
|
||||
return fmt.Sprintf("https://%s/v1beta1/publishers/google/models?%s", host, query.Encode()), nil
|
||||
}
|
||||
|
||||
func buildGeminiModelsURL(base string) string {
|
||||
normalized := strings.TrimRight(strings.TrimSpace(base), "/")
|
||||
if strings.HasSuffix(normalized, "/v1beta/models") {
|
||||
|
||||
@@ -0,0 +1,276 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type vertexModelsTokenCacheStub struct {
|
||||
token string
|
||||
calls int
|
||||
}
|
||||
|
||||
func (c *vertexModelsTokenCacheStub) GetAccessToken(context.Context, string) (string, error) {
|
||||
c.calls++
|
||||
if c.token == "" {
|
||||
return "", errors.New("missing token")
|
||||
}
|
||||
return c.token, nil
|
||||
}
|
||||
|
||||
func (c *vertexModelsTokenCacheStub) SetAccessToken(context.Context, string, string, time.Duration) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *vertexModelsTokenCacheStub) DeleteAccessToken(context.Context, string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *vertexModelsTokenCacheStub) AcquireRefreshLock(context.Context, string, time.Duration) (bool, error) {
|
||||
return false, nil
|
||||
}
|
||||
|
||||
func (c *vertexModelsTokenCacheStub) ReleaseRefreshLock(context.Context, string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func vertexModelSyncTestAccount(credentials map[string]any) *Account {
|
||||
merged := map[string]any{
|
||||
"service_account_json": map[string]any{
|
||||
"type": "service_account",
|
||||
"project_id": "proj",
|
||||
"private_key": "-----BEGIN PRIVATE KEY-----\nabc\n-----END PRIVATE KEY-----\n",
|
||||
"client_email": "svc@proj.iam.gserviceaccount.com",
|
||||
},
|
||||
}
|
||||
for key, value := range credentials {
|
||||
merged[key] = value
|
||||
}
|
||||
return &Account{
|
||||
ID: 173,
|
||||
Platform: PlatformGemini,
|
||||
Type: AccountTypeServiceAccount,
|
||||
Credentials: merged,
|
||||
}
|
||||
}
|
||||
|
||||
func vertexModelSyncTestService(upstream *httpUpstreamRecorder) *AccountTestService {
|
||||
return &AccountTestService{
|
||||
httpUpstream: upstream,
|
||||
cfg: upstreamModelSyncTestConfig(),
|
||||
geminiTokenProvider: NewGeminiTokenProvider(nil, &vertexModelsTokenCacheStub{token: "ya29.vertex-token"}, nil),
|
||||
}
|
||||
}
|
||||
|
||||
func vertexModelsJSONResponse(status int, body string) *http.Response {
|
||||
return &http.Response{
|
||||
StatusCode: status,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(body)),
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildVertexPublisherModelsURL(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
u, err := buildVertexPublisherModelsURL("global", "")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "https://aiplatform.googleapis.com/v1beta1/publishers/google/models?pageSize=300", u)
|
||||
|
||||
u, err = buildVertexPublisherModelsURL("us-central1", "next/page+token")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "https://us-central1-aiplatform.googleapis.com/v1beta1/publishers/google/models?pageSize=300&pageToken=next%2Fpage%2Btoken", u)
|
||||
|
||||
u, err = buildVertexPublisherModelsURL("", "")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "https://us-central1-aiplatform.googleapis.com/v1beta1/publishers/google/models?pageSize=300", u)
|
||||
|
||||
_, err = buildVertexPublisherModelsURL("evil.com/x", "")
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
// Scenario: Vertex SA 账号按账号区域分页读取发布方模型列表,只保留 gemini- 前缀模型。
|
||||
func TestFetchUpstreamSupportedModelsVertexServiceAccountPaginatesAndKeepsGemini(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
upstream := &httpUpstreamRecorder{responses: []*http.Response{
|
||||
vertexModelsJSONResponse(http.StatusOK, `{"publisherModels":[
|
||||
{"name":"publishers/google/models/gemini-3.8-flash"},
|
||||
{"name":"publishers/google/models/gemma4"},
|
||||
{"name":"publishers/google/models/veo-3.1-generate-001"}
|
||||
],"nextPageToken":"page-2"}`),
|
||||
vertexModelsJSONResponse(http.StatusOK, `{"publisherModels":[
|
||||
{"name":" publishers/google/models/gemini-2.5-flash "},
|
||||
{"name":"publishers/google/models/gemini-3.8-flash"},
|
||||
{"name":"gemini-without-publisher-prefix"},
|
||||
{"name":"publishers/anthropic/models/gemini-other-publisher"},
|
||||
{"name":"publishers/google/models/gemini-"}
|
||||
]}`),
|
||||
}}
|
||||
svc := vertexModelSyncTestService(upstream)
|
||||
|
||||
models, err := svc.FetchUpstreamSupportedModels(context.Background(), vertexModelSyncTestAccount(map[string]any{
|
||||
"location": "europe-west4",
|
||||
}))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []string{"gemini-2.5-flash", "gemini-3.8-flash"}, models)
|
||||
|
||||
require.Len(t, upstream.requests, 2)
|
||||
require.Equal(t, "https://europe-west4-aiplatform.googleapis.com/v1beta1/publishers/google/models?pageSize=300", upstream.requests[0].URL.String())
|
||||
require.Equal(t, "https://europe-west4-aiplatform.googleapis.com/v1beta1/publishers/google/models?pageSize=300&pageToken=page-2", upstream.requests[1].URL.String())
|
||||
for _, req := range upstream.requests {
|
||||
require.Equal(t, http.MethodGet, req.Method)
|
||||
require.Equal(t, "Bearer ya29.vertex-token", req.Header.Get("Authorization"))
|
||||
}
|
||||
}
|
||||
|
||||
func TestFetchUpstreamSupportedModelsVertexUpstreamErrorKeepsStatus(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
upstream := &httpUpstreamRecorder{resp: vertexModelsJSONResponse(http.StatusForbidden, `{"error":{"code":403,"status":"PERMISSION_DENIED"}}`)}
|
||||
svc := vertexModelSyncTestService(upstream)
|
||||
|
||||
_, err := svc.FetchUpstreamSupportedModels(context.Background(), vertexModelSyncTestAccount(map[string]any{"location": "global"}))
|
||||
var syncErr *UpstreamModelSyncError
|
||||
require.ErrorAs(t, err, &syncErr)
|
||||
require.Equal(t, UpstreamModelSyncErrorUpstream, syncErr.Kind)
|
||||
require.Equal(t, http.StatusForbidden, syncErr.StatusCode)
|
||||
require.Equal(t, "https://aiplatform.googleapis.com/v1beta1/publishers/google/models?pageSize=300", upstream.lastReq.URL.String())
|
||||
}
|
||||
|
||||
func TestFetchUpstreamSupportedModelsVertexWithoutGeminiModelsFails(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
upstream := &httpUpstreamRecorder{resp: vertexModelsJSONResponse(http.StatusOK, `{"publisherModels":[{"name":"publishers/google/models/gemma4"}]}`)}
|
||||
svc := vertexModelSyncTestService(upstream)
|
||||
|
||||
_, err := svc.FetchUpstreamSupportedModels(context.Background(), vertexModelSyncTestAccount(nil))
|
||||
var syncErr *UpstreamModelSyncError
|
||||
require.ErrorAs(t, err, &syncErr)
|
||||
require.Equal(t, UpstreamModelSyncErrorUpstream, syncErr.Kind)
|
||||
require.Equal(t, "Upstream returned no supported models", syncErr.SafeMessage())
|
||||
}
|
||||
|
||||
func TestFetchUpstreamSupportedModelsVertexStopsEndlessPagination(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
responses := make([]*http.Response, 0, vertexPublisherModelsMaxPages)
|
||||
for i := 0; i < vertexPublisherModelsMaxPages; i++ {
|
||||
responses = append(responses, vertexModelsJSONResponse(http.StatusOK, `{"publisherModels":[{"name":"publishers/google/models/gemini-3.8-flash"}],"nextPageToken":"again"}`))
|
||||
}
|
||||
upstream := &httpUpstreamRecorder{responses: responses}
|
||||
svc := vertexModelSyncTestService(upstream)
|
||||
|
||||
_, err := svc.FetchUpstreamSupportedModels(context.Background(), vertexModelSyncTestAccount(nil))
|
||||
var syncErr *UpstreamModelSyncError
|
||||
require.ErrorAs(t, err, &syncErr)
|
||||
require.Equal(t, "Upstream model list pagination did not finish", syncErr.SafeMessage())
|
||||
require.Len(t, upstream.requests, vertexPublisherModelsMaxPages)
|
||||
}
|
||||
|
||||
func TestFetchUpstreamSupportedModelsVertexAPIKeyUnsupported(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
upstream := &httpUpstreamRecorder{}
|
||||
svc := vertexModelSyncTestService(upstream)
|
||||
account := &Account{
|
||||
ID: 174,
|
||||
Platform: PlatformGemini,
|
||||
Type: AccountTypeServiceAccount,
|
||||
Credentials: map[string]any{
|
||||
"auth_mode": "apikey",
|
||||
"api_key": "AQ.test-key",
|
||||
},
|
||||
}
|
||||
|
||||
_, err := svc.FetchUpstreamSupportedModels(context.Background(), account)
|
||||
var syncErr *UpstreamModelSyncError
|
||||
require.ErrorAs(t, err, &syncErr)
|
||||
require.Equal(t, UpstreamModelSyncErrorUnsupported, syncErr.Kind)
|
||||
require.Empty(t, upstream.requests)
|
||||
}
|
||||
|
||||
func TestFetchUpstreamSupportedModelsVertexInvalidCredentialsIsConfigError(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
cache := &vertexModelsTokenCacheStub{token: "ya29.vertex-token"}
|
||||
upstream := &httpUpstreamRecorder{}
|
||||
svc := &AccountTestService{
|
||||
httpUpstream: upstream,
|
||||
cfg: upstreamModelSyncTestConfig(),
|
||||
geminiTokenProvider: NewGeminiTokenProvider(nil, cache, nil),
|
||||
}
|
||||
account := &Account{
|
||||
ID: 175,
|
||||
Platform: PlatformGemini,
|
||||
Type: AccountTypeServiceAccount,
|
||||
Credentials: map[string]any{"api_key": "not-a-service-account"},
|
||||
}
|
||||
|
||||
_, err := svc.FetchUpstreamSupportedModels(context.Background(), account)
|
||||
var syncErr *UpstreamModelSyncError
|
||||
require.ErrorAs(t, err, &syncErr)
|
||||
require.Equal(t, UpstreamModelSyncErrorConfiguration, syncErr.Kind)
|
||||
require.Equal(t, "Invalid Vertex service account credentials", syncErr.SafeMessage())
|
||||
require.Zero(t, cache.calls)
|
||||
require.Empty(t, upstream.requests)
|
||||
}
|
||||
|
||||
func TestFetchUpstreamSupportedModelsVertexInvalidLocationSkipsTokenExchange(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
cache := &vertexModelsTokenCacheStub{token: "ya29.vertex-token"}
|
||||
upstream := &httpUpstreamRecorder{}
|
||||
svc := &AccountTestService{
|
||||
httpUpstream: upstream,
|
||||
cfg: upstreamModelSyncTestConfig(),
|
||||
geminiTokenProvider: NewGeminiTokenProvider(nil, cache, nil),
|
||||
}
|
||||
|
||||
_, err := svc.FetchUpstreamSupportedModels(context.Background(), vertexModelSyncTestAccount(map[string]any{
|
||||
"location": "us-central1.evil.com",
|
||||
}))
|
||||
var syncErr *UpstreamModelSyncError
|
||||
require.ErrorAs(t, err, &syncErr)
|
||||
require.Equal(t, UpstreamModelSyncErrorConfiguration, syncErr.Kind)
|
||||
require.Equal(t, "Invalid Vertex location", syncErr.SafeMessage())
|
||||
require.Zero(t, cache.calls)
|
||||
require.Empty(t, upstream.requests)
|
||||
}
|
||||
|
||||
// Scenario: 翻页途中的 404 不能被当成列表端点不可用而回退到配置模型。
|
||||
func TestSyncUpstreamModelCatalogVertexMidPagination404DoesNotFallback(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
upstream := &httpUpstreamRecorder{responses: []*http.Response{
|
||||
vertexModelsJSONResponse(http.StatusOK, `{"publisherModels":[{"name":"publishers/google/models/gemini-a"}],"nextPageToken":"page-2"}`),
|
||||
vertexModelsJSONResponse(http.StatusNotFound, `{"error":{"code":404}}`),
|
||||
}}
|
||||
svc := vertexModelSyncTestService(upstream)
|
||||
|
||||
_, err := svc.SyncUpstreamModelCatalog(context.Background(), vertexModelSyncTestAccount(map[string]any{
|
||||
"model_mapping": map[string]any{"x": "gemini-mapped"},
|
||||
}))
|
||||
var syncErr *UpstreamModelSyncError
|
||||
require.ErrorAs(t, err, &syncErr)
|
||||
require.Equal(t, UpstreamModelSyncErrorUpstream, syncErr.Kind)
|
||||
require.Zero(t, syncErr.StatusCode)
|
||||
require.Len(t, upstream.requests, 2)
|
||||
}
|
||||
|
||||
func TestFetchUpstreamSupportedModelsVertexFirstPage404KeepsStatus(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
upstream := &httpUpstreamRecorder{resp: vertexModelsJSONResponse(http.StatusNotFound, `{"error":{"code":404}}`)}
|
||||
svc := vertexModelSyncTestService(upstream)
|
||||
|
||||
_, err := svc.FetchUpstreamSupportedModels(context.Background(), vertexModelSyncTestAccount(nil))
|
||||
require.True(t, upstreamModelListEndpointUnsupported(err))
|
||||
}
|
||||
@@ -29,6 +29,8 @@ const (
|
||||
vertexServiceAccountCacheSkew = 5 * time.Minute
|
||||
vertexLockWaitTime = 200 * time.Millisecond
|
||||
vertexAnthropicVersion = "vertex-2023-10-16"
|
||||
vertexGlobalHost = "aiplatform.googleapis.com"
|
||||
vertexAuthModeAPIKey = "apikey"
|
||||
)
|
||||
|
||||
var (
|
||||
@@ -58,6 +60,13 @@ func (a *Account) IsVertexServiceAccount() bool {
|
||||
return a != nil && a.Type == AccountTypeServiceAccount
|
||||
}
|
||||
|
||||
// IsVertexAPIKey 判断 Gemini 平台的 Vertex 账号是否使用 API Key 认证(credentials.auth_mode=apikey)。
|
||||
// Anthropic 平台的 Vertex 账号只支持 Service Account JSON。
|
||||
func (a *Account) IsVertexAPIKey() bool {
|
||||
return a.IsVertexServiceAccount() && a.Platform == PlatformGemini &&
|
||||
strings.TrimSpace(a.GetCredential("auth_mode")) == vertexAuthModeAPIKey
|
||||
}
|
||||
|
||||
func (a *Account) VertexProjectID() string {
|
||||
if a == nil {
|
||||
return ""
|
||||
@@ -283,6 +292,63 @@ func exchangeVertexServiceAccountToken(ctx context.Context, key *vertexServiceAc
|
||||
return parsed.AccessToken, ttl, nil
|
||||
}
|
||||
|
||||
// buildVertexGeminiAccountURL 按账号构造 Vertex Gemini 请求地址:未配置 project_id 的 API Key 账号
|
||||
// 走 Express mode 全局端点,其余账号走项目级端点。
|
||||
func buildVertexGeminiAccountURL(account *Account, model, action string, stream bool) (string, error) {
|
||||
if account.IsVertexAPIKey() && account.VertexProjectID() == "" {
|
||||
return buildVertexGeminiExpressURL(model, action, stream)
|
||||
}
|
||||
return buildVertexGeminiURL(account.VertexProjectID(), account.VertexLocation(model), model, action, stream)
|
||||
}
|
||||
|
||||
// buildVertexGeminiExpressURL 构造不带项目与区域的全局端点地址,仅 API Key 认证可用。
|
||||
func buildVertexGeminiExpressURL(model, action string, stream bool) (string, error) {
|
||||
model = strings.TrimSpace(model)
|
||||
action = strings.TrimSpace(action)
|
||||
if model == "" {
|
||||
return "", errors.New("vertex model is required")
|
||||
}
|
||||
if err := validateVertexGeminiAction(action); err != nil {
|
||||
return "", err
|
||||
}
|
||||
u := fmt.Sprintf("https://%s/v1/publishers/google/models/%s:%s", vertexGlobalHost, url.PathEscape(model), action)
|
||||
if stream {
|
||||
u += "?alt=sse"
|
||||
}
|
||||
return u, nil
|
||||
}
|
||||
|
||||
func validateVertexGeminiAction(action string) error {
|
||||
switch action {
|
||||
case "generateContent", "streamGenerateContent", "countTokens":
|
||||
return nil
|
||||
default:
|
||||
return fmt.Errorf("unsupported vertex gemini action: %s", action)
|
||||
}
|
||||
}
|
||||
|
||||
// setVertexGeminiAuth 给 Vertex Gemini 请求设置认证头:API Key 账号用 x-goog-api-key,
|
||||
// Service Account JSON 账号用换取的 Bearer access token。
|
||||
func setVertexGeminiAuth(ctx context.Context, req *http.Request, account *Account, tokenProvider *GeminiTokenProvider) error {
|
||||
if account.IsVertexAPIKey() {
|
||||
apiKey := strings.TrimSpace(account.GetCredential("api_key"))
|
||||
if apiKey == "" {
|
||||
return errors.New("vertex api_key not configured")
|
||||
}
|
||||
req.Header.Set("x-goog-api-key", apiKey)
|
||||
return nil
|
||||
}
|
||||
if tokenProvider == nil {
|
||||
return errors.New("gemini token provider not configured")
|
||||
}
|
||||
accessToken, err := tokenProvider.GetAccessToken(ctx, account)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
req.Header.Set("Authorization", "Bearer "+accessToken)
|
||||
return nil
|
||||
}
|
||||
|
||||
func buildVertexGeminiURL(projectID, location, model, action string, stream bool) (string, error) {
|
||||
projectID = strings.TrimSpace(projectID)
|
||||
location = strings.TrimSpace(location)
|
||||
@@ -300,14 +366,12 @@ func buildVertexGeminiURL(projectID, location, model, action string, stream bool
|
||||
if model == "" {
|
||||
return "", errors.New("vertex model is required")
|
||||
}
|
||||
switch action {
|
||||
case "generateContent", "streamGenerateContent", "countTokens":
|
||||
default:
|
||||
return "", fmt.Errorf("unsupported vertex gemini action: %s", action)
|
||||
if err := validateVertexGeminiAction(action); err != nil {
|
||||
return "", err
|
||||
}
|
||||
host := fmt.Sprintf("%s-aiplatform.googleapis.com", location)
|
||||
if location == "global" {
|
||||
host = "aiplatform.googleapis.com"
|
||||
host = vertexGlobalHost
|
||||
}
|
||||
u := fmt.Sprintf(
|
||||
"https://%s/v1/projects/%s/locations/%s/publishers/google/models/%s:%s",
|
||||
|
||||
@@ -30,6 +30,62 @@ func TestBuildVertexGeminiURLUsesGlobalEndpointHost(t *testing.T) {
|
||||
require.Equal(t, "https://aiplatform.googleapis.com/v1/projects/my-project/locations/global/publishers/google/models/gemini-3-flash-preview:streamGenerateContent?alt=sse", got)
|
||||
}
|
||||
|
||||
func TestIsVertexAPIKey(t *testing.T) {
|
||||
apiKeyCreds := map[string]any{"auth_mode": "apikey", "api_key": "k"}
|
||||
require.True(t, (&Account{Platform: PlatformGemini, Type: AccountTypeServiceAccount, Credentials: apiKeyCreds}).IsVertexAPIKey())
|
||||
require.False(t, (&Account{Platform: PlatformGemini, Type: AccountTypeServiceAccount, Credentials: map[string]any{"service_account_json": "{}"}}).IsVertexAPIKey())
|
||||
require.False(t, (&Account{Platform: PlatformAnthropic, Type: AccountTypeServiceAccount, Credentials: apiKeyCreds}).IsVertexAPIKey())
|
||||
require.False(t, (&Account{Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: apiKeyCreds}).IsVertexAPIKey())
|
||||
}
|
||||
|
||||
func TestBuildVertexGeminiAccountURL(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
credentials map[string]any
|
||||
want string
|
||||
}{
|
||||
{
|
||||
name: "api key without project uses express global endpoint",
|
||||
credentials: map[string]any{"auth_mode": "apikey", "api_key": "k", "location": "us-central1"},
|
||||
want: "https://aiplatform.googleapis.com/v1/publishers/google/models/gemini-3-flash-preview:streamGenerateContent?alt=sse",
|
||||
},
|
||||
{
|
||||
name: "api key with project uses project endpoint",
|
||||
credentials: map[string]any{"auth_mode": "apikey", "api_key": "k", "project_id": "my-project", "location": "us-central1"},
|
||||
want: "https://us-central1-aiplatform.googleapis.com/v1/projects/my-project/locations/us-central1/publishers/google/models/gemini-3-flash-preview:streamGenerateContent?alt=sse",
|
||||
},
|
||||
{
|
||||
name: "service account uses project endpoint",
|
||||
credentials: map[string]any{"project_id": "sa-project", "location": "global"},
|
||||
want: "https://aiplatform.googleapis.com/v1/projects/sa-project/locations/global/publishers/google/models/gemini-3-flash-preview:streamGenerateContent?alt=sse",
|
||||
},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
account := &Account{Platform: PlatformGemini, Type: AccountTypeServiceAccount, Credentials: tc.credentials}
|
||||
got, err := buildVertexGeminiAccountURL(account, "gemini-3-flash-preview", "streamGenerateContent", true)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, tc.want, got)
|
||||
})
|
||||
}
|
||||
|
||||
account := &Account{Platform: PlatformGemini, Type: AccountTypeServiceAccount, Credentials: map[string]any{"auth_mode": "apikey", "api_key": "k"}}
|
||||
_, err := buildVertexGeminiAccountURL(account, "gemini-3-flash-preview", "rawPredict", false)
|
||||
require.ErrorContains(t, err, "unsupported vertex gemini action")
|
||||
}
|
||||
|
||||
func TestSetVertexGeminiAuthUsesAPIKeyHeader(t *testing.T) {
|
||||
account := &Account{Platform: PlatformGemini, Type: AccountTypeServiceAccount, Credentials: map[string]any{"auth_mode": "apikey", "api_key": " vertex-key "}}
|
||||
req := httptest.NewRequest(http.MethodPost, "https://aiplatform.googleapis.com/v1/publishers/google/models/gemini-3-flash-preview:generateContent", nil)
|
||||
|
||||
require.NoError(t, setVertexGeminiAuth(context.Background(), req, account, nil))
|
||||
require.Equal(t, "vertex-key", req.Header.Get("x-goog-api-key"))
|
||||
require.Empty(t, req.Header.Get("Authorization"))
|
||||
|
||||
account.Credentials["api_key"] = ""
|
||||
require.ErrorContains(t, setVertexGeminiAuth(context.Background(), req, account, nil), "vertex api_key not configured")
|
||||
}
|
||||
|
||||
func TestBuildVertexAnthropicURL(t *testing.T) {
|
||||
got, err := buildVertexAnthropicURL("my-project", "us-east5", "claude-sonnet-4-5@20250929", false)
|
||||
require.NoError(t, err)
|
||||
|
||||
@@ -640,6 +640,13 @@ func ProvideIdempotencyCleanupService(repo IdempotencyRepository, cfg *config.Co
|
||||
return svc
|
||||
}
|
||||
|
||||
// ProvideGeminiCachedContentService creates GeminiCachedContentService and starts its mutation recovery loop.
|
||||
func ProvideGeminiCachedContentService(repo GeminiCachedContentRepository, compat *GeminiMessagesCompatService, cfg *config.Config) *GeminiCachedContentService {
|
||||
svc := NewGeminiCachedContentService(repo, compat, cfg)
|
||||
svc.Start()
|
||||
return svc
|
||||
}
|
||||
|
||||
// ProvideScheduledTestService creates ScheduledTestService.
|
||||
func ProvideScheduledTestService(
|
||||
planRepo ScheduledTestPlanRepository,
|
||||
@@ -885,6 +892,7 @@ var ProviderSet = wire.NewSet(
|
||||
ProvideGrokOAuthService,
|
||||
wire.Bind(new(GrokOAuthTokenService), new(*GrokOAuthService)),
|
||||
NewGeminiOAuthService,
|
||||
ProvideGeminiCachedContentService,
|
||||
NewGeminiQuotaService,
|
||||
NewCompositeTokenCacheInvalidator,
|
||||
wire.Bind(new(TokenCacheInvalidator), new(*CompositeTokenCacheInvalidator)),
|
||||
|
||||
@@ -42,6 +42,8 @@ type textRedactPatterns struct {
|
||||
var (
|
||||
reGOCSPX = regexp.MustCompile(`GOCSPX-[0-9A-Za-z_-]{24,}`)
|
||||
reAIza = regexp.MustCompile(`AIza[0-9A-Za-z_-]{35}`)
|
||||
// Gemini auth key:AQ.Ab 前缀 + 48 位,长度放宽以覆盖格式变体。
|
||||
reGeminiAuthKey = regexp.MustCompile(`AQ\.Ab[0-9A-Za-z_-]{40,}`)
|
||||
|
||||
defaultTextRedactPatterns = compileTextRedactPatterns(nil)
|
||||
extraTextPatternCache sync.Map // map[string]*textRedactPatterns
|
||||
@@ -99,6 +101,7 @@ func RedactText(input string, extraKeys ...string) string {
|
||||
out := input
|
||||
out = reGOCSPX.ReplaceAllString(out, "GOCSPX-***")
|
||||
out = reAIza.ReplaceAllString(out, "AIza***")
|
||||
out = reGeminiAuthKey.ReplaceAllString(out, "AQ.***")
|
||||
out = patterns.reJSONLike.ReplaceAllString(out, `$1***$3`)
|
||||
out = patterns.reQueryLike.ReplaceAllString(out, `$1=***`)
|
||||
out = patterns.rePlain.ReplaceAllString(out, `$1$2***`)
|
||||
|
||||
@@ -38,6 +38,28 @@ func TestRedactText_GOCSPX(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestRedactText_GeminiAPIKeys(t *testing.T) {
|
||||
legacyKey := "AI" + "za" + strings.Repeat("Ab1_-", 7)
|
||||
authKey := "AQ." + "Ab" + strings.Repeat("Xy9_-z", 8)
|
||||
in := "upstream rejected " + legacyKey + " and " + authKey + " for model gemini-3.1-pro"
|
||||
out := RedactText(in)
|
||||
if strings.Contains(out, legacyKey) || strings.Contains(out, authKey) {
|
||||
t.Fatalf("expected Gemini keys redacted, got %q", out)
|
||||
}
|
||||
for _, want := range []string{"AIza***", "AQ.***", "for model gemini-3.1-pro"} {
|
||||
if !strings.Contains(out, want) {
|
||||
t.Fatalf("expected %q in %q", want, out)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRedactText_GeminiAuthKeyPrefixAloneUnchanged(t *testing.T) {
|
||||
in := "key format AQ.Ab... is documented"
|
||||
if out := RedactText(in); out != in {
|
||||
t.Fatalf("expected short AQ. prefix unchanged, got %q", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRedactText_ExtraKeyCacheUsesNormalizedSortedKey(t *testing.T) {
|
||||
clearExtraTextPatternCache()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user