From 27600b1d2c9579f83abbb1ee469bfc45b908f840 Mon Sep 17 00:00:00 2001 From: Pluviobyte Date: Thu, 28 May 2026 05:40:50 +0000 Subject: [PATCH 01/36] fix(gateway): filter count_tokens generation fields Anthropic count_tokens rejects generation-only fields such as temperature, top_p, top_k, stream, and stop sequences. Passing the original messages payload through unchanged can turn otherwise valid requests into upstream 400 errors. Sanitize only the count_tokens upstream payload after the gateway's existing request normalization, preserving fields that existing compatibility paths rely on while removing parameters the count_tokens endpoint does not accept. Fixes #2764 Co-authored-by: Cursor --- ...teway_anthropic_apikey_passthrough_test.go | 60 +++++++++++++++++++ backend/internal/service/gateway_service.go | 21 +++++++ 2 files changed, 81 insertions(+) diff --git a/backend/internal/service/gateway_anthropic_apikey_passthrough_test.go b/backend/internal/service/gateway_anthropic_apikey_passthrough_test.go index 5cb03f3017ce..9062c517aa76 100644 --- a/backend/internal/service/gateway_anthropic_apikey_passthrough_test.go +++ b/backend/internal/service/gateway_anthropic_apikey_passthrough_test.go @@ -476,6 +476,66 @@ func TestGatewayService_AnthropicAPIKeyPassthrough_ModelMappingPreservesOtherFie require.Equal(t, int64(1024), gjson.GetBytes(sentBody, "max_tokens").Int(), "max_tokens 不应被修改") } +func TestGatewayService_AnthropicAPIKeyPassthrough_CountTokensFiltersGenerationFields(t *testing.T) { + gin.SetMode(gin.TestMode) + + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages/count_tokens", nil) + + body := []byte(`{"model":"claude-sonnet-4-20250514","system":[{"type":"text","text":"sys"}],"messages":[{"role":"user","content":"hello"}],"tools":[{"name":"tool","input_schema":{"type":"object"}}],"temperature":0.7,"top_p":0.9,"top_k":40,"stream":true,"stop_sequences":["END"],"max_tokens":1024,"thinking":{"type":"enabled","budget_tokens":5000}}`) + parsed := &ParsedRequest{ + Body: body, + Model: "claude-sonnet-4-20250514", + } + + upstreamRespBody := `{"input_tokens":42}` + upstream := &anthropicHTTPUpstreamRecorder{ + resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(upstreamRespBody)), + }, + } + + svc := &GatewayService{ + cfg: &config.Config{Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize}}, + httpUpstream: upstream, + rateLimitService: &RateLimitService{}, + } + + account := &Account{ + ID: 302, + Name: "count-token-filter-test", + Platform: PlatformAnthropic, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "upstream-key", + "base_url": "https://api.anthropic.com", + }, + Extra: map[string]any{"anthropic_passthrough": true}, + Status: StatusActive, + Schedulable: true, + } + + err := svc.ForwardCountTokens(context.Background(), c, account, parsed) + require.NoError(t, err) + + sentBody := upstream.lastBody + require.False(t, gjson.GetBytes(sentBody, "temperature").Exists()) + require.False(t, gjson.GetBytes(sentBody, "top_p").Exists()) + require.False(t, gjson.GetBytes(sentBody, "top_k").Exists()) + require.False(t, gjson.GetBytes(sentBody, "stream").Exists()) + require.False(t, gjson.GetBytes(sentBody, "stop_sequences").Exists()) + require.Equal(t, "claude-sonnet-4-20250514", gjson.GetBytes(sentBody, "model").String()) + require.Equal(t, "sys", gjson.GetBytes(sentBody, "system.0.text").String()) + require.Equal(t, "hello", gjson.GetBytes(sentBody, "messages.0.content").String()) + require.Equal(t, "tool", gjson.GetBytes(sentBody, "tools.0.name").String()) + require.Equal(t, int64(1024), gjson.GetBytes(sentBody, "max_tokens").Int()) + require.Equal(t, "enabled", gjson.GetBytes(sentBody, "thinking.type").String()) +} + // TestGatewayService_AnthropicAPIKeyPassthrough_EmptyModelSkipsMapping // 确保空模型名不会触发映射逻辑 func TestGatewayService_AnthropicAPIKeyPassthrough_EmptyModelSkipsMapping(t *testing.T) { diff --git a/backend/internal/service/gateway_service.go b/backend/internal/service/gateway_service.go index 4a8175a4c18f..a787e3ebb8a3 100644 --- a/backend/internal/service/gateway_service.go +++ b/backend/internal/service/gateway_service.go @@ -9311,6 +9311,7 @@ func (s *GatewayService) buildCountTokensRequestAnthropicAPIKeyPassthrough( } targetURL = validatedURL + "/v1/messages/count_tokens?beta=true" } + body = sanitizeCountTokensRequestBody(body) req, err := http.NewRequestWithContext(ctx, http.MethodPost, targetURL, bytes.NewReader(body)) if err != nil { @@ -9405,6 +9406,7 @@ func (s *GatewayService) buildCountTokensRequest(ctx context.Context, c *gin.Con if ctEnableCCH { body = signBillingHeaderCCH(body) } + body = sanitizeCountTokensRequestBody(body) req, err := http.NewRequestWithContext(ctx, "POST", targetURL, bytes.NewReader(body)) if err != nil { @@ -9501,6 +9503,25 @@ func (s *GatewayService) buildCountTokensRequest(ctx context.Context, c *gin.Con return req, nil } +func sanitizeCountTokensRequestBody(body []byte) []byte { + out := body + for _, path := range []string{ + "temperature", + "top_p", + "top_k", + "stream", + "stop_sequences", + "stop", + } { + if gjson.GetBytes(out, path).Exists() { + if next, ok := deleteJSONPathBytes(out, path); ok { + out = next + } + } + } + return out +} + // countTokensError 返回 count_tokens 错误响应 func (s *GatewayService) countTokensError(c *gin.Context, status int, errType, message string) { c.JSON(status, gin.H{ From 2bd3125d0fe23515fd42ee3c6651bbe38e47593c Mon Sep 17 00:00:00 2001 From: Wey Gu Date: Thu, 28 May 2026 22:08:02 +0800 Subject: [PATCH 02/36] Preserve usage request context --- backend/internal/handler/gateway_handler.go | 7 +-- .../gateway_handler_chat_completions.go | 2 +- .../handler/gateway_handler_responses.go | 2 +- .../internal/handler/gemini_v1beta_handler.go | 2 +- .../handler/openai_chat_completions.go | 2 +- backend/internal/handler/openai_embeddings.go | 2 +- .../handler/openai_gateway_handler.go | 44 +++++++++++++--- .../openai_gateway_usage_context_test.go | 41 +++++++++++++++ backend/internal/handler/openai_images.go | 2 +- .../handler/usage_record_submit_task_test.go | 24 ++++----- .../server/middleware/client_request_id.go | 6 ++- .../middleware/client_request_id_test.go | 50 +++++++++++++++++++ 12 files changed, 154 insertions(+), 30 deletions(-) create mode 100644 backend/internal/handler/openai_gateway_usage_context_test.go create mode 100644 backend/internal/server/middleware/client_request_id_test.go diff --git a/backend/internal/handler/gateway_handler.go b/backend/internal/handler/gateway_handler.go index 4695a7915ab3..a67491919ef8 100644 --- a/backend/internal/handler/gateway_handler.go +++ b/backend/internal/handler/gateway_handler.go @@ -510,7 +510,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) { // 使用量记录通过有界 worker 池提交，避免请求热路径创建无界 goroutine。 quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey) - h.submitUsageRecordTask(func(ctx context.Context) { + h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) { if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{ Result: result, ParsedRequest: parsedReq, @@ -905,7 +905,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) { // 使用量记录通过有界 worker 池提交，避免请求热路径创建无界 goroutine。 quotaPlatform := service.QuotaPlatform(c.Request.Context(), currentAPIKey) - h.submitUsageRecordTask(func(ctx context.Context) { + h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) { if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{ Result: result, ParsedRequest: parsedReq, @@ -2056,10 +2056,11 @@ func (h *GatewayHandler) maybeLogCompatibilityFallbackMetrics(reqLog *zap.Logger ) } -func (h *GatewayHandler) submitUsageRecordTask(task service.UsageRecordTask) { +func (h *GatewayHandler) submitUsageRecordTask(parent context.Context, task service.UsageRecordTask) { if task == nil { return } + task = wrapUsageRecordTaskContext(parent, task) if h.usageRecordWorkerPool != nil { h.usageRecordWorkerPool.Submit(task) return diff --git a/backend/internal/handler/gateway_handler_chat_completions.go b/backend/internal/handler/gateway_handler_chat_completions.go index acbdc261d206..daf6e6ea0415 100644 --- a/backend/internal/handler/gateway_handler_chat_completions.go +++ b/backend/internal/handler/gateway_handler_chat_completions.go @@ -292,7 +292,7 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) { upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform) quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey) - h.submitUsageRecordTask(func(ctx context.Context) { + h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) { if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{ Result: result, QuotaPlatform: quotaPlatform, diff --git a/backend/internal/handler/gateway_handler_responses.go b/backend/internal/handler/gateway_handler_responses.go index 6a083f317b91..f57b9989cf02 100644 --- a/backend/internal/handler/gateway_handler_responses.go +++ b/backend/internal/handler/gateway_handler_responses.go @@ -267,7 +267,7 @@ func (h *GatewayHandler) Responses(c *gin.Context) { upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform) quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey) - h.submitUsageRecordTask(func(ctx context.Context) { + h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) { if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{ Result: result, QuotaPlatform: quotaPlatform, diff --git a/backend/internal/handler/gemini_v1beta_handler.go b/backend/internal/handler/gemini_v1beta_handler.go index 27ea44044f26..0b33ca3e3efe 100644 --- a/backend/internal/handler/gemini_v1beta_handler.go +++ b/backend/internal/handler/gemini_v1beta_handler.go @@ -528,7 +528,7 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) { inboundEndpoint := GetInboundEndpoint(c) upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform) quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey) - h.submitUsageRecordTask(func(ctx context.Context) { + h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) { if err := h.gatewayService.RecordUsageWithLongContext(ctx, &service.RecordUsageLongContextInput{ Result: result, QuotaPlatform: quotaPlatform, diff --git a/backend/internal/handler/openai_chat_completions.go b/backend/internal/handler/openai_chat_completions.go index 17f0d47e458d..9f63ef1f7285 100644 --- a/backend/internal/handler/openai_chat_completions.go +++ b/backend/internal/handler/openai_chat_completions.go @@ -273,7 +273,7 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) { inboundEndpoint := GetInboundEndpoint(c) upstreamEndpoint := resolveRawCCUpstreamEndpoint(c, account) - h.submitOpenAIUsageRecordTask(result, func(ctx context.Context) { + h.submitOpenAIUsageRecordTask(c.Request.Context(), result, func(ctx context.Context) { if err := h.gatewayService.RecordUsage(ctx, &service.OpenAIRecordUsageInput{ Result: result, APIKey: apiKey, diff --git a/backend/internal/handler/openai_embeddings.go b/backend/internal/handler/openai_embeddings.go index bbb67044b683..b64ac41d5734 100644 --- a/backend/internal/handler/openai_embeddings.go +++ b/backend/internal/handler/openai_embeddings.go @@ -220,7 +220,7 @@ func (h *OpenAIGatewayHandler) Embeddings(c *gin.Context) { inboundEndpoint := GetInboundEndpoint(c) upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform) - h.submitOpenAIUsageRecordTask(result, func(ctx context.Context) { + h.submitOpenAIUsageRecordTask(c.Request.Context(), result, func(ctx context.Context) { if err := h.gatewayService.RecordUsage(ctx, &service.OpenAIRecordUsageInput{ Result: result, APIKey: apiKey, diff --git a/backend/internal/handler/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index a51eee86ebee..86503f301361 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -12,6 +12,7 @@ import ( "time" "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey" 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" @@ -46,6 +47,31 @@ func resolveOpenAIMessagesDispatchMappedModel(apiKey *service.APIKey, requestedM return strings.TrimSpace(apiKey.Group.ResolveMessagesDispatchModel(requestedModel)) } +func usageRecordContext(parent context.Context, base context.Context) context.Context { + if base == nil { + base = context.Background() + } + if parent == nil { + return base + } + if clientRequestID, _ := parent.Value(ctxkey.ClientRequestID).(string); strings.TrimSpace(clientRequestID) != "" { + base = context.WithValue(base, ctxkey.ClientRequestID, strings.TrimSpace(clientRequestID)) + } + if requestID, _ := parent.Value(ctxkey.RequestID).(string); strings.TrimSpace(requestID) != "" { + base = context.WithValue(base, ctxkey.RequestID, strings.TrimSpace(requestID)) + } + return base +} + +func wrapUsageRecordTaskContext(parent context.Context, task service.UsageRecordTask) service.UsageRecordTask { + if task == nil { + return nil + } + return func(ctx context.Context) { + task(usageRecordContext(parent, ctx)) + } +} + // NewOpenAIGatewayHandler creates a new OpenAIGatewayHandler func NewOpenAIGatewayHandler( gatewayService *service.OpenAIGatewayService, @@ -437,7 +463,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform) // 使用量记录通过有界 worker 池提交，避免请求热路径创建无界 goroutine。 - h.submitOpenAIUsageRecordTask(result, func(ctx context.Context) { + h.submitOpenAIUsageRecordTask(c.Request.Context(), result, func(ctx context.Context) { if err := h.gatewayService.RecordUsage(ctx, &service.OpenAIRecordUsageInput{ Result: result, APIKey: apiKey, @@ -821,7 +847,7 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) { inboundEndpoint := GetInboundEndpoint(c) upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform) - h.submitOpenAIUsageRecordTask(result, func(ctx context.Context) { + h.submitOpenAIUsageRecordTask(c.Request.Context(), result, func(ctx context.Context) { if err := h.gatewayService.RecordUsage(ctx, &service.OpenAIRecordUsageInput{ Result: result, APIKey: apiKey, @@ -1424,7 +1450,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, result.FirstTokenMs) inboundEndpoint := GetInboundEndpoint(c) upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform) - h.submitOpenAIUsageRecordTask(result, func(taskCtx context.Context) { + h.submitOpenAIUsageRecordTask(ctx, result, func(taskCtx context.Context) { if err := h.gatewayService.RecordUsage(taskCtx, &service.OpenAIRecordUsageInput{ Result: result, APIKey: apiKey, @@ -1609,10 +1635,11 @@ func getContextInt64(c *gin.Context, key string) (int64, bool) { } } -func (h *OpenAIGatewayHandler) submitUsageRecordTask(task service.UsageRecordTask) { +func (h *OpenAIGatewayHandler) submitUsageRecordTask(parent context.Context, task service.UsageRecordTask) { if task == nil { return } + task = wrapUsageRecordTaskContext(parent, task) if h.usageRecordWorkerPool != nil { h.usageRecordWorkerPool.Submit(task) return @@ -1631,18 +1658,19 @@ func (h *OpenAIGatewayHandler) submitUsageRecordTask(task service.UsageRecordTas task(ctx) } -func (h *OpenAIGatewayHandler) submitOpenAIUsageRecordTask(result *service.OpenAIForwardResult, task service.UsageRecordTask) { +func (h *OpenAIGatewayHandler) submitOpenAIUsageRecordTask(parent context.Context, result *service.OpenAIForwardResult, task service.UsageRecordTask) { if result != nil && result.ImageCount > 0 { - h.submitMandatoryUsageRecordTask(task) + h.submitMandatoryUsageRecordTask(parent, task) return } - h.submitUsageRecordTask(task) + h.submitUsageRecordTask(parent, task) } -func (h *OpenAIGatewayHandler) submitMandatoryUsageRecordTask(task service.UsageRecordTask) { +func (h *OpenAIGatewayHandler) submitMandatoryUsageRecordTask(parent context.Context, task service.UsageRecordTask) { if task == nil { return } + task = wrapUsageRecordTaskContext(parent, task) if h.usageRecordWorkerPool != nil { if mode := h.usageRecordWorkerPool.Submit(task); mode != service.UsageRecordSubmitModeDropped { return diff --git a/backend/internal/handler/openai_gateway_usage_context_test.go b/backend/internal/handler/openai_gateway_usage_context_test.go new file mode 100644 index 000000000000..7091c9c09944 --- /dev/null +++ b/backend/internal/handler/openai_gateway_usage_context_test.go @@ -0,0 +1,41 @@ +package handler + +import ( + "context" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey" + "github.com/stretchr/testify/require" +) + +func TestSubmitUsageRecordTaskCopiesRequestContext(t *testing.T) { + parent := context.WithValue(context.Background(), ctxkey.ClientRequestID, "client-request-123") + parent = context.WithValue(parent, ctxkey.RequestID, "request-456") + + var gotClientRequestID string + var gotRequestID string + h := &GatewayHandler{} + h.submitUsageRecordTask(parent, func(ctx context.Context) { + gotClientRequestID, _ = ctx.Value(ctxkey.ClientRequestID).(string) + gotRequestID, _ = ctx.Value(ctxkey.RequestID).(string) + }) + + require.Equal(t, "client-request-123", gotClientRequestID) + require.Equal(t, "request-456", gotRequestID) +} + +func TestOpenAISubmitUsageRecordTaskCopiesRequestContext(t *testing.T) { + parent := context.WithValue(context.Background(), ctxkey.ClientRequestID, "openai-client-request-123") + parent = context.WithValue(parent, ctxkey.RequestID, "openai-request-456") + + var gotClientRequestID string + var gotRequestID string + h := &OpenAIGatewayHandler{} + h.submitUsageRecordTask(parent, func(ctx context.Context) { + gotClientRequestID, _ = ctx.Value(ctxkey.ClientRequestID).(string) + gotRequestID, _ = ctx.Value(ctxkey.RequestID).(string) + }) + + require.Equal(t, "openai-client-request-123", gotClientRequestID) + require.Equal(t, "openai-request-456", gotRequestID) +} diff --git a/backend/internal/handler/openai_images.go b/backend/internal/handler/openai_images.go index bbb080149c2b..36339d4b7582 100644 --- a/backend/internal/handler/openai_images.go +++ b/backend/internal/handler/openai_images.go @@ -311,7 +311,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { if result != nil { upstreamModel = result.UpstreamModel } - h.submitMandatoryUsageRecordTask(func(ctx context.Context) { + h.submitMandatoryUsageRecordTask(c.Request.Context(), func(ctx context.Context) { if err := h.gatewayService.RecordUsage(ctx, &service.OpenAIRecordUsageInput{ Result: result, APIKey: apiKey, diff --git a/backend/internal/handler/usage_record_submit_task_test.go b/backend/internal/handler/usage_record_submit_task_test.go index e4c2837a592c..ebe5c3df8266 100644 --- a/backend/internal/handler/usage_record_submit_task_test.go +++ b/backend/internal/handler/usage_record_submit_task_test.go @@ -29,7 +29,7 @@ func TestGatewayHandlerSubmitUsageRecordTask_WithPool(t *testing.T) { h := &GatewayHandler{usageRecordWorkerPool: pool} done := make(chan struct{}) - h.submitUsageRecordTask(func(ctx context.Context) { + h.submitUsageRecordTask(context.Background(), func(ctx context.Context) { close(done) }) @@ -44,7 +44,7 @@ func TestGatewayHandlerSubmitUsageRecordTask_WithoutPoolSyncFallback(t *testing. h := &GatewayHandler{} var called atomic.Bool - h.submitUsageRecordTask(func(ctx context.Context) { + h.submitUsageRecordTask(context.Background(), func(ctx context.Context) { if _, ok := ctx.Deadline(); !ok { t.Fatal("expected deadline in fallback context") } @@ -57,7 +57,7 @@ func TestGatewayHandlerSubmitUsageRecordTask_WithoutPoolSyncFallback(t *testing. func TestGatewayHandlerSubmitUsageRecordTask_NilTask(t *testing.T) { h := &GatewayHandler{} require.NotPanics(t, func() { - h.submitUsageRecordTask(nil) + h.submitUsageRecordTask(context.Background(), nil) }) } @@ -66,12 +66,12 @@ func TestGatewayHandlerSubmitUsageRecordTask_WithoutPool_TaskPanicRecovered(t *t var called atomic.Bool require.NotPanics(t, func() { - h.submitUsageRecordTask(func(ctx context.Context) { + h.submitUsageRecordTask(context.Background(), func(ctx context.Context) { panic("usage task panic") }) }) - h.submitUsageRecordTask(func(ctx context.Context) { + h.submitUsageRecordTask(context.Background(), func(ctx context.Context) { called.Store(true) }) require.True(t, called.Load(), "panic 后后续任务应仍可执行") @@ -82,7 +82,7 @@ func TestOpenAIGatewayHandlerSubmitUsageRecordTask_WithPool(t *testing.T) { h := &OpenAIGatewayHandler{usageRecordWorkerPool: pool} done := make(chan struct{}) - h.submitUsageRecordTask(func(ctx context.Context) { + h.submitUsageRecordTask(context.Background(), func(ctx context.Context) { close(done) }) @@ -97,7 +97,7 @@ func TestOpenAIGatewayHandlerSubmitUsageRecordTask_WithoutPoolSyncFallback(t *te h := &OpenAIGatewayHandler{} var called atomic.Bool - h.submitUsageRecordTask(func(ctx context.Context) { + h.submitUsageRecordTask(context.Background(), func(ctx context.Context) { if _, ok := ctx.Deadline(); !ok { t.Fatal("expected deadline in fallback context") } @@ -110,7 +110,7 @@ func TestOpenAIGatewayHandlerSubmitUsageRecordTask_WithoutPoolSyncFallback(t *te func TestOpenAIGatewayHandlerSubmitUsageRecordTask_NilTask(t *testing.T) { h := &OpenAIGatewayHandler{} require.NotPanics(t, func() { - h.submitUsageRecordTask(nil) + h.submitUsageRecordTask(context.Background(), nil) }) } @@ -119,12 +119,12 @@ func TestOpenAIGatewayHandlerSubmitUsageRecordTask_WithoutPool_TaskPanicRecovere var called atomic.Bool require.NotPanics(t, func() { - h.submitUsageRecordTask(func(ctx context.Context) { + h.submitUsageRecordTask(context.Background(), func(ctx context.Context) { panic("usage task panic") }) }) - h.submitUsageRecordTask(func(ctx context.Context) { + h.submitUsageRecordTask(context.Background(), func(ctx context.Context) { called.Store(true) }) require.True(t, called.Load(), "panic 后后续任务应仍可执行") @@ -152,7 +152,7 @@ func TestOpenAIGatewayHandlerSubmitMandatoryUsageRecordTask_DroppedTaskSyncFallb pool.Submit(func(ctx context.Context) {}) var called atomic.Bool - h.submitMandatoryUsageRecordTask(func(ctx context.Context) { + h.submitMandatoryUsageRecordTask(context.Background(), func(ctx context.Context) { called.Store(true) }) close(release) @@ -182,7 +182,7 @@ func TestOpenAIGatewayHandlerSubmitOpenAIUsageRecordTask_ImageResultUsesMandator pool.Submit(func(ctx context.Context) {}) var called atomic.Bool - h.submitOpenAIUsageRecordTask(&service.OpenAIForwardResult{ImageCount: 1}, func(ctx context.Context) { + h.submitOpenAIUsageRecordTask(context.Background(), &service.OpenAIForwardResult{ImageCount: 1}, func(ctx context.Context) { called.Store(true) }) close(release) diff --git a/backend/internal/server/middleware/client_request_id.go b/backend/internal/server/middleware/client_request_id.go index 6838d6afd50c..5f886646eee9 100644 --- a/backend/internal/server/middleware/client_request_id.go +++ b/backend/internal/server/middleware/client_request_id.go @@ -11,6 +11,8 @@ import ( "go.uber.org/zap" ) +const clientRequestIDHeader = "X-Client-Request-ID" + // ClientRequestID ensures every request has a unique client_request_id in request.Context(). // // This is used by the Ops monitoring module for end-to-end request correlation. @@ -21,12 +23,14 @@ func ClientRequestID() gin.HandlerFunc { return } - if v := c.Request.Context().Value(ctxkey.ClientRequestID); v != nil { + if v, _ := c.Request.Context().Value(ctxkey.ClientRequestID).(string); strings.TrimSpace(v) != "" { + c.Header(clientRequestIDHeader, strings.TrimSpace(v)) c.Next() return } id := uuid.New().String() + c.Header(clientRequestIDHeader, id) ctx := context.WithValue(c.Request.Context(), ctxkey.ClientRequestID, id) requestLogger := logger.FromContext(ctx).With(zap.String("client_request_id", strings.TrimSpace(id))) ctx = logger.IntoContext(ctx, requestLogger) diff --git a/backend/internal/server/middleware/client_request_id_test.go b/backend/internal/server/middleware/client_request_id_test.go new file mode 100644 index 000000000000..394c16125d0f --- /dev/null +++ b/backend/internal/server/middleware/client_request_id_test.go @@ -0,0 +1,50 @@ +package middleware + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestClientRequestIDGeneratesAndExposesID(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(ClientRequestID()) + router.GET("/", func(c *gin.Context) { + value, _ := c.Request.Context().Value(ctxkey.ClientRequestID).(string) + c.String(http.StatusOK, value) + }) + + req := httptest.NewRequest(http.MethodGet, "/", nil) + w := httptest.NewRecorder() + + router.ServeHTTP(w, req) + + require.Equal(t, http.StatusOK, w.Code) + require.NotEmpty(t, w.Body.String()) + require.Equal(t, w.Body.String(), w.Header().Get(clientRequestIDHeader)) +} + +func TestClientRequestIDPreservesExistingContextID(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(ClientRequestID()) + router.GET("/", func(c *gin.Context) { + value, _ := c.Request.Context().Value(ctxkey.ClientRequestID).(string) + c.String(http.StatusOK, value) + }) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/", nil) + req = req.WithContext(context.WithValue(req.Context(), ctxkey.ClientRequestID, "existing-client-request-id")) + router.ServeHTTP(w, req) + + require.Equal(t, http.StatusOK, w.Code) + require.Equal(t, "existing-client-request-id", w.Body.String()) + require.Equal(t, "existing-client-request-id", w.Header().Get(clientRequestIDHeader)) +} From ead471d64bd7ebfc7769c36fa33d0d46bb015f80 Mon Sep 17 00:00:00 2001 From: wucm667 Date: Fri, 29 May 2026 10:38:00 +0800 Subject: [PATCH 03/36] =?UTF-8?q?feat(account):=20=E6=94=AF=E6=8C=81?= =?UTF-8?q?=E6=8C=89=205h/7d=20=E7=94=A8=E9=87=8F=E9=98=88=E5=80=BC?= =?UTF-8?q?=E8=87=AA=E5=8A=A8=E6=9A=82=E5=81=9C=E8=B4=A6=E5=8F=B7=E8=B0=83?= =?UTF-8?q?=E5=BA=A6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../internal/repository/scheduler_cache.go | 6 + .../repository/scheduler_cache_unit_test.go | 19 +++ .../service/openai_account_scheduler.go | 2 +- .../service/openai_account_scheduler_test.go | 111 ++++++++++++++ .../service/openai_gateway_service.go | 144 +++++++++++++++++- backend/internal/service/ops_settings.go | 19 +++ .../service/ops_settings_advanced_test.go | 16 ++ .../internal/service/ops_settings_models.go | 28 ++-- backend/internal/service/setting_service.go | 81 ++++++++++ .../components/account/EditAccountModal.vue | 66 ++++++-- .../__tests__/EditAccountModal.spec.ts | 22 +++ frontend/src/i18n/locales/en.ts | 3 + frontend/src/i18n/locales/zh.ts | 3 + 13 files changed, 496 insertions(+), 24 deletions(-) diff --git a/backend/internal/repository/scheduler_cache.go b/backend/internal/repository/scheduler_cache.go index ab01a8638ace..ec8c72dca978 100644 --- a/backend/internal/repository/scheduler_cache.go +++ b/backend/internal/repository/scheduler_cache.go @@ -548,6 +548,12 @@ func filterSchedulerExtra(extra map[string]any) map[string]any { "openai_ws_force_http", "openai_responses_mode", "openai_responses_supported", + "codex_5h_used_percent", + "codex_7d_used_percent", + "auto_pause_5h_threshold", + "auto_pause_7d_threshold", + "auto_pause_5h_limit", + "auto_pause_7d_limit", } filtered := make(map[string]any) for _, key := range keys { diff --git a/backend/internal/repository/scheduler_cache_unit_test.go b/backend/internal/repository/scheduler_cache_unit_test.go index 86de87c70995..fabc6bad242c 100644 --- a/backend/internal/repository/scheduler_cache_unit_test.go +++ b/backend/internal/repository/scheduler_cache_unit_test.go @@ -75,3 +75,22 @@ func TestBuildSchedulerMetadataAccount_KeepsSlimGroupMembership(t *testing.T) { require.Equal(t, int64(11), got.AccountGroups[1].GroupID) require.Nil(t, got.Groups) } + +func TestBuildSchedulerMetadataAccount_KeepsQuotaAutoPauseFields(t *testing.T) { + account := service.Account{ + ID: 88, + Extra: map[string]any{ + "codex_5h_used_percent": 12.34, + "codex_7d_used_percent": 56.78, + "auto_pause_5h_threshold": 0.95, + "auto_pause_7d_threshold": 0.96, + }, + } + + got := buildSchedulerMetadataAccount(account) + + require.Equal(t, 12.34, got.Extra["codex_5h_used_percent"]) + require.Equal(t, 56.78, got.Extra["codex_7d_used_percent"]) + require.Equal(t, 0.95, got.Extra["auto_pause_5h_threshold"]) + require.Equal(t, 0.96, got.Extra["auto_pause_7d_threshold"]) +} diff --git a/backend/internal/service/openai_account_scheduler.go b/backend/internal/service/openai_account_scheduler.go index 1eca08b128d1..fd28fa86c367 100644 --- a/backend/internal/service/openai_account_scheduler.go +++ b/backend/internal/service/openai_account_scheduler.go @@ -370,7 +370,6 @@ func (s *defaultOpenAIAccountScheduler) selectBySessionHash( _ = s.service.deleteStickySessionAccountID(ctx, req.GroupID, sessionHash) return nil, nil } - result, acquireErr := s.service.tryAcquireAccountSlot(ctx, accountID, account.Concurrency) if acquireErr == nil && result != nil && result.Acquired { _ = s.service.refreshStickySessionTTL(ctx, req.GroupID, sessionHash, s.service.openAIWSSessionStickyTTL()) @@ -1154,6 +1153,7 @@ func (s *OpenAIGatewayService) selectAccountWithScheduler( requiredImageCapability OpenAIImagesCapability, requireCompact bool, ) (*AccountSelectionResult, OpenAIAccountScheduleDecision, error) { + ctx = s.withOpenAIQuotaAutoPauseContext(ctx) decision := OpenAIAccountScheduleDecision{} scheduler := s.getOpenAIAccountScheduler(ctx) if scheduler == nil { diff --git a/backend/internal/service/openai_account_scheduler_test.go b/backend/internal/service/openai_account_scheduler_test.go index fedf7e9c11a6..37810870bc37 100644 --- a/backend/internal/service/openai_account_scheduler_test.go +++ b/backend/internal/service/openai_account_scheduler_test.go @@ -691,6 +691,117 @@ func TestOpenAIGatewayService_SelectAccountWithScheduler_SessionStickyRateLimite require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer) } +func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_AutoPauseBy5hThreshold(t *testing.T) { + ctx := context.Background() + primary := Account{ + ID: 35001, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 0, + Extra: map[string]any{ + "codex_5h_used_percent": 95.0, + "auto_pause_5h_threshold": 0.95, + "auto_pause_5h_limit": 100, + }, + } + secondary := Account{ID: 35002, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5} + svc := &OpenAIGatewayService{accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}}, cfg: &config.Config{}} + + account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.1", nil) + require.NoError(t, err) + require.NotNil(t, account) + require.Equal(t, int64(35002), account.ID) +} + +func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_AllowsBelow5hThreshold(t *testing.T) { + ctx := context.Background() + primary := Account{ + ID: 35101, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 0, + Extra: map[string]any{ + "codex_5h_used_percent": 80.0, + "auto_pause_5h_threshold": 0.95, + "auto_pause_5h_limit": 100, + }, + } + secondary := Account{ID: 35102, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5} + svc := &OpenAIGatewayService{accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}}, cfg: &config.Config{}} + + account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.1", nil) + require.NoError(t, err) + require.NotNil(t, account) + require.Equal(t, int64(35101), account.ID) +} + +func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_AutoPauseBy7dThreshold(t *testing.T) { + ctx := context.Background() + primary := Account{ + ID: 35201, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 0, + Extra: map[string]any{ + "codex_7d_used_percent": 95.0, + "auto_pause_7d_threshold": 0.95, + "auto_pause_7d_limit": 200, + }, + } + secondary := Account{ID: 35202, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5} + svc := &OpenAIGatewayService{accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}}, cfg: &config.Config{}} + + account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.1", nil) + require.NoError(t, err) + require.NotNil(t, account) + require.Equal(t, int64(35202), account.ID) +} + +func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_UnconfiguredThresholdKeepsLegacyBehavior(t *testing.T) { + ctx := context.Background() + primary := Account{ID: 35301, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, Extra: map[string]any{"codex_5h_used_percent": 99.0, "codex_7d_used_percent": 99.0}} + secondary := Account{ID: 35302, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5} + svc := &OpenAIGatewayService{accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}}, cfg: &config.Config{}} + + account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.1", nil) + require.NoError(t, err) + require.NotNil(t, account) + require.Equal(t, int64(35301), account.ID) +} + +func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_UsesGlobalDefaultThreshold(t *testing.T) { + ctx := withOpenAIQuotaAutoPauseSettings(context.Background(), OpsOpenAIAccountQuotaAutoPauseSettings{DefaultThreshold5h: 0.95}) + primary := Account{ + ID: 35401, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 0, + Extra: map[string]any{ + "codex_5h_used_percent": 95.0, + "auto_pause_5h_limit": 100, + }, + } + secondary := Account{ID: 35402, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5} + svc := &OpenAIGatewayService{accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}}, cfg: &config.Config{}} + + account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.1", nil) + require.NoError(t, err) + require.NotNil(t, account) + require.Equal(t, int64(35402), account.ID) +} + func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_SkipsFreshlyRateLimitedSnapshotCandidate(t *testing.T) { ctx := context.Background() groupID := int64(10102) diff --git a/backend/internal/service/openai_gateway_service.go b/backend/internal/service/openai_gateway_service.go index 623a64e980db..268f985cec80 100644 --- a/backend/internal/service/openai_gateway_service.go +++ b/backend/internal/service/openai_gateway_service.go @@ -1290,7 +1290,7 @@ func (s *OpenAIGatewayService) SelectAccountForModel(ctx context.Context, groupI // SelectAccountForModelWithExclusions selects an account supporting the requested model while excluding specified accounts. // SelectAccountForModelWithExclusions 选择支持指定模型的账号，同时排除指定的账号。 func (s *OpenAIGatewayService) SelectAccountForModelWithExclusions(ctx context.Context, groupID *int64, sessionHash string, requestedModel string, excludedIDs map[int64]struct{}) (*Account, error) { - return s.selectAccountForModelWithExclusions(ctx, groupID, sessionHash, requestedModel, excludedIDs, false, 0, "") + return s.selectAccountForModelWithExclusions(s.withOpenAIQuotaAutoPauseContext(ctx), groupID, sessionHash, requestedModel, excludedIDs, false, 0, "") } // noAvailableOpenAISelectionError builds the standard "no account available" error @@ -1327,6 +1327,18 @@ func isOpenAIAccountEligibleForRequest(ctx context.Context, account *Account, re if account == nil || !account.IsOpenAI() || !account.IsSchedulableForModelWithContext(ctx, requestedModel) { return false } + if paused, reason := shouldAutoPauseOpenAIAccountByQuota(ctx, account); paused { + slog.Info("account_auto_paused_by_quota", + "account_id", account.ID, + "usage_5h_percent", readOpenAIQuotaUsedPercent(account.Extra, "5h"), + "usage_7d_percent", readOpenAIQuotaUsedPercent(account.Extra, "7d"), + "threshold_type", reason.window, + "threshold", reason.threshold, + "limit", reason.limit, + "utilization", reason.utilization, + ) + return false + } if requestedModel != "" && !account.IsModelSupported(requestedModel) { return false } @@ -1339,6 +1351,134 @@ func isOpenAIAccountEligibleForRequest(ctx context.Context, account *Account, re return true } +type openAIQuotaAutoPauseDecision struct { + window string + threshold float64 + limit float64 + utilization float64 +} + +func shouldAutoPauseOpenAIAccountByQuota(ctx context.Context, account *Account) (bool, openAIQuotaAutoPauseDecision) { + if account == nil || !account.IsOpenAI() { + return false, openAIQuotaAutoPauseDecision{} + } + threshold5h, threshold7d := resolveOpenAIQuotaAutoPauseThresholds(ctx, account) + if threshold5h > 0 { + if utilization, limit, ok := resolveOpenAIQuotaUtilization(account.Extra, "5h"); ok { + if utilization >= threshold5h { + return true, openAIQuotaAutoPauseDecision{window: "5h", threshold: threshold5h, limit: limit, utilization: utilization} + } + } + } + if threshold7d > 0 { + if utilization, limit, ok := resolveOpenAIQuotaUtilization(account.Extra, "7d"); ok { + if utilization >= threshold7d { + return true, openAIQuotaAutoPauseDecision{window: "7d", threshold: threshold7d, limit: limit, utilization: utilization} + } + } + } + return false, openAIQuotaAutoPauseDecision{} +} + +func resolveOpenAIQuotaAutoPauseThresholds(ctx context.Context, account *Account) (float64, float64) { + threshold5h, _ := resolveAccountExtraNumber(account.Extra, "auto_pause_5h_threshold") + threshold7d, _ := resolveAccountExtraNumber(account.Extra, "auto_pause_7d_threshold") + threshold5h = clamp01(threshold5h) + threshold7d = clamp01(threshold7d) + if threshold5h > 0 && threshold7d > 0 { + return threshold5h, threshold7d + } + settings := openAIQuotaAutoPauseSettingsFromContext(ctx) + if threshold5h <= 0 { + threshold5h = clamp01(settings.DefaultThreshold5h) + } + if threshold7d <= 0 { + threshold7d = clamp01(settings.DefaultThreshold7d) + } + return threshold5h, threshold7d +} + +func resolveAccountExtraNumber(extra map[string]any, keys ...string) (float64, bool) { + if len(extra) == 0 { + return 0, false + } + for _, key := range keys { + value, ok := extra[key] + if !ok || value == nil { + continue + } + switch v := value.(type) { + case float64: + return v, true + case float32: + return float64(v), true + case int: + return float64(v), true + case int64: + return float64(v), true + case json.Number: + parsed, err := v.Float64() + if err == nil { + return parsed, true + } + case string: + parsed, err := strconv.ParseFloat(strings.TrimSpace(v), 64) + if err == nil { + return parsed, true + } + } + } + return 0, false +} + +func resolveOpenAIQuotaUtilization(extra map[string]any, window string) (float64, float64, bool) { + limitKeys := []string{"auto_pause_" + window + "_limit", "quota_" + window + "_limit", window + "_limit"} + if limit, ok := resolveAccountExtraNumber(extra, limitKeys...); ok && limit > 0 { + if usage, ok := resolveAccountExtraNumber(extra, "usage_"+window); ok && usage >= 0 { + return usage / limit, limit, true + } + } + usedPercent := readOpenAIQuotaUsedPercent(extra, window) + if usedPercent <= 0 { + return 0, 0, false + } + return usedPercent / 100, 100, true +} + +func readOpenAIQuotaUsedPercent(extra map[string]any, window string) float64 { + if len(extra) == 0 { + return 0 + } + if value, ok := resolveAccountExtraNumber(extra, "codex_"+window+"_used_percent"); ok { + return value + } + return 0 +} + +type openAIQuotaAutoPauseCtxKey struct{} + +func withOpenAIQuotaAutoPauseSettings(ctx context.Context, settings OpsOpenAIAccountQuotaAutoPauseSettings) context.Context { + if ctx == nil { + ctx = context.Background() + } + return context.WithValue(ctx, openAIQuotaAutoPauseCtxKey{}, settings) +} + +func openAIQuotaAutoPauseSettingsFromContext(ctx context.Context) OpsOpenAIAccountQuotaAutoPauseSettings { + if ctx == nil { + return OpsOpenAIAccountQuotaAutoPauseSettings{} + } + settings, _ := ctx.Value(openAIQuotaAutoPauseCtxKey{}).(OpsOpenAIAccountQuotaAutoPauseSettings) + return settings +} + +func (s *OpenAIGatewayService) withOpenAIQuotaAutoPauseContext(ctx context.Context) context.Context { + if s == nil || s.settingService == nil { + return ctx + } + return withOpenAIQuotaAutoPauseSettings(ctx, s.settingService.GetOpenAIQuotaAutoPauseSettings(ctx)) +} + // prioritizeOpenAICompactAccounts re-orders a slice so that accounts with known // compact support are tried first, followed by unknown, then explicitly unsupported. // The relative order within each tier is preserved. @@ -1587,7 +1727,7 @@ func (s *OpenAIGatewayService) isBetterAccount(candidate, current *Account) bool // SelectAccountWithLoadAwareness selects an account with load-awareness and wait plan. func (s *OpenAIGatewayService) SelectAccountWithLoadAwareness(ctx context.Context, groupID *int64, sessionHash string, requestedModel string, excludedIDs map[int64]struct{}) (*AccountSelectionResult, error) { - return s.selectAccountWithLoadAwareness(ctx, groupID, sessionHash, requestedModel, excludedIDs, false, "") + return s.selectAccountWithLoadAwareness(s.withOpenAIQuotaAutoPauseContext(ctx), groupID, sessionHash, requestedModel, excludedIDs, false, "") } func (s *OpenAIGatewayService) selectAccountWithLoadAwareness(ctx context.Context, groupID *int64, sessionHash string, requestedModel string, excludedIDs map[int64]struct{}, requireCompact bool, requiredCapability OpenAIEndpointCapability) (*AccountSelectionResult, error) { diff --git a/backend/internal/service/ops_settings.go b/backend/internal/service/ops_settings.go index 68c1d9ddea49..23e92e5af21b 100644 --- a/backend/internal/service/ops_settings.go +++ b/backend/internal/service/ops_settings.go @@ -369,6 +369,7 @@ func defaultOpsAdvancedSettings() *OpsAdvancedSettings { Aggregation: OpsAggregationSettings{ AggregationEnabled: false, }, + OpenAIAccountQuotaAutoPause: OpsOpenAIAccountQuotaAutoPauseSettings{}, IgnoreCountTokensErrors: true, // count_tokens 404 是预期行为，默认忽略 IgnoreContextCanceled: true, // Default to true - client disconnects are not errors IgnoreNoAvailableAccounts: false, // Default to false - this is a real routing issue @@ -384,6 +385,8 @@ func normalizeOpsAdvancedSettings(cfg *OpsAdvancedSettings) { if cfg == nil { return } + cfg.OpenAIAccountQuotaAutoPause.DefaultThreshold5h = clampOpsQuotaAutoPauseThreshold(cfg.OpenAIAccountQuotaAutoPause.DefaultThreshold5h) + cfg.OpenAIAccountQuotaAutoPause.DefaultThreshold7d = clampOpsQuotaAutoPauseThreshold(cfg.OpenAIAccountQuotaAutoPause.DefaultThreshold7d) cfg.DataRetention.CleanupSchedule = strings.TrimSpace(cfg.DataRetention.CleanupSchedule) if cfg.DataRetention.CleanupSchedule == "" { cfg.DataRetention.CleanupSchedule = opsCleanupDefaultSchedule @@ -405,6 +408,16 @@ func normalizeOpsAdvancedSettings(cfg *OpsAdvancedSettings) { } } +func clampOpsQuotaAutoPauseThreshold(value float64) float64 { + if value <= 0 { + return 0 + } + if value > 1 { + return 1 + } + return value +} + func validateOpsAdvancedSettings(cfg *OpsAdvancedSettings) error { if cfg == nil { return errors.New("invalid config") @@ -477,6 +490,12 @@ func (s *OpsService) UpdateOpsAdvancedSettings(ctx context.Context, cfg *OpsAdva if err := s.settingRepo.Set(ctx, SettingKeyOpsAdvancedSettings, string(raw)); err != nil { return nil, err } + cacheKey := openAIQuotaAutoPauseSettingsCacheKey(s.settingRepo) + openAIQuotaAutoPauseSettingsSF.Forget(cacheKey) + storeOpenAIQuotaAutoPauseSettingsCache(s.settingRepo, &cachedOpenAIQuotaAutoPauseSettings{ + settings: cfg.OpenAIAccountQuotaAutoPause, + expiresAt: time.Now().Add(openAIQuotaAutoPauseSettingsCacheTTL).UnixNano(), + }) // notify cleanup service to reload schedule/enabled. if s.cleanupReloader != nil { diff --git a/backend/internal/service/ops_settings_advanced_test.go b/backend/internal/service/ops_settings_advanced_test.go index 06cc545bb218..d8598fe0233f 100644 --- a/backend/internal/service/ops_settings_advanced_test.go +++ b/backend/internal/service/ops_settings_advanced_test.go @@ -4,6 +4,8 @@ import ( "context" "encoding/json" "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" ) func TestGetOpsAdvancedSettings_DefaultHidesOpenAITokenStats(t *testing.T) { @@ -95,3 +97,17 @@ func TestGetOpsAdvancedSettings_BackfillsNewDisplayFlagsFromDefaults(t *testing. t.Fatalf("DisplayAlertEvents = false, want true default backfill") } } + +func TestGetOpenAIQuotaAutoPauseSettings_ReadsDefaultsFromOpsAdvancedSettings(t *testing.T) { + repo := newRuntimeSettingRepoStub() + repo.values[SettingKeyOpsAdvancedSettings] = `{"openai_account_quota_auto_pause":{"default_threshold_5h":0.95,"default_threshold_7d":0.9}}` + svc := NewSettingService(repo, &config.Config{}) + + settings := svc.GetOpenAIQuotaAutoPauseSettings(context.Background()) + if settings.DefaultThreshold5h != 0.95 { + t.Fatalf("DefaultThreshold5h = %v, want 0.95", settings.DefaultThreshold5h) + } + if settings.DefaultThreshold7d != 0.9 { + t.Fatalf("DefaultThreshold7d = %v, want 0.9", settings.DefaultThreshold7d) + } +} diff --git a/backend/internal/service/ops_settings_models.go b/backend/internal/service/ops_settings_models.go index fa18b05fb7cb..4d459e2b9496 100644 --- a/backend/internal/service/ops_settings_models.go +++ b/backend/internal/service/ops_settings_models.go @@ -92,17 +92,23 @@ type OpsAlertRuntimeSettings struct { // OpsAdvancedSettings stores advanced ops configuration (data retention, aggregation). type OpsAdvancedSettings struct { - DataRetention OpsDataRetentionSettings `json:"data_retention"` - Aggregation OpsAggregationSettings `json:"aggregation"` - IgnoreCountTokensErrors bool `json:"ignore_count_tokens_errors"` - IgnoreContextCanceled bool `json:"ignore_context_canceled"` - IgnoreNoAvailableAccounts bool `json:"ignore_no_available_accounts"` - IgnoreInvalidApiKeyErrors bool `json:"ignore_invalid_api_key_errors"` - IgnoreInsufficientBalanceErrors bool `json:"ignore_insufficient_balance_errors"` - DisplayOpenAITokenStats bool `json:"display_openai_token_stats"` - DisplayAlertEvents bool `json:"display_alert_events"` - AutoRefreshEnabled bool `json:"auto_refresh_enabled"` - AutoRefreshIntervalSec int `json:"auto_refresh_interval_seconds"` + DataRetention OpsDataRetentionSettings `json:"data_retention"` + Aggregation OpsAggregationSettings `json:"aggregation"` + OpenAIAccountQuotaAutoPause OpsOpenAIAccountQuotaAutoPauseSettings `json:"openai_account_quota_auto_pause"` + IgnoreCountTokensErrors bool `json:"ignore_count_tokens_errors"` + IgnoreContextCanceled bool `json:"ignore_context_canceled"` + IgnoreNoAvailableAccounts bool `json:"ignore_no_available_accounts"` + IgnoreInvalidApiKeyErrors bool `json:"ignore_invalid_api_key_errors"` + IgnoreInsufficientBalanceErrors bool `json:"ignore_insufficient_balance_errors"` + DisplayOpenAITokenStats bool `json:"display_openai_token_stats"` + DisplayAlertEvents bool `json:"display_alert_events"` + AutoRefreshEnabled bool `json:"auto_refresh_enabled"` + AutoRefreshIntervalSec int `json:"auto_refresh_interval_seconds"` +} + +type OpsOpenAIAccountQuotaAutoPauseSettings struct { + DefaultThreshold5h float64 `json:"default_threshold_5h"` + DefaultThreshold7d float64 `json:"default_threshold_7d"` } type OpsDataRetentionSettings struct { diff --git a/backend/internal/service/setting_service.go b/backend/internal/service/setting_service.go index 08c0d045c6e0..40e98b88a5ba 100644 --- a/backend/internal/service/setting_service.go +++ b/backend/internal/service/setting_service.go @@ -14,6 +14,7 @@ import ( "sort" "strconv" "strings" + "sync" "sync/atomic" "time" @@ -137,6 +138,11 @@ type cachedOpenAICodexUserAgent struct { expiresAt int64 // unix nano } +type cachedOpenAIQuotaAutoPauseSettings struct { + settings OpsOpenAIAccountQuotaAutoPauseSettings + expiresAt int64 +} + const openAICodexUserAgentCacheTTL = 60 * time.Second const openAICodexUserAgentErrorTTL = 5 * time.Second const openAICodexUserAgentDBTimeout = 5 * time.Second @@ -152,6 +158,33 @@ const openAIAllowCodexPluginCacheTTL = 60 * time.Second const openAIAllowCodexPluginErrorTTL = 5 * time.Second const openAIAllowCodexPluginDBTimeout = 5 * time.Second +const openAIQuotaAutoPauseSettingsCacheTTL = 60 * time.Second +const openAIQuotaAutoPauseSettingsErrorTTL = 5 * time.Second +const openAIQuotaAutoPauseSettingsDBTimeout = 5 * time.Second + +var openAIQuotaAutoPauseSettingsCache sync.Map // map[string]*cachedOpenAIQuotaAutoPauseSettings +var openAIQuotaAutoPauseSettingsSF singleflight.Group + +func openAIQuotaAutoPauseSettingsCacheKey(repo SettingRepository) string { + if repo == nil { + return "nil" + } + return fmt.Sprintf("%T:%p", repo, repo) +} + +func loadOpenAIQuotaAutoPauseSettingsCache(repo SettingRepository) (*cachedOpenAIQuotaAutoPauseSettings, bool) { + value, ok := openAIQuotaAutoPauseSettingsCache.Load(openAIQuotaAutoPauseSettingsCacheKey(repo)) + if !ok || value == nil { + return nil, false + } + cached, ok := value.(*cachedOpenAIQuotaAutoPauseSettings) + return cached, ok && cached != nil +} + +func storeOpenAIQuotaAutoPauseSettingsCache(repo SettingRepository, cached *cachedOpenAIQuotaAutoPauseSettings) { + openAIQuotaAutoPauseSettingsCache.Store(openAIQuotaAutoPauseSettingsCacheKey(repo), cached) +} + // DefaultSubscriptionGroupReader validates group references used by default subscriptions. type DefaultSubscriptionGroupReader interface { GetByID(ctx context.Context, id int64) (*Group, error) @@ -2027,6 +2060,9 @@ func (s *SettingService) refreshCachedSettings(settings *SystemSettings) { enabled: settings.OpenAIAdvancedSchedulerEnabled, expiresAt: time.Now().Add(openAIAdvancedSchedulerSettingCacheTTL).UnixNano(), }) + cacheKey := openAIQuotaAutoPauseSettingsCacheKey(s.settingRepo) + openAIQuotaAutoPauseSettingsSF.Forget(cacheKey) + openAIQuotaAutoPauseSettingsCache.Delete(cacheKey) if s.cfg != nil { s.cfg.SetTrustForwardedIPForAPIKeyACL(settings.APIKeyACLTrustForwardedIP) } @@ -4448,6 +4484,51 @@ func (s *SettingService) GetClaudeCodeVersionBounds(ctx context.Context) (min, m return b.min, b.max } +func (s *SettingService) GetOpenAIQuotaAutoPauseSettings(ctx context.Context) OpsOpenAIAccountQuotaAutoPauseSettings { + if cached, ok := loadOpenAIQuotaAutoPauseSettingsCache(s.settingRepo); ok { + if time.Now().UnixNano() < cached.expiresAt { + return cached.settings + } + } + + cacheKey := openAIQuotaAutoPauseSettingsCacheKey(s.settingRepo) + result, _, _ := openAIQuotaAutoPauseSettingsSF.Do(cacheKey, func() (any, error) { + if cached, ok := loadOpenAIQuotaAutoPauseSettingsCache(s.settingRepo); ok { + if time.Now().UnixNano() < cached.expiresAt { + return cached.settings, nil + } + } + + settings := OpsOpenAIAccountQuotaAutoPauseSettings{} + ttl := openAIQuotaAutoPauseSettingsCacheTTL + if s != nil && s.settingRepo != nil { + dbCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), openAIQuotaAutoPauseSettingsDBTimeout) + defer cancel() + raw, err := s.settingRepo.GetValue(dbCtx, SettingKeyOpsAdvancedSettings) + if err == nil { + cfg := defaultOpsAdvancedSettings() + if strings.TrimSpace(raw) != "" { + if jsonErr := json.Unmarshal([]byte(raw), cfg); jsonErr == nil { + normalizeOpsAdvancedSettings(cfg) + } + } + settings = cfg.OpenAIAccountQuotaAutoPause + } else { + ttl = openAIQuotaAutoPauseSettingsErrorTTL + } + } + + storeOpenAIQuotaAutoPauseSettingsCache(s.settingRepo, &cachedOpenAIQuotaAutoPauseSettings{ + settings: settings, + expiresAt: time.Now().Add(ttl).UnixNano(), + }) + return settings, nil + }) + + settings, _ := result.(OpsOpenAIAccountQuotaAutoPauseSettings) + return settings +} + // GetRectifierSettings 获取请求整流器配置 func (s *SettingService) GetRectifierSettings(ctx context.Context) (*RectifierSettings, error) { value, err := s.settingRepo.GetValue(ctx, SettingKeyRectifierSettings) diff --git a/frontend/src/components/account/EditAccountModal.vue b/frontend/src/components/account/EditAccountModal.vue index 515ba58f7c7e..470c0bfbdf88 100644 --- a/frontend/src/components/account/EditAccountModal.vue +++ b/frontend/src/components/account/EditAccountModal.vue @@ -1787,6 +1787,38 @@ +
+
+ + +

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

+
+
+ + +

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

+
+
+
([]) const customErrorCodeInput = ref(null) const interceptWarmupRequests = ref(false) const autoPauseOnExpired = ref(false) +const autoPause5hThreshold = ref(null) +const autoPause7dThreshold = ref(null) const mixedScheduling = ref(false) // For antigravity accounts: enable mixed scheduling const allowOverages = ref(false) // For antigravity accounts: enable AI Credits overages const antigravityModelRestrictionMode = ref<'whitelist' | 'mapping'>('whitelist') @@ -2862,9 +2896,11 @@ const syncFormFromAccount = (newAccount: Account | null) => { // Load mixed scheduling setting (only for antigravity accounts) mixedScheduling.value = false allowOverages.value = false - const extra = newAccount.extra as Record | undefined - mixedScheduling.value = extra?.mixed_scheduling === true - allowOverages.value = extra?.allow_overages === true + const extra = newAccount.extra as Record | undefined + mixedScheduling.value = extra?.mixed_scheduling === true + allowOverages.value = extra?.allow_overages === true + autoPause5hThreshold.value = typeof extra?.auto_pause_5h_threshold === 'number' ? extra.auto_pause_5h_threshold * 100 : null + autoPause7dThreshold.value = typeof extra?.auto_pause_7d_threshold === 'number' ? extra.auto_pause_7d_threshold * 100 : null // Load OpenAI passthrough toggle (OpenAI OAuth/API Key) openaiPassthroughEnabled.value = false @@ -3987,9 +4023,9 @@ const handleSubmit = async () => { } // For OpenAI OAuth/API Key accounts, handle passthrough mode in extra - if (props.account.platform === 'openai' && (props.account.type === 'oauth' || props.account.type === 'apikey')) { - const currentExtra = (props.account.extra as Record) || {} - const newExtra: Record = { ...currentExtra } + if (props.account.platform === 'openai' && (props.account.type === 'oauth' || props.account.type === 'apikey')) { + const currentExtra = (props.account.extra as Record) || {} + const newExtra: Record = { ...currentExtra } const hadCodexCLIOnlyEnabled = currentExtra.codex_cli_only === true if (props.account.type === 'oauth') { newExtra.openai_oauth_responses_websockets_v2_mode = openaiOAuthResponsesWebSocketV2Mode.value @@ -4011,15 +4047,25 @@ const handleSubmit = async () => { } else { newExtra.openai_compact_mode = openAICompactMode.value } - if (props.account.type === 'apikey') { + if (props.account.type === 'apikey') { if (!openAITextGenerationCapabilityEnabled.value || openAIResponsesMode.value === 'auto') { delete newExtra.openai_responses_mode } else { newExtra.openai_responses_mode = openAIResponsesMode.value } - } - - delete newExtra.codex_image_generation_bridge_enabled + } + if (autoPause5hThreshold.value != null && autoPause5hThreshold.value > 0) { + newExtra.auto_pause_5h_threshold = autoPause5hThreshold.value / 100 + } else { + delete newExtra.auto_pause_5h_threshold + } + if (autoPause7dThreshold.value != null && autoPause7dThreshold.value > 0) { + newExtra.auto_pause_7d_threshold = autoPause7dThreshold.value / 100 + } else { + delete newExtra.auto_pause_7d_threshold + } + + delete newExtra.codex_image_generation_bridge_enabled if (codexImageGenerationBridgeMode.value === 'inherit') { delete newExtra.codex_image_generation_bridge } else { diff --git a/frontend/src/components/account/__tests__/EditAccountModal.spec.ts b/frontend/src/components/account/__tests__/EditAccountModal.spec.ts index 4561924f675d..6db6383129ad 100644 --- a/frontend/src/components/account/__tests__/EditAccountModal.spec.ts +++ b/frontend/src/components/account/__tests__/EditAccountModal.spec.ts @@ -330,6 +330,28 @@ describe('EditAccountModal', () => { ]) }) + it('submits OpenAI quota auto-pause thresholds in extra', async () => { + const account = buildAccount() + account.extra = { + auto_pause_5h_threshold: 0.9, + auto_pause_7d_threshold: 0.8 + } + updateAccountMock.mockReset() + checkMixedChannelRiskMock.mockReset() + checkMixedChannelRiskMock.mockResolvedValue({ has_risk: false }) + updateAccountMock.mockResolvedValue(account) + + const wrapper = mountModal(account) + + await wrapper.get('[data-testid="auto-pause-5h-threshold"]').setValue('95') + await wrapper.get('[data-testid="auto-pause-7d-threshold"]').setValue('96') + await wrapper.get('form#edit-account-form').trigger('submit.prevent') + + expect(updateAccountMock).toHaveBeenCalledTimes(1) + expect(updateAccountMock.mock.calls[0]?.[1]?.extra?.auto_pause_5h_threshold).toBe(0.95) + expect(updateAccountMock.mock.calls[0]?.[1]?.extra?.auto_pause_7d_threshold).toBe(0.96) + }) + it('keeps at least one OpenAI APIKey endpoint capability selected', async () => { const account = buildAccount() updateAccountMock.mockReset() diff --git a/frontend/src/i18n/locales/en.ts b/frontend/src/i18n/locales/en.ts index b2aeb2f83f58..fa2e5a920a9f 100644 --- a/frontend/src/i18n/locales/en.ts +++ b/frontend/src/i18n/locales/en.ts @@ -3475,6 +3475,9 @@ export default { 'When enabled, warmup requests like title generation will return mock responses without consuming upstream tokens', autoPauseOnExpired: 'Auto Pause On Expired', autoPauseOnExpiredDesc: 'When enabled, the account will auto pause scheduling after it expires', + autoPause5hThreshold: '5h Usage Threshold (%)', + autoPause7dThreshold: '7d Usage Threshold (%)', + autoPauseThresholdHint: 'Leave empty or set 0 to disable. Reaching the threshold only skips the account during scheduling and does not modify schedulable.', // Quota control (Anthropic OAuth/SetupToken only) quotaControl: { title: 'Quota Control', diff --git a/frontend/src/i18n/locales/zh.ts b/frontend/src/i18n/locales/zh.ts index 85d1feee4565..2364f9c4b956 100644 --- a/frontend/src/i18n/locales/zh.ts +++ b/frontend/src/i18n/locales/zh.ts @@ -3613,6 +3613,9 @@ export default { interceptWarmupRequestsDesc: '启用后，标题生成等预热请求将返回 mock 响应，不消耗上游 token', autoPauseOnExpired: '过期自动暂停调度', autoPauseOnExpiredDesc: '启用后，账号过期将自动暂停调度', + autoPause5hThreshold: '5h 用量阈值(%)', + autoPause7dThreshold: '7d 用量阈值(%)', + autoPauseThresholdHint: '填 0 或留空表示不启用；达到阈值后仅在调度时跳过账号，不修改 schedulable。', // Quota control (Anthropic OAuth/SetupToken only) quotaControl: { title: '配额控制', From 8b7a8227060e3eac20b5c0e331b50b2dc99f5e13 Mon Sep 17 00:00:00 2001 From: wucm667 Date: Fri, 29 May 2026 12:20:30 +0800 Subject: [PATCH 04/36] fix(account): address review on OpenAI quota auto-pause - gate previous_response_id sticky path with quota auto-pause check at both the snapshot and DB-recheck stages (previously bypassed, #1) - skip pausing when the usage window already reset to avoid a stale stuck-pause; carry codex_*_reset_at / reset_after_seconds / codex_usage_updated_at through the scheduler snapshot whitelist (#2) - remove the incomplete limit mode; percentage threshold only (#3) - add global default 5h/7d threshold inputs to the Ops settings dialog with validation and en/zh i18n (#4) - downgrade account_auto_paused_by_quota log from Info to Debug; it fires per-candidate on the scheduling hot path (#5) Co-Authored-By: Claude Opus 4.8 --- .../internal/repository/scheduler_cache.go | 7 +- .../repository/scheduler_cache_unit_test.go | 18 +++-- .../service/openai_account_scheduler_test.go | 58 +++++++++++++-- .../service/openai_gateway_service.go | 72 +++++++++++++------ .../service/openai_ws_account_sticky_test.go | 40 +++++++++++ .../internal/service/openai_ws_forwarder.go | 10 +++ frontend/src/api/admin/ops.ts | 6 ++ frontend/src/i18n/locales/en.ts | 8 ++- frontend/src/i18n/locales/zh.ts | 8 ++- .../ops/components/OpsSettingsDialog.vue | 65 +++++++++++++++++ 10 files changed, 257 insertions(+), 35 deletions(-) diff --git a/backend/internal/repository/scheduler_cache.go b/backend/internal/repository/scheduler_cache.go index ec8c72dca978..ff3c4301388b 100644 --- a/backend/internal/repository/scheduler_cache.go +++ b/backend/internal/repository/scheduler_cache.go @@ -550,10 +550,13 @@ func filterSchedulerExtra(extra map[string]any) map[string]any { "openai_responses_supported", "codex_5h_used_percent", "codex_7d_used_percent", + "codex_5h_reset_at", + "codex_7d_reset_at", + "codex_5h_reset_after_seconds", + "codex_7d_reset_after_seconds", + "codex_usage_updated_at", "auto_pause_5h_threshold", "auto_pause_7d_threshold", - "auto_pause_5h_limit", - "auto_pause_7d_limit", } filtered := make(map[string]any) for _, key := range keys { diff --git a/backend/internal/repository/scheduler_cache_unit_test.go b/backend/internal/repository/scheduler_cache_unit_test.go index fabc6bad242c..9e4ec23e9aa9 100644 --- a/backend/internal/repository/scheduler_cache_unit_test.go +++ b/backend/internal/repository/scheduler_cache_unit_test.go @@ -80,10 +80,15 @@ func TestBuildSchedulerMetadataAccount_KeepsQuotaAutoPauseFields(t *testing.T) { account := service.Account{ ID: 88, Extra: map[string]any{ - "codex_5h_used_percent": 12.34, - "codex_7d_used_percent": 56.78, - "auto_pause_5h_threshold": 0.95, - "auto_pause_7d_threshold": 0.96, + "codex_5h_used_percent": 12.34, + "codex_7d_used_percent": 56.78, + "codex_5h_reset_at": "2026-05-29T10:00:00Z", + "codex_7d_reset_at": "2026-06-01T10:00:00Z", + "codex_5h_reset_after_seconds": 300, + "codex_7d_reset_after_seconds": 600, + "codex_usage_updated_at": "2026-05-29T09:00:00Z", + "auto_pause_5h_threshold": 0.95, + "auto_pause_7d_threshold": 0.96, }, } @@ -91,6 +96,11 @@ func TestBuildSchedulerMetadataAccount_KeepsQuotaAutoPauseFields(t *testing.T) { require.Equal(t, 12.34, got.Extra["codex_5h_used_percent"]) require.Equal(t, 56.78, got.Extra["codex_7d_used_percent"]) + require.Equal(t, "2026-05-29T10:00:00Z", got.Extra["codex_5h_reset_at"]) + require.Equal(t, "2026-06-01T10:00:00Z", got.Extra["codex_7d_reset_at"]) + require.Equal(t, 300, got.Extra["codex_5h_reset_after_seconds"]) + require.Equal(t, 600, got.Extra["codex_7d_reset_after_seconds"]) + require.Equal(t, "2026-05-29T09:00:00Z", got.Extra["codex_usage_updated_at"]) require.Equal(t, 0.95, got.Extra["auto_pause_5h_threshold"]) require.Equal(t, 0.96, got.Extra["auto_pause_7d_threshold"]) } diff --git a/backend/internal/service/openai_account_scheduler_test.go b/backend/internal/service/openai_account_scheduler_test.go index 37810870bc37..531769a73ef6 100644 --- a/backend/internal/service/openai_account_scheduler_test.go +++ b/backend/internal/service/openai_account_scheduler_test.go @@ -704,7 +704,6 @@ func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_AutoPauseBy5hT Extra: map[string]any{ "codex_5h_used_percent": 95.0, "auto_pause_5h_threshold": 0.95, - "auto_pause_5h_limit": 100, }, } secondary := Account{ID: 35002, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5} @@ -729,7 +728,6 @@ func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_AllowsBelow5hT Extra: map[string]any{ "codex_5h_used_percent": 80.0, "auto_pause_5h_threshold": 0.95, - "auto_pause_5h_limit": 100, }, } secondary := Account{ID: 35102, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5} @@ -754,7 +752,6 @@ func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_AutoPauseBy7dT Extra: map[string]any{ "codex_7d_used_percent": 95.0, "auto_pause_7d_threshold": 0.95, - "auto_pause_7d_limit": 200, }, } secondary := Account{ID: 35202, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5} @@ -790,7 +787,6 @@ func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_UsesGlobalDefa Priority: 0, Extra: map[string]any{ "codex_5h_used_percent": 95.0, - "auto_pause_5h_limit": 100, }, } secondary := Account{ID: 35402, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5} @@ -802,6 +798,60 @@ func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_UsesGlobalDefa require.Equal(t, int64(35402), account.ID) } +func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_StaleUsageWindowResetSkipsPause(t *testing.T) { + ctx := context.Background() + // Usage is over threshold but the window's reset time has already passed, so the + // cached percentage is stale (the real window rolled over) and the account must NOT + // stay paused — otherwise it could be skipped forever with no traffic to refresh it. + primary := Account{ + ID: 35501, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 0, + Extra: map[string]any{ + "codex_5h_used_percent": 99.0, + "auto_pause_5h_threshold": 0.95, + "codex_5h_reset_at": time.Now().Add(-time.Minute).Format(time.RFC3339), + }, + } + secondary := Account{ID: 35502, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5} + svc := &OpenAIGatewayService{accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}}, cfg: &config.Config{}} + + account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.1", nil) + require.NoError(t, err) + require.NotNil(t, account) + require.Equal(t, int64(35501), account.ID) +} + +func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_FreshUsageWindowStillPauses(t *testing.T) { + ctx := context.Background() + // Same as above but the window has not reset yet, so the account stays paused. + primary := Account{ + ID: 35601, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 0, + Extra: map[string]any{ + "codex_5h_used_percent": 99.0, + "auto_pause_5h_threshold": 0.95, + "codex_5h_reset_at": time.Now().Add(time.Hour).Format(time.RFC3339), + }, + } + secondary := Account{ID: 35602, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5} + svc := &OpenAIGatewayService{accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}}, cfg: &config.Config{}} + + account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.1", nil) + require.NoError(t, err) + require.NotNil(t, account) + require.Equal(t, int64(35602), account.ID) +} + func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_SkipsFreshlyRateLimitedSnapshotCandidate(t *testing.T) { ctx := context.Background() groupID := int64(10102) diff --git a/backend/internal/service/openai_gateway_service.go b/backend/internal/service/openai_gateway_service.go index 268f985cec80..e2534cc2bbf9 100644 --- a/backend/internal/service/openai_gateway_service.go +++ b/backend/internal/service/openai_gateway_service.go @@ -1328,13 +1328,12 @@ func isOpenAIAccountEligibleForRequest(ctx context.Context, account *Account, re return false } if paused, reason := shouldAutoPauseOpenAIAccountByQuota(ctx, account); paused { - slog.Info("account_auto_paused_by_quota", + // Debug level: this fires per-candidate on the scheduling hot path, so Info + // would amplify into log spam once several accounts cross the threshold. + slog.Debug("account_auto_paused_by_quota", "account_id", account.ID, - "usage_5h_percent", readOpenAIQuotaUsedPercent(account.Extra, "5h"), - "usage_7d_percent", readOpenAIQuotaUsedPercent(account.Extra, "7d"), - "threshold_type", reason.window, + "window", reason.window, "threshold", reason.threshold, - "limit", reason.limit, "utilization", reason.utilization, ) return false @@ -1354,7 +1353,6 @@ func isOpenAIAccountEligibleForRequest(ctx context.Context, account *Account, re type openAIQuotaAutoPauseDecision struct { window string threshold float64 - limit float64 utilization float64 } @@ -1363,18 +1361,15 @@ func shouldAutoPauseOpenAIAccountByQuota(ctx context.Context, account *Account) return false, openAIQuotaAutoPauseDecision{} } threshold5h, threshold7d := resolveOpenAIQuotaAutoPauseThresholds(ctx, account) + now := time.Now() if threshold5h > 0 { - if utilization, limit, ok := resolveOpenAIQuotaUtilization(account.Extra, "5h"); ok { - if utilization >= threshold5h { - return true, openAIQuotaAutoPauseDecision{window: "5h", threshold: threshold5h, limit: limit, utilization: utilization} - } + if utilization, ok := resolveOpenAIQuotaUtilization(account.Extra, "5h", now); ok && utilization >= threshold5h { + return true, openAIQuotaAutoPauseDecision{window: "5h", threshold: threshold5h, utilization: utilization} } } if threshold7d > 0 { - if utilization, limit, ok := resolveOpenAIQuotaUtilization(account.Extra, "7d"); ok { - if utilization >= threshold7d { - return true, openAIQuotaAutoPauseDecision{window: "7d", threshold: threshold7d, limit: limit, utilization: utilization} - } + if utilization, ok := resolveOpenAIQuotaUtilization(account.Extra, "7d", now); ok && utilization >= threshold7d { + return true, openAIQuotaAutoPauseDecision{window: "7d", threshold: threshold7d, utilization: utilization} } } return false, openAIQuotaAutoPauseDecision{} @@ -1431,18 +1426,49 @@ func resolveAccountExtraNumber(extra map[string]any, keys ...string) (float64, b return 0, false } -func resolveOpenAIQuotaUtilization(extra map[string]any, window string) (float64, float64, bool) { - limitKeys := []string{"auto_pause_" + window + "_limit", "quota_" + window + "_limit", window + "_limit"} - if limit, ok := resolveAccountExtraNumber(extra, limitKeys...); ok && limit > 0 { - if usage, ok := resolveAccountExtraNumber(extra, "usage_"+window); ok && usage >= 0 { - return usage / limit, limit, true - } - } +// resolveOpenAIQuotaUtilization returns the current utilization ratio (0..1) for the +// given Codex usage window. ok=false means there is no usable signal to pause on: +// either no snapshot exists, or the window has already rolled over so the cached +// percentage is stale. The stale guard matters because a paused account stops +// receiving requests, so its snapshot is never refreshed from upstream headers — +// without this check an old used_percent would keep the account paused forever even +// after the real window reset. +func resolveOpenAIQuotaUtilization(extra map[string]any, window string, now time.Time) (float64, bool) { usedPercent := readOpenAIQuotaUsedPercent(extra, window) if usedPercent <= 0 { - return 0, 0, false + return 0, false + } + if openAIQuotaWindowReset(extra, window, now) { + return 0, false + } + return usedPercent / 100, true +} + +// openAIQuotaWindowReset reports whether the Codex usage window's reset time has +// already passed relative to now. It prefers the absolute codex__reset_at +// timestamp and falls back to codex__reset_after_seconds anchored at +// codex_usage_updated_at, mirroring AccountUsageService's window-progress logic. +func openAIQuotaWindowReset(extra map[string]any, window string, now time.Time) bool { + if len(extra) == 0 { + return false + } + if resetAtRaw, ok := extra["codex_"+window+"_reset_at"]; ok { + if resetAt, err := parseTime(fmt.Sprint(resetAtRaw)); err == nil { + return !now.Before(resetAt) + } + } + resetAfter := parseExtraInt(extra["codex_"+window+"_reset_after_seconds"]) + if resetAfter <= 0 { + return false + } + base := now + if updatedRaw, ok := extra["codex_usage_updated_at"]; ok { + if updatedAt, err := parseTime(fmt.Sprint(updatedRaw)); err == nil { + base = updatedAt + } } - return usedPercent / 100, 100, true + resetAt := base.Add(time.Duration(resetAfter) * time.Second) + return !now.Before(resetAt) } func readOpenAIQuotaUsedPercent(extra map[string]any, window string) float64 { diff --git a/backend/internal/service/openai_ws_account_sticky_test.go b/backend/internal/service/openai_ws_account_sticky_test.go index c8b28a465b43..6fc44298a400 100644 --- a/backend/internal/service/openai_ws_account_sticky_test.go +++ b/backend/internal/service/openai_ws_account_sticky_test.go @@ -48,6 +48,46 @@ func TestOpenAIGatewayService_SelectAccountByPreviousResponseID_Hit(t *testing.T } } +func TestOpenAIGatewayService_SelectAccountByPreviousResponseID_QuotaAutoPausedMiss(t *testing.T) { + ctx := context.Background() + groupID := int64(23) + account := Account{ + ID: 77, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 2, + Extra: map[string]any{ + "openai_apikey_responses_websockets_v2_enabled": true, + "codex_5h_used_percent": 96.0, + "auto_pause_5h_threshold": 0.95, + }, + } + cache := &stubGatewayCache{} + store := NewOpenAIWSStateStore(cache) + cfg := newOpenAIWSV2TestConfig() + svc := &OpenAIGatewayService{ + accountRepo: stubOpenAIAccountRepo{accounts: []Account{account}}, + cache: cache, + cfg: cfg, + concurrencyService: NewConcurrencyService(stubConcurrencyCache{}), + openaiWSStateStore: store, + } + + require.NoError(t, store.BindResponseAccount(ctx, groupID, "resp_prev_quota", account.ID, time.Hour)) + + selection, err := svc.SelectAccountByPreviousResponseID(ctx, &groupID, "resp_prev_quota", "gpt-5.1", nil, false) + require.NoError(t, err) + require.Nil(t, selection, "超过 5h 配额阈值的账号不应继续命中 previous_response_id 粘连") + + // Auto-pause is transient, so the binding is preserved: the chain can resume on the + // same account once the quota window resets. + boundAccountID, getErr := store.GetResponseAccount(ctx, groupID, "resp_prev_quota") + require.NoError(t, getErr) + require.Equal(t, account.ID, boundAccountID) +} + func TestOpenAIGatewayService_SelectAccountByPreviousResponseID_RateLimitedMiss(t *testing.T) { ctx := context.Background() groupID := int64(23) diff --git a/backend/internal/service/openai_ws_forwarder.go b/backend/internal/service/openai_ws_forwarder.go index 6eea01911701..878ff4869936 100644 --- a/backend/internal/service/openai_ws_forwarder.go +++ b/backend/internal/service/openai_ws_forwarder.go @@ -4045,6 +4045,13 @@ func (s *OpenAIGatewayService) selectAccountByPreviousResponseIDForCapability( if !account.SupportsOpenAIEndpointCapability(requiredCapability) { return nil, nil } + // Quota auto-pause must also gate the previous_response_id sticky path; otherwise an + // account over its 5h/7d threshold keeps serving the same response chain even though + // normal scheduling skips it. Pause is transient, so fall through to normal scheduling + // without deleting the binding (the window may reset before the next turn). + if paused, _ := shouldAutoPauseOpenAIAccountByQuota(ctx, account); paused { + return nil, nil + } if s.schedulerSnapshot != nil && s.accountRepo != nil { latest, latestErr := s.accountRepo.GetByID(ctx, account.ID) if latestErr != nil || latest == nil { @@ -4061,6 +4068,9 @@ func (s *OpenAIGatewayService) selectAccountByPreviousResponseIDForCapability( if !latest.SupportsOpenAIEndpointCapability(requiredCapability) { return nil, nil } + if paused, _ := shouldAutoPauseOpenAIAccountByQuota(ctx, latest); paused { + return nil, nil + } if s.isOpenAIAccountRuntimeBlocked(latest) { _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID) return nil, nil diff --git a/frontend/src/api/admin/ops.ts b/frontend/src/api/admin/ops.ts index 6923566867a9..847fc8c92bc7 100644 --- a/frontend/src/api/admin/ops.ts +++ b/frontend/src/api/admin/ops.ts @@ -778,9 +778,15 @@ export interface OpsAlertRuntimeSettings { thresholds: OpsMetricThresholds // 指标阈值配置 } +export interface OpsOpenAIAccountQuotaAutoPauseSettings { + default_threshold_5h: number // 0~1，0 表示不启用全局默认 5h 阈值 + default_threshold_7d: number // 0~1，0 表示不启用全局默认 7d 阈值 +} + export interface OpsAdvancedSettings { data_retention: OpsDataRetentionSettings aggregation: OpsAggregationSettings + openai_account_quota_auto_pause: OpsOpenAIAccountQuotaAutoPauseSettings ignore_count_tokens_errors: boolean ignore_context_canceled: boolean ignore_no_available_accounts: boolean diff --git a/frontend/src/i18n/locales/en.ts b/frontend/src/i18n/locales/en.ts index fa2e5a920a9f..8ab90961e7c4 100644 --- a/frontend/src/i18n/locales/en.ts +++ b/frontend/src/i18n/locales/en.ts @@ -5193,6 +5193,11 @@ export default { aggregation: 'Pre-aggregation Tasks', enableAggregation: 'Enable Pre-aggregation', aggregationHint: 'Pre-aggregation improves query performance for long time windows', + openaiQuotaAutoPause: 'OpenAI Account Quota Auto-pause', + openaiQuotaAutoPauseHint: 'When an OpenAI account reaches its 5h / 7d usage threshold, the scheduler skips it automatically and resumes once the window rolls over. Per-account thresholds take precedence over this global default.', + openaiQuotaAutoPauseDefault5h: 'Default 5h usage threshold (%)', + openaiQuotaAutoPauseDefault7d: 'Default 7d usage threshold (%)', + openaiQuotaAutoPauseThresholdHint: 'Value 0-100; leave blank or 0 to disable the global default threshold.', errorFiltering: 'Error Filtering', ignoreCountTokensErrors: 'Ignore count_tokens errors', ignoreCountTokensErrorsHint: 'When enabled, errors from count_tokens requests will not be written to the error log.', @@ -5223,7 +5228,8 @@ export default { slaMinPercentRange: 'SLA minimum percentage must be between 0 and 100', ttftP99MaxRange: 'TTFT P99 maximum must be a number ≥ 0', requestErrorRateMaxRange: 'Request error rate maximum must be between 0 and 100', - upstreamErrorRateMaxRange: 'Upstream error rate maximum must be between 0 and 100' + upstreamErrorRateMaxRange: 'Upstream error rate maximum must be between 0 and 100', + openaiQuotaAutoPauseRange: 'OpenAI quota auto-pause threshold must be between 0 and 100' } }, concurrency: { diff --git a/frontend/src/i18n/locales/zh.ts b/frontend/src/i18n/locales/zh.ts index 2364f9c4b956..4f1d1f13975e 100644 --- a/frontend/src/i18n/locales/zh.ts +++ b/frontend/src/i18n/locales/zh.ts @@ -5352,6 +5352,11 @@ export default { aggregation: '预聚合任务', enableAggregation: '启用预聚合任务', aggregationHint: '预聚合可提升长时间窗口查询性能', + openaiQuotaAutoPause: 'OpenAI 账号配额自动暂停', + openaiQuotaAutoPauseHint: '当 OpenAI 账号 5h / 7d 用量达到阈值时，调度会自动跳过该账号；窗口滚动后自动恢复。账号级阈值优先于此全局默认值。', + openaiQuotaAutoPauseDefault5h: '默认 5h 用量阈值 (%)', + openaiQuotaAutoPauseDefault7d: '默认 7d 用量阈值 (%)', + openaiQuotaAutoPauseThresholdHint: '取值 0-100，留空或 0 表示不启用全局默认阈值。', errorFiltering: '错误过滤', ignoreCountTokensErrors: '忽略 count_tokens 错误', ignoreCountTokensErrorsHint: '启用后，count_tokens 请求的错误将不会写入错误日志。', @@ -5383,7 +5388,8 @@ export default { slaMinPercentRange: 'SLA最低百分比必须在0-100之间', ttftP99MaxRange: 'TTFT P99最大值必须大于等于0', requestErrorRateMaxRange: '请求错误率最大值必须在0-100之间', - upstreamErrorRateMaxRange: '上游错误率最大值必须在0-100之间' + upstreamErrorRateMaxRange: '上游错误率最大值必须在0-100之间', + openaiQuotaAutoPauseRange: 'OpenAI 配额自动暂停阈值必须在 0-100 之间' } }, concurrency: { diff --git a/frontend/src/views/admin/ops/components/OpsSettingsDialog.vue b/frontend/src/views/admin/ops/components/OpsSettingsDialog.vue index 5dba5b1dc4c3..bfb7a65f85cf 100644 --- a/frontend/src/views/admin/ops/components/OpsSettingsDialog.vue +++ b/frontend/src/views/admin/ops/components/OpsSettingsDialog.vue @@ -50,6 +50,10 @@ async function loadAllSettings() { runtimeSettings.value = runtime emailConfig.value = email advancedSettings.value = advanced + // 兼容旧 payload：后端未返回该字段时补默认值，保证表单可绑定 + if (advancedSettings.value && !advancedSettings.value.openai_account_quota_auto_pause) { + advancedSettings.value.openai_account_quota_auto_pause = { default_threshold_5h: 0, default_threshold_7d: 0 } + } // 如果后端返回了阈值，使用后端的值；否则保持默认值 if (thresholds && Object.keys(thresholds).length > 0) { metricThresholds.value = { @@ -119,6 +123,28 @@ function removeRecipient(target: 'alert' | 'report', email: string) { if (idx >= 0) list.splice(idx, 1) } +// OpenAI 账号配额自动暂停：后端按 0~1 分数存储，UI 按百分比(0~100)展示 +const quotaAutoPause5hPercent = computed({ + get() { + const v = advancedSettings.value?.openai_account_quota_auto_pause?.default_threshold_5h + return v && v > 0 ? Math.round(v * 1000) / 10 : null + }, + set(val) { + if (!advancedSettings.value?.openai_account_quota_auto_pause) return + advancedSettings.value.openai_account_quota_auto_pause.default_threshold_5h = val != null && val > 0 ? val / 100 : 0 + } +}) +const quotaAutoPause7dPercent = computed({ + get() { + const v = advancedSettings.value?.openai_account_quota_auto_pause?.default_threshold_7d + return v && v > 0 ? Math.round(v * 1000) / 10 : null + }, + set(val) { + if (!advancedSettings.value?.openai_account_quota_auto_pause) return + advancedSettings.value.openai_account_quota_auto_pause.default_threshold_7d = val != null && val > 0 ? val / 100 : 0 + } +}) + // 验证 const validation = computed(() => { const errors: string[] = [] @@ -145,6 +171,11 @@ const validation = computed(() => { if (hourly_metrics_retention_days < 0 || hourly_metrics_retention_days > 365) { errors.push(t('admin.ops.settings.validation.retentionDaysRange')) } + + const { default_threshold_5h, default_threshold_7d } = advancedSettings.value.openai_account_quota_auto_pause + if (default_threshold_5h < 0 || default_threshold_5h > 1 || default_threshold_7d < 0 || default_threshold_7d > 1) { + errors.push(t('admin.ops.settings.validation.openaiQuotaAutoPauseRange')) + } } // 验证指标阈值 @@ -473,6 +504,40 @@ async function saveAllSettings() {
+ +
+
{{ t('admin.ops.settings.openaiQuotaAutoPause') }}
+

{{ t('admin.ops.settings.openaiQuotaAutoPauseHint') }}

+ +
+
+ + +
+
+ + +
+
+

{{ t('admin.ops.settings.openaiQuotaAutoPauseThresholdHint') }}

+
+
{{ t('admin.ops.settings.errorFiltering') }}
From eba2046320ec643e7976296aabbc283f05aeaa52 Mon Sep 17 00:00:00 2001 From: xiaoyiluck666 <83876597+xiaoyiluck666@users.noreply.github.com> Date: Fri, 29 May 2026 13:34:10 +0800 Subject: [PATCH 05/36] fix: enrich OpenAI OAuth token refresh --- backend/cmd/server/wire_gen.go | 2 +- .../internal/service/openai_oauth_service.go | 54 +++++++++++------ .../openai_oauth_service_refresh_test.go | 7 +++ .../service/openai_privacy_service.go | 58 +++++++++++++++++++ .../service/openai_subscription_test.go | 42 ++++++++++++++ backend/internal/service/wire.go | 13 ++++- 6 files changed, 157 insertions(+), 19 deletions(-) create mode 100644 backend/internal/service/openai_subscription_test.go diff --git a/backend/cmd/server/wire_gen.go b/backend/cmd/server/wire_gen.go index 465f5e25a712..441bcd679206 100644 --- a/backend/cmd/server/wire_gen.go +++ b/backend/cmd/server/wire_gen.go @@ -137,7 +137,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { httpUpstream := repository.NewHTTPUpstream(configConfig) deferredService := service.ProvideDeferredService(accountRepository, timingWheelService) openAIOAuthClient := repository.NewOpenAIOAuthClient() - openAIOAuthService := service.NewOpenAIOAuthService(proxyRepository, openAIOAuthClient) + openAIOAuthService := service.ProvideOpenAIOAuthService(proxyRepository, openAIOAuthClient, privacyClientFactory) oAuthRefreshAPI := service.ProvideOAuthRefreshAPI(accountRepository, geminiTokenCache) openAITokenProvider := service.ProvideOpenAITokenProvider(accountRepository, geminiTokenCache, openAIOAuthService, oAuthRefreshAPI) channelRepository := repository.NewChannelRepository(db) diff --git a/backend/internal/service/openai_oauth_service.go b/backend/internal/service/openai_oauth_service.go index dc094d43ce2c..0ee357a97c05 100644 --- a/backend/internal/service/openai_oauth_service.go +++ b/backend/internal/service/openai_oauth_service.go @@ -278,11 +278,29 @@ func (s *OpenAIOAuthService) enrichTokenInfo(ctx context.Context, tokenInfo *Ope tokenInfo.Email = info.Email } } + if strings.TrimSpace(tokenInfo.SubscriptionExpiresAt) == "" { + if expiresAt := fetchChatGPTSubscriptionExpiresAt(ctx, s.privacyClientFactory, tokenInfo.AccessToken, proxyURL, resolveChatGPTSubscriptionAccountID(tokenInfo, orgID)); expiresAt != "" { + tokenInfo.SubscriptionExpiresAt = expiresAt + } + } // 尝试设置隐私（关闭训练数据共享），best-effort tokenInfo.PrivacyMode = disableOpenAITraining(ctx, s.privacyClientFactory, tokenInfo.AccessToken, proxyURL) } +func resolveChatGPTSubscriptionAccountID(tokenInfo *OpenAITokenInfo, orgID string) string { + for _, candidate := range []string{ + tokenInfo.ChatGPTAccountID, + tokenInfo.OrganizationID, + orgID, + } { + if trimmed := strings.TrimSpace(candidate); trimmed != "" { + return trimmed + } + } + return "" +} + // RefreshAccountToken refreshes token for an OpenAI OAuth account func (s *OpenAIOAuthService) RefreshAccountToken(ctx context.Context, account *Account) (*OpenAITokenInfo, error) { if account.Platform != PlatformOpenAI { @@ -292,38 +310,40 @@ func (s *OpenAIOAuthService) RefreshAccountToken(ctx context.Context, account *A return nil, infraerrors.New(http.StatusBadRequest, "OPENAI_OAUTH_INVALID_ACCOUNT_TYPE", "account is not an OAuth account") } + var proxyURL string + if account.ProxyID != nil { + proxy, err := s.proxyRepo.GetByID(ctx, *account.ProxyID) + if err == nil && proxy != nil { + proxyURL = proxy.URL() + } + } + refreshToken := account.GetCredential("refresh_token") if refreshToken == "" { accessToken := account.GetCredential("access_token") if accessToken != "" { tokenInfo := &OpenAITokenInfo{ - AccessToken: accessToken, - RefreshToken: "", - IDToken: account.GetCredential("id_token"), - ClientID: account.GetCredential("client_id"), - Email: account.GetCredential("email"), - ChatGPTAccountID: account.GetCredential("chatgpt_account_id"), - ChatGPTUserID: account.GetCredential("chatgpt_user_id"), - OrganizationID: account.GetCredential("organization_id"), - PlanType: account.GetCredential("plan_type"), + AccessToken: accessToken, + RefreshToken: "", + IDToken: account.GetCredential("id_token"), + ClientID: account.GetCredential("client_id"), + Email: account.GetCredential("email"), + ChatGPTAccountID: account.GetCredential("chatgpt_account_id"), + ChatGPTUserID: account.GetCredential("chatgpt_user_id"), + OrganizationID: account.GetCredential("organization_id"), + PlanType: account.GetCredential("plan_type"), + SubscriptionExpiresAt: account.GetCredential("subscription_expires_at"), } if expiresAt := account.GetCredentialAsTime("expires_at"); expiresAt != nil { tokenInfo.ExpiresAt = expiresAt.Unix() tokenInfo.ExpiresIn = int64(time.Until(*expiresAt).Seconds()) } + s.enrichTokenInfo(ctx, tokenInfo, proxyURL) return tokenInfo, nil } return nil, infraerrors.New(http.StatusBadRequest, "OPENAI_OAUTH_NO_REFRESH_TOKEN", "no refresh token available") } - var proxyURL string - if account.ProxyID != nil { - proxy, err := s.proxyRepo.GetByID(ctx, *account.ProxyID) - if err == nil && proxy != nil { - proxyURL = proxy.URL() - } - } - clientID := account.GetCredential("client_id") return s.RefreshTokenWithClientID(ctx, refreshToken, proxyURL, clientID) } diff --git a/backend/internal/service/openai_oauth_service_refresh_test.go b/backend/internal/service/openai_oauth_service_refresh_test.go index 84b68ea643a6..75588c8db4d0 100644 --- a/backend/internal/service/openai_oauth_service_refresh_test.go +++ b/backend/internal/service/openai_oauth_service_refresh_test.go @@ -8,6 +8,7 @@ import ( "time" "github.com/Wei-Shaw/sub2api/internal/pkg/openai" + "github.com/imroc/req/v3" "github.com/stretchr/testify/require" ) @@ -32,6 +33,11 @@ func (s *openaiOAuthClientRefreshStub) RefreshTokenWithClientID(ctx context.Cont func TestOpenAIOAuthService_RefreshAccountToken_NoRefreshTokenUsesExistingAccessToken(t *testing.T) { client := &openaiOAuthClientRefreshStub{} svc := NewOpenAIOAuthService(nil, client) + var privacyClientCalls int32 + svc.SetPrivacyClientFactory(func(proxyURL string) (*req.Client, error) { + atomic.AddInt32(&privacyClientCalls, 1) + return nil, errors.New("stop before request") + }) expiresAt := time.Now().Add(30 * time.Minute).UTC().Format(time.RFC3339) account := &Account{ @@ -51,6 +57,7 @@ func TestOpenAIOAuthService_RefreshAccountToken_NoRefreshTokenUsesExistingAccess require.Equal(t, "existing-access-token", info.AccessToken) require.Equal(t, "client-id-1", info.ClientID) require.Zero(t, atomic.LoadInt32(&client.refreshCalls), "existing access token should be reused without calling refresh") + require.Positive(t, atomic.LoadInt32(&privacyClientCalls), "existing access token should still run enrichment") } func TestOpenAITokenRefresher_NeedsRefresh_SkipsAccountWithoutRefreshToken(t *testing.T) { diff --git a/backend/internal/service/openai_privacy_service.go b/backend/internal/service/openai_privacy_service.go index da6dbefc93f6..99cbb7267e46 100644 --- a/backend/internal/service/openai_privacy_service.go +++ b/backend/internal/service/openai_privacy_service.go @@ -95,6 +95,8 @@ type ChatGPTAccountInfo struct { const chatGPTAccountsCheckURL = "https://chatgpt.com/backend-api/accounts/check/v4-2023-04-27" +var chatGPTSubscriptionsURL = "https://chatgpt.com/backend-api/subscriptions" + // fetchChatGPTAccountInfo calls ChatGPT backend-api to get account info (plan_type, etc.). // Used as fallback when id_token doesn't contain these fields (e.g., Mobile RT). // orgID is used to match the correct account when multiple accounts exist (e.g., personal + team). @@ -199,6 +201,62 @@ func fetchChatGPTAccountInfo(ctx context.Context, clientFactory PrivacyClientFac return info } +// fetchChatGPTSubscriptionExpiresAt reads the lightweight subscription endpoint used by +// ChatGPT/Codex clients. Some Plus accounts no longer expose entitlement.expires_at in +// accounts/check, but this endpoint still returns active_until. +func fetchChatGPTSubscriptionExpiresAt(ctx context.Context, clientFactory PrivacyClientFactory, accessToken, proxyURL, accountID string) string { + accountID = strings.TrimSpace(accountID) + if accessToken == "" || accountID == "" || clientFactory == nil { + return "" + } + + ctx, cancel := context.WithTimeout(ctx, 15*time.Second) + defer cancel() + + client, err := clientFactory(proxyURL) + if err != nil { + slog.Debug("chatgpt_subscription_client_error", "error", err.Error()) + return "" + } + + var result struct { + PlanType string `json:"plan_type"` + ActiveUntil string `json:"active_until"` + WillRenew bool `json:"will_renew"` + ID string `json:"id"` + } + resp, err := client.R(). + SetContext(ctx). + SetHeader("Authorization", "Bearer "+accessToken). + SetHeader("Origin", "https://chatgpt.com"). + SetHeader("Referer", "https://chatgpt.com/"). + SetHeader("Accept", "application/json"). + SetSuccessResult(&result). + SetQueryParam("account_id", accountID). + Get(chatGPTSubscriptionsURL) + if err != nil { + slog.Debug("chatgpt_subscription_request_error", "error", err.Error()) + return "" + } + if !resp.IsSuccessState() { + slog.Debug("chatgpt_subscription_failed", "status", resp.StatusCode, "body", truncate(resp.String(), 200)) + return "" + } + + activeUntil := strings.TrimSpace(result.ActiveUntil) + if activeUntil == "" { + slog.Debug("chatgpt_subscription_no_active_until", "plan_type", result.PlanType, "has_subscription_id", strings.TrimSpace(result.ID) != "", "will_renew", result.WillRenew) + return "" + } + if _, err := time.Parse(time.RFC3339, activeUntil); err != nil { + slog.Debug("chatgpt_subscription_bad_active_until", "active_until", activeUntil, "error", err.Error()) + return "" + } + + slog.Info("chatgpt_subscription_success", "plan_type", result.PlanType, "subscription_expires_at", activeUntil, "account_id", accountID) + return activeUntil +} + // fillAccountInfo 从单个 account 对象中提取 plan_type 和 subscription_expires_at func fillAccountInfo(info *ChatGPTAccountInfo, acct map[string]any) { info.PlanType = extractPlanType(acct) diff --git a/backend/internal/service/openai_subscription_test.go b/backend/internal/service/openai_subscription_test.go new file mode 100644 index 000000000000..89df54dbacb3 --- /dev/null +++ b/backend/internal/service/openai_subscription_test.go @@ -0,0 +1,42 @@ +package service + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/imroc/req/v3" + "github.com/stretchr/testify/require" +) + +func TestFetchChatGPTSubscriptionExpiresAt(t *testing.T) { + const wantExpiresAt = "2026-06-10T02:52:15Z" + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + require.Equal(t, "/backend-api/subscriptions", r.URL.Path) + require.Equal(t, "acc_123", r.URL.Query().Get("account_id")) + require.Equal(t, "Bearer access-token", r.Header.Get("Authorization")) + + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]any{ + "plan_type": "plus", + "active_until": wantExpiresAt, + "will_renew": true, + "id": "sub_123", + }) + })) + defer server.Close() + + oldURL := chatGPTSubscriptionsURL + chatGPTSubscriptionsURL = server.URL + "/backend-api/subscriptions" + t.Cleanup(func() { chatGPTSubscriptionsURL = oldURL }) + + got := fetchChatGPTSubscriptionExpiresAt(context.Background(), func(proxyURL string) (*req.Client, error) { + return req.C().SetTimeout(5 * time.Second), nil + }, "access-token", "", "acc_123") + + require.Equal(t, wantExpiresAt, got) +} diff --git a/backend/internal/service/wire.go b/backend/internal/service/wire.go index b22e10ae382b..e0c9f5910acd 100644 --- a/backend/internal/service/wire.go +++ b/backend/internal/service/wire.go @@ -45,6 +45,17 @@ func ProvideOAuthRefreshAPI(accountRepo AccountRepository, tokenCache GeminiToke return NewOAuthRefreshAPI(accountRepo, tokenCache) } +// ProvideOpenAIOAuthService creates OpenAIOAuthService with privacy/account enrichment support. +func ProvideOpenAIOAuthService( + proxyRepo ProxyRepository, + oauthClient OpenAIOAuthClient, + privacyClientFactory PrivacyClientFactory, +) *OpenAIOAuthService { + svc := NewOpenAIOAuthService(proxyRepo, oauthClient) + svc.SetPrivacyClientFactory(privacyClientFactory) + return svc +} + // ProvideTokenRefreshService creates and starts TokenRefreshService func ProvideTokenRefreshService( accountRepo AccountRepository, @@ -461,7 +472,7 @@ var ProviderSet = wire.NewSet( NewOpenAIGatewayService, wire.Bind(new(AccountRuntimeBlocker), new(*OpenAIGatewayService)), NewOAuthService, - NewOpenAIOAuthService, + ProvideOpenAIOAuthService, NewGeminiOAuthService, NewGeminiQuotaService, NewCompositeTokenCacheInvalidator, From c9caadb3782a4e658ce46789df092d0b15d58172 Mon Sep 17 00:00:00 2001 From: wucm667 Date: Fri, 29 May 2026 14:32:45 +0800 Subject: [PATCH 06/36] fix(account): address second-round review on quota auto-pause - TopK initial filter now drops quota-paused accounts: fold the quota check into isAccountRequestCompatible so session-hash, TopK pool, and per-candidate rechecks all skip paused accounts. Previously the candidate pool was built without the quota check, so paused accounts could fill TopK and leave the scheduler returning "no available accounts" even with healthy ones available. - Add per-account explicit disable flags auto_pause_5h_disabled / auto_pause_7d_disabled with toggles in EditAccountModal. Without these, leaving the account threshold blank silently falls back to the global default, so admins could not exempt a single account once a global default existed. Disable is per-window: an account can opt out of 5h auto-pause while still honoring 7d. Schedule snapshot whitelist includes the new fields, i18n EN/ZH updated, threshold-hint text revised to explain "blank = global default". - Move quota auto-pause settings off the request hot path: replace the per-repo TTL+singleflight sync DB read with a per-SettingService stale-while-revalidate in-memory snapshot. Get is non-blocking (atomic.Pointer load + async refresh on staleness); writes via UpdateOpsAdvancedSettings push directly into the cache through an injected sink; wire warms the cache at startup. Adds Warm (sync) for tests/init and SetOpenAIQuotaAutoPauseSettings (sink target). Co-Authored-By: Claude Opus 4.7 --- backend/cmd/server/wire_gen.go | 2 +- .../internal/repository/scheduler_cache.go | 2 + .../repository/scheduler_cache_unit_test.go | 4 + .../service/openai_account_scheduler.go | 7 + .../service/openai_account_scheduler_test.go | 136 ++++++++++++++ .../service/openai_gateway_service.go | 44 ++++- backend/internal/service/ops_service.go | 15 ++ backend/internal/service/ops_settings.go | 12 +- .../service/ops_settings_advanced_test.go | 50 +++++- backend/internal/service/setting_service.go | 170 +++++++++++------- backend/internal/service/wire.go | 42 ++++- .../components/account/EditAccountModal.vue | 60 +++++++ .../__tests__/EditAccountModal.spec.ts | 21 +++ frontend/src/i18n/locales/en.ts | 5 +- frontend/src/i18n/locales/zh.ts | 5 +- 15 files changed, 502 insertions(+), 73 deletions(-) diff --git a/backend/cmd/server/wire_gen.go b/backend/cmd/server/wire_gen.go index 465f5e25a712..6e8be8fce9a1 100644 --- a/backend/cmd/server/wire_gen.go +++ b/backend/cmd/server/wire_gen.go @@ -195,7 +195,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { gatewayService := service.NewGatewayService(accountRepository, groupRepository, usageLogRepository, usageBillingRepository, userRepository, userSubscriptionRepository, userGroupRateRepository, gatewayCache, configConfig, schedulerSnapshotService, concurrencyService, billingService, rateLimitService, billingCacheService, identityService, httpUpstream, deferredService, claudeTokenProvider, sessionLimitCache, rpmCache, digestSessionStore, settingService, tlsFingerprintProfileService, channelService, modelPricingResolver, balanceNotifyService, serviceUserPlatformQuotaRepository) geminiMessagesCompatService := service.NewGeminiMessagesCompatService(accountRepository, groupRepository, gatewayCache, schedulerSnapshotService, geminiTokenProvider, rateLimitService, httpUpstream, antigravityGatewayService, configConfig) opsSystemLogSink := service.ProvideOpsSystemLogSink(opsRepository) - opsService := service.NewOpsService(opsRepository, settingRepository, configConfig, accountRepository, userRepository, concurrencyService, gatewayService, openAIGatewayService, geminiMessagesCompatService, antigravityGatewayService, opsSystemLogSink) + opsService := service.ProvideOpsService(opsRepository, settingRepository, configConfig, accountRepository, userRepository, concurrencyService, gatewayService, openAIGatewayService, geminiMessagesCompatService, antigravityGatewayService, opsSystemLogSink, settingService) encryptionKey, err := payment.ProvideEncryptionKey(configConfig) if err != nil { return nil, err diff --git a/backend/internal/repository/scheduler_cache.go b/backend/internal/repository/scheduler_cache.go index ff3c4301388b..cf19dedaa0b1 100644 --- a/backend/internal/repository/scheduler_cache.go +++ b/backend/internal/repository/scheduler_cache.go @@ -557,6 +557,8 @@ func filterSchedulerExtra(extra map[string]any) map[string]any { "codex_usage_updated_at", "auto_pause_5h_threshold", "auto_pause_7d_threshold", + "auto_pause_5h_disabled", + "auto_pause_7d_disabled", } filtered := make(map[string]any) for _, key := range keys { diff --git a/backend/internal/repository/scheduler_cache_unit_test.go b/backend/internal/repository/scheduler_cache_unit_test.go index 9e4ec23e9aa9..a4667591722a 100644 --- a/backend/internal/repository/scheduler_cache_unit_test.go +++ b/backend/internal/repository/scheduler_cache_unit_test.go @@ -89,6 +89,8 @@ func TestBuildSchedulerMetadataAccount_KeepsQuotaAutoPauseFields(t *testing.T) { "codex_usage_updated_at": "2026-05-29T09:00:00Z", "auto_pause_5h_threshold": 0.95, "auto_pause_7d_threshold": 0.96, + "auto_pause_5h_disabled": true, + "auto_pause_7d_disabled": false, }, } @@ -103,4 +105,6 @@ func TestBuildSchedulerMetadataAccount_KeepsQuotaAutoPauseFields(t *testing.T) { require.Equal(t, "2026-05-29T09:00:00Z", got.Extra["codex_usage_updated_at"]) require.Equal(t, 0.95, got.Extra["auto_pause_5h_threshold"]) require.Equal(t, 0.96, got.Extra["auto_pause_7d_threshold"]) + require.Equal(t, true, got.Extra["auto_pause_5h_disabled"]) + require.Equal(t, false, got.Extra["auto_pause_7d_disabled"]) } diff --git a/backend/internal/service/openai_account_scheduler.go b/backend/internal/service/openai_account_scheduler.go index fd28fa86c367..47a8142a8683 100644 --- a/backend/internal/service/openai_account_scheduler.go +++ b/backend/internal/service/openai_account_scheduler.go @@ -974,6 +974,13 @@ func (s *defaultOpenAIAccountScheduler) isAccountRequestCompatible(ctx context.C if s != nil && s.service != nil && s.service.isOpenAIAccountRuntimeBlocked(account) { return false } + // Quota auto-pause must be evaluated during the initial filter too. Without it the + // TopK candidate pool can be filled with paused accounts and the later fresh/DB + // rechecks won't reach healthy accounts that fell outside TopK — manifesting as + // "no available accounts" even though healthy ones exist. + if paused, _ := shouldAutoPauseOpenAIAccountByQuota(ctx, account); paused { + return false + } if req.RequestedModel != "" && !account.IsModelSupported(req.RequestedModel) { return false } diff --git a/backend/internal/service/openai_account_scheduler_test.go b/backend/internal/service/openai_account_scheduler_test.go index 531769a73ef6..da5f0a6644bf 100644 --- a/backend/internal/service/openai_account_scheduler_test.go +++ b/backend/internal/service/openai_account_scheduler_test.go @@ -798,6 +798,63 @@ func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_UsesGlobalDefa require.Equal(t, int64(35402), account.ID) } +// Regression: a per-account explicit-disable flag exempts the account from auto-pause +// even when a global default threshold is set. Without this, "leave threshold blank" +// silently falls back to global default and admins have no way to whitelist a single +// account. +func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_PerAccountDisableOverridesGlobalDefault(t *testing.T) { + ctx := withOpenAIQuotaAutoPauseSettings(context.Background(), OpsOpenAIAccountQuotaAutoPauseSettings{DefaultThreshold5h: 0.95}) + // Account has high usage AND no per-account threshold (would normally fall back to + // the global default and get paused), but the explicit disable flag is set. + primary := Account{ + ID: 35701, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 0, + Extra: map[string]any{ + "codex_5h_used_percent": 99.0, + "auto_pause_5h_disabled": true, + }, + } + secondary := Account{ID: 35702, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5} + svc := &OpenAIGatewayService{accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}}, cfg: &config.Config{}} + + account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.1", nil) + require.NoError(t, err) + require.NotNil(t, account) + require.Equal(t, int64(35701), account.ID) +} + +// Disable is per-window: disabling only 5h must still allow 7d auto-pause to fire. +func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_PerWindowDisableScoped(t *testing.T) { + ctx := context.Background() + primary := Account{ + ID: 35801, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 0, + Extra: map[string]any{ + "codex_5h_used_percent": 99.0, + "codex_7d_used_percent": 99.0, + "auto_pause_5h_disabled": true, + "auto_pause_7d_threshold": 0.95, + }, + } + secondary := Account{ID: 35802, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5} + svc := &OpenAIGatewayService{accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}}, cfg: &config.Config{}} + + account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.1", nil) + require.NoError(t, err) + require.NotNil(t, account) + require.Equal(t, int64(35802), account.ID, "7d auto-pause must still fire even though 5h is disabled") +} + func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_StaleUsageWindowResetSkipsPause(t *testing.T) { ctx := context.Background() // Usage is over threshold but the window's reset time has already passed, so the @@ -1399,6 +1456,85 @@ func TestOpenAIGatewayService_SelectAccountWithScheduler_LoadBalanceTopKFallback } } +// Regression: TopK initial filter must drop quota-auto-paused accounts. Otherwise +// the candidate pool is filled with paused accounts, healthy accounts fall outside +// TopK, and the scheduler returns "no available accounts" even though healthy ones +// exist. +func TestOpenAIGatewayService_SelectAccountWithScheduler_LoadBalanceTopKExcludesQuotaPaused(t *testing.T) { + ctx := context.Background() + groupID := int64(110) + accounts := []Account{ + { + ID: 37001, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 0, + Extra: map[string]any{ + "codex_5h_used_percent": 96.0, + "auto_pause_5h_threshold": 0.95, + }, + }, + { + ID: 37002, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 5, + }, + } + + cfg := &config.Config{} + cfg.Gateway.OpenAIWS.LBTopK = 1 // TopK=1 makes the bug fatal: paused account would crowd out the healthy one entirely + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Priority = 0.4 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Load = 1.0 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Queue = 1.0 + + concurrencyCache := schedulerTestConcurrencyCache{ + loadMap: map[int64]*AccountLoadInfo{ + 37001: {AccountID: 37001, LoadRate: 5, WaitingCount: 0}, + 37002: {AccountID: 37002, LoadRate: 5, WaitingCount: 0}, + }, + acquireResults: map[int64]bool{ + 37002: true, + }, + } + + svc := &OpenAIGatewayService{ + accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts}, + cache: &schedulerTestGatewayCache{}, + cfg: cfg, + rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"), + concurrencyService: NewConcurrencyService(concurrencyCache), + } + + selection, decision, err := svc.SelectAccountWithScheduler( + ctx, + &groupID, + "", + "", + "gpt-5.1", + nil, + OpenAIUpstreamTransportAny, + false, + ) + require.NoError(t, err) + require.NotNil(t, selection) + require.NotNil(t, selection.Account) + require.Equal(t, int64(37002), selection.Account.ID) + require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer) + // Only the healthy account should ever enter the candidate pool; the paused one + // must be filtered out at the initial-filter stage. + require.Equal(t, 1, decision.CandidateCount) + if selection.ReleaseFunc != nil { + selection.ReleaseFunc() + } +} + func TestOpenAIGatewayService_OpenAIAccountSchedulerMetrics(t *testing.T) { ctx := context.Background() groupID := int64(12) diff --git a/backend/internal/service/openai_gateway_service.go b/backend/internal/service/openai_gateway_service.go index e2534cc2bbf9..b1ae6a9a7fc7 100644 --- a/backend/internal/service/openai_gateway_service.go +++ b/backend/internal/service/openai_gateway_service.go @@ -1360,14 +1360,21 @@ func shouldAutoPauseOpenAIAccountByQuota(ctx context.Context, account *Account) if account == nil || !account.IsOpenAI() { return false, openAIQuotaAutoPauseDecision{} } + // Per-account explicit-disable flags must take precedence over the global default. + // Without these, leaving the account threshold blank means "use global default", + // so an admin has no way to exempt a single account from auto-pause once a global + // default exists. The disable flag is per-window so an account can opt out of + // only 5h or only 7d auto-pause. + disabled5h := resolveAccountExtraBool(account.Extra, "auto_pause_5h_disabled") + disabled7d := resolveAccountExtraBool(account.Extra, "auto_pause_7d_disabled") threshold5h, threshold7d := resolveOpenAIQuotaAutoPauseThresholds(ctx, account) now := time.Now() - if threshold5h > 0 { + if !disabled5h && threshold5h > 0 { if utilization, ok := resolveOpenAIQuotaUtilization(account.Extra, "5h", now); ok && utilization >= threshold5h { return true, openAIQuotaAutoPauseDecision{window: "5h", threshold: threshold5h, utilization: utilization} } } - if threshold7d > 0 { + if !disabled7d && threshold7d > 0 { if utilization, ok := resolveOpenAIQuotaUtilization(account.Extra, "7d", now); ok && utilization >= threshold7d { return true, openAIQuotaAutoPauseDecision{window: "7d", threshold: threshold7d, utilization: utilization} } @@ -1375,6 +1382,39 @@ func shouldAutoPauseOpenAIAccountByQuota(ctx context.Context, account *Account) return false, openAIQuotaAutoPauseDecision{} } +// resolveAccountExtraBool reads a bool-like value from account extra, tolerating +// the few shapes JSON unmarshalling may produce (real bool, "true"/"false" +// strings, 0/1 numbers). +func resolveAccountExtraBool(extra map[string]any, key string) bool { + if len(extra) == 0 { + return false + } + value, ok := extra[key] + if !ok || value == nil { + return false + } + switch v := value.(type) { + case bool: + return v + case string: + parsed, err := strconv.ParseBool(strings.TrimSpace(v)) + return err == nil && parsed + case float64: + return v != 0 + case float32: + return v != 0 + case int: + return v != 0 + case int64: + return v != 0 + case json.Number: + if i, err := v.Int64(); err == nil { + return i != 0 + } + } + return false +} + func resolveOpenAIQuotaAutoPauseThresholds(ctx context.Context, account *Account) (float64, float64) { threshold5h, _ := resolveAccountExtraNumber(account.Extra, "auto_pause_5h_threshold") threshold7d, _ := resolveAccountExtraNumber(account.Extra, "auto_pause_7d_threshold") diff --git a/backend/internal/service/ops_service.go b/backend/internal/service/ops_service.go index 1cea72fae2ce..2d7c5bd42d7c 100644 --- a/backend/internal/service/ops_service.go +++ b/backend/internal/service/ops_service.go @@ -41,6 +41,11 @@ type OpsService struct { // cleanupReloader 由 wire 在 OpsCleanupService 构造完成后通过 SetCleanupReloader 注入。 // 解耦避免 OpsService -> OpsCleanupService 的硬依赖（cleanup 也读 settings，会循环）。 cleanupReloader CleanupReloader + + // quotaAutoPauseSink 由 wire 注入（通常是 SettingService.SetOpenAIQuotaAutoPauseSettings）。 + // UpdateOpsAdvancedSettings 写入新配置后调用，把最新的 quota auto-pause 全局默认阈值 + // 立即同步到调度热路径读取的内存缓存，避免下次请求才能感知新值。 + quotaAutoPauseSink func(OpsOpenAIAccountQuotaAutoPauseSettings) } // CleanupReloader 由 OpsCleanupService 实现。 @@ -57,6 +62,16 @@ func (s *OpsService) SetCleanupReloader(r CleanupReloader) { s.cleanupReloader = r } +// SetOpenAIQuotaAutoPauseSettingsSink 由 wire 注入，把最新的 quota auto-pause 全局默认 +// 阈值 push 到调度热路径读取的内存缓存。同 SetCleanupReloader 的解耦目的：避免 OpsService +// 持有 *SettingService 引入循环依赖。 +func (s *OpsService) SetOpenAIQuotaAutoPauseSettingsSink(sink func(OpsOpenAIAccountQuotaAutoPauseSettings)) { + if s == nil { + return + } + s.quotaAutoPauseSink = sink +} + func NewOpsService( opsRepo OpsRepository, settingRepo SettingRepository, diff --git a/backend/internal/service/ops_settings.go b/backend/internal/service/ops_settings.go index 23e92e5af21b..472f4e3207aa 100644 --- a/backend/internal/service/ops_settings.go +++ b/backend/internal/service/ops_settings.go @@ -490,12 +490,12 @@ func (s *OpsService) UpdateOpsAdvancedSettings(ctx context.Context, cfg *OpsAdva if err := s.settingRepo.Set(ctx, SettingKeyOpsAdvancedSettings, string(raw)); err != nil { return nil, err } - cacheKey := openAIQuotaAutoPauseSettingsCacheKey(s.settingRepo) - openAIQuotaAutoPauseSettingsSF.Forget(cacheKey) - storeOpenAIQuotaAutoPauseSettingsCache(s.settingRepo, &cachedOpenAIQuotaAutoPauseSettings{ - settings: cfg.OpenAIAccountQuotaAutoPause, - expiresAt: time.Now().Add(openAIQuotaAutoPauseSettingsCacheTTL).UnixNano(), - }) + // Push the new quota auto-pause settings straight into the in-memory cache that + // the OpenAI scheduling hot path reads, so the next request observes the new value + // without waiting for the background refresher's TTL. + if s.quotaAutoPauseSink != nil { + s.quotaAutoPauseSink(cfg.OpenAIAccountQuotaAutoPause) + } // notify cleanup service to reload schedule/enabled. if s.cleanupReloader != nil { diff --git a/backend/internal/service/ops_settings_advanced_test.go b/backend/internal/service/ops_settings_advanced_test.go index d8598fe0233f..62803f948b9b 100644 --- a/backend/internal/service/ops_settings_advanced_test.go +++ b/backend/internal/service/ops_settings_advanced_test.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "testing" + "time" "github.com/Wei-Shaw/sub2api/internal/config" ) @@ -103,11 +104,58 @@ func TestGetOpenAIQuotaAutoPauseSettings_ReadsDefaultsFromOpsAdvancedSettings(t repo.values[SettingKeyOpsAdvancedSettings] = `{"openai_account_quota_auto_pause":{"default_threshold_5h":0.95,"default_threshold_7d":0.9}}` svc := NewSettingService(repo, &config.Config{}) - settings := svc.GetOpenAIQuotaAutoPauseSettings(context.Background()) + // Warm the in-memory cache synchronously so the assertion below is deterministic. + // GetOpenAIQuotaAutoPauseSettings is non-blocking on the hot path (returns the + // cached value, refreshes asynchronously); for tests and startup, Warm is the + // synchronous entry point that guarantees a populated cache. + settings := svc.WarmOpenAIQuotaAutoPauseSettings(context.Background()) if settings.DefaultThreshold5h != 0.95 { t.Fatalf("DefaultThreshold5h = %v, want 0.95", settings.DefaultThreshold5h) } if settings.DefaultThreshold7d != 0.9 { t.Fatalf("DefaultThreshold7d = %v, want 0.9", settings.DefaultThreshold7d) } + + // Subsequent Get must hit the warm cache and return the same value without any DB + // access — that's the hot-path invariant. + cached := svc.GetOpenAIQuotaAutoPauseSettings(context.Background()) + if cached.DefaultThreshold5h != 0.95 || cached.DefaultThreshold7d != 0.9 { + t.Fatalf("cached read = %+v, want {0.95, 0.9}", cached) + } +} + +// Hot-path invariant: a Get with cold cache must return immediately (zero defaults) +// rather than blocking on the DB. The async refresher will populate the cache for +// subsequent calls. +func TestGetOpenAIQuotaAutoPauseSettings_ColdCacheNonBlocking(t *testing.T) { + repo := newRuntimeSettingRepoStub() + repo.values[SettingKeyOpsAdvancedSettings] = `{"openai_account_quota_auto_pause":{"default_threshold_5h":0.7}}` + svc := NewSettingService(repo, &config.Config{}) + + start := time.Now() + settings := svc.GetOpenAIQuotaAutoPauseSettings(context.Background()) + elapsed := time.Since(start) + if elapsed > 50*time.Millisecond { + t.Fatalf("cold-cache Get must be non-blocking, took %v", elapsed) + } + // Cold cache means we get zero defaults (the async refresh hasn't completed yet). + if settings.DefaultThreshold5h != 0 || settings.DefaultThreshold7d != 0 { + t.Fatalf("cold-cache Get = %+v, want zeroes", settings) + } +} + +// Explicit cache write (e.g. from UpdateOpsAdvancedSettings) must be visible on the +// very next read without any DB roundtrip. +func TestSetOpenAIQuotaAutoPauseSettings_VisibleImmediately(t *testing.T) { + svc := NewSettingService(newRuntimeSettingRepoStub(), &config.Config{}) + + svc.SetOpenAIQuotaAutoPauseSettings(OpsOpenAIAccountQuotaAutoPauseSettings{ + DefaultThreshold5h: 0.88, + DefaultThreshold7d: 0.77, + }) + + got := svc.GetOpenAIQuotaAutoPauseSettings(context.Background()) + if got.DefaultThreshold5h != 0.88 || got.DefaultThreshold7d != 0.77 { + t.Fatalf("after Set, Get = %+v, want {0.88, 0.77}", got) + } } diff --git a/backend/internal/service/setting_service.go b/backend/internal/service/setting_service.go index 40e98b88a5ba..98acdb80f56d 100644 --- a/backend/internal/service/setting_service.go +++ b/backend/internal/service/setting_service.go @@ -14,7 +14,6 @@ import ( "sort" "strconv" "strings" - "sync" "sync/atomic" "time" @@ -162,28 +161,7 @@ const openAIQuotaAutoPauseSettingsCacheTTL = 60 * time.Second const openAIQuotaAutoPauseSettingsErrorTTL = 5 * time.Second const openAIQuotaAutoPauseSettingsDBTimeout = 5 * time.Second -var openAIQuotaAutoPauseSettingsCache sync.Map // map[string]*cachedOpenAIQuotaAutoPauseSettings -var openAIQuotaAutoPauseSettingsSF singleflight.Group - -func openAIQuotaAutoPauseSettingsCacheKey(repo SettingRepository) string { - if repo == nil { - return "nil" - } - return fmt.Sprintf("%T:%p", repo, repo) -} - -func loadOpenAIQuotaAutoPauseSettingsCache(repo SettingRepository) (*cachedOpenAIQuotaAutoPauseSettings, bool) { - value, ok := openAIQuotaAutoPauseSettingsCache.Load(openAIQuotaAutoPauseSettingsCacheKey(repo)) - if !ok || value == nil { - return nil, false - } - cached, ok := value.(*cachedOpenAIQuotaAutoPauseSettings) - return cached, ok && cached != nil -} - -func storeOpenAIQuotaAutoPauseSettingsCache(repo SettingRepository, cached *cachedOpenAIQuotaAutoPauseSettings) { - openAIQuotaAutoPauseSettingsCache.Store(openAIQuotaAutoPauseSettingsCacheKey(repo), cached) -} +const openAIQuotaAutoPauseSettingsRefreshKey = "openai_quota_auto_pause_settings" // DefaultSubscriptionGroupReader validates group references used by default subscriptions. type DefaultSubscriptionGroupReader interface { @@ -209,6 +187,15 @@ type SettingService struct { openAICodexUASF singleflight.Group openAIAllowCodexPluginCache atomic.Value // *cachedOpenAIAllowCodexPlugin openAIAllowCodexPluginSF singleflight.Group + + // openAIQuotaAutoPauseSettingsCache holds the most recently observed quota auto-pause + // settings. GetOpenAIQuotaAutoPauseSettings reads this atomic.Value on the request hot + // path without ever blocking on the DB; when the cached entry expires, a background + // goroutine refreshes it via openAIQuotaAutoPauseSettingsSF (stale-while-revalidate). + // This per-service field also gives tests natural isolation — each SettingService + // instance owns its own cache, no shared package-level state. + openAIQuotaAutoPauseSettingsCache atomic.Value // *cachedOpenAIQuotaAutoPauseSettings + openAIQuotaAutoPauseSettingsSF singleflight.Group } // DefaultPlatformQuotaSetting 单 platform 三档限额（nil = 沿用上层；0 = 显式禁用；>0 = 上限） @@ -2060,9 +2047,17 @@ func (s *SettingService) refreshCachedSettings(settings *SystemSettings) { enabled: settings.OpenAIAdvancedSchedulerEnabled, expiresAt: time.Now().Add(openAIAdvancedSchedulerSettingCacheTTL).UnixNano(), }) - cacheKey := openAIQuotaAutoPauseSettingsCacheKey(s.settingRepo) - openAIQuotaAutoPauseSettingsSF.Forget(cacheKey) - openAIQuotaAutoPauseSettingsCache.Delete(cacheKey) + // Invalidate the quota auto-pause cache and let the next read trigger a fresh load. + // We can't know from here whether ops_advanced_settings was also touched, so be + // defensive: store an expired entry — GetOpenAIQuotaAutoPauseSettings will serve + // stale and kick off an async refresh, never blocking the request that follows. + s.openAIQuotaAutoPauseSettingsSF.Forget(openAIQuotaAutoPauseSettingsRefreshKey) + if cached, _ := s.openAIQuotaAutoPauseSettingsCache.Load().(*cachedOpenAIQuotaAutoPauseSettings); cached != nil { + s.openAIQuotaAutoPauseSettingsCache.Store(&cachedOpenAIQuotaAutoPauseSettings{ + settings: cached.settings, + expiresAt: 0, + }) + } if s.cfg != nil { s.cfg.SetTrustForwardedIPForAPIKeyACL(settings.APIKeyACLTrustForwardedIP) } @@ -4484,49 +4479,104 @@ func (s *SettingService) GetClaudeCodeVersionBounds(ctx context.Context) (min, m return b.min, b.max } +// GetOpenAIQuotaAutoPauseSettings returns the current global default quota auto-pause +// settings. It is invoked on the OpenAI scheduling hot path (once per request) and is +// therefore designed to never block on the DB: +// +// - Fresh cached value → returned immediately. +// - Stale or empty cache → the last known value is returned, and a background +// goroutine refreshes the cache via singleflight (stale-while-revalidate). +// - First call with no cache yet → zero defaults are returned and the same async +// refresh is kicked off; the next call gets the freshly populated value. +// +// Callers that need the freshly persisted value synchronously (tests, post-update +// confirmation, optional startup warm-up) should call WarmOpenAIQuotaAutoPauseSettings. func (s *SettingService) GetOpenAIQuotaAutoPauseSettings(ctx context.Context) OpsOpenAIAccountQuotaAutoPauseSettings { - if cached, ok := loadOpenAIQuotaAutoPauseSettingsCache(s.settingRepo); ok { - if time.Now().UnixNano() < cached.expiresAt { - return cached.settings - } + if s == nil { + return OpsOpenAIAccountQuotaAutoPauseSettings{} + } + cached, _ := s.openAIQuotaAutoPauseSettingsCache.Load().(*cachedOpenAIQuotaAutoPauseSettings) + now := time.Now().UnixNano() + if cached != nil && now < cached.expiresAt { + return cached.settings + } + // Stale or unset: trigger background refresh without blocking this request. + // singleflight.DoChan dedupes concurrent refreshes; we deliberately ignore the + // returned channel — the result is observable via the atomic cache. + s.openAIQuotaAutoPauseSettingsSF.DoChan(openAIQuotaAutoPauseSettingsRefreshKey, func() (any, error) { + s.refreshOpenAIQuotaAutoPauseSettings(context.Background()) + return nil, nil + }) + if cached != nil { + return cached.settings // serve stale value while revalidating } + return OpsOpenAIAccountQuotaAutoPauseSettings{} +} - cacheKey := openAIQuotaAutoPauseSettingsCacheKey(s.settingRepo) - result, _, _ := openAIQuotaAutoPauseSettingsSF.Do(cacheKey, func() (any, error) { - if cached, ok := loadOpenAIQuotaAutoPauseSettingsCache(s.settingRepo); ok { - if time.Now().UnixNano() < cached.expiresAt { - return cached.settings, nil - } - } +// WarmOpenAIQuotaAutoPauseSettings synchronously loads the quota auto-pause settings +// into the in-memory cache. Useful for application startup (so the first request hits +// a warm cache) and for tests that need deterministic reads immediately after +// constructing the service. +func (s *SettingService) WarmOpenAIQuotaAutoPauseSettings(ctx context.Context) OpsOpenAIAccountQuotaAutoPauseSettings { + if s == nil { + return OpsOpenAIAccountQuotaAutoPauseSettings{} + } + s.refreshOpenAIQuotaAutoPauseSettings(ctx) + cached, _ := s.openAIQuotaAutoPauseSettingsCache.Load().(*cachedOpenAIQuotaAutoPauseSettings) + if cached == nil { + return OpsOpenAIAccountQuotaAutoPauseSettings{} + } + return cached.settings +} - settings := OpsOpenAIAccountQuotaAutoPauseSettings{} - ttl := openAIQuotaAutoPauseSettingsCacheTTL - if s != nil && s.settingRepo != nil { - dbCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), openAIQuotaAutoPauseSettingsDBTimeout) - defer cancel() - raw, err := s.settingRepo.GetValue(dbCtx, SettingKeyOpsAdvancedSettings) - if err == nil { - cfg := defaultOpsAdvancedSettings() - if strings.TrimSpace(raw) != "" { - if jsonErr := json.Unmarshal([]byte(raw), cfg); jsonErr == nil { - normalizeOpsAdvancedSettings(cfg) - } - } - settings = cfg.OpenAIAccountQuotaAutoPause - } else { - ttl = openAIQuotaAutoPauseSettingsErrorTTL +// refreshOpenAIQuotaAutoPauseSettings reads the latest settings from the DB and stores +// them into the in-memory cache. On error it stores the prior value (or zero defaults +// if nothing is cached yet) with the shorter error TTL so the next refresh comes +// sooner. Always uses its own timeout-bounded context to keep refresh latency +// predictable regardless of the caller. +func (s *SettingService) refreshOpenAIQuotaAutoPauseSettings(ctx context.Context) { + if s == nil || s.settingRepo == nil { + return + } + dbCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), openAIQuotaAutoPauseSettingsDBTimeout) + defer cancel() + + settings := OpsOpenAIAccountQuotaAutoPauseSettings{} + ttl := openAIQuotaAutoPauseSettingsCacheTTL + raw, err := s.settingRepo.GetValue(dbCtx, SettingKeyOpsAdvancedSettings) + if err == nil { + cfg := defaultOpsAdvancedSettings() + if strings.TrimSpace(raw) != "" { + if jsonErr := json.Unmarshal([]byte(raw), cfg); jsonErr == nil { + normalizeOpsAdvancedSettings(cfg) } } + settings = cfg.OpenAIAccountQuotaAutoPause + } else if !errors.Is(err, ErrSettingNotFound) { + // Real error: keep serving prior value but refresh sooner. + if prior, _ := s.openAIQuotaAutoPauseSettingsCache.Load().(*cachedOpenAIQuotaAutoPauseSettings); prior != nil { + settings = prior.settings + } + ttl = openAIQuotaAutoPauseSettingsErrorTTL + } - storeOpenAIQuotaAutoPauseSettingsCache(s.settingRepo, &cachedOpenAIQuotaAutoPauseSettings{ - settings: settings, - expiresAt: time.Now().Add(ttl).UnixNano(), - }) - return settings, nil + s.openAIQuotaAutoPauseSettingsCache.Store(&cachedOpenAIQuotaAutoPauseSettings{ + settings: settings, + expiresAt: time.Now().Add(ttl).UnixNano(), }) +} - settings, _ := result.(OpsOpenAIAccountQuotaAutoPauseSettings) - return settings +// SetOpenAIQuotaAutoPauseSettings writes the given settings directly into the in-memory +// cache. Called from settings-write code paths so that the next read reflects the new +// value immediately, without waiting for the background refresh. +func (s *SettingService) SetOpenAIQuotaAutoPauseSettings(settings OpsOpenAIAccountQuotaAutoPauseSettings) { + if s == nil { + return + } + s.openAIQuotaAutoPauseSettingsCache.Store(&cachedOpenAIQuotaAutoPauseSettings{ + settings: settings, + expiresAt: time.Now().Add(openAIQuotaAutoPauseSettingsCacheTTL).UnixNano(), + }) } // GetRectifierSettings 获取请求整流器配置 diff --git a/backend/internal/service/wire.go b/backend/internal/service/wire.go index b22e10ae382b..d3e4ce51bc7e 100644 --- a/backend/internal/service/wire.go +++ b/backend/internal/service/wire.go @@ -396,6 +396,46 @@ func ProvideBackupService( return svc } +// ProvideOpsService constructs OpsService and wires the SettingService-backed quota +// auto-pause cache sink. Mirrors the SetCleanupReloader pattern: OpsService doesn't +// hold a *SettingService reference, but wire injects a tiny callback so writes to +// ops_advanced_settings immediately propagate into the scheduler hot-path cache. +func ProvideOpsService( + opsRepo OpsRepository, + settingRepo SettingRepository, + cfg *config.Config, + accountRepo AccountRepository, + userRepo UserRepository, + concurrencyService *ConcurrencyService, + gatewayService *GatewayService, + openAIGatewayService *OpenAIGatewayService, + geminiCompatService *GeminiMessagesCompatService, + antigravityGatewayService *AntigravityGatewayService, + systemLogSink *OpsSystemLogSink, + settingService *SettingService, +) *OpsService { + svc := NewOpsService( + opsRepo, + settingRepo, + cfg, + accountRepo, + userRepo, + concurrencyService, + gatewayService, + openAIGatewayService, + geminiCompatService, + antigravityGatewayService, + systemLogSink, + ) + if settingService != nil { + svc.SetOpenAIQuotaAutoPauseSettingsSink(settingService.SetOpenAIQuotaAutoPauseSettings) + // Optional warm-up so the first scheduled request after process start observes + // a populated cache rather than zero defaults. Best-effort, sync-bounded. + settingService.WarmOpenAIQuotaAutoPauseSettings(context.Background()) + } + return svc +} + // ProvideSettingService wires SettingService with group reader and proxy repo. func ProvideSettingService(settingRepo SettingRepository, groupRepo GroupRepository, proxyRepo ProxyRepository, cfg *config.Config) *SettingService { svc := NewSettingService(settingRepo, cfg) @@ -481,7 +521,7 @@ var ProviderSet = wire.NewSet( NewDataManagementService, ProvideBackupService, ProvideOpsSystemLogSink, - NewOpsService, + ProvideOpsService, ProvideOpsMetricsCollector, ProvideOpsAggregationService, ProvideOpsAlertEvaluatorService, diff --git a/frontend/src/components/account/EditAccountModal.vue b/frontend/src/components/account/EditAccountModal.vue index 470c0bfbdf88..8dc85d0eb841 100644 --- a/frontend/src/components/account/EditAccountModal.vue +++ b/frontend/src/components/account/EditAccountModal.vue @@ -1791,6 +1791,28 @@ v-if="account?.platform === 'openai'" class="border-t border-gray-200 pt-4 dark:border-dark-600 space-y-4" > +
+
+ + +
+

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

+

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

+
+
+ + +
+

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

+

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

@@ -2481,6 +2527,8 @@ const interceptWarmupRequests = ref(false) const autoPauseOnExpired = ref(false) const autoPause5hThreshold = ref(null) const autoPause7dThreshold = ref(null) +const autoPause5hDisabled = ref(false) +const autoPause7dDisabled = ref(false) const mixedScheduling = ref(false) // For antigravity accounts: enable mixed scheduling const allowOverages = ref(false) // For antigravity accounts: enable AI Credits overages const antigravityModelRestrictionMode = ref<'whitelist' | 'mapping'>('whitelist') @@ -2901,6 +2949,8 @@ const syncFormFromAccount = (newAccount: Account | null) => { allowOverages.value = extra?.allow_overages === true autoPause5hThreshold.value = typeof extra?.auto_pause_5h_threshold === 'number' ? extra.auto_pause_5h_threshold * 100 : null autoPause7dThreshold.value = typeof extra?.auto_pause_7d_threshold === 'number' ? extra.auto_pause_7d_threshold * 100 : null + autoPause5hDisabled.value = extra?.auto_pause_5h_disabled === true + autoPause7dDisabled.value = extra?.auto_pause_7d_disabled === true // Load OpenAI passthrough toggle (OpenAI OAuth/API Key) openaiPassthroughEnabled.value = false @@ -4064,6 +4114,16 @@ const handleSubmit = async () => { } else { delete newExtra.auto_pause_7d_threshold } + if (autoPause5hDisabled.value) { + newExtra.auto_pause_5h_disabled = true + } else { + delete newExtra.auto_pause_5h_disabled + } + if (autoPause7dDisabled.value) { + newExtra.auto_pause_7d_disabled = true + } else { + delete newExtra.auto_pause_7d_disabled + } delete newExtra.codex_image_generation_bridge_enabled if (codexImageGenerationBridgeMode.value === 'inherit') { diff --git a/frontend/src/components/account/__tests__/EditAccountModal.spec.ts b/frontend/src/components/account/__tests__/EditAccountModal.spec.ts index 6db6383129ad..f4865de9e6c5 100644 --- a/frontend/src/components/account/__tests__/EditAccountModal.spec.ts +++ b/frontend/src/components/account/__tests__/EditAccountModal.spec.ts @@ -352,6 +352,27 @@ describe('EditAccountModal', () => { expect(updateAccountMock.mock.calls[0]?.[1]?.extra?.auto_pause_7d_threshold).toBe(0.96) }) + it('submits OpenAI quota auto-pause disable flag in extra', async () => { + // Toggling the per-account disable flag must persist as auto_pause_5h_disabled + // so an admin can exempt one account from auto-pause even when a global default + // threshold is configured (otherwise leaving the threshold blank would silently + // fall back to the global default). + const account = buildAccount() + updateAccountMock.mockReset() + checkMixedChannelRiskMock.mockReset() + checkMixedChannelRiskMock.mockResolvedValue({ has_risk: false }) + updateAccountMock.mockResolvedValue(account) + + const wrapper = mountModal(account) + + await wrapper.get('[data-testid="auto-pause-5h-disabled"]').trigger('click') + await wrapper.get('form#edit-account-form').trigger('submit.prevent') + + expect(updateAccountMock).toHaveBeenCalledTimes(1) + expect(updateAccountMock.mock.calls[0]?.[1]?.extra?.auto_pause_5h_disabled).toBe(true) + expect(updateAccountMock.mock.calls[0]?.[1]?.extra?.auto_pause_7d_disabled).toBeUndefined() + }) + it('keeps at least one OpenAI APIKey endpoint capability selected', async () => { const account = buildAccount() updateAccountMock.mockReset() diff --git a/frontend/src/i18n/locales/en.ts b/frontend/src/i18n/locales/en.ts index 8ab90961e7c4..6735029c7c3a 100644 --- a/frontend/src/i18n/locales/en.ts +++ b/frontend/src/i18n/locales/en.ts @@ -3477,7 +3477,10 @@ export default { autoPauseOnExpiredDesc: 'When enabled, the account will auto pause scheduling after it expires', autoPause5hThreshold: '5h Usage Threshold (%)', autoPause7dThreshold: '7d Usage Threshold (%)', - autoPauseThresholdHint: 'Leave empty or set 0 to disable. Reaching the threshold only skips the account during scheduling and does not modify schedulable.', + autoPauseThresholdHint: 'Leave empty or set 0 to use the global default threshold (configured in Ops settings); set a value to override the global default. Reaching the threshold only skips the account during scheduling and does not modify schedulable.', + autoPause5hDisabled: 'Disable 5h auto-pause', + autoPause7dDisabled: 'Disable 7d auto-pause', + autoPauseDisabledHint: 'When enabled, this account is never auto-paused (even if a global default threshold is configured).', // Quota control (Anthropic OAuth/SetupToken only) quotaControl: { title: 'Quota Control', diff --git a/frontend/src/i18n/locales/zh.ts b/frontend/src/i18n/locales/zh.ts index 4f1d1f13975e..abb8dff730e7 100644 --- a/frontend/src/i18n/locales/zh.ts +++ b/frontend/src/i18n/locales/zh.ts @@ -3615,7 +3615,10 @@ export default { autoPauseOnExpiredDesc: '启用后，账号过期将自动暂停调度', autoPause5hThreshold: '5h 用量阈值(%)', autoPause7dThreshold: '7d 用量阈值(%)', - autoPauseThresholdHint: '填 0 或留空表示不启用；达到阈值后仅在调度时跳过账号，不修改 schedulable。', + autoPauseThresholdHint: '留空或填 0 表示使用全局默认阈值（在运维设置中配置）；填具体值则覆盖全局默认。达到阈值后仅在调度时跳过账号，不修改 schedulable。', + autoPause5hDisabled: '禁用 5h 自动暂停', + autoPause7dDisabled: '禁用 7d 自动暂停', + autoPauseDisabledHint: '开启后该账号永不进入自动暂停（即使全局默认阈值已配置）。', // Quota control (Anthropic OAuth/SetupToken only) quotaControl: { title: '配额控制', From 0a521f09fbc481ccdd2a0d6bccfaa8f80df220fc Mon Sep 17 00:00:00 2001 From: Pluviobyte Date: Fri, 29 May 2026 06:46:49 +0000 Subject: [PATCH 07/36] fix(gemini): close tool_use block before text in messages streaming When the Gemini->Anthropic streaming bridge for the /v1/messages endpoint receives a functionCall part followed by a text part, the text branch in handleStreamingResponse opened a new text content block without closing the already-open tool_use block. The tool block's content_block_stop was only emitted at end-of-stream, after the text block's content_block_start, so the Anthropic SSE stream contained overlapping/unterminated content blocks. Clients that assemble messages by block index (e.g. Claude Code) can drop the tool input or mis-parse the response. The functionCall branch already closes an open text block before opening a tool block, and the chat-completions sibling closes the tool block in its text branch via closeOpenTool(). This applies the same symmetric handling to the messages variant: close any open tool_use block (resetting openToolIndex/openToolName/ seenToolJSON) before starting text. Adds a regression test that replays a tool->text Gemini stream and asserts the Anthropic content-block lifecycle never overlaps. --- .../service/gemini_messages_compat_service.go | 16 +++ .../gemini_messages_compat_service_test.go | 105 ++++++++++++++++++ 2 files changed, 121 insertions(+) diff --git a/backend/internal/service/gemini_messages_compat_service.go b/backend/internal/service/gemini_messages_compat_service.go index 516556ca7de8..64f19b2ec758 100644 --- a/backend/internal/service/gemini_messages_compat_service.go +++ b/backend/internal/service/gemini_messages_compat_service.go @@ -2031,6 +2031,22 @@ func (s *GeminiMessagesCompatService) handleStreamingResponse(c *gin.Context, re parts := extractGeminiParts(geminiResp) for _, part := range parts { if text, ok := part["text"].(string); ok && text != "" { + // Close an open tool_use block before starting text, mirroring + // the functionCall branch (which closes open text blocks) and + // the chat-completions sibling's closeOpenTool(). Otherwise a + // tool→text sequence keeps the tool_use block open while the + // text block starts, emitting overlapping Anthropic content + // blocks that violate the SSE contract. + if openToolIndex >= 0 { + writeSSE(c.Writer, "content_block_stop", map[string]any{ + "type": "content_block_stop", + "index": openToolIndex, + }) + openToolIndex = -1 + openToolName = "" + seenToolJSON = "" + } + delta, newSeen := computeGeminiTextDelta(seenText, text) seenText = newSeen if delta == "" { diff --git a/backend/internal/service/gemini_messages_compat_service_test.go b/backend/internal/service/gemini_messages_compat_service_test.go index d0560344dc47..79db633aae0f 100644 --- a/backend/internal/service/gemini_messages_compat_service_test.go +++ b/backend/internal/service/gemini_messages_compat_service_test.go @@ -832,3 +832,108 @@ func TestParseGeminiRateLimitResetTime(t *testing.T) { }) } } + +// TestGeminiMessagesHandleStreamingResponse_ClosesToolBlockBeforeText guards the +// tool→text ordering in the Gemini→Anthropic (messages) streaming bridge. When +// Gemini emits a functionCall part followed by a text part, the tool_use content +// block must be closed before the text block opens; otherwise the Anthropic SSE +// stream contains overlapping content blocks. The chat-completions sibling +// already enforces this via closeOpenTool(). +func TestGeminiMessagesHandleStreamingResponse_ClosesToolBlockBeforeText(t *testing.T) { + gin.SetMode(gin.TestMode) + + upstreamBody := `data: {"candidates":[{"content":{"parts":[{"functionCall":{"name":"get_weather","args":{"city":"SF"}}}]}}]}` + "\n\n" + + `data: {"candidates":[{"content":{"parts":[{"text":"All done."}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":5,"candidatesTokenCount":3}}` + "\n\n" + + "data: [DONE]\n\n" + + resp := &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(strings.NewReader(upstreamBody)), + } + + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + + svc := &GeminiMessagesCompatService{} + result, err := svc.handleStreamingResponse(c, resp, time.Now(), "claude-3-5-sonnet") + require.NoError(t, err) + require.NotNil(t, result) + + events := parseAnthropicContentBlockEvents(t, rec.Body.String()) + + // Anthropic allows at most one content block open at a time: every + // content_block_start must be matched by a content_block_stop before the + // next start. Replay the lifecycle and assert there is no overlap. + open := -1 + blockTypes := map[int]string{} + textStarted := false + toolClosed := false + toolClosedBeforeText := false + for _, ev := range events { + switch ev.event { + case "content_block_start": + require.Equalf(t, -1, open, + "content block %d opened while block %d was still open (overlapping blocks)", ev.index, open) + open = ev.index + blockTypes[ev.index] = ev.blockType + if ev.blockType == "text" { + textStarted = true + if toolClosed { + toolClosedBeforeText = true + } + } + case "content_block_stop": + require.Equalf(t, open, ev.index, + "content_block_stop index %d does not match the open block %d", ev.index, open) + if blockTypes[ev.index] == "tool_use" { + toolClosed = true + } + open = -1 + } + } + + require.True(t, textStarted, "expected a text content block to be emitted after the tool call") + require.True(t, toolClosedBeforeText, "tool_use block must be closed before the text block starts") + require.Equal(t, -1, open, "stream ended with a content block still open") +} + +type anthropicContentBlockEvent struct { + event string + index int + blockType string +} + +// parseAnthropicContentBlockEvents extracts content_block_start/stop events (with +// their index and, for starts, the content block type) from an Anthropic SSE body. +func parseAnthropicContentBlockEvents(t *testing.T, raw string) []anthropicContentBlockEvent { + t.Helper() + var events []anthropicContentBlockEvent + for _, chunk := range strings.Split(raw, "\n\n") { + var eventName, dataLine string + for _, line := range strings.Split(chunk, "\n") { + switch { + case strings.HasPrefix(line, "event:"): + eventName = strings.TrimSpace(strings.TrimPrefix(line, "event:")) + case strings.HasPrefix(line, "data:"): + dataLine = strings.TrimSpace(strings.TrimPrefix(line, "data:")) + } + } + if eventName != "content_block_start" && eventName != "content_block_stop" { + continue + } + var payload struct { + Index int `json:"index"` + ContentBlock struct { + Type string `json:"type"` + } `json:"content_block"` + } + require.NoError(t, json.Unmarshal([]byte(dataLine), &payload)) + events = append(events, anthropicContentBlockEvent{ + event: eventName, + index: payload.Index, + blockType: payload.ContentBlock.Type, + }) + } + return events +} From 68901cbfff783af794d96028be8dad3e532c0fe7 Mon Sep 17 00:00:00 2001 From: shaw Date: Fri, 29 May 2026 16:29:29 +0800 Subject: [PATCH 08/36] chore(pricing): update model pricing metadata --- .../model_prices_and_context_window.json | 4727 ++++++++--------- 1 file changed, 2288 insertions(+), 2439 deletions(-) diff --git a/backend/resources/model-pricing/model_prices_and_context_window.json b/backend/resources/model-pricing/model_prices_and_context_window.json index 3cae8c8b58dc..e88ed2da2215 100644 --- a/backend/resources/model-pricing/model_prices_and_context_window.json +++ b/backend/resources/model-pricing/model_prices_and_context_window.json @@ -1,135 +1,4 @@ { - "claude-3-5-haiku-20241022": { - "cache_creation_input_token_cost": 1e-06, - "cache_creation_input_token_cost_above_1hr": 6e-06, - "cache_read_input_token_cost": 8e-08, - "deprecation_date": "2025-10-01", - "input_cost_per_token": 8e-07, - "litellm_provider": "anthropic", - "max_input_tokens": 200000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 4e-06, - "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 - }, - "supports_assistant_prefill": true, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true, - "tool_use_system_prompt_tokens": 264 - }, - "claude-3-5-haiku-latest": { - "cache_creation_input_token_cost": 1.25e-06, - "cache_creation_input_token_cost_above_1hr": 6e-06, - "cache_read_input_token_cost": 1e-07, - "deprecation_date": "2025-10-01", - "input_cost_per_token": 1e-06, - "litellm_provider": "anthropic", - "max_input_tokens": 200000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 5e-06, - "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 - }, - "supports_assistant_prefill": true, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true, - "tool_use_system_prompt_tokens": 264 - }, - "claude-3-5-sonnet-20240620": { - "cache_creation_input_token_cost": 3.75e-06, - "cache_creation_input_token_cost_above_1hr": 6e-06, - "cache_read_input_token_cost": 3e-07, - "deprecation_date": "2025-06-01", - "input_cost_per_token": 3e-06, - "litellm_provider": "anthropic", - "max_input_tokens": 200000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 1.5e-05, - "supports_assistant_prefill": true, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 - }, - "claude-3-5-sonnet-20241022": { - "cache_creation_input_token_cost": 3.75e-06, - "cache_creation_input_token_cost_above_1hr": 6e-06, - "cache_read_input_token_cost": 3e-07, - "deprecation_date": "2025-10-01", - "input_cost_per_token": 3e-06, - "litellm_provider": "anthropic", - "max_input_tokens": 200000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 1.5e-05, - "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 - }, - "supports_assistant_prefill": true, - "supports_computer_use": true, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true, - "tool_use_system_prompt_tokens": 159 - }, - "claude-3-5-sonnet-latest": { - "cache_creation_input_token_cost": 3.75e-06, - "cache_creation_input_token_cost_above_1hr": 6e-06, - "cache_read_input_token_cost": 3e-07, - "deprecation_date": "2025-06-01", - "input_cost_per_token": 3e-06, - "litellm_provider": "anthropic", - "max_input_tokens": 200000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 1.5e-05, - "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 - }, - "supports_assistant_prefill": true, - "supports_computer_use": true, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true, - "tool_use_system_prompt_tokens": 159 - }, "claude-3-7-sonnet-20250219": { "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, @@ -159,34 +28,6 @@ "supports_web_search": true, "tool_use_system_prompt_tokens": 159 }, - "claude-3-7-sonnet-latest": { - "cache_creation_input_token_cost": 3.75e-06, - "cache_creation_input_token_cost_above_1hr": 6e-06, - "cache_read_input_token_cost": 3e-07, - "deprecation_date": "2025-06-01", - "input_cost_per_token": 3e-06, - "litellm_provider": "anthropic", - "max_input_tokens": 200000, - "max_output_tokens": 64000, - "max_tokens": 64000, - "mode": "chat", - "output_cost_per_token": 1.5e-05, - "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 - }, - "supports_assistant_prefill": true, - "supports_computer_use": true, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 - }, "claude-3-haiku-20240307": { "cache_creation_input_token_cost": 3e-07, "cache_creation_input_token_cost_above_1hr": 6e-06, @@ -226,28 +67,9 @@ "supports_vision": true, "tool_use_system_prompt_tokens": 395 }, - "claude-3-opus-latest": { - "cache_creation_input_token_cost": 1.875e-05, - "cache_creation_input_token_cost_above_1hr": 6e-06, - "cache_read_input_token_cost": 1.5e-06, - "deprecation_date": "2025-03-01", - "input_cost_per_token": 1.5e-05, - "litellm_provider": "anthropic", - "max_input_tokens": 200000, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 7.5e-05, - "supports_assistant_prefill": true, - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 395 - }, "claude-4-opus-20250514": { "cache_creation_input_token_cost": 1.875e-05, + "cache_creation_input_token_cost_above_1hr": 3e-05, "cache_read_input_token_cost": 1.5e-06, "input_cost_per_token": 1.5e-05, "litellm_provider": "anthropic", @@ -274,6 +96,7 @@ }, "claude-4-sonnet-20250514": { "cache_creation_input_token_cost": 3.75e-06, + "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, "cache_read_input_token_cost": 3e-07, "cache_read_input_token_cost_above_200k_tokens": 6e-07, @@ -447,6 +270,7 @@ "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -474,6 +298,7 @@ "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -485,18 +310,14 @@ "claude-opus-4-6": { "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, - "cache_creation_input_token_cost_above_200k_tokens": 1.25e-05, "cache_read_input_token_cost": 5e-07, - "cache_read_input_token_cost_above_200k_tokens": 1e-06, "input_cost_per_token": 5e-06, - "input_cost_per_token_above_200k_tokens": 1e-05, "litellm_provider": "anthropic", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 2.5e-05, - "output_cost_per_token_above_200k_tokens": 3.75e-05, "provider_specific_entry": { "fast": 6.0, "us": 1.1 @@ -506,9 +327,13 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": true, + "supports_output_config": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -520,18 +345,14 @@ "claude-opus-4-6-20260205": { "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, - "cache_creation_input_token_cost_above_200k_tokens": 1.25e-05, "cache_read_input_token_cost": 5e-07, - "cache_read_input_token_cost_above_200k_tokens": 1e-06, "input_cost_per_token": 5e-06, - "input_cost_per_token_above_200k_tokens": 1e-05, "litellm_provider": "anthropic", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 2.5e-05, - "output_cost_per_token_above_200k_tokens": 3.75e-05, "provider_specific_entry": { "fast": 6.0, "us": 1.1 @@ -541,9 +362,13 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": true, + "supports_output_config": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -555,18 +380,14 @@ "claude-opus-4-6-thinking": { "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, - "cache_creation_input_token_cost_above_200k_tokens": 1.25e-05, "cache_read_input_token_cost": 5e-07, - "cache_read_input_token_cost_above_200k_tokens": 1e-06, "input_cost_per_token": 5e-06, - "input_cost_per_token_above_200k_tokens": 1e-05, "litellm_provider": "anthropic", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 2.5e-05, - "output_cost_per_token_above_200k_tokens": 3.75e-05, "provider_specific_entry": { "fast": 6.0, "us": 1.1 @@ -576,15 +397,127 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": true, + "supports_output_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "tool_use_system_prompt_tokens": 346 + }, + "claude-opus-4-7": { + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "anthropic", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "provider_specific_entry": { + "fast": 6.0, + "us": 1.1 + }, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": true, + "supports_output_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "tool_use_system_prompt_tokens": 346 + }, + "claude-opus-4-7-20260416": { + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "anthropic", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "provider_specific_entry": { + "fast": 6.0, + "us": 1.1 + }, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": true, + "supports_output_config": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "tool_use_system_prompt_tokens": 346 + }, + "claude-opus-4-8": { + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "anthropic", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "provider_specific_entry": { + "fast": 6.0, + "us": 1.1 + }, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": true, + "supports_output_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, "tool_use_system_prompt_tokens": 346 }, "claude-sonnet-4-20250514": { @@ -621,6 +554,7 @@ }, "claude-sonnet-4-5": { "cache_creation_input_token_cost": 3.75e-06, + "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, "cache_read_input_token_cost": 3e-07, "cache_read_input_token_cost_above_200k_tokens": 6e-07, @@ -651,6 +585,7 @@ }, "claude-sonnet-4-5-20250929": { "cache_creation_input_token_cost": 3.75e-06, + "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, "cache_read_input_token_cost": 3e-07, "cache_read_input_token_cost_above_200k_tokens": 6e-07, @@ -682,6 +617,7 @@ }, "claude-sonnet-4-5-20250929-v1:0": { "cache_creation_input_token_cost": 3.75e-06, + "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, "cache_read_input_token_cost": 3e-07, "cache_read_input_token_cost_above_200k_tokens": 6e-07, @@ -707,26 +643,27 @@ }, "claude-sonnet-4-6": { "cache_creation_input_token_cost": 3.75e-06, - "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, + "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, - "cache_read_input_token_cost_above_200k_tokens": 6e-07, "input_cost_per_token": 3e-06, - "input_cost_per_token_above_200k_tokens": 6e-06, "litellm_provider": "anthropic", - "max_input_tokens": 200000, + "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 1.5e-05, - "output_cost_per_token_above_200k_tokens": 2.25e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "supports_adaptive_thinking": true, "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": true, + "supports_output_config": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -735,6 +672,54 @@ "supports_vision": true, "tool_use_system_prompt_tokens": 346 }, + "codex-auto-review": { + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "cache_read_input_token_cost_flex": 2.5e-07, + "cache_read_input_token_cost_priority": 1e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_token_above_272k_tokens": 1e-05, + "input_cost_per_token_batches": 2.5e-06, + "input_cost_per_token_flex": 2.5e-06, + "input_cost_per_token_priority": 1e-05, + "litellm_provider": "openai", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, + "output_cost_per_token_batches": 1.5e-05, + "output_cost_per_token_flex": 1.5e-05, + "output_cost_per_token_priority": 6e-05, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_service_tier": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "deepseek-chat": { "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 2.8e-07, @@ -792,493 +777,250 @@ "supports_reasoning": true, "supports_tool_choice": true }, - "gemini-1.0-pro": { - "input_cost_per_character": 1.25e-07, - "input_cost_per_image": 0.0025, - "input_cost_per_token": 5e-07, - "input_cost_per_video_per_second": 0.002, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 32760, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 3.75e-07, - "output_cost_per_token": 1.5e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#google_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_tool_choice": true - }, - "gemini-1.0-pro-001": { - "deprecation_date": "2025-04-09", - "input_cost_per_character": 1.25e-07, - "input_cost_per_image": 0.0025, - "input_cost_per_token": 5e-07, - "input_cost_per_video_per_second": 0.002, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 32760, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 3.75e-07, - "output_cost_per_token": 1.5e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_tool_choice": true - }, - "gemini-1.0-pro-002": { - "deprecation_date": "2025-04-09", - "input_cost_per_character": 1.25e-07, - "input_cost_per_image": 0.0025, - "input_cost_per_token": 5e-07, - "input_cost_per_video_per_second": 0.002, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 32760, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 3.75e-07, - "output_cost_per_token": 1.5e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_tool_choice": true - }, - "gemini-1.0-pro-vision": { - "input_cost_per_image": 0.0025, - "input_cost_per_token": 5e-07, - "litellm_provider": "vertex_ai-vision-models", - "max_images_per_prompt": 16, - "max_input_tokens": 16384, - "max_output_tokens": 2048, - "max_tokens": 2048, - "max_video_length": 2, - "max_videos_per_prompt": 1, - "mode": "chat", - "output_cost_per_token": 1.5e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gemini-1.0-pro-vision-001": { - "deprecation_date": "2025-04-09", - "input_cost_per_image": 0.0025, - "input_cost_per_token": 5e-07, - "litellm_provider": "vertex_ai-vision-models", - "max_images_per_prompt": 16, - "max_input_tokens": 16384, - "max_output_tokens": 2048, - "max_tokens": 2048, - "max_video_length": 2, - "max_videos_per_prompt": 1, - "mode": "chat", - "output_cost_per_token": 1.5e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gemini-1.0-ultra": { - "input_cost_per_character": 1.25e-07, - "input_cost_per_image": 0.0025, - "input_cost_per_token": 5e-07, - "input_cost_per_video_per_second": 0.002, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 8192, - "max_output_tokens": 2048, - "max_tokens": 2048, - "mode": "chat", - "output_cost_per_character": 3.75e-07, - "output_cost_per_token": 1.5e-06, - "source": "As of Jun, 2024. There is no available doc on vertex ai pricing gemini-1.0-ultra-001. Using gemini-1.0-pro pricing. Got max_tokens info here: https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_tool_choice": true - }, - "gemini-1.0-ultra-001": { - "input_cost_per_character": 1.25e-07, - "input_cost_per_image": 0.0025, - "input_cost_per_token": 5e-07, - "input_cost_per_video_per_second": 0.002, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 8192, - "max_output_tokens": 2048, - "max_tokens": 2048, - "mode": "chat", - "output_cost_per_character": 3.75e-07, - "output_cost_per_token": 1.5e-06, - "source": "As of Jun, 2024. There is no available doc on vertex ai pricing gemini-1.0-ultra-001. Using gemini-1.0-pro pricing. Got max_tokens info here: https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_tool_choice": true - }, - "gemini-1.5-flash": { - "deprecation_date": "2025-09-29", - "input_cost_per_audio_per_second": 2e-06, - "input_cost_per_audio_per_second_above_128k_tokens": 4e-06, - "input_cost_per_character": 1.875e-08, - "input_cost_per_character_above_128k_tokens": 2.5e-07, - "input_cost_per_image": 2e-05, - "input_cost_per_image_above_128k_tokens": 4e-05, - "input_cost_per_token": 7.5e-08, - "input_cost_per_token_above_128k_tokens": 1e-06, - "input_cost_per_video_per_second": 2e-05, - "input_cost_per_video_per_second_above_128k_tokens": 4e-05, + "gemini-2.0-flash": { + "cache_read_input_token_cost": 2.5e-08, + "deprecation_date": "2026-06-01", + "input_cost_per_audio_token": 7e-07, + "input_cost_per_token": 1e-07, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, - "max_input_tokens": 1000000, + "max_input_tokens": 1048576, "max_output_tokens": 8192, "max_pdf_size_mb": 30, "max_tokens": 8192, "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_character": 7.5e-08, - "output_cost_per_character_above_128k_tokens": 1.5e-07, - "output_cost_per_token": 3e-07, - "output_cost_per_token_above_128k_tokens": 6e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", + "output_cost_per_token": 4e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.035, + "search_context_size_low": 0.035, + "search_context_size_medium": 0.035 + }, + "source": "https://ai.google.dev/pricing#2_0flash", + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_audio_input": true, + "supports_audio_output": true, "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_url_context": true, + "supports_vision": true, + "supports_web_search": true }, - "gemini-1.5-flash-001": { - "deprecation_date": "2025-05-24", - "input_cost_per_audio_per_second": 2e-06, - "input_cost_per_audio_per_second_above_128k_tokens": 4e-06, - "input_cost_per_character": 1.875e-08, - "input_cost_per_character_above_128k_tokens": 2.5e-07, - "input_cost_per_image": 2e-05, - "input_cost_per_image_above_128k_tokens": 4e-05, - "input_cost_per_token": 7.5e-08, - "input_cost_per_token_above_128k_tokens": 1e-06, - "input_cost_per_video_per_second": 2e-05, - "input_cost_per_video_per_second_above_128k_tokens": 4e-05, + "gemini-2.0-flash-001": { + "cache_read_input_token_cost": 3.75e-08, + "deprecation_date": "2026-06-01", + "input_cost_per_audio_token": 1e-06, + "input_cost_per_token": 1.5e-07, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, - "max_input_tokens": 1000000, + "max_input_tokens": 1048576, "max_output_tokens": 8192, "max_pdf_size_mb": 30, "max_tokens": 8192, "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_character": 7.5e-08, - "output_cost_per_character_above_128k_tokens": 1.5e-07, - "output_cost_per_token": 3e-07, - "output_cost_per_token_above_128k_tokens": 6e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", + "output_cost_per_token": 6e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.035, + "search_context_size_low": 0.035, + "search_context_size_medium": 0.035 + }, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_audio_output": true, "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, + "gemini-2.0-flash-exp-image-generation": { + "input_cost_per_token": 0.0, + "litellm_provider": "gemini", + "max_images_per_prompt": 3000, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "image_generation", + "output_cost_per_image": 0.039, + "output_cost_per_token": 0.0, + "source": "https://ai.google.dev/pricing", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], "supports_vision": true }, - "gemini-1.5-flash-002": { - "deprecation_date": "2025-09-24", - "input_cost_per_audio_per_second": 2e-06, - "input_cost_per_audio_per_second_above_128k_tokens": 4e-06, - "input_cost_per_character": 1.875e-08, - "input_cost_per_character_above_128k_tokens": 2.5e-07, - "input_cost_per_image": 2e-05, - "input_cost_per_image_above_128k_tokens": 4e-05, + "gemini-2.0-flash-lite": { + "cache_read_input_token_cost": 1.875e-08, + "deprecation_date": "2026-06-01", + "input_cost_per_audio_token": 7.5e-08, "input_cost_per_token": 7.5e-08, - "input_cost_per_token_above_128k_tokens": 1e-06, - "input_cost_per_video_per_second": 2e-05, - "input_cost_per_video_per_second_above_128k_tokens": 4e-05, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, "max_input_tokens": 1048576, "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, + "max_pdf_size_mb": 50, "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_character": 7.5e-08, - "output_cost_per_character_above_128k_tokens": 1.5e-07, "output_cost_per_token": 3e-07, - "output_cost_per_token_above_128k_tokens": 6e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#gemini-1.5-flash", + "search_context_cost_per_query": { + "search_context_size_high": 0.035, + "search_context_size_low": 0.035, + "search_context_size_medium": 0.035 + }, + "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#gemini-2.0-flash", + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_output": true, "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, - "gemini-1.5-flash-exp-0827": { - "deprecation_date": "2025-09-29", - "input_cost_per_audio_per_second": 2e-06, - "input_cost_per_audio_per_second_above_128k_tokens": 4e-06, - "input_cost_per_character": 1.875e-08, - "input_cost_per_character_above_128k_tokens": 2.5e-07, - "input_cost_per_image": 2e-05, - "input_cost_per_image_above_128k_tokens": 4e-05, - "input_cost_per_token": 4.688e-09, - "input_cost_per_token_above_128k_tokens": 1e-06, - "input_cost_per_video_per_second": 2e-05, - "input_cost_per_video_per_second_above_128k_tokens": 4e-05, + "gemini-2.0-flash-lite-001": { + "cache_read_input_token_cost": 1.875e-08, + "deprecation_date": "2026-06-01", + "input_cost_per_audio_token": 7.5e-08, + "input_cost_per_token": 7.5e-08, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, - "max_input_tokens": 1000000, + "max_input_tokens": 1048576, "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, + "max_pdf_size_mb": 50, "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_character": 1.875e-08, - "output_cost_per_character_above_128k_tokens": 3.75e-08, - "output_cost_per_token": 4.6875e-09, - "output_cost_per_token_above_128k_tokens": 9.375e-09, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", + "output_cost_per_token": 3e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.035, + "search_context_size_low": 0.035, + "search_context_size_medium": 0.035 + }, + "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#gemini-2.0-flash", + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_output": true, "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, - "gemini-1.5-flash-preview-0514": { - "deprecation_date": "2025-09-29", - "input_cost_per_audio_per_second": 2e-06, - "input_cost_per_audio_per_second_above_128k_tokens": 4e-06, - "input_cost_per_character": 1.875e-08, - "input_cost_per_character_above_128k_tokens": 2.5e-07, - "input_cost_per_image": 2e-05, - "input_cost_per_image_above_128k_tokens": 4e-05, - "input_cost_per_token": 7.5e-08, - "input_cost_per_token_above_128k_tokens": 1e-06, - "input_cost_per_video_per_second": 2e-05, - "input_cost_per_video_per_second_above_128k_tokens": 4e-05, + "gemini-2.5-computer-use-preview-10-2025": { + "input_cost_per_token": 1.25e-06, + "input_cost_per_token_above_200k_tokens": 2.5e-06, "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, - "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, + "max_input_tokens": 128000, + "max_output_tokens": 64000, + "max_tokens": 64000, "mode": "chat", - "output_cost_per_character": 1.875e-08, - "output_cost_per_character_above_128k_tokens": 3.75e-08, - "output_cost_per_token": 4.6875e-09, - "output_cost_per_token_above_128k_tokens": 9.375e-09, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gemini-1.5-pro": { - "deprecation_date": "2025-09-29", - "input_cost_per_audio_per_second": 3.125e-05, - "input_cost_per_audio_per_second_above_128k_tokens": 6.25e-05, - "input_cost_per_character": 3.125e-07, - "input_cost_per_character_above_128k_tokens": 6.25e-07, - "input_cost_per_image": 0.00032875, - "input_cost_per_image_above_128k_tokens": 0.0006575, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_128k_tokens": 2.5e-06, - "input_cost_per_video_per_second": 0.00032875, - "input_cost_per_video_per_second_above_128k_tokens": 0.0006575, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 2097152, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 1.25e-06, - "output_cost_per_character_above_128k_tokens": 2.5e-06, - "output_cost_per_token": 5e-06, - "output_cost_per_token_above_128k_tokens": 1e-05, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gemini-1.5-pro-001": { - "deprecation_date": "2025-05-24", - "input_cost_per_audio_per_second": 3.125e-05, - "input_cost_per_audio_per_second_above_128k_tokens": 6.25e-05, - "input_cost_per_character": 3.125e-07, - "input_cost_per_character_above_128k_tokens": 6.25e-07, - "input_cost_per_image": 0.00032875, - "input_cost_per_image_above_128k_tokens": 0.0006575, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_128k_tokens": 2.5e-06, - "input_cost_per_video_per_second": 0.00032875, - "input_cost_per_video_per_second_above_128k_tokens": 0.0006575, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 1.25e-06, - "output_cost_per_character_above_128k_tokens": 2.5e-06, - "output_cost_per_token": 5e-06, - "output_cost_per_token_above_128k_tokens": 1e-05, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gemini-1.5-pro-002": { - "deprecation_date": "2025-09-24", - "input_cost_per_audio_per_second": 3.125e-05, - "input_cost_per_audio_per_second_above_128k_tokens": 6.25e-05, - "input_cost_per_character": 3.125e-07, - "input_cost_per_character_above_128k_tokens": 6.25e-07, - "input_cost_per_image": 0.00032875, - "input_cost_per_image_above_128k_tokens": 0.0006575, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_128k_tokens": 2.5e-06, - "input_cost_per_video_per_second": 0.00032875, - "input_cost_per_video_per_second_above_128k_tokens": 0.0006575, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 2097152, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 1.25e-06, - "output_cost_per_character_above_128k_tokens": 2.5e-06, - "output_cost_per_token": 5e-06, - "output_cost_per_token_above_128k_tokens": 1e-05, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#gemini-1.5-pro", + "output_cost_per_token": 1e-05, + "output_cost_per_token_above_200k_tokens": 1.5e-05, + "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/computer-use", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true }, - "gemini-1.5-pro-preview-0215": { - "deprecation_date": "2025-09-29", - "input_cost_per_audio_per_second": 3.125e-05, - "input_cost_per_audio_per_second_above_128k_tokens": 6.25e-05, - "input_cost_per_character": 3.125e-07, - "input_cost_per_character_above_128k_tokens": 6.25e-07, - "input_cost_per_image": 0.00032875, - "input_cost_per_image_above_128k_tokens": 0.0006575, - "input_cost_per_token": 7.8125e-08, - "input_cost_per_token_above_128k_tokens": 1.5625e-07, - "input_cost_per_video_per_second": 0.00032875, - "input_cost_per_video_per_second_above_128k_tokens": 0.0006575, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 1.25e-06, - "output_cost_per_character_above_128k_tokens": 2.5e-06, - "output_cost_per_token": 3.125e-07, - "output_cost_per_token_above_128k_tokens": 6.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, - "gemini-1.5-pro-preview-0409": { - "deprecation_date": "2025-09-29", - "input_cost_per_audio_per_second": 3.125e-05, - "input_cost_per_audio_per_second_above_128k_tokens": 6.25e-05, - "input_cost_per_character": 3.125e-07, - "input_cost_per_character_above_128k_tokens": 6.25e-07, - "input_cost_per_image": 0.00032875, - "input_cost_per_image_above_128k_tokens": 0.0006575, - "input_cost_per_token": 7.8125e-08, - "input_cost_per_token_above_128k_tokens": 1.5625e-07, - "input_cost_per_video_per_second": 0.00032875, - "input_cost_per_video_per_second_above_128k_tokens": 0.0006575, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 1.25e-06, - "output_cost_per_character_above_128k_tokens": 2.5e-06, - "output_cost_per_token": 3.125e-07, - "output_cost_per_token_above_128k_tokens": 6.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_tool_choice": true - }, - "gemini-1.5-pro-preview-0514": { - "deprecation_date": "2025-09-29", - "input_cost_per_audio_per_second": 3.125e-05, - "input_cost_per_audio_per_second_above_128k_tokens": 6.25e-05, - "input_cost_per_character": 3.125e-07, - "input_cost_per_character_above_128k_tokens": 6.25e-07, - "input_cost_per_image": 0.00032875, - "input_cost_per_image_above_128k_tokens": 0.0006575, - "input_cost_per_token": 7.8125e-08, - "input_cost_per_token_above_128k_tokens": 1.5625e-07, - "input_cost_per_video_per_second": 0.00032875, - "input_cost_per_video_per_second_above_128k_tokens": 0.0006575, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 1.25e-06, - "output_cost_per_character_above_128k_tokens": 2.5e-06, - "output_cost_per_token": 3.125e-07, - "output_cost_per_token_above_128k_tokens": 6.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, - "gemini-2.0-flash": { - "cache_read_input_token_cost": 2.5e-08, - "deprecation_date": "2026-03-31", - "input_cost_per_audio_token": 7e-07, - "input_cost_per_token": 1e-07, + "gemini-2.5-flash": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_audio_token": 1e-06, + "input_cost_per_token": 3e-07, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, "max_input_tokens": 1048576, - "max_output_tokens": 8192, + "max_output_tokens": 65535, "max_pdf_size_mb": 30, - "max_tokens": 8192, + "max_tokens": 65535, "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_token": 4e-07, - "source": "https://ai.google.dev/pricing#2_0flash", + "output_cost_per_reasoning_token": 2.5e-06, + "output_cost_per_token": 2.5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.035, + "search_context_size_low": 0.035, + "search_context_size_medium": 0.035 + }, + "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], "supported_modalities": [ "text", "image", @@ -1286,39 +1028,48 @@ "video" ], "supported_output_modalities": [ - "text", - "image" + "text" ], - "supports_audio_input": true, - "supports_audio_output": true, + "supports_audio_output": false, "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_reasoning": true, "supports_response_schema": true, + "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, "supports_vision": true, "supports_web_search": true }, - "gemini-2.0-flash-001": { - "cache_read_input_token_cost": 3.75e-08, - "deprecation_date": "2026-03-31", + "gemini-2.5-flash-image": { + "cache_read_input_token_cost": 3e-08, "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 1.5e-07, + "input_cost_per_token": 3e-07, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 8192, + "max_input_tokens": 32768, + "max_output_tokens": 32768, "max_pdf_size_mb": 30, - "max_tokens": 8192, + "max_tokens": 32768, "max_video_length": 1, "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 6e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "mode": "image_generation", + "output_cost_per_image": 0.039, + "output_cost_per_image_token": 3e-05, + "output_cost_per_reasoning_token": 2.5e-06, + "output_cost_per_token": 2.5e-06, + "rpm": 100000, + "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-2.5-flash-image", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], "supported_modalities": [ "text", "image", @@ -1329,44 +1080,48 @@ "text", "image" ], - "supports_audio_output": true, + "supports_audio_output": false, "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, + "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, + "supports_url_context": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": false, + "tpm": 8000000 }, - "gemini-2.0-flash-exp": { - "cache_read_input_token_cost": 3.75e-08, - "input_cost_per_audio_per_second": 0, - "input_cost_per_audio_per_second_above_128k_tokens": 0, - "input_cost_per_character": 0, - "input_cost_per_character_above_128k_tokens": 0, - "input_cost_per_image": 0, - "input_cost_per_image_above_128k_tokens": 0, - "input_cost_per_token": 1.5e-07, - "input_cost_per_token_above_128k_tokens": 0, - "input_cost_per_video_per_second": 0, - "input_cost_per_video_per_second_above_128k_tokens": 0, + "gemini-2.5-flash-lite": { + "cache_read_input_token_cost": 1e-08, + "input_cost_per_audio_token": 3e-07, + "input_cost_per_token": 1e-07, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, "max_input_tokens": 1048576, - "max_output_tokens": 8192, + "max_output_tokens": 65535, "max_pdf_size_mb": 30, - "max_tokens": 8192, + "max_tokens": 65535, "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_character": 0, - "output_cost_per_character_above_128k_tokens": 0, - "output_cost_per_token": 6e-07, - "output_cost_per_token_above_128k_tokens": 0, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "output_cost_per_reasoning_token": 4e-07, + "output_cost_per_token": 4e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.035, + "search_context_size_low": 0.035, + "search_context_size_medium": 0.035 + }, + "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], "supported_modalities": [ "text", "image", @@ -1374,57 +1129,51 @@ "video" ], "supported_output_modalities": [ - "text", - "image" + "text" ], - "supports_audio_output": true, + "supports_audio_output": false, "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_reasoning": true, "supports_response_schema": true, + "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, + "supports_url_context": true, "supports_vision": true, "supports_web_search": true }, - "gemini-2.0-flash-exp-image-generation": { - "input_cost_per_token": 0.0, - "litellm_provider": "gemini", - "max_images_per_prompt": 3000, - "max_input_tokens": 32768, - "max_output_tokens": 32768, - "max_tokens": 32768, - "mode": "image_generation", - "output_cost_per_image": 0.039, - "output_cost_per_token": 0.0, - "source": "https://ai.google.dev/pricing", - "supported_modalities": [ - "text", - "image" - ], - "supported_output_modalities": [ - "text", - "image" - ], - "supports_vision": true - }, - "gemini-2.0-flash-lite": { - "cache_read_input_token_cost": 1.875e-08, - "deprecation_date": "2026-03-31", - "input_cost_per_audio_token": 7.5e-08, - "input_cost_per_token": 7.5e-08, + "gemini-2.5-flash-lite-preview-06-17": { + "cache_read_input_token_cost": 2.5e-08, + "deprecation_date": "2025-11-18", + "input_cost_per_audio_token": 5e-07, + "input_cost_per_token": 1e-07, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, "max_input_tokens": 1048576, - "max_output_tokens": 8192, - "max_pdf_size_mb": 50, + "max_output_tokens": 65535, + "max_pdf_size_mb": 30, + "max_tokens": 65535, "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_token": 3e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#gemini-2.0-flash", + "output_cost_per_reasoning_token": 4e-07, + "output_cost_per_token": 4e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.035, + "search_context_size_low": 0.035, + "search_context_size_medium": 0.035 + }, + "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], "supported_modalities": [ "text", "image", @@ -1434,34 +1183,48 @@ "supported_output_modalities": [ "text" ], - "supports_audio_output": true, + "supports_audio_output": false, "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, + "supports_url_context": true, "supports_vision": true, "supports_web_search": true }, - "gemini-2.0-flash-lite-001": { - "cache_read_input_token_cost": 1.875e-08, - "deprecation_date": "2026-03-31", - "input_cost_per_audio_token": 7.5e-08, - "input_cost_per_token": 7.5e-08, + "gemini-2.5-flash-lite-preview-09-2025": { + "cache_read_input_token_cost": 1e-08, + "input_cost_per_audio_token": 3e-07, + "input_cost_per_token": 1e-07, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, "max_input_tokens": 1048576, - "max_output_tokens": 8192, - "max_pdf_size_mb": 50, + "max_output_tokens": 65535, + "max_pdf_size_mb": 30, + "max_tokens": 65535, "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_token": 3e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#gemini-2.0-flash", - "supported_modalities": [ + "output_cost_per_reasoning_token": 4e-07, + "output_cost_per_token": 4e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.035, + "search_context_size_low": 0.035, + "search_context_size_medium": 0.035 + }, + "source": "https://developers.googleblog.com/en/continuing-to-bring-you-our-latest-models-with-an-improved-gemini-2-5-flash-and-flash-lite-release/", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ "text", "image", "audio", @@ -1470,181 +1233,119 @@ "supported_output_modalities": [ "text" ], - "supports_audio_output": true, + "supports_audio_output": false, "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, + "supports_url_context": true, "supports_vision": true, "supports_web_search": true }, - "gemini-2.0-flash-live-preview-04-09": { - "cache_read_input_token_cost": 7.5e-08, - "input_cost_per_audio_token": 3e-06, - "input_cost_per_image": 3e-06, - "input_cost_per_token": 5e-07, - "input_cost_per_video_per_second": 3e-06, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, + "gemini-2.5-flash-native-audio-latest": { + "input_cost_per_audio_token": 1e-06, + "input_cost_per_token": 3e-07, + "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, + "max_output_tokens": 8192, + "max_tokens": 8192, "mode": "chat", - "output_cost_per_audio_token": 1.2e-05, - "output_cost_per_token": 2e-06, - "rpm": 10, - "source": "https://cloud.google.com/vertex-ai/docs/generative-ai/model-reference/gemini#gemini-2-0-flash-live-preview-04-09", + "output_cost_per_token": 2.5e-06, + "source": "https://ai.google.dev/pricing", "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions" + "/v1/realtime" ], "supported_modalities": [ "text", - "image", - "audio", - "video" + "audio" ], "supported_output_modalities": [ "text", "audio" ], - "supports_audio_output": true, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 250000 + "supports_audio_input": true, + "supports_audio_output": true }, - "gemini-2.0-flash-preview-image-generation": { - "cache_read_input_token_cost": 2.5e-08, - "deprecation_date": "2025-11-14", - "input_cost_per_audio_token": 7e-07, - "input_cost_per_token": 1e-07, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, + "gemini-2.5-flash-native-audio-preview-09-2025": { + "input_cost_per_audio_token": 1e-06, + "input_cost_per_token": 3e-07, + "litellm_provider": "gemini", "max_input_tokens": 1048576, "max_output_tokens": 8192, - "max_pdf_size_mb": 30, "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_token": 4e-07, - "source": "https://ai.google.dev/pricing#2_0flash", + "output_cost_per_token": 2.5e-06, + "source": "https://ai.google.dev/pricing", + "supported_endpoints": [ + "/v1/realtime" + ], "supported_modalities": [ "text", - "image", - "audio", - "video" + "audio" ], "supported_output_modalities": [ "text", - "image" + "audio" ], "supports_audio_input": true, - "supports_audio_output": true, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true + "supports_audio_output": true }, - "gemini-2.0-flash-thinking-exp": { - "cache_read_input_token_cost": 0.0, - "deprecation_date": "2025-12-02", - "input_cost_per_audio_per_second": 0, - "input_cost_per_audio_per_second_above_128k_tokens": 0, - "input_cost_per_character": 0, - "input_cost_per_character_above_128k_tokens": 0, - "input_cost_per_image": 0, - "input_cost_per_image_above_128k_tokens": 0, - "input_cost_per_token": 0, - "input_cost_per_token_above_128k_tokens": 0, - "input_cost_per_video_per_second": 0, - "input_cost_per_video_per_second_above_128k_tokens": 0, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, + "gemini-2.5-flash-native-audio-preview-12-2025": { + "input_cost_per_audio_token": 1e-06, + "input_cost_per_token": 3e-07, + "litellm_provider": "gemini", "max_input_tokens": 1048576, "max_output_tokens": 8192, - "max_pdf_size_mb": 30, "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_character": 0, - "output_cost_per_character_above_128k_tokens": 0, - "output_cost_per_token": 0, - "output_cost_per_token_above_128k_tokens": 0, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#gemini-2.0-flash", + "output_cost_per_token": 2.5e-06, + "source": "https://ai.google.dev/pricing", + "supported_endpoints": [ + "/v1/realtime" + ], "supported_modalities": [ "text", - "image", - "audio", - "video" + "audio" ], "supported_output_modalities": [ "text", - "image" + "audio" ], - "supports_audio_output": true, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true + "supports_audio_input": true, + "supports_audio_output": true }, - "gemini-2.0-flash-thinking-exp-01-21": { - "cache_read_input_token_cost": 0.0, - "deprecation_date": "2025-12-02", - "input_cost_per_audio_per_second": 0, - "input_cost_per_audio_per_second_above_128k_tokens": 0, - "input_cost_per_character": 0, - "input_cost_per_character_above_128k_tokens": 0, - "input_cost_per_image": 0, - "input_cost_per_image_above_128k_tokens": 0, - "input_cost_per_token": 0, - "input_cost_per_token_above_128k_tokens": 0, - "input_cost_per_video_per_second": 0, - "input_cost_per_video_per_second_above_128k_tokens": 0, + "gemini-2.5-flash-preview-09-2025": { + "cache_read_input_token_cost": 7.5e-08, + "input_cost_per_audio_token": 1e-06, + "input_cost_per_token": 3e-07, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, "max_input_tokens": 1048576, - "max_output_tokens": 65536, + "max_output_tokens": 65535, "max_pdf_size_mb": 30, - "max_tokens": 65536, + "max_tokens": 65535, "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_character": 0, - "output_cost_per_character_above_128k_tokens": 0, - "output_cost_per_token": 0, - "output_cost_per_token_above_128k_tokens": 0, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#gemini-2.0-flash", + "output_cost_per_reasoning_token": 2.5e-06, + "output_cost_per_token": 2.5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.035, + "search_context_size_low": 0.035, + "search_context_size_medium": 0.035 + }, + "source": "https://developers.googleblog.com/en/continuing-to-bring-you-our-latest-models-with-an-improved-gemini-2-5-flash-and-flash-lite-release/", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], "supported_modalities": [ "text", "image", @@ -1652,37 +1353,55 @@ "video" ], "supported_output_modalities": [ - "text", - "image" + "text" ], "supports_audio_output": false, - "supports_function_calling": false, + "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, - "supports_response_schema": false, + "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, + "supports_url_context": true, "supports_vision": true, "supports_web_search": true }, - "gemini-2.0-pro-exp-02-05": { - "cache_read_input_token_cost": 3.125e-07, + "gemini-2.5-flash-preview-tts": { + "input_cost_per_token": 3e-07, + "litellm_provider": "gemini", + "mode": "audio_speech", + "output_cost_per_token": 2.5e-06, + "source": "https://ai.google.dev/pricing", + "supported_endpoints": [ + "/v1/audio/speech" + ] + }, + "gemini-2.5-pro": { + "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, + "cache_read_input_token_cost": 1.25e-07, + "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, "input_cost_per_token": 1.25e-06, "input_cost_per_token_above_200k_tokens": 2.5e-06, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, - "max_input_tokens": 2097152, - "max_output_tokens": 8192, + "max_input_tokens": 1048576, + "max_output_tokens": 65535, "max_pdf_size_mb": 30, - "max_tokens": 8192, + "max_tokens": 65535, "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_above_200k_tokens": 1.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.035, + "search_context_size_low": 0.035, + "search_context_size_medium": 0.035 + }, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -1699,45 +1418,65 @@ ], "supports_audio_input": true, "supports_function_calling": true, - "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_reasoning": true, "supports_response_schema": true, + "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true, "supports_web_search": true }, - "gemini-2.5-computer-use-preview-10-2025": { + "gemini-2.5-pro-preview-tts": { + "cache_read_input_token_cost": 1.25e-07, + "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, + "input_cost_per_audio_token": 7e-07, "input_cost_per_token": 1.25e-06, "input_cost_per_token_above_200k_tokens": 2.5e-06, "litellm_provider": "vertex_ai-language-models", + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, - "max_input_tokens": 128000, - "max_output_tokens": 64000, - "max_tokens": 64000, + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_pdf_size_mb": 30, + "max_tokens": 65535, + "max_video_length": 1, + "max_videos_per_prompt": 10, "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_above_200k_tokens": 1.5e-05, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/computer-use", + "search_context_cost_per_query": { + "search_context_size_high": 0.035, + "search_context_size_low": 0.035, + "search_context_size_medium": 0.035 + }, + "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-2.5-pro-preview", "supported_modalities": [ - "text", - "image" + "text" ], "supported_output_modalities": [ - "text" + "audio" ], - "supports_computer_use": true, + "supports_audio_output": false, "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, - "gemini-2.5-flash": { - "cache_read_input_token_cost": 3e-08, + "gemini-3-flash": { + "cache_read_input_token_cost": 5e-08, + "cache_read_input_token_cost_priority": 9e-08, "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, + "input_cost_per_audio_token_priority": 1.8e-06, + "input_cost_per_token": 5e-07, + "input_cost_per_token_priority": 9e-07, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, @@ -1749,9 +1488,15 @@ "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_reasoning_token": 2.5e-06, - "output_cost_per_token": 2.5e-06, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", + "output_cost_per_reasoning_token": 3e-06, + "output_cost_per_token": 3e-06, + "output_cost_per_token_priority": 5.4e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://ai.google.dev/pricing/gemini-3", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -1768,38 +1513,47 @@ ], "supports_audio_output": false, "supports_function_calling": true, + "supports_native_streaming": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "web_search_billing_unit": "per_query" }, - "gemini-2.5-flash-image": { - "cache_read_input_token_cost": 3e-08, + "gemini-3-flash-preview": { + "cache_read_input_token_cost": 5e-08, + "cache_read_input_token_cost_priority": 9e-08, "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, + "input_cost_per_audio_token_priority": 1.8e-06, + "input_cost_per_token": 5e-07, + "input_cost_per_token_priority": 9e-07, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, - "max_input_tokens": 32768, - "max_output_tokens": 32768, + "max_input_tokens": 1048576, + "max_output_tokens": 65535, "max_pdf_size_mb": 30, - "max_tokens": 32768, + "max_tokens": 65535, "max_video_length": 1, "max_videos_per_prompt": 10, - "mode": "image_generation", - "output_cost_per_image": 0.039, - "output_cost_per_image_token": 3e-05, - "output_cost_per_reasoning_token": 2.5e-06, - "output_cost_per_token": 2.5e-06, - "rpm": 100000, - "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-2.5-flash-image", + "mode": "chat", + "output_cost_per_reasoning_token": 3e-06, + "output_cost_per_token": 3e-06, + "output_cost_per_token_priority": 5.4e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://ai.google.dev/pricing/gemini-3", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -1812,45 +1566,43 @@ "video" ], "supported_output_modalities": [ - "text", - "image" + "text" ], "supports_audio_output": false, "supports_function_calling": true, + "supports_native_streaming": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_reasoning": true, "supports_response_schema": true, + "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, "supports_vision": true, - "supports_web_search": false, - "tpm": 8000000 + "supports_web_search": true, + "web_search_billing_unit": "per_query" }, - "gemini-2.5-flash-image-preview": { - "cache_read_input_token_cost": 7.5e-08, - "deprecation_date": "2026-01-15", - "input_cost_per_audio_token": 1e-06, - "input_cost_per_image_token": 3e-07, - "input_cost_per_token": 3e-07, + "gemini-3-pro-image-preview": { + "input_cost_per_image": 0.0011, + "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, + "max_input_tokens": 65536, + "max_output_tokens": 32768, + "max_tokens": 32768, "mode": "image_generation", - "output_cost_per_image": 0.039, - "output_cost_per_image_token": 3e-05, - "output_cost_per_reasoning_token": 3e-05, - "output_cost_per_token": 3e-05, - "rpm": 100000, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", + "output_cost_per_image": 0.134, + "output_cost_per_image_token": 0.00012, + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_batches": 6e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -1858,31 +1610,33 @@ ], "supported_modalities": [ "text", - "image", - "audio", - "video" + "image" ], "supported_output_modalities": [ "text", "image" ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, + "supports_function_calling": false, "supports_prompt_caching": true, "supports_response_schema": true, + "supports_service_tier": true, "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, "supports_vision": true, "supports_web_search": true, - "tpm": 8000000 + "web_search_billing_unit": "per_query" }, - "gemini-2.5-flash-lite": { - "cache_read_input_token_cost": 1e-08, - "input_cost_per_audio_token": 3e-07, - "input_cost_per_token": 1e-07, + "gemini-3-pro-preview": { + "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, + "cache_read_input_token_cost_priority": 3.6e-07, + "deprecation_date": "2026-03-26", + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, + "input_cost_per_token_batches": 1e-06, + "input_cost_per_token_priority": 3.6e-06, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, @@ -1894,9 +1648,17 @@ "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_reasoning_token": 4e-07, - "output_cost_per_token": 4e-07, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, + "output_cost_per_token_batches": 6e-06, + "output_cost_per_token_priority": 2.16e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -1911,38 +1673,130 @@ "supported_output_modalities": [ "text" ], - "supports_audio_output": false, + "supports_audio_input": true, "supports_function_calling": true, - "supports_parallel_function_calling": true, + "supports_native_streaming": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_url_context": true, + "supports_video_input": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "web_search_billing_unit": "per_query" }, - "gemini-2.5-flash-lite-preview-06-17": { + "gemini-3.1-flash-image": { + "input_cost_per_image": 0.00056, + "input_cost_per_token": 5e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "image_generation", + "output_cost_per_image": 0.0672, + "output_cost_per_image_token": 6e-05, + "output_cost_per_token": 3e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": true, + "supports_web_search": true, + "web_search_billing_unit": "per_query" + }, + "gemini-3.1-flash-image-preview": { + "input_cost_per_image": 0.00056, + "input_cost_per_token": 5e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "image_generation", + "output_cost_per_image": 0.0672, + "output_cost_per_image_token": 6e-05, + "output_cost_per_token": 3e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": true, + "supports_web_search": true, + "web_search_billing_unit": "per_query" + }, + "gemini-3.1-flash-lite": { "cache_read_input_token_cost": 2.5e-08, - "deprecation_date": "2025-11-18", + "cache_read_input_token_cost_batches": 1.25e-08, + "cache_read_input_token_cost_flex": 1.25e-08, + "cache_read_input_token_cost_per_audio_token": 5e-08, + "cache_read_input_token_cost_priority": 4.5e-08, "input_cost_per_audio_token": 5e-07, - "input_cost_per_token": 1e-07, + "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, + "input_cost_per_token_flex": 1.25e-07, + "input_cost_per_token_priority": 4.5e-07, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, "max_input_tokens": 1048576, - "max_output_tokens": 65535, + "max_output_tokens": 65536, "max_pdf_size_mb": 30, - "max_tokens": 65535, + "max_tokens": 65536, "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_reasoning_token": 4e-07, - "output_cost_per_token": 4e-07, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", + "output_cost_per_reasoning_token": 1.5e-06, + "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, + "output_cost_per_token_flex": 7.5e-07, + "output_cost_per_token_priority": 2.7e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3.1-flash-lite", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -1957,37 +1811,50 @@ "supported_output_modalities": [ "text" ], + "supports_audio_input": true, "supports_audio_output": false, + "supports_code_execution": true, + "supports_file_search": true, "supports_function_calling": true, + "supports_native_streaming": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, + "supports_video_input": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "web_search_billing_unit": "per_query" }, - "gemini-2.5-flash-lite-preview-09-2025": { - "cache_read_input_token_cost": 1e-08, - "input_cost_per_audio_token": 3e-07, - "input_cost_per_token": 1e-07, + "gemini-3.1-flash-lite-preview": { + "cache_read_input_token_cost": 2.5e-08, + "cache_read_input_token_cost_per_audio_token": 5e-08, + "input_cost_per_audio_token": 5e-07, + "input_cost_per_token": 2.5e-07, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, "max_input_tokens": 1048576, - "max_output_tokens": 65535, + "max_output_tokens": 65536, "max_pdf_size_mb": 30, - "max_tokens": 65535, + "max_tokens": 65536, "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_reasoning_token": 4e-07, - "output_cost_per_token": 4e-07, - "source": "https://developers.googleblog.com/en/continuing-to-bring-you-our-latest-models-with-an-improved-gemini-2-5-flash-and-flash-lite-release/", + "output_cost_per_reasoning_token": 1.5e-06, + "output_cost_per_token": 1.5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://ai.google.dev/gemini-api/docs/models", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -2002,109 +1869,92 @@ "supported_output_modalities": [ "text" ], + "supports_audio_input": true, "supports_audio_output": false, + "supports_code_execution": true, + "supports_file_search": true, "supports_function_calling": true, + "supports_native_streaming": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, + "supports_video_input": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "web_search_billing_unit": "per_query" }, - "gemini-2.5-flash-native-audio-latest": { - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, + "gemini-3.1-flash-live-preview": { + "input_cost_per_audio_token": 3e-06, + "input_cost_per_image_token": 1e-06, + "input_cost_per_token": 7.5e-07, + "input_cost_per_video_per_second": 3.3333333333333335e-05, "litellm_provider": "gemini", - "max_input_tokens": 1048576, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_input_tokens": 131072, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", - "output_cost_per_token": 2.5e-06, - "source": "https://ai.google.dev/pricing", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 4.5e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/realtime" ], "supported_modalities": [ "text", - "audio" + "image", + "audio", + "video" ], "supported_output_modalities": [ "text", "audio" ], "supports_audio_input": true, - "supports_audio_output": true + "supports_audio_output": true, + "supports_function_calling": true, + "supports_vision": true, + "supports_web_search": true }, - "gemini-2.5-flash-native-audio-preview-09-2025": { - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, - "litellm_provider": "gemini", - "max_input_tokens": 1048576, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 2.5e-06, - "source": "https://ai.google.dev/pricing", - "supported_endpoints": [ - "/v1/realtime" - ], - "supported_modalities": [ - "text", - "audio" - ], - "supported_output_modalities": [ - "text", - "audio" - ], - "supports_audio_input": true, - "supports_audio_output": true - }, - "gemini-2.5-flash-native-audio-preview-12-2025": { - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, - "litellm_provider": "gemini", - "max_input_tokens": 1048576, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 2.5e-06, - "source": "https://ai.google.dev/pricing", - "supported_endpoints": [ - "/v1/realtime" - ], - "supported_modalities": [ - "text", - "audio" - ], - "supported_output_modalities": [ - "text", - "audio" - ], - "supports_audio_input": true, - "supports_audio_output": true - }, - "gemini-2.5-flash-preview-04-17": { - "cache_read_input_token_cost": 3.75e-08, - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 1.5e-07, + "gemini-3.1-pro-high": { + "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, + "cache_read_input_token_cost_priority": 3.6e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, + "input_cost_per_token_batches": 1e-06, + "input_cost_per_token_priority": 3.6e-06, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, "max_input_tokens": 1048576, - "max_output_tokens": 65535, + "max_output_tokens": 65536, "max_pdf_size_mb": 30, - "max_tokens": 65535, + "max_tokens": 65536, "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_reasoning_token": 3.5e-06, - "output_cost_per_token": 6e-07, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", + "output_cost_per_image": 0.00012, + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, + "output_cost_per_token_batches": 6e-06, + "output_cost_per_token_priority": 2.16e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -2119,37 +1969,56 @@ "supported_output_modalities": [ "text" ], - "supports_audio_output": false, + "supports_audio_input": true, "supports_function_calling": true, - "supports_parallel_function_calling": true, + "supports_native_streaming": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "web_search_billing_unit": "per_query" }, - "gemini-2.5-flash-preview-05-20": { - "cache_read_input_token_cost": 7.5e-08, - "deprecation_date": "2025-11-18", - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, + "gemini-3.1-pro-low": { + "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, + "cache_read_input_token_cost_priority": 3.6e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, + "input_cost_per_token_batches": 1e-06, + "input_cost_per_token_priority": 3.6e-06, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, "max_input_tokens": 1048576, - "max_output_tokens": 65535, + "max_output_tokens": 65536, "max_pdf_size_mb": 30, - "max_tokens": 65535, + "max_tokens": 65536, "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_reasoning_token": 2.5e-06, - "output_cost_per_token": 2.5e-06, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", + "output_cost_per_image": 0.00012, + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, + "output_cost_per_token_batches": 6e-06, + "output_cost_per_token_priority": 2.16e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -2164,37 +2033,56 @@ "supported_output_modalities": [ "text" ], - "supports_audio_output": false, + "supports_audio_input": true, "supports_function_calling": true, - "supports_parallel_function_calling": true, + "supports_native_streaming": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, + "supports_video_input": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "web_search_billing_unit": "per_query" }, - "gemini-2.5-flash-preview-09-2025": { - "cache_read_input_token_cost": 7.5e-08, - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, + "gemini-3.1-pro-preview": { + "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, + "cache_read_input_token_cost_priority": 3.6e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, + "input_cost_per_token_batches": 1e-06, + "input_cost_per_token_priority": 3.6e-06, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, "max_input_tokens": 1048576, - "max_output_tokens": 65535, + "max_output_tokens": 65536, "max_pdf_size_mb": 30, - "max_tokens": 65535, + "max_tokens": 65536, "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_reasoning_token": 2.5e-06, - "output_cost_per_token": 2.5e-06, - "source": "https://developers.googleblog.com/en/continuing-to-bring-you-our-latest-models-with-an-improved-gemini-2-5-flash-and-flash-lite-release/", + "output_cost_per_image": 0.00012, + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, + "output_cost_per_token_batches": 6e-06, + "output_cost_per_token_priority": 2.16e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -2209,52 +2097,54 @@ "supported_output_modalities": [ "text" ], - "supports_audio_output": false, + "supports_audio_input": true, "supports_function_calling": true, - "supports_parallel_function_calling": true, + "supports_native_streaming": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, + "supports_video_input": true, "supports_vision": true, - "supports_web_search": true - }, - "gemini-2.5-flash-preview-tts": { - "input_cost_per_token": 3e-07, - "litellm_provider": "gemini", - "mode": "audio_speech", - "output_cost_per_token": 2.5e-06, - "source": "https://ai.google.dev/pricing", - "supported_endpoints": [ - "/v1/audio/speech" - ] + "supports_web_search": true, + "web_search_billing_unit": "per_query" }, - "gemini-2.5-pro": { + "gemini-3.1-pro-preview-customtools": { "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, "max_input_tokens": 1048576, - "max_output_tokens": 65535, + "max_output_tokens": 65536, "max_pdf_size_mb": 30, - "max_tokens": 65535, + "max_tokens": 65536, "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "output_cost_per_image": 0.00012, + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_batches": 6e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", "supported_endpoints": [ "/v1/chat/completions", - "/v1/completions" + "/v1/completions", + "/v1/batch" ], "supported_modalities": [ "text", @@ -2267,21 +2157,26 @@ ], "supports_audio_input": true, "supports_function_calling": true, + "supports_native_streaming": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, + "supports_url_context": true, "supports_video_input": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "web_search_billing_unit": "per_query" }, - "gemini-2.5-pro-exp-03-25": { - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, + "gemini-3.5-flash": { + "cache_read_input_token_cost": 1.5e-07, + "cache_read_input_token_cost_priority": 2.7e-07, + "input_cost_per_audio_token": 1e-06, + "input_cost_per_audio_token_priority": 1.8e-06, + "input_cost_per_token": 1.5e-06, + "input_cost_per_token_priority": 2.7e-06, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, @@ -2293,12 +2188,19 @@ "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "output_cost_per_reasoning_token": 9e-06, + "output_cost_per_token": 9e-06, + "output_cost_per_token_priority": 1.62e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://ai.google.dev/pricing/gemini-3", "supported_endpoints": [ "/v1/chat/completions", - "/v1/completions" + "/v1/completions", + "/v1/batch" ], "supported_modalities": [ "text", @@ -2310,25 +2212,67 @@ "text" ], "supports_audio_input": true, + "supports_audio_output": false, "supports_function_calling": true, + "supports_native_streaming": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_reasoning": true, "supports_response_schema": true, + "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, + "supports_url_context": true, "supports_video_input": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "web_search_billing_unit": "per_query" }, - "gemini-2.5-pro-preview-03-25": { - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, - "deprecation_date": "2025-12-02", - "input_cost_per_audio_token": 1.25e-06, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, - "litellm_provider": "vertex_ai-language-models", + "gemini-embedding-001": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "vertex_ai-embedding-models", + "max_input_tokens": 2048, + "max_tokens": 2048, + "mode": "embedding", + "output_cost_per_token": 0, + "output_vector_size": 3072, + "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models" + }, + "gemini-embedding-2": { + "input_cost_per_audio_per_second": 0.00016, + "input_cost_per_image": 0.00012, + "input_cost_per_token": 2e-07, + "input_cost_per_video_per_second": 0.00079, + "litellm_provider": "vertex_ai-embedding-models", + "max_input_tokens": 8192, + "max_tokens": 8192, + "mode": "embedding", + "output_cost_per_token": 0, + "output_vector_size": 3072, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "supports_multimodal": true, + "uses_embed_content": true + }, + "gemini-embedding-2-preview": { + "input_cost_per_audio_per_second": 0.00016, + "input_cost_per_image": 0.00012, + "input_cost_per_token": 2e-07, + "input_cost_per_video_per_second": 0.00079, + "litellm_provider": "vertex_ai-embedding-models", + "max_input_tokens": 8192, + "max_tokens": 8192, + "mode": "embedding", + "output_cost_per_token": 0, + "output_vector_size": 3072, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "uses_embed_content": true + }, + "gemini-exp-1206": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_audio_token": 1e-06, + "input_cost_per_token": 3e-07, + "litellm_provider": "gemini", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, @@ -2339,10 +2283,16 @@ "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", - "supported_endpoints": [ + "output_cost_per_reasoning_token": 2.5e-06, + "output_cost_per_token": 2.5e-06, + "rpm": 100000, + "search_context_cost_per_query": { + "search_context_size_high": 0.035, + "search_context_size_low": 0.035, + "search_context_size_medium": 0.035 + }, + "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", + "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", "/v1/batch" @@ -2365,17 +2315,29 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, + "supports_url_context": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "tpm": 8000000 }, - "gemini-2.5-pro-preview-05-06": { - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, - "deprecation_date": "2025-12-02", - "input_cost_per_audio_token": 1.25e-06, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, + "gemini-flash-experimental": { + "input_cost_per_character": 0, + "input_cost_per_token": 0, "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 1000000, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "embedding", + "output_cost_per_token": 0, + "output_vector_size": 3072, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "uses_embed_content": true + }, + "gemini-flash-latest": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_audio_token": 1e-06, + "input_cost_per_token": 3e-07, + "litellm_provider": "gemini", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, @@ -2386,8 +2348,14 @@ "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, + "output_cost_per_reasoning_token": 2.5e-06, + "output_cost_per_token": 2.5e-06, + "rpm": 100000, + "search_context_cost_per_query": { + "search_context_size_high": 0.035, + "search_context_size_low": 0.035, + "search_context_size_medium": 0.035 + }, "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", "supported_endpoints": [ "/v1/chat/completions", @@ -2403,9 +2371,6 @@ "supported_output_modalities": [ "text" ], - "supported_regions": [ - "global" - ], "supports_audio_output": false, "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -2415,16 +2380,16 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, + "supports_url_context": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "tpm": 8000000 }, - "gemini-2.5-pro-preview-06-05": { - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, - "input_cost_per_audio_token": 1.25e-06, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, - "litellm_provider": "vertex_ai-language-models", + "gemini-flash-lite-latest": { + "cache_read_input_token_cost": 1e-08, + "input_cost_per_audio_token": 3e-07, + "input_cost_per_token": 1e-07, + "litellm_provider": "gemini", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, @@ -2435,9 +2400,15 @@ "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", + "output_cost_per_reasoning_token": 4e-07, + "output_cost_per_token": 4e-07, + "rpm": 15, + "search_context_cost_per_query": { + "search_context_size_high": 0.035, + "search_context_size_low": 0.035, + "search_context_size_medium": 0.035 + }, + "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-lite", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -2461,52 +2432,15 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, + "supports_url_context": true, "supports_vision": true, - "supports_web_search": true - }, - "gemini-2.5-pro-preview-tts": { - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, - "input_cost_per_audio_token": 7e-07, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, - "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-2.5-pro-preview", - "supported_modalities": [ - "text" - ], - "supported_output_modalities": [ - "audio" - ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "tpm": 250000 }, - "gemini-3-flash-preview": { - "cache_read_input_token_cost": 5e-08, - "cache_read_input_token_cost_priority": 9e-08, - "input_cost_per_audio_token": 1e-06, - "input_cost_per_audio_token_priority": 1.8e-06, - "input_cost_per_token": 5e-07, - "input_cost_per_token_priority": 9e-07, + "gemini-live-2.5-flash-preview-native-audio-09-2025": { + "cache_read_input_token_cost": 7.5e-08, + "input_cost_per_audio_token": 3e-06, + "input_cost_per_token": 3e-07, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, @@ -2517,15 +2451,17 @@ "max_tokens": 65535, "max_video_length": 1, "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_reasoning_token": 3e-06, - "output_cost_per_token": 3e-06, - "output_cost_per_token_priority": 5.4e-06, - "source": "https://ai.google.dev/pricing/gemini-3", + "mode": "realtime", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 2e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.035, + "search_context_size_low": 0.035, + "search_context_size_medium": 0.035 + }, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" + "/vertex_ai/live" ], "supported_modalities": [ "text", @@ -2534,69 +2470,28 @@ "video" ], "supported_output_modalities": [ - "text" + "text", + "audio" ], - "supports_audio_output": false, + "supports_audio_input": true, + "supports_audio_output": true, "supports_function_calling": true, - "supports_native_streaming": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, - "supports_reasoning": true, "supports_response_schema": true, - "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, "supports_vision": true, "supports_web_search": true }, - "gemini-3-pro-image-preview": { - "input_cost_per_image": 0.0011, - "input_cost_per_token": 2e-06, - "input_cost_per_token_batches": 1e-06, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 65536, - "max_output_tokens": 32768, - "max_tokens": 32768, - "mode": "image_generation", - "output_cost_per_image": 0.134, - "output_cost_per_image_token": 0.00012, - "output_cost_per_token": 1.2e-05, - "output_cost_per_token_batches": 6e-06, - "source": "https://ai.google.dev/gemini-api/docs/pricing", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" - ], - "supported_modalities": [ - "text", - "image" - ], - "supported_output_modalities": [ - "text", - "image" - ], - "supports_function_calling": false, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_vision": true, - "supports_web_search": true - }, - "gemini-3-pro-preview": { - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, - "cache_read_input_token_cost": 2e-07, - "cache_read_input_token_cost_above_200k_tokens": 4e-07, - "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, - "cache_read_input_token_cost_priority": 3.6e-07, - "input_cost_per_token": 2e-06, - "input_cost_per_token_above_200k_tokens": 4e-06, - "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, - "input_cost_per_token_batches": 1e-06, - "input_cost_per_token_priority": 3.6e-06, - "litellm_provider": "vertex_ai-language-models", + "gemini-pro-latest": { + "cache_read_input_token_cost": 1.25e-07, + "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, + "input_cost_per_token": 1.25e-06, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "litellm_provider": "gemini", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, @@ -2607,16 +2502,18 @@ "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_token": 1.2e-05, - "output_cost_per_token_above_200k_tokens": 1.8e-05, - "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, - "output_cost_per_token_batches": 6e-06, - "output_cost_per_token_priority": 2.16e-05, + "output_cost_per_token": 1e-05, + "output_cost_per_token_above_200k_tokens": 1.5e-05, + "rpm": 2000, + "search_context_cost_per_query": { + "search_context_size_high": 0.035, + "search_context_size_low": 0.035, + "search_context_size_medium": 0.035 + }, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", - "/v1/completions", - "/v1/batch" + "/v1/completions" ], "supported_modalities": [ "text", @@ -2629,576 +2526,43 @@ ], "supports_audio_input": true, "supports_function_calling": true, - "supports_native_streaming": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, - "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "tpm": 800000 }, - "gemini-3.1-pro-high": { - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, - "cache_read_input_token_cost": 2e-07, - "cache_read_input_token_cost_above_200k_tokens": 4e-07, - "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, - "cache_read_input_token_cost_priority": 3.6e-07, - "input_cost_per_token": 2e-06, - "input_cost_per_token_above_200k_tokens": 4e-06, - "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, - "input_cost_per_token_batches": 1e-06, - "input_cost_per_token_priority": 3.6e-06, + "gemini-robotics-er-1.5-preview": { + "cache_read_input_token_cost": 0, + "input_cost_per_audio_token": 1e-06, + "input_cost_per_token": 3e-07, "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, "max_input_tokens": 1048576, - "max_output_tokens": 65536, - "max_pdf_size_mb": 30, - "max_tokens": 65536, - "max_video_length": 1, - "max_videos_per_prompt": 10, + "max_output_tokens": 65535, + "max_tokens": 65535, "mode": "chat", - "output_cost_per_image": 0.00012, - "output_cost_per_token": 1.2e-05, - "output_cost_per_token_above_200k_tokens": 1.8e-05, - "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, - "output_cost_per_token_batches": 6e-06, - "output_cost_per_token_priority": 2.16e-05, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "output_cost_per_reasoning_token": 2.5e-06, + "output_cost_per_token": 2.5e-06, + "source": "https://ai.google.dev/gemini-api/docs/models#gemini-robotics-er-1-5-preview", "supported_endpoints": [ "/v1/chat/completions", - "/v1/completions", - "/v1/batch" + "/v1/completions" ], "supported_modalities": [ "text", "image", - "audio", - "video" + "video", + "audio" ], "supported_output_modalities": [ "text" ], - "supports_audio_input": true, - "supports_function_calling": true, - "supports_native_streaming": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_service_tier": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, - "supports_video_input": true, - "supports_vision": true, - "supports_web_search": true - }, - "gemini-3.1-pro-low": { - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, - "cache_read_input_token_cost": 2e-07, - "cache_read_input_token_cost_above_200k_tokens": 4e-07, - "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, - "cache_read_input_token_cost_priority": 3.6e-07, - "input_cost_per_token": 2e-06, - "input_cost_per_token_above_200k_tokens": 4e-06, - "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, - "input_cost_per_token_batches": 1e-06, - "input_cost_per_token_priority": 3.6e-06, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65536, - "max_pdf_size_mb": 30, - "max_tokens": 65536, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_image": 0.00012, - "output_cost_per_token": 1.2e-05, - "output_cost_per_token_above_200k_tokens": 1.8e-05, - "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, - "output_cost_per_token_batches": 6e-06, - "output_cost_per_token_priority": 2.16e-05, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_input": true, - "supports_function_calling": true, - "supports_native_streaming": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_service_tier": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, - "supports_video_input": true, - "supports_vision": true, - "supports_web_search": true - }, - "gemini-3.1-pro-preview": { - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, - "cache_read_input_token_cost": 2e-07, - "cache_read_input_token_cost_above_200k_tokens": 4e-07, - "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, - "cache_read_input_token_cost_priority": 3.6e-07, - "input_cost_per_token": 2e-06, - "input_cost_per_token_above_200k_tokens": 4e-06, - "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, - "input_cost_per_token_batches": 1e-06, - "input_cost_per_token_priority": 3.6e-06, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65536, - "max_pdf_size_mb": 30, - "max_tokens": 65536, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_image": 0.00012, - "output_cost_per_token": 1.2e-05, - "output_cost_per_token_above_200k_tokens": 1.8e-05, - "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, - "output_cost_per_token_batches": 6e-06, - "output_cost_per_token_priority": 2.16e-05, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_input": true, - "supports_function_calling": true, - "supports_native_streaming": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_service_tier": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, - "supports_video_input": true, - "supports_vision": true, - "supports_web_search": true - }, - "gemini-3.1-pro-preview-customtools": { - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, - "cache_read_input_token_cost": 2e-07, - "cache_read_input_token_cost_above_200k_tokens": 4e-07, - "input_cost_per_token": 2e-06, - "input_cost_per_token_above_200k_tokens": 4e-06, - "input_cost_per_token_batches": 1e-06, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65536, - "max_pdf_size_mb": 30, - "max_tokens": 65536, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_image": 0.00012, - "output_cost_per_token": 1.2e-05, - "output_cost_per_token_above_200k_tokens": 1.8e-05, - "output_cost_per_token_batches": 6e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_input": true, - "supports_function_calling": true, - "supports_native_streaming": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, - "supports_video_input": true, - "supports_vision": true, - "supports_web_search": true - }, - "gemini-embedding-001": { - "input_cost_per_token": 1.5e-07, - "litellm_provider": "vertex_ai-embedding-models", - "max_input_tokens": 2048, - "max_tokens": 2048, - "mode": "embedding", - "output_cost_per_token": 0, - "output_vector_size": 3072, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models" - }, - "gemini-exp-1206": { - "cache_read_input_token_cost": 3e-08, - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_reasoning_token": 2.5e-06, - "output_cost_per_token": 2.5e-06, - "rpm": 100000, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 8000000 - }, - "gemini-flash-experimental": { - "input_cost_per_character": 0, - "input_cost_per_token": 0, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 0, - "output_cost_per_token": 0, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/gemini-experimental", - "supports_function_calling": false, - "supports_parallel_function_calling": true, - "supports_tool_choice": true - }, - "gemini-flash-latest": { - "cache_read_input_token_cost": 3e-08, - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_reasoning_token": 2.5e-06, - "output_cost_per_token": 2.5e-06, - "rpm": 100000, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 8000000 - }, - "gemini-flash-lite-latest": { - "cache_read_input_token_cost": 1e-08, - "input_cost_per_audio_token": 3e-07, - "input_cost_per_token": 1e-07, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_reasoning_token": 4e-07, - "output_cost_per_token": 4e-07, - "rpm": 15, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-lite", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 250000 - }, - "gemini-live-2.5-flash-preview-native-audio-09-2025": { - "cache_read_input_token_cost": 7.5e-08, - "input_cost_per_audio_token": 3e-06, - "input_cost_per_token": 3e-07, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_audio_token": 1.2e-05, - "output_cost_per_token": 2e-06, - "source": "https://ai.google.dev/gemini-api/docs/pricing", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text", - "audio" - ], - "supports_audio_input": true, - "supports_audio_output": true, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, - "supports_vision": true, - "supports_web_search": true - }, - "gemini-pro": { - "input_cost_per_character": 1.25e-07, - "input_cost_per_image": 0.0025, - "input_cost_per_token": 5e-07, - "input_cost_per_video_per_second": 0.002, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 32760, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 3.75e-07, - "output_cost_per_token": 1.5e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_tool_choice": true - }, - "gemini-pro-experimental": { - "input_cost_per_character": 0, - "input_cost_per_token": 0, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 0, - "output_cost_per_token": 0, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/gemini-experimental", - "supports_function_calling": false, - "supports_parallel_function_calling": true, - "supports_tool_choice": true - }, - "gemini-pro-latest": { - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, - "rpm": 2000, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_input": true, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_video_input": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 800000 - }, - "gemini-pro-vision": { - "input_cost_per_image": 0.0025, - "input_cost_per_token": 5e-07, - "litellm_provider": "vertex_ai-vision-models", - "max_images_per_prompt": 16, - "max_input_tokens": 16384, - "max_output_tokens": 2048, - "max_tokens": 2048, - "max_video_length": 2, - "max_videos_per_prompt": 1, - "mode": "chat", - "output_cost_per_token": 1.5e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gemini-robotics-er-1.5-preview": { - "cache_read_input_token_cost": 0, - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, - "mode": "chat", - "output_cost_per_reasoning_token": 2.5e-06, - "output_cost_per_token": 2.5e-06, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-robotics-er-1-5-preview", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions" - ], - "supported_modalities": [ - "text", - "image", - "video", - "audio" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_output": false, + "supports_audio_output": false, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": false, @@ -3236,31 +2600,6 @@ "supports_system_messages": true, "supports_tool_choice": true }, - "gpt-3.5-turbo-0301": { - "input_cost_per_token": 1.5e-06, - "litellm_provider": "openai", - "max_input_tokens": 4097, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 2e-06, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, - "gpt-3.5-turbo-0613": { - "input_cost_per_token": 1.5e-06, - "litellm_provider": "openai", - "max_input_tokens": 4097, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 2e-06, - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, "gpt-3.5-turbo-1106": { "deprecation_date": "2026-09-28", "input_cost_per_token": 1e-06, @@ -3288,18 +2627,6 @@ "supports_system_messages": true, "supports_tool_choice": true }, - "gpt-3.5-turbo-16k-0613": { - "input_cost_per_token": 3e-06, - "litellm_provider": "openai", - "max_input_tokens": 16385, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 4e-06, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, "gpt-3.5-turbo-instruct": { "input_cost_per_token": 1.5e-06, "litellm_provider": "text-completion-openai", @@ -3347,6 +2674,7 @@ "supports_tool_choice": true }, "gpt-4-0314": { + "deprecation_date": "2026-03-26", "input_cost_per_token": 3e-05, "litellm_provider": "openai", "max_input_tokens": 8192, @@ -3354,7 +2682,6 @@ "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 6e-05, - "supports_prompt_caching": true, "supports_system_messages": true, "supports_tool_choice": true }, @@ -3387,57 +2714,6 @@ "supports_system_messages": true, "supports_tool_choice": true }, - "gpt-4-1106-vision-preview": { - "deprecation_date": "2024-12-06", - "input_cost_per_token": 1e-05, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 3e-05, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gpt-4-32k": { - "input_cost_per_token": 6e-05, - "litellm_provider": "openai", - "max_input_tokens": 32768, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 0.00012, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, - "gpt-4-32k-0314": { - "input_cost_per_token": 6e-05, - "litellm_provider": "openai", - "max_input_tokens": 32768, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 0.00012, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, - "gpt-4-32k-0613": { - "input_cost_per_token": 6e-05, - "litellm_provider": "openai", - "max_input_tokens": 32768, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 0.00012, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, "gpt-4-turbo": { "input_cost_per_token": 1e-05, "litellm_provider": "openai", @@ -3454,23 +2730,7 @@ "supports_tool_choice": true, "supports_vision": true }, - "gpt-4-turbo-2024-04-09": { - "input_cost_per_token": 1e-05, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 3e-05, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gpt-4-turbo-preview": { + "gpt-4-turbo-2024-04-09": { "input_cost_per_token": 1e-05, "litellm_provider": "openai", "max_input_tokens": 128000, @@ -3483,10 +2743,10 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_system_messages": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_vision": true }, - "gpt-4-vision-preview": { - "deprecation_date": "2024-12-06", + "gpt-4-turbo-preview": { "input_cost_per_token": 1e-05, "litellm_provider": "openai", "max_input_tokens": 128000, @@ -3494,11 +2754,12 @@ "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 3e-05, + "supports_function_calling": true, + "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true + "supports_tool_choice": true }, "gpt-4.1": { "cache_read_input_token_cost": 5e-07, @@ -3535,7 +2796,8 @@ "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "gpt-4.1-2025-04-14": { "cache_read_input_token_cost": 5e-07, @@ -3569,7 +2831,8 @@ "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "gpt-4.1-mini": { "cache_read_input_token_cost": 1e-07, @@ -3606,7 +2869,8 @@ "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "gpt-4.1-mini-2025-04-14": { "cache_read_input_token_cost": 1e-07, @@ -3640,7 +2904,8 @@ "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "gpt-4.1-nano": { "cache_read_input_token_cost": 2.5e-08, @@ -3713,47 +2978,6 @@ "supports_tool_choice": true, "supports_vision": true }, - "gpt-4.5-preview": { - "cache_read_input_token_cost": 3.75e-05, - "input_cost_per_token": 7.5e-05, - "input_cost_per_token_batches": 3.75e-05, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 16384, - "max_tokens": 16384, - "mode": "chat", - "output_cost_per_token": 0.00015, - "output_cost_per_token_batches": 7.5e-05, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gpt-4.5-preview-2025-02-27": { - "cache_read_input_token_cost": 3.75e-05, - "deprecation_date": "2025-07-14", - "input_cost_per_token": 7.5e-05, - "input_cost_per_token_batches": 3.75e-05, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 16384, - "max_tokens": 16384, - "mode": "chat", - "output_cost_per_token": 0.00015, - "output_cost_per_token_batches": 7.5e-05, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true - }, "gpt-4o": { "cache_read_input_token_cost": 1.25e-06, "cache_read_input_token_cost_priority": 2.125e-06, @@ -3857,23 +3081,6 @@ "supports_system_messages": true, "supports_tool_choice": true }, - "gpt-4o-audio-preview-2024-10-01": { - "input_cost_per_audio_token": 4e-05, - "input_cost_per_token": 2.5e-06, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 16384, - "max_tokens": 16384, - "mode": "chat", - "output_cost_per_audio_token": 8e-05, - "output_cost_per_token": 1e-05, - "supports_audio_input": true, - "supports_audio_output": true, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, "gpt-4o-audio-preview-2024-12-17": { "input_cost_per_audio_token": 4e-05, "input_cost_per_token": 2.5e-06, @@ -4056,17 +3263,207 @@ "supports_vision": true, "supports_web_search": true }, - "gpt-4o-mini-search-preview-2025-03-11": { - "cache_read_input_token_cost": 7.5e-08, - "input_cost_per_token": 1.5e-07, - "input_cost_per_token_batches": 7.5e-08, + "gpt-4o-mini-search-preview-2025-03-11": { + "cache_read_input_token_cost": 7.5e-08, + "input_cost_per_token": 1.5e-07, + "input_cost_per_token_batches": 7.5e-08, + "litellm_provider": "openai", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 6e-07, + "output_cost_per_token_batches": 3e-07, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "gpt-4o-mini-transcribe": { + "input_cost_per_audio_token": 1.25e-06, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "openai", + "max_input_tokens": 16000, + "max_output_tokens": 2000, + "mode": "audio_transcription", + "output_cost_per_token": 5e-06, + "supported_endpoints": [ + "/v1/audio/transcriptions" + ] + }, + "gpt-4o-mini-transcribe-2025-03-20": { + "input_cost_per_audio_token": 1.25e-06, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "openai", + "max_input_tokens": 16000, + "max_output_tokens": 2000, + "mode": "audio_transcription", + "output_cost_per_token": 5e-06, + "supported_endpoints": [ + "/v1/audio/transcriptions" + ] + }, + "gpt-4o-mini-transcribe-2025-12-15": { + "input_cost_per_audio_token": 1.25e-06, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "openai", + "max_input_tokens": 16000, + "max_output_tokens": 2000, + "mode": "audio_transcription", + "output_cost_per_token": 5e-06, + "supported_endpoints": [ + "/v1/audio/transcriptions" + ] + }, + "gpt-4o-mini-tts": { + "input_cost_per_token": 2.5e-06, + "litellm_provider": "openai", + "mode": "audio_speech", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_second": 0.00025, + "output_cost_per_token": 1e-05, + "supported_endpoints": [ + "/v1/audio/speech" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "audio" + ] + }, + "gpt-4o-mini-tts-2025-03-20": { + "input_cost_per_token": 2.5e-06, + "litellm_provider": "openai", + "mode": "audio_speech", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_second": 0.00025, + "output_cost_per_token": 1e-05, + "supported_endpoints": [ + "/v1/audio/speech" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "audio" + ] + }, + "gpt-4o-mini-tts-2025-12-15": { + "input_cost_per_token": 2.5e-06, + "litellm_provider": "openai", + "mode": "audio_speech", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_second": 0.00025, + "output_cost_per_token": 1e-05, + "supported_endpoints": [ + "/v1/audio/speech" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "audio" + ] + }, + "gpt-4o-realtime-preview": { + "cache_read_input_token_cost": 2.5e-06, + "input_cost_per_audio_token": 4e-05, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "max_input_tokens": 128000, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "chat", + "output_cost_per_audio_token": 8e-05, + "output_cost_per_token": 2e-05, + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "gpt-4o-realtime-preview-2024-12-17": { + "cache_read_input_token_cost": 2.5e-06, + "input_cost_per_audio_token": 4e-05, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "max_input_tokens": 128000, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "chat", + "output_cost_per_audio_token": 8e-05, + "output_cost_per_token": 2e-05, + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "gpt-4o-realtime-preview-2025-06-03": { + "cache_read_input_token_cost": 2.5e-06, + "input_cost_per_audio_token": 4e-05, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "max_input_tokens": 128000, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "chat", + "output_cost_per_audio_token": 8e-05, + "output_cost_per_token": 2e-05, + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "gpt-4o-search-preview": { + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 2.5e-06, + "input_cost_per_token_batches": 1.25e-06, + "litellm_provider": "openai", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1e-05, + "output_cost_per_token_batches": 5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.05, + "search_context_size_low": 0.03, + "search_context_size_medium": 0.035 + }, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, + "gpt-4o-search-preview-2025-03-11": { + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 2.5e-06, + "input_cost_per_token_batches": 1.25e-06, "litellm_provider": "openai", "max_input_tokens": 128000, "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", - "output_cost_per_token": 6e-07, - "output_cost_per_token_batches": 3e-07, + "output_cost_per_token": 1e-05, + "output_cost_per_token_batches": 5e-06, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -4076,186 +3473,469 @@ "supports_tool_choice": true, "supports_vision": true }, - "gpt-4o-mini-transcribe": { - "input_cost_per_audio_token": 3e-06, - "input_cost_per_token": 1.25e-06, + "gpt-4o-transcribe": { + "input_cost_per_audio_token": 2.5e-06, + "input_cost_per_token": 2.5e-06, "litellm_provider": "openai", "max_input_tokens": 16000, "max_output_tokens": 2000, "mode": "audio_transcription", - "output_cost_per_token": 5e-06, + "output_cost_per_token": 1e-05, "supported_endpoints": [ "/v1/audio/transcriptions" ] }, - "gpt-4o-mini-transcribe-2025-03-20": { - "input_cost_per_audio_token": 3e-06, - "input_cost_per_token": 1.25e-06, + "gpt-4o-transcribe-diarize": { + "input_cost_per_audio_token": 2.5e-06, + "input_cost_per_token": 2.5e-06, "litellm_provider": "openai", "max_input_tokens": 16000, "max_output_tokens": 2000, "mode": "audio_transcription", - "output_cost_per_token": 5e-06, + "output_cost_per_token": 1e-05, "supported_endpoints": [ "/v1/audio/transcriptions" ] }, - "gpt-4o-mini-transcribe-2025-12-15": { - "input_cost_per_audio_token": 3e-06, + "gpt-5": { + "cache_read_input_token_cost": 1.25e-07, + "cache_read_input_token_cost_flex": 6.25e-08, + "cache_read_input_token_cost_priority": 2.5e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_flex": 6.25e-07, + "input_cost_per_token_priority": 2.5e-06, "litellm_provider": "openai", - "max_input_tokens": 16000, - "max_output_tokens": 2000, - "mode": "audio_transcription", - "output_cost_per_token": 5e-06, + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "output_cost_per_token_flex": 5e-06, + "output_cost_per_token_priority": 2e-05, "supported_endpoints": [ - "/v1/audio/transcriptions" - ] + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_service_tier": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": false }, - "gpt-4o-mini-tts": { - "input_cost_per_token": 2.5e-06, + "gpt-5-2025-08-07": { + "cache_read_input_token_cost": 1.25e-07, + "cache_read_input_token_cost_flex": 6.25e-08, + "cache_read_input_token_cost_priority": 2.5e-07, + "input_cost_per_token": 1.25e-06, + "input_cost_per_token_flex": 6.25e-07, + "input_cost_per_token_priority": 2.5e-06, "litellm_provider": "openai", - "mode": "audio_speech", - "output_cost_per_audio_token": 1.2e-05, - "output_cost_per_second": 0.00025, + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", "output_cost_per_token": 1e-05, + "output_cost_per_token_flex": 5e-06, + "output_cost_per_token_priority": 2e-05, "supported_endpoints": [ - "/v1/audio/speech" + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" ], "supported_modalities": [ "text", - "audio" + "image" ], "supported_output_modalities": [ - "audio" - ] + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_service_tier": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": false }, - "gpt-4o-mini-tts-2025-03-20": { - "input_cost_per_token": 2.5e-06, + "gpt-5-chat": { + "cache_read_input_token_cost": 1.25e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", - "mode": "audio_speech", - "output_cost_per_audio_token": 1.2e-05, - "output_cost_per_second": 0.00025, + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1e-05, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": false, + "supports_minimal_reasoning_effort": true, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": false, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": false, + "supports_vision": true, + "supports_xhigh_reasoning_effort": false + }, + "gpt-5-chat-latest": { + "cache_read_input_token_cost": 1.25e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "openai", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1e-05, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": false, + "supports_minimal_reasoning_effort": true, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": false, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": false, + "supports_vision": true, + "supports_xhigh_reasoning_effort": false + }, + "gpt-5-codex": { + "cache_read_input_token_cost": 1.25e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "openai", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", "output_cost_per_token": 1e-05, "supported_endpoints": [ - "/v1/audio/speech" + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": false + }, + "gpt-5-mini": { + "cache_read_input_token_cost": 2.5e-08, + "cache_read_input_token_cost_flex": 1.25e-08, + "cache_read_input_token_cost_priority": 4.5e-08, + "input_cost_per_token": 2.5e-07, + "input_cost_per_token_flex": 1.25e-07, + "input_cost_per_token_priority": 4.5e-07, + "litellm_provider": "openai", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2e-06, + "output_cost_per_token_flex": 1e-06, + "output_cost_per_token_priority": 3.6e-06, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_service_tier": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": false + }, + "gpt-5-mini-2025-08-07": { + "cache_read_input_token_cost": 2.5e-08, + "cache_read_input_token_cost_flex": 1.25e-08, + "cache_read_input_token_cost_priority": 4.5e-08, + "input_cost_per_token": 2.5e-07, + "input_cost_per_token_flex": 1.25e-07, + "input_cost_per_token_priority": 4.5e-07, + "litellm_provider": "openai", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2e-06, + "output_cost_per_token_flex": 1e-06, + "output_cost_per_token_priority": 3.6e-06, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" ], "supported_modalities": [ "text", - "audio" + "image" ], "supported_output_modalities": [ - "audio" - ] + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_service_tier": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": false }, - "gpt-4o-mini-tts-2025-12-15": { - "input_cost_per_token": 2.5e-06, + "gpt-5-nano": { + "cache_read_input_token_cost": 5e-09, + "cache_read_input_token_cost_flex": 2.5e-09, + "input_cost_per_token": 5e-08, + "input_cost_per_token_flex": 2.5e-08, + "input_cost_per_token_priority": 2.5e-06, "litellm_provider": "openai", - "mode": "audio_speech", - "output_cost_per_audio_token": 1.2e-05, - "output_cost_per_second": 0.00025, - "output_cost_per_token": 1e-05, + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-07, + "output_cost_per_token_flex": 2e-07, "supported_endpoints": [ - "/v1/audio/speech" + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" ], "supported_modalities": [ "text", - "audio" + "image" ], "supported_output_modalities": [ - "audio" - ] - }, - "gpt-4o-realtime-preview": { - "cache_read_input_token_cost": 2.5e-06, - "input_cost_per_audio_token": 4e-05, - "input_cost_per_token": 5e-06, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_audio_token": 8e-05, - "output_cost_per_token": 2e-05, - "supports_audio_input": true, - "supports_audio_output": true, + "text" + ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, "supports_system_messages": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": false }, - "gpt-4o-realtime-preview-2024-10-01": { - "cache_creation_input_audio_token_cost": 2e-05, - "cache_read_input_token_cost": 2.5e-06, - "input_cost_per_audio_token": 0.0001, - "input_cost_per_token": 5e-06, + "gpt-5-nano-2025-08-07": { + "cache_read_input_token_cost": 5e-09, + "cache_read_input_token_cost_flex": 2.5e-09, + "input_cost_per_token": 5e-08, + "input_cost_per_token_flex": 2.5e-08, "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 4096, - "max_tokens": 4096, + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, "mode": "chat", - "output_cost_per_audio_token": 0.0002, - "output_cost_per_token": 2e-05, - "supports_audio_input": true, - "supports_audio_output": true, + "output_cost_per_token": 4e-07, + "output_cost_per_token_flex": 2e-07, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, "supports_system_messages": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": false }, - "gpt-4o-realtime-preview-2024-12-17": { - "cache_read_input_token_cost": 2.5e-06, - "input_cost_per_audio_token": 4e-05, - "input_cost_per_token": 5e-06, + "gpt-5-pro": { + "input_cost_per_token": 1.5e-05, + "input_cost_per_token_batches": 7.5e-06, "litellm_provider": "openai", "max_input_tokens": 128000, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_audio_token": 8e-05, - "output_cost_per_token": 2e-05, - "supports_audio_input": true, - "supports_audio_output": true, + "max_output_tokens": 272000, + "max_tokens": 272000, + "mode": "responses", + "output_cost_per_token": 0.00012, + "output_cost_per_token_batches": 6e-05, + "supported_endpoints": [ + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, + "supports_native_streaming": false, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, "supports_system_messages": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": false }, - "gpt-4o-realtime-preview-2025-06-03": { - "cache_read_input_token_cost": 2.5e-06, - "input_cost_per_audio_token": 4e-05, - "input_cost_per_token": 5e-06, + "gpt-5-pro-2025-10-06": { + "input_cost_per_token": 1.5e-05, + "input_cost_per_token_batches": 7.5e-06, "litellm_provider": "openai", "max_input_tokens": 128000, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_audio_token": 8e-05, - "output_cost_per_token": 2e-05, - "supports_audio_input": true, - "supports_audio_output": true, + "max_output_tokens": 272000, + "max_tokens": 272000, + "mode": "responses", + "output_cost_per_token": 0.00012, + "output_cost_per_token_batches": 6e-05, + "supported_endpoints": [ + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, + "supports_native_streaming": false, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, "supports_system_messages": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": false }, - "gpt-4o-search-preview": { - "cache_read_input_token_cost": 1.25e-06, - "input_cost_per_token": 2.5e-06, - "input_cost_per_token_batches": 1.25e-06, + "gpt-5-search-api": { + "cache_read_input_token_cost": 1.25e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 16384, - "max_tokens": 16384, + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 1e-05, - "output_cost_per_token_batches": 5e-06, - "search_context_cost_per_query": { - "search_context_size_high": 0.05, - "search_context_size_low": 0.03, - "search_context_size_medium": 0.035 - }, "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -4263,58 +3943,34 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "supports_xhigh_reasoning_effort": false }, - "gpt-4o-search-preview-2025-03-11": { - "cache_read_input_token_cost": 1.25e-06, - "input_cost_per_token": 2.5e-06, - "input_cost_per_token_batches": 1.25e-06, + "gpt-5-search-api-2025-10-14": { + "cache_read_input_token_cost": 1.25e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 16384, - "max_tokens": 16384, + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 1e-05, - "output_cost_per_token_batches": 5e-06, "supports_function_calling": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true - }, - "gpt-4o-transcribe": { - "input_cost_per_audio_token": 6e-06, - "input_cost_per_token": 2.5e-06, - "litellm_provider": "openai", - "max_input_tokens": 16000, - "max_output_tokens": 2000, - "mode": "audio_transcription", - "output_cost_per_token": 1e-05, - "supported_endpoints": [ - "/v1/audio/transcriptions" - ] - }, - "gpt-4o-transcribe-diarize": { - "input_cost_per_audio_token": 6e-06, - "input_cost_per_token": 2.5e-06, - "litellm_provider": "openai", - "max_input_tokens": 16000, - "max_output_tokens": 2000, - "mode": "audio_transcription", - "output_cost_per_token": 1e-05, - "supported_endpoints": [ - "/v1/audio/transcriptions" - ] + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": false }, - "gpt-5": { + "gpt-5.1": { "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_flex": 6.25e-08, "cache_read_input_token_cost_priority": 2.5e-07, "input_cost_per_token": 1.25e-06, - "input_cost_per_token_flex": 6.25e-07, "input_cost_per_token_priority": 2.5e-06, "litellm_provider": "openai", "max_input_tokens": 272000, @@ -4322,11 +3978,9 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 1e-05, - "output_cost_per_token_flex": 5e-06, "output_cost_per_token_priority": 2e-05, "supported_endpoints": [ "/v1/chat/completions", - "/v1/batch", "/v1/responses" ], "supported_modalities": [ @@ -4334,10 +3988,13 @@ "image" ], "supported_output_modalities": [ - "text" + "text", + "image" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, "supports_native_streaming": true, + "supports_none_reasoning_effort": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -4346,14 +4003,14 @@ "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": false }, - "gpt-5-2025-08-07": { + "gpt-5.1-2025-11-13": { "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_flex": 6.25e-08, "cache_read_input_token_cost_priority": 2.5e-07, "input_cost_per_token": 1.25e-06, - "input_cost_per_token_flex": 6.25e-07, "input_cost_per_token_priority": 2.5e-06, "litellm_provider": "openai", "max_input_tokens": 272000, @@ -4361,11 +4018,9 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 1e-05, - "output_cost_per_token_flex": 5e-06, "output_cost_per_token_priority": 2e-05, "supported_endpoints": [ "/v1/chat/completions", - "/v1/batch", "/v1/responses" ], "supported_modalities": [ @@ -4373,10 +4028,13 @@ "image" ], "supported_output_modalities": [ - "text" + "text", + "image" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, "supports_native_streaming": true, + "supports_none_reasoning_effort": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -4385,20 +4043,24 @@ "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": false }, - "gpt-5-chat": { + "gpt-5.1-chat-latest": { "cache_read_input_token_cost": 1.25e-07, + "cache_read_input_token_cost_priority": 2.5e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_priority": 2.5e-06, "litellm_provider": "openai", "max_input_tokens": 128000, "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 1e-05, + "output_cost_per_token_priority": 2e-05, "supported_endpoints": [ "/v1/chat/completions", - "/v1/batch", "/v1/responses" ], "supported_modalities": [ @@ -4406,42 +4068,13 @@ "image" ], "supported_output_modalities": [ - "text" - ], - "supports_function_calling": false, - "supports_native_streaming": true, - "supports_parallel_function_calling": false, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": false, - "supports_vision": true - }, - "gpt-5-chat-latest": { - "cache_read_input_token_cost": 1.25e-07, - "input_cost_per_token": 1.25e-06, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 16384, - "max_tokens": 16384, - "mode": "chat", - "output_cost_per_token": 1e-05, - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/batch", - "/v1/responses" - ], - "supported_modalities": [ "text", "image" ], - "supported_output_modalities": [ - "text" - ], "supports_function_calling": false, + "supports_minimal_reasoning_effort": true, "supports_native_streaming": true, + "supports_none_reasoning_effort": true, "supports_parallel_function_calling": false, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -4449,17 +4082,22 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": false, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": false }, - "gpt-5-codex": { + "gpt-5.1-codex": { "cache_read_input_token_cost": 1.25e-07, + "cache_read_input_token_cost_priority": 2.5e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_priority": 2.5e-06, "litellm_provider": "openai", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", "output_cost_per_token": 1e-05, + "output_cost_per_token_priority": 2e-05, "supported_endpoints": [ "/v1/responses" ], @@ -4471,7 +4109,9 @@ "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, "supports_native_streaming": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -4479,26 +4119,20 @@ "supports_response_schema": true, "supports_system_messages": false, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": false }, - "gpt-5-mini": { - "cache_read_input_token_cost": 2.5e-08, - "cache_read_input_token_cost_flex": 1.25e-08, - "cache_read_input_token_cost_priority": 4.5e-08, - "input_cost_per_token": 2.5e-07, - "input_cost_per_token_flex": 1.25e-07, - "input_cost_per_token_priority": 4.5e-07, + "gpt-5.1-codex-max": { + "cache_read_input_token_cost": 1.25e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, - "mode": "chat", - "output_cost_per_token": 2e-06, - "output_cost_per_token_flex": 1e-06, - "output_cost_per_token_priority": 3.6e-06, + "mode": "responses", + "output_cost_per_token": 1e-05, "supported_endpoints": [ - "/v1/chat/completions", - "/v1/batch", "/v1/responses" ], "supported_modalities": [ @@ -4509,35 +4143,33 @@ "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, "supports_native_streaming": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, - "supports_service_tier": true, - "supports_system_messages": true, + "supports_system_messages": false, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true }, - "gpt-5-mini-2025-08-07": { + "gpt-5.1-codex-mini": { "cache_read_input_token_cost": 2.5e-08, - "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, "input_cost_per_token": 2.5e-07, - "input_cost_per_token_flex": 1.25e-07, "input_cost_per_token_priority": 4.5e-07, "litellm_provider": "openai", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, - "mode": "chat", + "mode": "responses", "output_cost_per_token": 2e-06, - "output_cost_per_token_flex": 1e-06, "output_cost_per_token_priority": 3.6e-06, "supported_endpoints": [ - "/v1/chat/completions", - "/v1/batch", "/v1/responses" ], "supported_modalities": [ @@ -4548,30 +4180,32 @@ "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, "supports_native_streaming": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, - "supports_service_tier": true, - "supports_system_messages": true, + "supports_system_messages": false, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": false }, - "gpt-5-nano": { - "cache_read_input_token_cost": 5e-09, - "cache_read_input_token_cost_flex": 2.5e-09, - "input_cost_per_token": 5e-08, - "input_cost_per_token_flex": 2.5e-08, - "input_cost_per_token_priority": 2.5e-06, + "gpt-5.2": { + "cache_read_input_token_cost": 1.75e-07, + "cache_read_input_token_cost_priority": 3.5e-07, + "input_cost_per_token": 1.75e-06, + "input_cost_per_token_priority": 3.5e-06, "litellm_provider": "openai", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 4e-07, - "output_cost_per_token_flex": 2e-07, + "output_cost_per_token": 1.4e-05, + "output_cost_per_token_priority": 2.8e-05, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -4582,31 +4216,37 @@ "image" ], "supported_output_modalities": [ - "text" + "text", + "image" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, "supports_native_streaming": true, + "supports_none_reasoning_effort": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true }, - "gpt-5-nano-2025-08-07": { - "cache_read_input_token_cost": 5e-09, - "cache_read_input_token_cost_flex": 2.5e-09, - "input_cost_per_token": 5e-08, - "input_cost_per_token_flex": 2.5e-08, + "gpt-5.2-2025-12-11": { + "cache_read_input_token_cost": 1.75e-07, + "cache_read_input_token_cost_priority": 3.5e-07, + "input_cost_per_token": 1.75e-06, + "input_cost_per_token_priority": 3.5e-06, "litellm_provider": "openai", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 4e-07, - "output_cost_per_token_flex": 2e-07, + "output_cost_per_token": 1.4e-05, + "output_cost_per_token_priority": 2.8e-05, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -4617,31 +4257,39 @@ "image" ], "supported_output_modalities": [ - "text" + "text", + "image" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, "supports_native_streaming": true, + "supports_none_reasoning_effort": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true }, - "gpt-5-pro": { - "input_cost_per_token": 1.5e-05, - "input_cost_per_token_batches": 7.5e-06, + "gpt-5.2-chat-latest": { + "cache_read_input_token_cost": 1.75e-07, + "cache_read_input_token_cost_priority": 3.5e-07, + "input_cost_per_token": 1.75e-06, + "input_cost_per_token_priority": 3.5e-06, "litellm_provider": "openai", "max_input_tokens": 128000, - "max_output_tokens": 272000, - "max_tokens": 272000, - "mode": "responses", - "output_cost_per_token": 0.00012, - "output_cost_per_token_batches": 6e-05, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.4e-05, + "output_cost_per_token_priority": 2.8e-05, "supported_endpoints": [ - "/v1/batch", + "/v1/chat/completions", "/v1/responses" ], "supported_modalities": [ @@ -4652,7 +4300,9 @@ "text" ], "supports_function_calling": true, - "supports_native_streaming": false, + "supports_minimal_reasoning_effort": true, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -4661,20 +4311,22 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "supports_xhigh_reasoning_effort": false }, - "gpt-5-pro-2025-10-06": { - "input_cost_per_token": 1.5e-05, - "input_cost_per_token_batches": 7.5e-06, + "gpt-5.2-codex": { + "cache_read_input_token_cost": 1.75e-07, + "cache_read_input_token_cost_priority": 3.5e-07, + "input_cost_per_token": 1.75e-06, + "input_cost_per_token_priority": 3.5e-06, "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 272000, - "max_tokens": 272000, + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, "mode": "responses", - "output_cost_per_token": 0.00012, - "output_cost_per_token_batches": 6e-05, + "output_cost_per_token": 1.4e-05, + "output_cost_per_token_priority": 2.8e-05, "supported_endpoints": [ - "/v1/batch", "/v1/responses" ], "supported_modalities": [ @@ -4685,69 +4337,30 @@ "text" ], "supports_function_calling": true, - "supports_native_streaming": false, + "supports_minimal_reasoning_effort": true, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true - }, - "gpt-5-search-api": { - "cache_read_input_token_cost": 1.25e-07, - "input_cost_per_token": 1.25e-06, - "litellm_provider": "openai", - "max_input_tokens": 272000, - "max_output_tokens": 128000, - "max_tokens": 128000, - "mode": "chat", - "output_cost_per_token": 1e-05, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true - }, - "gpt-5-search-api-2025-10-14": { - "cache_read_input_token_cost": 1.25e-07, - "input_cost_per_token": 1.25e-06, - "litellm_provider": "openai", - "max_input_tokens": 272000, - "max_output_tokens": 128000, - "max_tokens": 128000, - "mode": "chat", - "output_cost_per_token": 1e-05, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, + "supports_system_messages": false, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true }, - "gpt-5.1": { - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_priority": 2.5e-07, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_priority": 2.5e-06, + "gpt-5.2-pro": { + "input_cost_per_token": 2.1e-05, "litellm_provider": "openai", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, - "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_priority": 2e-05, + "mode": "responses", + "output_cost_per_token": 0.000168, "supported_endpoints": [ - "/v1/chat/completions", + "/v1/batch", "/v1/responses" ], "supported_modalities": [ @@ -4755,35 +4368,33 @@ "image" ], "supported_output_modalities": [ - "text", - "image" + "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, "supports_native_streaming": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, - "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true }, - "gpt-5.1-2025-11-13": { - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_priority": 2.5e-07, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_priority": 2.5e-06, + "gpt-5.2-pro-2025-12-11": { + "input_cost_per_token": 2.1e-05, "litellm_provider": "openai", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, - "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_priority": 2e-05, + "mode": "responses", + "output_cost_per_token": 0.000168, "supported_endpoints": [ - "/v1/chat/completions", + "/v1/batch", "/v1/responses" ], "supported_modalities": [ @@ -4791,33 +4402,35 @@ "image" ], "supported_output_modalities": [ - "text", - "image" + "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, "supports_native_streaming": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, - "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true }, - "gpt-5.1-chat-latest": { - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_priority": 2.5e-07, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_priority": 2.5e-06, + "gpt-5.3-chat-latest": { + "cache_read_input_token_cost": 1.75e-07, + "cache_read_input_token_cost_priority": 3.5e-07, + "input_cost_per_token": 1.75e-06, + "input_cost_per_token_priority": 3.5e-06, "litellm_provider": "openai", "max_input_tokens": 128000, "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_priority": 2e-05, + "output_cost_per_token": 1.4e-05, + "output_cost_per_token_priority": 2.8e-05, "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -4827,32 +4440,35 @@ "image" ], "supported_output_modalities": [ - "text", - "image" + "text" ], - "supports_function_calling": false, + "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, "supports_native_streaming": true, - "supports_parallel_function_calling": false, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, - "supports_tool_choice": false, - "supports_vision": true + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": false }, - "gpt-5.1-codex": { - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_priority": 2.5e-07, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_priority": 2.5e-06, + "gpt-5.3-codex": { + "cache_read_input_token_cost": 1.75e-07, + "cache_read_input_token_cost_priority": 3.5e-07, + "input_cost_per_token": 1.75e-06, + "input_cost_per_token_priority": 3.5e-06, "litellm_provider": "openai", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", - "output_cost_per_token": 1e-05, - "output_cost_per_token_priority": 2e-05, + "output_cost_per_token": 1.4e-05, + "output_cost_per_token_priority": 2.8e-05, "supported_endpoints": [ "/v1/responses" ], @@ -4864,7 +4480,9 @@ "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, "supports_native_streaming": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -4872,17 +4490,22 @@ "supports_response_schema": true, "supports_system_messages": false, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": false }, - "gpt-5.1-codex-max": { - "cache_read_input_token_cost": 1.25e-07, - "input_cost_per_token": 1.25e-06, + "gpt-5.3-codex-spark": { + "cache_read_input_token_cost": 1.75e-07, + "cache_read_input_token_cost_priority": 3.5e-07, + "input_cost_per_token": 1.75e-06, + "input_cost_per_token_priority": 3.5e-06, "litellm_provider": "openai", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", - "output_cost_per_token": 1e-05, + "output_cost_per_token": 1.4e-05, + "output_cost_per_token_priority": 2.8e-05, "supported_endpoints": [ "/v1/responses" ], @@ -4894,7 +4517,9 @@ "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, "supports_native_streaming": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -4902,21 +4527,33 @@ "supports_response_schema": true, "supports_system_messages": false, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": false }, - "gpt-5.1-codex-mini": { - "cache_read_input_token_cost": 2.5e-08, - "cache_read_input_token_cost_priority": 4.5e-08, - "input_cost_per_token": 2.5e-07, - "input_cost_per_token_priority": 4.5e-07, + "gpt-5.4": { + "cache_read_input_token_cost": 2.5e-07, + "cache_read_input_token_cost_above_272k_tokens": 5e-07, + "cache_read_input_token_cost_flex": 1.3e-07, + "cache_read_input_token_cost_priority": 5e-07, + "input_cost_per_token": 2.5e-06, + "input_cost_per_token_above_272k_tokens": 5e-06, + "input_cost_per_token_batches": 1.25e-06, + "input_cost_per_token_flex": 1.25e-06, + "input_cost_per_token_priority": 5e-06, "litellm_provider": "openai", - "max_input_tokens": 272000, + "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, - "mode": "responses", - "output_cost_per_token": 2e-06, - "output_cost_per_token_priority": 3.6e-06, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "output_cost_per_token_above_272k_tokens": 2.25e-05, + "output_cost_per_token_batches": 7.5e-06, + "output_cost_per_token_flex": 7.5e-06, + "output_cost_per_token_priority": 3e-05, "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", "/v1/responses" ], "supported_modalities": [ @@ -4927,28 +4564,40 @@ "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, "supports_native_streaming": true, + "supports_none_reasoning_effort": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, - "supports_system_messages": false, + "supports_service_tier": true, + "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_xhigh_reasoning_effort": true }, - "gpt-5.2": { - "cache_read_input_token_cost": 1.75e-07, - "cache_read_input_token_cost_priority": 3.5e-07, - "input_cost_per_token": 1.75e-06, - "input_cost_per_token_priority": 3.5e-06, + "gpt-5.4-2026-03-05": { + "cache_read_input_token_cost": 2.5e-07, + "cache_read_input_token_cost_above_272k_tokens": 5e-07, + "cache_read_input_token_cost_flex": 1.3e-07, + "cache_read_input_token_cost_priority": 5e-07, + "input_cost_per_token": 2.5e-06, + "input_cost_per_token_above_272k_tokens": 5e-06, + "input_cost_per_token_batches": 1.25e-06, + "input_cost_per_token_flex": 1.25e-06, + "input_cost_per_token_priority": 5e-06, "litellm_provider": "openai", - "max_input_tokens": 272000, + "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 1.4e-05, - "output_cost_per_token_priority": 2.8e-05, + "output_cost_per_token": 1.5e-05, + "output_cost_per_token_above_272k_tokens": 2.25e-05, + "output_cost_per_token_batches": 7.5e-06, + "output_cost_per_token_flex": 7.5e-06, + "output_cost_per_token_priority": 3e-05, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -4959,8 +4608,7 @@ "image" ], "supported_output_modalities": [ - "text", - "image" + "text" ], "supports_function_calling": true, "supports_native_streaming": true, @@ -4974,18 +4622,24 @@ "supports_tool_choice": true, "supports_vision": true }, - "gpt-5.2-2025-12-11": { - "cache_read_input_token_cost": 1.75e-07, - "cache_read_input_token_cost_priority": 3.5e-07, - "input_cost_per_token": 1.75e-06, - "input_cost_per_token_priority": 3.5e-06, + "gpt-5.4-mini": { + "cache_read_input_token_cost": 7.5e-08, + "cache_read_input_token_cost_batches": 3.75e-08, + "cache_read_input_token_cost_flex": 3.75e-08, + "cache_read_input_token_cost_priority": 1.5e-07, + "input_cost_per_token": 7.5e-07, + "input_cost_per_token_batches": 3.75e-07, + "input_cost_per_token_flex": 3.75e-07, + "input_cost_per_token_priority": 1.5e-06, "litellm_provider": "openai", - "max_input_tokens": 272000, + "max_input_tokens": 400000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 1.4e-05, - "output_cost_per_token_priority": 2.8e-05, + "output_cost_per_token": 4.5e-06, + "output_cost_per_token_batches": 2.25e-06, + "output_cost_per_token_flex": 2.25e-06, + "output_cost_per_token_priority": 9e-06, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -4996,11 +4650,12 @@ "image" ], "supported_output_modalities": [ - "text", - "image" + "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, + "supports_none_reasoning_effort": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -5009,22 +4664,31 @@ "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true }, - "gpt-5.2-chat-latest": { - "cache_read_input_token_cost": 1.75e-07, - "cache_read_input_token_cost_priority": 3.5e-07, - "input_cost_per_token": 1.75e-06, - "input_cost_per_token_priority": 3.5e-06, + "gpt-5.4-mini-2026-03-17": { + "cache_read_input_token_cost": 7.5e-08, + "cache_read_input_token_cost_batches": 3.75e-08, + "cache_read_input_token_cost_flex": 3.75e-08, + "cache_read_input_token_cost_priority": 1.5e-07, + "input_cost_per_token": 7.5e-07, + "input_cost_per_token_batches": 3.75e-07, + "input_cost_per_token_flex": 3.75e-07, + "input_cost_per_token_priority": 1.5e-06, "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 16384, - "max_tokens": 16384, + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 1.4e-05, - "output_cost_per_token_priority": 2.8e-05, + "output_cost_per_token": 4.5e-06, + "output_cost_per_token_batches": 2.25e-06, + "output_cost_per_token_flex": 2.25e-06, + "output_cost_per_token_priority": 9e-06, "supported_endpoints": [ "/v1/chat/completions", + "/v1/batch", "/v1/responses" ], "supported_modalities": [ @@ -5035,29 +4699,39 @@ "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, + "supports_none_reasoning_effort": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true }, - "gpt-5.2-codex": { - "cache_read_input_token_cost": 1.75e-07, - "cache_read_input_token_cost_priority": 3.5e-07, - "input_cost_per_token": 1.75e-06, - "input_cost_per_token_priority": 3.5e-06, + "gpt-5.4-nano": { + "cache_read_input_token_cost": 2e-08, + "cache_read_input_token_cost_batches": 1e-08, + "cache_read_input_token_cost_flex": 1e-08, + "input_cost_per_token": 2e-07, + "input_cost_per_token_batches": 1e-07, + "input_cost_per_token_flex": 1e-07, "litellm_provider": "openai", - "max_input_tokens": 272000, + "max_input_tokens": 400000, "max_output_tokens": 128000, "max_tokens": 128000, - "mode": "responses", - "output_cost_per_token": 1.4e-05, - "output_cost_per_token_priority": 2.8e-05, + "mode": "chat", + "output_cost_per_token": 1.25e-06, + "output_cost_per_token_batches": 6.25e-07, + "output_cost_per_token_flex": 6.25e-07, "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", "/v1/responses" ], "supported_modalities": [ @@ -5068,25 +4742,38 @@ "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, + "supports_none_reasoning_effort": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, - "supports_system_messages": false, + "supports_service_tier": true, + "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true }, - "gpt-5.2-pro": { - "input_cost_per_token": 2.1e-05, + "gpt-5.4-nano-2026-03-17": { + "cache_read_input_token_cost": 2e-08, + "cache_read_input_token_cost_batches": 1e-08, + "cache_read_input_token_cost_flex": 1e-08, + "input_cost_per_token": 2e-07, + "input_cost_per_token_batches": 1e-07, + "input_cost_per_token_flex": 1e-07, "litellm_provider": "openai", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, - "mode": "responses", - "output_cost_per_token": 0.000168, + "mode": "chat", + "output_cost_per_token": 1.25e-06, + "output_cost_per_token_batches": 6.25e-07, + "output_cost_per_token_flex": 6.25e-07, "supported_endpoints": [ + "/v1/chat/completions", "/v1/batch", "/v1/responses" ], @@ -5098,28 +4785,40 @@ "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, + "supports_none_reasoning_effort": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true }, - "gpt-5.2-pro-2025-12-11": { - "input_cost_per_token": 2.1e-05, + "gpt-5.4-pro": { + "cache_read_input_token_cost": 3e-06, + "cache_read_input_token_cost_above_272k_tokens": 6e-06, + "input_cost_per_token": 3e-05, + "input_cost_per_token_above_272k_tokens": 6e-05, + "input_cost_per_token_batches": 1.5e-05, + "input_cost_per_token_flex": 1.5e-05, "litellm_provider": "openai", - "max_input_tokens": 272000, + "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", - "output_cost_per_token": 0.000168, + "output_cost_per_token": 0.00018, + "output_cost_per_token_above_272k_tokens": 0.00027, + "output_cost_per_token_batches": 9e-05, + "output_cost_per_token_flex": 9e-05, "supported_endpoints": [ - "/v1/batch", - "/v1/responses" + "/v1/responses", + "/v1/batch" ], "supported_modalities": [ "text", @@ -5129,61 +4828,87 @@ "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, "supports_native_streaming": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, - "supports_response_schema": true, + "supports_response_schema": false, + "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true }, - "gpt-5.4": { - "cache_read_input_token_cost": 2.5e-07, - "input_cost_per_token": 2.5e-06, + "gpt-5.4-pro-2026-03-05": { + "cache_read_input_token_cost": 3e-06, + "cache_read_input_token_cost_above_272k_tokens": 6e-06, + "input_cost_per_token": 3e-05, + "input_cost_per_token_above_272k_tokens": 6e-05, + "input_cost_per_token_batches": 1.5e-05, + "input_cost_per_token_flex": 1.5e-05, "litellm_provider": "openai", "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, - "mode": "chat", - "output_cost_per_token": 1.5e-05, + "mode": "responses", + "output_cost_per_token": 0.00018, + "output_cost_per_token_above_272k_tokens": 0.00027, + "output_cost_per_token_batches": 9e-05, + "output_cost_per_token_flex": 9e-05, "supported_endpoints": [ - "/v1/chat/completions", - "/v1/responses" + "/v1/responses", + "/v1/batch" ], "supported_modalities": [ "text", "image" ], "supported_output_modalities": [ - "text", - "image" + "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, "supports_native_streaming": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true }, - "codex-auto-review": { - "cache_read_input_token_cost": 2.5e-07, - "input_cost_per_token": 2.5e-06, + "gpt-5.5": { + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "cache_read_input_token_cost_flex": 2.5e-07, + "cache_read_input_token_cost_priority": 1e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_token_above_272k_tokens": 1e-05, + "input_cost_per_token_batches": 2.5e-06, + "input_cost_per_token_flex": 2.5e-06, + "input_cost_per_token_priority": 1e-05, "litellm_provider": "openai", "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 1.5e-05, + "output_cost_per_token": 3e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, + "output_cost_per_token_batches": 1.5e-05, + "output_cost_per_token_flex": 1.5e-05, + "output_cost_per_token_priority": 6e-05, "supported_endpoints": [ "/v1/chat/completions", + "/v1/batch", "/v1/responses" ], "supported_modalities": [ @@ -5191,11 +4916,12 @@ "image" ], "supported_output_modalities": [ - "text", - "image" + "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, + "supports_none_reasoning_effort": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -5204,17 +4930,30 @@ "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true }, - "gpt-5.4-mini": { - "cache_read_input_token_cost": 7.5e-08, - "input_cost_per_token": 7.5e-07, + "gpt-5.5-2026-04-23": { + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "cache_read_input_token_cost_flex": 2.5e-07, + "cache_read_input_token_cost_priority": 1e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_token_above_272k_tokens": 1e-05, + "input_cost_per_token_batches": 2.5e-06, + "input_cost_per_token_flex": 2.5e-06, + "input_cost_per_token_priority": 1e-05, "litellm_provider": "openai", - "max_input_tokens": 400000, + "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 4.5e-06, + "output_cost_per_token": 3e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, + "output_cost_per_token_batches": 1.5e-05, + "output_cost_per_token_flex": 1.5e-05, + "output_cost_per_token_priority": 6e-05, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -5228,7 +4967,9 @@ "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, + "supports_none_reasoning_effort": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -5237,21 +4978,29 @@ "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true }, - "gpt-5.4-nano": { - "cache_read_input_token_cost": 2e-08, - "input_cost_per_token": 2e-07, + "gpt-5.5-pro": { + "cache_read_input_token_cost": 3e-06, + "cache_read_input_token_cost_above_272k_tokens": 6e-06, + "input_cost_per_token": 3e-05, + "input_cost_per_token_above_272k_tokens": 6e-05, + "input_cost_per_token_batches": 1.5e-05, + "input_cost_per_token_flex": 1.5e-05, "litellm_provider": "openai", - "max_input_tokens": 400000, + "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, - "mode": "chat", - "output_cost_per_token": 1.25e-06, + "mode": "responses", + "output_cost_per_token": 0.00018, + "output_cost_per_token_above_272k_tokens": 0.00027, + "output_cost_per_token_batches": 9e-05, + "output_cost_per_token_flex": 9e-05, "supported_endpoints": [ - "/v1/chat/completions", - "/v1/batch", - "/v1/responses" + "/v1/responses", + "/v1/batch" ], "supported_modalities": [ "text", @@ -5261,30 +5010,41 @@ "text" ], "supports_function_calling": true, + "supports_low_reasoning_effort": false, + "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, - "supports_response_schema": true, + "supports_response_schema": false, + "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true }, - "gpt-5.3-codex": { - "cache_read_input_token_cost": 1.75e-07, - "cache_read_input_token_cost_priority": 3.5e-07, - "input_cost_per_token": 1.75e-06, - "input_cost_per_token_priority": 3.5e-06, + "gpt-5.5-pro-2026-04-23": { + "cache_read_input_token_cost": 3e-06, + "cache_read_input_token_cost_above_272k_tokens": 6e-06, + "input_cost_per_token": 3e-05, + "input_cost_per_token_above_272k_tokens": 6e-05, + "input_cost_per_token_batches": 1.5e-05, + "input_cost_per_token_flex": 1.5e-05, "litellm_provider": "openai", - "max_input_tokens": 272000, + "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", - "output_cost_per_token": 1.4e-05, - "output_cost_per_token_priority": 2.8e-05, + "output_cost_per_token": 0.00018, + "output_cost_per_token_above_272k_tokens": 0.00027, + "output_cost_per_token_batches": 9e-05, + "output_cost_per_token_flex": 9e-05, "supported_endpoints": [ - "/v1/responses" + "/v1/responses", + "/v1/batch" ], "supported_modalities": [ "text", @@ -5294,15 +5054,21 @@ "text" ], "supports_function_calling": true, + "supports_low_reasoning_effort": false, + "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": false, + "supports_response_schema": false, + "supports_service_tier": true, + "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true }, "gpt-audio": { "input_cost_per_audio_token": 3.2e-05, @@ -5340,6 +5106,39 @@ "supports_tool_choice": true, "supports_vision": false }, + "gpt-audio-1.5": { + "input_cost_per_audio_token": 3.2e-05, + "input_cost_per_token": 2.5e-06, + "litellm_provider": "openai", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_audio_token": 6.4e-05, + "output_cost_per_token": 1e-05, + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": false, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, "gpt-audio-2025-08-28": { "input_cost_per_audio_token": 3.2e-05, "input_cost_per_token": 2.5e-06, @@ -5540,6 +5339,38 @@ "supports_pdf_input": true, "supports_vision": true }, + "gpt-image-2": { + "cache_read_input_image_token_cost": 2e-06, + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_image_token": 8e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "output_cost_per_image_token": 3e-05, + "output_cost_per_token": 1e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_pdf_input": true, + "supports_vision": true + }, + "gpt-image-2-2026-04-21": { + "cache_read_input_image_token_cost": 2e-06, + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_image_token": 8e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "output_cost_per_image_token": 3e-05, + "output_cost_per_token": 1e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_pdf_input": true, + "supports_vision": true + }, "gpt-realtime": { "cache_creation_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, @@ -5572,6 +5403,70 @@ "supports_system_messages": true, "supports_tool_choice": true }, + "gpt-realtime-1.5": { + "cache_creation_input_audio_token_cost": 4e-07, + "cache_read_input_token_cost": 4e-07, + "input_cost_per_audio_token": 3.2e-05, + "input_cost_per_image": 5e-06, + "input_cost_per_token": 4e-06, + "litellm_provider": "openai", + "max_input_tokens": 32000, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "chat", + "output_cost_per_audio_token": 6.4e-05, + "output_cost_per_token": 1.6e-05, + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "text", + "image", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "gpt-realtime-2": { + "cache_creation_input_audio_token_cost": 4e-07, + "cache_read_input_token_cost": 4e-07, + "input_cost_per_audio_token": 3.2e-05, + "input_cost_per_image": 5e-06, + "input_cost_per_token": 4e-06, + "litellm_provider": "openai", + "max_input_tokens": 32000, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "chat", + "output_cost_per_audio_token": 6.4e-05, + "output_cost_per_token": 1.6e-05, + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "text", + "image", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, "gpt-realtime-2025-08-28": { "cache_creation_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, @@ -5720,62 +5615,6 @@ "supports_tool_choice": true, "supports_vision": true }, - "o1-mini": { - "cache_read_input_token_cost": 5.5e-07, - "input_cost_per_token": 1.1e-06, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 65536, - "max_tokens": 65536, - "mode": "chat", - "output_cost_per_token": 4.4e-06, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_vision": true - }, - "o1-mini-2024-09-12": { - "cache_read_input_token_cost": 1.5e-06, - "deprecation_date": "2025-10-27", - "input_cost_per_token": 3e-06, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 65536, - "max_tokens": 65536, - "mode": "chat", - "output_cost_per_token": 1.2e-05, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_vision": true - }, - "o1-preview": { - "cache_read_input_token_cost": 7.5e-06, - "input_cost_per_token": 1.5e-05, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 32768, - "max_tokens": 32768, - "mode": "chat", - "output_cost_per_token": 6e-05, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_vision": true - }, - "o1-preview-2024-09-12": { - "cache_read_input_token_cost": 7.5e-06, - "input_cost_per_token": 1.5e-05, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 32768, - "max_tokens": 32768, - "mode": "chat", - "output_cost_per_token": 6e-05, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_vision": true - }, "o1-pro": { "input_cost_per_token": 0.00015, "input_cost_per_token_batches": 7.5e-05, @@ -5876,7 +5715,8 @@ "supports_response_schema": true, "supports_service_tier": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "o3-2025-04-16": { "cache_read_input_token_cost": 5e-07, @@ -5908,7 +5748,8 @@ "supports_response_schema": true, "supports_service_tier": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "o3-deep-research": { "cache_read_input_token_cost": 2.5e-06, @@ -5941,7 +5782,8 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "o3-deep-research-2025-06-26": { "cache_read_input_token_cost": 2.5e-06, @@ -5974,7 +5816,8 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "o3-mini": { "cache_read_input_token_cost": 5.5e-07, @@ -6038,7 +5881,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "o3-pro-2025-06-10": { "input_cost_per_token": 2e-05, @@ -6068,7 +5912,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "o4-mini": { "cache_read_input_token_cost": 2.75e-07, @@ -6093,7 +5938,8 @@ "supports_response_schema": true, "supports_service_tier": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "o4-mini-2025-04-16": { "cache_read_input_token_cost": 2.75e-07, @@ -6112,7 +5958,8 @@ "supports_response_schema": true, "supports_service_tier": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "o4-mini-deep-research": { "cache_read_input_token_cost": 5e-07, @@ -6145,7 +5992,8 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "o4-mini-deep-research-2025-06-26": { "cache_read_input_token_cost": 5e-07, @@ -6178,6 +6026,7 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true } } From 7321e4dea807651dd6fc309eee3d057ba75e7e70 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" <41898282+github-actions[bot]@users.noreply.github.com> Date: Fri, 29 May 2026 08:50:18 +0000 Subject: [PATCH 09/36] chore: sync VERSION to 0.1.133 [skip ci] --- backend/cmd/server/VERSION | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/backend/cmd/server/VERSION b/backend/cmd/server/VERSION index 7b9dfc4dbf97..56ebc9e5b0e6 100644 --- a/backend/cmd/server/VERSION +++ b/backend/cmd/server/VERSION @@ -1 +1 @@ -0.1.132 +0.1.133 From 06fca662735302def3b78bcb47f1650dd09350d8 Mon Sep 17 00:00:00 2001 From: DaydreamCoding <22166516+DaydreamCoding@users.noreply.github.com> Date: Fri, 29 May 2026 09:39:02 +0800 Subject: [PATCH 10/36] =?UTF-8?q?feat(quota):=20sentinel=20=E5=9B=9E?= =?UTF-8?q?=E5=A1=AB=E6=B6=88=E9=99=A4=E6=97=A0=E9=85=8D=E9=A2=9D=E8=A1=8C?= =?UTF-8?q?=E7=94=A8=E6=88=B7=20preflight=20=E6=AF=8F=E8=AF=B7=E6=B1=82?= =?UTF-8?q?=E5=9B=9E=E6=BA=90=20DB?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 无 user×platform 配额行的用户,preflight 每次 cache MISS 后回源 DB 查得"无行" 却不缓存该结论,导致每请求一次 DB 往返。本 PR 回填 sentinel 占位 entry,使后续 请求命中 Redis 后稳定判"无 limit",TTL 内不再查 DB。 - config: 加 UserPlatformQuotaSentinelTTLSeconds(默认 3600s,短于普通 quota cache 的 86400s 以控 Redis 内存) - metrics: 加 userPlatformQuotaSentinelSetCacheErrorTotal,并入 GatewayUserPlatformQuotaIncrStats 暴露 - billing_cache: checkUserPlatformQuotaEligibility 在 cache MISS + DB 无行且 cacheErr==nil 时回填 sentinel(三 limit nil、三 window_start non-nil、SchemaV1); TTL<=0 fallback 1h 防 EXPIRE 立即删 key 击穿;SET 失败 fail-open + 计 metric - billing_cache: HIT 路径对 sentinel(三 limit nil)跳过 windowExpired refresh, 避免短 sentinel TTL 被误升级为 86400s 有配额 limit 的用户 enforcement 行为不变(rec!=nil 不回填、isSentinel=false 不跳过 refresh)。 测试:扩展 fakeFullCache 夹具(setCalls/lastSetTTL/getErr/setErr);新增回填正确性 / Redis-GET-故障不回填 / SET-失败 fail-open / sentinel 跨窗口不 refresh 四个单测。 go build、quota+billing unit、三态 go vet 全绿。 Co-Authored-By: Claude Opus 4.8 (1M context) --- backend/internal/config/config.go | 4 + .../internal/service/billing_cache_service.go | 34 ++++- ..._cache_service_user_platform_quota_test.go | 122 +++++++++++++++++- backend/internal/service/gateway_service.go | 13 +- 4 files changed, 167 insertions(+), 6 deletions(-) diff --git a/backend/internal/config/config.go b/backend/internal/config/config.go index f689c2a908e6..dcbf30b460d0 100644 --- a/backend/internal/config/config.go +++ b/backend/internal/config/config.go @@ -649,6 +649,9 @@ type BillingConfig struct { // - billing_cache_service.checkUserPlatformQuotaEligibility 首次缓存装载 // 读写两端必须共用同一 TTL，避免缓存生命周期不一致导致 quota 计数漂移。 UserPlatformQuotaCacheTTLSeconds int `mapstructure:"user_platform_quota_cache_ttl_seconds"` + // UserPlatformQuotaSentinelTTLSeconds sentinel(无 limit 占位)entry 的 TTL, + // 显著短于 quota cache 默认 86400s 以控 Redis 内存;默认 3600=1h。 + UserPlatformQuotaSentinelTTLSeconds int `mapstructure:"user_platform_quota_sentinel_ttl_seconds"` } type CircuitBreakerConfig struct { @@ -1571,6 +1574,7 @@ func setDefaults() { viper.SetDefault("billing.circuit_breaker.reset_timeout_seconds", 30) viper.SetDefault("billing.circuit_breaker.half_open_requests", 3) viper.SetDefault("billing.user_platform_quota_cache_ttl_seconds", 86400) + viper.SetDefault("billing.user_platform_quota_sentinel_ttl_seconds", 3600) // Turnstile viper.SetDefault("turnstile.required", false) diff --git a/backend/internal/service/billing_cache_service.go b/backend/internal/service/billing_cache_service.go index 2b7c06baa287..8a5172f4d6c9 100644 --- a/backend/internal/service/billing_cache_service.go +++ b/backend/internal/service/billing_cache_service.go @@ -1096,7 +1096,12 @@ func (s *BillingCacheService) checkUserPlatformQuotaEligibility( // 超时 50ms:覆盖正常路径与可接受抖动;Redis 异常时 hot path 不阻塞超过此值。 // 用 context.Background()+短超时,避免请求 ctx 取消导致刷新丢失。 // 显式 setCancel()(而非 defer):缩短 context 生命周期,避免 defer 延迟到函数返回。 - if windowExpired && s.cache != nil { + // isSentinel 判定「该 entry 无任何 limit」,涵盖两类,跨窗口命中时都跳过 refresh: + // 1) A3 回填的 sentinel(DB 无行,短 TTL):refresh 会把短 TTL 误升级为 86400s,有害; + // 2) DB 有行但三 limit 全未配置的用户(TTL 86400s):refresh 纯属无意义(TTL 升级本身无害)。 + // 两类的 enforcement(下方 limit!=nil 比较)都因 limit 全 nil 永远放行,跳过 refresh 均正确。 + isSentinel := entry.DailyLimitUSD == nil && entry.WeeklyLimitUSD == nil && entry.MonthlyLimitUSD == nil + if windowExpired && s.cache != nil && !isSentinel { refreshed := &UserPlatformQuotaCacheEntry{ DailyUsageUSD: dailyUsage, WeeklyUsageUSD: weeklyUsage, @@ -1159,6 +1164,33 @@ func (s *BillingCacheService) checkUserPlatformQuotaEligibility( } rec, _ := v.(*UserPlatformQuotaRecord) if rec == nil { + // 仅在 cache 可用且本次 GET 未出错时回填 sentinel:Redis GET 故障(cacheErr!=nil) + // 时不回填,与下方 line ~1201 "Redis 故障时 fail-open:不回填" 保持一致, + // 避免在 Redis 异常期做一次注定失败的 SET。 + if s.cache != nil && cacheErr == nil { + now := time.Now() + startOfDay := timezone.StartOfDay(now) + startOfWeek := timezone.StartOfWeek(now) + sentinel := &UserPlatformQuotaCacheEntry{ + SchemaVersion: UserPlatformQuotaCacheSchemaV1, + DailyWindowStart: &startOfDay, + WeeklyWindowStart: &startOfWeek, + MonthlyWindowStart: &now, + // limits 全 nil, usage 全 0(零值) + } + sentinelTTL := time.Duration(s.cfg.Billing.UserPlatformQuotaSentinelTTLSeconds) * time.Second + if sentinelTTL <= 0 { + // 防御:TTL<=0 时 Redis EXPIRE 会立即删除整个 key(见 billing_cache.go 的 pipe.Expire), + // sentinel 不持久化 → 每请求击穿 DB。配置缺失/误配为 0 时 fallback 到 1h。 + sentinelTTL = time.Hour + } + setCtx, setCancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + if setErr := s.cache.SetUserPlatformQuotaCache(setCtx, userID, platform, sentinel, sentinelTTL); setErr != nil { + userPlatformQuotaSentinelSetCacheErrorTotal.Add(1) + logger.LegacyPrintf("service.billing_cache", "Warning: set sentinel quota cache failed user=%d platform=%s: %v", userID, platform, setErr) + } + setCancel() + } return nil } diff --git a/backend/internal/service/billing_cache_service_user_platform_quota_test.go b/backend/internal/service/billing_cache_service_user_platform_quota_test.go index 57697ddb05b8..674aa9a5fa85 100644 --- a/backend/internal/service/billing_cache_service_user_platform_quota_test.go +++ b/backend/internal/service/billing_cache_service_user_platform_quota_test.go @@ -95,6 +95,10 @@ type fakeFullCache struct { mu sync.Mutex entry *UserPlatformQuotaCacheEntry deleteCalls int + setCalls int // SetUserPlatformQuotaCache 调用次数 + lastSetTTL time.Duration // 最近一次 Set 的 ttl + getErr error // 非 nil 时 Get 先返回 (nil,false,getErr) + setErr error // 非 nil 时 Set 返回该 err(setCalls 仍+1) } // getDeleteCalls 线程安全地读取 deleteCalls。 @@ -111,19 +115,41 @@ func (f *fakeFullCache) getEntry() *UserPlatformQuotaCacheEntry { return f.entry } +// getSetCalls 线程安全地读取 setCalls。 +func (f *fakeFullCache) getSetCalls() int { + f.mu.Lock() + defer f.mu.Unlock() + return f.setCalls +} + +// getLastSetTTL 线程安全地读取 lastSetTTL。 +func (f *fakeFullCache) getLastSetTTL() time.Duration { + f.mu.Lock() + defer f.mu.Unlock() + return f.lastSetTTL +} + func (f *fakeFullCache) GetUserPlatformQuotaCache(_ context.Context, _ int64, _ string) (*UserPlatformQuotaCacheEntry, bool, error) { f.mu.Lock() defer f.mu.Unlock() + if f.getErr != nil { + return nil, false, f.getErr + } if f.entry == nil { return nil, false, nil } return f.entry, true, nil } -func (f *fakeFullCache) SetUserPlatformQuotaCache(_ context.Context, _ int64, _ string, e *UserPlatformQuotaCacheEntry, _ time.Duration) error { +func (f *fakeFullCache) SetUserPlatformQuotaCache(_ context.Context, _ int64, _ string, e *UserPlatformQuotaCacheEntry, ttl time.Duration) error { f.mu.Lock() defer f.mu.Unlock() + f.setCalls++ + if f.setErr != nil { + return f.setErr + } f.entry = e + f.lastSetTTL = ttl return nil } @@ -593,3 +619,97 @@ func TestMonthlyQuotaWindowExpired_BoundaryTable(t *testing.T) { }) } } + +// TestCheckUserPlatformQuotaEligibility_NoRow_WritesSentinel 验证: +// cache MISS + DB 无行时,回填 sentinel entry(三 limit 全 nil,三 window_start 全 non-nil), +// TTL = UserPlatformQuotaSentinelTTLSeconds,函数返回 nil(fail-open)。 +func TestCheckUserPlatformQuotaEligibility_NoRow_WritesSentinel(t *testing.T) { + repo := &fakeQuotaRepo{rec: nil} // DB 无行 + cache := &fakeFullCache{} // entry=nil → Get 返回 MISS + svc := newServiceForPreflight(t, repo, cache) + svc.cfg.Billing.UserPlatformQuotaSentinelTTLSeconds = 3600 + + if err := svc.checkUserPlatformQuotaEligibility(context.Background(), 1, "anthropic"); err != nil { + t.Fatalf("expected nil (fail-open), got %v", err) + } + if cache.getSetCalls() != 1 { + t.Fatalf("expected 1 SetUserPlatformQuotaCache call for sentinel, got %d", cache.getSetCalls()) + } + sentinel := cache.getEntry() + if sentinel == nil { + t.Fatal("expected sentinel entry backfilled") + } + if sentinel.DailyLimitUSD != nil || sentinel.WeeklyLimitUSD != nil || sentinel.MonthlyLimitUSD != nil { + t.Errorf("sentinel must have all-nil limits") + } + if sentinel.DailyWindowStart == nil || sentinel.WeeklyWindowStart == nil || sentinel.MonthlyWindowStart == nil { + t.Errorf("sentinel must have non-nil window_start to avoid refresh churn") + } + if sentinel.SchemaVersion != UserPlatformQuotaCacheSchemaV1 { + t.Errorf("sentinel schema = %d, want V1", sentinel.SchemaVersion) + } + if cache.getLastSetTTL() != 3600*time.Second { + t.Errorf("sentinel ttl = %v, want 3600s", cache.getLastSetTTL()) + } +} + +// TestCheckUserPlatformQuotaEligibility_RedisGetError_NoSentinelBackfill 验证: +// Redis GET 故障(cacheErr!=nil)+ DB 无行时,不应回填 sentinel(与 "Redis 故障时不回填" 一致),且 fail-open。 +func TestCheckUserPlatformQuotaEligibility_RedisGetError_NoSentinelBackfill(t *testing.T) { + repo := &fakeQuotaRepo{rec: nil} + cache := &fakeFullCache{getErr: errors.New("redis get down")} + svc := newServiceForPreflight(t, repo, cache) + svc.cfg.Billing.UserPlatformQuotaSentinelTTLSeconds = 3600 + + if err := svc.checkUserPlatformQuotaEligibility(context.Background(), 1, "anthropic"); err != nil { + t.Fatalf("redis 故障应 fail-open, got %v", err) + } + if cache.getSetCalls() != 0 { + t.Errorf("redis-get-error 时不应回填 sentinel, got %d set calls", cache.getSetCalls()) + } +} + +// TestCheckUserPlatformQuotaEligibility_NoRow_SentinelSetFailsFailOpen 验证: +// sentinel SET 失败时 fail-open(返回 nil)且计 metric。 +func TestCheckUserPlatformQuotaEligibility_NoRow_SentinelSetFailsFailOpen(t *testing.T) { + before := userPlatformQuotaSentinelSetCacheErrorTotal.Load() + repo := &fakeQuotaRepo{rec: nil} + cache := &fakeFullCache{setErr: errors.New("redis set timeout")} + svc := newServiceForPreflight(t, repo, cache) + svc.cfg.Billing.UserPlatformQuotaSentinelTTLSeconds = 3600 + + if err := svc.checkUserPlatformQuotaEligibility(context.Background(), 1, "anthropic"); err != nil { + t.Fatalf("sentinel set 失败应 fail-open, got %v", err) + } + if cache.getSetCalls() != 1 { + t.Errorf("应尝试 set sentinel 恰好一次, got %d", cache.getSetCalls()) + } + if got := userPlatformQuotaSentinelSetCacheErrorTotal.Load() - before; got != 1 { + t.Errorf("set 失败应使 metric +1, got delta %d", got) + } +} + +// TestCheckUserPlatformQuotaEligibility_SentinelCrossDay_NoRefresh 验证: +// 命中 sentinel(三 limit 全 nil)且跨窗口(daily/weekly 过期)时,不应触发 refresh SetCache +// (否则会把短 sentinel TTL 误升级为 quota cache 默认 86400s)。 +func TestCheckUserPlatformQuotaEligibility_SentinelCrossDay_NoRefresh(t *testing.T) { + yesterday := timezone.StartOfDay(time.Now().AddDate(0, 0, -1)) + lastWeek := timezone.StartOfWeek(time.Now().AddDate(0, 0, -7)) + monthAgoOK := time.Now().AddDate(0, 0, -5) // <30d, monthly 不过期 + sentinel := &UserPlatformQuotaCacheEntry{ + SchemaVersion: UserPlatformQuotaCacheSchemaV1, + DailyWindowStart: &yesterday, // 跨日 → daily windowExpired = true + WeeklyWindowStart: &lastWeek, // 跨周 → weekly windowExpired = true + MonthlyWindowStart: &monthAgoOK, + // limits 全 nil → sentinel + } + cache := &fakeFullCache{entry: sentinel} // entry 非 nil → Get HIT + svc := newServiceForPreflight(t, &fakeQuotaRepo{}, cache) + + if err := svc.checkUserPlatformQuotaEligibility(context.Background(), 1, "anthropic"); err != nil { + t.Fatalf("sentinel = no limit, expected nil, got %v", err) + } + if cache.getSetCalls() != 0 { + t.Errorf("sentinel cross-window must NOT trigger refresh SetCache, got %d calls", cache.getSetCalls()) + } +} diff --git a/backend/internal/service/gateway_service.go b/backend/internal/service/gateway_service.go index 94197f37c83c..effa803aa755 100644 --- a/backend/internal/service/gateway_service.go +++ b/backend/internal/service/gateway_service.go @@ -105,6 +105,9 @@ var ( // （applyUsageBilling 在 repo==nil 时 fallback）路径下的失败次数； // 与 DB Incr 失败分开计数，便于区分"主路径暂时故障"vs"基础设施长期未配齐"。 userPlatformQuotaDBIncrLegacyErrorTotal atomic.Int64 + // userPlatformQuotaSentinelSetCacheErrorTotal 统计 checkUserPlatformQuotaEligibility + // 在 DB 无行时回填 sentinel cache entry 写 Redis 失败的次数（phase A）。 + userPlatformQuotaSentinelSetCacheErrorTotal atomic.Int64 ) func GatewayWindowCostPrefetchStats() (cacheHit, cacheMiss, batchSQL, fallback, errCount int64) { @@ -127,13 +130,15 @@ func GatewayModelsListCacheStats() (cacheHit, cacheMiss, store int64) { return modelsListCacheHitTotal.Load(), modelsListCacheMissTotal.Load(), modelsListCacheStoreTotal.Load() } -// GatewayUserPlatformQuotaIncrStats 返回 (mainPathErr, legacyPathErr)。 +// GatewayUserPlatformQuotaIncrStats 返回 (mainPathErr, legacyPathErr, sentinelSetErr)。 // mainPathErr：finalizePostUsageBilling 异步 goroutine 写 DB 失败累计次数； -// legacyPathErr：postUsageBilling fallback 路径写 DB 失败累计次数。 +// legacyPathErr：postUsageBilling fallback 路径写 DB 失败累计次数； +// sentinelSetErr：DB 无行时回填 sentinel cache entry 写 Redis 失败累计次数。 // ops 监控面板可以按"持续上升斜率"做告警阈值。 -func GatewayUserPlatformQuotaIncrStats() (mainPathErr, legacyPathErr int64) { +func GatewayUserPlatformQuotaIncrStats() (mainPathErr, legacyPathErr, sentinelSetErr int64) { return userPlatformQuotaDBIncrErrorTotal.Load(), - userPlatformQuotaDBIncrLegacyErrorTotal.Load() + userPlatformQuotaDBIncrLegacyErrorTotal.Load(), + userPlatformQuotaSentinelSetCacheErrorTotal.Load() } func openAIStreamEventIsTerminal(data string) bool { From f7f5e33830bcf4d8998544cfb1d726c4c436af1f Mon Sep 17 00:00:00 2001 From: DaydreamCoding <22166516+DaydreamCoding@users.noreply.github.com> Date: Fri, 29 May 2026 13:11:36 +0800 Subject: [PATCH 11/36] =?UTF-8?q?feat(quota):=20user=C3=97platform=20?= =?UTF-8?q?=E9=85=8D=E9=A2=9D=20DB=20=E5=86=99=E8=81=9A=E5=90=88=20flusher?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Redis 同步权威 + DB 镜像,不在进程内维护 delta: - 写入点 HasUserPlatformQuotaLimit 守卫:无 limit 跳过 Redis 写与持久化 - 累加 usage 的 Lua 在 flusher_enabled 时 SADD 脏集 billing:upq:dirty - UserPlatformQuotaUsageFlusher 定时 SPOP 脏集 → 批量 HGETALL 读当前窗口 usage 快照 → BatchSnapshotUsage 绝对值 UPSERT 覆盖 DB(去 SELECT FOR UPDATE 行锁) → 失败 SADD 回 / FK(23503)整批丢弃 - flusher 单批 clamp 到 ≤6000,保证一次 flush 只生成一条 UPSERT(单事务原子) - flusher_enabled 默认 false(降级=旧异步直写 DB) 效果:DB 写连接从 O(QPS) 收敛到 O(副本)。 循环依赖:service 层独立 Snapshot/FK 类型,repository adapter 转换 + %w 映射 FK error。 admin reset/upsert 后失效 cache(脏残留被 flusher 当 MISS 跳过)。 健壮性与可观测性: - flusher_enabled=false 时 Start 不注册定时器;flush_interval_ms 非法回退 2s - Readd 回填失败单独计 dirty_lost(不再误记 dirty_readd)并 ALERT;脏集 Readd 补兜底 TTL - 单 tick 达 max batches 上限仍有积压时记 log - admin 失效 cache 失败升级为 ALERT(提示 enforcement 可能延迟至 sentinel TTL) - BatchGet 单条命令失败 / usage 字段损坏均记 log,避免静默以 0 覆写 DB 三态 go vet + 单测/集成测全绿。 已知取舍(默认 flusher_enabled=false 不触发): - FK 整批丢弃牵连同批正常 key(活跃 key 靠下次 SADD+绝对值快照自愈;Redis 仍权威) - admin reset/upsert 直写 DB 与 flusher 异步刷存在覆盖竞态:flusher 持旧快照在途时可能覆盖 admin 刚写值(limit 列不受影响;usage 有 preflight windowExpired 兜底;低频)。彻底消除需 version OCC。 Co-Authored-By: Claude Opus 4.8 --- backend/cmd/server/wire.go | 7 + backend/cmd/server/wire_gen.go | 10 +- backend/cmd/server/wire_gen_test.go | 1 + backend/internal/config/config.go | 10 + .../internal/handler/admin/user_handler.go | 4 +- backend/internal/repository/billing_cache.go | 210 +++++-- .../billing_cache_user_platform_quota_test.go | 6 +- .../user_platform_quota_adapter_test.go | 3 + .../repository/user_platform_quota_repo.go | 92 ++++ ...er_platform_quota_repo_integration_test.go | 98 ++++ .../user_platform_quota_service_adapter.go | 46 ++ .../service/admin_service_delete_test.go | 14 +- .../auth_service_platform_quota_test.go | 4 + .../service/auth_service_register_test.go | 4 + .../internal/service/billing_cache_service.go | 20 +- ...billing_cache_service_singleflight_test.go | 14 +- .../service/billing_cache_service_test.go | 14 +- ..._cache_service_user_platform_quota_test.go | 135 ++++- backend/internal/service/billing_service.go | 14 +- backend/internal/service/gateway_service.go | 96 ++-- .../service/user_platform_quota_flusher.go | 267 +++++++++ .../user_platform_quota_flusher_test.go | 511 ++++++++++++++++++ .../service/user_platform_quota_port.go | 19 + backend/internal/service/user_service_test.go | 14 +- backend/internal/service/wire.go | 8 + 25 files changed, 1515 insertions(+), 106 deletions(-) create mode 100644 backend/internal/service/user_platform_quota_flusher.go create mode 100644 backend/internal/service/user_platform_quota_flusher_test.go diff --git a/backend/cmd/server/wire.go b/backend/cmd/server/wire.go index 9bfa27174db6..b474cfa1aad5 100644 --- a/backend/cmd/server/wire.go +++ b/backend/cmd/server/wire.go @@ -98,6 +98,7 @@ func provideCleanup( backupSvc *service.BackupService, paymentOrderExpiry *service.PaymentOrderExpiryService, channelMonitorRunner *service.ChannelMonitorRunner, + quotaFlusher *service.UserPlatformQuotaUsageFlusher, ) func() { return func() { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) @@ -246,6 +247,12 @@ func provideCleanup( } return nil }}, + {"UserPlatformQuotaUsageFlusher", func() error { + if quotaFlusher != nil { + quotaFlusher.Stop() + } + return nil + }}, } infraSteps := []cleanupStep{ diff --git a/backend/cmd/server/wire_gen.go b/backend/cmd/server/wire_gen.go index 6e8be8fce9a1..9b059c5d823a 100644 --- a/backend/cmd/server/wire_gen.go +++ b/backend/cmd/server/wire_gen.go @@ -269,7 +269,8 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { scheduledTestRunnerService := service.ProvideScheduledTestRunnerService(scheduledTestPlanRepository, scheduledTestService, accountTestService, rateLimitService, configConfig) paymentOrderExpiryService := service.ProvidePaymentOrderExpiryService(paymentService) channelMonitorRunner := service.ProvideChannelMonitorRunner(channelMonitorService, settingService) - v := provideCleanup(client, redisClient, opsMetricsCollector, opsAggregationService, opsAlertEvaluatorService, opsCleanupService, opsScheduledReportService, opsSystemLogSink, schedulerSnapshotService, tokenRefreshService, accountExpiryService, subscriptionExpiryService, usageCleanupService, idempotencyCleanupService, pricingService, emailQueueService, billingCacheService, usageRecordWorkerPool, subscriptionService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, openAIGatewayService, scheduledTestRunnerService, backupService, paymentOrderExpiryService, channelMonitorRunner) + userPlatformQuotaUsageFlusher := service.ProvideUserPlatformQuotaUsageFlusher(configConfig, billingCache, serviceUserPlatformQuotaRepository, timingWheelService) + v := provideCleanup(client, redisClient, opsMetricsCollector, opsAggregationService, opsAlertEvaluatorService, opsCleanupService, opsScheduledReportService, opsSystemLogSink, schedulerSnapshotService, tokenRefreshService, accountExpiryService, subscriptionExpiryService, usageCleanupService, idempotencyCleanupService, pricingService, emailQueueService, billingCacheService, usageRecordWorkerPool, subscriptionService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, openAIGatewayService, scheduledTestRunnerService, backupService, paymentOrderExpiryService, channelMonitorRunner, userPlatformQuotaUsageFlusher) application := &Application{ Server: httpServer, Cleanup: v, @@ -324,6 +325,7 @@ func provideCleanup( backupSvc *service.BackupService, paymentOrderExpiry *service.PaymentOrderExpiryService, channelMonitorRunner *service.ChannelMonitorRunner, + quotaFlusher *service.UserPlatformQuotaUsageFlusher, ) func() { return func() { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) @@ -471,6 +473,12 @@ func provideCleanup( } return nil }}, + {"UserPlatformQuotaUsageFlusher", func() error { + if quotaFlusher != nil { + quotaFlusher.Stop() + } + return nil + }}, } infraSteps := []cleanupStep{ diff --git a/backend/cmd/server/wire_gen_test.go b/backend/cmd/server/wire_gen_test.go index a44b2d5caacd..7f4e47731765 100644 --- a/backend/cmd/server/wire_gen_test.go +++ b/backend/cmd/server/wire_gen_test.go @@ -77,6 +77,7 @@ func TestProvideCleanup_WithMinimalDependencies_NoPanic(t *testing.T) { nil, // backupSvc nil, // paymentOrderExpiry nil, // channelMonitorRunner + nil, // quotaFlusher ) require.NotPanics(t, func() { diff --git a/backend/internal/config/config.go b/backend/internal/config/config.go index dcbf30b460d0..df9dcefc6017 100644 --- a/backend/internal/config/config.go +++ b/backend/internal/config/config.go @@ -1094,6 +1094,13 @@ type DatabaseConfig struct { ConnMaxLifetimeMinutes int `mapstructure:"conn_max_lifetime_minutes"` // ConnMaxIdleTimeMinutes: 空闲连接最大存活时间，及时释放不活跃连接 ConnMaxIdleTimeMinutes int `mapstructure:"conn_max_idle_time_minutes"` + // UserPlatformQuotaFlusherEnabled: 是否启用 user×platform 配额写聚合 flusher + UserPlatformQuotaFlusherEnabled bool `mapstructure:"user_platform_quota_flusher_enabled"` + // UserPlatformQuotaFlushIntervalMs: flusher 刷写间隔（毫秒） + UserPlatformQuotaFlushIntervalMs int `mapstructure:"user_platform_quota_flush_interval_ms"` + // UserPlatformQuotaFlushBatchSize: flusher 单批最大条数 + // 建议 ≤ 6000（单条 UPSERT 原子上限） + UserPlatformQuotaFlushBatchSize int `mapstructure:"user_platform_quota_flush_batch_size"` } func (d *DatabaseConfig) DSN() string { @@ -1661,6 +1668,9 @@ func setDefaults() { viper.SetDefault("database.max_idle_conns", 128) viper.SetDefault("database.conn_max_lifetime_minutes", 30) viper.SetDefault("database.conn_max_idle_time_minutes", 5) + viper.SetDefault("database.user_platform_quota_flusher_enabled", false) + viper.SetDefault("database.user_platform_quota_flush_interval_ms", 2000) + viper.SetDefault("database.user_platform_quota_flush_batch_size", 1000) // Redis viper.SetDefault("redis.host", "localhost") diff --git a/backend/internal/handler/admin/user_handler.go b/backend/internal/handler/admin/user_handler.go index 32a216929e87..6c0a02ffb768 100644 --- a/backend/internal/handler/admin/user_handler.go +++ b/backend/internal/handler/admin/user_handler.go @@ -743,7 +743,7 @@ func (h *UserHandler) UpdateUserPlatformQuotas(c *gin.Context) { if h.billingCache != nil { for _, p := range service.AllowedQuotaPlatforms { if err := h.billingCache.DeleteUserPlatformQuotaCache(ctx, userID, p); err != nil { - slog.Warn("quota cache invalidation failed", "user_id", userID, "platform", p, "err", err) + slog.Error("ALERT: quota cache invalidation failed after UpsertForUser; limit 生效可能延迟至 sentinel TTL(最长 1h),需人工确认或重试失效", "user_id", userID, "platform", p, "err", err) } } } @@ -827,7 +827,7 @@ func (h *UserHandler) ResetUserPlatformQuotaWindow(c *gin.Context) { if h.billingCache != nil { if err := h.billingCache.DeleteUserPlatformQuotaCache(ctx, userID, req.Platform); err != nil { - slog.Warn("quota cache invalidation failed", "user_id", userID, "platform", req.Platform, "err", err) + slog.Error("ALERT: quota cache invalidation failed after ResetExpiredWindow; 窗口重置可能延迟至 sentinel TTL(最长 1h)", "user_id", userID, "platform", req.Platform, "err", err) } } diff --git a/backend/internal/repository/billing_cache.go b/backend/internal/repository/billing_cache.go index 60dae9540d34..de229da9c351 100644 --- a/backend/internal/repository/billing_cache.go +++ b/backend/internal/repository/billing_cache.go @@ -7,6 +7,7 @@ import ( "log" "math/rand/v2" "strconv" + "strings" "time" "github.com/Wei-Shaw/sub2api/internal/service" @@ -338,38 +339,26 @@ func userPlatformQuotaCacheKey(userID int64, platform string) string { return fmt.Sprintf("billing:user_platform_quota:%d:%s", userID, platform) } -func (c *billingCache) GetUserPlatformQuotaCache(ctx context.Context, userID int64, platform string) (*service.UserPlatformQuotaCacheEntry, bool, error) { - key := userPlatformQuotaCacheKey(userID, platform) - fields := []string{ - "daily_usage", "weekly_usage", "monthly_usage", "version", "schema_version", - "daily_limit", "weekly_limit", "monthly_limit", - "daily_window_start", "weekly_window_start", "monthly_window_start", - } - vals, err := c.rdb.HMGet(ctx, key, fields...).Result() - if err != nil { - return nil, false, err - } - // 前4个全为nil → key 不存在 - if vals[0] == nil && vals[1] == nil && vals[2] == nil && vals[3] == nil { - return nil, false, nil +// parseUserPlatformQuotaHash 将 Redis HGETALL 返回的 map[string]string 反序列化为 +// *service.UserPlatformQuotaCacheEntry。空 map（key 不存在）返回 nil。 +// GetUserPlatformQuotaCache 和 BatchGetUserPlatformQuotaCache 共用此函数，确保解析逻辑一致。 +func parseUserPlatformQuotaHash(m map[string]string) *service.UserPlatformQuotaCacheEntry { + if len(m) == 0 { + return nil } - parseFloat := func(v any) float64 { - if v == nil { + parseFloat := func(s string) float64 { + if s == "" { return 0 } - s, ok := v.(string) - if !ok { + f, err := strconv.ParseFloat(s, 64) + if err != nil { + log.Printf("billing_cache: corrupt quota usage field %q (using 0): %v", s, err) return 0 } - f, _ := strconv.ParseFloat(s, 64) return f } - parseFloatPtr := func(v any) *float64 { - if v == nil { - return nil - } - s, ok := v.(string) - if !ok || s == "" { + parseFloatPtr := func(s string) *float64 { + if s == "" { return nil } f, err := strconv.ParseFloat(s, 64) @@ -378,12 +367,8 @@ func (c *billingCache) GetUserPlatformQuotaCache(ctx context.Context, userID int } return &f } - parseTimePtr := func(v any) *time.Time { - if v == nil { - return nil - } - s, ok := v.(string) - if !ok || s == "" { + parseTimePtr := func(s string) *time.Time { + if s == "" { return nil } n, err := strconv.ParseInt(s, 10, 64) @@ -393,30 +378,37 @@ func (c *billingCache) GetUserPlatformQuotaCache(ctx context.Context, userID int t := time.Unix(n, 0).UTC() return &t } - parseInt64 := func(v any) int64 { - if v == nil { - return 0 - } - s, ok := v.(string) - if !ok { - return 0 - } + parseInt64 := func(s string) int64 { n, _ := strconv.ParseInt(s, 10, 64) return n } return &service.UserPlatformQuotaCacheEntry{ - DailyUsageUSD: parseFloat(vals[0]), - WeeklyUsageUSD: parseFloat(vals[1]), - MonthlyUsageUSD: parseFloat(vals[2]), - Version: parseInt64(vals[3]), - SchemaVersion: parseInt64(vals[4]), - DailyLimitUSD: parseFloatPtr(vals[5]), - WeeklyLimitUSD: parseFloatPtr(vals[6]), - MonthlyLimitUSD: parseFloatPtr(vals[7]), - DailyWindowStart: parseTimePtr(vals[8]), - WeeklyWindowStart: parseTimePtr(vals[9]), - MonthlyWindowStart: parseTimePtr(vals[10]), - }, true, nil + DailyUsageUSD: parseFloat(m["daily_usage"]), + WeeklyUsageUSD: parseFloat(m["weekly_usage"]), + MonthlyUsageUSD: parseFloat(m["monthly_usage"]), + Version: parseInt64(m["version"]), + SchemaVersion: parseInt64(m["schema_version"]), + DailyLimitUSD: parseFloatPtr(m["daily_limit"]), + WeeklyLimitUSD: parseFloatPtr(m["weekly_limit"]), + MonthlyLimitUSD: parseFloatPtr(m["monthly_limit"]), + DailyWindowStart: parseTimePtr(m["daily_window_start"]), + WeeklyWindowStart: parseTimePtr(m["weekly_window_start"]), + MonthlyWindowStart: parseTimePtr(m["monthly_window_start"]), + } +} + +func (c *billingCache) GetUserPlatformQuotaCache(ctx context.Context, userID int64, platform string) (*service.UserPlatformQuotaCacheEntry, bool, error) { + key := userPlatformQuotaCacheKey(userID, platform) + m, err := c.rdb.HGetAll(ctx, key).Result() + if err != nil { + return nil, false, err + } + entry := parseUserPlatformQuotaHash(m) + if entry == nil { + // 空 map → key 不存在 → MISS + return nil, false, nil + } + return entry, true, nil } func (c *billingCache) SetUserPlatformQuotaCache(ctx context.Context, userID int64, platform string, entry *service.UserPlatformQuotaCacheEntry, ttl time.Duration) error { @@ -468,9 +460,12 @@ func (c *billingCache) DeleteUserPlatformQuotaCache(ctx context.Context, userID // SetCache 重建为新版 entry —— 若此处仍累加，上层覆盖时会丢失这部分增量，导致 Redis usage 比真实偏小。 // key 不存在同样跳过（由下次 SetCache 重建）。 // KEYS[1] = hash key +// KEYS[2] = 脏集 key（dirty set） // ARGV[1] = cost (string float) // ARGV[2] = ttl seconds // ARGV[3] = expected schema_version (Go 侧 UserPlatformQuotaCacheSchemaV1) +// ARGV[4] = dirty set member（空串则不 SADD） +// ARGV[5] = 脏集兜底 TTL 秒 const updateUserPlatformQuotaUsageScript = ` if redis.call("EXISTS", KEYS[1]) == 0 then return 0 @@ -484,18 +479,125 @@ redis.call("HINCRBYFLOAT", KEYS[1], "weekly_usage", ARGV[1]) redis.call("HINCRBYFLOAT", KEYS[1], "monthly_usage", ARGV[1]) redis.call("HINCRBY", KEYS[1], "version", 1) redis.call("EXPIRE", KEYS[1], ARGV[2]) +if ARGV[4] ~= "" then + redis.call("SADD", KEYS[2], ARGV[4]) + redis.call("EXPIRE", KEYS[2], ARGV[5]) +end return 1 ` -func (c *billingCache) IncrUserPlatformQuotaUsageCache(ctx context.Context, userID int64, platform string, cost float64, ttl time.Duration) error { - key := userPlatformQuotaCacheKey(userID, platform) - _, err := c.rdb.Eval(ctx, updateUserPlatformQuotaUsageScript, []string{key}, +// userPlatformQuotaDirtySetKey 返回脏集（dirty set）的 Redis key。 +// 使用与 userPlatformQuotaCacheKey 相同的前缀 "billing:"。 +func userPlatformQuotaDirtySetKey() string { return "billing:" + "upq:dirty" } + +// userPlatformQuotaDirtyTTLSeconds 脏集兜底 TTL（秒）：初始 SADD（Lua）与 Readd 共用， +// 确保 flusher 长期停摆时脏集最终过期；正常运行因持续 SADD 不断续期。 +const userPlatformQuotaDirtyTTLSeconds = 86400 + +// userPlatformQuotaDirtyMember 构造脏集成员字符串 "userID:platform"。 +func userPlatformQuotaDirtyMember(userID int64, platform string) string { + return strconv.FormatInt(userID, 10) + ":" + platform +} + +func (c *billingCache) IncrUserPlatformQuotaUsageCache(ctx context.Context, userID int64, platform string, cost float64, ttl time.Duration, markDirty bool) error { + member := "" + if markDirty { + member = userPlatformQuotaDirtyMember(userID, platform) + } + _, err := c.rdb.Eval(ctx, updateUserPlatformQuotaUsageScript, + []string{userPlatformQuotaCacheKey(userID, platform), userPlatformQuotaDirtySetKey()}, strconv.FormatFloat(cost, 'f', -1, 64), int(ttl.Seconds()), service.UserPlatformQuotaCacheSchemaV1, + member, + userPlatformQuotaDirtyTTLSeconds, ).Result() if err != nil && !errors.Is(err, redis.Nil) { return err } return nil } + +// parseUserPlatformQuotaDirtyMember 将脏集成员字符串 "userID:platform" 解析为 +// service.UserPlatformQuotaKey。解析失败返回 ok=false。 +func parseUserPlatformQuotaDirtyMember(m string) (service.UserPlatformQuotaKey, bool) { + parts := strings.SplitN(m, ":", 2) + if len(parts) != 2 { + return service.UserPlatformQuotaKey{}, false + } + uid, err := strconv.ParseInt(parts[0], 10, 64) + if err != nil { + return service.UserPlatformQuotaKey{}, false + } + return service.UserPlatformQuotaKey{UserID: uid, Platform: parts[1]}, true +} + +// PopDirtyUserPlatformQuotaKeys 从脏集随机弹出最多 n 个 key。 +// 脏集为空时返回 (nil, nil)。 +func (c *billingCache) PopDirtyUserPlatformQuotaKeys(ctx context.Context, n int) ([]service.UserPlatformQuotaKey, error) { + members, err := c.rdb.SPopN(ctx, userPlatformQuotaDirtySetKey(), int64(n)).Result() + if err != nil { + if errors.Is(err, redis.Nil) { + return nil, nil + } + return nil, err + } + keys := make([]service.UserPlatformQuotaKey, 0, len(members)) + for _, m := range members { + k, ok := parseUserPlatformQuotaDirtyMember(m) + if !ok { + log.Printf("billing_cache: skipping invalid dirty member %q", m) + continue + } + keys = append(keys, k) + } + return keys, nil +} + +// ReaddDirtyUserPlatformQuotaKeys 将 keys 重新加入脏集（flush 失败时回填）。 +// 通过 pipeline 同时执行 SAdd + Expire，确保 Readd 后脏集具有兜底 TTL。 +// 空切片时直接返回 nil。 +func (c *billingCache) ReaddDirtyUserPlatformQuotaKeys(ctx context.Context, keys []service.UserPlatformQuotaKey) error { + if len(keys) == 0 { + return nil + } + dirtyKey := userPlatformQuotaDirtySetKey() + members := make([]any, len(keys)) + for i, k := range keys { + members[i] = userPlatformQuotaDirtyMember(k.UserID, k.Platform) + } + pipe := c.rdb.Pipeline() + pipe.SAdd(ctx, dirtyKey, members...) + pipe.Expire(ctx, dirtyKey, userPlatformQuotaDirtyTTLSeconds*time.Second) + _, err := pipe.Exec(ctx) + return err +} + +// BatchGetUserPlatformQuotaCache 通过 Pipeline 批量 HGETALL 获取多个 user×platform 的 +// quota cache。返回切片与 keys 顺序、长度对齐；MISS 或解析失败位置返回 nil。 +func (c *billingCache) BatchGetUserPlatformQuotaCache(ctx context.Context, keys []service.UserPlatformQuotaKey) ([]*service.UserPlatformQuotaCacheEntry, error) { + if len(keys) == 0 { + return nil, nil + } + pipe := c.rdb.Pipeline() + cmds := make([]*redis.MapStringStringCmd, len(keys)) + for i, k := range keys { + cmds[i] = pipe.HGetAll(ctx, userPlatformQuotaCacheKey(k.UserID, k.Platform)) + } + if _, err := pipe.Exec(ctx); err != nil && !errors.Is(err, redis.Nil) { + return nil, err + } + results := make([]*service.UserPlatformQuotaCacheEntry, len(keys)) + for i, cmd := range cmds { + m, err := cmd.Result() + if err != nil { + if !errors.Is(err, redis.Nil) { + log.Printf("billing_cache: BatchGet HGETALL cmd[%d] failed: %v (skip, self-heal)", i, err) + } + // 单个命令失败 → 对应位置 nil，继续 + continue + } + results[i] = parseUserPlatformQuotaHash(m) + } + return results, nil +} diff --git a/backend/internal/repository/billing_cache_user_platform_quota_test.go b/backend/internal/repository/billing_cache_user_platform_quota_test.go index 8d49fd31cfde..15b185e77693 100644 --- a/backend/internal/repository/billing_cache_user_platform_quota_test.go +++ b/backend/internal/repository/billing_cache_user_platform_quota_test.go @@ -88,7 +88,7 @@ func TestUserPlatformQuotaCache_NilLimitSetThenGet(t *testing.T) { func TestUserPlatformQuotaCache_IncrMissIsNoop(t *testing.T) { c, _ := newMiniRedisCache(t) - if err := c.IncrUserPlatformQuotaUsageCache(context.Background(), 1, "openai", 0.5, time.Minute); err != nil { + if err := c.IncrUserPlatformQuotaUsageCache(context.Background(), 1, "openai", 0.5, time.Minute, false); err != nil { t.Fatal(err) } _, ok, _ := c.GetUserPlatformQuotaCache(context.Background(), 1, "openai") @@ -105,10 +105,10 @@ func TestUserPlatformQuotaCache_IncrHitAccumulates(t *testing.T) { Version: 1, SchemaVersion: service.UserPlatformQuotaCacheSchemaV1, }, time.Minute) - if err := c.IncrUserPlatformQuotaUsageCache(ctx, 1, "openai", 0.5, time.Minute); err != nil { + if err := c.IncrUserPlatformQuotaUsageCache(ctx, 1, "openai", 0.5, time.Minute, false); err != nil { t.Fatal(err) } - if err := c.IncrUserPlatformQuotaUsageCache(ctx, 1, "openai", 0.25, time.Minute); err != nil { + if err := c.IncrUserPlatformQuotaUsageCache(ctx, 1, "openai", 0.25, time.Minute, false); err != nil { t.Fatal(err) } got, _, _ := c.GetUserPlatformQuotaCache(ctx, 1, "openai") diff --git a/backend/internal/repository/user_platform_quota_adapter_test.go b/backend/internal/repository/user_platform_quota_adapter_test.go index a55d2e9c0330..f31defe5d55a 100644 --- a/backend/internal/repository/user_platform_quota_adapter_test.go +++ b/backend/internal/repository/user_platform_quota_adapter_test.go @@ -38,6 +38,9 @@ func (f *fakeRepoForAdapter) UpsertForUser(_ context.Context, userID int64, reco f.upsertCalledWith = records return f.upsertErr } +func (f *fakeRepoForAdapter) BatchSnapshotUsage(_ context.Context, _ []UserPlatformQuotaSnapshot, _ time.Time) error { + return nil +} func TestGenericAdapter_UpsertForUser_ForwardsRecords(t *testing.T) { fake := &fakeRepoForAdapter{} diff --git a/backend/internal/repository/user_platform_quota_repo.go b/backend/internal/repository/user_platform_quota_repo.go index 1e2e7f51bd48..ccba2330ced2 100644 --- a/backend/internal/repository/user_platform_quota_repo.go +++ b/backend/internal/repository/user_platform_quota_repo.go @@ -2,6 +2,7 @@ package repository import ( "context" + "errors" "fmt" "strings" "time" @@ -9,6 +10,7 @@ import ( dbent "github.com/Wei-Shaw/sub2api/ent" "github.com/Wei-Shaw/sub2api/ent/userplatformquota" "github.com/Wei-Shaw/sub2api/internal/pkg/timezone" + "github.com/lib/pq" ) // UserPlatformQuotaRecord 是 repository 层的传输结构体， @@ -30,6 +32,22 @@ type UserPlatformQuotaRecord struct { // ErrUserPlatformQuotaNotFound 用于 ResetExpiredWindow 等需要"必须命中已有记录"的方法。 var ErrUserPlatformQuotaNotFound = fmt.Errorf("user platform quota record not found") +// ErrUserPlatformQuotaFKViolation 当批量 UPSERT 中存在 user_id 不在 users 表的记录时返回。 +var ErrUserPlatformQuotaFKViolation = errors.New("user platform quota snapshot FK violation") + +// UserPlatformQuotaSnapshot 是 BatchSnapshotUsage 的输入结构体， +// 表示 Redis 当前窗口快照（用于绝对值覆盖写入 DB）。 +type UserPlatformQuotaSnapshot struct { + UserID int64 + Platform string + DailyUsageUSD float64 + WeeklyUsageUSD float64 + MonthlyUsageUSD float64 + DailyWindowStart time.Time + WeeklyWindowStart time.Time + MonthlyWindowStart time.Time +} + // UserPlatformQuotaRepository 定义用户平台配额的数据访问接口。 type UserPlatformQuotaRepository interface { // BulkInsertInitial 幂等批量插入初始配额记录（ON CONFLICT DO NOTHING）。 @@ -44,6 +62,10 @@ type UserPlatformQuotaRepository interface { ResetExpiredWindow(ctx context.Context, userID int64, platform string, window string, newStart time.Time) error // UpsertForUser 全量替换该用户所有平台限额配置（详见 service.UserPlatformQuotaRepository.UpsertForUser）。 UpsertForUser(ctx context.Context, userID int64, records []UserPlatformQuotaRecord) error + // BatchSnapshotUsage 用一条多行 UPSERT 把整批 usage 以绝对值覆盖写入(非累加)。 + // usage/window_start 直接取 EXCLUDED(Redis 当前窗口快照),无 CASE。整批共用 now 作 created/updated_at。 + // 要求 snapshots 内 (user,platform) 不重复。FK 违反返回 ErrUserPlatformQuotaFKViolation。 + BatchSnapshotUsage(ctx context.Context, snapshots []UserPlatformQuotaSnapshot, now time.Time) error } type userPlatformQuotaRepository struct { @@ -414,3 +436,73 @@ func insertLimitsRow(ctx context.Context, client *dbent.Client, userID int64, re } return nil } + +// batchRows 是 BatchSnapshotUsage 每批最大行数（9 参/行 × 6000 ≈ 54000 参,低于 Postgres 65535 上限）。 +const batchRows = 6000 + +// BatchSnapshotUsage 用一条多行 UPSERT 把整批 usage 以绝对值覆盖写入（非累加）。 +// 每批最多 batchRows 行；$1=now 共用；每行 8 个 per-row 参（user_id, platform, 3×usage, 3×window_start）。 +// FK 违反（user_id 不存在）返回 ErrUserPlatformQuotaFKViolation。 +// +// 注意:snapshots 超过 batchRows 会分多条 SQL 执行且【非单事务】——若某子批 FK 失败, +// 先前子批已写入无法回滚。调用方(flusher)应保证单次 batchSize ≤ batchRows +// (默认 flush_batch_size=1000 < 6000,安全)。 +// 另注:启用 flusher 后,本绝对值覆盖与 admin 直写 DB(ResetExpiredWindow/UpsertForUser)存在覆盖竞态, +// 详见 service/user_platform_quota_flusher.go 中 flushOneBatch 的"已知竞态"注释。 +func (r *userPlatformQuotaRepository) BatchSnapshotUsage(ctx context.Context, snapshots []UserPlatformQuotaSnapshot, now time.Time) error { + if len(snapshots) == 0 { + return nil + } + + client := clientFromContext(ctx, r.client) + + for start := 0; start < len(snapshots); start += batchRows { + end := start + batchRows + if end > len(snapshots) { + end = len(snapshots) + } + batch := snapshots[start:end] + + var sb strings.Builder + _, _ = sb.WriteString( + "INSERT INTO user_platform_quotas" + + " (user_id, platform, daily_usage_usd, weekly_usage_usd, monthly_usage_usd," + + " daily_window_start, weekly_window_start, monthly_window_start, created_at, updated_at)" + + " VALUES ") + + // $1 = now（共用）；每行 8 个 per-row 参，从 $2 起连续编号。 + args := []any{now} + for i, s := range batch { + if i > 0 { + _, _ = sb.WriteString(",") + } + b := len(args) // 当前 per-row 第一个参数的 0-based 索引，实际占位符 = b+1 + fmt.Fprintf(&sb, "($%d,$%d,$%d,$%d,$%d,$%d,$%d,$%d,$1,$1)", + b+1, b+2, b+3, b+4, b+5, b+6, b+7, b+8) + args = append(args, + s.UserID, s.Platform, + s.DailyUsageUSD, s.WeeklyUsageUSD, s.MonthlyUsageUSD, + s.DailyWindowStart, s.WeeklyWindowStart, s.MonthlyWindowStart, + ) + } + + _, _ = sb.WriteString( + " ON CONFLICT (user_id, platform) WHERE deleted_at IS NULL DO UPDATE SET" + + " daily_usage_usd = EXCLUDED.daily_usage_usd," + + " weekly_usage_usd = EXCLUDED.weekly_usage_usd," + + " monthly_usage_usd = EXCLUDED.monthly_usage_usd," + + " daily_window_start = EXCLUDED.daily_window_start," + + " weekly_window_start = EXCLUDED.weekly_window_start," + + " monthly_window_start = EXCLUDED.monthly_window_start," + + " updated_at = EXCLUDED.updated_at") + + if _, err := client.ExecContext(ctx, sb.String(), args...); err != nil { + var pqErr *pq.Error + if errors.As(err, &pqErr) && pqErr.Code == "23503" { + return ErrUserPlatformQuotaFKViolation + } + return err + } + } + return nil +} diff --git a/backend/internal/repository/user_platform_quota_repo_integration_test.go b/backend/internal/repository/user_platform_quota_repo_integration_test.go index f02eeaa900cc..39e2f6e0ba54 100644 --- a/backend/internal/repository/user_platform_quota_repo_integration_test.go +++ b/backend/internal/repository/user_platform_quota_repo_integration_test.go @@ -267,3 +267,101 @@ func TestUserPlatformQuotaRepository_ResetExpiredWindow_NotFoundReturnsSentinel( require.True(t, errors.Is(err, ErrUserPlatformQuotaNotFound), "expected ErrUserPlatformQuotaNotFound, got %v", err) } + +// TestBatchSnapshotUsage_InsertOverwriteMultiKey 验证 BatchSnapshotUsage 的绝对值覆盖语义： +// 1. 首批插入 2 条（不同 user），验证 daily 等于首批值； +// 2. 对同一 key 传不同值，验证 daily 等于新值（绝对覆盖，非累加）。 +func TestBatchSnapshotUsage_InsertOverwriteMultiKey(t *testing.T) { + ctx := context.Background() + // BatchSnapshotUsage 不开事务（直接写），使用独立 client 保证跨调用可见性。 + client := testEntClient(t) + + userID1 := mustCreateUserForQuota(t, client) + userID2 := mustCreateUserForQuota(t, client) + + repo := NewUserPlatformQuotaRepository(client) + + now := time.Date(2026, 5, 29, 12, 0, 0, 0, time.UTC) + dailyStart := time.Date(2026, 5, 29, 0, 0, 0, 0, time.UTC) + weeklyStart := time.Date(2026, 5, 25, 0, 0, 0, 0, time.UTC) // 当周一 + monthlyStart := time.Date(2026, 5, 1, 0, 0, 0, 0, time.UTC) + + // ── 第一批：插入 2 行 ────────────────────────────────────────────────────── + firstBatch := []UserPlatformQuotaSnapshot{ + { + UserID: userID1, + Platform: "anthropic", + DailyUsageUSD: 1.0, + WeeklyUsageUSD: 3.0, + MonthlyUsageUSD: 5.0, + DailyWindowStart: dailyStart, + WeeklyWindowStart: weeklyStart, + MonthlyWindowStart: monthlyStart, + }, + { + UserID: userID2, + Platform: "openai", + DailyUsageUSD: 2.0, + WeeklyUsageUSD: 4.0, + MonthlyUsageUSD: 6.0, + DailyWindowStart: dailyStart, + WeeklyWindowStart: weeklyStart, + MonthlyWindowStart: monthlyStart, + }, + } + require.NoError(t, repo.BatchSnapshotUsage(ctx, firstBatch, now), "first batch upsert") + + // 验证首批 daily 值 + rec1, err := repo.GetByUserPlatform(ctx, userID1, "anthropic") + require.NoError(t, err) + require.NotNil(t, rec1, "user1/anthropic should exist after first batch") + require.InDelta(t, 1.0, rec1.DailyUsageUSD, 1e-9, "user1 daily after first batch") + require.InDelta(t, 3.0, rec1.WeeklyUsageUSD, 1e-9, "user1 weekly after first batch") + require.InDelta(t, 5.0, rec1.MonthlyUsageUSD, 1e-9, "user1 monthly after first batch") + + rec2, err := repo.GetByUserPlatform(ctx, userID2, "openai") + require.NoError(t, err) + require.NotNil(t, rec2, "user2/openai should exist after first batch") + require.InDelta(t, 2.0, rec2.DailyUsageUSD, 1e-9, "user2 daily after first batch") + + // ── 第二批：对同一 key 传不同值，验证绝对覆盖（非累加）────────────────── + now2 := now.Add(5 * time.Minute) + secondBatch := []UserPlatformQuotaSnapshot{ + { + UserID: userID1, + Platform: "anthropic", + DailyUsageUSD: 9.9, // 新值，不是 1.0+9.9=10.9 + WeeklyUsageUSD: 19.9, // 新值，不是 3.0+19.9=22.9 + MonthlyUsageUSD: 29.9, // 新值 + DailyWindowStart: dailyStart, + WeeklyWindowStart: weeklyStart, + MonthlyWindowStart: monthlyStart, + }, + { + UserID: userID2, + Platform: "openai", + DailyUsageUSD: 8.8, + WeeklyUsageUSD: 18.8, + MonthlyUsageUSD: 28.8, + DailyWindowStart: dailyStart, + WeeklyWindowStart: weeklyStart, + MonthlyWindowStart: monthlyStart, + }, + } + require.NoError(t, repo.BatchSnapshotUsage(ctx, secondBatch, now2), "second batch upsert") + + // 验证第二批覆盖：daily 应为新值，不是累加 + rec1After, err := repo.GetByUserPlatform(ctx, userID1, "anthropic") + require.NoError(t, err) + require.NotNil(t, rec1After) + require.InDelta(t, 9.9, rec1After.DailyUsageUSD, 1e-9, "user1 daily must be overwritten to 9.9 (not accumulated)") + require.InDelta(t, 19.9, rec1After.WeeklyUsageUSD, 1e-9, "user1 weekly must be overwritten to 19.9") + require.InDelta(t, 29.9, rec1After.MonthlyUsageUSD, 1e-9, "user1 monthly must be overwritten to 29.9") + + rec2After, err := repo.GetByUserPlatform(ctx, userID2, "openai") + require.NoError(t, err) + require.NotNil(t, rec2After) + require.InDelta(t, 8.8, rec2After.DailyUsageUSD, 1e-9, "user2 daily must be overwritten to 8.8 (not accumulated)") + require.InDelta(t, 18.8, rec2After.WeeklyUsageUSD, 1e-9, "user2 weekly must be overwritten to 18.8") + require.InDelta(t, 28.8, rec2After.MonthlyUsageUSD, 1e-9, "user2 monthly must be overwritten to 28.8") +} diff --git a/backend/internal/repository/user_platform_quota_service_adapter.go b/backend/internal/repository/user_platform_quota_service_adapter.go index 7495cd26ba2b..5240bb5412d1 100644 --- a/backend/internal/repository/user_platform_quota_service_adapter.go +++ b/backend/internal/repository/user_platform_quota_service_adapter.go @@ -94,6 +94,29 @@ func (a *userPlatformQuotaServiceAdapter) ResetExpiredWindow(ctx context.Context return err } +// BatchSnapshotUsage 转换 []service.UserPlatformQuotaSnapshot → []UserPlatformQuotaSnapshot， +// 调底层 repo，并将 repository FK sentinel 包装为 service sentinel。 +func (a *userPlatformQuotaServiceAdapter) BatchSnapshotUsage(ctx context.Context, snapshots []service.UserPlatformQuotaSnapshot, now time.Time) error { + repoSnaps := make([]UserPlatformQuotaSnapshot, len(snapshots)) + for i, s := range snapshots { + repoSnaps[i] = UserPlatformQuotaSnapshot{ + UserID: s.UserID, + Platform: s.Platform, + DailyUsageUSD: s.DailyUsageUSD, + WeeklyUsageUSD: s.WeeklyUsageUSD, + MonthlyUsageUSD: s.MonthlyUsageUSD, + DailyWindowStart: s.DailyWindowStart, + WeeklyWindowStart: s.WeeklyWindowStart, + MonthlyWindowStart: s.MonthlyWindowStart, + } + } + err := a.inner.BatchSnapshotUsage(ctx, repoSnaps, now) + if errors.Is(err, ErrUserPlatformQuotaFKViolation) { + return fmt.Errorf("%w: %v", service.ErrUserPlatformQuotaFKViolation, err) + } + return err +} + // genericUserPlatformQuotaAdapter 通过通用接口适配（用于测试 fake 或非标准实现）。 type genericUserPlatformQuotaAdapter struct { inner UserPlatformQuotaRepository @@ -167,6 +190,29 @@ func (a *genericUserPlatformQuotaAdapter) ResetExpiredWindow(ctx context.Context return err } +// BatchSnapshotUsage 转换 []service.UserPlatformQuotaSnapshot → []UserPlatformQuotaSnapshot（通用 adapter）， +// 并将 repository FK sentinel 包装为 service sentinel。 +func (a *genericUserPlatformQuotaAdapter) BatchSnapshotUsage(ctx context.Context, snapshots []service.UserPlatformQuotaSnapshot, now time.Time) error { + repoSnaps := make([]UserPlatformQuotaSnapshot, len(snapshots)) + for i, s := range snapshots { + repoSnaps[i] = UserPlatformQuotaSnapshot{ + UserID: s.UserID, + Platform: s.Platform, + DailyUsageUSD: s.DailyUsageUSD, + WeeklyUsageUSD: s.WeeklyUsageUSD, + MonthlyUsageUSD: s.MonthlyUsageUSD, + DailyWindowStart: s.DailyWindowStart, + WeeklyWindowStart: s.WeeklyWindowStart, + MonthlyWindowStart: s.MonthlyWindowStart, + } + } + err := a.inner.BatchSnapshotUsage(ctx, repoSnaps, now) + if errors.Is(err, ErrUserPlatformQuotaFKViolation) { + return fmt.Errorf("%w: %v", service.ErrUserPlatformQuotaFKViolation, err) + } + return err +} + // toServiceRecord 将 repository.UserPlatformQuotaRecord 转换为 service.UserPlatformQuotaRecord。 func toServiceRecord(rec *UserPlatformQuotaRecord) *service.UserPlatformQuotaRecord { return &service.UserPlatformQuotaRecord{ diff --git a/backend/internal/service/admin_service_delete_test.go b/backend/internal/service/admin_service_delete_test.go index d01b11e67241..2aae73a96489 100644 --- a/backend/internal/service/admin_service_delete_test.go +++ b/backend/internal/service/admin_service_delete_test.go @@ -471,10 +471,22 @@ func (s *billingCacheStub) DeleteUserPlatformQuotaCache(ctx context.Context, use panic("unexpected DeleteUserPlatformQuotaCache call") } -func (s *billingCacheStub) IncrUserPlatformQuotaUsageCache(ctx context.Context, userID int64, platform string, cost float64, ttl time.Duration) error { +func (s *billingCacheStub) IncrUserPlatformQuotaUsageCache(ctx context.Context, userID int64, platform string, cost float64, ttl time.Duration, markDirty bool) error { panic("unexpected IncrUserPlatformQuotaUsageCache call") } +func (s *billingCacheStub) PopDirtyUserPlatformQuotaKeys(ctx context.Context, n int) ([]UserPlatformQuotaKey, error) { + panic("unexpected PopDirtyUserPlatformQuotaKeys call") +} + +func (s *billingCacheStub) ReaddDirtyUserPlatformQuotaKeys(ctx context.Context, keys []UserPlatformQuotaKey) error { + panic("unexpected ReaddDirtyUserPlatformQuotaKeys call") +} + +func (s *billingCacheStub) BatchGetUserPlatformQuotaCache(ctx context.Context, keys []UserPlatformQuotaKey) ([]*UserPlatformQuotaCacheEntry, error) { + panic("unexpected BatchGetUserPlatformQuotaCache call") +} + func waitForInvalidations(t *testing.T, ch <-chan subscriptionInvalidateCall, expected int) []subscriptionInvalidateCall { t.Helper() calls := make([]subscriptionInvalidateCall, 0, expected) diff --git a/backend/internal/service/auth_service_platform_quota_test.go b/backend/internal/service/auth_service_platform_quota_test.go index f58dc48c35a2..46069814b04c 100644 --- a/backend/internal/service/auth_service_platform_quota_test.go +++ b/backend/internal/service/auth_service_platform_quota_test.go @@ -43,6 +43,10 @@ func (f *fakeInsertRecorder) ResetExpiredWindow(_ context.Context, _ int64, _ st return nil } +func (f *fakeInsertRecorder) BatchSnapshotUsage(_ context.Context, _ []UserPlatformQuotaSnapshot, _ time.Time) error { + return nil +} + func TestSnapshotPlatformQuotaDefaults_PassesToRepoBulkInsert(t *testing.T) { fakeRepo := &fakeInsertRecorder{} s := &AuthService{userPlatformQuotaRepo: fakeRepo} diff --git a/backend/internal/service/auth_service_register_test.go b/backend/internal/service/auth_service_register_test.go index a7c0d260c872..2ee9f21a562c 100644 --- a/backend/internal/service/auth_service_register_test.go +++ b/backend/internal/service/auth_service_register_test.go @@ -105,6 +105,10 @@ func (s *userPlatformQuotaRepoStub) ResetExpiredWindow(context.Context, int64, s panic("unexpected ResetExpiredWindow call") } +func (s *userPlatformQuotaRepoStub) BatchSnapshotUsage(_ context.Context, _ []UserPlatformQuotaSnapshot, _ time.Time) error { + return nil +} + func (s *defaultSubscriptionAssignerStub) AssignOrExtendSubscription(_ context.Context, input *AssignSubscriptionInput) (*UserSubscription, bool, error) { if input != nil { s.calls = append(s.calls, *input) diff --git a/backend/internal/service/billing_cache_service.go b/backend/internal/service/billing_cache_service.go index 8a5172f4d6c9..b734fab13ee8 100644 --- a/backend/internal/service/billing_cache_service.go +++ b/backend/internal/service/billing_cache_service.go @@ -689,7 +689,8 @@ func (s *BillingCacheService) IncrementUserPlatformQuotaUsage(userID int64, plat ctx, cancel := context.WithTimeout(context.Background(), cacheWriteTimeout) defer cancel() ttl := time.Duration(s.cfg.Billing.UserPlatformQuotaCacheTTLSeconds) * time.Second - if err := s.cache.IncrUserPlatformQuotaUsageCache(ctx, userID, platform, cost, ttl); err != nil { + markDirty := s.cfg.Database.UserPlatformQuotaFlusherEnabled + if err := s.cache.IncrUserPlatformQuotaUsageCache(ctx, userID, platform, cost, ttl, markDirty); err != nil { logger.LegacyPrintf("service.billing_cache", "ALERT: incr user platform quota cache failed user=%d platform=%s cost=%f: %v", userID, platform, cost, err) @@ -1310,3 +1311,20 @@ func monthlyQuotaWindowExpired(start *time.Time, now time.Time) bool { } return now.Sub(*start) >= 30*24*time.Hour } + +// HasUserPlatformQuotaLimit 判断该 user×platform 是否设了任一非 nil limit。 +// 写入点守卫:无 limit 直接跳过 Redis 写 + 脏集标记,消除无谓写入。 +// fail-safe:任何不确定(simple 模式除外)都返回 true 维持写入。 +func (s *BillingCacheService) HasUserPlatformQuotaLimit(ctx context.Context, userID int64, platform string) bool { + if s.cfg.RunMode == config.RunModeSimple { + return false + } + if s.cache == nil { + return true + } + entry, ok, err := s.cache.GetUserPlatformQuotaCache(ctx, userID, platform) + if err != nil || !ok || entry == nil { + return true + } + return entry.DailyLimitUSD != nil || entry.WeeklyLimitUSD != nil || entry.MonthlyLimitUSD != nil +} diff --git a/backend/internal/service/billing_cache_service_singleflight_test.go b/backend/internal/service/billing_cache_service_singleflight_test.go index b443d97eaa0b..235b13a6c86c 100644 --- a/backend/internal/service/billing_cache_service_singleflight_test.go +++ b/backend/internal/service/billing_cache_service_singleflight_test.go @@ -79,10 +79,22 @@ func (s *billingCacheMissStub) DeleteUserPlatformQuotaCache(ctx context.Context, return nil } -func (s *billingCacheMissStub) IncrUserPlatformQuotaUsageCache(ctx context.Context, userID int64, platform string, cost float64, ttl time.Duration) error { +func (s *billingCacheMissStub) IncrUserPlatformQuotaUsageCache(ctx context.Context, userID int64, platform string, cost float64, ttl time.Duration, markDirty bool) error { return nil } +func (s *billingCacheMissStub) PopDirtyUserPlatformQuotaKeys(ctx context.Context, n int) ([]UserPlatformQuotaKey, error) { + return nil, nil +} + +func (s *billingCacheMissStub) ReaddDirtyUserPlatformQuotaKeys(ctx context.Context, keys []UserPlatformQuotaKey) error { + return nil +} + +func (s *billingCacheMissStub) BatchGetUserPlatformQuotaCache(ctx context.Context, keys []UserPlatformQuotaKey) ([]*UserPlatformQuotaCacheEntry, error) { + return nil, nil +} + type balanceLoadUserRepoStub struct { mockUserRepo calls atomic.Int64 diff --git a/backend/internal/service/billing_cache_service_test.go b/backend/internal/service/billing_cache_service_test.go index bcd086fa5c1c..c344b4173849 100644 --- a/backend/internal/service/billing_cache_service_test.go +++ b/backend/internal/service/billing_cache_service_test.go @@ -80,10 +80,22 @@ func (b *billingCacheWorkerStub) DeleteUserPlatformQuotaCache(ctx context.Contex return nil } -func (b *billingCacheWorkerStub) IncrUserPlatformQuotaUsageCache(ctx context.Context, userID int64, platform string, cost float64, ttl time.Duration) error { +func (b *billingCacheWorkerStub) IncrUserPlatformQuotaUsageCache(ctx context.Context, userID int64, platform string, cost float64, ttl time.Duration, markDirty bool) error { return nil } +func (b *billingCacheWorkerStub) PopDirtyUserPlatformQuotaKeys(ctx context.Context, n int) ([]UserPlatformQuotaKey, error) { + return nil, nil +} + +func (b *billingCacheWorkerStub) ReaddDirtyUserPlatformQuotaKeys(ctx context.Context, keys []UserPlatformQuotaKey) error { + return nil +} + +func (b *billingCacheWorkerStub) BatchGetUserPlatformQuotaCache(ctx context.Context, keys []UserPlatformQuotaKey) ([]*UserPlatformQuotaCacheEntry, error) { + return nil, nil +} + func TestBillingCacheServiceQueueHighLoad(t *testing.T) { cache := &billingCacheWorkerStub{} svc := NewBillingCacheService(cache, nil, nil, nil, nil, nil, &config.Config{}, nil) diff --git a/backend/internal/service/billing_cache_service_user_platform_quota_test.go b/backend/internal/service/billing_cache_service_user_platform_quota_test.go index 674aa9a5fa85..a82c0e00d75c 100644 --- a/backend/internal/service/billing_cache_service_user_platform_quota_test.go +++ b/backend/internal/service/billing_cache_service_user_platform_quota_test.go @@ -20,14 +20,15 @@ type fakeIncrCache struct { } type incrCall struct { - userID int64 - platform string - cost float64 - ttl time.Duration + userID int64 + platform string + cost float64 + ttl time.Duration + markDirty bool } -func (f *fakeIncrCache) IncrUserPlatformQuotaUsageCache(ctx context.Context, userID int64, platform string, cost float64, ttl time.Duration) error { - f.calls = append(f.calls, incrCall{userID, platform, cost, ttl}) +func (f *fakeIncrCache) IncrUserPlatformQuotaUsageCache(ctx context.Context, userID int64, platform string, cost float64, ttl time.Duration, markDirty bool) error { + f.calls = append(f.calls, incrCall{userID, platform, cost, ttl, markDirty}) return nil } @@ -49,10 +50,10 @@ func TestIncrementUserPlatformQuotaUsage_SyncCallsCache(t *testing.T) { if len(fake.calls) != 2 { t.Fatalf("expected 2 incr calls, got %d", len(fake.calls)) } - if fake.calls[0] != (incrCall{101, "anthropic", 0.25, 120 * time.Second}) { + if fake.calls[0] != (incrCall{userID: 101, platform: "anthropic", cost: 0.25, ttl: 120 * time.Second, markDirty: false}) { t.Errorf("call[0] = %+v", fake.calls[0]) } - if fake.calls[1] != (incrCall{101, "openai", 0.50, 120 * time.Second}) { + if fake.calls[1] != (incrCall{userID: 101, platform: "openai", cost: 0.50, ttl: 120 * time.Second, markDirty: false}) { t.Errorf("call[1] = %+v", fake.calls[1]) } } @@ -88,7 +89,11 @@ func (f *fakeQuotaRepo) ResetExpiredWindow(_ context.Context, _ int64, _ string, return nil } -// fakeFullCache 同时支持 Get + Set + Incr + Delete。 +func (f *fakeQuotaRepo) BatchSnapshotUsage(_ context.Context, _ []UserPlatformQuotaSnapshot, _ time.Time) error { + return nil +} + +// fakeFullCache 同时支持 Get + Set + Incr + Delete + Pop/Readd/BatchGet（脏集读写）。 // mu 保护 entry 和 deleteCalls，防止异步 goroutine 与主 goroutine 之间的 data race。 type fakeFullCache struct { BillingCache @@ -99,6 +104,8 @@ type fakeFullCache struct { lastSetTTL time.Duration // 最近一次 Set 的 ttl getErr error // 非 nil 时 Get 先返回 (nil,false,getErr) setErr error // 非 nil 时 Set 返回该 err(setCalls 仍+1) + // dirty 模拟脏集，供 flusher 测试使用。 + dirty map[UserPlatformQuotaKey]struct{} } // getDeleteCalls 线程安全地读取 deleteCalls。 @@ -161,6 +168,48 @@ func (f *fakeFullCache) DeleteUserPlatformQuotaCache(_ context.Context, _ int64, return nil } +func (f *fakeFullCache) PopDirtyUserPlatformQuotaKeys(_ context.Context, n int) ([]UserPlatformQuotaKey, error) { + f.mu.Lock() + defer f.mu.Unlock() + if len(f.dirty) == 0 { + return nil, nil + } + keys := make([]UserPlatformQuotaKey, 0, n) + for k := range f.dirty { + if len(keys) >= n { + break + } + keys = append(keys, k) + delete(f.dirty, k) + } + return keys, nil +} + +func (f *fakeFullCache) ReaddDirtyUserPlatformQuotaKeys(_ context.Context, keys []UserPlatformQuotaKey) error { + f.mu.Lock() + defer f.mu.Unlock() + if f.dirty == nil { + f.dirty = make(map[UserPlatformQuotaKey]struct{}) + } + for _, k := range keys { + f.dirty[k] = struct{}{} + } + return nil +} + +// BatchGetUserPlatformQuotaCache 对每个 key 返回 f.entry（MISS → nil）， +// 保持与输入 keys 顺序/长度对齐。注意此处所有 key 共享同一个 entry， +// 仅用于测试场景。 +func (f *fakeFullCache) BatchGetUserPlatformQuotaCache(_ context.Context, keys []UserPlatformQuotaKey) ([]*UserPlatformQuotaCacheEntry, error) { + f.mu.Lock() + defer f.mu.Unlock() + results := make([]*UserPlatformQuotaCacheEntry, len(keys)) + for i := range keys { + results[i] = f.entry + } + return results, nil +} + func newServiceForPreflight(t *testing.T, repo UserPlatformQuotaRepository, cache BillingCache) *BillingCacheService { t.Helper() cfg := &config.Config{} @@ -713,3 +762,71 @@ func TestCheckUserPlatformQuotaEligibility_SentinelCrossDay_NoRefresh(t *testing t.Errorf("sentinel cross-window must NOT trigger refresh SetCache, got %d calls", cache.getSetCalls()) } } + +// ── TestHasUserPlatformQuotaLimit ──────────────────────────────────────────── + +func TestHasUserPlatformQuotaLimit(t *testing.T) { + daily := 5.0 + + tests := []struct { + name string + setup func() *BillingCacheService + want bool + }{ + { + name: "has_limit", + setup: func() *BillingCacheService { + entry := &UserPlatformQuotaCacheEntry{DailyLimitUSD: &daily} + svc := newServiceForPreflight(t, &fakeQuotaRepo{}, &fakeFullCache{entry: entry}) + return svc + }, + want: true, + }, + { + name: "sentinel_no_limit", + setup: func() *BillingCacheService { + entry := &UserPlatformQuotaCacheEntry{} // 三个 limit 字段全 nil + svc := newServiceForPreflight(t, &fakeQuotaRepo{}, &fakeFullCache{entry: entry}) + return svc + }, + want: false, + }, + { + name: "cache_miss", + setup: func() *BillingCacheService { + // entry==nil → GetUserPlatformQuotaCache 返回 (nil,false,nil) + svc := newServiceForPreflight(t, &fakeQuotaRepo{}, &fakeFullCache{}) + return svc + }, + want: true, // fail-safe + }, + { + name: "redis_err", + setup: func() *BillingCacheService { + svc := newServiceForPreflight(t, &fakeQuotaRepo{}, &fakeFullCache{getErr: errors.New("redis down")}) + return svc + }, + want: true, // fail-safe + }, + { + name: "simple_mode", + setup: func() *BillingCacheService { + entry := &UserPlatformQuotaCacheEntry{DailyLimitUSD: &daily} + svc := newServiceForPreflight(t, &fakeQuotaRepo{}, &fakeFullCache{entry: entry}) + svc.cfg.RunMode = config.RunModeSimple + return svc + }, + want: false, // simple 模式始终跳过 + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + svc := tt.setup() + got := svc.HasUserPlatformQuotaLimit(context.Background(), 1, "anthropic") + if got != tt.want { + t.Errorf("HasUserPlatformQuotaLimit() = %v, want %v", got, tt.want) + } + }) + } +} diff --git a/backend/internal/service/billing_service.go b/backend/internal/service/billing_service.go index 940a827df52e..6b1438e8d915 100644 --- a/backend/internal/service/billing_service.go +++ b/backend/internal/service/billing_service.go @@ -21,6 +21,12 @@ type APIKeyRateLimitCacheData struct { Window7d int64 `json:"window_7d"` } +// UserPlatformQuotaKey 标识一个 user×platform，用于脏集出入与批量读。 +type UserPlatformQuotaKey struct { + UserID int64 + Platform string +} + // UserPlatformQuotaCacheEntry Redis hash 反序列化结果。 // // SchemaVersion 用于向后兼容： @@ -72,7 +78,13 @@ type BillingCache interface { SetUserPlatformQuotaCache(ctx context.Context, userID int64, platform string, entry *UserPlatformQuotaCacheEntry, ttl time.Duration) error DeleteUserPlatformQuotaCache(ctx context.Context, userID int64, platform string) error // IncrUserPlatformQuotaUsageCache 在缓存命中时累加用量；缓存未命中（key 不存在）静默返回 nil。 - IncrUserPlatformQuotaUsageCache(ctx context.Context, userID int64, platform string, cost float64, ttl time.Duration) error + // markDirty=true 时将该 key 的 member 写入 Redis 脏集，供 flusher 批量回写 DB。 + IncrUserPlatformQuotaUsageCache(ctx context.Context, userID int64, platform string, cost float64, ttl time.Duration, markDirty bool) error + + // 脏集读写，供 flusher 使用。 + PopDirtyUserPlatformQuotaKeys(ctx context.Context, n int) ([]UserPlatformQuotaKey, error) + ReaddDirtyUserPlatformQuotaKeys(ctx context.Context, keys []UserPlatformQuotaKey) error + BatchGetUserPlatformQuotaCache(ctx context.Context, keys []UserPlatformQuotaKey) ([]*UserPlatformQuotaCacheEntry, error) } // ModelPricing 模型价格配置（per-token价格，与LiteLLM格式一致） diff --git a/backend/internal/service/gateway_service.go b/backend/internal/service/gateway_service.go index effa803aa755..f807f3ec6532 100644 --- a/backend/internal/service/gateway_service.go +++ b/backend/internal/service/gateway_service.go @@ -96,11 +96,13 @@ var ( modelsListCacheMissTotal atomic.Int64 modelsListCacheStoreTotal atomic.Int64 + // Deprecated: flusher_enabled=true 后不再增长(仅 flag=false 降级直写路径使用);新主路径见 FlusherMetrics。remove after 2026-09。 // userPlatformQuotaDBIncrErrorTotal 统计 finalizePostUsageBilling 异步 goroutine // 中 IncrementUsageWithReset 失败次数。Redis 已成功累加 + DB 写失败意味着 // Redis cache TTL 过期或被清后该笔 cost 会丢失（与实际消费偏差）。 // oncall 通过 GatewayUserPlatformQuotaIncrStats() 暴露给 ops 面板做阈值告警。 userPlatformQuotaDBIncrErrorTotal atomic.Int64 + // Deprecated: flusher_enabled=true 后不再增长(仅 flag=false 降级直写路径使用);新主路径见 FlusherMetrics。remove after 2026-09。 // userPlatformQuotaDBIncrLegacyErrorTotal 统计 legacy postUsageBilling // （applyUsageBilling 在 repo==nil 时 fallback）路径下的失败次数； // 与 DB Incr 失败分开计数，便于区分"主路径暂时故障"vs"基础设施长期未配齐"。 @@ -141,6 +143,23 @@ func GatewayUserPlatformQuotaIncrStats() (mainPathErr, legacyPathErr, sentinelSe userPlatformQuotaSentinelSetCacheErrorTotal.Load() } +// GatewayUserPlatformQuotaFlusherStats 暴露 flusher 运行指标供 ops/health 面板查询。 +func GatewayUserPlatformQuotaFlusherStats(f *UserPlatformQuotaUsageFlusher) map[string]int64 { + if f == nil || f.metrics == nil { + return nil + } + m := f.metrics + return map[string]int64{ + "flush_success": m.FlushSuccessTotal.Load(), + "flush_error": m.FlushErrorTotal.Load(), + "flush_batch_size": m.FlushBatchSizeTotal.Load(), + "flush_latency_ms_max": m.FlushLatencyMsMax.Load(), + "dirty_readd": m.DirtyReaddTotal.Load(), + "dirty_lost": m.DirtyLostTotal.Load(), + "flush_fk_violation": m.FlushFKViolationTotal.Load(), + } +} + func openAIStreamEventIsTerminal(data string) bool { trimmed := strings.TrimSpace(data) if trimmed == "" { @@ -8234,18 +8253,23 @@ func postUsageBilling(ctx context.Context, p *postUsageBillingParams, deps *bill } } - // Platform quota DB-only 累加（与 finalizePostUsageBilling 行为对齐的兜底）： - // - 仅对 standard（余额）模式生效；订阅模式豁免 - // - 直接走 DB，不经 Redis Incr 队列：legacy 路径在 repo==nil（仓库未注入） - // 时被触发，此时整套 billing repo 都不可用，没有"双队列"风险 - // - 失败仅记 ALERT log + counter，不阻断主扣费流程；与正常路径一致 - // - // 历史背景：原 legacy path 完全跳过此累加，导致部署中如果 repo 偶然为 nil - // 时用户消费可绕过 platform quota，存在静默资金风险。 + // Platform quota 累加（legacy 兜底路径）：仅对 standard（余额）模式生效；订阅模式豁免；仅对有 limit 的用户写 + // - HasUserPlatformQuotaLimit 守卫:与正常路径对齐，无 limit 公司跳过 + // - 新增 Redis 同步写:enforcement 走 Redis，legacy 路径也必须同步写，否则 preflight 看不到消费 + // - flusher_enabled=false（降级）:保留原有同步直写 DB + // - flusher_enabled=true:跳过直写 DB，由 flusher 异步批量刷（markDirty 在 IncrementUserPlatformQuotaUsage 内部完成） + // - 失败仅记 ALERT log + counter，不阻断主扣费流程 if !p.IsSubscriptionBill && p.Platform != "" && cost.ActualCost > 0 && p.User != nil && deps.userPlatformQuotaRepo != nil { - if err := deps.userPlatformQuotaRepo.IncrementUsageWithReset(billingCtx, p.User.ID, p.Platform, cost.ActualCost, time.Now().UTC()); err != nil { - userPlatformQuotaDBIncrLegacyErrorTotal.Add(1) - logger.LegacyPrintf("service.gateway", "ALERT: legacy incr user platform quota DB failed user=%d platform=%s cost=%f: %v", p.User.ID, p.Platform, cost.ActualCost, err) + if deps.billingCacheService.HasUserPlatformQuotaLimit(billingCtx, p.User.ID, p.Platform) { + deps.billingCacheService.IncrementUserPlatformQuotaUsage(p.User.ID, p.Platform, cost.ActualCost) + if deps.cfg == nil || !deps.cfg.Database.UserPlatformQuotaFlusherEnabled { + // 降级路径:flusher 未启用时保留原有同步直写 DB + if err := deps.userPlatformQuotaRepo.IncrementUsageWithReset(billingCtx, p.User.ID, p.Platform, cost.ActualCost, time.Now().UTC()); err != nil { + userPlatformQuotaDBIncrLegacyErrorTotal.Add(1) + logger.LegacyPrintf("service.gateway", "ALERT: legacy incr user platform quota DB failed user=%d platform=%s cost=%f: %v", p.User.ID, p.Platform, cost.ActualCost, err) + } + } + // flusher_enabled=true:不直写 DB，flusher 异步批量刷 } } @@ -8395,30 +8419,38 @@ func finalizePostUsageBilling(ctx context.Context, p *postUsageBillingParams, de deps.deferredService.ScheduleLastUsedUpdate(p.Account.ID) - // Platform quota 累加：仅在 standard（余额）模式生效；订阅模式豁免 - // Redis 同步写 + DB 异步持久化: + // Platform quota 累加：仅在 standard（余额）模式生效；订阅模式豁免；仅对有 limit 的用户写 + // Redis 同步写 + DB 异步持久化（flag=false 降级）或 flusher 异步刷（flag=true）: + // - HasUserPlatformQuotaLimit 守卫:无 limit 的公司跳过,避免无效写入 + 浪费 Redis 容量 // - Redis 同步:确保下次 preflight 立即看到最新 usage,把 TOCTOU 超支窗口 // 限制在并发 in-flight 请求数量内（旧实现的异步入队会让超支无限累积直到 worker 处理） - // - DB 异步:在独立 goroutine 中走 detached context,失败用 ALERT log 触发 oncall 对账 + // - DB 异步(flusher_enabled=false):在独立 goroutine 中走 detached context,失败用 ALERT log 触发 oncall 对账 + // - flusher_enabled=true:不直写 DB,由 flusher 异步批量刷（markDirty 已在 IncrementUserPlatformQuotaUsage 内部完成） if !p.IsSubscriptionBill && p.Platform != "" && p.Cost.ActualCost > 0 && p.User != nil && deps.userPlatformQuotaRepo != nil { - deps.billingCacheService.IncrementUserPlatformQuotaUsage(p.User.ID, p.Platform, p.Cost.ActualCost) - dbCtx, dbCancel := detachUpstreamContext(ctx) - userID, platform, cost := p.User.ID, p.Platform, p.Cost.ActualCost - go func() { - defer func() { - if r := recover(); r != nil { - logger.LegacyPrintf("service.gateway", "ALERT: panic in user platform quota incr goroutine user=%d platform=%s: %v", userID, platform, r) - } - }() - defer dbCancel() - if err := deps.userPlatformQuotaRepo.IncrementUsageWithReset(dbCtx, userID, platform, cost, time.Now().UTC()); err != nil { - // 失败计数器:暴露给 GatewayUserPlatformQuotaIncrStats(),由 ops 面板做斜率告警。 - userPlatformQuotaDBIncrErrorTotal.Add(1) - // ALERT 级别:DB 持久化失败意味着 Redis cache 失效后该笔 cost 永久丢失, - // 用户配额视图与实际消费会偏差,oncall 需要据此对账或人工补录。 - logger.LegacyPrintf("service.gateway", "ALERT: incr user platform quota DB failed user=%d platform=%s cost=%f: %v", userID, platform, cost, err) + if deps.billingCacheService.HasUserPlatformQuotaLimit(ctx, p.User.ID, p.Platform) { + deps.billingCacheService.IncrementUserPlatformQuotaUsage(p.User.ID, p.Platform, p.Cost.ActualCost) + if deps.cfg == nil || !deps.cfg.Database.UserPlatformQuotaFlusherEnabled { + // 降级路径:flusher 未启用时保留原有异步直写 DB + dbCtx, dbCancel := detachUpstreamContext(ctx) + userID, platform, cost := p.User.ID, p.Platform, p.Cost.ActualCost + go func() { + defer func() { + if r := recover(); r != nil { + logger.LegacyPrintf("service.gateway", "ALERT: panic in user platform quota incr goroutine user=%d platform=%s: %v", userID, platform, r) + } + }() + defer dbCancel() + if err := deps.userPlatformQuotaRepo.IncrementUsageWithReset(dbCtx, userID, platform, cost, time.Now().UTC()); err != nil { + // 失败计数器:暴露给 GatewayUserPlatformQuotaIncrStats(),由 ops 面板做斜率告警。 + userPlatformQuotaDBIncrErrorTotal.Add(1) + // ALERT 级别:DB 持久化失败意味着 Redis cache 失效后该笔 cost 永久丢失, + // 用户配额视图与实际消费会偏差,oncall 需要据此对账或人工补录。 + logger.LegacyPrintf("service.gateway", "ALERT: incr user platform quota DB failed user=%d platform=%s cost=%f: %v", userID, platform, cost, err) + } + }() } - }() + // flusher_enabled=true:不直写 DB,flusher 异步批量刷 + } } // Notification checks run async — all parameters are already captured, @@ -8533,6 +8565,7 @@ type billingDeps struct { deferredService *DeferredService balanceNotifyService *BalanceNotifyService userPlatformQuotaRepo UserPlatformQuotaRepository + cfg *config.Config } func (s *GatewayService) billingDeps() *billingDeps { @@ -8544,6 +8577,7 @@ func (s *GatewayService) billingDeps() *billingDeps { deferredService: s.deferredService, balanceNotifyService: s.balanceNotifyService, userPlatformQuotaRepo: s.userPlatformQuotaRepo, + cfg: s.cfg, } } diff --git a/backend/internal/service/user_platform_quota_flusher.go b/backend/internal/service/user_platform_quota_flusher.go new file mode 100644 index 000000000000..3ee23d2c9372 --- /dev/null +++ b/backend/internal/service/user_platform_quota_flusher.go @@ -0,0 +1,267 @@ +package service + +import ( + "context" + "errors" + "sync/atomic" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" +) + +// quotaDirtyCache 是 flusher 依赖的窄接口（来自 BillingCache）。 +type quotaDirtyCache interface { + PopDirtyUserPlatformQuotaKeys(ctx context.Context, n int) ([]UserPlatformQuotaKey, error) + ReaddDirtyUserPlatformQuotaKeys(ctx context.Context, keys []UserPlatformQuotaKey) error + BatchGetUserPlatformQuotaCache(ctx context.Context, keys []UserPlatformQuotaKey) ([]*UserPlatformQuotaCacheEntry, error) +} + +// quotaSnapshotWriter 是 flusher 依赖的 DB 写入窄接口。 +// 使用 service 层的 UserPlatformQuotaSnapshot，避免与 repository 包形成循环依赖； +// 实际实现由 repository adapter 在 B7 注入。 +type quotaSnapshotWriter interface { + BatchSnapshotUsage(ctx context.Context, snapshots []UserPlatformQuotaSnapshot, now time.Time) error +} + +// FlusherMetrics 记录 flusher 运行时指标（原子量，零值可用）。 +type FlusherMetrics struct { + FlushSuccessTotal atomic.Int64 + FlushErrorTotal atomic.Int64 + FlushBatchSizeTotal atomic.Int64 + FlushLatencyMsMax atomic.Int64 + DirtyReaddTotal atomic.Int64 + // DirtyLostTotal：Readd 失败导致脏 key 丢失——已 SPOP+主操作失败+Readd 也失败； + // Redis 仍权威，活跃 key 下次 SADD 自愈。 + DirtyLostTotal atomic.Int64 + FlushFKViolationTotal atomic.Int64 +} + +// flusherMaxBatchesPerTick 单次 tick 最多消费的批数，防止 tick 执行时间过长。 +const flusherMaxBatchesPerTick = 16 + +// maxFlushBatchSize 限制单批行数,必须 ≤ repository.BatchSnapshotUsage 的 batchRows(6000), +// 以保证单次 flush 的 snapshots 仅生成一条 UPSERT(单事务原子)。两处需手动保持一致。 +const maxFlushBatchSize = 6000 + +// defaultFlushBatchSize 是配置 flush_batch_size 非法(≤0)时的回退值。 +const defaultFlushBatchSize = 1000 + +// UserPlatformQuotaUsageFlusher 将 Redis 脏集快照定期批量写入 DB。 +// 不维护任何 delta/in-process 状态；每批读取 Redis 当前绝对值覆盖写入。 +type UserPlatformQuotaUsageFlusher struct { + cache quotaDirtyCache + quotaRepo quotaSnapshotWriter + timingWheel *TimingWheelService + // enabled 对应 flusher_enabled 配置；false 时 Start() 不注册定时器。 + enabled bool + interval time.Duration + batchSize int + flushTimeout time.Duration + metrics *FlusherMetrics + stopped atomic.Bool +} + +// NewUserPlatformQuotaUsageFlusher 创建 UserPlatformQuotaUsageFlusher。 +// cache(BillingCache) 隐式满足 quotaDirtyCache；quotaRepo(UserPlatformQuotaRepository) 隐式满足 quotaSnapshotWriter。 +func NewUserPlatformQuotaUsageFlusher(cfg *config.Config, cache BillingCache, quotaRepo UserPlatformQuotaRepository, tw *TimingWheelService) *UserPlatformQuotaUsageFlusher { + batchSize := cfg.Database.UserPlatformQuotaFlushBatchSize + if batchSize <= 0 { + batchSize = defaultFlushBatchSize + } + if batchSize > maxFlushBatchSize { + logger.LegacyPrintf("quota_flusher", + "[QuotaFlusher] flush_batch_size %d 超过上限 %d,已 clamp(避免 BatchSnapshotUsage 多子批非原子)", + cfg.Database.UserPlatformQuotaFlushBatchSize, maxFlushBatchSize) + batchSize = maxFlushBatchSize + } + interval := time.Duration(cfg.Database.UserPlatformQuotaFlushIntervalMs) * time.Millisecond + if interval <= 0 { + logger.LegacyPrintf("quota_flusher", "[QuotaFlusher] flush_interval_ms %d 非法,回退 2000ms", cfg.Database.UserPlatformQuotaFlushIntervalMs) + interval = 2 * time.Second + } + return &UserPlatformQuotaUsageFlusher{ + cache: cache, + quotaRepo: quotaRepo, + timingWheel: tw, + enabled: cfg.Database.UserPlatformQuotaFlusherEnabled, + interval: interval, + batchSize: batchSize, + flushTimeout: 3 * time.Second, + metrics: &FlusherMetrics{}, + } +} + +// updateLatencyMax 用 CAS 单调更新最大延迟。 +func (s *UserPlatformQuotaUsageFlusher) updateLatencyMax(ms int64) { + for { + old := s.metrics.FlushLatencyMsMax.Load() + if ms <= old { + return + } + if s.metrics.FlushLatencyMsMax.CompareAndSwap(old, ms) { + return + } + } +} + +// readdOrCountLost 尝试把 keys 回填脏集：成功计 DirtyReaddTotal，失败计 DirtyLostTotal 并 ALERT。 +func (s *UserPlatformQuotaUsageFlusher) readdOrCountLost(ctx context.Context, keys []UserPlatformQuotaKey, stage string) { + if err := s.cache.ReaddDirtyUserPlatformQuotaKeys(ctx, keys); err != nil { + s.metrics.DirtyLostTotal.Add(int64(len(keys))) + logger.LegacyPrintf("quota_flusher", "[QuotaFlusher] ALERT: Readd after %s failed, %d keys 丢出脏集(DB 镜像缺这批,Redis 仍权威,活跃 key 下次 SADD 自愈): %v", stage, len(keys), err) + return + } + s.metrics.DirtyReaddTotal.Add(int64(len(keys))) +} + +// flushOneBatch 处理单批：Pop → BatchGet → 组装 snaps → BatchSnapshotUsage。 +// 返回 (shouldContinue bool)：false 表示本轮循环应停止（空集/错误/最后一批）。 +// 每次调用独立创建带 timeout 的 ctx 并 defer cancel，不会在循环中累积泄漏。 +func (s *UserPlatformQuotaUsageFlusher) flushOneBatch(parentCtx context.Context) bool { + ctx, cancel := context.WithTimeout(parentCtx, s.flushTimeout) + defer cancel() + + // 1. Pop 脏集 + keys, err := s.cache.PopDirtyUserPlatformQuotaKeys(ctx, s.batchSize) + if err != nil { + s.metrics.FlushErrorTotal.Add(1) + logger.LegacyPrintf("quota_flusher", "[QuotaFlusher] PopDirty error: %v", err) + return false + } + if len(keys) == 0 { + // 脏集已空 + return false + } + + // 2. 批量读 Redis 快照 + entries, err := s.cache.BatchGetUserPlatformQuotaCache(ctx, keys) + if err != nil { + s.metrics.FlushErrorTotal.Add(1) + s.readdOrCountLost(ctx, keys, "BatchGet") + logger.LegacyPrintf("quota_flusher", "[QuotaFlusher] BatchGet error: %v", err) + return false + } + + // 3. 组装 snapshots（MISS 或任一 WindowStart==nil → 跳过） + snaps := make([]UserPlatformQuotaSnapshot, 0, len(keys)) + for i, key := range keys { + e := entries[i] + if e == nil { + continue + } + if e.DailyWindowStart == nil || e.WeeklyWindowStart == nil || e.MonthlyWindowStart == nil { + continue + } + snaps = append(snaps, UserPlatformQuotaSnapshot{ + UserID: key.UserID, + Platform: key.Platform, + DailyUsageUSD: e.DailyUsageUSD, + WeeklyUsageUSD: e.WeeklyUsageUSD, + MonthlyUsageUSD: e.MonthlyUsageUSD, + DailyWindowStart: *e.DailyWindowStart, + WeeklyWindowStart: *e.WeeklyWindowStart, + MonthlyWindowStart: *e.MonthlyWindowStart, + }) + } + + // 4. 全部 MISS/异常跳过时 + if len(snaps) == 0 { + // 若 Pop 数量已不满一批，表示脏集将空，停止 + if len(keys) < s.batchSize { + return false + } + // 否则继续下一批（可能还有更多脏 key） + return true + } + + // 已知竞态(admin 写 × flusher 刷,仅 flusher_enabled=true 时存在): + // admin ResetExpiredWindow/UpsertForUser 是"先写 DB 再 DeleteCache"。若本批已 SPOP + BatchGet + // 读到旧 usage 快照(此刻 member 已离开脏集),而 admin 随后写 DB、本行 UPSERT 又在 admin 写之后落库, + // 则旧快照会覆盖 admin 刚写入的值;DeleteCache 后 Redis MISS,下次 preflight 从 DB 重载被覆盖的旧值。 + // 因 member 已被 SPOP,admin 侧 SREM/清脏标记无法拦截本批(故未做)。影响有限,暂列为已知取舍: + // - UpsertForUser 改 limit,而本 UPSERT 不写 limit 列 → limit 配置不受影响; + // - ResetExpiredWindow 改 usage,但 preflight windowExpired 会在窗口真正过期时自愈重置, + // 仅"强制重置未过期窗口"且与本批精确交错时短暂失效; + // - 低频 admin 操作 + 默认 flusher_enabled=false。彻底消除需 version OCC(DB 加 version 列条件 UPSERT), + // 成本高;启用 flusher 后如需强一致再评估。 + + // 5. 写入 DB + start := time.Now() + writeErr := s.quotaRepo.BatchSnapshotUsage(ctx, snaps, time.Now().UTC()) + s.updateLatencyMax(time.Since(start).Milliseconds()) + + if writeErr != nil { + if errors.Is(writeErr, ErrUserPlatformQuotaFKViolation) { + // 注意:PG FK violation 是整条 INSERT 回滚 → 整批(含同批正常用户)均未写入 DB, + // 且这些 key 已被 SPOP 出脏集、此处不 Readd。活跃 key 会在下次请求重新 SADD, + // flusher 读 Redis 当前累计绝对值刷库即自愈;低活跃 key 这轮 DB usage 偏低 + // (Redis 仍是 enforcement 权威,不受影响;DB 仅展示)。已删用户边角的接受取舍,不做逐行重试。 + // FK 违反：用户已被删除，直接丢弃不 Readd + s.metrics.FlushFKViolationTotal.Add(1) + s.metrics.FlushErrorTotal.Add(1) + logger.LegacyPrintf("quota_flusher", "[QuotaFlusher] FK violation (dropped %d snaps): %v", len(snaps), writeErr) + } else { + // 其他错误：回填脏集，保留下次重试 + s.metrics.FlushErrorTotal.Add(1) + s.readdOrCountLost(ctx, keys, "BatchSnapshotUsage") + logger.LegacyPrintf("quota_flusher", "[QuotaFlusher] BatchSnapshotUsage error: %v", writeErr) + } + return false + } + + // 6. 成功 + s.metrics.FlushSuccessTotal.Add(1) + s.metrics.FlushBatchSizeTotal.Add(int64(len(snaps))) + + // 若 Pop 数量不满一批，脏集已空，停止 + if len(keys) < s.batchSize { + return false + } + return true +} + +// flush 执行一次完整的 flush，循环消费至脏集空或达到 maxBatchesPerTick。 +func (s *UserPlatformQuotaUsageFlusher) flush() { + if s == nil { + return + } + parentCtx := context.Background() + for b := 0; b < flusherMaxBatchesPerTick; b++ { + if !s.flushOneBatch(parentCtx) { + return + } + } + // 连续消费满 flusherMaxBatchesPerTick 批仍未取空脏集:本 tick 主动让出,剩余积压留待下一 tick。 + // 记一条 log 便于 oncall 发现 distinct 活跃 key 远超 maxBatchesPerTick×batchSize(DB 镜像延迟上升); + // 可配合 Redis SCARD billing:upq:dirty 观察脏集存量。 + logger.LegacyPrintf("quota_flusher", + "[QuotaFlusher] 单 tick 达到 max batches 上限(%d × batchSize=%d),脏集仍非空,积压顺延至下一 tick", + flusherMaxBatchesPerTick, s.batchSize) +} + +// tick 是 TimingWheel 回调。若 flusher 已停止则直接返回。 +func (s *UserPlatformQuotaUsageFlusher) tick() { + if s == nil || s.stopped.Load() { + return + } + s.flush() +} + +// Start 注册定时 tick。flusher_enabled=false 时直接返回，不注册定时器。 +func (s *UserPlatformQuotaUsageFlusher) Start() { + if s == nil || !s.enabled { + return + } + s.timingWheel.ScheduleRecurring("deferred:platform_quota", s.interval, s.tick) +} + +// Stop 停止 flusher：标记 stopped → Cancel 定时器 → 执行最后一次 flush。 +func (s *UserPlatformQuotaUsageFlusher) Stop() { + if s == nil { + return + } + s.stopped.Store(true) + s.timingWheel.Cancel("deferred:platform_quota") + s.flush() +} diff --git a/backend/internal/service/user_platform_quota_flusher_test.go b/backend/internal/service/user_platform_quota_flusher_test.go new file mode 100644 index 000000000000..4f7344819ae0 --- /dev/null +++ b/backend/internal/service/user_platform_quota_flusher_test.go @@ -0,0 +1,511 @@ +package service + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" +) + +// --------------------------------------------------------------------------- +// Mock: quotaDirtyCache +// --------------------------------------------------------------------------- + +type mockQuotaDirtyCache struct { + // popSequence: 第 0 次 Pop 返回 popSequence[0]，之后返回 nil（空集） + popSequence [][]UserPlatformQuotaKey + popCallIdx int + + // getEntries: BatchGetUserPlatformQuotaCache 返回的 entries（与 keys 对齐） + getEntries []*UserPlatformQuotaCacheEntry + getErr error + + // readdCalled: 记录 Readd 收到的 keys（累积所有次调用） + readdCalled [][]UserPlatformQuotaKey + readdErr error +} + +func (m *mockQuotaDirtyCache) PopDirtyUserPlatformQuotaKeys(_ context.Context, _ int) ([]UserPlatformQuotaKey, error) { + if m.popCallIdx < len(m.popSequence) { + keys := m.popSequence[m.popCallIdx] + m.popCallIdx++ + return keys, nil + } + // 超出序列 → 空集（模拟脏集已清空） + return nil, nil +} + +func (m *mockQuotaDirtyCache) ReaddDirtyUserPlatformQuotaKeys(_ context.Context, keys []UserPlatformQuotaKey) error { + m.readdCalled = append(m.readdCalled, keys) + return m.readdErr +} + +func (m *mockQuotaDirtyCache) BatchGetUserPlatformQuotaCache(_ context.Context, _ []UserPlatformQuotaKey) ([]*UserPlatformQuotaCacheEntry, error) { + if m.getErr != nil { + return nil, m.getErr + } + return m.getEntries, nil +} + +// --------------------------------------------------------------------------- +// Mock: quotaSnapshotWriter +// --------------------------------------------------------------------------- + +type mockQuotaSnapshotWriter struct { + receivedSnaps []UserPlatformQuotaSnapshot + returnErr error +} + +func (m *mockQuotaSnapshotWriter) BatchSnapshotUsage(_ context.Context, snaps []UserPlatformQuotaSnapshot, _ time.Time) error { + m.receivedSnaps = append(m.receivedSnaps, snaps...) + return m.returnErr +} + +// --------------------------------------------------------------------------- +// Helper: 构造窗口起始时间（非 nil） +// --------------------------------------------------------------------------- + +func flusherPtrTime(t time.Time) *time.Time { return &t } + +func makeEntry(daily, weekly, monthly float64) *UserPlatformQuotaCacheEntry { + now := time.Now().UTC() + return &UserPlatformQuotaCacheEntry{ + DailyUsageUSD: daily, + WeeklyUsageUSD: weekly, + MonthlyUsageUSD: monthly, + DailyWindowStart: flusherPtrTime(now), + WeeklyWindowStart: flusherPtrTime(now), + MonthlyWindowStart: flusherPtrTime(now), + } +} + +// --------------------------------------------------------------------------- +// newTestFlusher: 直接构造 struct（跳过构造函数，B7 才注入） +// --------------------------------------------------------------------------- + +func newTestFlusher(cache quotaDirtyCache, writer quotaSnapshotWriter) *UserPlatformQuotaUsageFlusher { + return &UserPlatformQuotaUsageFlusher{ + cache: cache, + quotaRepo: writer, + timingWheel: nil, // 单测不启动 TimingWheel + interval: 5 * time.Second, + batchSize: 100, + flushTimeout: 5 * time.Second, + metrics: &FlusherMetrics{}, + } +} + +// --------------------------------------------------------------------------- +// 场景 1: PopSnapshotUpsert — 2 key + 2 个含 window 的 entry → writer 收 2 行 +// --------------------------------------------------------------------------- + +func TestFlusher_PopSnapshotUpsert(t *testing.T) { + keys := []UserPlatformQuotaKey{ + {UserID: 1, Platform: "anthropic"}, + {UserID: 2, Platform: "openai"}, + } + cache := &mockQuotaDirtyCache{ + popSequence: [][]UserPlatformQuotaKey{keys}, // 第 1 次返回 keys，之后空 + getEntries: []*UserPlatformQuotaCacheEntry{ + makeEntry(1.0, 2.0, 3.0), + makeEntry(4.0, 5.0, 6.0), + }, + } + writer := &mockQuotaSnapshotWriter{} + f := newTestFlusher(cache, writer) + + f.flush() + + if len(writer.receivedSnaps) != 2 { + t.Fatalf("expected 2 snaps, got %d", len(writer.receivedSnaps)) + } + if f.metrics.FlushBatchSizeTotal.Load() != 2 { + t.Errorf("FlushBatchSizeTotal = %d, want 2", f.metrics.FlushBatchSizeTotal.Load()) + } + if f.metrics.FlushSuccessTotal.Load() != 1 { + t.Errorf("FlushSuccessTotal = %d, want 1", f.metrics.FlushSuccessTotal.Load()) + } + if f.metrics.FlushErrorTotal.Load() != 0 { + t.Errorf("FlushErrorTotal = %d, want 0", f.metrics.FlushErrorTotal.Load()) + } +} + +// --------------------------------------------------------------------------- +// 场景 2: MissKeySkipped — 2 key，BatchGet 返回 [entry, nil] → 只刷 1 行，nil 跳过，不 Readd +// --------------------------------------------------------------------------- + +func TestFlusher_MissKeySkipped(t *testing.T) { + keys := []UserPlatformQuotaKey{ + {UserID: 1, Platform: "anthropic"}, + {UserID: 2, Platform: "openai"}, + } + cache := &mockQuotaDirtyCache{ + popSequence: [][]UserPlatformQuotaKey{keys}, + getEntries: []*UserPlatformQuotaCacheEntry{ + makeEntry(1.0, 2.0, 3.0), + nil, // MISS + }, + } + writer := &mockQuotaSnapshotWriter{} + f := newTestFlusher(cache, writer) + + f.flush() + + if len(writer.receivedSnaps) != 1 { + t.Fatalf("expected 1 snap, got %d", len(writer.receivedSnaps)) + } + if writer.receivedSnaps[0].UserID != 1 { + t.Errorf("expected snap for UserID=1, got %d", writer.receivedSnaps[0].UserID) + } + if len(cache.readdCalled) != 0 { + t.Errorf("Readd should NOT be called on MISS, got %d calls", len(cache.readdCalled)) + } + if f.metrics.FlushSuccessTotal.Load() != 1 { + t.Errorf("FlushSuccessTotal = %d, want 1", f.metrics.FlushSuccessTotal.Load()) + } +} + +// --------------------------------------------------------------------------- +// 场景 3: UpsertFailReadds — writer 返普通 error → keys 被 Readd，FlushErrorTotal=1，DirtyReaddTotal=len +// --------------------------------------------------------------------------- + +func TestFlusher_UpsertFailReadds(t *testing.T) { + keys := []UserPlatformQuotaKey{ + {UserID: 1, Platform: "anthropic"}, + {UserID: 2, Platform: "openai"}, + } + cache := &mockQuotaDirtyCache{ + popSequence: [][]UserPlatformQuotaKey{keys}, + getEntries: []*UserPlatformQuotaCacheEntry{ + makeEntry(1.0, 2.0, 3.0), + makeEntry(4.0, 5.0, 6.0), + }, + } + writeErr := errors.New("db connection timeout") + writer := &mockQuotaSnapshotWriter{returnErr: writeErr} + f := newTestFlusher(cache, writer) + + f.flush() + + if f.metrics.FlushErrorTotal.Load() != 1 { + t.Errorf("FlushErrorTotal = %d, want 1", f.metrics.FlushErrorTotal.Load()) + } + if len(cache.readdCalled) == 0 { + t.Fatal("Readd should be called after write error") + } + totalReadd := 0 + for _, rk := range cache.readdCalled { + totalReadd += len(rk) + } + if totalReadd != len(keys) { + t.Errorf("DirtyReaddTotal (from Readd calls) = %d, want %d", totalReadd, len(keys)) + } + if f.metrics.DirtyReaddTotal.Load() != int64(len(keys)) { + t.Errorf("DirtyReaddTotal metric = %d, want %d", f.metrics.DirtyReaddTotal.Load(), len(keys)) + } + if f.metrics.FlushSuccessTotal.Load() != 0 { + t.Errorf("FlushSuccessTotal = %d, want 0", f.metrics.FlushSuccessTotal.Load()) + } +} + +// --------------------------------------------------------------------------- +// 场景 4: FKViolationDropsNoReadd — writer 返 ErrUserPlatformQuotaFKViolation → 不 Readd，FlushFKViolationTotal=1 +// --------------------------------------------------------------------------- + +func TestFlusher_FKViolationDropsNoReadd(t *testing.T) { + keys := []UserPlatformQuotaKey{ + {UserID: 999, Platform: "anthropic"}, + } + cache := &mockQuotaDirtyCache{ + popSequence: [][]UserPlatformQuotaKey{keys}, + getEntries: []*UserPlatformQuotaCacheEntry{ + makeEntry(1.0, 2.0, 3.0), + }, + } + writer := &mockQuotaSnapshotWriter{returnErr: ErrUserPlatformQuotaFKViolation} + f := newTestFlusher(cache, writer) + + f.flush() + + if f.metrics.FlushFKViolationTotal.Load() != 1 { + t.Errorf("FlushFKViolationTotal = %d, want 1", f.metrics.FlushFKViolationTotal.Load()) + } + if f.metrics.FlushErrorTotal.Load() != 1 { + t.Errorf("FlushErrorTotal = %d, want 1", f.metrics.FlushErrorTotal.Load()) + } + if len(cache.readdCalled) != 0 { + t.Errorf("Readd should NOT be called for FK violation (drop), got %d calls", len(cache.readdCalled)) + } + if f.metrics.DirtyReaddTotal.Load() != 0 { + t.Errorf("DirtyReaddTotal = %d, want 0 (FK violation drops)", f.metrics.DirtyReaddTotal.Load()) + } +} + +// --------------------------------------------------------------------------- +// 场景 5: NilSafe — var f *UserPlatformQuotaUsageFlusher; f.flush(); f.Stop() 不 panic +// --------------------------------------------------------------------------- + +func TestFlusher_NilSafe(t *testing.T) { + var f *UserPlatformQuotaUsageFlusher + // 下面两行不应 panic + f.flush() + f.Stop() +} + +// --------------------------------------------------------------------------- +// 场景 6: StopPreventsFlush — stopped=true 后 tick() 不调 flush（writer 没收到 snaps） +// --------------------------------------------------------------------------- + +func TestFlusher_StopPreventsFlush(t *testing.T) { + keys := []UserPlatformQuotaKey{ + {UserID: 1, Platform: "anthropic"}, + } + cache := &mockQuotaDirtyCache{ + popSequence: [][]UserPlatformQuotaKey{keys}, + getEntries: []*UserPlatformQuotaCacheEntry{ + makeEntry(1.0, 2.0, 3.0), + }, + } + writer := &mockQuotaSnapshotWriter{} + f := newTestFlusher(cache, writer) + + // 标记为已停止 + f.stopped.Store(true) + + // tick 应该直接返回，不触发 flush + f.tick() + + if len(writer.receivedSnaps) != 0 { + t.Errorf("expected 0 snaps after stop, got %d", len(writer.receivedSnaps)) + } + if cache.popCallIdx != 0 { + t.Errorf("Pop should not be called after stop, popCallIdx = %d", cache.popCallIdx) + } +} + +// --------------------------------------------------------------------------- +// 场景 B13-1: ZeroPercentCompany — 0% 公司脏集恒空，flusher 空跑无 DB 写 +// +// 模拟几乎没有用户配置 quota limit 的公司：脏集始终为空（popSequence 为空切片）， +// Pop 每次返回空集。flush() 应早退，不写 DB、不计成功、不 Readd。 +// --------------------------------------------------------------------------- + +func TestScenario_ZeroPercentCompany(t *testing.T) { + cache := &mockQuotaDirtyCache{ + // popSequence 为空 → Pop 超出序列 → 始终返回 nil（空集） + popSequence: [][]UserPlatformQuotaKey{}, + } + writer := &mockQuotaSnapshotWriter{} + f := newTestFlusher(cache, writer) + + f.flush() + + if len(writer.receivedSnaps) != 0 { + t.Errorf("0%% company: expected 0 snaps, got %d", len(writer.receivedSnaps)) + } + if f.metrics.FlushBatchSizeTotal.Load() != 0 { + t.Errorf("0%% company: FlushBatchSizeTotal = %d, want 0", f.metrics.FlushBatchSizeTotal.Load()) + } + if f.metrics.FlushSuccessTotal.Load() != 0 { + t.Errorf("0%% company: FlushSuccessTotal = %d, want 0 (empty-set early return)", f.metrics.FlushSuccessTotal.Load()) + } + if f.metrics.FlushErrorTotal.Load() != 0 { + t.Errorf("0%% company: FlushErrorTotal = %d, want 0", f.metrics.FlushErrorTotal.Load()) + } + if len(cache.readdCalled) != 0 { + t.Errorf("0%% company: Readd should never be called, got %d calls", len(cache.readdCalled)) + } +} + +// --------------------------------------------------------------------------- +// P1: IntervalFallback — flush_interval_ms ≤0 时回退 2s；正常值保留 +// --------------------------------------------------------------------------- + +func TestNewUserPlatformQuotaUsageFlusher_IntervalFallback(t *testing.T) { + cases := []struct { + name string + inMs int + wantDu time.Duration + }{ + {"零值回退 2s", 0, 2 * time.Second}, + {"负数回退 2s", -100, 2 * time.Second}, + {"正常 2000ms 保留", 2000, 2 * time.Second}, + {"正常 500ms 保留", 500, 500 * time.Millisecond}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + cfg := &config.Config{} + cfg.Database.UserPlatformQuotaFlushIntervalMs = tc.inMs + f := NewUserPlatformQuotaUsageFlusher(cfg, nil, nil, nil) + if f.interval != tc.wantDu { + t.Fatalf("interval = %v, want %v", f.interval, tc.wantDu) + } + }) + } +} + +// --------------------------------------------------------------------------- +// P1: EnabledField — flusher_enabled 配置正确写入 f.enabled +// --------------------------------------------------------------------------- + +func TestNewUserPlatformQuotaUsageFlusher_EnabledField(t *testing.T) { + for _, enabled := range []bool{true, false} { + cfg := &config.Config{} + cfg.Database.UserPlatformQuotaFlusherEnabled = enabled + cfg.Database.UserPlatformQuotaFlushIntervalMs = 500 + f := NewUserPlatformQuotaUsageFlusher(cfg, nil, nil, nil) + if f.enabled != enabled { + t.Errorf("enabled = %v, want %v", f.enabled, enabled) + } + } +} + +// --------------------------------------------------------------------------- +// P2: ReaddFailCounts — BatchGet 失败 + Readd 失败 → DirtyLostTotal 增、DirtyReaddTotal 不变 +// BatchGet 失败 + Readd 成功 → DirtyReaddTotal 增、DirtyLostTotal 不变 +// --------------------------------------------------------------------------- + +func TestFlusher_ReaddFailCounts(t *testing.T) { + keys := []UserPlatformQuotaKey{ + {UserID: 10, Platform: "anthropic"}, + {UserID: 11, Platform: "openai"}, + } + + t.Run("Readd 失败计 DirtyLostTotal", func(t *testing.T) { + cache := &mockQuotaDirtyCache{ + popSequence: [][]UserPlatformQuotaKey{keys}, + getErr: errors.New("redis timeout"), // 触发 BatchGet 失败路径 + readdErr: errors.New("redis connection refused"), // Readd 也失败 + } + f := newTestFlusher(cache, &mockQuotaSnapshotWriter{}) + + f.flush() + + if f.metrics.DirtyLostTotal.Load() != int64(len(keys)) { + t.Errorf("DirtyLostTotal = %d, want %d", f.metrics.DirtyLostTotal.Load(), len(keys)) + } + if f.metrics.DirtyReaddTotal.Load() != 0 { + t.Errorf("DirtyReaddTotal = %d, want 0 (Readd 失败不应计入)", f.metrics.DirtyReaddTotal.Load()) + } + }) + + t.Run("Readd 成功计 DirtyReaddTotal", func(t *testing.T) { + cache := &mockQuotaDirtyCache{ + popSequence: [][]UserPlatformQuotaKey{keys}, + getErr: errors.New("redis timeout"), // 触发 BatchGet 失败路径 + readdErr: nil, // Readd 成功 + } + f := newTestFlusher(cache, &mockQuotaSnapshotWriter{}) + + f.flush() + + if f.metrics.DirtyReaddTotal.Load() != int64(len(keys)) { + t.Errorf("DirtyReaddTotal = %d, want %d", f.metrics.DirtyReaddTotal.Load(), len(keys)) + } + if f.metrics.DirtyLostTotal.Load() != 0 { + t.Errorf("DirtyLostTotal = %d, want 0 (Readd 成功不应计 lost)", f.metrics.DirtyLostTotal.Load()) + } + }) +} + +// --------------------------------------------------------------------------- +// ClampsBatchSize — NewUserPlatformQuotaUsageFlusher 构造时按 +// [defaultFlushBatchSize, maxFlushBatchSize] 区间 clamp batchSize +// --------------------------------------------------------------------------- + +func TestNewUserPlatformQuotaUsageFlusher_ClampsBatchSize(t *testing.T) { + cases := []struct { + name string + in int + want int + }{ + {"超上限被 clamp", 7000, maxFlushBatchSize}, + {"恰好上限保留", maxFlushBatchSize, maxFlushBatchSize}, + {"零回退默认", 0, defaultFlushBatchSize}, + {"负数回退默认", -5, defaultFlushBatchSize}, + {"正常值保留", 500, 500}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + cfg := &config.Config{} + cfg.Database.UserPlatformQuotaFlushBatchSize = tc.in + f := NewUserPlatformQuotaUsageFlusher(cfg, nil, nil, nil) + if f.batchSize != tc.want { + t.Fatalf("batchSize = %d, want %d", f.batchSize, tc.want) + } + }) + } +} + +// --------------------------------------------------------------------------- +// 场景 B13-2: NinetyPercentCompany — 90% 公司大量用户配 limit，一批 5 key 批量刷库 +// +// 模拟大量用户配置了 quota limit 的公司：脏集第一次 Pop 返回 5 个不同用户的 key， +// 之后返回空集（避免 flush 循环）。flush() 应构造 5 条 snapshot 写入 DB， +// 断言绝对值语义（snap 的 DailyUsageUSD 等于 entry 的值）、metrics 正确、不 Readd。 +// --------------------------------------------------------------------------- + +func TestScenario_NinetyPercentCompany(t *testing.T) { + keys := []UserPlatformQuotaKey{ + {UserID: 101, Platform: "anthropic"}, + {UserID: 102, Platform: "anthropic"}, + {UserID: 103, Platform: "openai"}, + {UserID: 104, Platform: "openai"}, + {UserID: 105, Platform: "anthropic"}, + } + entries := []*UserPlatformQuotaCacheEntry{ + makeEntry(1.1, 2.2, 3.3), + makeEntry(4.4, 5.5, 6.6), + makeEntry(7.7, 8.8, 9.9), + makeEntry(0.5, 1.0, 1.5), + makeEntry(10.0, 20.0, 30.0), + } + cache := &mockQuotaDirtyCache{ + // 第 1 次 Pop 返回 5 keys，之后返回空集（防止 flush 无限循环） + popSequence: [][]UserPlatformQuotaKey{keys}, + getEntries: entries, + } + writer := &mockQuotaSnapshotWriter{} + f := newTestFlusher(cache, writer) + + f.flush() + + // 应收到 5 条 snapshot + if len(writer.receivedSnaps) != 5 { + t.Fatalf("90%% company: expected 5 snaps, got %d", len(writer.receivedSnaps)) + } + + // 验证绝对值语义：第 1 条 snap 的各窗口 usage 应等于 entries[0] 的值 + snap0 := writer.receivedSnaps[0] + entry0 := entries[0] + if snap0.DailyUsageUSD != entry0.DailyUsageUSD { + t.Errorf("snap[0].DailyUsageUSD = %v, want %v", snap0.DailyUsageUSD, entry0.DailyUsageUSD) + } + if snap0.WeeklyUsageUSD != entry0.WeeklyUsageUSD { + t.Errorf("snap[0].WeeklyUsageUSD = %v, want %v", snap0.WeeklyUsageUSD, entry0.WeeklyUsageUSD) + } + if snap0.MonthlyUsageUSD != entry0.MonthlyUsageUSD { + t.Errorf("snap[0].MonthlyUsageUSD = %v, want %v", snap0.MonthlyUsageUSD, entry0.MonthlyUsageUSD) + } + + // FlushBatchSizeTotal 应为 5（本批 keys 数量） + if f.metrics.FlushBatchSizeTotal.Load() != 5 { + t.Errorf("90%% company: FlushBatchSizeTotal = %d, want 5", f.metrics.FlushBatchSizeTotal.Load()) + } + // FlushSuccessTotal 应为 1（1 个批次写成功） + if f.metrics.FlushSuccessTotal.Load() != 1 { + t.Errorf("90%% company: FlushSuccessTotal = %d, want 1", f.metrics.FlushSuccessTotal.Load()) + } + // 无错误、无 Readd + if f.metrics.FlushErrorTotal.Load() != 0 { + t.Errorf("90%% company: FlushErrorTotal = %d, want 0", f.metrics.FlushErrorTotal.Load()) + } + if f.metrics.DirtyReaddTotal.Load() != 0 { + t.Errorf("90%% company: DirtyReaddTotal = %d, want 0", f.metrics.DirtyReaddTotal.Load()) + } + if len(cache.readdCalled) != 0 { + t.Errorf("90%% company: Readd should not be called, got %d calls", len(cache.readdCalled)) + } +} diff --git a/backend/internal/service/user_platform_quota_port.go b/backend/internal/service/user_platform_quota_port.go index cb09542a421f..0f88eda4c3b7 100644 --- a/backend/internal/service/user_platform_quota_port.go +++ b/backend/internal/service/user_platform_quota_port.go @@ -11,6 +11,23 @@ import ( // handler 只需引用 service 包，无需直接依赖 repository 包。 var ErrUserPlatformQuotaNotFound = errors.New("user platform quota not found") +// ErrUserPlatformQuotaFKViolation service 层 sentinel：批量 snapshot UPSERT 时存在 +// user_id 不在 users 表的记录（外键违反）。adapter 负责将 repository 层同名 sentinel 包装为此错误。 +var ErrUserPlatformQuotaFKViolation = errors.New("user platform quota snapshot FK violation") + +// UserPlatformQuotaSnapshot 是 service 层 flusher 向 DB 写入快照时使用的传输结构。 +// 字段语义与 repository.UserPlatformQuotaSnapshot 完全对应，由 adapter 负责转换。 +type UserPlatformQuotaSnapshot struct { + UserID int64 + Platform string + DailyUsageUSD float64 + WeeklyUsageUSD float64 + MonthlyUsageUSD float64 + DailyWindowStart time.Time + WeeklyWindowStart time.Time + MonthlyWindowStart time.Time +} + // UserPlatformQuotaRecord service 层传输结构体（与 repository 层解耦）。 type UserPlatformQuotaRecord struct { UserID int64 @@ -47,4 +64,6 @@ type UserPlatformQuotaRepository interface { // ResetExpiredWindow 重置指定窗口（"daily"|"weekly"|"monthly"）的用量与起始时间。 // 未命中活跃记录时返回（service-side wrapper of repository.ErrUserPlatformQuotaNotFound）。 ResetExpiredWindow(ctx context.Context, userID int64, platform string, window string, newStart time.Time) error + // BatchSnapshotUsage 绝对值覆盖写入整批 usage 快照。FK 违反返回 ErrUserPlatformQuotaFKViolation。 + BatchSnapshotUsage(ctx context.Context, snapshots []UserPlatformQuotaSnapshot, now time.Time) error } diff --git a/backend/internal/service/user_service_test.go b/backend/internal/service/user_service_test.go index 19aec5d312be..1a18e70ac04f 100644 --- a/backend/internal/service/user_service_test.go +++ b/backend/internal/service/user_service_test.go @@ -327,10 +327,22 @@ func (m *mockBillingCache) DeleteUserPlatformQuotaCache(context.Context, int64, return nil } -func (m *mockBillingCache) IncrUserPlatformQuotaUsageCache(context.Context, int64, string, float64, time.Duration) error { +func (m *mockBillingCache) IncrUserPlatformQuotaUsageCache(context.Context, int64, string, float64, time.Duration, bool) error { return nil } +func (m *mockBillingCache) PopDirtyUserPlatformQuotaKeys(context.Context, int) ([]UserPlatformQuotaKey, error) { + return nil, nil +} + +func (m *mockBillingCache) ReaddDirtyUserPlatformQuotaKeys(context.Context, []UserPlatformQuotaKey) error { + return nil +} + +func (m *mockBillingCache) BatchGetUserPlatformQuotaCache(context.Context, []UserPlatformQuotaKey) ([]*UserPlatformQuotaCacheEntry, error) { + return nil, nil +} + // --- 测试 --- func TestUpdateBalance_Success(t *testing.T) { diff --git a/backend/internal/service/wire.go b/backend/internal/service/wire.go index d3e4ce51bc7e..19bd841d4d24 100644 --- a/backend/internal/service/wire.go +++ b/backend/internal/service/wire.go @@ -571,8 +571,16 @@ var ProviderSet = wire.NewSet( ProvideChannelMonitorService, ProvideChannelMonitorRunner, NewChannelMonitorRequestTemplateService, + ProvideUserPlatformQuotaUsageFlusher, ) +// ProvideUserPlatformQuotaUsageFlusher 创建并启动 UserPlatformQuotaUsageFlusher。 +func ProvideUserPlatformQuotaUsageFlusher(cfg *config.Config, cache BillingCache, quotaRepo UserPlatformQuotaRepository, tw *TimingWheelService) *UserPlatformQuotaUsageFlusher { + svc := NewUserPlatformQuotaUsageFlusher(cfg, cache, quotaRepo, tw) + svc.Start() + return svc +} + // ProvidePaymentConfigService wraps NewPaymentConfigService to accept the named // payment.EncryptionKey type instead of raw []byte, avoiding Wire ambiguity. func ProvidePaymentConfigService(entClient *dbent.Client, settingRepo SettingRepository, key payment.EncryptionKey) *PaymentConfigService { From 5fd9a35093dffe8abc3b72ed6b53b5653b50ee98 Mon Sep 17 00:00:00 2001 From: DaydreamCoding <22166516+DaydreamCoding@users.noreply.github.com> Date: Fri, 29 May 2026 18:47:27 +0800 Subject: [PATCH 12/36] =?UTF-8?q?test(pricing):=20=E4=BF=AE=E5=A4=8D=20cod?= =?UTF-8?q?ex-auto-review=20=E5=AE=9A=E4=BB=B7=E6=B5=8B=E8=AF=95=E6=96=AD?= =?UTF-8?q?=E8=A8=80?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 68901cbff 批量同步定价数据后，codex-auto-review 的 input_cost_per_token 从 2.5e-6 更新为 5e-6，output 从 1.5e-5 更新为 3e-5，cache_read 从 2.5e-7 更新为 5e-7。测试断言需要同步更新以匹配当前定价数据。 Co-Authored-By: Claude Opus 4.8 (1M context) --- backend/internal/service/pricing_service_test.go | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/backend/internal/service/pricing_service_test.go b/backend/internal/service/pricing_service_test.go index cc8b120a1aa7..f4252f954048 100644 --- a/backend/internal/service/pricing_service_test.go +++ b/backend/internal/service/pricing_service_test.go @@ -124,9 +124,9 @@ func TestDefaultPricingIncludesCodexAutoReview(t *testing.T) { got := svc.GetModelPricing("codex-auto-review") require.NotNil(t, got) - require.InDelta(t, 2.5e-6, got.InputCostPerToken, 1e-12) - require.InDelta(t, 1.5e-5, got.OutputCostPerToken, 1e-12) - require.InDelta(t, 2.5e-7, got.CacheReadInputTokenCost, 1e-12) + require.InDelta(t, 5e-6, got.InputCostPerToken, 1e-12) + require.InDelta(t, 3e-5, got.OutputCostPerToken, 1e-12) + require.InDelta(t, 5e-7, got.CacheReadInputTokenCost, 1e-12) } func TestGetModelPricing_Gpt54MiniUsesDedicatedStaticFallbackWhenRemoteMissing(t *testing.T) { From c256a5441a80a5aade3e54c255867c17db8bfa8c Mon Sep 17 00:00:00 2001 From: wucm667 Date: Fri, 29 May 2026 20:57:29 +0800 Subject: [PATCH 13/36] =?UTF-8?q?feat(admin):=20=E8=B4=A6=E5=8F=B7?= =?UTF-8?q?=E7=94=A8=E9=87=8F=E7=AA=97=E5=8F=A3=205h/7d=20=E5=A2=9E?= =?UTF-8?q?=E5=8A=A0=E8=AF=B4=E6=98=8E=20tooltip?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 在账号管理列表"用量窗口"列表头增加一个说明性 HelpTooltip， 解释 5h / 7d 是上游账号（如 OpenAI ChatGPT、Claude）官方的滚动 用量窗口限制，由上游设定、非 sub2api 配置、与映射模型无关，且 窗口滚动到期后自动重置、无法在 sub2api 端解除。 复用现有 HelpTooltip 组件（teleport 到 body，避免表格裁剪）， 单个 ⓘ 图标置于列表头，避免每行重复。新增 i18n key admin.accounts.usageWindowsHint（zh/en 同步）。纯展示说明， 不改用量计算与后端逻辑。 Co-Authored-By: Claude Opus 4.8 --- frontend/src/i18n/locales/en.ts | 1 + frontend/src/i18n/locales/zh.ts | 1 + frontend/src/views/admin/AccountsView.vue | 7 + .../AccountsView.usageWindowsHint.spec.ts | 164 ++++++++++++++++++ 4 files changed, 173 insertions(+) create mode 100644 frontend/src/views/admin/__tests__/AccountsView.usageWindowsHint.spec.ts diff --git a/frontend/src/i18n/locales/en.ts b/frontend/src/i18n/locales/en.ts index 6735029c7c3a..2030ddf0510f 100644 --- a/frontend/src/i18n/locales/en.ts +++ b/frontend/src/i18n/locales/en.ts @@ -3105,6 +3105,7 @@ export default { expiresAt: 'Expires At', actions: 'Actions' }, + usageWindowsHint: '"5h / 7d" are the upstream account\'s official rolling usage windows (e.g. OpenAI ChatGPT, Claude). They are imposed by the upstream provider on the account itself — not configured by sub2api, and unrelated to the models you map. Usage resets automatically once each window rolls over, and the limit cannot be lifted from within sub2api.', allPrivacyModes: 'All Privacy States', privacyUnset: 'Unset', privacyTrainingOff: 'Training data sharing disabled', diff --git a/frontend/src/i18n/locales/zh.ts b/frontend/src/i18n/locales/zh.ts index abb8dff730e7..6106aa5a501c 100644 --- a/frontend/src/i18n/locales/zh.ts +++ b/frontend/src/i18n/locales/zh.ts @@ -3143,6 +3143,7 @@ export default { expiresAt: '过期时间', actions: '操作' }, + usageWindowsHint: '“5h / 7d”是上游账号（如 OpenAI ChatGPT、Claude）官方的滚动用量窗口限制，由上游对账号设定，并非 sub2api 配置，也与你映射的模型无关。窗口滚动到期后用量会自动重置，无法在 sub2api 端解除该限制。', allPrivacyModes: '全部Privacy状态', privacyUnset: '未设置', privacyTrainingOff: '已关闭训练数据共享', diff --git a/frontend/src/views/admin/AccountsView.vue b/frontend/src/views/admin/AccountsView.vue index c602225c2a6b..04b46a8de04d 100644 --- a/frontend/src/views/admin/AccountsView.vue +++ b/frontend/src/views/admin/AccountsView.vue @@ -273,6 +273,12 @@ + diff --git a/frontend/src/components/admin/usage/__tests__/UsageFilters.spec.ts b/frontend/src/components/admin/usage/__tests__/UsageFilters.spec.ts new file mode 100644 index 000000000000..6f1780d18b73 --- /dev/null +++ b/frontend/src/components/admin/usage/__tests__/UsageFilters.spec.ts @@ -0,0 +1,168 @@ +import { describe, expect, it, vi, beforeEach, afterEach } from 'vitest' +import { mount, flushPromises } from '@vue/test-utils' + +import UsageFilters from '../UsageFilters.vue' + +// --- i18n messages (only what UsageFilters needs) --- +const messages: Record = { + 'admin.usage.userDeletedBadge': 'deleted', + 'admin.usage.userFilter': 'User', + 'admin.usage.searchUserPlaceholder': 'Search user...', + 'usage.apiKeyFilter': 'API Key', + 'admin.usage.searchApiKeyPlaceholder': 'Search API key...', + 'usage.model': 'Model', + 'admin.usage.allModels': 'All Models', + 'admin.usage.account': 'Account', + 'admin.usage.searchAccountPlaceholder': 'Search account...', + 'usage.type': 'Type', + 'admin.usage.allTypes': 'All Types', + 'usage.ws': 'WS', + 'usage.stream': 'Stream', + 'usage.sync': 'Sync', + 'admin.usage.billingType': 'Billing Type', + 'admin.usage.allBillingTypes': 'All Billing Types', + 'admin.usage.billingTypeBalance': 'Balance', + 'admin.usage.billingTypeSubscription': 'Subscription', + 'admin.usage.billingMode': 'Billing Mode', + 'admin.usage.allBillingModes': 'All Billing Modes', + 'admin.usage.billingModeToken': 'Token', + 'admin.usage.billingModePerRequest': 'Per Request', + 'admin.usage.billingModeImage': 'Image', + 'admin.usage.group': 'Group', + 'admin.usage.allGroups': 'All Groups', + 'common.refresh': 'Refresh', + 'common.reset': 'Reset', + 'admin.usage.cleanup.button': 'Cleanup', + 'usage.exportExcel': 'Export', +} + +// Mock vue-i18n +vi.mock('vue-i18n', async () => { + const actual = await vi.importActual('vue-i18n') + return { + ...actual, + useI18n: () => ({ + t: (key: string) => messages[key] ?? key, + }), + } +}) + +// Mock the admin API module — we control searchUsers return value per test +const mockSearchUsers = vi.fn() +const mockSearchApiKeys = vi.fn().mockResolvedValue([]) + +vi.mock('@/api/admin', () => ({ + adminAPI: { + usage: { + searchUsers: (...args: any[]) => mockSearchUsers(...args), + searchApiKeys: (...args: any[]) => mockSearchApiKeys(...args), + }, + groups: { + list: vi.fn().mockResolvedValue({ items: [] }), + }, + dashboard: { + getModelStats: vi.fn().mockResolvedValue({ models: [] }), + }, + accounts: { + list: vi.fn().mockResolvedValue({ items: [] }), + }, + }, +})) + +// Default props helper +const defaultFilters = () => ({ + user_id: undefined, + api_key_id: undefined, + account_id: undefined, + model: null, + request_type: null, + billing_type: null, + billing_mode: null, + group_id: null, + start_date: '', + end_date: '', +}) + +function mountFilters(filters = defaultFilters()) { + return mount(UsageFilters, { + props: { + modelValue: filters, + exporting: false, + startDate: '2026-05-01', + endDate: '2026-05-28', + showActions: false, + }, + global: { + stubs: { + Select: true, + Teleport: true, + }, + }, + }) +} + +describe('UsageFilters — user search dropdown', () => { + beforeEach(() => { + vi.useFakeTimers() + mockSearchUsers.mockReset() + mockSearchApiKeys.mockResolvedValue([]) + }) + + afterEach(() => { + vi.useRealTimers() + }) + + it('(a) labels deleted users with the i18n badge and (b) sorts active users before deleted ones, (c) selection sets user_id', async () => { + // Arrange: mock returns deleted FIRST (proves sorting re-orders to active-first) + mockSearchUsers.mockResolvedValue([ + { id: 2, email: 'gone@test.com', deleted: true }, + { id: 1, email: 'active@test.com', deleted: false }, + ]) + + const wrapper = mountFilters() + + // Trigger focus (sets showUserDropdown = true) then input (fires debounceUserSearch) + const input = wrapper.find('input[type="text"]') + await input.trigger('focus') + await input.setValue('test') + await input.trigger('input') + + // Advance debounce timer (300ms) then flush the resolved promise + vi.advanceTimersByTime(300) + await flushPromises() + + // --- (b) Sort: active user should appear BEFORE deleted user --- + // Check the underlying component state via rendered DOM order + const buttons = wrapper.findAll('.usage-filter-dropdown button[type="button"]') + const emailTexts = buttons.map((b) => b.text()) + + // active@test.com should be listed first + const activeIdx = emailTexts.findIndex((t) => t.includes('active@test.com')) + const deletedIdx = emailTexts.findIndex((t) => t.includes('gone@test.com')) + expect(activeIdx).toBeGreaterThanOrEqual(0) + expect(deletedIdx).toBeGreaterThanOrEqual(0) + expect(activeIdx).toBeLessThan(deletedIdx) + + // --- (a) Label: deleted user's button shows the badge text --- + const deletedButton = buttons[deletedIdx] + expect(deletedButton.text()).toContain('deleted') + + // active user's button does NOT show the badge text + const activeButton = buttons[activeIdx] + expect(activeButton.text()).not.toContain('deleted') + + // --- (c) Selection: clicking active user button sets filters.user_id --- + await activeButton.trigger('click') + await flushPromises() + + // The component emits 'update:modelValue' or modifies filters.user_id via toRef + // selectUser sets filters.value.user_id = u.id and emits 'change' + const changeEmits = wrapper.emitted('change') + expect(changeEmits).toBeTruthy() + expect(changeEmits!.length).toBeGreaterThan(0) + + // Also confirm user_id was set by checking the emitted change came through + // (the component uses toRef so modelValue is mutated in place and 'change' is emitted) + expect(wrapper.props('modelValue').user_id).toBe(1) + }) +}) diff --git a/frontend/src/components/admin/usage/__tests__/UsageTable.spec.ts b/frontend/src/components/admin/usage/__tests__/UsageTable.spec.ts index ece0dbda6e41..7856d19a6773 100644 --- a/frontend/src/components/admin/usage/__tests__/UsageTable.spec.ts +++ b/frontend/src/components/admin/usage/__tests__/UsageTable.spec.ts @@ -5,6 +5,7 @@ import { nextTick } from 'vue' import UsageTable from '../UsageTable.vue' const messages: Record = { + 'admin.usage.userDeletedBadge': 'Deleted', 'usage.costDetails': 'Cost Breakdown', 'admin.usage.inputCost': 'Input Cost', 'admin.usage.outputCost': 'Output Cost', @@ -321,3 +322,91 @@ describe('admin UsageTable tooltip', () => { expect(text).not.toContain('(2K)') }) }) + +// A DataTable stub that also renders cell-user, so the deleted badge can be asserted. +const DataTableStubWithUser = { + props: ['data'], + template: ` +
+
+ + + + + +
+
+ `, +} + +describe('admin UsageTable deleted-user badge', () => { + it('renders deleted badge for a soft-deleted user row', () => { + const row = { + request_id: 'req-deleted-user-1', + model: 'claude-3', + user_id: 2, + user: { id: 2, email: 'd@test.com', deleted_at: '2026-05-28T00:00:00Z' }, + actual_cost: 0, + total_cost: 0, + input_cost: 0, + output_cost: 0, + rate_multiplier: 1, + input_tokens: 1, + output_tokens: 1, + } + + const wrapper = mount(UsageTable, { + props: { + data: [row], + loading: false, + columns: [{ key: 'user', label: 'User' }], + }, + global: { + stubs: { + DataTable: DataTableStubWithUser, + EmptyState: true, + Icon: true, + Teleport: true, + }, + }, + }) + + expect(wrapper.text()).toContain('Deleted') + expect(wrapper.text()).toContain('d@test.com') + }) + + it('does NOT render deleted badge for an active user row', () => { + const row = { + request_id: 'req-active-user-1', + model: 'claude-3', + user_id: 3, + user: { id: 3, email: 'active@test.com', deleted_at: null }, + actual_cost: 0, + total_cost: 0, + input_cost: 0, + output_cost: 0, + rate_multiplier: 1, + input_tokens: 1, + output_tokens: 1, + } + + const wrapper = mount(UsageTable, { + props: { + data: [row], + loading: false, + columns: [{ key: 'user', label: 'User' }], + }, + global: { + stubs: { + DataTable: DataTableStubWithUser, + EmptyState: true, + Icon: true, + Teleport: true, + }, + }, + }) + + expect(wrapper.text()).not.toContain('Deleted') + expect(wrapper.text()).toContain('active@test.com') + }) +}) diff --git a/frontend/src/components/admin/user/UserBalanceHistoryModal.vue b/frontend/src/components/admin/user/UserBalanceHistoryModal.vue index 6d48ed77420e..4bce43ad5379 100644 --- a/frontend/src/components/admin/user/UserBalanceHistoryModal.vue +++ b/frontend/src/components/admin/user/UserBalanceHistoryModal.vue @@ -13,6 +13,9 @@

{{ user.email }}

+ + {{ t('admin.usage.userDeletedBadge') }} + { const handleUserClick = async (userId: number) => { try { - const user = await adminAPI.users.getById(userId) + const user = await adminAPI.users.getById(userId, true) balanceHistoryUser.value = user showBalanceHistoryModal.value = true } catch { diff --git a/frontend/src/views/admin/__tests__/UsageView.spec.ts b/frontend/src/views/admin/__tests__/UsageView.spec.ts index 1a5d285a8603..9cdb2a3c8b9c 100644 --- a/frontend/src/views/admin/__tests__/UsageView.spec.ts +++ b/frontend/src/views/admin/__tests__/UsageView.spec.ts @@ -84,6 +84,10 @@ vi.mock('vue-router', () => ({ const AppLayoutStub = { template: '
' } const UsageFiltersStub = { template: '
' } +const UsageTableStub = { + emits: ['userClick'], + template: '
', +} const ModelDistributionChartStub = { props: ['metric'], emits: ['update:metric'], @@ -194,3 +198,59 @@ describe('admin UsageView distribution metric toggles', () => { expect(getSnapshotV2).toHaveBeenCalledTimes(1) }) }) + +describe('admin UsageView handleUserClick', () => { + beforeEach(() => { + vi.useFakeTimers() + list.mockReset() + getStats.mockReset() + getSnapshotV2.mockReset() + getById.mockReset() + + list.mockResolvedValue({ items: [], total: 0, pages: 0 }) + getStats.mockResolvedValue({ + total_requests: 0, total_input_tokens: 0, total_output_tokens: 0, + total_cache_tokens: 0, total_tokens: 0, total_cost: 0, total_actual_cost: 0, average_duration_ms: 0, + }) + getSnapshotV2.mockResolvedValue({ trend: [], models: [], groups: [] }) + }) + + afterEach(() => { + vi.useRealTimers() + }) + + it('opens user via include_deleted when clicking a usage row user', async () => { + getById.mockResolvedValue({ id: 2, email: 'd@test.com', deleted_at: '2026-05-28T00:00:00Z' }) + + const wrapper = mount(UsageView, { + global: { + stubs: { + AppLayout: AppLayoutStub, + UsageStatsCards: true, + UsageFilters: UsageFiltersStub, + UsageTable: UsageTableStub, + UsageExportProgress: true, + UsageCleanupDialog: true, + UserBalanceHistoryModal: true, + AuditLogModal: true, + Pagination: true, + Select: true, + DateRangePicker: true, + Icon: true, + TokenUsageTrend: true, + ModelDistributionChart: true, + GroupDistributionChart: true, + EndpointDistributionChart: true, + }, + }, + }) + + vi.advanceTimersByTime(120) + await flushPromises() + + await wrapper.find('[data-test="usage-table"] .user-click').trigger('click') + await flushPromises() + + expect(getById).toHaveBeenCalledWith(2, true) + }) +}) From bf24b611399fd43d34a8d1801deac30a7b41fe4e Mon Sep 17 00:00:00 2001 From: DaydreamCoding <22166516+DaydreamCoding@users.noreply.github.com> Date: Sat, 30 May 2026 18:02:13 +0800 Subject: [PATCH 16/36] =?UTF-8?q?perf(usage):=20=E4=BC=98=E5=8C=96=20/admi?= =?UTF-8?q?n/usage=20=E6=89=93=E5=BC=80=E9=80=9F=E5=BA=A6=E4=B8=8E?= =?UTF-8?q?=E5=88=B7=E6=96=B0=E5=93=8D=E5=BA=94?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 根因:页面 mount 并发 6 个请求,其中 5 个在原始 usage_logs 上 live 聚合, 且有 1 个重复 getModelStats。优化(不引入预聚合): 前端 - UsageView 经 :model-options 下传 model 列表,移除 UsageFilters 重复的 getModelStats(mount 请求 6→5,少一次 usage_logs 全表 GROUP BY model) - 刷新/换筛选保留旧模型数据(invalidateModelStatsCache 只失效标记不清空数据), 图表不再闪空,刷新期间页面保持可交互 后端 - GetStatsWithFilters 4 条聚合 errgroup 并行(仅 *sql.DB 连接池路径,ent.Tx 顺序回退以保事务内不并发),endpoint 明细 best-effort;抑制取消级联噪声日志 - /admin/usage/stats 复用 dashboard 的 newSnapshotCache 30s 处理器层缓存,按 filters+窗口为 key;前端手动刷新带 nocache=1 强制回源(刷新=最新) 注:自合并提交 4c8396c 迁移而来,仅取"列表打开速度与刷新响应"部分;原提交的 "审计查看弹窗大 body 渲染"改动(AuditLogModal / audit-log-format / 审计 i18n) 依赖尚未迁移的审计功能(d8389ade),本次已排除。 Co-Authored-By: Claude Opus 4.8 (1M context) --- .../internal/handler/admin/usage_handler.go | 22 ++++- .../handler/admin/usage_query_cache.go | 62 ++++++++++++ .../handler/admin/usage_query_cache_test.go | 28 ++++++ backend/internal/repository/usage_log_repo.go | 98 +++++++++++++------ .../usage_log_repo_stats_integration_test.go | 51 ++++++++++ frontend/src/api/admin/usage.ts | 1 + .../components/admin/usage/UsageFilters.vue | 27 ++--- .../usage/__tests__/UsageFilters.spec.ts | 45 +++++++-- frontend/src/views/admin/UsageView.vue | 28 ++++-- .../views/admin/__tests__/UsageView.spec.ts | 37 ++++++- 10 files changed, 325 insertions(+), 74 deletions(-) create mode 100644 backend/internal/handler/admin/usage_query_cache.go create mode 100644 backend/internal/handler/admin/usage_query_cache_test.go create mode 100644 backend/internal/repository/usage_log_repo_stats_integration_test.go diff --git a/backend/internal/handler/admin/usage_handler.go b/backend/internal/handler/admin/usage_handler.go index 2ded3c1d839d..11a4aeb8f601 100644 --- a/backend/internal/handler/admin/usage_handler.go +++ b/backend/internal/handler/admin/usage_handler.go @@ -325,10 +325,24 @@ func (h *UsageHandler) Stats(c *gin.Context) { EndTime: &endTime, } - stats, err := h.usageService.GetStatsWithFilters(c.Request.Context(), filters) - if err != nil { - response.ErrorFrom(c, err) - return + var stats *usagestats.UsageStats + // nocache: 绕过缓存直接回源,刷新者本人拿最新;不回写缓存(管理台"我刷新我自己拿最新"语义,非全局失效)。 + if parseBoolQueryWithDefault(c.Query("nocache"), false) { + s, err := h.usageService.GetStatsWithFilters(c.Request.Context(), filters) + if err != nil { + response.ErrorFrom(c, err) + return + } + stats = s + c.Header("X-Usage-Stats-Cache", "bypass") + } else { + s, hit, err := h.getStatsCached(c.Request.Context(), filters) + if err != nil { + response.ErrorFrom(c, err) + return + } + stats = s + c.Header("X-Usage-Stats-Cache", cacheStatusValue(hit)) } response.Success(c, stats) diff --git a/backend/internal/handler/admin/usage_query_cache.go b/backend/internal/handler/admin/usage_query_cache.go new file mode 100644 index 000000000000..b288a95ba445 --- /dev/null +++ b/backend/internal/handler/admin/usage_query_cache.go @@ -0,0 +1,62 @@ +package admin + +import ( + "context" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/usagestats" +) + +// 与 dashboard 查询缓存同款:30s TTL 进程内缓存,仅服务 /admin/usage/stats 读路径。 +var usageStatsCache = newSnapshotCache(30 * time.Second) + +type usageStatsCacheKeyData struct { + StartTime string `json:"start_time"` + EndTime string `json:"end_time"` + UserID int64 `json:"user_id"` + APIKeyID int64 `json:"api_key_id"` + AccountID int64 `json:"account_id"` + GroupID int64 `json:"group_id"` + Model string `json:"model"` + BillingMode string `json:"billing_mode"` + RequestType *int16 `json:"request_type"` + Stream *bool `json:"stream"` + BillingType *int8 `json:"billing_type"` +} + +func usageStatsCacheKey(filters usagestats.UsageLogFilters) string { + start := "" + if filters.StartTime != nil { + start = filters.StartTime.UTC().Format(time.RFC3339) + } + end := "" + if filters.EndTime != nil { + end = filters.EndTime.UTC().Format(time.RFC3339) + } + return mustMarshalDashboardCacheKey(usageStatsCacheKeyData{ + StartTime: start, + EndTime: end, + UserID: filters.UserID, + APIKeyID: filters.APIKeyID, + AccountID: filters.AccountID, + GroupID: filters.GroupID, + Model: filters.Model, + BillingMode: filters.BillingMode, + RequestType: filters.RequestType, + Stream: filters.Stream, + BillingType: filters.BillingType, + }) +} + +// getStatsCached 命中则返回缓存,未命中则回源 usageService 并写缓存。 +func (h *UsageHandler) getStatsCached(ctx context.Context, filters usagestats.UsageLogFilters) (*usagestats.UsageStats, bool, error) { + key := usageStatsCacheKey(filters) + entry, hit, err := usageStatsCache.GetOrLoad(key, func() (any, error) { + return h.usageService.GetStatsWithFilters(ctx, filters) + }) + if err != nil { + return nil, hit, err + } + stats, err := snapshotPayloadAs[*usagestats.UsageStats](entry.Payload) + return stats, hit, err +} diff --git a/backend/internal/handler/admin/usage_query_cache_test.go b/backend/internal/handler/admin/usage_query_cache_test.go new file mode 100644 index 000000000000..857e507a58dc --- /dev/null +++ b/backend/internal/handler/admin/usage_query_cache_test.go @@ -0,0 +1,28 @@ +package admin + +import ( + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/usagestats" + "github.com/stretchr/testify/require" +) + +func TestUsageStatsCacheKey_StableAndDistinct(t *testing.T) { + start := time.Date(2026, 5, 29, 0, 0, 0, 0, time.UTC) + end := time.Date(2026, 5, 31, 0, 0, 0, 0, time.UTC) + base := usagestats.UsageLogFilters{StartTime: &start, EndTime: &end, Model: "claude-3"} + + k1 := usageStatsCacheKey(base) + k2 := usageStatsCacheKey(base) + require.NotEmpty(t, k1) + require.Equal(t, k1, k2, "same filters must produce same key") + + other := base + other.Model = "gpt-4o" + require.NotEqual(t, k1, usageStatsCacheKey(other), "different model must change key") + + withUser := base + withUser.UserID = 7 + require.NotEqual(t, k1, usageStatsCacheKey(withUser), "different user must change key") +} diff --git a/backend/internal/repository/usage_log_repo.go b/backend/internal/repository/usage_log_repo.go index 1835ebe511eb..b0992dae107f 100644 --- a/backend/internal/repository/usage_log_repo.go +++ b/backend/internal/repository/usage_log_repo.go @@ -27,6 +27,7 @@ import ( "github.com/Wei-Shaw/sub2api/internal/service" "github.com/lib/pq" gocache "github.com/patrickmn/go-cache" + "golang.org/x/sync/errgroup" ) const usageLogSelectColumns = "id, user_id, api_key_id, account_id, request_id, model, requested_model, upstream_model, group_id, subscription_id, input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, cache_creation_5m_tokens, cache_creation_1h_tokens, image_output_tokens, image_output_cost, input_cost, output_cost, cache_creation_cost, cache_read_cost, total_cost, actual_cost, rate_multiplier, account_rate_multiplier, billing_type, request_type, stream, openai_ws_mode, duration_ms, first_token_ms, user_agent, ip_address, image_count, image_size, image_input_size, image_output_size, image_size_source, image_size_breakdown, service_tier, reasoning_effort, inbound_endpoint, upstream_endpoint, cache_ttl_overridden, channel_id, model_mapping_chain, billing_tier, billing_mode, account_stats_cost, created_at" @@ -3538,24 +3539,6 @@ func (r *usageLogRepository) GetStatsWithFilters(ctx context.Context, filters Us stats := &UsageStats{} var totalAccountCost float64 - if err := scanSingleRow( - ctx, - r.sql, - query, - args, - &stats.TotalRequests, - &stats.TotalInputTokens, - &stats.TotalOutputTokens, - &stats.TotalCacheTokens, - &stats.TotalCost, - &stats.TotalActualCost, - &totalAccountCost, - &stats.AverageDurationMs, - ); err != nil { - return nil, err - } - stats.TotalAccountCost = &totalAccountCost - stats.TotalTokens = stats.TotalInputTokens + stats.TotalOutputTokens + stats.TotalCacheTokens start := time.Unix(0, 0).UTC() if filters.StartTime != nil { @@ -3566,21 +3549,76 @@ func (r *usageLogRepository) GetStatsWithFilters(ctx context.Context, filters Us end = *filters.EndTime } - endpoints, endpointErr := r.GetEndpointStatsWithFilters(ctx, start, end, filters.UserID, filters.APIKeyID, filters.AccountID, filters.GroupID, filters.Model, filters.RequestType, filters.Stream, filters.BillingType) - if endpointErr != nil { - logger.LegacyPrintf("repository.usage_log", "GetEndpointStatsWithFilters failed in GetStatsWithFilters: %v", endpointErr) - endpoints = []EndpointStat{} + var endpoints, upstreamEndpoints, endpointPaths []EndpointStat + + // 汇总查询:失败即致命。 + runSummary := func(c context.Context) error { + return scanSingleRow( + c, r.sql, query, args, + &stats.TotalRequests, + &stats.TotalInputTokens, + &stats.TotalOutputTokens, + &stats.TotalCacheTokens, + &stats.TotalCost, + &stats.TotalActualCost, + &totalAccountCost, + &stats.AverageDurationMs, + ) } - upstreamEndpoints, upstreamEndpointErr := r.GetUpstreamEndpointStatsWithFilters(ctx, start, end, filters.UserID, filters.APIKeyID, filters.AccountID, filters.GroupID, filters.Model, filters.RequestType, filters.Stream, filters.BillingType) - if upstreamEndpointErr != nil { - logger.LegacyPrintf("repository.usage_log", "GetUpstreamEndpointStatsWithFilters failed in GetStatsWithFilters: %v", upstreamEndpointErr) - upstreamEndpoints = []EndpointStat{} + // endpoint 明细:best-effort(失败 log + 返空),不致命。 + runEndpoints := func(c context.Context) { + res, err := r.GetEndpointStatsWithFilters(c, start, end, filters.UserID, filters.APIKeyID, filters.AccountID, filters.GroupID, filters.Model, filters.RequestType, filters.Stream, filters.BillingType) + if err != nil { + if !errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded) { + logger.LegacyPrintf("repository.usage_log", "GetEndpointStatsWithFilters failed in GetStatsWithFilters: %v", err) + } + res = []EndpointStat{} + } + endpoints = res + } + runUpstream := func(c context.Context) { + res, err := r.GetUpstreamEndpointStatsWithFilters(c, start, end, filters.UserID, filters.APIKeyID, filters.AccountID, filters.GroupID, filters.Model, filters.RequestType, filters.Stream, filters.BillingType) + if err != nil { + if !errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded) { + logger.LegacyPrintf("repository.usage_log", "GetUpstreamEndpointStatsWithFilters failed in GetStatsWithFilters: %v", err) + } + res = []EndpointStat{} + } + upstreamEndpoints = res + } + runPaths := func(c context.Context) { + res, err := r.getEndpointPathStatsWithFilters(c, start, end, filters.UserID, filters.APIKeyID, filters.AccountID, filters.GroupID, filters.Model, filters.RequestType, filters.Stream, filters.BillingType) + if err != nil { + if !errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded) { + logger.LegacyPrintf("repository.usage_log", "getEndpointPathStatsWithFilters failed in GetStatsWithFilters: %v", err) + } + res = []EndpointStat{} + } + endpointPaths = res } - endpointPaths, endpointPathErr := r.getEndpointPathStatsWithFilters(ctx, start, end, filters.UserID, filters.APIKeyID, filters.AccountID, filters.GroupID, filters.Model, filters.RequestType, filters.Stream, filters.BillingType) - if endpointPathErr != nil { - logger.LegacyPrintf("repository.usage_log", "getEndpointPathStatsWithFilters failed in GetStatsWithFilters: %v", endpointPathErr) - endpointPaths = []EndpointStat{} + + if r.db != nil { + // 生产路径:r.sql 是 *sql.DB 连接池,可并发。4 条查询并行,延迟取最大值。 + g, gctx := errgroup.WithContext(ctx) + g.Go(func() error { return runSummary(gctx) }) + g.Go(func() error { runEndpoints(gctx); return nil }) + g.Go(func() error { runUpstream(gctx); return nil }) + g.Go(func() error { runPaths(gctx); return nil }) + if err := g.Wait(); err != nil { + return nil, err + } + } else { + // 事务路径(ent.Tx 不能并发查询):顺序执行,行为与重构前一致。 + if err := runSummary(ctx); err != nil { + return nil, err + } + runEndpoints(ctx) + runUpstream(ctx) + runPaths(ctx) } + + stats.TotalAccountCost = &totalAccountCost + stats.TotalTokens = stats.TotalInputTokens + stats.TotalOutputTokens + stats.TotalCacheTokens stats.Endpoints = endpoints stats.UpstreamEndpoints = upstreamEndpoints stats.EndpointPaths = endpointPaths diff --git a/backend/internal/repository/usage_log_repo_stats_integration_test.go b/backend/internal/repository/usage_log_repo_stats_integration_test.go new file mode 100644 index 000000000000..09ac2aee1417 --- /dev/null +++ b/backend/internal/repository/usage_log_repo_stats_integration_test.go @@ -0,0 +1,51 @@ +//go:build integration + +package repository + +import ( + "context" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/usagestats" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/stretchr/testify/require" +) + +func TestUsageLog_GetStatsWithFilters_AggregatesAndEndpoints(t *testing.T) { + ctx := context.Background() + tx := testEntTx(t) + client := tx.Client() + repo := newUsageLogRepositoryWithSQL(client, tx) + + user := mustCreateUser(t, client, &service.User{Email: "stats@test.com"}) + apiKey := mustCreateApiKey(t, client, &service.APIKey{UserID: user.ID, Key: "sk-stats-1", Name: "k"}) + account := mustCreateAccount(t, client, &service.Account{Name: "acc-stats"}) + + now := time.Now().UTC() + inboundEndpoint := "/v1/messages" + upstreamEndpoint := "/v1/responses" + for i := 0; i < 3; i++ { + _, err := repo.Create(ctx, &service.UsageLog{ + UserID: user.ID, APIKeyID: apiKey.ID, AccountID: account.ID, + Model: "claude-3", InputTokens: 2, OutputTokens: 3, + TotalCost: 0.5, ActualCost: 0.4, CreatedAt: now, + InboundEndpoint: &inboundEndpoint, UpstreamEndpoint: &upstreamEndpoint, + }) + require.NoError(t, err) + } + + start := now.Add(-1 * time.Hour) + end := now.Add(1 * time.Hour) + // 按本测试创建的 user 维度过滤:集成库为共享实例,其它用 testEntClient 的兄弟测试会留下 + // 已提交的 usage_log 行(含零 token 的失败请求),不限定 user 会把它们计入 TotalRequests。 + stats, err := repo.GetStatsWithFilters(ctx, usagestats.UsageLogFilters{UserID: user.ID, StartTime: &start, EndTime: &end}) + require.NoError(t, err) + require.Equal(t, int64(3), stats.TotalRequests) + require.Equal(t, int64(6), stats.TotalInputTokens) + require.Equal(t, int64(9), stats.TotalOutputTokens) + require.InDelta(t, 1.2, stats.TotalActualCost, 1e-9) + require.NotEmpty(t, stats.Endpoints) + require.NotEmpty(t, stats.UpstreamEndpoints) + require.NotEmpty(t, stats.EndpointPaths) +} diff --git a/frontend/src/api/admin/usage.ts b/frontend/src/api/admin/usage.ts index c3a38c0a2cc1..d933ac6351f0 100644 --- a/frontend/src/api/admin/usage.ts +++ b/frontend/src/api/admin/usage.ts @@ -121,6 +121,7 @@ export async function getStats(params: { start_date?: string end_date?: string timezone?: string + nocache?: number }): Promise { const { data } = await apiClient.get('/admin/usage/stats', { params diff --git a/frontend/src/components/admin/usage/UsageFilters.vue b/frontend/src/components/admin/usage/UsageFilters.vue index 40f2c9f82d3d..a800f1903d1b 100644 --- a/frontend/src/components/admin/usage/UsageFilters.vue +++ b/frontend/src/components/admin/usage/UsageFilters.vue @@ -168,7 +168,7 @@ diff --git a/frontend/src/views/admin/__tests__/UsageView.spec.ts b/frontend/src/views/admin/__tests__/UsageView.spec.ts index 9cdb2a3c8b9c..8c644a7536d2 100644 --- a/frontend/src/views/admin/__tests__/UsageView.spec.ts +++ b/frontend/src/views/admin/__tests__/UsageView.spec.ts @@ -3,7 +3,7 @@ import { flushPromises, mount } from '@vue/test-utils' import UsageView from '../UsageView.vue' -const { list, getStats, getSnapshotV2, getById } = vi.hoisted(() => { +const { list, getStats, getSnapshotV2, getById, getModelStats } = vi.hoisted(() => { vi.stubGlobal('localStorage', { getItem: vi.fn(() => null), setItem: vi.fn(), @@ -15,6 +15,7 @@ const { list, getStats, getSnapshotV2, getById } = vi.hoisted(() => { getStats: vi.fn(), getSnapshotV2: vi.fn(), getById: vi.fn(), + getModelStats: vi.fn(), } }) @@ -40,6 +41,7 @@ vi.mock('@/api/admin', () => ({ }, dashboard: { getSnapshotV2, + getModelStats, }, users: { getById, @@ -116,6 +118,7 @@ describe('admin UsageView distribution metric toggles', () => { getStats.mockReset() getSnapshotV2.mockReset() getById.mockReset() + getModelStats.mockReset() list.mockResolvedValue({ items: [], @@ -137,12 +140,44 @@ describe('admin UsageView distribution metric toggles', () => { models: [], groups: [], }) + getModelStats.mockResolvedValue({ models: [] }) }) afterEach(() => { vi.useRealTimers() }) + it('keeps previous model stats visible during refresh until new data arrives', async () => { + // 首次加载返回 A + getModelStats.mockResolvedValueOnce({ models: [{ model: 'A', total_tokens: 10 }] }) + + const wrapper = mount(UsageView, { + global: { stubs: { + AppLayout: AppLayoutStub, UsageStatsCards: true, UsageFilters: UsageFiltersStub, + UsageTable: true, UsageExportProgress: true, UsageCleanupDialog: true, + UserBalanceHistoryModal: true, AuditLogModal: true, Pagination: true, Select: true, + DateRangePicker: true, Icon: true, TokenUsageTrend: true, + ModelDistributionChart: ModelDistributionChartStub, GroupDistributionChart: GroupDistributionChartStub, + EndpointDistributionChart: true, + } }, + }) + vi.advanceTimersByTime(120) + await flushPromises() + expect((wrapper.vm as any).requestedModelStats).toEqual([{ model: 'A', total_tokens: 10 }]) + + // 刷新:让第二次 getModelStats 处于 pending,断言旧数据 A 仍在(不被清空成 []) + let resolveSecond: (v: any) => void = () => {} + getModelStats.mockReturnValueOnce(new Promise((res) => { resolveSecond = res })) + ;(wrapper.vm as any).refreshData() + await flushPromises() + expect((wrapper.vm as any).requestedModelStats).toEqual([{ model: 'A', total_tokens: 10 }]) + + // 新数据到达后替换为 B + resolveSecond({ models: [{ model: 'B', total_tokens: 20 }] }) + await flushPromises() + expect((wrapper.vm as any).requestedModelStats).toEqual([{ model: 'B', total_tokens: 20 }]) + }) + it('keeps model and group metric toggles independent without refetching chart data', async () => { const wrapper = mount(UsageView, { global: { From d8cbf9ab5cb9c5ec4d3511e63d642384a6c60e62 Mon Sep 17 00:00:00 2001 From: name <136912576+is7Qin@users.noreply.github.com> Date: Fri, 29 May 2026 21:05:47 +0800 Subject: [PATCH 17/36] refactor(gateway): introduce request body refs --- backend/internal/handler/gateway_handler.go | 12 +++-- .../gateway_handler_chat_completions.go | 5 +- .../handler/gateway_handler_responses.go | 5 +- .../internal/handler/gemini_v1beta_handler.go | 2 +- ...teway_anthropic_apikey_passthrough_test.go | 16 +++---- backend/internal/service/gateway_request.go | 46 +++++++++++++++---- .../internal/service/gateway_request_test.go | 44 +++++++++--------- backend/internal/service/gateway_service.go | 12 ++--- .../service/gateway_service_benchmark_test.go | 2 +- .../service/gateway_websearch_emulation.go | 2 +- .../service/generate_session_hash_test.go | 10 ++-- 11 files changed, 95 insertions(+), 61 deletions(-) diff --git a/backend/internal/handler/gateway_handler.go b/backend/internal/handler/gateway_handler.go index eb5c4a42b393..79bed8b93156 100644 --- a/backend/internal/handler/gateway_handler.go +++ b/backend/internal/handler/gateway_handler.go @@ -154,7 +154,8 @@ func (h *GatewayHandler) Messages(c *gin.Context) { setOpsRequestContext(c, "", false) - parsedReq, err := service.ParseGatewayRequest(body, domain.PlatformAnthropic) + bodyRef := service.NewRequestBodyRef(body) + parsedReq, err := service.ParseGatewayRequest(bodyRef, domain.PlatformAnthropic) if err != nil { h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body") return @@ -746,11 +747,11 @@ func (h *GatewayHandler) Messages(c *gin.Context) { // 应用渠道模型映射到请求 if channelMapping.Mapped { parsedReq.Model = channelMapping.MappedModel - parsedReq.Body = h.gatewayService.ReplaceModelInBody(parsedReq.Body, channelMapping.MappedModel) + parsedReq.Body.Replace(h.gatewayService.ReplaceModelInBody(parsedReq.Body.Bytes(), channelMapping.MappedModel)) } // Bedrock CC 兼容：渠道模型映射后，清理 Anthropic API 专有字段、注入 Bedrock 必需字段 - parsedReq.Body = h.gatewayService.ApplyBedrockCCCompat(c.Request.Context(), parsedReq.Body, parsedReq.Model, account, apiKey.GroupID) - body = parsedReq.Body + parsedReq.Body.Replace(h.gatewayService.ApplyBedrockCCCompat(c.Request.Context(), parsedReq.Body.Bytes(), parsedReq.Model, account, apiKey.GroupID)) + body = parsedReq.Body.Bytes() // 转发请求 - 根据账号平台分流 c.Set("parsed_request", parsedReq) @@ -1683,7 +1684,8 @@ func (h *GatewayHandler) CountTokens(c *gin.Context) { setOpsRequestContext(c, "", false) - parsedReq, err := service.ParseGatewayRequest(body, domain.PlatformAnthropic) + bodyRef := service.NewRequestBodyRef(body) + parsedReq, err := service.ParseGatewayRequest(bodyRef, domain.PlatformAnthropic) if err != nil { h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body") return diff --git a/backend/internal/handler/gateway_handler_chat_completions.go b/backend/internal/handler/gateway_handler_chat_completions.go index daf6e6ea0415..719700aa9f63 100644 --- a/backend/internal/handler/gateway_handler_chat_completions.go +++ b/backend/internal/handler/gateway_handler_chat_completions.go @@ -151,9 +151,10 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) { } // Parse request for session hash - parsedReq, _ := service.ParseGatewayRequest(body, "chat_completions") + bodyRef := service.NewRequestBodyRef(body) + parsedReq, _ := service.ParseGatewayRequest(bodyRef, "chat_completions") if parsedReq == nil { - parsedReq = &service.ParsedRequest{Model: reqModel, Stream: reqStream, Body: body} + parsedReq = &service.ParsedRequest{Model: reqModel, Stream: reqStream, Body: bodyRef} } parsedReq.SessionContext = &service.SessionContext{ ClientIP: ip.GetClientIP(c), diff --git a/backend/internal/handler/gateway_handler_responses.go b/backend/internal/handler/gateway_handler_responses.go index f57b9989cf02..49f80d19c149 100644 --- a/backend/internal/handler/gateway_handler_responses.go +++ b/backend/internal/handler/gateway_handler_responses.go @@ -156,9 +156,10 @@ func (h *GatewayHandler) Responses(c *gin.Context) { } // Parse request for session hash - parsedReq, _ := service.ParseGatewayRequest(body, "responses") + bodyRef := service.NewRequestBodyRef(body) + parsedReq, _ := service.ParseGatewayRequest(bodyRef, "responses") if parsedReq == nil { - parsedReq = &service.ParsedRequest{Model: reqModel, Stream: reqStream, Body: body} + parsedReq = &service.ParsedRequest{Model: reqModel, Stream: reqStream, Body: bodyRef} } parsedReq.SessionContext = &service.SessionContext{ ClientIP: ip.GetClientIP(c), diff --git a/backend/internal/handler/gemini_v1beta_handler.go b/backend/internal/handler/gemini_v1beta_handler.go index 0b33ca3e3efe..5d8e6fa890a8 100644 --- a/backend/internal/handler/gemini_v1beta_handler.go +++ b/backend/internal/handler/gemini_v1beta_handler.go @@ -262,7 +262,7 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) { sessionHash := extractGeminiCLISessionHash(c, body) if sessionHash == "" { // Fallback: 使用通用的会话哈希生成逻辑（适用于其他客户端） - parsedReq, _ := service.ParseGatewayRequest(body, domain.PlatformGemini) + parsedReq, _ := service.ParseGatewayRequest(service.NewRequestBodyRef(body), domain.PlatformGemini) if parsedReq != nil { parsedReq.SessionContext = &service.SessionContext{ ClientIP: ip.GetClientIP(c), diff --git a/backend/internal/service/gateway_anthropic_apikey_passthrough_test.go b/backend/internal/service/gateway_anthropic_apikey_passthrough_test.go index 9062c517aa76..a67a3dc22f77 100644 --- a/backend/internal/service/gateway_anthropic_apikey_passthrough_test.go +++ b/backend/internal/service/gateway_anthropic_apikey_passthrough_test.go @@ -112,7 +112,7 @@ func TestGatewayService_AnthropicAPIKeyPassthrough_ForwardStreamPreservesBodyAnd body := []byte(`{"model":"claude-3-7-sonnet-20250219","stream":true,"system":[{"type":"text","text":"x-anthropic-billing-header keep"}],"messages":[{"role":"user","content":[{"type":"text","text":"hello"}]}]}`) parsed := &ParsedRequest{ - Body: body, + Body: NewRequestBodyRef(body), Model: "claude-3-7-sonnet-20250219", Stream: true, } @@ -202,7 +202,7 @@ func TestGatewayService_AnthropicAPIKeyPassthrough_ForwardCountTokensPreservesBo body := []byte(`{"model":"claude-3-5-sonnet-latest","messages":[{"role":"user","content":[{"type":"text","text":"hello"}]}],"thinking":{"type":"enabled"}}`) parsed := &ParsedRequest{ - Body: body, + Body: NewRequestBodyRef(body), Model: "claude-3-5-sonnet-latest", } @@ -344,7 +344,7 @@ func TestGatewayService_AnthropicAPIKeyPassthrough_ModelMappingEdgeCases(t *test body := []byte(`{"model":"` + tt.model + `","messages":[{"role":"user","content":[{"type":"text","text":"hello"}]}]}`) parsed := &ParsedRequest{ - Body: body, + Body: NewRequestBodyRef(body), Model: tt.model, } @@ -429,7 +429,7 @@ func TestGatewayService_AnthropicAPIKeyPassthrough_ModelMappingPreservesOtherFie // 包含复杂字段的请求体：system、thinking、messages body := []byte(`{"model":"claude-sonnet-4-20250514","system":[{"type":"text","text":"You are a helpful assistant."}],"messages":[{"role":"user","content":[{"type":"text","text":"hello world"}]}],"thinking":{"type":"enabled","budget_tokens":5000},"max_tokens":1024}`) parsed := &ParsedRequest{ - Body: body, + Body: NewRequestBodyRef(body), Model: "claude-sonnet-4-20250514", } @@ -485,7 +485,7 @@ func TestGatewayService_AnthropicAPIKeyPassthrough_CountTokensFiltersGenerationF body := []byte(`{"model":"claude-sonnet-4-20250514","system":[{"type":"text","text":"sys"}],"messages":[{"role":"user","content":"hello"}],"tools":[{"name":"tool","input_schema":{"type":"object"}}],"temperature":0.7,"top_p":0.9,"top_k":40,"stream":true,"stop_sequences":["END"],"max_tokens":1024,"thinking":{"type":"enabled","budget_tokens":5000}}`) parsed := &ParsedRequest{ - Body: body, + Body: NewRequestBodyRef(body), Model: "claude-sonnet-4-20250514", } @@ -547,7 +547,7 @@ func TestGatewayService_AnthropicAPIKeyPassthrough_EmptyModelSkipsMapping(t *tes body := []byte(`{"messages":[{"role":"user","content":"hello"}]}`) parsed := &ParsedRequest{ - Body: body, + Body: NewRequestBodyRef(body), Model: "", // 空模型 } @@ -636,7 +636,7 @@ func TestGatewayService_AnthropicAPIKeyPassthrough_CountTokens404PassthroughNotE c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages/count_tokens", nil) body := []byte(`{"model":"claude-sonnet-4-5-20250929","messages":[{"role":"user","content":"hi"}]}`) - parsed := &ParsedRequest{Body: body, Model: "claude-sonnet-4-5-20250929"} + parsed := &ParsedRequest{Body: NewRequestBodyRef(body), Model: "claude-sonnet-4-5-20250929"} upstream := &anthropicHTTPUpstreamRecorder{ resp: &http.Response{ @@ -767,7 +767,7 @@ func TestGatewayService_AnthropicOAuth_ForwardPreservesBillingHeaderSystemBlock( c, _ := gin.CreateTestContext(rec) c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) - parsed, err := ParseGatewayRequest([]byte(tt.body), PlatformAnthropic) + parsed, err := ParseGatewayRequest(NewRequestBodyRef([]byte(tt.body)), PlatformAnthropic) require.NoError(t, err) upstream := &anthropicHTTPUpstreamRecorder{ diff --git a/backend/internal/service/gateway_request.go b/backend/internal/service/gateway_request.go index 91f7601cb1ac..819bb0a882cc 100644 --- a/backend/internal/service/gateway_request.go +++ b/backend/internal/service/gateway_request.go @@ -51,6 +51,35 @@ type SessionContext struct { APIKeyID int64 } +type RequestBodyRef struct { + data []byte +} + +func NewRequestBodyRef(data []byte) *RequestBodyRef { + return &RequestBodyRef{data: data} +} + +func (b *RequestBodyRef) Bytes() []byte { + if b == nil { + return nil + } + return b.data +} + +func (b *RequestBodyRef) Len() int { + if b == nil { + return 0 + } + return len(b.data) +} + +func (b *RequestBodyRef) Replace(data []byte) { + if b == nil { + return + } + b.data = data +} + // ParsedRequest 保存网关请求的预解析结果 // // 性能优化说明： @@ -64,7 +93,7 @@ type SessionContext struct { // 2. 将解析结果 ParsedRequest 传递给 Service 层 // 3. 避免重复 json.Unmarshal，减少 CPU 和内存开销 type ParsedRequest struct { - Body []byte // 原始请求体（保留用于转发） + Body *RequestBodyRef // 原始请求体引用（保留用于转发） Model string // 请求的模型名称 Stream bool // 是否为流式请求 MetadataUserID string // metadata.user_id（用于会话亲和） @@ -130,17 +159,18 @@ func normalizeSessionUserAgentFallback(raw string) string { // ParseGatewayRequest 解析网关请求体并返回结构化结果。 // protocol 指定请求协议格式（domain.PlatformAnthropic / domain.PlatformGemini）， // 不同协议使用不同的 system/messages 字段名。 -func ParseGatewayRequest(body []byte, protocol string) (*ParsedRequest, error) { +func ParseGatewayRequest(body *RequestBodyRef, protocol string) (*ParsedRequest, error) { + bodyBytes := body.Bytes() // 保持与旧实现一致：请求体必须是合法 JSON。 // 注意：gjson.GetBytes 对非法 JSON 不会报错，因此需要显式校验。 - if !gjson.ValidBytes(body) { + if !gjson.ValidBytes(bodyBytes) { return nil, fmt.Errorf("invalid json") } // 性能： // - gjson.GetBytes 会把匹配的 Raw/Str 安全复制成 string（对于巨大 messages 会产生额外拷贝）。 // - 这里将 body 通过 unsafe 零拷贝视为 string，仅在本函数内使用，且 body 不会被修改。 - jsonStr := *(*string)(unsafe.Pointer(&body)) + jsonStr := *(*string)(unsafe.Pointer(&bodyBytes)) parsed := &ParsedRequest{ Body: body, @@ -197,7 +227,7 @@ func ParseGatewayRequest(body []byte, protocol string) (*ParsedRequest, error) { // Gemini 原生格式: systemInstruction.parts / contents if sysParts := gjson.Get(jsonStr, "systemInstruction.parts"); sysParts.Exists() && sysParts.IsArray() { var parts []any - if err := json.Unmarshal(sliceRawFromBody(body, sysParts), &parts); err != nil { + if err := json.Unmarshal(sliceRawFromBody(bodyBytes, sysParts), &parts); err != nil { return nil, err } parsed.System = parts @@ -205,7 +235,7 @@ func ParseGatewayRequest(body []byte, protocol string) (*ParsedRequest, error) { if contents := gjson.Get(jsonStr, "contents"); contents.Exists() && contents.IsArray() { var msgs []any - if err := json.Unmarshal(sliceRawFromBody(body, contents), &msgs); err != nil { + if err := json.Unmarshal(sliceRawFromBody(bodyBytes, contents), &msgs); err != nil { return nil, err } parsed.Messages = msgs @@ -224,7 +254,7 @@ func ParseGatewayRequest(body []byte, protocol string) (*ParsedRequest, error) { parsed.System = sys.String() default: var system any - if err := json.Unmarshal(sliceRawFromBody(body, sys), &system); err != nil { + if err := json.Unmarshal(sliceRawFromBody(bodyBytes, sys), &system); err != nil { return nil, err } parsed.System = system @@ -233,7 +263,7 @@ func ParseGatewayRequest(body []byte, protocol string) (*ParsedRequest, error) { if msgs := gjson.Get(jsonStr, "messages"); msgs.Exists() && msgs.IsArray() { var messages []any - if err := json.Unmarshal(sliceRawFromBody(body, msgs), &messages); err != nil { + if err := json.Unmarshal(sliceRawFromBody(bodyBytes, msgs), &messages); err != nil { return nil, err } parsed.Messages = messages diff --git a/backend/internal/service/gateway_request_test.go b/backend/internal/service/gateway_request_test.go index 045dc66c15bd..d415b871a0fc 100644 --- a/backend/internal/service/gateway_request_test.go +++ b/backend/internal/service/gateway_request_test.go @@ -14,7 +14,7 @@ import ( func TestParseGatewayRequest(t *testing.T) { body := []byte(`{"model":"claude-3-7-sonnet","stream":true,"metadata":{"user_id":"session_123e4567-e89b-12d3-a456-426614174000"},"system":[{"type":"text","text":"hello","cache_control":{"type":"ephemeral"}}],"messages":[{"content":"hi"}]}`) - parsed, err := ParseGatewayRequest(body, "") + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), "") require.NoError(t, err) require.Equal(t, "claude-3-7-sonnet", parsed.Model) require.True(t, parsed.Stream) @@ -27,7 +27,7 @@ func TestParseGatewayRequest(t *testing.T) { func TestParseGatewayRequest_ThinkingEnabled(t *testing.T) { body := []byte(`{"model":"claude-sonnet-4-5","thinking":{"type":"enabled"},"messages":[{"content":"hi"}]}`) - parsed, err := ParseGatewayRequest(body, "") + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), "") require.NoError(t, err) require.Equal(t, "claude-sonnet-4-5", parsed.Model) require.True(t, parsed.ThinkingEnabled) @@ -35,7 +35,7 @@ func TestParseGatewayRequest_ThinkingEnabled(t *testing.T) { func TestParseGatewayRequest_ThinkingAdaptiveEnabled(t *testing.T) { body := []byte(`{"model":"claude-sonnet-4-5","thinking":{"type":"adaptive"},"messages":[{"content":"hi"}]}`) - parsed, err := ParseGatewayRequest(body, "") + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), "") require.NoError(t, err) require.Equal(t, "claude-sonnet-4-5", parsed.Model) require.True(t, parsed.ThinkingEnabled) @@ -43,21 +43,21 @@ func TestParseGatewayRequest_ThinkingAdaptiveEnabled(t *testing.T) { func TestParseGatewayRequest_MaxTokens(t *testing.T) { body := []byte(`{"model":"claude-haiku-4-5","max_tokens":1}`) - parsed, err := ParseGatewayRequest(body, "") + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), "") require.NoError(t, err) require.Equal(t, 1, parsed.MaxTokens) } func TestParseGatewayRequest_MaxTokensNonIntegralIgnored(t *testing.T) { body := []byte(`{"model":"claude-haiku-4-5","max_tokens":1.5}`) - parsed, err := ParseGatewayRequest(body, "") + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), "") require.NoError(t, err) require.Equal(t, 0, parsed.MaxTokens) } func TestParseGatewayRequest_SystemNull(t *testing.T) { body := []byte(`{"model":"claude-3","system":null}`) - parsed, err := ParseGatewayRequest(body, "") + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), "") require.NoError(t, err) // 显式传入 system:null 也应视为“字段已存在”，避免默认 system 被注入。 require.True(t, parsed.HasSystem) @@ -66,13 +66,13 @@ func TestParseGatewayRequest_SystemNull(t *testing.T) { func TestParseGatewayRequest_InvalidModelType(t *testing.T) { body := []byte(`{"model":123}`) - _, err := ParseGatewayRequest(body, "") + _, err := ParseGatewayRequest(NewRequestBodyRef(body), "") require.Error(t, err) } func TestParseGatewayRequest_InvalidStreamType(t *testing.T) { body := []byte(`{"stream":"true"}`) - _, err := ParseGatewayRequest(body, "") + _, err := ParseGatewayRequest(NewRequestBodyRef(body), "") require.Error(t, err) } @@ -86,7 +86,7 @@ func TestParseGatewayRequest_GeminiContents(t *testing.T) { {"role": "user", "parts": [{"text": "How are you?"}]} ] }`) - parsed, err := ParseGatewayRequest(body, domain.PlatformGemini) + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), domain.PlatformGemini) require.NoError(t, err) require.Len(t, parsed.Messages, 3, "should parse contents as Messages") require.False(t, parsed.HasSystem, "Gemini format should not set HasSystem") @@ -102,7 +102,7 @@ func TestParseGatewayRequest_GeminiSystemInstruction(t *testing.T) { {"role": "user", "parts": [{"text": "Hello"}]} ] }`) - parsed, err := ParseGatewayRequest(body, domain.PlatformGemini) + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), domain.PlatformGemini) require.NoError(t, err) require.NotNil(t, parsed.System, "should parse systemInstruction.parts as System") parts, ok := parsed.System.([]any) @@ -119,7 +119,7 @@ func TestParseGatewayRequest_GeminiWithModel(t *testing.T) { "model": "gemini-2.5-pro", "contents": [{"role": "user", "parts": [{"text": "test"}]}] }`) - parsed, err := ParseGatewayRequest(body, domain.PlatformGemini) + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), domain.PlatformGemini) require.NoError(t, err) require.Equal(t, "gemini-2.5-pro", parsed.Model) require.Len(t, parsed.Messages, 1) @@ -132,7 +132,7 @@ func TestParseGatewayRequest_GeminiIgnoresAnthropicFields(t *testing.T) { "messages": [{"role": "user", "content": "ignored"}], "contents": [{"role": "user", "parts": [{"text": "real content"}]}] }`) - parsed, err := ParseGatewayRequest(body, domain.PlatformGemini) + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), domain.PlatformGemini) require.NoError(t, err) require.False(t, parsed.HasSystem, "Gemini protocol should not parse Anthropic system field") require.Nil(t, parsed.System, "no systemInstruction = nil System") @@ -141,14 +141,14 @@ func TestParseGatewayRequest_GeminiIgnoresAnthropicFields(t *testing.T) { func TestParseGatewayRequest_GeminiEmptyContents(t *testing.T) { body := []byte(`{"contents": []}`) - parsed, err := ParseGatewayRequest(body, domain.PlatformGemini) + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), domain.PlatformGemini) require.NoError(t, err) require.Empty(t, parsed.Messages) } func TestParseGatewayRequest_GeminiNoContents(t *testing.T) { body := []byte(`{"model": "gemini-2.5-flash"}`) - parsed, err := ParseGatewayRequest(body, domain.PlatformGemini) + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), domain.PlatformGemini) require.NoError(t, err) require.Nil(t, parsed.Messages) require.Equal(t, "gemini-2.5-flash", parsed.Model) @@ -162,7 +162,7 @@ func TestParseGatewayRequest_AnthropicIgnoresGeminiFields(t *testing.T) { "contents": [{"role": "user", "parts": [{"text": "ignored"}]}], "systemInstruction": {"parts": [{"text": "ignored"}]} }`) - parsed, err := ParseGatewayRequest(body, domain.PlatformAnthropic) + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), domain.PlatformAnthropic) require.NoError(t, err) require.True(t, parsed.HasSystem) require.Equal(t, "real system", parsed.System) @@ -897,7 +897,7 @@ func TestParseGatewayRequest_TypeValidation(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - _, err := ParseGatewayRequest([]byte(tt.body), "") + _, err := ParseGatewayRequest(NewRequestBodyRef([]byte(tt.body)), "") if tt.wantErr { require.Error(t, err) if tt.errSubstr != "" { @@ -959,7 +959,7 @@ func TestParseGatewayRequest_OptionalFieldsMissing(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - parsed, err := ParseGatewayRequest([]byte(tt.body), "") + parsed, err := ParseGatewayRequest(NewRequestBodyRef([]byte(tt.body)), "") require.NoError(t, err) require.Equal(t, tt.wantModel, parsed.Model) @@ -1023,7 +1023,7 @@ func TestParseGatewayRequest_MaxTokensBoundary(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - parsed, err := ParseGatewayRequest([]byte(tt.body), "") + parsed, err := ParseGatewayRequest(NewRequestBodyRef([]byte(tt.body)), "") if tt.wantErr { require.Error(t, err) return @@ -1040,7 +1040,7 @@ func TestParseGatewayRequest_MaxTokensBoundary(t *testing.T) { // 核心路径：先 Unmarshal 到 map[string]any，再逐字段提取。 func parseGatewayRequestOld(body []byte, protocol string) (*ParsedRequest, error) { parsed := &ParsedRequest{ - Body: body, + Body: NewRequestBodyRef(body), } var req map[string]any @@ -1151,7 +1151,7 @@ func BenchmarkParseGatewayRequest_New_Small(b *testing.B) { b.SetBytes(int64(len(data))) b.ResetTimer() for i := 0; i < b.N; i++ { - _, _ = ParseGatewayRequest(data, "") + _, _ = ParseGatewayRequest(NewRequestBodyRef(data), "") } } @@ -1203,7 +1203,7 @@ func TestParseGatewayRequest_OutputEffort(t *testing.T) { } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - parsed, err := ParseGatewayRequest([]byte(tt.body), "") + parsed, err := ParseGatewayRequest(NewRequestBodyRef([]byte(tt.body)), "") require.NoError(t, err) require.Equal(t, tt.wantEffort, parsed.OutputEffort) }) @@ -1245,6 +1245,6 @@ func BenchmarkParseGatewayRequest_New_Large(b *testing.B) { b.SetBytes(int64(len(data))) b.ResetTimer() for i := 0; i < b.N; i++ { - _, _ = ParseGatewayRequest(data, "") + _, _ = ParseGatewayRequest(NewRequestBodyRef(data), "") } } diff --git a/backend/internal/service/gateway_service.go b/backend/internal/service/gateway_service.go index f807f3ec6532..8f55bf13c6d3 100644 --- a/backend/internal/service/gateway_service.go +++ b/backend/internal/service/gateway_service.go @@ -4400,12 +4400,12 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A } // Web Search 模拟：纯 web_search 请求时，直接调用搜索 API 构造响应 - if account != nil && s.shouldEmulateWebSearch(ctx, account, parsed.GroupID, parsed.Body) { + if account != nil && s.shouldEmulateWebSearch(ctx, account, parsed.GroupID, parsed.Body.Bytes()) { return s.handleWebSearchEmulation(ctx, c, account, parsed) } if account != nil && account.IsAnthropicAPIKeyPassthroughEnabled() { - passthroughBody := parsed.Body + passthroughBody := parsed.Body.Bytes() passthroughModel := parsed.Model if passthroughModel != "" { if mappedModel := account.GetMappedModel(passthroughModel); mappedModel != passthroughModel { @@ -4441,7 +4441,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A c.Set(betaPolicyFilterSetKey, filterSet) } - body := parsed.Body + body := parsed.Body.Bytes() reqModel := parsed.Model reqStream := parsed.Stream originalModel := reqModel @@ -5735,7 +5735,7 @@ func (s *GatewayService) forwardBedrock( ) (*ForwardResult, error) { reqModel := parsed.Model reqStream := parsed.Stream - body := parsed.Body + body := parsed.Body.Bytes() region := bedrockRuntimeRegion(account) mappedModel, ok := ResolveBedrockModelID(account, reqModel) @@ -9172,7 +9172,7 @@ func (s *GatewayService) ForwardCountTokens(ctx context.Context, c *gin.Context, } if account != nil && account.IsAnthropicAPIKeyPassthroughEnabled() { - passthroughBody := parsed.Body + passthroughBody := parsed.Body.Bytes() if reqModel := parsed.Model; reqModel != "" { if mappedModel := account.GetMappedModel(reqModel); mappedModel != reqModel { passthroughBody = s.replaceModelInBody(passthroughBody, mappedModel) @@ -9188,7 +9188,7 @@ func (s *GatewayService) ForwardCountTokens(ctx context.Context, c *gin.Context, return nil } - body := parsed.Body + body := parsed.Body.Bytes() reqModel := parsed.Model // Pre-filter: strip empty text blocks to prevent upstream 400. diff --git a/backend/internal/service/gateway_service_benchmark_test.go b/backend/internal/service/gateway_service_benchmark_test.go index c9c4d3dd8bab..5637680b3f07 100644 --- a/backend/internal/service/gateway_service_benchmark_test.go +++ b/backend/internal/service/gateway_service_benchmark_test.go @@ -14,7 +14,7 @@ func BenchmarkGenerateSessionHash_Metadata(b *testing.B) { b.ReportAllocs() for i := 0; i < b.N; i++ { - parsed, err := ParseGatewayRequest(body, "") + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), "") if err != nil { b.Fatalf("解析请求失败: %v", err) } diff --git a/backend/internal/service/gateway_websearch_emulation.go b/backend/internal/service/gateway_websearch_emulation.go index a42b5585e323..2f9c8e0c258b 100644 --- a/backend/internal/service/gateway_websearch_emulation.go +++ b/backend/internal/service/gateway_websearch_emulation.go @@ -150,7 +150,7 @@ func (s *GatewayService) handleWebSearchEmulation( parsed.OnUpstreamAccepted() } - query := extractSearchQueryFromBody(parsed.Body) + query := extractSearchQueryFromBody(parsed.Body.Bytes()) if query == "" { return nil, fmt.Errorf("web search emulation: no query found in messages") } diff --git a/backend/internal/service/generate_session_hash_test.go b/backend/internal/service/generate_session_hash_test.go index 39679c3da9cf..8f3258b72c71 100644 --- a/backend/internal/service/generate_session_hash_test.go +++ b/backend/internal/service/generate_session_hash_test.go @@ -1198,7 +1198,7 @@ func TestGenerateSessionHash_GeminiMultiTurnHashNotSticky(t *testing.T) { hashes := make([]string, 3) for i, body := range [][]byte{round1Body, round2Body, round3Body} { - parsed, err := ParseGatewayRequest(body, "gemini") + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), "gemini") require.NoError(t, err) parsed.SessionContext = ctx hashes[i] = svc.GenerateSessionHash(parsed) @@ -1211,7 +1211,7 @@ func TestGenerateSessionHash_GeminiMultiTurnHashNotSticky(t *testing.T) { require.NotEqual(t, hashes[0], hashes[2], "round 1 vs 3 hash should differ") // 同一轮重试应产生相同 hash - parsed1Again, err := ParseGatewayRequest(round2Body, "gemini") + parsed1Again, err := ParseGatewayRequest(NewRequestBodyRef(round2Body), "gemini") require.NoError(t, err) parsed1Again.SessionContext = ctx h2Again := svc.GenerateSessionHash(parsed1Again) @@ -1234,7 +1234,7 @@ func TestGenerateSessionHash_GeminiEndToEnd(t *testing.T) { ] }`) - parsed, err := ParseGatewayRequest(body, "gemini") + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), "gemini") require.NoError(t, err) parsed.SessionContext = &SessionContext{ ClientIP: "10.0.0.1", @@ -1246,7 +1246,7 @@ func TestGenerateSessionHash_GeminiEndToEnd(t *testing.T) { require.NotEmpty(t, h, "end-to-end Gemini flow should produce a hash") // 同一请求再次解析应产生相同 hash - parsed2, err := ParseGatewayRequest(body, "gemini") + parsed2, err := ParseGatewayRequest(NewRequestBodyRef(body), "gemini") require.NoError(t, err) parsed2.SessionContext = &SessionContext{ ClientIP: "10.0.0.1", @@ -1258,7 +1258,7 @@ func TestGenerateSessionHash_GeminiEndToEnd(t *testing.T) { require.Equal(t, h, h2, "same request should produce same hash") // 不同用户发送相同请求应产生不同 hash - parsed3, err := ParseGatewayRequest(body, "gemini") + parsed3, err := ParseGatewayRequest(NewRequestBodyRef(body), "gemini") require.NoError(t, err) parsed3.SessionContext = &SessionContext{ ClientIP: "10.0.0.2", From b1c4be4ac81519df4c562f9563e7be8fe94ada2a Mon Sep 17 00:00:00 2001 From: name <136912576+is7Qin@users.noreply.github.com> Date: Sat, 30 May 2026 00:37:34 +0800 Subject: [PATCH 18/36] refactor(gateway): remove parsed request object graphs Keep large gateway payloads as raw body ranges and bind OpenAI parsed-body caches to the body bytes so failover and mapping do not reuse stale mutable state. --- backend/internal/handler/gateway_handler.go | 4 +- backend/internal/handler/gateway_helper.go | 14 +- .../handler/gateway_helper_hotpath_test.go | 11 +- .../handler/openai_gateway_handler.go | 2 +- backend/internal/service/anthropic_session.go | 45 +- .../service/anthropic_session_test.go | 148 +-- .../service/gateway_oauth_metadata_test.go | 2 - backend/internal/service/gateway_request.go | 196 +++- .../internal/service/gateway_request_test.go | 68 +- backend/internal/service/gateway_service.go | 193 ++-- .../service/gateway_service_benchmark_test.go | 24 +- .../service/generate_session_hash_test.go | 1009 +++-------------- .../service/openai_gateway_service.go | 69 +- .../openai_gateway_service_hotpath_test.go | 20 +- .../service/openai_oauth_passthrough_test.go | 52 + .../service/user_msg_queue_service.go | 52 +- 16 files changed, 690 insertions(+), 1219 deletions(-) diff --git a/backend/internal/handler/gateway_handler.go b/backend/internal/handler/gateway_handler.go index 79bed8b93156..10d67ba32d4d 100644 --- a/backend/internal/handler/gateway_handler.go +++ b/backend/internal/handler/gateway_handler.go @@ -747,10 +747,10 @@ func (h *GatewayHandler) Messages(c *gin.Context) { // 应用渠道模型映射到请求 if channelMapping.Mapped { parsedReq.Model = channelMapping.MappedModel - parsedReq.Body.Replace(h.gatewayService.ReplaceModelInBody(parsedReq.Body.Bytes(), channelMapping.MappedModel)) + parsedReq.ReplaceBody(h.gatewayService.ReplaceModelInBody(parsedReq.Body.Bytes(), channelMapping.MappedModel)) } // Bedrock CC 兼容：渠道模型映射后，清理 Anthropic API 专有字段、注入 Bedrock 必需字段 - parsedReq.Body.Replace(h.gatewayService.ApplyBedrockCCCompat(c.Request.Context(), parsedReq.Body.Bytes(), parsedReq.Model, account, apiKey.GroupID)) + parsedReq.ReplaceBody(h.gatewayService.ApplyBedrockCCCompat(c.Request.Context(), parsedReq.Body.Bytes(), parsedReq.Model, account, apiKey.GroupID)) body = parsedReq.Body.Bytes() // 转发请求 - 根据账号平台分流 diff --git a/backend/internal/handler/gateway_helper.go b/backend/internal/handler/gateway_helper.go index e48975027830..52362dd11db5 100644 --- a/backend/internal/handler/gateway_helper.go +++ b/backend/internal/handler/gateway_helper.go @@ -74,8 +74,12 @@ func claudeCodeBodyMapFromParsedRequest(parsedReq *service.ParsedRequest) map[st bodyMap := map[string]any{ "model": parsedReq.Model, } - if parsedReq.System != nil || parsedReq.HasSystem { - bodyMap["system"] = parsedReq.System + if parsedReq.HasSystem { + if system, ok := parsedReq.SystemValue(); ok { + bodyMap["system"] = system + } else { + bodyMap["system"] = nil + } } if parsedReq.MetadataUserID != "" { bodyMap["metadata"] = map[string]any{"user_id": parsedReq.MetadataUserID} @@ -87,10 +91,8 @@ func claudeCodeBodyMapFromContextCache(c *gin.Context) map[string]any { if c == nil { return nil } - if cached, ok := c.Get(service.OpenAIParsedRequestBodyKey); ok { - if bodyMap, ok := cached.(map[string]any); ok { - return bodyMap - } + if bodyMap := service.CachedOpenAIParsedRequestBody(c); bodyMap != nil { + return bodyMap } if cached, ok := c.Get(claudeCodeParsedRequestContextKey); ok { switch v := cached.(type) { diff --git a/backend/internal/handler/gateway_helper_hotpath_test.go b/backend/internal/handler/gateway_helper_hotpath_test.go index d57c396cc1e4..d973bedbc8cd 100644 --- a/backend/internal/handler/gateway_helper_hotpath_test.go +++ b/backend/internal/handler/gateway_helper_hotpath_test.go @@ -185,13 +185,8 @@ func TestSetClaudeCodeClientContext_ReuseParsedRequestAndContextCache(t *testing c.Request.Header.Set("anthropic-beta", "message-batches-2024-09-24") c.Request.Header.Set("anthropic-version", "2023-06-01") - parsedReq := &service.ParsedRequest{ - Model: "claude-3-5-sonnet-20241022", - System: []any{ - map[string]any{"text": "You are Claude Code, Anthropic's official CLI for Claude."}, - }, - MetadataUserID: "user_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa_account__session_aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa", - } + parsedReq, err := service.ParseGatewayRequest(service.NewRequestBodyRef(validClaudeCodeBodyJSON()), "") + require.NoError(t, err) // body 非法 JSON，如果函数复用 parsedReq 成功则仍应判定为 Claude Code。 SetClaudeCodeClientContext(c, []byte(`{invalid`), parsedReq) @@ -204,7 +199,7 @@ func TestSetClaudeCodeClientContext_ReuseParsedRequestAndContextCache(t *testing c.Request.Header.Set("X-App", "claude-code") c.Request.Header.Set("anthropic-beta", "message-batches-2024-09-24") c.Request.Header.Set("anthropic-version", "2023-06-01") - c.Set(service.OpenAIParsedRequestBodyKey, map[string]any{ + service.CacheOpenAIParsedRequestBody(c, []byte(`{invalid`), map[string]any{ "model": "claude-3-5-sonnet-20241022", "system": []any{ map[string]any{"text": "You are Claude Code, Anthropic's official CLI for Claude."}, diff --git a/backend/internal/handler/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index 620c686179e3..a131cbd909d2 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -954,7 +954,7 @@ func (h *OpenAIGatewayHandler) validateFunctionCallOutputRequest(c *gin.Context, return true } - c.Set(service.OpenAIParsedRequestBodyKey, reqBody) + service.CacheOpenAIParsedRequestBody(c, body, reqBody) validation := service.ValidateFunctionCallOutputContext(reqBody) if !validation.HasFunctionCallOutput { return true diff --git a/backend/internal/service/anthropic_session.go b/backend/internal/service/anthropic_session.go index 26544c68cd48..bca8cc7ffba7 100644 --- a/backend/internal/service/anthropic_session.go +++ b/backend/internal/service/anthropic_session.go @@ -4,6 +4,8 @@ import ( "encoding/json" "strings" "time" + + "github.com/tidwall/gjson" ) // Anthropic 会话 Fallback 相关常量 @@ -30,30 +32,39 @@ func BuildAnthropicDigestChain(parsed *ParsedRequest) string { var parts []string - // 1. system prompt - if parsed.System != nil { - systemData, _ := json.Marshal(parsed.System) - if len(systemData) > 0 && string(systemData) != "null" { - parts = append(parts, "s:"+shortHash(systemData)) - } + if systemRaw := parsed.SystemRaw(); len(systemRaw) > 0 && string(systemRaw) != "null" { + parts = append(parts, "s:"+shortHash(canonicalAnthropicDigestJSON(systemRaw))) } - // 2. messages - for _, msg := range parsed.Messages { - msgMap, ok := msg.(map[string]any) - if !ok { - continue - } - role, _ := msgMap["role"].(string) - prefix := rolePrefix(role) - content := msgMap["content"] - contentData, _ := json.Marshal(content) - parts = append(parts, prefix+":"+shortHash(contentData)) + messages := parsed.MessagesRaw() + if len(messages) > 0 { + gjson.ParseBytes(messages).ForEach(func(_, msg gjson.Result) bool { + prefix := rolePrefix(msg.Get("role").String()) + content := msg.Get("content") + parts = append(parts, prefix+":"+shortHash(canonicalAnthropicDigestJSON([]byte(content.Raw)))) + return true + }) } return strings.Join(parts, "-") } +// canonicalAnthropicDigestJSON 保持 digest 对 JSON key 顺序和空白不敏感。 +func canonicalAnthropicDigestJSON(raw []byte) []byte { + if len(raw) == 0 { + return raw + } + var value any + if err := json.Unmarshal(raw, &value); err != nil { + return raw + } + canonical, err := json.Marshal(value) + if err != nil { + return raw + } + return canonical +} + // rolePrefix 将 Anthropic 的 role 映射为单字符前缀 func rolePrefix(role string) string { switch role { diff --git a/backend/internal/service/anthropic_session_test.go b/backend/internal/service/anthropic_session_test.go index 10406643bda8..4d88fc8b3419 100644 --- a/backend/internal/service/anthropic_session_test.go +++ b/backend/internal/service/anthropic_session_test.go @@ -1,3 +1,5 @@ +//go:build unit + package service import ( @@ -5,6 +7,15 @@ import ( "testing" ) +func mustParseAnthropicDigestRequest(t *testing.T, body string) *ParsedRequest { + t.Helper() + parsed, err := ParseGatewayRequest(NewRequestBodyRef([]byte(body)), "") + if err != nil { + t.Fatalf("ParseGatewayRequest failed: %v", err) + } + return parsed +} + func TestBuildAnthropicDigestChain_NilRequest(t *testing.T) { result := BuildAnthropicDigestChain(nil) if result != "" { @@ -13,9 +24,7 @@ func TestBuildAnthropicDigestChain_NilRequest(t *testing.T) { } func TestBuildAnthropicDigestChain_EmptyMessages(t *testing.T) { - parsed := &ParsedRequest{ - Messages: []any{}, - } + parsed := mustParseAnthropicDigestRequest(t, `{"messages":[]}`) result := BuildAnthropicDigestChain(parsed) if result != "" { t.Errorf("expected empty string for empty messages, got: %s", result) @@ -23,11 +32,7 @@ func TestBuildAnthropicDigestChain_EmptyMessages(t *testing.T) { } func TestBuildAnthropicDigestChain_SingleUserMessage(t *testing.T) { - parsed := &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - } + parsed := mustParseAnthropicDigestRequest(t, `{"messages":[{"role":"user","content":"hello"}]}`) result := BuildAnthropicDigestChain(parsed) parts := splitChain(result) if len(parts) != 1 { @@ -39,12 +44,7 @@ func TestBuildAnthropicDigestChain_SingleUserMessage(t *testing.T) { } func TestBuildAnthropicDigestChain_UserAndAssistant(t *testing.T) { - parsed := &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - map[string]any{"role": "assistant", "content": "hi there"}, - }, - } + parsed := mustParseAnthropicDigestRequest(t, `{"messages":[{"role":"user","content":"hello"},{"role":"assistant","content":"hi there"}]}`) result := BuildAnthropicDigestChain(parsed) parts := splitChain(result) if len(parts) != 2 { @@ -59,12 +59,7 @@ func TestBuildAnthropicDigestChain_UserAndAssistant(t *testing.T) { } func TestBuildAnthropicDigestChain_WithSystemString(t *testing.T) { - parsed := &ParsedRequest{ - System: "You are a helpful assistant", - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - } + parsed := mustParseAnthropicDigestRequest(t, `{"system":"You are a helpful assistant","messages":[{"role":"user","content":"hello"}]}`) result := BuildAnthropicDigestChain(parsed) parts := splitChain(result) if len(parts) != 2 { @@ -79,14 +74,7 @@ func TestBuildAnthropicDigestChain_WithSystemString(t *testing.T) { } func TestBuildAnthropicDigestChain_WithSystemContentBlocks(t *testing.T) { - parsed := &ParsedRequest{ - System: []any{ - map[string]any{"type": "text", "text": "You are a helpful assistant"}, - }, - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - } + parsed := mustParseAnthropicDigestRequest(t, `{"system":[{"type":"text","text":"You are a helpful assistant"}],"messages":[{"role":"user","content":"hello"}]}`) result := BuildAnthropicDigestChain(parsed) parts := splitChain(result) if len(parts) != 2 { @@ -100,74 +88,33 @@ func TestBuildAnthropicDigestChain_WithSystemContentBlocks(t *testing.T) { func TestBuildAnthropicDigestChain_ConversationPrefixRelationship(t *testing.T) { // 核心测试：验证对话增长时链的前缀关系 // 上一轮的完整链一定是下一轮链的前缀 - system := "You are a helpful assistant" - - // 第 1 轮: system + user - round1 := &ParsedRequest{ - System: system, - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - } + round1 := mustParseAnthropicDigestRequest(t, `{"system":"You are a helpful assistant","messages":[{"role":"user","content":"hello"}]}`) chain1 := BuildAnthropicDigestChain(round1) - // 第 2 轮: system + user + assistant + user - round2 := &ParsedRequest{ - System: system, - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - map[string]any{"role": "assistant", "content": "hi there"}, - map[string]any{"role": "user", "content": "how are you?"}, - }, - } + round2 := mustParseAnthropicDigestRequest(t, `{"system":"You are a helpful assistant","messages":[{"role":"user","content":"hello"},{"role":"assistant","content":"hi there"},{"role":"user","content":"how are you?"}]}`) chain2 := BuildAnthropicDigestChain(round2) - // 第 3 轮: system + user + assistant + user + assistant + user - round3 := &ParsedRequest{ - System: system, - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - map[string]any{"role": "assistant", "content": "hi there"}, - map[string]any{"role": "user", "content": "how are you?"}, - map[string]any{"role": "assistant", "content": "I'm doing well"}, - map[string]any{"role": "user", "content": "great"}, - }, - } + round3 := mustParseAnthropicDigestRequest(t, `{"system":"You are a helpful assistant","messages":[{"role":"user","content":"hello"},{"role":"assistant","content":"hi there"},{"role":"user","content":"how are you?"},{"role":"assistant","content":"I'm doing well"},{"role":"user","content":"great"}]}`) chain3 := BuildAnthropicDigestChain(round3) t.Logf("Chain1: %s", chain1) t.Logf("Chain2: %s", chain2) t.Logf("Chain3: %s", chain3) - // chain1 是 chain2 的前缀 if !strings.HasPrefix(chain2, chain1) { t.Errorf("chain1 should be prefix of chain2:\n chain1: %s\n chain2: %s", chain1, chain2) } - - // chain2 是 chain3 的前缀 if !strings.HasPrefix(chain3, chain2) { t.Errorf("chain2 should be prefix of chain3:\n chain2: %s\n chain3: %s", chain2, chain3) } - - // chain1 也是 chain3 的前缀（传递性） if !strings.HasPrefix(chain3, chain1) { t.Errorf("chain1 should be prefix of chain3:\n chain1: %s\n chain3: %s", chain1, chain3) } } func TestBuildAnthropicDigestChain_DifferentSystemProducesDifferentChain(t *testing.T) { - parsed1 := &ParsedRequest{ - System: "System A", - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - } - parsed2 := &ParsedRequest{ - System: "System B", - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - } + parsed1 := mustParseAnthropicDigestRequest(t, `{"system":"System A","messages":[{"role":"user","content":"hello"}]}`) + parsed2 := mustParseAnthropicDigestRequest(t, `{"system":"System B","messages":[{"role":"user","content":"hello"}]}`) chain1 := BuildAnthropicDigestChain(parsed1) chain2 := BuildAnthropicDigestChain(parsed2) @@ -176,7 +123,6 @@ func TestBuildAnthropicDigestChain_DifferentSystemProducesDifferentChain(t *test t.Error("Different system prompts should produce different chains") } - // 但 user 部分的 hash 应该相同 parts1 := splitChain(chain1) parts2 := splitChain(chain2) if parts1[1] != parts2[1] { @@ -185,20 +131,8 @@ func TestBuildAnthropicDigestChain_DifferentSystemProducesDifferentChain(t *test } func TestBuildAnthropicDigestChain_DifferentContentProducesDifferentChain(t *testing.T) { - parsed1 := &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - map[string]any{"role": "assistant", "content": "ORIGINAL reply"}, - map[string]any{"role": "user", "content": "next"}, - }, - } - parsed2 := &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - map[string]any{"role": "assistant", "content": "TAMPERED reply"}, - map[string]any{"role": "user", "content": "next"}, - }, - } + parsed1 := mustParseAnthropicDigestRequest(t, `{"messages":[{"role":"user","content":"hello"},{"role":"assistant","content":"ORIGINAL reply"},{"role":"user","content":"next"}]}`) + parsed2 := mustParseAnthropicDigestRequest(t, `{"messages":[{"role":"user","content":"hello"},{"role":"assistant","content":"TAMPERED reply"},{"role":"user","content":"next"}]}`) chain1 := BuildAnthropicDigestChain(parsed1) chain2 := BuildAnthropicDigestChain(parsed2) @@ -209,24 +143,16 @@ func TestBuildAnthropicDigestChain_DifferentContentProducesDifferentChain(t *tes parts1 := splitChain(chain1) parts2 := splitChain(chain2) - // 第一个 user message hash 应该相同 if parts1[0] != parts2[0] { t.Error("First user message hash should be the same") } - // assistant reply hash 应该不同 if parts1[1] == parts2[1] { t.Error("Assistant reply hash should differ") } } func TestBuildAnthropicDigestChain_Deterministic(t *testing.T) { - parsed := &ParsedRequest{ - System: "test system", - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - map[string]any{"role": "assistant", "content": "hi"}, - }, - } + parsed := mustParseAnthropicDigestRequest(t, `{"system":"test system","messages":[{"role":"user","content":"hello"},{"role":"assistant","content":"hi"}]}`) chain1 := BuildAnthropicDigestChain(parsed) chain2 := BuildAnthropicDigestChain(parsed) @@ -236,6 +162,18 @@ func TestBuildAnthropicDigestChain_Deterministic(t *testing.T) { } } +func TestBuildAnthropicDigestChain_CanonicalJSON(t *testing.T) { + parsed1 := mustParseAnthropicDigestRequest(t, `{"system":[{"type":"text","text":"system"}],"messages":[{"role":"user","content":{"type":"text","text":"hello"}}]}`) + parsed2 := mustParseAnthropicDigestRequest(t, `{"system":[{"text":"system","type":"text"}],"messages":[{"role":"user","content":{"text":"hello","type":"text"}}]}`) + + chain1 := BuildAnthropicDigestChain(parsed1) + chain2 := BuildAnthropicDigestChain(parsed2) + + if chain1 != chain2 { + t.Errorf("semantically equivalent JSON should produce same chain: %s vs %s", chain1, chain2) + } +} + func TestGenerateAnthropicDigestSessionKey(t *testing.T) { tests := []struct { name string @@ -278,7 +216,6 @@ func TestGenerateAnthropicDigestSessionKey(t *testing.T) { }) } - // 验证不同 uuid 产生不同 sessionKey t.Run("different uuid different key", func(t *testing.T) { hash := "sameprefix123456" result1 := GenerateAnthropicDigestSessionKey(hash, "uuid0001-session-a") @@ -297,18 +234,7 @@ func TestAnthropicSessionTTL(t *testing.T) { } func TestBuildAnthropicDigestChain_ContentBlocks(t *testing.T) { - // 测试 content 为 content blocks 数组的情况 - parsed := &ParsedRequest{ - Messages: []any{ - map[string]any{ - "role": "user", - "content": []any{ - map[string]any{"type": "text", "text": "describe this image"}, - map[string]any{"type": "image", "source": map[string]any{"type": "base64"}}, - }, - }, - }, - } + parsed := mustParseAnthropicDigestRequest(t, `{"messages":[{"role":"user","content":[{"type":"text","text":"describe this image"},{"type":"image","source":{"type":"base64"}}]}]}`) result := BuildAnthropicDigestChain(parsed) parts := splitChain(result) if len(parts) != 1 { diff --git a/backend/internal/service/gateway_oauth_metadata_test.go b/backend/internal/service/gateway_oauth_metadata_test.go index ed6f1887145a..b172dc6eebdb 100644 --- a/backend/internal/service/gateway_oauth_metadata_test.go +++ b/backend/internal/service/gateway_oauth_metadata_test.go @@ -14,8 +14,6 @@ func TestBuildOAuthMetadataUserID_FallbackWithoutAccountUUID(t *testing.T) { Model: "claude-sonnet-4-5", Stream: true, MetadataUserID: "", - System: nil, - Messages: nil, } account := &Account{ diff --git a/backend/internal/service/gateway_request.go b/backend/internal/service/gateway_request.go index 819bb0a882cc..65718c27148f 100644 --- a/backend/internal/service/gateway_request.go +++ b/backend/internal/service/gateway_request.go @@ -51,6 +51,12 @@ type SessionContext struct { APIKeyID int64 } +type jsonRange struct { + start int // 原始请求体中的起始偏移（闭区间） + end int // 原始请求体中的结束偏移（开区间） + kind gjson.Type // JSON 值类型，用于调用方做轻量分支 +} + type RequestBodyRef struct { data []byte } @@ -80,6 +86,76 @@ func (b *RequestBodyRef) Replace(data []byte) { b.data = data } +func missingJSONRange() jsonRange { + return jsonRange{start: -1, end: -1} +} + +func rangeFromResult(r gjson.Result) jsonRange { + if r.Raw == "" || r.Index <= 0 { + return missingJSONRange() + } + end := r.Index + len(r.Raw) + if end < r.Index { + return missingJSONRange() + } + return jsonRange{start: r.Index, end: end, kind: r.Type} +} + +func (r jsonRange) exists() bool { + return r.start >= 0 && r.end >= r.start +} + +func clearGatewayRequestRanges(parsed *ParsedRequest) { + if parsed == nil { + return + } + parsed.HasSystem = false + parsed.systemRange = missingJSONRange() + parsed.messagesRange = missingJSONRange() +} + +func setGatewayRequestRanges(parsed *ParsedRequest, protocol string, jsonStr string) { + if parsed == nil { + return + } + switch protocol { + case domain.PlatformGemini: + if sysParts := gjson.Get(jsonStr, "systemInstruction.parts"); sysParts.Exists() && sysParts.IsArray() { + parsed.systemRange = rangeFromResult(sysParts) + } + if contents := gjson.Get(jsonStr, "contents"); contents.Exists() && contents.IsArray() { + parsed.messagesRange = rangeFromResult(contents) + } + default: + if sys := gjson.Get(jsonStr, "system"); sys.Exists() { + parsed.HasSystem = true + parsed.systemRange = rangeFromResult(sys) + } + if msgs := gjson.Get(jsonStr, "messages"); msgs.Exists() && msgs.IsArray() { + parsed.messagesRange = rangeFromResult(msgs) + } + } +} + +func refreshGatewayRequestRanges(parsed *ParsedRequest, protocol string) error { + if parsed == nil { + return fmt.Errorf("empty request body") + } + clearGatewayRequestRanges(parsed) + if parsed.Body == nil { + return fmt.Errorf("empty request body") + } + + bodyBytes := parsed.Body.Bytes() + if !gjson.ValidBytes(bodyBytes) { + return fmt.Errorf("invalid json") + } + + jsonStr := *(*string)(unsafe.Pointer(&bodyBytes)) + setGatewayRequestRanges(parsed, protocol, jsonStr) + return nil +} + // ParsedRequest 保存网关请求的预解析结果 // // 性能优化说明： @@ -93,18 +169,20 @@ func (b *RequestBodyRef) Replace(data []byte) { // 2. 将解析结果 ParsedRequest 传递给 Service 层 // 3. 避免重复 json.Unmarshal，减少 CPU 和内存开销 type ParsedRequest struct { - Body *RequestBodyRef // 原始请求体引用（保留用于转发） + Body *RequestBodyRef // 原始请求体引用（保留用于转发）；替换内容请走 ReplaceBody Model string // 请求的模型名称 Stream bool // 是否为流式请求 MetadataUserID string // metadata.user_id（用于会话亲和） - System any // system 字段内容 - Messages []any // messages 数组 HasSystem bool // 是否包含 system 字段（包含 null 也视为显式传入） ThinkingEnabled bool // 是否开启 thinking（部分平台会影响最终模型名） OutputEffort string // output_config.effort（Claude API 的推理强度控制） MaxTokens int // max_tokens 值（用于探测请求拦截） SessionContext *SessionContext // 可选：请求上下文区分因子（nil 时行为不变） + protocol string // 当前 Body 的协议格式，用于 Body 替换后刷新 raw range + systemRange jsonRange // system/systemInstruction.parts 的 raw JSON 范围，绑定 Body 当前内容 + messagesRange jsonRange // messages/contents 的 raw JSON 范围，绑定 Body 当前内容 + // GroupID 请求所属分组 ID（来自 API Key） GroupID *int64 @@ -173,7 +251,10 @@ func ParseGatewayRequest(body *RequestBodyRef, protocol string) (*ParsedRequest, jsonStr := *(*string)(unsafe.Pointer(&bodyBytes)) parsed := &ParsedRequest{ - Body: body, + Body: body, + protocol: protocol, + systemRange: missingJSONRange(), + messagesRange: missingJSONRange(), } // --- gjson 提取简单字段（避免完整 Unmarshal） --- @@ -219,58 +300,71 @@ func ParseGatewayRequest(body *RequestBodyRef, protocol string) (*ParsedRequest, } // --- system/messages 提取 --- - // 避免把整个 body Unmarshal 到 map（会产生大量 map/接口分配）。 - // 使用 gjson 抽取目标字段的 Raw，再对该子树进行 Unmarshal。 + // 只保存大字段 raw range，不默认反序列化成 []any/map[string]any 对象图。 + setGatewayRequestRanges(parsed, protocol, jsonStr) - switch protocol { - case domain.PlatformGemini: - // Gemini 原生格式: systemInstruction.parts / contents - if sysParts := gjson.Get(jsonStr, "systemInstruction.parts"); sysParts.Exists() && sysParts.IsArray() { - var parts []any - if err := json.Unmarshal(sliceRawFromBody(bodyBytes, sysParts), &parts); err != nil { - return nil, err - } - parsed.System = parts - } + return parsed, nil +} - if contents := gjson.Get(jsonStr, "contents"); contents.Exists() && contents.IsArray() { - var msgs []any - if err := json.Unmarshal(sliceRawFromBody(bodyBytes, contents), &msgs); err != nil { - return nil, err - } - parsed.Messages = msgs - } - default: - // Anthropic / OpenAI 格式: system / messages - // system 字段只要存在就视为显式提供（即使为 null）， - // 以避免客户端传 null 时被默认 system 误注入。 - if sys := gjson.Get(jsonStr, "system"); sys.Exists() { - parsed.HasSystem = true - switch sys.Type { - case gjson.Null: - parsed.System = nil - case gjson.String: - // 与 encoding/json 的 Unmarshal 行为一致：返回解码后的字符串。 - parsed.System = sys.String() - default: - var system any - if err := json.Unmarshal(sliceRawFromBody(bodyBytes, sys), &system); err != nil { - return nil, err - } - parsed.System = system - } - } +func (p *ParsedRequest) raw(r jsonRange) []byte { + if p == nil || p.Body == nil || !r.exists() { + return nil + } + body := p.Body.Bytes() + if r.end > len(body) { + return nil + } + return body[r.start:r.end] +} - if msgs := gjson.Get(jsonStr, "messages"); msgs.Exists() && msgs.IsArray() { - var messages []any - if err := json.Unmarshal(sliceRawFromBody(bodyBytes, msgs), &messages); err != nil { - return nil, err - } - parsed.Messages = messages - } +func (p *ParsedRequest) SystemRaw() []byte { + return p.raw(p.systemRange) +} + +func (p *ParsedRequest) MessagesRaw() []byte { + return p.raw(p.messagesRange) +} + +func (p *ParsedRequest) DecodeSystem(dst any) error { + raw := p.SystemRaw() + if len(raw) == 0 { + return nil } + return json.Unmarshal(raw, dst) +} - return parsed, nil +func (p *ParsedRequest) DecodeMessages(dst any) error { + raw := p.MessagesRaw() + if len(raw) == 0 { + return nil + } + return json.Unmarshal(raw, dst) +} + +func (p *ParsedRequest) SystemValue() (any, bool) { + raw := p.SystemRaw() + if len(raw) == 0 { + return nil, false + } + var system any + if err := json.Unmarshal(raw, &system); err != nil { + return nil, false + } + return system, true +} + +func (p *ParsedRequest) ReplaceBody(data []byte) { + if p == nil { + return + } + if p.Body == nil { + p.Body = NewRequestBodyRef(data) + } else { + p.Body.Replace(data) + } + if err := refreshGatewayRequestRanges(p, p.protocol); err != nil { + clearGatewayRequestRanges(p) + } } // sliceRawFromBody 返回 Result.Raw 对应的原始字节切片。 diff --git a/backend/internal/service/gateway_request_test.go b/backend/internal/service/gateway_request_test.go index d415b871a0fc..288c031c37d4 100644 --- a/backend/internal/service/gateway_request_test.go +++ b/backend/internal/service/gateway_request_test.go @@ -10,6 +10,7 @@ import ( "github.com/Wei-Shaw/sub2api/internal/domain" "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" ) func TestParseGatewayRequest(t *testing.T) { @@ -20,8 +21,8 @@ func TestParseGatewayRequest(t *testing.T) { require.True(t, parsed.Stream) require.Equal(t, "session_123e4567-e89b-12d3-a456-426614174000", parsed.MetadataUserID) require.True(t, parsed.HasSystem) - require.NotNil(t, parsed.System) - require.Len(t, parsed.Messages, 1) + require.NotEmpty(t, parsed.SystemRaw()) + require.NotEmpty(t, parsed.MessagesRaw()) require.False(t, parsed.ThinkingEnabled) } @@ -61,7 +62,7 @@ func TestParseGatewayRequest_SystemNull(t *testing.T) { require.NoError(t, err) // 显式传入 system:null 也应视为“字段已存在”，避免默认 system 被注入。 require.True(t, parsed.HasSystem) - require.Nil(t, parsed.System) + require.Equal(t, []byte("null"), parsed.SystemRaw()) } func TestParseGatewayRequest_InvalidModelType(t *testing.T) { @@ -88,9 +89,9 @@ func TestParseGatewayRequest_GeminiContents(t *testing.T) { }`) parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), domain.PlatformGemini) require.NoError(t, err) - require.Len(t, parsed.Messages, 3, "should parse contents as Messages") + require.Len(t, gjson.ParseBytes(parsed.MessagesRaw()).Array(), 3, "should parse contents as Messages") require.False(t, parsed.HasSystem, "Gemini format should not set HasSystem") - require.Nil(t, parsed.System, "no systemInstruction means nil System") + require.Nil(t, parsed.SystemRaw(), "no systemInstruction means nil System") } func TestParseGatewayRequest_GeminiSystemInstruction(t *testing.T) { @@ -104,14 +105,11 @@ func TestParseGatewayRequest_GeminiSystemInstruction(t *testing.T) { }`) parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), domain.PlatformGemini) require.NoError(t, err) - require.NotNil(t, parsed.System, "should parse systemInstruction.parts as System") - parts, ok := parsed.System.([]any) - require.True(t, ok) - require.Len(t, parts, 1) - partMap, ok := parts[0].(map[string]any) - require.True(t, ok) - require.Equal(t, "You are a helpful assistant.", partMap["text"]) - require.Len(t, parsed.Messages, 1) + system := gjson.ParseBytes(parsed.SystemRaw()) + require.True(t, system.IsArray(), "should parse systemInstruction.parts as System") + require.Len(t, system.Array(), 1) + require.Equal(t, "You are a helpful assistant.", system.Get("0.text").String()) + require.Len(t, gjson.ParseBytes(parsed.MessagesRaw()).Array(), 1) } func TestParseGatewayRequest_GeminiWithModel(t *testing.T) { @@ -122,7 +120,7 @@ func TestParseGatewayRequest_GeminiWithModel(t *testing.T) { parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), domain.PlatformGemini) require.NoError(t, err) require.Equal(t, "gemini-2.5-pro", parsed.Model) - require.Len(t, parsed.Messages, 1) + require.Len(t, gjson.ParseBytes(parsed.MessagesRaw()).Array(), 1) } func TestParseGatewayRequest_GeminiIgnoresAnthropicFields(t *testing.T) { @@ -135,22 +133,22 @@ func TestParseGatewayRequest_GeminiIgnoresAnthropicFields(t *testing.T) { parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), domain.PlatformGemini) require.NoError(t, err) require.False(t, parsed.HasSystem, "Gemini protocol should not parse Anthropic system field") - require.Nil(t, parsed.System, "no systemInstruction = nil System") - require.Len(t, parsed.Messages, 1, "should use contents, not messages") + require.Nil(t, parsed.SystemRaw(), "no systemInstruction = nil System") + require.Len(t, gjson.ParseBytes(parsed.MessagesRaw()).Array(), 1, "should use contents, not messages") } func TestParseGatewayRequest_GeminiEmptyContents(t *testing.T) { body := []byte(`{"contents": []}`) parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), domain.PlatformGemini) require.NoError(t, err) - require.Empty(t, parsed.Messages) + require.Empty(t, gjson.ParseBytes(parsed.MessagesRaw()).Array()) } func TestParseGatewayRequest_GeminiNoContents(t *testing.T) { body := []byte(`{"model": "gemini-2.5-flash"}`) parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), domain.PlatformGemini) require.NoError(t, err) - require.Nil(t, parsed.Messages) + require.Nil(t, parsed.MessagesRaw()) require.Equal(t, "gemini-2.5-flash", parsed.Model) } @@ -165,11 +163,10 @@ func TestParseGatewayRequest_AnthropicIgnoresGeminiFields(t *testing.T) { parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), domain.PlatformAnthropic) require.NoError(t, err) require.True(t, parsed.HasSystem) - require.Equal(t, "real system", parsed.System) - require.Len(t, parsed.Messages, 1) - msg, ok := parsed.Messages[0].(map[string]any) - require.True(t, ok) - require.Equal(t, "real content", msg["content"]) + require.Equal(t, "real system", gjson.ParseBytes(parsed.SystemRaw()).String()) + messages := gjson.ParseBytes(parsed.MessagesRaw()).Array() + require.Len(t, messages, 1) + require.Equal(t, "real content", messages[0].Get("content").String()) } func TestFilterThinkingBlocks(t *testing.T) { @@ -970,10 +967,10 @@ func TestParseGatewayRequest_OptionalFieldsMissing(t *testing.T) { require.Equal(t, tt.wantMaxTokens, parsed.MaxTokens) if tt.wantMessagesNil { - require.Nil(t, parsed.Messages) + require.Nil(t, parsed.MessagesRaw()) } if tt.wantMessagesLen > 0 { - require.Len(t, parsed.Messages, tt.wantMessagesLen) + require.Len(t, gjson.ParseBytes(parsed.MessagesRaw()).Array(), tt.wantMessagesLen) } }) } @@ -1087,25 +1084,8 @@ func parseGatewayRequestOld(body []byte, protocol string) (*ParsedRequest, error } } - // system / messages（按协议分支） - switch protocol { - case domain.PlatformGemini: - if sysInst, ok := req["systemInstruction"].(map[string]any); ok { - if parts, ok := sysInst["parts"].([]any); ok { - parsed.System = parts - } - } - if contents, ok := req["contents"].([]any); ok { - parsed.Messages = contents - } - default: - if system, ok := req["system"]; ok { - parsed.HasSystem = true - parsed.System = system - } - if messages, ok := req["messages"].([]any); ok { - parsed.Messages = messages - } + if err := refreshGatewayRequestRanges(parsed, protocol); err != nil { + return nil, err } return parsed, nil diff --git a/backend/internal/service/gateway_service.go b/backend/internal/service/gateway_service.go index 8f55bf13c6d3..438592e3f45b 100644 --- a/backend/internal/service/gateway_service.go +++ b/backend/internal/service/gateway_service.go @@ -748,31 +748,10 @@ func (s *GatewayService) GenerateSessionHash(parsed *ParsedRequest) string { _, _ = combined.WriteString(strconv.FormatInt(parsed.SessionContext.APIKeyID, 10)) _, _ = combined.WriteString("|") } - if parsed.System != nil { - systemText := s.extractTextFromSystem(parsed.System) - if systemText != "" { - _, _ = combined.WriteString(systemText) - } - } - for _, msg := range parsed.Messages { - if m, ok := msg.(map[string]any); ok { - if content, exists := m["content"]; exists { - // Anthropic: messages[].content - if msgText := s.extractTextFromContent(content); msgText != "" { - _, _ = combined.WriteString(msgText) - } - } else if parts, ok := m["parts"].([]any); ok { - // Gemini: contents[].parts[].text - for _, part := range parts { - if partMap, ok := part.(map[string]any); ok { - if text, ok := partMap["text"].(string); ok { - _, _ = combined.WriteString(text) - } - } - } - } - } + if systemText := extractTextFromSystemRaw(parsed.SystemRaw()); systemText != "" { + _, _ = combined.WriteString(systemText) } + appendMessageTextsFromRaw(&combined, parsed.MessagesRaw()) if combined.Len() > 0 { hash := s.hashContent(combined.String()) slog.Info("sticky.hash_source", @@ -847,80 +826,125 @@ func (s *GatewayService) extractCacheableContent(parsed *ParsedRequest) string { return "" } - var builder strings.Builder - - // 检查 system 中的 cacheable 内容 - if system, ok := parsed.System.([]any); ok { - for _, part := range system { - if partMap, ok := part.(map[string]any); ok { - if cc, ok := partMap["cache_control"].(map[string]any); ok { - if cc["type"] == "ephemeral" { - if text, ok := partMap["text"].(string); ok { - _, _ = builder.WriteString(text) - } - } - } - } - } + systemText := extractCacheableTextFromSystemRaw(parsed.SystemRaw()) + if messageText := extractCacheableTextFromMessagesRaw(parsed.MessagesRaw()); messageText != "" { + return messageText } - systemText := builder.String() + return systemText +} - // 检查 messages 中的 cacheable 内容 - for _, msg := range parsed.Messages { - if msgMap, ok := msg.(map[string]any); ok { - if msgContent, ok := msgMap["content"].([]any); ok { - for _, part := range msgContent { - if partMap, ok := part.(map[string]any); ok { - if cc, ok := partMap["cache_control"].(map[string]any); ok { - if cc["type"] == "ephemeral" { - return s.extractTextFromContent(msgMap["content"]) - } - } - } - } - } +func extractTextFromSystemRaw(raw []byte) string { + system := gjson.ParseBytes(raw) + switch system.Type { + case gjson.String: + return system.String() + case gjson.JSON: + if !system.IsArray() { + return "" } + var builder strings.Builder + system.ForEach(func(_, part gjson.Result) bool { + if text := part.Get("text").String(); text != "" { + _, _ = builder.WriteString(text) + } + return true + }) + return builder.String() } - - return systemText + return "" } -func (s *GatewayService) extractTextFromSystem(system any) string { - switch v := system.(type) { - case string: - return v - case []any: - var texts []string - for _, part := range v { - if partMap, ok := part.(map[string]any); ok { - if text, ok := partMap["text"].(string); ok { - texts = append(texts, text) +func extractTextFromContentRaw(content gjson.Result) string { + switch content.Type { + case gjson.String: + return content.String() + case gjson.JSON: + if !content.IsArray() { + return "" + } + var builder strings.Builder + content.ForEach(func(_, part gjson.Result) bool { + if part.Get("type").String() == "text" { + if text := part.Get("text").String(); text != "" { + _, _ = builder.WriteString(text) } } - } - return strings.Join(texts, "") + return true + }) + return builder.String() } return "" } -func (s *GatewayService) extractTextFromContent(content any) string { - switch v := content.(type) { - case string: - return v - case []any: - var texts []string - for _, part := range v { - if partMap, ok := part.(map[string]any); ok { - if partMap["type"] == "text" { - if text, ok := partMap["text"].(string); ok { - texts = append(texts, text) - } +func appendMessageTextsFromRaw(builder *strings.Builder, raw []byte) { + if builder == nil || len(raw) == 0 { + return + } + messages := gjson.ParseBytes(raw) + if !messages.IsArray() { + return + } + messages.ForEach(func(_, msg gjson.Result) bool { + if content := msg.Get("content"); content.Exists() { + _, _ = builder.WriteString(extractTextFromContentRaw(content)) + return true + } + parts := msg.Get("parts") + if parts.IsArray() { + parts.ForEach(func(_, part gjson.Result) bool { + if text := part.Get("text").String(); text != "" { + _, _ = builder.WriteString(text) } + return true + }) + } + return true + }) +} + +func extractCacheableTextFromSystemRaw(raw []byte) string { + system := gjson.ParseBytes(raw) + if !system.IsArray() { + return "" + } + var builder strings.Builder + system.ForEach(func(_, part gjson.Result) bool { + if part.Get("cache_control.type").String() == "ephemeral" { + if text := part.Get("text").String(); text != "" { + _, _ = builder.WriteString(text) } } - return strings.Join(texts, "") + return true + }) + return builder.String() +} + +func extractCacheableTextFromMessagesRaw(raw []byte) string { + messages := gjson.ParseBytes(raw) + if !messages.IsArray() { + return "" } - return "" + var text string + messages.ForEach(func(_, msg gjson.Result) bool { + content := msg.Get("content") + if !content.IsArray() { + return true + } + found := false + content.ForEach(func(_, part gjson.Result) bool { + if part.Get("cache_control.type").String() == "ephemeral" { + found = true + return false + } + return true + }) + if found { + text = extractTextFromContentRaw(content) + return false + } + return true + }) + return text } func (s *GatewayService) hashContent(content string) string { @@ -1284,7 +1308,7 @@ func (s *GatewayService) applyClaudeCodeOAuthMimicryToBody( systemRewritten := false if !strings.Contains(strings.ToLower(model), "haiku") { - body = rewriteSystemForNonClaudeCode(body, systemRaw) + body = rewriteSystemForNonClaudeCode(body, normalizeSystemParam(systemRaw)) systemRewritten = true } @@ -4474,7 +4498,8 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A // Parrot 的 transform_request 从不检查客户端 system 内容，直接覆盖。 systemRewritten := false if !strings.Contains(strings.ToLower(reqModel), "haiku") { - body = rewriteSystemForNonClaudeCode(body, parsed.System) + systemRaw, _ := parsed.SystemValue() + body = rewriteSystemForNonClaudeCode(body, systemRaw) systemRewritten = true } diff --git a/backend/internal/service/gateway_service_benchmark_test.go b/backend/internal/service/gateway_service_benchmark_test.go index 5637680b3f07..8b30cb249e3d 100644 --- a/backend/internal/service/gateway_service_benchmark_test.go +++ b/backend/internal/service/gateway_service_benchmark_test.go @@ -2,6 +2,7 @@ package service import ( "strconv" + "strings" "testing" ) @@ -34,17 +35,20 @@ func BenchmarkExtractCacheableContent_System(b *testing.B) { } func buildSystemCacheableRequest(parts int) *ParsedRequest { - systemParts := make([]any, 0, parts) + var builder strings.Builder + builder.WriteString(`{"system":[`) for i := 0; i < parts; i++ { - systemParts = append(systemParts, map[string]any{ - "text": "system_part_" + strconv.Itoa(i), - "cache_control": map[string]any{ - "type": "ephemeral", - }, - }) + if i > 0 { + builder.WriteByte(',') + } + builder.WriteString(`{"text":"system_part_`) + builder.WriteString(strconv.Itoa(i)) + builder.WriteString(`","cache_control":{"type":"ephemeral"}}`) } - return &ParsedRequest{ - System: systemParts, - HasSystem: true, + builder.WriteString(`]}`) + parsed, err := ParseGatewayRequest(NewRequestBodyRef([]byte(builder.String())), "") + if err != nil { + panic(err) } + return parsed } diff --git a/backend/internal/service/generate_session_hash_test.go b/backend/internal/service/generate_session_hash_test.go index 8f3258b72c71..5ed3f0ae99d5 100644 --- a/backend/internal/service/generate_session_hash_test.go +++ b/backend/internal/service/generate_session_hash_test.go @@ -3,12 +3,67 @@ package service import ( + "encoding/json" "testing" + "github.com/Wei-Shaw/sub2api/internal/domain" "github.com/stretchr/testify/require" ) -// ============ 基础优先级测试 ============ +func mustParseSessionHashRequest(t *testing.T, body string, ctx *SessionContext) *ParsedRequest { + t.Helper() + parsed, err := ParseGatewayRequest(NewRequestBodyRef([]byte(body)), domain.PlatformAnthropic) + require.NoError(t, err) + parsed.SessionContext = ctx + return parsed +} + +func mustParseGeminiSessionHashRequest(t *testing.T, body string, ctx *SessionContext) *ParsedRequest { + t.Helper() + parsed, err := ParseGatewayRequest(NewRequestBodyRef([]byte(body)), domain.PlatformGemini) + require.NoError(t, err) + parsed.SessionContext = ctx + return parsed +} + +func anthropicSessionBody(system any, messages []any, metadataUserID string) string { + body := map[string]any{} + if system != nil { + body["system"] = system + } + if messages != nil { + body["messages"] = messages + } + if metadataUserID != "" { + body["metadata"] = map[string]any{"user_id": metadataUserID} + } + data, _ := json.Marshal(body) + return string(data) +} + +func geminiSessionBody(systemParts []any, contents []any) string { + body := map[string]any{} + if systemParts != nil { + body["systemInstruction"] = map[string]any{"parts": systemParts} + } + if contents != nil { + body["contents"] = contents + } + data, _ := json.Marshal(body) + return string(data) +} + +func msg(role string, content any) map[string]any { + return map[string]any{"role": role, "content": content} +} + +func geminiMsg(role string, texts ...string) map[string]any { + parts := make([]any, 0, len(texts)) + for _, text := range texts { + parts = append(parts, map[string]any{"text": text}) + } + return map[string]any{"role": role, "parts": parts} +} func TestGenerateSessionHash_NilParsedRequest(t *testing.T) { svc := &GatewayService{} @@ -22,37 +77,17 @@ func TestGenerateSessionHash_EmptyRequest(t *testing.T) { func TestGenerateSessionHash_MetadataHasHighestPriority(t *testing.T) { svc := &GatewayService{} - - parsed := &ParsedRequest{ - MetadataUserID: "user_a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2_account__session_123e4567-e89b-12d3-a456-426614174000", - System: "You are a helpful assistant.", - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - } + metadata := "user_a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2_account__session_123e4567-e89b-12d3-a456-426614174000" + parsed := mustParseSessionHashRequest(t, anthropicSessionBody("You are a helpful assistant.", []any{msg("user", "hello")}, metadata), nil) hash := svc.GenerateSessionHash(parsed) require.Equal(t, "123e4567-e89b-12d3-a456-426614174000", hash, "metadata session_id should have highest priority") } -// ============ System + Messages 基础测试 ============ - func TestGenerateSessionHash_SystemPlusMessages(t *testing.T) { svc := &GatewayService{} - - withSystem := &ParsedRequest{ - System: "You are a helpful assistant.", - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - } - withoutSystem := &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - } + withSystem := mustParseSessionHashRequest(t, anthropicSessionBody("You are a helpful assistant.", []any{msg("user", "hello")}, ""), nil) + withoutSystem := mustParseSessionHashRequest(t, anthropicSessionBody(nil, []any{msg("user", "hello")}, ""), nil) h1 := svc.GenerateSessionHash(withSystem) h2 := svc.GenerateSessionHash(withoutSystem) @@ -63,32 +98,16 @@ func TestGenerateSessionHash_SystemPlusMessages(t *testing.T) { func TestGenerateSessionHash_SystemOnlyProducesHash(t *testing.T) { svc := &GatewayService{} + parsed := mustParseSessionHashRequest(t, anthropicSessionBody("You are a helpful assistant.", nil, ""), nil) - parsed := &ParsedRequest{ - System: "You are a helpful assistant.", - HasSystem: true, - } hash := svc.GenerateSessionHash(parsed) require.NotEmpty(t, hash, "system prompt alone should produce a hash as part of full digest") } func TestGenerateSessionHash_DifferentSystemsSameMessages(t *testing.T) { svc := &GatewayService{} - - parsed1 := &ParsedRequest{ - System: "You are assistant A.", - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - } - parsed2 := &ParsedRequest{ - System: "You are assistant B.", - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - } + parsed1 := mustParseSessionHashRequest(t, anthropicSessionBody("You are assistant A.", []any{msg("user", "hello")}, ""), nil) + parsed2 := mustParseSessionHashRequest(t, anthropicSessionBody("You are assistant B.", []any{msg("user", "hello")}, ""), nil) h1 := svc.GenerateSessionHash(parsed1) h2 := svc.GenerateSessionHash(parsed2) @@ -97,16 +116,8 @@ func TestGenerateSessionHash_DifferentSystemsSameMessages(t *testing.T) { func TestGenerateSessionHash_SameSystemSameMessages(t *testing.T) { svc := &GatewayService{} - mk := func() *ParsedRequest { - return &ParsedRequest{ - System: "You are a helpful assistant.", - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - map[string]any{"role": "assistant", "content": "hi"}, - }, - } + return mustParseSessionHashRequest(t, anthropicSessionBody("You are a helpful assistant.", []any{msg("user", "hello"), msg("assistant", "hi")}, ""), nil) } h1 := svc.GenerateSessionHash(mk()) @@ -116,53 +127,19 @@ func TestGenerateSessionHash_SameSystemSameMessages(t *testing.T) { func TestGenerateSessionHash_DifferentMessagesProduceDifferentHash(t *testing.T) { svc := &GatewayService{} - - parsed1 := &ParsedRequest{ - System: "You are a helpful assistant.", - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "help me with Go"}, - }, - } - parsed2 := &ParsedRequest{ - System: "You are a helpful assistant.", - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "help me with Python"}, - }, - } + parsed1 := mustParseSessionHashRequest(t, anthropicSessionBody("You are a helpful assistant.", []any{msg("user", "help me with Go")}, ""), nil) + parsed2 := mustParseSessionHashRequest(t, anthropicSessionBody("You are a helpful assistant.", []any{msg("user", "help me with Python")}, ""), nil) h1 := svc.GenerateSessionHash(parsed1) h2 := svc.GenerateSessionHash(parsed2) require.NotEqual(t, h1, h2, "same system but different messages should produce different hashes") } -// ============ SessionContext 核心测试 ============ - func TestGenerateSessionHash_DifferentSessionContextProducesDifferentHash(t *testing.T) { svc := &GatewayService{} - - // 相同消息 + 不同 SessionContext → 不同 hash（解决碰撞问题的核心场景） - parsed1 := &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - SessionContext: &SessionContext{ - ClientIP: "192.168.1.1", - UserAgent: "Mozilla/5.0", - APIKeyID: 100, - }, - } - parsed2 := &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - SessionContext: &SessionContext{ - ClientIP: "10.0.0.1", - UserAgent: "curl/7.0", - APIKeyID: 200, - }, - } + body := anthropicSessionBody(nil, []any{msg("user", "hello")}, "") + parsed1 := mustParseSessionHashRequest(t, body, &SessionContext{ClientIP: "192.168.1.1", UserAgent: "Mozilla/5.0", APIKeyID: 100}) + parsed2 := mustParseSessionHashRequest(t, body, &SessionContext{ClientIP: "10.0.0.1", UserAgent: "curl/7.0", APIKeyID: 200}) h1 := svc.GenerateSessionHash(parsed1) h2 := svc.GenerateSessionHash(parsed2) @@ -173,19 +150,9 @@ func TestGenerateSessionHash_DifferentSessionContextProducesDifferentHash(t *tes func TestGenerateSessionHash_SameSessionContextProducesSameHash(t *testing.T) { svc := &GatewayService{} - - mk := func() *ParsedRequest { - return &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - SessionContext: &SessionContext{ - ClientIP: "192.168.1.1", - UserAgent: "Mozilla/5.0", - APIKeyID: 100, - }, - } - } + ctx := &SessionContext{ClientIP: "192.168.1.1", UserAgent: "Mozilla/5.0", APIKeyID: 100} + body := anthropicSessionBody(nil, []any{msg("user", "hello")}, "") + mk := func() *ParsedRequest { return mustParseSessionHashRequest(t, body, ctx) } h1 := svc.GenerateSessionHash(mk()) h2 := svc.GenerateSessionHash(mk()) @@ -194,35 +161,17 @@ func TestGenerateSessionHash_SameSessionContextProducesSameHash(t *testing.T) { func TestGenerateSessionHash_MetadataOverridesSessionContext(t *testing.T) { svc := &GatewayService{} - - parsed := &ParsedRequest{ - MetadataUserID: "user_a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2_account__session_123e4567-e89b-12d3-a456-426614174000", - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - SessionContext: &SessionContext{ - ClientIP: "192.168.1.1", - UserAgent: "Mozilla/5.0", - APIKeyID: 100, - }, - } + metadata := "user_a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2_account__session_123e4567-e89b-12d3-a456-426614174000" + parsed := mustParseSessionHashRequest(t, anthropicSessionBody(nil, []any{msg("user", "hello")}, metadata), &SessionContext{ClientIP: "192.168.1.1", UserAgent: "Mozilla/5.0", APIKeyID: 100}) hash := svc.GenerateSessionHash(parsed) - require.Equal(t, "123e4567-e89b-12d3-a456-426614174000", hash, - "metadata session_id should take priority over SessionContext") + require.Equal(t, "123e4567-e89b-12d3-a456-426614174000", hash, "metadata session_id should take priority over SessionContext") } func TestGenerateSessionHash_MetadataJSON_HasHighestPriority(t *testing.T) { svc := &GatewayService{} - - parsed := &ParsedRequest{ - MetadataUserID: `{"device_id":"a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2","account_uuid":"","session_id":"c72554f2-1234-5678-abcd-123456789abc"}`, - System: "You are a helpful assistant.", - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - } + metadata := `{"device_id":"a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2","account_uuid":"","session_id":"c72554f2-1234-5678-abcd-123456789abc"}` + parsed := mustParseSessionHashRequest(t, anthropicSessionBody("You are a helpful assistant.", []any{msg("user", "hello")}, metadata), nil) hash := svc.GenerateSessionHash(parsed) require.Equal(t, "c72554f2-1234-5678-abcd-123456789abc", hash, "JSON format metadata session_id should have highest priority") @@ -230,69 +179,25 @@ func TestGenerateSessionHash_MetadataJSON_HasHighestPriority(t *testing.T) { func TestGenerateSessionHash_NilSessionContextBackwardCompatible(t *testing.T) { svc := &GatewayService{} - - withCtx := &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - SessionContext: nil, - } - withoutCtx := &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - } + body := anthropicSessionBody(nil, []any{msg("user", "hello")}, "") + withCtx := mustParseSessionHashRequest(t, body, nil) + withoutCtx := mustParseSessionHashRequest(t, body, nil) h1 := svc.GenerateSessionHash(withCtx) h2 := svc.GenerateSessionHash(withoutCtx) require.Equal(t, h1, h2, "nil SessionContext should produce same hash as no SessionContext") } -// ============ 多轮连续会话测试 ============ - func TestGenerateSessionHash_ContinuousConversation_HashChangesWithMessages(t *testing.T) { svc := &GatewayService{} - ctx := &SessionContext{ClientIP: "1.2.3.4", UserAgent: "test", APIKeyID: 1} - - // 模拟连续会话：每增加一轮对话，hash 应该不同（内容累积变化） - round1 := &ParsedRequest{ - System: "You are a helpful assistant.", - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - SessionContext: ctx, - } - - round2 := &ParsedRequest{ - System: "You are a helpful assistant.", - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - map[string]any{"role": "assistant", "content": "Hi there!"}, - map[string]any{"role": "user", "content": "How are you?"}, - }, - SessionContext: ctx, - } - - round3 := &ParsedRequest{ - System: "You are a helpful assistant.", - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - map[string]any{"role": "assistant", "content": "Hi there!"}, - map[string]any{"role": "user", "content": "How are you?"}, - map[string]any{"role": "assistant", "content": "I'm doing well!"}, - map[string]any{"role": "user", "content": "Tell me a joke"}, - }, - SessionContext: ctx, - } + round1 := mustParseSessionHashRequest(t, anthropicSessionBody("You are a helpful assistant.", []any{msg("user", "hello")}, ""), ctx) + round2 := mustParseSessionHashRequest(t, anthropicSessionBody("You are a helpful assistant.", []any{msg("user", "hello"), msg("assistant", "Hi there!"), msg("user", "How are you?")}, ""), ctx) + round3 := mustParseSessionHashRequest(t, anthropicSessionBody("You are a helpful assistant.", []any{msg("user", "hello"), msg("assistant", "Hi there!"), msg("user", "How are you?"), msg("assistant", "I'm doing well!"), msg("user", "Tell me a joke")}, ""), ctx) h1 := svc.GenerateSessionHash(round1) h2 := svc.GenerateSessionHash(round2) h3 := svc.GenerateSessionHash(round3) - require.NotEmpty(t, h1) require.NotEmpty(t, h2) require.NotEmpty(t, h3) @@ -303,62 +208,20 @@ func TestGenerateSessionHash_ContinuousConversation_HashChangesWithMessages(t *t func TestGenerateSessionHash_ContinuousConversation_SameRoundSameHash(t *testing.T) { svc := &GatewayService{} - ctx := &SessionContext{ClientIP: "1.2.3.4", UserAgent: "test", APIKeyID: 1} - - // 同一轮对话重复请求（如重试）应产生相同 hash - mk := func() *ParsedRequest { - return &ParsedRequest{ - System: "You are a helpful assistant.", - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - map[string]any{"role": "assistant", "content": "Hi there!"}, - map[string]any{"role": "user", "content": "How are you?"}, - }, - SessionContext: ctx, - } - } + body := anthropicSessionBody("You are a helpful assistant.", []any{msg("user", "hello"), msg("assistant", "Hi there!"), msg("user", "How are you?")}, "") + mk := func() *ParsedRequest { return mustParseSessionHashRequest(t, body, ctx) } h1 := svc.GenerateSessionHash(mk()) h2 := svc.GenerateSessionHash(mk()) require.Equal(t, h1, h2, "same conversation state should produce identical hash on retry") } -// ============ 消息回退测试 ============ - func TestGenerateSessionHash_MessageRollback(t *testing.T) { svc := &GatewayService{} - ctx := &SessionContext{ClientIP: "1.2.3.4", UserAgent: "test", APIKeyID: 1} - - // 模拟消息回退：用户删掉最后一轮再重发 - original := &ParsedRequest{ - System: "System prompt", - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "msg1"}, - map[string]any{"role": "assistant", "content": "reply1"}, - map[string]any{"role": "user", "content": "msg2"}, - map[string]any{"role": "assistant", "content": "reply2"}, - map[string]any{"role": "user", "content": "msg3"}, - }, - SessionContext: ctx, - } - - // 回退到 msg2 后，用新的 msg3 替代 - rollback := &ParsedRequest{ - System: "System prompt", - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "msg1"}, - map[string]any{"role": "assistant", "content": "reply1"}, - map[string]any{"role": "user", "content": "msg2"}, - map[string]any{"role": "assistant", "content": "reply2"}, - map[string]any{"role": "user", "content": "different msg3"}, - }, - SessionContext: ctx, - } + original := mustParseSessionHashRequest(t, anthropicSessionBody("System prompt", []any{msg("user", "msg1"), msg("assistant", "reply1"), msg("user", "msg2"), msg("assistant", "reply2"), msg("user", "msg3")}, ""), ctx) + rollback := mustParseSessionHashRequest(t, anthropicSessionBody("System prompt", []any{msg("user", "msg1"), msg("assistant", "reply1"), msg("user", "msg2"), msg("assistant", "reply2"), msg("user", "different msg3")}, ""), ctx) hOrig := svc.GenerateSessionHash(original) hRollback := svc.GenerateSessionHash(rollback) @@ -367,58 +230,19 @@ func TestGenerateSessionHash_MessageRollback(t *testing.T) { func TestGenerateSessionHash_MessageRollbackSameContent(t *testing.T) { svc := &GatewayService{} - ctx := &SessionContext{ClientIP: "1.2.3.4", UserAgent: "test", APIKeyID: 1} - - // 回退后重新发送相同内容 → 相同 hash（合理的粘性恢复） - mk := func() *ParsedRequest { - return &ParsedRequest{ - System: "System prompt", - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "msg1"}, - map[string]any{"role": "assistant", "content": "reply1"}, - map[string]any{"role": "user", "content": "msg2"}, - }, - SessionContext: ctx, - } - } + body := anthropicSessionBody("System prompt", []any{msg("user", "msg1"), msg("assistant", "reply1"), msg("user", "msg2")}, "") + mk := func() *ParsedRequest { return mustParseSessionHashRequest(t, body, ctx) } h1 := svc.GenerateSessionHash(mk()) h2 := svc.GenerateSessionHash(mk()) require.Equal(t, h1, h2, "rollback and resend same content should produce same hash") } -// ============ 相同 System、不同用户消息 ============ - func TestGenerateSessionHash_SameSystemDifferentUsers(t *testing.T) { svc := &GatewayService{} - - // 两个不同用户使用相同 system prompt 但发送不同消息 - user1 := &ParsedRequest{ - System: "You are a code reviewer.", - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "Review this Go code"}, - }, - SessionContext: &SessionContext{ - ClientIP: "1.1.1.1", - UserAgent: "vscode", - APIKeyID: 1, - }, - } - user2 := &ParsedRequest{ - System: "You are a code reviewer.", - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "Review this Python code"}, - }, - SessionContext: &SessionContext{ - ClientIP: "2.2.2.2", - UserAgent: "vscode", - APIKeyID: 2, - }, - } + user1 := mustParseSessionHashRequest(t, anthropicSessionBody("You are a code reviewer.", []any{msg("user", "Review this Go code")}, ""), &SessionContext{ClientIP: "1.1.1.1", UserAgent: "vscode", APIKeyID: 1}) + user2 := mustParseSessionHashRequest(t, anthropicSessionBody("You are a code reviewer.", []any{msg("user", "Review this Python code")}, ""), &SessionContext{ClientIP: "2.2.2.2", UserAgent: "vscode", APIKeyID: 2}) h1 := svc.GenerateSessionHash(user1) h2 := svc.GenerateSessionHash(user2) @@ -427,55 +251,20 @@ func TestGenerateSessionHash_SameSystemDifferentUsers(t *testing.T) { func TestGenerateSessionHash_SameSystemSameMessageDifferentContext(t *testing.T) { svc := &GatewayService{} - - // 这是修复的核心场景：两个不同用户发送完全相同的 system + messages（如 "hello"） - // 有了 SessionContext 后应该产生不同 hash - user1 := &ParsedRequest{ - System: "You are a helpful assistant.", - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - SessionContext: &SessionContext{ - ClientIP: "1.1.1.1", - UserAgent: "Mozilla/5.0", - APIKeyID: 10, - }, - } - user2 := &ParsedRequest{ - System: "You are a helpful assistant.", - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - SessionContext: &SessionContext{ - ClientIP: "2.2.2.2", - UserAgent: "Mozilla/5.0", - APIKeyID: 20, - }, - } + body := anthropicSessionBody("You are a helpful assistant.", []any{msg("user", "hello")}, "") + user1 := mustParseSessionHashRequest(t, body, &SessionContext{ClientIP: "1.1.1.1", UserAgent: "Mozilla/5.0", APIKeyID: 10}) + user2 := mustParseSessionHashRequest(t, body, &SessionContext{ClientIP: "2.2.2.2", UserAgent: "Mozilla/5.0", APIKeyID: 20}) h1 := svc.GenerateSessionHash(user1) h2 := svc.GenerateSessionHash(user2) require.NotEqual(t, h1, h2, "CRITICAL: same system+messages but different users should get different hashes") } -// ============ SessionContext 各字段独立影响测试 ============ - func TestGenerateSessionHash_SessionContext_IPDifference(t *testing.T) { svc := &GatewayService{} - + body := anthropicSessionBody(nil, []any{msg("user", "test")}, "") base := func(ip string) *ParsedRequest { - return &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "test"}, - }, - SessionContext: &SessionContext{ - ClientIP: ip, - UserAgent: "same-ua", - APIKeyID: 1, - }, - } + return mustParseSessionHashRequest(t, body, &SessionContext{ClientIP: ip, UserAgent: "same-ua", APIKeyID: 1}) } h1 := svc.GenerateSessionHash(base("1.1.1.1")) @@ -485,18 +274,9 @@ func TestGenerateSessionHash_SessionContext_IPDifference(t *testing.T) { func TestGenerateSessionHash_SessionContext_UADifference(t *testing.T) { svc := &GatewayService{} - + body := anthropicSessionBody(nil, []any{msg("user", "test")}, "") base := func(ua string) *ParsedRequest { - return &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "test"}, - }, - SessionContext: &SessionContext{ - ClientIP: "1.1.1.1", - UserAgent: ua, - APIKeyID: 1, - }, - } + return mustParseSessionHashRequest(t, body, &SessionContext{ClientIP: "1.1.1.1", UserAgent: ua, APIKeyID: 1}) } h1 := svc.GenerateSessionHash(base("Mozilla/5.0")) @@ -506,18 +286,9 @@ func TestGenerateSessionHash_SessionContext_UADifference(t *testing.T) { func TestGenerateSessionHash_SessionContext_UAVersionNoiseIgnored(t *testing.T) { svc := &GatewayService{} - + body := anthropicSessionBody(nil, []any{msg("user", "test")}, "") base := func(ua string) *ParsedRequest { - return &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "test"}, - }, - SessionContext: &SessionContext{ - ClientIP: "1.1.1.1", - UserAgent: ua, - APIKeyID: 1, - }, - } + return mustParseSessionHashRequest(t, body, &SessionContext{ClientIP: "1.1.1.1", UserAgent: ua, APIKeyID: 1}) } h1 := svc.GenerateSessionHash(base("Mozilla/5.0 codex_cli_rs/0.1.0")) @@ -527,18 +298,9 @@ func TestGenerateSessionHash_SessionContext_UAVersionNoiseIgnored(t *testing.T) func TestGenerateSessionHash_SessionContext_FreeformUAVersionNoiseIgnored(t *testing.T) { svc := &GatewayService{} - + body := anthropicSessionBody(nil, []any{msg("user", "test")}, "") base := func(ua string) *ParsedRequest { - return &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "test"}, - }, - SessionContext: &SessionContext{ - ClientIP: "1.1.1.1", - UserAgent: ua, - APIKeyID: 1, - }, - } + return mustParseSessionHashRequest(t, body, &SessionContext{ClientIP: "1.1.1.1", UserAgent: ua, APIKeyID: 1}) } h1 := svc.GenerateSessionHash(base("Codex CLI 0.1.0")) @@ -548,18 +310,9 @@ func TestGenerateSessionHash_SessionContext_FreeformUAVersionNoiseIgnored(t *tes func TestGenerateSessionHash_SessionContext_APIKeyIDDifference(t *testing.T) { svc := &GatewayService{} - + body := anthropicSessionBody(nil, []any{msg("user", "test")}, "") base := func(keyID int64) *ParsedRequest { - return &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "test"}, - }, - SessionContext: &SessionContext{ - ClientIP: "1.1.1.1", - UserAgent: "same-ua", - APIKeyID: keyID, - }, - } + return mustParseSessionHashRequest(t, body, &SessionContext{ClientIP: "1.1.1.1", UserAgent: "same-ua", APIKeyID: keyID}) } h1 := svc.GenerateSessionHash(base(1)) @@ -567,24 +320,12 @@ func TestGenerateSessionHash_SessionContext_APIKeyIDDifference(t *testing.T) { require.NotEqual(t, h1, h2, "different APIKeyID should produce different hash") } -// ============ 多用户并发相同消息场景 ============ - func TestGenerateSessionHash_MultipleUsersSameFirstMessage(t *testing.T) { svc := &GatewayService{} - - // 模拟 5 个不同用户同时发送 "hello" → 应该产生 5 个不同的 hash hashes := make(map[string]bool) + body := anthropicSessionBody(nil, []any{msg("user", "hello")}, "") for i := 0; i < 5; i++ { - parsed := &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - SessionContext: &SessionContext{ - ClientIP: "192.168.1." + string(rune('1'+i)), - UserAgent: "client-" + string(rune('A'+i)), - APIKeyID: int64(i + 1), - }, - } + parsed := mustParseSessionHashRequest(t, body, &SessionContext{ClientIP: "192.168.1." + string(rune('1'+i)), UserAgent: "client-" + string(rune('A'+i)), APIKeyID: int64(i + 1)}) h := svc.GenerateSessionHash(parsed) require.NotEmpty(t, h) require.False(t, hashes[h], "hash collision detected for user %d", i) @@ -593,134 +334,56 @@ func TestGenerateSessionHash_MultipleUsersSameFirstMessage(t *testing.T) { require.Len(t, hashes, 5, "5 different users should produce 5 unique hashes") } -// ============ 连续会话粘性：多轮对话同一用户 ============ - func TestGenerateSessionHash_SameUserGrowingConversation(t *testing.T) { svc := &GatewayService{} - ctx := &SessionContext{ClientIP: "1.2.3.4", UserAgent: "browser", APIKeyID: 42} - - // 模拟同一用户的连续会话，每轮 hash 不同但同用户重试保持一致 - messages := []map[string]any{ - {"role": "user", "content": "msg1"}, - {"role": "assistant", "content": "reply1"}, - {"role": "user", "content": "msg2"}, - {"role": "assistant", "content": "reply2"}, - {"role": "user", "content": "msg3"}, - {"role": "assistant", "content": "reply3"}, - {"role": "user", "content": "msg4"}, + messages := []any{ + msg("user", "msg1"), msg("assistant", "reply1"), msg("user", "msg2"), msg("assistant", "reply2"), + msg("user", "msg3"), msg("assistant", "reply3"), msg("user", "msg4"), } prevHash := "" for round := 1; round <= len(messages); round += 2 { - // 构建前 round 条消息 - msgs := make([]any, round) - for j := 0; j < round; j++ { - msgs[j] = messages[j] - } - parsed := &ParsedRequest{ - System: "System", - HasSystem: true, - Messages: msgs, - SessionContext: ctx, - } + parsed := mustParseSessionHashRequest(t, anthropicSessionBody("System", messages[:round], ""), ctx) h := svc.GenerateSessionHash(parsed) require.NotEmpty(t, h, "round %d hash should not be empty", round) - if prevHash != "" { require.NotEqual(t, prevHash, h, "round %d hash should differ from previous round", round) } prevHash = h - - // 同一轮重试应该相同 h2 := svc.GenerateSessionHash(parsed) require.Equal(t, h, h2, "retry of round %d should produce same hash", round) } } -// ============ 多轮消息内容结构化测试 ============ - func TestGenerateSessionHash_MultipleUserMessages(t *testing.T) { svc := &GatewayService{} - ctx := &SessionContext{ClientIP: "1.2.3.4", UserAgent: "test", APIKeyID: 1} - - // 5 条用户消息（无 assistant 回复） - parsed := &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "first"}, - map[string]any{"role": "user", "content": "second"}, - map[string]any{"role": "user", "content": "third"}, - map[string]any{"role": "user", "content": "fourth"}, - map[string]any{"role": "user", "content": "fifth"}, - }, - SessionContext: ctx, - } - + parsed := mustParseSessionHashRequest(t, anthropicSessionBody(nil, []any{msg("user", "first"), msg("user", "second"), msg("user", "third"), msg("user", "fourth"), msg("user", "fifth")}, ""), ctx) h := svc.GenerateSessionHash(parsed) require.NotEmpty(t, h) - // 修改中间一条消息应该改变 hash - parsed2 := &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "first"}, - map[string]any{"role": "user", "content": "CHANGED"}, - map[string]any{"role": "user", "content": "third"}, - map[string]any{"role": "user", "content": "fourth"}, - map[string]any{"role": "user", "content": "fifth"}, - }, - SessionContext: ctx, - } - + parsed2 := mustParseSessionHashRequest(t, anthropicSessionBody(nil, []any{msg("user", "first"), msg("user", "CHANGED"), msg("user", "third"), msg("user", "fourth"), msg("user", "fifth")}, ""), ctx) h2 := svc.GenerateSessionHash(parsed2) require.NotEqual(t, h, h2, "changing any message should change the hash") } func TestGenerateSessionHash_MessageOrderMatters(t *testing.T) { svc := &GatewayService{} - ctx := &SessionContext{ClientIP: "1.2.3.4", UserAgent: "test", APIKeyID: 1} - - parsed1 := &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "alpha"}, - map[string]any{"role": "user", "content": "beta"}, - }, - SessionContext: ctx, - } - parsed2 := &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "beta"}, - map[string]any{"role": "user", "content": "alpha"}, - }, - SessionContext: ctx, - } + parsed1 := mustParseSessionHashRequest(t, anthropicSessionBody(nil, []any{msg("user", "alpha"), msg("user", "beta")}, ""), ctx) + parsed2 := mustParseSessionHashRequest(t, anthropicSessionBody(nil, []any{msg("user", "beta"), msg("user", "alpha")}, ""), ctx) h1 := svc.GenerateSessionHash(parsed1) h2 := svc.GenerateSessionHash(parsed2) require.NotEqual(t, h1, h2, "message order should affect the hash") } -// ============ 复杂内容格式测试 ============ - func TestGenerateSessionHash_StructuredContent(t *testing.T) { svc := &GatewayService{} - ctx := &SessionContext{ClientIP: "1.2.3.4", UserAgent: "test", APIKeyID: 1} - - // 结构化 content（数组形式） - parsed := &ParsedRequest{ - Messages: []any{ - map[string]any{ - "role": "user", - "content": []any{ - map[string]any{"type": "text", "text": "Look at this"}, - map[string]any{"type": "text", "text": "And this too"}, - }, - }, - }, - SessionContext: ctx, - } + content := []any{map[string]any{"type": "text", "text": "Look at this"}, map[string]any{"type": "text", "text": "And this too"}} + parsed := mustParseSessionHashRequest(t, anthropicSessionBody(nil, []any{msg("user", content)}, ""), ctx) h := svc.GenerateSessionHash(parsed) require.NotEmpty(t, h, "structured content should produce a hash") @@ -728,100 +391,37 @@ func TestGenerateSessionHash_StructuredContent(t *testing.T) { func TestGenerateSessionHash_ArraySystemPrompt(t *testing.T) { svc := &GatewayService{} - ctx := &SessionContext{ClientIP: "1.2.3.4", UserAgent: "test", APIKeyID: 1} - - // 数组格式的 system prompt - parsed := &ParsedRequest{ - System: []any{ - map[string]any{"type": "text", "text": "You are a helpful assistant."}, - map[string]any{"type": "text", "text": "Be concise."}, - }, - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - SessionContext: ctx, - } + system := []any{map[string]any{"type": "text", "text": "You are a helpful assistant."}, map[string]any{"type": "text", "text": "Be concise."}} + parsed := mustParseSessionHashRequest(t, anthropicSessionBody(system, []any{msg("user", "hello")}, ""), ctx) h := svc.GenerateSessionHash(parsed) require.NotEmpty(t, h, "array system prompt should produce a hash") } -// ============ SessionContext 与 cache_control 优先级 ============ - func TestGenerateSessionHash_CacheControlOverridesSessionContext(t *testing.T) { svc := &GatewayService{} - - // 当有 cache_control: ephemeral 时，使用第 2 级优先级 - // SessionContext 不应影响结果 - parsed1 := &ParsedRequest{ - System: []any{ - map[string]any{ - "type": "text", - "text": "You are a tool-specific assistant.", - "cache_control": map[string]any{"type": "ephemeral"}, - }, - }, - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - SessionContext: &SessionContext{ - ClientIP: "1.1.1.1", - UserAgent: "ua1", - APIKeyID: 100, - }, - } - parsed2 := &ParsedRequest{ - System: []any{ - map[string]any{ - "type": "text", - "text": "You are a tool-specific assistant.", - "cache_control": map[string]any{"type": "ephemeral"}, - }, - }, - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - SessionContext: &SessionContext{ - ClientIP: "2.2.2.2", - UserAgent: "ua2", - APIKeyID: 200, - }, - } + system := []any{map[string]any{"type": "text", "text": "You are a tool-specific assistant.", "cache_control": map[string]any{"type": "ephemeral"}}} + body := anthropicSessionBody(system, []any{msg("user", "hello")}, "") + parsed1 := mustParseSessionHashRequest(t, body, &SessionContext{ClientIP: "1.1.1.1", UserAgent: "ua1", APIKeyID: 100}) + parsed2 := mustParseSessionHashRequest(t, body, &SessionContext{ClientIP: "2.2.2.2", UserAgent: "ua2", APIKeyID: 200}) h1 := svc.GenerateSessionHash(parsed1) h2 := svc.GenerateSessionHash(parsed2) require.Equal(t, h1, h2, "cache_control ephemeral has higher priority, SessionContext should not affect result") } -// ============ 边界情况 ============ - func TestGenerateSessionHash_EmptyMessages(t *testing.T) { svc := &GatewayService{} + parsed := mustParseSessionHashRequest(t, anthropicSessionBody(nil, []any{}, ""), &SessionContext{ClientIP: "1.1.1.1", UserAgent: "test", APIKeyID: 1}) - parsed := &ParsedRequest{ - Messages: []any{}, - SessionContext: &SessionContext{ - ClientIP: "1.1.1.1", - UserAgent: "test", - APIKeyID: 1, - }, - } - - // 空 messages + 只有 SessionContext 时，combined.Len() > 0 因为有 context 写入 h := svc.GenerateSessionHash(parsed) require.NotEmpty(t, h, "empty messages with SessionContext should still produce a hash from context") } func TestGenerateSessionHash_EmptyMessagesNoContext(t *testing.T) { svc := &GatewayService{} - - parsed := &ParsedRequest{ - Messages: []any{}, - } + parsed := mustParseSessionHashRequest(t, anthropicSessionBody(nil, []any{}, ""), nil) h := svc.GenerateSessionHash(parsed) require.Empty(t, h, "empty messages without SessionContext should produce empty hash") @@ -829,98 +429,37 @@ func TestGenerateSessionHash_EmptyMessagesNoContext(t *testing.T) { func TestGenerateSessionHash_SessionContextWithEmptyFields(t *testing.T) { svc := &GatewayService{} - - // SessionContext 字段为空字符串和零值时仍应影响 hash - withEmptyCtx := &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "test"}, - }, - SessionContext: &SessionContext{ - ClientIP: "", - UserAgent: "", - APIKeyID: 0, - }, - } - withoutCtx := &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "test"}, - }, - } + body := anthropicSessionBody(nil, []any{msg("user", "test")}, "") + withEmptyCtx := mustParseSessionHashRequest(t, body, &SessionContext{ClientIP: "", UserAgent: "", APIKeyID: 0}) + withoutCtx := mustParseSessionHashRequest(t, body, nil) h1 := svc.GenerateSessionHash(withEmptyCtx) h2 := svc.GenerateSessionHash(withoutCtx) - // 有 SessionContext（即使字段为空）仍然会写入分隔符 "::" 等 require.NotEqual(t, h1, h2, "empty-field SessionContext should still differ from nil SessionContext") } -// ============ 长对话历史测试 ============ - func TestGenerateSessionHash_LongConversation(t *testing.T) { svc := &GatewayService{} - ctx := &SessionContext{ClientIP: "1.2.3.4", UserAgent: "test", APIKeyID: 1} - - // 构建 20 轮对话 messages := make([]any, 0, 40) for i := 0; i < 20; i++ { - messages = append(messages, map[string]any{ - "role": "user", - "content": "user message " + string(rune('A'+i)), - }) - messages = append(messages, map[string]any{ - "role": "assistant", - "content": "assistant reply " + string(rune('A'+i)), - }) - } - - parsed := &ParsedRequest{ - System: "System prompt", - HasSystem: true, - Messages: messages, - SessionContext: ctx, + messages = append(messages, msg("user", "user message "+string(rune('A'+i)))) + messages = append(messages, msg("assistant", "assistant reply "+string(rune('A'+i)))) } + parsed := mustParseSessionHashRequest(t, anthropicSessionBody("System prompt", messages, ""), ctx) h := svc.GenerateSessionHash(parsed) require.NotEmpty(t, h) - // 再加一轮应该不同 - moreMessages := make([]any, len(messages)+2) - copy(moreMessages, messages) - moreMessages[len(messages)] = map[string]any{"role": "user", "content": "one more"} - moreMessages[len(messages)+1] = map[string]any{"role": "assistant", "content": "ok"} - - parsed2 := &ParsedRequest{ - System: "System prompt", - HasSystem: true, - Messages: moreMessages, - SessionContext: ctx, - } - + moreMessages := append(append([]any{}, messages...), msg("user", "one more"), msg("assistant", "ok")) + parsed2 := mustParseSessionHashRequest(t, anthropicSessionBody("System prompt", moreMessages, ""), ctx) h2 := svc.GenerateSessionHash(parsed2) require.NotEqual(t, h, h2, "adding more messages to long conversation should change hash") } -// ============ Gemini 原生格式 session hash 测试 ============ - func TestGenerateSessionHash_GeminiContentsProducesHash(t *testing.T) { svc := &GatewayService{} - - // Gemini 格式: contents[].parts[].text - parsed := &ParsedRequest{ - Messages: []any{ - map[string]any{ - "role": "user", - "parts": []any{ - map[string]any{"text": "Hello from Gemini"}, - }, - }, - }, - SessionContext: &SessionContext{ - ClientIP: "1.2.3.4", - UserAgent: "gemini-cli", - APIKeyID: 1, - }, - } + parsed := mustParseGeminiSessionHashRequest(t, geminiSessionBody(nil, []any{geminiMsg("user", "Hello from Gemini")}), &SessionContext{ClientIP: "1.2.3.4", UserAgent: "gemini-cli", APIKeyID: 1}) h := svc.GenerateSessionHash(parsed) require.NotEmpty(t, h, "Gemini contents with parts should produce a non-empty hash") @@ -928,31 +467,9 @@ func TestGenerateSessionHash_GeminiContentsProducesHash(t *testing.T) { func TestGenerateSessionHash_GeminiDifferentContentsDifferentHash(t *testing.T) { svc := &GatewayService{} - ctx := &SessionContext{ClientIP: "1.2.3.4", UserAgent: "gemini-cli", APIKeyID: 1} - - parsed1 := &ParsedRequest{ - Messages: []any{ - map[string]any{ - "role": "user", - "parts": []any{ - map[string]any{"text": "Hello"}, - }, - }, - }, - SessionContext: ctx, - } - parsed2 := &ParsedRequest{ - Messages: []any{ - map[string]any{ - "role": "user", - "parts": []any{ - map[string]any{"text": "Goodbye"}, - }, - }, - }, - SessionContext: ctx, - } + parsed1 := mustParseGeminiSessionHashRequest(t, geminiSessionBody(nil, []any{geminiMsg("user", "Hello")}), ctx) + parsed2 := mustParseGeminiSessionHashRequest(t, geminiSessionBody(nil, []any{geminiMsg("user", "Goodbye")}), ctx) h1 := svc.GenerateSessionHash(parsed1) h2 := svc.GenerateSessionHash(parsed2) @@ -961,28 +478,9 @@ func TestGenerateSessionHash_GeminiDifferentContentsDifferentHash(t *testing.T) func TestGenerateSessionHash_GeminiSameContentsSameHash(t *testing.T) { svc := &GatewayService{} - ctx := &SessionContext{ClientIP: "1.2.3.4", UserAgent: "gemini-cli", APIKeyID: 1} - - mk := func() *ParsedRequest { - return &ParsedRequest{ - Messages: []any{ - map[string]any{ - "role": "user", - "parts": []any{ - map[string]any{"text": "Hello"}, - }, - }, - map[string]any{ - "role": "model", - "parts": []any{ - map[string]any{"text": "Hi there!"}, - }, - }, - }, - SessionContext: ctx, - } - } + body := geminiSessionBody(nil, []any{geminiMsg("user", "Hello"), geminiMsg("model", "Hi there!")}) + mk := func() *ParsedRequest { return mustParseGeminiSessionHashRequest(t, body, ctx) } h1 := svc.GenerateSessionHash(mk()) h2 := svc.GenerateSessionHash(mk()) @@ -991,36 +489,9 @@ func TestGenerateSessionHash_GeminiSameContentsSameHash(t *testing.T) { func TestGenerateSessionHash_GeminiMultiTurnHashChanges(t *testing.T) { svc := &GatewayService{} - ctx := &SessionContext{ClientIP: "1.2.3.4", UserAgent: "gemini-cli", APIKeyID: 1} - - round1 := &ParsedRequest{ - Messages: []any{ - map[string]any{ - "role": "user", - "parts": []any{map[string]any{"text": "hello"}}, - }, - }, - SessionContext: ctx, - } - - round2 := &ParsedRequest{ - Messages: []any{ - map[string]any{ - "role": "user", - "parts": []any{map[string]any{"text": "hello"}}, - }, - map[string]any{ - "role": "model", - "parts": []any{map[string]any{"text": "Hi!"}}, - }, - map[string]any{ - "role": "user", - "parts": []any{map[string]any{"text": "How are you?"}}, - }, - }, - SessionContext: ctx, - } + round1 := mustParseGeminiSessionHashRequest(t, geminiSessionBody(nil, []any{geminiMsg("user", "hello")}), ctx) + round2 := mustParseGeminiSessionHashRequest(t, geminiSessionBody(nil, []any{geminiMsg("user", "hello"), geminiMsg("model", "Hi!"), geminiMsg("user", "How are you?")}), ctx) h1 := svc.GenerateSessionHash(round1) h2 := svc.GenerateSessionHash(round2) @@ -1031,34 +502,9 @@ func TestGenerateSessionHash_GeminiMultiTurnHashChanges(t *testing.T) { func TestGenerateSessionHash_GeminiDifferentUsersSameContentDifferentHash(t *testing.T) { svc := &GatewayService{} - - // 核心场景：两个不同用户发送相同 Gemini 格式消息应得到不同 hash - user1 := &ParsedRequest{ - Messages: []any{ - map[string]any{ - "role": "user", - "parts": []any{map[string]any{"text": "hello"}}, - }, - }, - SessionContext: &SessionContext{ - ClientIP: "1.1.1.1", - UserAgent: "gemini-cli", - APIKeyID: 10, - }, - } - user2 := &ParsedRequest{ - Messages: []any{ - map[string]any{ - "role": "user", - "parts": []any{map[string]any{"text": "hello"}}, - }, - }, - SessionContext: &SessionContext{ - ClientIP: "2.2.2.2", - UserAgent: "gemini-cli", - APIKeyID: 20, - }, - } + body := geminiSessionBody(nil, []any{geminiMsg("user", "hello")}) + user1 := mustParseGeminiSessionHashRequest(t, body, &SessionContext{ClientIP: "1.1.1.1", UserAgent: "gemini-cli", APIKeyID: 10}) + user2 := mustParseGeminiSessionHashRequest(t, body, &SessionContext{ClientIP: "2.2.2.2", UserAgent: "gemini-cli", APIKeyID: 20}) h1 := svc.GenerateSessionHash(user1) h2 := svc.GenerateSessionHash(user2) @@ -1067,31 +513,9 @@ func TestGenerateSessionHash_GeminiDifferentUsersSameContentDifferentHash(t *tes func TestGenerateSessionHash_GeminiSystemInstructionAffectsHash(t *testing.T) { svc := &GatewayService{} - ctx := &SessionContext{ClientIP: "1.2.3.4", UserAgent: "gemini-cli", APIKeyID: 1} - - // systemInstruction 经 ParseGatewayRequest 解析后存入 parsed.System - withSys := &ParsedRequest{ - System: []any{ - map[string]any{"text": "You are a coding assistant."}, - }, - Messages: []any{ - map[string]any{ - "role": "user", - "parts": []any{map[string]any{"text": "hello"}}, - }, - }, - SessionContext: ctx, - } - withoutSys := &ParsedRequest{ - Messages: []any{ - map[string]any{ - "role": "user", - "parts": []any{map[string]any{"text": "hello"}}, - }, - }, - SessionContext: ctx, - } + withSys := mustParseGeminiSessionHashRequest(t, geminiSessionBody([]any{map[string]any{"text": "You are a coding assistant."}}, []any{geminiMsg("user", "hello")}), ctx) + withoutSys := mustParseGeminiSessionHashRequest(t, geminiSessionBody(nil, []any{geminiMsg("user", "hello")}), ctx) h1 := svc.GenerateSessionHash(withSys) h2 := svc.GenerateSessionHash(withoutSys) @@ -1100,64 +524,21 @@ func TestGenerateSessionHash_GeminiSystemInstructionAffectsHash(t *testing.T) { func TestGenerateSessionHash_GeminiMultiPartMessage(t *testing.T) { svc := &GatewayService{} - ctx := &SessionContext{ClientIP: "1.2.3.4", UserAgent: "gemini-cli", APIKeyID: 1} - - // 多 parts 的消息 - parsed := &ParsedRequest{ - Messages: []any{ - map[string]any{ - "role": "user", - "parts": []any{ - map[string]any{"text": "Part 1"}, - map[string]any{"text": "Part 2"}, - map[string]any{"text": "Part 3"}, - }, - }, - }, - SessionContext: ctx, - } - + parsed := mustParseGeminiSessionHashRequest(t, geminiSessionBody(nil, []any{geminiMsg("user", "Part 1", "Part 2", "Part 3")}), ctx) h := svc.GenerateSessionHash(parsed) require.NotEmpty(t, h, "multi-part Gemini message should produce a hash") - // 不同内容的多 parts - parsed2 := &ParsedRequest{ - Messages: []any{ - map[string]any{ - "role": "user", - "parts": []any{ - map[string]any{"text": "Part 1"}, - map[string]any{"text": "CHANGED"}, - map[string]any{"text": "Part 3"}, - }, - }, - }, - SessionContext: ctx, - } - + parsed2 := mustParseGeminiSessionHashRequest(t, geminiSessionBody(nil, []any{geminiMsg("user", "Part 1", "CHANGED", "Part 3")}), ctx) h2 := svc.GenerateSessionHash(parsed2) require.NotEqual(t, h, h2, "changing a part should change the hash") } func TestGenerateSessionHash_GeminiNonTextPartsIgnored(t *testing.T) { svc := &GatewayService{} - ctx := &SessionContext{ClientIP: "1.2.3.4", UserAgent: "gemini-cli", APIKeyID: 1} - - // 含非 text 类型 parts（如 inline_data），应被跳过但不报错 - parsed := &ParsedRequest{ - Messages: []any{ - map[string]any{ - "role": "user", - "parts": []any{ - map[string]any{"text": "Describe this image"}, - map[string]any{"inline_data": map[string]any{"mime_type": "image/png", "data": "base64..."}}, - }, - }, - }, - SessionContext: ctx, - } + content := []any{map[string]any{"role": "user", "parts": []any{map[string]any{"text": "Describe this image"}, map[string]any{"inline_data": map[string]any{"mime_type": "image/png", "data": "base64..."}}}}} + parsed := mustParseGeminiSessionHashRequest(t, geminiSessionBody(nil, content), ctx) h := svc.GenerateSessionHash(parsed) require.NotEmpty(t, h, "Gemini message with mixed parts should still produce a hash from text parts") @@ -1165,107 +546,41 @@ func TestGenerateSessionHash_GeminiNonTextPartsIgnored(t *testing.T) { func TestGenerateSessionHash_GeminiMultiTurnHashNotSticky(t *testing.T) { svc := &GatewayService{} - ctx := &SessionContext{ClientIP: "10.0.0.1", UserAgent: "gemini-cli", APIKeyID: 42} + rounds := []string{ + geminiSessionBody([]any{map[string]any{"text": "You are a coding assistant."}}, []any{geminiMsg("user", "Write a Go function")}), + geminiSessionBody([]any{map[string]any{"text": "You are a coding assistant."}}, []any{geminiMsg("user", "Write a Go function"), geminiMsg("model", "func hello() {}"), geminiMsg("user", "Add error handling")}), + geminiSessionBody([]any{map[string]any{"text": "You are a coding assistant."}}, []any{geminiMsg("user", "Write a Go function"), geminiMsg("model", "func hello() {}"), geminiMsg("user", "Add error handling"), geminiMsg("model", "func hello() error { return nil }"), geminiMsg("user", "Now add tests")}), + } - // 模拟同一 Gemini 会话的三轮请求，每轮 contents 累积增长。 - // 验证预期行为：每轮 hash 都不同，即 GenerateSessionHash 不具备跨轮粘性。 - // 这是 by-design 的——Gemini 的跨轮粘性由 Digest Fallback（BuildGeminiDigestChain）负责。 - round1Body := []byte(`{ - "systemInstruction": {"parts": [{"text": "You are a coding assistant."}]}, - "contents": [ - {"role": "user", "parts": [{"text": "Write a Go function"}]} - ] - }`) - round2Body := []byte(`{ - "systemInstruction": {"parts": [{"text": "You are a coding assistant."}]}, - "contents": [ - {"role": "user", "parts": [{"text": "Write a Go function"}]}, - {"role": "model", "parts": [{"text": "func hello() {}"}]}, - {"role": "user", "parts": [{"text": "Add error handling"}]} - ] - }`) - round3Body := []byte(`{ - "systemInstruction": {"parts": [{"text": "You are a coding assistant."}]}, - "contents": [ - {"role": "user", "parts": [{"text": "Write a Go function"}]}, - {"role": "model", "parts": [{"text": "func hello() {}"}]}, - {"role": "user", "parts": [{"text": "Add error handling"}]}, - {"role": "model", "parts": [{"text": "func hello() error { return nil }"}]}, - {"role": "user", "parts": [{"text": "Now add tests"}]} - ] - }`) - - hashes := make([]string, 3) - for i, body := range [][]byte{round1Body, round2Body, round3Body} { - parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), "gemini") - require.NoError(t, err) - parsed.SessionContext = ctx + hashes := make([]string, len(rounds)) + for i, body := range rounds { + parsed := mustParseGeminiSessionHashRequest(t, body, ctx) hashes[i] = svc.GenerateSessionHash(parsed) require.NotEmpty(t, hashes[i], "round %d hash should not be empty", i+1) } - - // 每轮 hash 都不同——这是预期行为 require.NotEqual(t, hashes[0], hashes[1], "round 1 vs 2 hash should differ (contents grow)") require.NotEqual(t, hashes[1], hashes[2], "round 2 vs 3 hash should differ (contents grow)") require.NotEqual(t, hashes[0], hashes[2], "round 1 vs 3 hash should differ") - // 同一轮重试应产生相同 hash - parsed1Again, err := ParseGatewayRequest(NewRequestBodyRef(round2Body), "gemini") - require.NoError(t, err) - parsed1Again.SessionContext = ctx - h2Again := svc.GenerateSessionHash(parsed1Again) + parsedAgain := mustParseGeminiSessionHashRequest(t, rounds[1], ctx) + h2Again := svc.GenerateSessionHash(parsedAgain) require.Equal(t, hashes[1], h2Again, "retry of same round should produce same hash") } func TestGenerateSessionHash_GeminiEndToEnd(t *testing.T) { svc := &GatewayService{} - - // 端到端测试：模拟 ParseGatewayRequest + GenerateSessionHash 完整流程 - body := []byte(`{ - "model": "gemini-2.5-pro", - "systemInstruction": { - "parts": [{"text": "You are a coding assistant."}] - }, - "contents": [ - {"role": "user", "parts": [{"text": "Write a Go function"}]}, - {"role": "model", "parts": [{"text": "Here is a function..."}]}, - {"role": "user", "parts": [{"text": "Now add error handling"}]} - ] - }`) - - parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), "gemini") - require.NoError(t, err) - parsed.SessionContext = &SessionContext{ - ClientIP: "10.0.0.1", - UserAgent: "gemini-cli/1.0", - APIKeyID: 42, - } + body := geminiSessionBody([]any{map[string]any{"text": "You are a coding assistant."}}, []any{geminiMsg("user", "Write a Go function"), geminiMsg("model", "Here is a function..."), geminiMsg("user", "Now add error handling")}) + parsed := mustParseGeminiSessionHashRequest(t, body, &SessionContext{ClientIP: "10.0.0.1", UserAgent: "gemini-cli/1.0", APIKeyID: 42}) h := svc.GenerateSessionHash(parsed) require.NotEmpty(t, h, "end-to-end Gemini flow should produce a hash") - // 同一请求再次解析应产生相同 hash - parsed2, err := ParseGatewayRequest(NewRequestBodyRef(body), "gemini") - require.NoError(t, err) - parsed2.SessionContext = &SessionContext{ - ClientIP: "10.0.0.1", - UserAgent: "gemini-cli/1.0", - APIKeyID: 42, - } - + parsed2 := mustParseGeminiSessionHashRequest(t, body, &SessionContext{ClientIP: "10.0.0.1", UserAgent: "gemini-cli/1.0", APIKeyID: 42}) h2 := svc.GenerateSessionHash(parsed2) require.Equal(t, h, h2, "same request should produce same hash") - // 不同用户发送相同请求应产生不同 hash - parsed3, err := ParseGatewayRequest(NewRequestBodyRef(body), "gemini") - require.NoError(t, err) - parsed3.SessionContext = &SessionContext{ - ClientIP: "10.0.0.2", - UserAgent: "gemini-cli/1.0", - APIKeyID: 99, - } - + parsed3 := mustParseGeminiSessionHashRequest(t, body, &SessionContext{ClientIP: "10.0.0.2", UserAgent: "gemini-cli/1.0", APIKeyID: 99}) h3 := svc.GenerateSessionHash(parsed3) require.NotEqual(t, h, h3, "different user with same Gemini request should get different hash") } diff --git a/backend/internal/service/openai_gateway_service.go b/backend/internal/service/openai_gateway_service.go index cd5a40157c9e..f8c3403120b6 100644 --- a/backend/internal/service/openai_gateway_service.go +++ b/backend/internal/service/openai_gateway_service.go @@ -36,6 +36,13 @@ import ( "go.uber.org/zap" ) +// openAIParsedRequestBodyCache 绑定 body 指纹，避免 handler 预解析的旧 body 污染后续 forwardBody。 +type openAIParsedRequestBodyCache struct { + bodyHash uint64 + bodyLen int + reqBody map[string]any +} + const ( // ChatGPT internal API for OAuth accounts chatgptCodexURL = "https://chatgpt.com/backend-api/codex/responses" @@ -2284,6 +2291,20 @@ func (s *OpenAIGatewayService) shouldFailoverOpenAIUpstreamResponse(statusCode i return isOpenAITransientProcessingError(statusCode, upstreamMsg, upstreamBody) } +func marshalOpenAIUpstreamJSON(v any) ([]byte, error) { + var buf bytes.Buffer + enc := json.NewEncoder(&buf) + enc.SetEscapeHTML(false) + if err := enc.Encode(v); err != nil { + return nil, err + } + out := buf.Bytes() + if len(out) > 0 && out[len(out)-1] == '\n' { + out = out[:len(out)-1] + } + return out, nil +} + func (s *OpenAIGatewayService) handleFailoverSideEffects(ctx context.Context, resp *http.Response, account *Account, requestedModel ...string) { body, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) if len(requestedModel) > 0 { @@ -2561,7 +2582,11 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco // 确保高版本模型向低版本模型映射不报错 if !SupportsVerbosity(upstreamModel) { if text, ok := reqBody["text"].(map[string]any); ok { - delete(text, "verbosity") + if _, exists := text["verbosity"]; exists { + delete(text, "verbosity") + bodyModified = true + markPatchDelete("text.verbosity") + } } } } @@ -2762,7 +2787,7 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco } if !serializedByPatch { var marshalErr error - body, marshalErr = json.Marshal(reqBody) + body, marshalErr = marshalOpenAIUpstreamJSON(reqBody) if marshalErr != nil { return nil, fmt.Errorf("serialize request body: %w", marshalErr) } @@ -3043,7 +3068,7 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco upstreamCode := extractUpstreamErrorCode(respBody) if !httpInvalidEncryptedContentRetryTried && resp.StatusCode == http.StatusBadRequest && upstreamCode == "invalid_encrypted_content" { if trimOpenAIEncryptedReasoningItems(reqBody) { - body, err = json.Marshal(reqBody) + body, err = marshalOpenAIUpstreamJSON(reqBody) if err != nil { return nil, fmt.Errorf("serialize invalid_encrypted_content retry body: %w", err) } @@ -3074,6 +3099,8 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco }) s.handleFailoverSideEffects(ctx, resp, account, upstreamModel) + // reqBody 会被本次账号尝试原地修改，failover 前必须释放，避免下一账号复用脏 map。 + releaseOpenAIParsedRequestBody(c) return nil, &UpstreamFailoverError{ StatusCode: resp.StatusCode, ResponseBody: respBody, @@ -6711,7 +6738,7 @@ func sanitizeEmptyBase64InputImagesInOpenAIBody(body []byte) ([]byte, bool, erro if !sanitizeEmptyBase64InputImagesInOpenAIRequestBodyMap(reqBody) { return body, false, nil } - normalized, err := json.Marshal(reqBody) + normalized, err := marshalOpenAIUpstreamJSON(reqBody) if err != nil { return body, false, fmt.Errorf("serialize sanitized request body: %w", err) } @@ -6817,10 +6844,13 @@ func isEmptyBase64DataURI(raw string) bool { } func getOpenAIRequestBodyMap(c *gin.Context, body []byte) (map[string]any, error) { + // 同一个 gin.Context 内 failover/渠道映射可能传入新 body，缓存必须先校验 body 指纹。 + bodyHash := xxhash.Sum64(body) + bodyLen := len(body) if c != nil { if cached, ok := c.Get(OpenAIParsedRequestBodyKey); ok { - if reqBody, ok := cached.(map[string]any); ok && reqBody != nil { - return reqBody, nil + if cache, ok := cached.(openAIParsedRequestBodyCache); ok && cache.reqBody != nil && cache.bodyLen == bodyLen && cache.bodyHash == bodyHash { + return cache.reqBody, nil } } } @@ -6830,11 +6860,36 @@ func getOpenAIRequestBodyMap(c *gin.Context, body []byte) (map[string]any, error return nil, fmt.Errorf("parse request: %w", err) } if c != nil { - c.Set(OpenAIParsedRequestBodyKey, reqBody) + c.Set(OpenAIParsedRequestBodyKey, openAIParsedRequestBodyCache{bodyHash: bodyHash, bodyLen: bodyLen, reqBody: reqBody}) } return reqBody, nil } +// CacheOpenAIParsedRequestBody 仅缓存与当前 body 绑定的解析结果。 +func CacheOpenAIParsedRequestBody(c *gin.Context, body []byte, reqBody map[string]any) { + if c == nil || reqBody == nil { + return + } + c.Set(OpenAIParsedRequestBodyKey, openAIParsedRequestBodyCache{ + bodyHash: xxhash.Sum64(body), + bodyLen: len(body), + reqBody: reqBody, + }) +} + +// CachedOpenAIParsedRequestBody 只给同请求内不关心 body 参数的轻量识别逻辑使用。 +func CachedOpenAIParsedRequestBody(c *gin.Context) map[string]any { + if c == nil { + return nil + } + if cached, ok := c.Get(OpenAIParsedRequestBodyKey); ok { + if cache, ok := cached.(openAIParsedRequestBodyCache); ok { + return cache.reqBody + } + } + return nil +} + func releaseOpenAIParsedRequestBody(c *gin.Context) { if c == nil { return diff --git a/backend/internal/service/openai_gateway_service_hotpath_test.go b/backend/internal/service/openai_gateway_service_hotpath_test.go index 234dee00cf86..2ff72e2fa418 100644 --- a/backend/internal/service/openai_gateway_service_hotpath_test.go +++ b/backend/internal/service/openai_gateway_service_hotpath_test.go @@ -112,7 +112,7 @@ func TestGetOpenAIRequestBodyMap_UsesContextCache(t *testing.T) { c, _ := gin.CreateTestContext(rec) cached := map[string]any{"model": "cached-model", "stream": true} - c.Set(OpenAIParsedRequestBodyKey, cached) + CacheOpenAIParsedRequestBody(c, []byte(`{invalid-json`), cached) got, err := getOpenAIRequestBodyMap(c, []byte(`{invalid-json`)) require.NoError(t, err) @@ -134,11 +134,19 @@ func TestGetOpenAIRequestBodyMap_WriteBackContextCache(t *testing.T) { require.NoError(t, err) require.Equal(t, "gpt-5", got["model"]) - cached, ok := c.Get(OpenAIParsedRequestBodyKey) - require.True(t, ok) - cachedMap, ok := cached.(map[string]any) - require.True(t, ok) - require.Equal(t, got, cachedMap) + require.Equal(t, got, CachedOpenAIParsedRequestBody(c)) +} + +func TestGetOpenAIRequestBodyMap_IgnoresCacheForDifferentBody(t *testing.T) { + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + + CacheOpenAIParsedRequestBody(c, []byte(`{"model":"cached-model"}`), map[string]any{"model": "cached-model"}) + + got, err := getOpenAIRequestBodyMap(c, []byte(`{"model":"forward-model"}`)) + require.NoError(t, err) + require.Equal(t, "forward-model", got["model"]) } func TestSanitizeEmptyBase64InputImagesInOpenAIRequestBodyMap(t *testing.T) { diff --git a/backend/internal/service/openai_oauth_passthrough_test.go b/backend/internal/service/openai_oauth_passthrough_test.go index 398cbb850b48..2710c696e69a 100644 --- a/backend/internal/service/openai_oauth_passthrough_test.go +++ b/backend/internal/service/openai_oauth_passthrough_test.go @@ -14,6 +14,7 @@ import ( "github.com/Wei-Shaw/sub2api/internal/config" "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "github.com/Wei-Shaw/sub2api/internal/pkg/openai_compat" "github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" @@ -101,6 +102,57 @@ func TestOpenAIGatewayService_ResponsesUnknownModelDoesNotFallbackToGPT54(t *tes require.True(t, rec.Code >= http.StatusBadRequest) } +func TestOpenAIGatewayService_NativeResponsesBodyModificationPreservesHTMLChars(t *testing.T) { + gin.SetMode(gin.TestMode) + + payloadText := strings.Repeat(`&value`, 128) + originalBody := []byte(fmt.Sprintf(`{"model":"gpt-5.5","stream":false,"max_output_tokens":100,"previous_response_id":"resp_prev","input":[{"type":"message","role":"user","content":[{"type":"input_text","text":%q}]}]}`, payloadText)) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(originalBody)) + c.Request.Header.Set("Content-Type", "application/json") + + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusBadRequest, + Header: http.Header{"Content-Type": []string{"application/json"}, "x-request-id": []string{"rid_native_reencode"}}, + Body: io.NopCloser(strings.NewReader(`{"error":{"type":"invalid_request_error","message":"stop after capture"}}`)), + }} + svc := &OpenAIGatewayService{ + cfg: &config.Config{Security: config.SecurityConfig{URLAllowlist: config.URLAllowlistConfig{ + Enabled: false, + AllowInsecureHTTP: true, + }}}, + httpUpstream: upstream, + } + account := &Account{ + ID: 456, + Name: "openai-apikey", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "sk-test", + "base_url": "http://upstream.example", + }, + Extra: map[string]any{ + openai_compat.ExtraKeyResponsesMode: string(openai_compat.ResponsesSupportModeAuto), + openai_compat.ExtraKeyResponsesSupported: true, + }, + Status: StatusActive, + Schedulable: true, + } + + result, err := svc.Forward(context.Background(), c, account, originalBody) + require.Error(t, err) + require.Nil(t, result) + require.NotNil(t, upstream.lastReq) + require.Equal(t, "http://upstream.example/v1/responses", upstream.lastReq.URL.String()) + require.Contains(t, string(upstream.lastBody), payloadText) + require.NotContains(t, string(upstream.lastBody), `\\u003c`) + require.NotContains(t, string(upstream.lastBody), `\\u003e`) + require.NotContains(t, string(upstream.lastBody), `\\u0026`) +} + func TestOpenAIGatewayService_OAuthMessagesBridgeDoesNotInjectDefaultInstructions(t *testing.T) { gin.SetMode(gin.TestMode) diff --git a/backend/internal/service/user_msg_queue_service.go b/backend/internal/service/user_msg_queue_service.go index a0ce95a8638e..f3f105ac9386 100644 --- a/backend/internal/service/user_msg_queue_service.go +++ b/backend/internal/service/user_msg_queue_service.go @@ -12,6 +12,7 @@ import ( "github.com/Wei-Shaw/sub2api/internal/config" "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "github.com/tidwall/gjson" ) // UserMsgQueueCache 用户消息串行队列 Redis 缓存接口 @@ -62,43 +63,48 @@ func NewUserMessageQueueService(cache UserMsgQueueCache, rpmCache RPMCache, cfg // 2. 最后一条消息 role == "user" // 3. 最后一条消息 content（如果是数组）中不含 type:"tool_result" / "tool_use_result" func IsRealUserMessage(parsed *ParsedRequest) bool { - if parsed == nil || len(parsed.Messages) == 0 { + if parsed == nil { return false } - - lastMsg := parsed.Messages[len(parsed.Messages)-1] - msgMap, ok := lastMsg.(map[string]any) - if !ok { + messagesRaw := parsed.MessagesRaw() + if len(messagesRaw) == 0 { return false } - role, _ := msgMap["role"].(string) - if role != "user" { + messages := gjson.ParseBytes(messagesRaw) + if !messages.IsArray() { return false } - - // 检查 content 是否包含 tool_result 类型 - content, ok := msgMap["content"] - if !ok { - return true // 没有 content 字段，视为普通用户消息 + lastMsg := gjson.Result{} + messages.ForEach(func(_, msg gjson.Result) bool { + lastMsg = msg + return true + }) + if !lastMsg.Exists() || !lastMsg.IsObject() { + return false + } + if lastMsg.Get("role").String() != "user" { + return false } - contentArr, ok := content.([]any) - if !ok { - return true // content 不是数组（可能是 string），视为普通用户消息 + content := lastMsg.Get("content") + if !content.Exists() { + return true + } + if !content.IsArray() { + return true } - for _, item := range contentArr { - itemMap, ok := item.(map[string]any) - if !ok { - continue - } - itemType, _ := itemMap["type"].(string) + isReal := true + content.ForEach(func(_, item gjson.Result) bool { + itemType := item.Get("type").String() if itemType == "tool_result" || itemType == "tool_use_result" { + isReal = false return false } - } - return true + return true + }) + return isReal } // TryAcquire 尝试立即获取串行锁 From 619e5ae6193703465eaa65a0ada3c045c5bacb2e Mon Sep 17 00:00:00 2001 From: name <136912576+is7Qin@users.noreply.github.com> Date: Sat, 30 May 2026 02:40:52 +0800 Subject: [PATCH 19/36] refactor(gateway): isolate anthropic body rewrites Keep Anthropic request body rewrites attempt-local and synchronize the accepted wire body only after upstream success so failover and retry paths do not reuse stale parsed state. --- backend/internal/handler/gateway_handler.go | 41 ++-- ...teway_anthropic_apikey_passthrough_test.go | 4 +- ...y_anthropic_vertex_service_account_test.go | 6 +- .../gateway_context_management_test.go | 16 +- .../gateway_forward_as_chat_completions.go | 2 +- .../service/gateway_forward_as_responses.go | 2 +- backend/internal/service/gateway_request.go | 153 +++++++------- backend/internal/service/gateway_service.go | 194 ++++++++++++++---- 8 files changed, 277 insertions(+), 141 deletions(-) diff --git a/backend/internal/handler/gateway_handler.go b/backend/internal/handler/gateway_handler.go index 10d67ba32d4d..c9b3297e1475 100644 --- a/backend/internal/handler/gateway_handler.go +++ b/backend/internal/handler/gateway_handler.go @@ -563,6 +563,12 @@ func (h *GatewayHandler) Messages(c *gin.Context) { retryWithFallback := false for { + attemptParsedReq, err := parsedReq.CloneForBody(body) + if err != nil { + h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body") + return + } + // 选择支持该模型的账号 reqLog.Info("sticky.selecting_account", zap.String("session_key", sessionKey), @@ -694,7 +700,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) { // ===== 用户消息串行队列 START ===== var queueRelease func() - umqMode := h.getUserMsgQueueMode(account, parsedReq) + umqMode := h.getUserMsgQueueMode(account, attemptParsedReq) switch umqMode { case config.UMQModeSerialize: @@ -741,20 +747,26 @@ func (h *GatewayHandler) Messages(c *gin.Context) { // 用 wrapReleaseOnDone 确保 context 取消时自动释放（仅 serialize 模式有 queueRelease） queueRelease = wrapReleaseOnDone(c.Request.Context(), queueRelease) // 注入回调到 ParsedRequest：使用外层 wrapper 以便提前清理 AfterFunc - parsedReq.OnUpstreamAccepted = queueRelease + attemptParsedReq.OnUpstreamAccepted = queueRelease // ===== 用户消息串行队列 END ===== - // 应用渠道模型映射到请求 + // 渠道模型映射只作用于本次账号尝试，避免 failover 后污染原始 ParsedRequest。 if channelMapping.Mapped { - parsedReq.Model = channelMapping.MappedModel - parsedReq.ReplaceBody(h.gatewayService.ReplaceModelInBody(parsedReq.Body.Bytes(), channelMapping.MappedModel)) + attemptParsedReq.Model = channelMapping.MappedModel + if err := attemptParsedReq.ReplaceBody(h.gatewayService.ReplaceModelInBody(attemptParsedReq.Body.Bytes(), channelMapping.MappedModel)); err != nil { + h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body") + return + } } // Bedrock CC 兼容：渠道模型映射后，清理 Anthropic API 专有字段、注入 Bedrock 必需字段 - parsedReq.ReplaceBody(h.gatewayService.ApplyBedrockCCCompat(c.Request.Context(), parsedReq.Body.Bytes(), parsedReq.Model, account, apiKey.GroupID)) - body = parsedReq.Body.Bytes() + if err := attemptParsedReq.ReplaceBody(h.gatewayService.ApplyBedrockCCCompat(c.Request.Context(), attemptParsedReq.Body.Bytes(), attemptParsedReq.Model, account, apiKey.GroupID)); err != nil { + h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body") + return + } + attemptBody := attemptParsedReq.Body.Bytes() // 转发请求 - 根据账号平台分流 - c.Set("parsed_request", parsedReq) + c.Set("parsed_request", attemptParsedReq) var result *service.ForwardResult requestCtx := c.Request.Context() if fs.SwitchCount > 0 { @@ -763,9 +775,9 @@ func (h *GatewayHandler) Messages(c *gin.Context) { // 记录 Forward 前已写入字节数，Forward 后若增加则说明 SSE 内容已发，禁止 failover writerSizeBeforeForward := c.Writer.Size() if account.Platform == service.PlatformAntigravity && account.Type != service.AccountTypeAPIKey { - result, err = h.antigravityGatewayService.Forward(requestCtx, c, account, body, hasBoundSession) + result, err = h.antigravityGatewayService.Forward(requestCtx, c, account, attemptBody, hasBoundSession) } else { - result, err = h.gatewayService.Forward(requestCtx, c, account, parsedReq) + result, err = h.gatewayService.Forward(requestCtx, c, account, attemptParsedReq) } // 兜底释放串行锁（正常情况已通过回调提前释放） @@ -773,7 +785,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) { queueRelease() } // 清理回调引用，防止 failover 重试时旧回调被错误调用 - parsedReq.OnUpstreamAccepted = nil + attemptParsedReq.OnUpstreamAccepted = nil if accountReleaseFunc != nil { accountReleaseFunc() @@ -896,12 +908,13 @@ func (h *GatewayHandler) Messages(c *gin.Context) { // 捕获请求信息（用于异步记录，避免在 goroutine 中访问 gin.Context） userAgent := c.GetHeader("User-Agent") clientIP := ip.GetClientIP(c) - requestPayloadHash := service.HashUsageRequestPayload(body) + // Forward 内部可能继续改写 body，usage 去重指纹必须使用最终上游接受的当前 body。 + requestPayloadHash := service.HashUsageRequestPayload(attemptParsedReq.Body.Bytes()) inboundEndpoint := GetInboundEndpoint(c) upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform) if result.ReasoningEffort == nil { - result.ReasoningEffort = service.NormalizeClaudeOutputEffort(parsedReq.OutputEffort) + result.ReasoningEffort = service.NormalizeClaudeOutputEffort(attemptParsedReq.OutputEffort) } // 使用量记录通过有界 worker 池提交，避免请求热路径创建无界 goroutine。 @@ -909,7 +922,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) { h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) { if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{ Result: result, - ParsedRequest: parsedReq, + ParsedRequest: attemptParsedReq, QuotaPlatform: quotaPlatform, APIKey: currentAPIKey, User: currentAPIKey.User, diff --git a/backend/internal/service/gateway_anthropic_apikey_passthrough_test.go b/backend/internal/service/gateway_anthropic_apikey_passthrough_test.go index a67a3dc22f77..e2da89b5db87 100644 --- a/backend/internal/service/gateway_anthropic_apikey_passthrough_test.go +++ b/backend/internal/service/gateway_anthropic_apikey_passthrough_test.go @@ -713,7 +713,7 @@ func TestGatewayService_AnthropicAPIKeyPassthrough_BuildRequestRejectsInvalidBas }, } - _, err := svc.buildUpstreamRequestAnthropicAPIKeyPassthrough(context.Background(), c, account, []byte(`{}`), "k") + _, _, err := svc.buildUpstreamRequestAnthropicAPIKeyPassthrough(context.Background(), c, account, []byte(`{}`), "k") require.Error(t, err) } @@ -738,7 +738,7 @@ func TestGatewayService_AnthropicOAuth_NotAffectedByAPIKeyPassthroughToggle(t *t require.False(t, account.IsAnthropicAPIKeyPassthroughEnabled()) - req, err := svc.buildUpstreamRequest(context.Background(), c, account, []byte(`{"model":"claude-3-7-sonnet-20250219"}`), "oauth-token", "oauth", "claude-3-7-sonnet-20250219", true, false) + req, _, err := svc.buildUpstreamRequest(context.Background(), c, account, []byte(`{"model":"claude-3-7-sonnet-20250219"}`), "oauth-token", "oauth", "claude-3-7-sonnet-20250219", true, false) require.NoError(t, err) require.Equal(t, "Bearer oauth-token", getHeaderRaw(req.Header, "authorization")) require.Contains(t, getHeaderRaw(req.Header, "anthropic-beta"), claude.BetaOAuth, "OAuth 链路仍应按原逻辑补齐 oauth beta") diff --git a/backend/internal/service/gateway_anthropic_vertex_service_account_test.go b/backend/internal/service/gateway_anthropic_vertex_service_account_test.go index 2f42b0ab4a9c..be8c58675f12 100644 --- a/backend/internal/service/gateway_anthropic_vertex_service_account_test.go +++ b/backend/internal/service/gateway_anthropic_vertex_service_account_test.go @@ -35,7 +35,7 @@ func TestGatewayService_BuildAnthropicVertexServiceAccountRequest(t *testing.T) body := []byte(`{"model":"claude-sonnet-4-5","stream":false,"max_tokens":32,"messages":[{"role":"user","content":"hello"}]}`) svc := &GatewayService{} - req, err := svc.buildUpstreamRequest( + req, _, err := svc.buildUpstreamRequest( context.Background(), c, account, @@ -87,7 +87,7 @@ func TestGatewayService_BuildAnthropicVertexServiceAccount_StripsContextManageme body := []byte(`{"model":"claude-haiku-4-5","context_management":{"edits":[{"type":"clear_thinking_20251015","keep":"all"}]},"messages":[{"role":"user","content":"hi"}]}`) svc := &GatewayService{} - req, err := svc.buildUpstreamRequest( + req, _, err := svc.buildUpstreamRequest( context.Background(), c, account, body, "vertex-token", "service_account", "claude-haiku-4-5@20251001", false, false, ) @@ -117,7 +117,7 @@ func TestGatewayService_BuildAnthropicVertexServiceAccount_PreservesContextManag body := []byte(`{"model":"claude-sonnet-4-6","context_management":{"edits":[{"type":"clear_thinking_20251015"}]},"messages":[]}`) svc := &GatewayService{} - req, err := svc.buildUpstreamRequest( + req, _, err := svc.buildUpstreamRequest( context.Background(), c, account, body, "vertex-token", "service_account", "claude-sonnet-4-6@20260218", false, false, ) diff --git a/backend/internal/service/gateway_context_management_test.go b/backend/internal/service/gateway_context_management_test.go index c2263bdc02b4..51b12809e794 100644 --- a/backend/internal/service/gateway_context_management_test.go +++ b/backend/internal/service/gateway_context_management_test.go @@ -364,7 +364,7 @@ func TestBuildUpstreamRequestAnthropicAPIKeyPassthrough_StripsContextManagementW body := []byte(`{"model":"claude-haiku-4-5","context_management":{"edits":[{"type":"clear_thinking_20251015"}]},"messages":[]}`) svc := &GatewayService{cfg: &config.Config{}} - req, err := svc.buildUpstreamRequestAnthropicAPIKeyPassthrough( + req, _, err := svc.buildUpstreamRequestAnthropicAPIKeyPassthrough( context.Background(), c, newAnthropicAPIKeyPassthroughAccountForBetaTest(), body, "token", ) require.NoError(t, err) @@ -381,7 +381,7 @@ func TestBuildUpstreamRequestAnthropicAPIKeyPassthrough_PreservesContextManageme body := []byte(`{"model":"claude-haiku-4-5","context_management":{"edits":[{"type":"clear_thinking_20251015"}]},"messages":[]}`) svc := &GatewayService{cfg: &config.Config{}} - req, err := svc.buildUpstreamRequestAnthropicAPIKeyPassthrough( + req, _, err := svc.buildUpstreamRequestAnthropicAPIKeyPassthrough( context.Background(), c, newAnthropicAPIKeyPassthroughAccountForBetaTest(), body, "token", ) require.NoError(t, err) @@ -427,7 +427,7 @@ func TestBuildUpstreamRequest_OAuthMimicHaiku_StripsContextManagementEndToEnd(t // body 必须 strip。 body := []byte(`{"model":"claude-haiku-4-5","context_management":{"edits":[{"type":"clear_thinking_20251015"}]},"messages":[]}`) svc := &GatewayService{cfg: &config.Config{}} - req, err := svc.buildUpstreamRequest( + req, _, err := svc.buildUpstreamRequest( context.Background(), c, account, body, "oauth-tok", "oauth", "claude-haiku-4-5", false, true, // mimicClaudeCode=true ) @@ -457,7 +457,7 @@ func TestBuildUpstreamRequest_OAuthMimicNonHaiku_PreservesContextManagementEndTo // body 保留。 body := []byte(`{"model":"claude-sonnet-4-6","context_management":{"edits":[{"type":"clear_thinking_20251015"}]},"messages":[]}`) svc := &GatewayService{cfg: &config.Config{}} - req, err := svc.buildUpstreamRequest( + req, _, err := svc.buildUpstreamRequest( context.Background(), c, account, body, "oauth-tok", "oauth", "claude-sonnet-4-6", false, true, ) @@ -488,7 +488,7 @@ func TestBuildUpstreamRequest_OAuthTransparentHaikuWithRealCCBeta_PreservesField } body := []byte(`{"model":"claude-haiku-4-5","context_management":{"edits":[{"type":"clear_thinking_20251015","keep":"all"}]},"messages":[]}`) svc := &GatewayService{cfg: &config.Config{}} - req, err := svc.buildUpstreamRequest( + req, _, err := svc.buildUpstreamRequest( context.Background(), c, account, body, "oauth-tok", "oauth", "claude-haiku-4-5", false, false, // mimicClaudeCode=false（真 CC） ) @@ -580,7 +580,7 @@ func TestBuildCountTokensRequest_OAuthMimicHaiku_PreservesContextManagementEndTo } body := []byte(`{"model":"claude-haiku-4-5","context_management":{"edits":[{"type":"clear_thinking_20251015"}]},"messages":[]}`) svc := &GatewayService{cfg: &config.Config{}} - req, err := svc.buildCountTokensRequest( + req, _, err := svc.buildCountTokensRequest( context.Background(), c, account, body, "oauth-tok", "oauth", "claude-haiku-4-5", true, // mimicClaudeCode=true ) @@ -611,7 +611,7 @@ func TestBuildCountTokensRequest_APIKeyHaiku_StripsContextManagementEndToEnd(t * } body := []byte(`{"model":"claude-haiku-4-5","context_management":{"edits":[]},"messages":[]}`) svc := &GatewayService{cfg: &config.Config{}} - req, err := svc.buildCountTokensRequest( + req, _, err := svc.buildCountTokensRequest( context.Background(), c, account, body, "sk-ant-xxx", "apikey", "claude-haiku-4-5", false, ) @@ -655,7 +655,7 @@ func TestBuildUpstreamRequest_APIKeyHaikuWithContextManagement_StripsField(t *te } body := []byte(`{"model":"claude-haiku-4-5","context_management":{"edits":[]},"messages":[]}`) svc := &GatewayService{cfg: &config.Config{}} - req, err := svc.buildUpstreamRequest( + req, _, err := svc.buildUpstreamRequest( context.Background(), c, account, body, "sk-ant-xxx", "apikey", "claude-haiku-4-5", false, false, ) diff --git a/backend/internal/service/gateway_forward_as_chat_completions.go b/backend/internal/service/gateway_forward_as_chat_completions.go index eaf67fab0ec8..729483e3be95 100644 --- a/backend/internal/service/gateway_forward_as_chat_completions.go +++ b/backend/internal/service/gateway_forward_as_chat_completions.go @@ -119,7 +119,7 @@ func (s *GatewayService) ForwardAsChatCompletions( // 10. Build upstream request upstreamCtx, releaseUpstreamCtx := detachStreamUpstreamContext(ctx, reqStream) - upstreamReq, err := s.buildUpstreamRequest(upstreamCtx, c, account, anthropicBody, token, tokenType, mappedModel, reqStream, shouldMimicClaudeCode) + upstreamReq, _, err := s.buildUpstreamRequest(upstreamCtx, c, account, anthropicBody, token, tokenType, mappedModel, reqStream, shouldMimicClaudeCode) releaseUpstreamCtx() if err != nil { return nil, fmt.Errorf("build upstream request: %w", err) diff --git a/backend/internal/service/gateway_forward_as_responses.go b/backend/internal/service/gateway_forward_as_responses.go index c55a5a984116..3baea018bd32 100644 --- a/backend/internal/service/gateway_forward_as_responses.go +++ b/backend/internal/service/gateway_forward_as_responses.go @@ -116,7 +116,7 @@ func (s *GatewayService) ForwardAsResponses( // 10. Build upstream request upstreamCtx, releaseUpstreamCtx := detachStreamUpstreamContext(ctx, reqStream) - upstreamReq, err := s.buildUpstreamRequest(upstreamCtx, c, account, anthropicBody, token, tokenType, mappedModel, reqStream, shouldMimicClaudeCode) + upstreamReq, _, err := s.buildUpstreamRequest(upstreamCtx, c, account, anthropicBody, token, tokenType, mappedModel, reqStream, shouldMimicClaudeCode) releaseUpstreamCtx() if err != nil { return nil, fmt.Errorf("build upstream request: %w", err) diff --git a/backend/internal/service/gateway_request.go b/backend/internal/service/gateway_request.go index 65718c27148f..dc59611cead5 100644 --- a/backend/internal/service/gateway_request.go +++ b/backend/internal/service/gateway_request.go @@ -105,6 +105,22 @@ func (r jsonRange) exists() bool { return r.start >= 0 && r.end >= r.start } +// clearGatewayRequestDerivedState 清空绑定当前 body 的轻量派生字段，防止 ReplaceBody 后读到旧值。 +func clearGatewayRequestDerivedState(parsed *ParsedRequest) { + if parsed == nil { + return + } + parsed.Model = "" + parsed.Stream = false + parsed.MetadataUserID = "" + parsed.HasSystem = false + parsed.ThinkingEnabled = false + parsed.OutputEffort = "" + parsed.MaxTokens = 0 + parsed.systemRange = missingJSONRange() + parsed.messagesRange = missingJSONRange() +} + func clearGatewayRequestRanges(parsed *ParsedRequest) { if parsed == nil { return @@ -137,12 +153,9 @@ func setGatewayRequestRanges(parsed *ParsedRequest, protocol string, jsonStr str } } -func refreshGatewayRequestRanges(parsed *ParsedRequest, protocol string) error { - if parsed == nil { - return fmt.Errorf("empty request body") - } - clearGatewayRequestRanges(parsed) - if parsed.Body == nil { +// parseGatewayRequestCurrentBody 只做标量和 raw range 轻量解析，不恢复 system/messages 对象图。 +func parseGatewayRequestCurrentBody(parsed *ParsedRequest, protocol string) error { + if parsed == nil || parsed.Body == nil { return fmt.Errorf("empty request body") } @@ -151,11 +164,51 @@ func refreshGatewayRequestRanges(parsed *ParsedRequest, protocol string) error { return fmt.Errorf("invalid json") } + // 只在当前函数内零拷贝读取 JSON 字段；ReplaceBody 后必须重新进入本函数刷新派生状态。 jsonStr := *(*string)(unsafe.Pointer(&bodyBytes)) + clearGatewayRequestDerivedState(parsed) + parsed.protocol = protocol + + modelResult := gjson.Get(jsonStr, "model") + if modelResult.Exists() { + if modelResult.Type != gjson.String { + return fmt.Errorf("invalid model field type") + } + parsed.Model = modelResult.String() + } + + streamResult := gjson.Get(jsonStr, "stream") + if streamResult.Exists() { + if streamResult.Type != gjson.True && streamResult.Type != gjson.False { + return fmt.Errorf("invalid stream field type") + } + parsed.Stream = streamResult.Bool() + } + + parsed.MetadataUserID = gjson.Get(jsonStr, "metadata.user_id").String() + + thinkingType := gjson.Get(jsonStr, "thinking.type").String() + parsed.ThinkingEnabled = thinkingType == "enabled" || thinkingType == "adaptive" + + parsed.OutputEffort = strings.TrimSpace(gjson.Get(jsonStr, "output_config.effort").String()) + + maxTokensResult := gjson.Get(jsonStr, "max_tokens") + if maxTokensResult.Exists() && maxTokensResult.Type == gjson.Number { + f := maxTokensResult.Float() + if !math.IsNaN(f) && !math.IsInf(f, 0) && f == math.Trunc(f) && + f <= float64(math.MaxInt) && f >= float64(math.MinInt) { + parsed.MaxTokens = int(f) + } + } + setGatewayRequestRanges(parsed, protocol, jsonStr) return nil } +func refreshGatewayRequestRanges(parsed *ParsedRequest, protocol string) error { + return parseGatewayRequestCurrentBody(parsed, protocol) +} + // ParsedRequest 保存网关请求的预解析结果 // // 性能优化说明： @@ -238,71 +291,10 @@ func normalizeSessionUserAgentFallback(raw string) string { // protocol 指定请求协议格式（domain.PlatformAnthropic / domain.PlatformGemini）， // 不同协议使用不同的 system/messages 字段名。 func ParseGatewayRequest(body *RequestBodyRef, protocol string) (*ParsedRequest, error) { - bodyBytes := body.Bytes() - // 保持与旧实现一致：请求体必须是合法 JSON。 - // 注意：gjson.GetBytes 对非法 JSON 不会报错，因此需要显式校验。 - if !gjson.ValidBytes(bodyBytes) { - return nil, fmt.Errorf("invalid json") + parsed := &ParsedRequest{Body: body} + if err := parseGatewayRequestCurrentBody(parsed, protocol); err != nil { + return nil, err } - - // 性能： - // - gjson.GetBytes 会把匹配的 Raw/Str 安全复制成 string（对于巨大 messages 会产生额外拷贝）。 - // - 这里将 body 通过 unsafe 零拷贝视为 string，仅在本函数内使用，且 body 不会被修改。 - jsonStr := *(*string)(unsafe.Pointer(&bodyBytes)) - - parsed := &ParsedRequest{ - Body: body, - protocol: protocol, - systemRange: missingJSONRange(), - messagesRange: missingJSONRange(), - } - - // --- gjson 提取简单字段（避免完整 Unmarshal） --- - - // model: 需要严格类型校验，非 string 返回错误 - modelResult := gjson.Get(jsonStr, "model") - if modelResult.Exists() { - if modelResult.Type != gjson.String { - return nil, fmt.Errorf("invalid model field type") - } - parsed.Model = modelResult.String() - } - - // stream: 需要严格类型校验，非 bool 返回错误 - streamResult := gjson.Get(jsonStr, "stream") - if streamResult.Exists() { - if streamResult.Type != gjson.True && streamResult.Type != gjson.False { - return nil, fmt.Errorf("invalid stream field type") - } - parsed.Stream = streamResult.Bool() - } - - // metadata.user_id: 直接路径提取，不需要严格类型校验 - parsed.MetadataUserID = gjson.Get(jsonStr, "metadata.user_id").String() - - // thinking.type: enabled/adaptive 都视为开启 - thinkingType := gjson.Get(jsonStr, "thinking.type").String() - if thinkingType == "enabled" || thinkingType == "adaptive" { - parsed.ThinkingEnabled = true - } - - // output_config.effort: Claude API 的推理强度控制参数 - parsed.OutputEffort = strings.TrimSpace(gjson.Get(jsonStr, "output_config.effort").String()) - - // max_tokens: 仅接受整数值 - maxTokensResult := gjson.Get(jsonStr, "max_tokens") - if maxTokensResult.Exists() && maxTokensResult.Type == gjson.Number { - f := maxTokensResult.Float() - if !math.IsNaN(f) && !math.IsInf(f, 0) && f == math.Trunc(f) && - f <= float64(math.MaxInt) && f >= float64(math.MinInt) { - parsed.MaxTokens = int(f) - } - } - - // --- system/messages 提取 --- - // 只保存大字段 raw range，不默认反序列化成 []any/map[string]any 对象图。 - setGatewayRequestRanges(parsed, protocol, jsonStr) - return parsed, nil } @@ -353,9 +345,24 @@ func (p *ParsedRequest) SystemValue() (any, bool) { return system, true } -func (p *ParsedRequest) ReplaceBody(data []byte) { +// CloneForBody 为单次账号尝试创建独立 body 视图，避免 failover 复用已改写的 ParsedRequest。 +func (p *ParsedRequest) CloneForBody(body []byte) (*ParsedRequest, error) { if p == nil { - return + return nil, fmt.Errorf("parse request: empty request") + } + clone := *p + clone.Body = NewRequestBodyRef(body) + clone.OnUpstreamAccepted = nil + if err := refreshGatewayRequestRanges(&clone, clone.protocol); err != nil { + return nil, err + } + return &clone, nil +} + +// ReplaceBody 统一刷新当前 body 和 raw range，保证后续 helper 读取的是最新请求体。 +func (p *ParsedRequest) ReplaceBody(data []byte) error { + if p == nil { + return fmt.Errorf("parse request: empty request") } if p.Body == nil { p.Body = NewRequestBodyRef(data) @@ -364,7 +371,9 @@ func (p *ParsedRequest) ReplaceBody(data []byte) { } if err := refreshGatewayRequestRanges(p, p.protocol); err != nil { clearGatewayRequestRanges(p) + return err } + return nil } // sliceRawFromBody 返回 Result.Raw 对应的原始字节切片。 diff --git a/backend/internal/service/gateway_service.go b/backend/internal/service/gateway_service.go index 438592e3f45b..01b541e217a5 100644 --- a/backend/internal/service/gateway_service.go +++ b/backend/internal/service/gateway_service.go @@ -4440,6 +4440,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A } return s.forwardAnthropicAPIKeyPassthroughWithInput(ctx, c, account, anthropicPassthroughForwardInput{ Body: passthroughBody, + Parsed: parsed, RequestModel: passthroughModel, OriginalModel: parsed.Model, RequestStream: parsed.Stream, @@ -4466,6 +4467,13 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A } body := parsed.Body.Bytes() + replaceBody := func(next []byte) error { + if err := parsed.ReplaceBody(next); err != nil { + return fmt.Errorf("rewrite request body: %w", err) + } + body = parsed.Body.Bytes() + return nil + } reqModel := parsed.Model reqStream := parsed.Stream originalModel := reqModel @@ -4499,7 +4507,9 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A systemRewritten := false if !strings.Contains(strings.ToLower(reqModel), "haiku") { systemRaw, _ := parsed.SystemValue() - body = rewriteSystemForNonClaudeCode(body, systemRaw) + if err := replaceBody(rewriteSystemForNonClaudeCode(body, systemRaw)); err != nil { + return nil, err + } systemRewritten = true } @@ -4521,22 +4531,34 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A } } - body, reqModel = normalizeClaudeOAuthRequestBody(body, reqModel, normalizeOpts) + var normalizedBody []byte + normalizedBody, reqModel = normalizeClaudeOAuthRequestBody(body, reqModel, normalizeOpts) + if err := replaceBody(normalizedBody); err != nil { + return nil, err + } // D/E/F: 可选 messages cache 策略 + 工具名混淆 + tools[-1] 断点 // 与 forward_as_chat_completions / forward_as_responses 路径对齐， // 原生 /v1/messages 路径也走同一套可配置字段级改写。 - body = s.rewriteMessageCacheControlIfEnabled(ctx, body) + if err := replaceBody(s.rewriteMessageCacheControlIfEnabled(ctx, body)); err != nil { + return nil, err + } if rw := buildToolNameRewriteFromBody(body); rw != nil { - body = applyToolNameRewriteToBody(body, rw) + if err := replaceBody(applyToolNameRewriteToBody(body, rw)); err != nil { + return nil, err + } c.Set(toolNameRewriteKey, rw) } else { - body = applyToolsLastCacheBreakpoint(body) + if err := replaceBody(applyToolsLastCacheBreakpoint(body)); err != nil { + return nil, err + } } } // 强制执行 cache_control 块数量限制（最多 4 个） - body = enforceCacheControlLimit(body) + if err := replaceBody(enforceCacheControlLimit(body)); err != nil { + return nil, err + } // 应用模型映射： // - APIKey 账号：使用账号级别的显式映射（如果配置），否则透传原始模型名 @@ -4570,13 +4592,18 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A } if mappedModel != reqModel { // 替换请求体中的模型名 - body = s.replaceModelInBody(body, mappedModel) + if err := replaceBody(s.replaceModelInBody(body, mappedModel)); err != nil { + return nil, err + } reqModel = mappedModel + parsed.Model = mappedModel logger.LegacyPrintf("service.gateway", "Model mapping applied: %s -> %s (account: %s, source=%s)", originalModel, mappedModel, account.Name, mappingSource) } if s.shouldInjectAnthropicCacheTTL1h(ctx, account) { - body = injectAnthropicCacheControlTTL1h(body) + if err := replaceBody(injectAnthropicCacheControlTTL1h(body)); err != nil { + return nil, err + } } // 获取凭证 @@ -4600,19 +4627,24 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A logger.LegacyPrintf("service.gateway", "[Forward] Using account: ID=%d Name=%s Platform=%s Type=%s TLSFingerprint=%v Proxy=%s", account.ID, account.Name, account.Platform, account.Type, tlsProfile, proxyURL) // Pre-filter: strip empty text blocks (including nested in tool_result) to prevent upstream 400. - body = StripEmptyTextBlocks(body) + if err := replaceBody(StripEmptyTextBlocks(body)); err != nil { + return nil, err + } // 重试循环 var resp *http.Response + lastWireBody := body retryStart := time.Now() for attempt := 1; attempt <= maxRetryAttempts; attempt++ { // 构建上游请求（每次重试需要重新构建，因为请求体需要重新读取） upstreamCtx, releaseUpstreamCtx := detachStreamUpstreamContext(ctx, reqStream) - upstreamReq, err := s.buildUpstreamRequest(upstreamCtx, c, account, body, token, tokenType, reqModel, reqStream, shouldMimicClaudeCode) + upstreamReq, wireBody, err := s.buildUpstreamRequest(upstreamCtx, c, account, body, token, tokenType, reqModel, reqStream, shouldMimicClaudeCode) releaseUpstreamCtx() if err != nil { return nil, err } + // 记录本次实际发送的 wire body；只有请求成功后才写回 ParsedRequest，避免 400 retry 基于已签名 CCH 再改写。 + lastWireBody = wireBody // 发送请求 resp, err = s.httpUpstream.DoWithTLS(upstreamReq, proxyURL, account.ID, account.Concurrency, tlsProfile) @@ -4690,12 +4722,17 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A filteredBody := FilterThinkingBlocksForRetry(body) retryCtx, releaseRetryCtx := detachStreamUpstreamContext(ctx, reqStream) - retryReq, buildErr := s.buildUpstreamRequest(retryCtx, c, account, filteredBody, token, tokenType, reqModel, reqStream, shouldMimicClaudeCode) + retryReq, retryWireBody, buildErr := s.buildUpstreamRequest(retryCtx, c, account, filteredBody, token, tokenType, reqModel, reqStream, shouldMimicClaudeCode) releaseRetryCtx() if buildErr == nil { retryResp, retryErr := s.httpUpstream.DoWithTLS(retryReq, proxyURL, account.ID, account.Concurrency, tlsProfile) if retryErr == nil { if retryResp.StatusCode < 400 { + // 重试请求被上游接受后同步 ParsedRequest，保证 usage/日志看到真实请求体。 + if err := replaceBody(retryWireBody); err != nil { + _ = retryResp.Body.Close() + return nil, err + } logger.LegacyPrintf("service.gateway", "Account %d: thinking block retry succeeded (blocks downgraded)", account.ID) resp = retryResp break @@ -4725,11 +4762,18 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A logger.LegacyPrintf("service.gateway", "Account %d: signature retry still failing and looks tool-related, retrying with tool blocks downgraded", account.ID) filteredBody2 := FilterSignatureSensitiveBlocksForRetry(body) retryCtx2, releaseRetryCtx2 := detachStreamUpstreamContext(ctx, reqStream) - retryReq2, buildErr2 := s.buildUpstreamRequest(retryCtx2, c, account, filteredBody2, token, tokenType, reqModel, reqStream, shouldMimicClaudeCode) + retryReq2, retryWireBody2, buildErr2 := s.buildUpstreamRequest(retryCtx2, c, account, filteredBody2, token, tokenType, reqModel, reqStream, shouldMimicClaudeCode) releaseRetryCtx2() if buildErr2 == nil { retryResp2, retryErr2 := s.httpUpstream.DoWithTLS(retryReq2, proxyURL, account.ID, account.Concurrency, tlsProfile) if retryErr2 == nil { + if retryResp2.StatusCode < 400 { + // 二阶段工具块降级成功时也必须更新当前 body。 + if err := replaceBody(retryWireBody2); err != nil { + _ = retryResp2.Body.Close() + return nil, err + } + } resp = retryResp2 break } @@ -4796,11 +4840,18 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A if applied && time.Since(retryStart) < maxRetryElapsed { logger.LegacyPrintf("service.gateway", "Account %d: detected budget_tokens constraint error, retrying with rectified budget (budget_tokens=%d, max_tokens=%d)", account.ID, BudgetRectifyBudgetTokens, BudgetRectifyMaxTokens) budgetRetryCtx, releaseBudgetRetryCtx := detachStreamUpstreamContext(ctx, reqStream) - budgetRetryReq, buildErr := s.buildUpstreamRequest(budgetRetryCtx, c, account, rectifiedBody, token, tokenType, reqModel, reqStream, shouldMimicClaudeCode) + budgetRetryReq, budgetWireBody, buildErr := s.buildUpstreamRequest(budgetRetryCtx, c, account, rectifiedBody, token, tokenType, reqModel, reqStream, shouldMimicClaudeCode) releaseBudgetRetryCtx() if buildErr == nil { budgetRetryResp, retryErr := s.httpUpstream.DoWithTLS(budgetRetryReq, proxyURL, account.ID, account.Concurrency, tlsProfile) if retryErr == nil { + if budgetRetryResp.StatusCode < 400 { + // budget 修正请求成功后，ParsedRequest 也要描述被接受的修正版。 + if err := replaceBody(budgetWireBody); err != nil { + _ = budgetRetryResp.Body.Close() + return nil, err + } + } resp = budgetRetryResp break } @@ -4997,6 +5048,13 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A // 处理正常响应 + if !bytes.Equal(lastWireBody, body) { + // 成功后再同步最终 wire body，避免失败重试从已签名 CCH 的 body 继续派生。 + if err := replaceBody(lastWireBody); err != nil { + return nil, err + } + } + // 触发上游接受回调（提前释放串行锁，不等流完成） if parsed.OnUpstreamAccepted != nil { parsed.OnUpstreamAccepted() @@ -5039,6 +5097,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A type anthropicPassthroughForwardInput struct { Body []byte + Parsed *ParsedRequest RequestModel string OriginalModel string RequestStream bool @@ -5091,16 +5150,29 @@ func (s *GatewayService) forwardAnthropicAPIKeyPassthroughWithInput( } // Pre-filter: strip empty text blocks (including nested in tool_result) to prevent upstream 400. input.Body = StripEmptyTextBlocks(input.Body) + if input.Parsed != nil { + // 透传分支也会改写实际 wire body，成功 usage hash 依赖这里同步当前 body。 + if err := input.Parsed.ReplaceBody(input.Body); err != nil { + return nil, err + } + } var resp *http.Response retryStart := time.Now() for attempt := 1; attempt <= maxRetryAttempts; attempt++ { upstreamCtx, releaseUpstreamCtx := detachStreamUpstreamContext(ctx, input.RequestStream) - upstreamReq, err := s.buildUpstreamRequestAnthropicAPIKeyPassthrough(upstreamCtx, c, account, input.Body, token) + upstreamReq, wireBody, err := s.buildUpstreamRequestAnthropicAPIKeyPassthrough(upstreamCtx, c, account, input.Body, token) releaseUpstreamCtx() if err != nil { return nil, err } + if input.Parsed != nil && !bytes.Equal(wireBody, input.Body) { + // build 阶段会按 beta 能力清理 body，发送前同步到 ParsedRequest 当前视图。 + if err := input.Parsed.ReplaceBody(wireBody); err != nil { + return nil, err + } + input.Body = input.Parsed.Body.Bytes() + } resp, err = s.httpUpstream.DoWithTLS(upstreamReq, proxyURL, account.ID, account.Concurrency, s.tlsFPProfileService.ResolveTLSProfile(account)) if err != nil { @@ -5292,13 +5364,13 @@ func (s *GatewayService) buildUpstreamRequestAnthropicAPIKeyPassthrough( account *Account, body []byte, token string, -) (*http.Request, error) { +) (*http.Request, []byte, error) { targetURL := claudeAPIURL baseURL := account.GetBaseURL() if baseURL != "" { validatedURL, err := s.validateUpstreamBaseURL(baseURL) if err != nil { - return nil, err + return nil, nil, err } targetURL = validatedURL + "/v1/messages?beta=true" } @@ -5316,7 +5388,7 @@ func (s *GatewayService) buildUpstreamRequestAnthropicAPIKeyPassthrough( req, err := http.NewRequestWithContext(ctx, http.MethodPost, targetURL, bytes.NewReader(body)) if err != nil { - return nil, err + return nil, nil, err } if c != nil && c.Request != nil { @@ -5346,7 +5418,7 @@ func (s *GatewayService) buildUpstreamRequestAnthropicAPIKeyPassthrough( setHeaderRaw(req.Header, "anthropic-version", "2023-06-01") } - return req, nil + return req, body, nil } func (s *GatewayService) handleStreamingResponseAnthropicAPIKeyPassthrough( @@ -5828,6 +5900,11 @@ func (s *GatewayService) forwardBedrock( return s.handleBedrockUpstreamErrors(ctx, resp, c, account) } + // Bedrock 分支绕过通用 Forward 成功路径，这里保持上游接受回调语义一致。 + if parsed.OnUpstreamAccepted != nil { + parsed.OnUpstreamAccepted() + } + // 响应处理 var usage *ClaudeUsage var firstTokenMs *int @@ -6104,9 +6181,10 @@ func (s *GatewayService) handleBedrockNonStreamingResponse( return usage, nil } -func (s *GatewayService) buildUpstreamRequest(ctx context.Context, c *gin.Context, account *Account, body []byte, token, tokenType, modelID string, reqStream bool, mimicClaudeCode bool) (*http.Request, error) { +func (s *GatewayService) buildUpstreamRequest(ctx context.Context, c *gin.Context, account *Account, body []byte, token, tokenType, modelID string, reqStream bool, mimicClaudeCode bool) (*http.Request, []byte, error) { if account.Platform == PlatformAnthropic && account.Type == AccountTypeServiceAccount { - return s.buildUpstreamRequestAnthropicVertex(ctx, c, account, body, token, modelID, reqStream) + req, err := s.buildUpstreamRequestAnthropicVertex(ctx, c, account, body, token, modelID, reqStream) + return req, body, err } // 确定目标URL @@ -6116,18 +6194,18 @@ func (s *GatewayService) buildUpstreamRequest(ctx context.Context, c *gin.Contex if baseURL != "" { validatedURL, err := s.validateUpstreamBaseURL(baseURL) if err != nil { - return nil, err + return nil, nil, err } targetURL = validatedURL + "/v1/messages?beta=true" } } else if account.IsCustomBaseURLEnabled() { customURL := account.GetCustomBaseURL() if customURL == "" { - return nil, fmt.Errorf("custom_base_url is enabled but not configured for account %d", account.ID) + return nil, nil, fmt.Errorf("custom_base_url is enabled but not configured for account %d", account.ID) } validatedURL, err := s.validateUpstreamBaseURL(customURL) if err != nil { - return nil, err + return nil, nil, err } targetURL = s.buildCustomRelayURL(validatedURL, "/v1/messages", account) } @@ -6202,7 +6280,7 @@ func (s *GatewayService) buildUpstreamRequest(ctx context.Context, c *gin.Contex req, err := http.NewRequestWithContext(ctx, "POST", targetURL, bytes.NewReader(body)) if err != nil { - return nil, err + return nil, nil, err } // 设置认证头（保持原始大小写） @@ -6287,7 +6365,7 @@ func (s *GatewayService) buildUpstreamRequest(ctx context.Context, c *gin.Contex logClaudeMimicDebug(req, body, account, tokenType, mimicClaudeCode) } - return req, nil + return req, body, nil } func (s *GatewayService) buildUpstreamRequestAnthropicVertex( @@ -9214,23 +9292,42 @@ func (s *GatewayService) ForwardCountTokens(ctx context.Context, c *gin.Context, } body := parsed.Body.Bytes() + replaceBody := func(next []byte) error { + if err := parsed.ReplaceBody(next); err != nil { + return fmt.Errorf("rewrite count_tokens body: %w", err) + } + body = parsed.Body.Bytes() + return nil + } reqModel := parsed.Model // Pre-filter: strip empty text blocks to prevent upstream 400. - body = StripEmptyTextBlocks(body) + if err := replaceBody(StripEmptyTextBlocks(body)); err != nil { + return err + } isClaudeCodeCT := IsClaudeCodeClient(ctx) || isClaudeCodeClient(c.GetHeader("User-Agent"), parsed.MetadataUserID) shouldMimicClaudeCode := account.IsOAuth() && !isClaudeCodeCT if shouldMimicClaudeCode { normalizeOpts := claudeOAuthNormalizeOptions{stripSystemCacheControl: true} - body, reqModel = normalizeClaudeOAuthRequestBody(body, reqModel, normalizeOpts) + var normalizedBody []byte + normalizedBody, reqModel = normalizeClaudeOAuthRequestBody(body, reqModel, normalizeOpts) + if err := replaceBody(normalizedBody); err != nil { + return err + } - body = s.rewriteMessageCacheControlIfEnabled(ctx, body) + if err := replaceBody(s.rewriteMessageCacheControlIfEnabled(ctx, body)); err != nil { + return err + } if rw := buildToolNameRewriteFromBody(body); rw != nil { - body = applyToolNameRewriteToBody(body, rw) + if err := replaceBody(applyToolNameRewriteToBody(body, rw)); err != nil { + return err + } } else { - body = applyToolsLastCacheBreakpoint(body) + if err := replaceBody(applyToolsLastCacheBreakpoint(body)); err != nil { + return err + } } } @@ -9261,9 +9358,13 @@ func (s *GatewayService) ForwardCountTokens(ctx context.Context, c *gin.Context, } } if mappedModel != reqModel { - body = s.replaceModelInBody(body, mappedModel) + originalReqModel := reqModel + if err := replaceBody(s.replaceModelInBody(body, mappedModel)); err != nil { + return err + } reqModel = mappedModel - logger.LegacyPrintf("service.gateway", "CountTokens model mapping applied: %s -> %s (account: %s, source=%s)", parsed.Model, mappedModel, account.Name, mappingSource) + parsed.Model = mappedModel + logger.LegacyPrintf("service.gateway", "CountTokens model mapping applied: %s -> %s (account: %s, source=%s)", originalReqModel, mappedModel, account.Name, mappingSource) } } @@ -9275,11 +9376,13 @@ func (s *GatewayService) ForwardCountTokens(ctx context.Context, c *gin.Context, } // 构建上游请求 - upstreamReq, err := s.buildCountTokensRequest(ctx, c, account, body, token, tokenType, reqModel, shouldMimicClaudeCode) + upstreamReq, wireBody, err := s.buildCountTokensRequest(ctx, c, account, body, token, tokenType, reqModel, shouldMimicClaudeCode) if err != nil { s.countTokensError(c, http.StatusInternalServerError, "api_error", "Failed to build request") return err } + // 先记录首发 wire body；如果后面进入 400 retry，retry 会基于未签名的逻辑 body 重新构建。 + acceptedWireBody := wireBody // 获取代理URL（自定义 base URL 模式下，proxy 通过 buildCustomRelayURL 作为查询参数传递） proxyURL := "" @@ -9315,10 +9418,14 @@ func (s *GatewayService) ForwardCountTokens(ctx context.Context, c *gin.Context, logger.LegacyPrintf("service.gateway", "Account %d: detected thinking block signature error on count_tokens, retrying with filtered thinking blocks", account.ID) filteredBody := FilterThinkingBlocksForRetry(body) - retryReq, buildErr := s.buildCountTokensRequest(ctx, c, account, filteredBody, token, tokenType, reqModel, shouldMimicClaudeCode) + retryReq, retryWireBody, buildErr := s.buildCountTokensRequest(ctx, c, account, filteredBody, token, tokenType, reqModel, shouldMimicClaudeCode) if buildErr == nil { retryResp, retryErr := s.httpUpstream.DoWithTLS(retryReq, proxyURL, account.ID, account.Concurrency, s.tlsFPProfileService.ResolveTLSProfile(account)) if retryErr == nil { + if retryResp.StatusCode < 400 { + // count_tokens 签名重试成功后记录最终 wire body，错误响应仍保留原 body 便于后续处理。 + acceptedWireBody = retryWireBody + } resp = retryResp respBody, err = ReadUpstreamResponseBody(resp.Body, s.cfg, c, countTokensTooLarge) _ = resp.Body.Close() @@ -9332,6 +9439,13 @@ func (s *GatewayService) ForwardCountTokens(ctx context.Context, c *gin.Context, } } + if resp.StatusCode < 400 && !bytes.Equal(acceptedWireBody, body) { + // count_tokens 成功后再同步最终 wire body，避免 retry 从已签名 body 派生。 + if err := replaceBody(acceptedWireBody); err != nil { + return err + } + } + // 处理错误响应 if resp.StatusCode >= 400 { // 标记账号状态（429/529等） @@ -9558,7 +9672,7 @@ func (s *GatewayService) buildCountTokensRequestAnthropicAPIKeyPassthrough( } // buildCountTokensRequest 构建 count_tokens 上游请求 -func (s *GatewayService) buildCountTokensRequest(ctx context.Context, c *gin.Context, account *Account, body []byte, token, tokenType, modelID string, mimicClaudeCode bool) (*http.Request, error) { +func (s *GatewayService) buildCountTokensRequest(ctx context.Context, c *gin.Context, account *Account, body []byte, token, tokenType, modelID string, mimicClaudeCode bool) (*http.Request, []byte, error) { // 确定目标 URL targetURL := claudeAPICountTokensURL if account.Type == AccountTypeAPIKey { @@ -9566,18 +9680,18 @@ func (s *GatewayService) buildCountTokensRequest(ctx context.Context, c *gin.Con if baseURL != "" { validatedURL, err := s.validateUpstreamBaseURL(baseURL) if err != nil { - return nil, err + return nil, nil, err } targetURL = validatedURL + "/v1/messages/count_tokens?beta=true" } } else if account.IsCustomBaseURLEnabled() { customURL := account.GetCustomBaseURL() if customURL == "" { - return nil, fmt.Errorf("custom_base_url is enabled but not configured for account %d", account.ID) + return nil, nil, fmt.Errorf("custom_base_url is enabled but not configured for account %d", account.ID) } validatedURL, err := s.validateUpstreamBaseURL(customURL) if err != nil { - return nil, err + return nil, nil, err } targetURL = s.buildCustomRelayURL(validatedURL, "/v1/messages/count_tokens", account) } @@ -9633,7 +9747,7 @@ func (s *GatewayService) buildCountTokensRequest(ctx context.Context, c *gin.Con req, err := http.NewRequestWithContext(ctx, "POST", targetURL, bytes.NewReader(body)) if err != nil { - return nil, err + return nil, nil, err } // 设置认证头（保持原始大小写） @@ -9697,7 +9811,7 @@ func (s *GatewayService) buildCountTokensRequest(ctx context.Context, c *gin.Con logClaudeMimicDebug(req, body, account, tokenType, mimicClaudeCode) } - return req, nil + return req, body, nil } func sanitizeCountTokensRequestBody(body []byte) []byte { From 2caee9d884224c43e5c2af1c992ccde99b51d2c7 Mon Sep 17 00:00:00 2001 From: name <136912576+is7Qin@users.noreply.github.com> Date: Sat, 30 May 2026 15:52:33 +0800 Subject: [PATCH 20/36] refactor(gateway): snapshot usage worker inputs --- backend/internal/handler/gateway_handler.go | 10 ++++---- .../internal/handler/gemini_v1beta_handler.go | 4 +++- .../handler/openai_gateway_handler.go | 6 ++++- backend/internal/handler/openai_images.go | 17 ++++++------- backend/internal/service/gateway_service.go | 24 +++++++------------ 5 files changed, 31 insertions(+), 30 deletions(-) diff --git a/backend/internal/handler/gateway_handler.go b/backend/internal/handler/gateway_handler.go index c9b3297e1475..3bd5e82e7b35 100644 --- a/backend/internal/handler/gateway_handler.go +++ b/backend/internal/handler/gateway_handler.go @@ -510,11 +510,12 @@ func (h *GatewayHandler) Messages(c *gin.Context) { } // 使用量记录通过有界 worker 池提交，避免请求热路径创建无界 goroutine。 + // ForceCacheBilling 提前拍成标量，避免 worker 闭包保活 failover 状态里的响应体。 + forceCacheBilling := fs.ForceCacheBilling quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey) h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) { if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{ Result: result, - ParsedRequest: parsedReq, QuotaPlatform: quotaPlatform, APIKey: apiKey, User: apiKey.User, @@ -525,7 +526,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) { UserAgent: userAgent, IPAddress: clientIP, RequestPayloadHash: requestPayloadHash, - ForceCacheBilling: fs.ForceCacheBilling, + ForceCacheBilling: forceCacheBilling, APIKeyService: h.apiKeyService, ChannelUsageFields: channelMapping.ToUsageFields(reqModel, result.UpstreamModel), }); err != nil { @@ -918,11 +919,12 @@ func (h *GatewayHandler) Messages(c *gin.Context) { } // 使用量记录通过有界 worker 池提交，避免请求热路径创建无界 goroutine。 + // ForceCacheBilling 提前拍成标量，避免 worker 闭包保活 failover 状态里的响应体。 + forceCacheBilling := fs.ForceCacheBilling quotaPlatform := service.QuotaPlatform(c.Request.Context(), currentAPIKey) h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) { if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{ Result: result, - ParsedRequest: attemptParsedReq, QuotaPlatform: quotaPlatform, APIKey: currentAPIKey, User: currentAPIKey.User, @@ -933,7 +935,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) { UserAgent: userAgent, IPAddress: clientIP, RequestPayloadHash: requestPayloadHash, - ForceCacheBilling: fs.ForceCacheBilling, + ForceCacheBilling: forceCacheBilling, APIKeyService: h.apiKeyService, ChannelUsageFields: channelMapping.ToUsageFields(reqModel, result.UpstreamModel), }); err != nil { diff --git a/backend/internal/handler/gemini_v1beta_handler.go b/backend/internal/handler/gemini_v1beta_handler.go index 5d8e6fa890a8..024f8727a220 100644 --- a/backend/internal/handler/gemini_v1beta_handler.go +++ b/backend/internal/handler/gemini_v1beta_handler.go @@ -527,6 +527,8 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) { requestPayloadHash := service.HashUsageRequestPayload(body) inboundEndpoint := GetInboundEndpoint(c) upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform) + // ForceCacheBilling 提前拍成标量，避免 worker 闭包保活 failover 状态里的响应体。 + forceCacheBilling := fs.ForceCacheBilling quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey) h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) { if err := h.gatewayService.RecordUsageWithLongContext(ctx, &service.RecordUsageLongContextInput{ @@ -543,7 +545,7 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) { RequestPayloadHash: requestPayloadHash, LongContextThreshold: 200000, // Gemini 200K 阈值 LongContextMultiplier: 2.0, // 超出部分双倍计费 - ForceCacheBilling: fs.ForceCacheBilling, + ForceCacheBilling: forceCacheBilling, APIKeyService: h.apiKeyService, ChannelUsageFields: channelMapping.ToUsageFields(reqModel, result.UpstreamModel), }); err != nil { diff --git a/backend/internal/handler/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index a131cbd909d2..ac312dbbabea 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -1379,6 +1379,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { zap.Int("candidate_count", scheduleDecision.CandidateCount), ) + var requestPayloadHash string hooks := &service.OpenAIWSIngressHooks{ InitialRequestModel: reqModel, BeforeRequest: func(turn int, payload []byte, originalModel string) error { @@ -1464,7 +1465,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { UpstreamEndpoint: upstreamEndpoint, UserAgent: userAgent, IPAddress: clientIP, - RequestPayloadHash: service.HashUsageRequestPayload(firstMessage), + RequestPayloadHash: requestPayloadHash, APIKeyService: h.apiKeyService, ChannelUsageFields: channelMappingWS.ToUsageFields(reqModel, result.UpstreamModel), }); err != nil { @@ -1484,6 +1485,9 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { wsFirstMessage = h.gatewayService.ReplaceModelInBody(firstMessage, channelMappingWS.MappedModel) } + // WebSocket 首包可能很大，hash 必须在 hooks 外算成字符串，避免 AfterTurn 闭包保活请求体。 + requestPayloadHash = service.HashUsageRequestPayload(wsFirstMessage) + if err := h.gatewayService.ProxyResponsesWebSocketFromClient(ctx, c, wsConn, account, token, wsFirstMessage, hooks); err != nil { var failoverErr *service.UpstreamFailoverError if errors.As(err, &failoverErr) { diff --git a/backend/internal/handler/openai_images.go b/backend/internal/handler/openai_images.go index 36339d4b7582..1e3b530601e3 100644 --- a/backend/internal/handler/openai_images.go +++ b/backend/internal/handler/openai_images.go @@ -73,9 +73,10 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", err.Error()) return } + requestModel := parsed.Model reqLog = reqLog.With( - zap.String("model", parsed.Model), + zap.String("model", requestModel), zap.Bool("stream", parsed.Stream), zap.Bool("multipart", parsed.Multipart), zap.String("capability", string(parsed.RequiredCapability)), @@ -85,7 +86,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { h.errorResponse(c, http.StatusForbidden, "permission_error", service.ImageGenerationPermissionMessage()) return } - if decision := h.checkContentModeration(c, reqLog, apiKey, subject, service.ContentModerationProtocolOpenAIImages, parsed.Model, parsed.ModerationBody()); decision != nil && decision.Blocked { + if decision := h.checkContentModeration(c, reqLog, apiKey, subject, service.ContentModerationProtocolOpenAIImages, requestModel, parsed.ModerationBody()); decision != nil && decision.Blocked { h.errorResponse(c, contentModerationStatus(decision), contentModerationErrorCode(decision), decision.Message) return } @@ -98,13 +99,13 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { } if parsed.Multipart { - setOpsRequestContext(c, parsed.Model, parsed.Stream) + setOpsRequestContext(c, requestModel, parsed.Stream) } else { - setOpsRequestContext(c, parsed.Model, parsed.Stream) + setOpsRequestContext(c, requestModel, parsed.Stream) } setOpsEndpointContext(c, "", int16(service.RequestTypeFromLegacy(parsed.Stream, false))) - channelMapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(c.Request.Context(), apiKey.GroupID, parsed.Model) + channelMapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(c.Request.Context(), apiKey.GroupID, requestModel) if h.errorPassthroughService != nil { service.BindErrorPassthroughService(c, h.errorPassthroughService) @@ -147,7 +148,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { c.Request.Context(), apiKey.GroupID, sessionHash, - parsed.Model, + requestModel, failedAccountIDs, parsed.RequiredCapability, ) @@ -324,14 +325,14 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { IPAddress: clientIP, RequestPayloadHash: requestPayloadHash, APIKeyService: h.apiKeyService, - ChannelUsageFields: channelMapping.ToUsageFields(parsed.Model, upstreamModel), + ChannelUsageFields: channelMapping.ToUsageFields(requestModel, upstreamModel), }); err != nil { logger.L().With( zap.String("component", "handler.openai_gateway.images"), zap.Int64("user_id", subject.UserID), zap.Int64("api_key_id", apiKey.ID), zap.Any("group_id", apiKey.GroupID), - zap.String("model", parsed.Model), + zap.String("model", requestModel), zap.Int64("account_id", account.ID), ).Error("openai.images.record_usage_failed", zap.Error(err)) } diff --git a/backend/internal/service/gateway_service.go b/backend/internal/service/gateway_service.go index 01b541e217a5..78a1f4bdabbd 100644 --- a/backend/internal/service/gateway_service.go +++ b/backend/internal/service/gateway_service.go @@ -4729,6 +4729,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A if retryErr == nil { if retryResp.StatusCode < 400 { // 重试请求被上游接受后同步 ParsedRequest，保证 usage/日志看到真实请求体。 + lastWireBody = retryWireBody if err := replaceBody(retryWireBody); err != nil { _ = retryResp.Body.Close() return nil, err @@ -4769,6 +4770,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A if retryErr2 == nil { if retryResp2.StatusCode < 400 { // 二阶段工具块降级成功时也必须更新当前 body。 + lastWireBody = retryWireBody2 if err := replaceBody(retryWireBody2); err != nil { _ = retryResp2.Body.Close() return nil, err @@ -4847,6 +4849,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A if retryErr == nil { if budgetRetryResp.StatusCode < 400 { // budget 修正请求成功后，ParsedRequest 也要描述被接受的修正版。 + lastWireBody = budgetWireBody if err := replaceBody(budgetWireBody); err != nil { _ = budgetRetryResp.Body.Close() return nil, err @@ -8228,10 +8231,10 @@ func (s *GatewayService) getUserGroupRateMultiplier(ctx context.Context, userID, return resolver.Resolve(ctx, userID, groupID, groupDefaultMultiplier) } -// RecordUsageInput 记录使用量的输入参数 +// RecordUsageInput 记录使用量的输入参数。 +// 异步 worker 只接收计费所需快照，不能持有 ParsedRequest/RequestBodyRef 这类大请求体引用。 type RecordUsageInput struct { Result *ForwardResult - ParsedRequest *ParsedRequest APIKey *APIKey User *User Account *Account @@ -8709,15 +8712,8 @@ func writeUsageLogBestEffort(ctx context.Context, repo UsageLogRepository, usage } } -// recordUsageOpts 内部选项，参数化 RecordUsage 与 RecordUsageWithLongContext 的差异点。 +// recordUsageOpts 内部选项，参数化普通计费与长上下文计费的差异点。 type recordUsageOpts struct { - // Claude Max 策略所需的 ParsedRequest（可选，仅 Claude 路径传入） - ParsedRequest *ParsedRequest - - // EnableClaudePath 启用 Claude 路径特有逻辑： - // - Claude Max 缓存计费策略 - EnableClaudePath bool - // 长上下文计费（仅 Gemini 路径需要） LongContextThreshold int LongContextMultiplier float64 @@ -8740,9 +8736,7 @@ func (s *GatewayService) RecordUsage(ctx context.Context, input *RecordUsageInpu APIKeyService: input.APIKeyService, QuotaPlatform: input.QuotaPlatform, ChannelUsageFields: input.ChannelUsageFields, - }, &recordUsageOpts{ - EnableClaudePath: true, - }) + }, &recordUsageOpts{}) } // RecordUsageLongContextInput 记录使用量的输入参数（支持长上下文双倍计费） @@ -8808,9 +8802,7 @@ type recordUsageCoreInput struct { } // recordUsageCore 是 RecordUsage 和 RecordUsageWithLongContext 的统一实现。 -// opts 中的字段控制两者之间的差异行为： -// - ParsedRequest != nil → 启用 Claude Max 缓存计费策略 -// - LongContextThreshold > 0 → Token 计费回退走 CalculateCostWithLongContext +// LongContextThreshold > 0 时 Token 计费回退走 CalculateCostWithLongContext。 func (s *GatewayService) recordUsageCore(ctx context.Context, input *recordUsageCoreInput, opts *recordUsageOpts) error { result := input.Result apiKey := input.APIKey From 5a3e193b5371b7fd89f1b23c30c69357fe9e8d8c Mon Sep 17 00:00:00 2001 From: name <136912576+is7Qin@users.noreply.github.com> Date: Sat, 30 May 2026 16:52:15 +0800 Subject: [PATCH 21/36] refactor(gateway): shorten OpenAI request body retention --- backend/internal/handler/gateway_helper.go | 27 -------- .../handler/gateway_helper_hotpath_test.go | 20 +----- .../handler/openai_gateway_handler.go | 11 +--- .../service/openai_gateway_service.go | 62 +---------------- .../openai_gateway_service_hotpath_test.go | 32 +-------- .../service/openai_tool_continuation.go | 63 +++++++++++++++++- .../service/openai_tool_continuation_test.go | 66 +++++++++++++++++++ 7 files changed, 137 insertions(+), 144 deletions(-) diff --git a/backend/internal/handler/gateway_helper.go b/backend/internal/handler/gateway_helper.go index 52362dd11db5..4b6a47eb2fa7 100644 --- a/backend/internal/handler/gateway_helper.go +++ b/backend/internal/handler/gateway_helper.go @@ -18,18 +18,12 @@ import ( // claudeCodeValidator is a singleton validator for Claude Code client detection var claudeCodeValidator = service.NewClaudeCodeValidator() -const claudeCodeParsedRequestContextKey = "claude_code_parsed_request" - // SetClaudeCodeClientContext 检查请求是否来自 Claude Code 客户端，并设置到 context 中 // 返回更新后的 context func SetClaudeCodeClientContext(c *gin.Context, body []byte, parsedReq *service.ParsedRequest) { if c == nil || c.Request == nil { return } - if parsedReq != nil { - c.Set(claudeCodeParsedRequestContextKey, parsedReq) - } - ua := c.GetHeader("User-Agent") // Fast path：非 Claude CLI UA 直接判定 false，避免热路径二次 JSON 反序列化。 if !claudeCodeValidator.ValidateUserAgent(ua) { @@ -45,9 +39,6 @@ func SetClaudeCodeClientContext(c *gin.Context, body []byte, parsedReq *service. } else { // 仅在确认为 Claude CLI 且 messages 路径时再做 body 解析。 bodyMap := claudeCodeBodyMapFromParsedRequest(parsedReq) - if bodyMap == nil { - bodyMap = claudeCodeBodyMapFromContextCache(c) - } if bodyMap == nil && len(body) > 0 { _ = json.Unmarshal(body, &bodyMap) } @@ -87,24 +78,6 @@ func claudeCodeBodyMapFromParsedRequest(parsedReq *service.ParsedRequest) map[st return bodyMap } -func claudeCodeBodyMapFromContextCache(c *gin.Context) map[string]any { - if c == nil { - return nil - } - if bodyMap := service.CachedOpenAIParsedRequestBody(c); bodyMap != nil { - return bodyMap - } - if cached, ok := c.Get(claudeCodeParsedRequestContextKey); ok { - switch v := cached.(type) { - case *service.ParsedRequest: - return claudeCodeBodyMapFromParsedRequest(v) - case service.ParsedRequest: - return claudeCodeBodyMapFromParsedRequest(&v) - } - } - return nil -} - // 并发槽位等待相关常量 // // 性能优化说明： diff --git a/backend/internal/handler/gateway_helper_hotpath_test.go b/backend/internal/handler/gateway_helper_hotpath_test.go index d973bedbc8cd..a6b6a429e9b7 100644 --- a/backend/internal/handler/gateway_helper_hotpath_test.go +++ b/backend/internal/handler/gateway_helper_hotpath_test.go @@ -177,7 +177,7 @@ func TestSetClaudeCodeClientContext_FastPathAndStrictPath(t *testing.T) { }) } -func TestSetClaudeCodeClientContext_ReuseParsedRequestAndContextCache(t *testing.T) { +func TestSetClaudeCodeClientContext_ReuseParsedRequest(t *testing.T) { t.Run("reuse parsed request without body unmarshal", func(t *testing.T) { c, _ := newHelperTestContext(http.MethodPost, "/v1/messages") c.Request.Header.Set("User-Agent", "claude-cli/1.0.1") @@ -192,24 +192,6 @@ func TestSetClaudeCodeClientContext_ReuseParsedRequestAndContextCache(t *testing SetClaudeCodeClientContext(c, []byte(`{invalid`), parsedReq) require.True(t, service.IsClaudeCodeClient(c.Request.Context())) }) - - t.Run("reuse context cache without body unmarshal", func(t *testing.T) { - c, _ := newHelperTestContext(http.MethodPost, "/v1/messages") - c.Request.Header.Set("User-Agent", "claude-cli/1.0.1") - c.Request.Header.Set("X-App", "claude-code") - c.Request.Header.Set("anthropic-beta", "message-batches-2024-09-24") - c.Request.Header.Set("anthropic-version", "2023-06-01") - service.CacheOpenAIParsedRequestBody(c, []byte(`{invalid`), map[string]any{ - "model": "claude-3-5-sonnet-20241022", - "system": []any{ - map[string]any{"text": "You are Claude Code, Anthropic's official CLI for Claude."}, - }, - "metadata": map[string]any{"user_id": "user_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa_account__session_aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa"}, - }) - - SetClaudeCodeClientContext(c, []byte(`{invalid`), nil) - require.True(t, service.IsClaudeCodeClient(c.Request.Context())) - }) } func TestWaitForSlotWithPingTimeout_AccountAndUserAcquire(t *testing.T) { diff --git a/backend/internal/handler/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index ac312dbbabea..0aa477b08cd8 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -948,19 +948,12 @@ func (h *OpenAIGatewayHandler) validateFunctionCallOutputRequest(c *gin.Context, return true } - var reqBody map[string]any - if err := json.Unmarshal(body, &reqBody); err != nil { - // 保持原有容错语义：解析失败时跳过预校验，沿用后续上游校验结果。 - return true - } - - service.CacheOpenAIParsedRequestBody(c, body, reqBody) - validation := service.ValidateFunctionCallOutputContext(reqBody) + validation := service.ValidateFunctionCallOutputContextBytes(body) if !validation.HasFunctionCallOutput { return true } - previousResponseID, _ := reqBody["previous_response_id"].(string) + previousResponseID := gjson.GetBytes(body, "previous_response_id").String() if strings.TrimSpace(previousResponseID) != "" || validation.HasToolCallContext { return true } diff --git a/backend/internal/service/openai_gateway_service.go b/backend/internal/service/openai_gateway_service.go index f8c3403120b6..62182655c722 100644 --- a/backend/internal/service/openai_gateway_service.go +++ b/backend/internal/service/openai_gateway_service.go @@ -36,13 +36,6 @@ import ( "go.uber.org/zap" ) -// openAIParsedRequestBodyCache 绑定 body 指纹，避免 handler 预解析的旧 body 污染后续 forwardBody。 -type openAIParsedRequestBodyCache struct { - bodyHash uint64 - bodyLen int - reqBody map[string]any -} - const ( // ChatGPT internal API for OAuth accounts chatgptCodexURL = "https://chatgpt.com/backend-api/codex/responses" @@ -53,8 +46,6 @@ const ( // codex_cli_only 拒绝时单个请求头日志长度上限（字符） codexCLIOnlyHeaderValueMaxBytes = 256 - // OpenAIParsedRequestBodyKey 缓存 handler 侧已解析的请求体，避免重复解析。 - OpenAIParsedRequestBodyKey = "openai_parsed_request_body" // OpenAI WS Mode 失败后的重连次数上限（不含首次尝试）。 // 与 Codex 客户端保持一致：失败后最多重连 5 次。 openAIWSReconnectRetryLimit = 5 @@ -3099,8 +3090,6 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco }) s.handleFailoverSideEffects(ctx, resp, account, upstreamModel) - // reqBody 会被本次账号尝试原地修改，failover 前必须释放，避免下一账号复用脏 map。 - releaseOpenAIParsedRequestBody(c) return nil, &UpstreamFailoverError{ StatusCode: resp.StatusCode, ResponseBody: respBody, @@ -3113,7 +3102,8 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco reasoningEffort := extractOpenAIReasoningEffort(reqBody, originalModel) serviceTier := extractOpenAIServiceTier(reqBody) - releaseOpenAIParsedRequestBody(c) + // 上游接受后只保留计费需要的标量，避免响应处理期间继续保活完整 input/tools map。 + reqBody = nil // Handle normal response var usage *OpenAIUsage @@ -6843,60 +6833,14 @@ func isEmptyBase64DataURI(raw string) bool { return strings.TrimSpace(strings.TrimPrefix(rest, "base64,")) == "" } -func getOpenAIRequestBodyMap(c *gin.Context, body []byte) (map[string]any, error) { - // 同一个 gin.Context 内 failover/渠道映射可能传入新 body，缓存必须先校验 body 指纹。 - bodyHash := xxhash.Sum64(body) - bodyLen := len(body) - if c != nil { - if cached, ok := c.Get(OpenAIParsedRequestBodyKey); ok { - if cache, ok := cached.(openAIParsedRequestBodyCache); ok && cache.reqBody != nil && cache.bodyLen == bodyLen && cache.bodyHash == bodyHash { - return cache.reqBody, nil - } - } - } - +func getOpenAIRequestBodyMap(_ *gin.Context, body []byte) (map[string]any, error) { var reqBody map[string]any if err := json.Unmarshal(body, &reqBody); err != nil { return nil, fmt.Errorf("parse request: %w", err) } - if c != nil { - c.Set(OpenAIParsedRequestBodyKey, openAIParsedRequestBodyCache{bodyHash: bodyHash, bodyLen: bodyLen, reqBody: reqBody}) - } return reqBody, nil } -// CacheOpenAIParsedRequestBody 仅缓存与当前 body 绑定的解析结果。 -func CacheOpenAIParsedRequestBody(c *gin.Context, body []byte, reqBody map[string]any) { - if c == nil || reqBody == nil { - return - } - c.Set(OpenAIParsedRequestBodyKey, openAIParsedRequestBodyCache{ - bodyHash: xxhash.Sum64(body), - bodyLen: len(body), - reqBody: reqBody, - }) -} - -// CachedOpenAIParsedRequestBody 只给同请求内不关心 body 参数的轻量识别逻辑使用。 -func CachedOpenAIParsedRequestBody(c *gin.Context) map[string]any { - if c == nil { - return nil - } - if cached, ok := c.Get(OpenAIParsedRequestBodyKey); ok { - if cache, ok := cached.(openAIParsedRequestBodyCache); ok { - return cache.reqBody - } - } - return nil -} - -func releaseOpenAIParsedRequestBody(c *gin.Context) { - if c == nil { - return - } - delete(c.Keys, OpenAIParsedRequestBodyKey) -} - func extractOpenAIReasoningEffort(reqBody map[string]any, requestedModel string) *string { if value, present := getOpenAIReasoningEffortFromReqBody(reqBody); present { if value == "" { diff --git a/backend/internal/service/openai_gateway_service_hotpath_test.go b/backend/internal/service/openai_gateway_service_hotpath_test.go index 2ff72e2fa418..af0d21c4a176 100644 --- a/backend/internal/service/openai_gateway_service_hotpath_test.go +++ b/backend/internal/service/openai_gateway_service_hotpath_test.go @@ -106,26 +106,13 @@ func TestExtractOpenAIReasoningEffortFromBody(t *testing.T) { } } -func TestGetOpenAIRequestBodyMap_UsesContextCache(t *testing.T) { - gin.SetMode(gin.TestMode) - rec := httptest.NewRecorder() - c, _ := gin.CreateTestContext(rec) - - cached := map[string]any{"model": "cached-model", "stream": true} - CacheOpenAIParsedRequestBody(c, []byte(`{invalid-json`), cached) - - got, err := getOpenAIRequestBodyMap(c, []byte(`{invalid-json`)) - require.NoError(t, err) - require.Equal(t, cached, got) -} - -func TestGetOpenAIRequestBodyMap_ParseErrorWithoutCache(t *testing.T) { +func TestGetOpenAIRequestBodyMap_ParseError(t *testing.T) { _, err := getOpenAIRequestBodyMap(nil, []byte(`{invalid-json`)) require.Error(t, err) require.Contains(t, err.Error(), "parse request") } -func TestGetOpenAIRequestBodyMap_WriteBackContextCache(t *testing.T) { +func TestGetOpenAIRequestBodyMap_DoesNotWriteContextCache(t *testing.T) { gin.SetMode(gin.TestMode) rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) @@ -133,20 +120,7 @@ func TestGetOpenAIRequestBodyMap_WriteBackContextCache(t *testing.T) { got, err := getOpenAIRequestBodyMap(c, []byte(`{"model":"gpt-5","stream":true}`)) require.NoError(t, err) require.Equal(t, "gpt-5", got["model"]) - - require.Equal(t, got, CachedOpenAIParsedRequestBody(c)) -} - -func TestGetOpenAIRequestBodyMap_IgnoresCacheForDifferentBody(t *testing.T) { - gin.SetMode(gin.TestMode) - rec := httptest.NewRecorder() - c, _ := gin.CreateTestContext(rec) - - CacheOpenAIParsedRequestBody(c, []byte(`{"model":"cached-model"}`), map[string]any{"model": "cached-model"}) - - got, err := getOpenAIRequestBodyMap(c, []byte(`{"model":"forward-model"}`)) - require.NoError(t, err) - require.Equal(t, "forward-model", got["model"]) + require.Empty(t, c.Keys) } func TestSanitizeEmptyBase64InputImagesInOpenAIRequestBodyMap(t *testing.T) { diff --git a/backend/internal/service/openai_tool_continuation.go b/backend/internal/service/openai_tool_continuation.go index 7d503f5a457b..36064549cd85 100644 --- a/backend/internal/service/openai_tool_continuation.go +++ b/backend/internal/service/openai_tool_continuation.go @@ -1,6 +1,10 @@ package service -import "strings" +import ( + "strings" + + "github.com/tidwall/gjson" +) // ToolContinuationSignals 聚合工具续链相关信号，避免重复遍历 input。 type ToolContinuationSignals struct { @@ -150,6 +154,63 @@ func AnalyzeToolContinuationSignals(reqBody map[string]any) ToolContinuationSign return signals } +// ValidateFunctionCallOutputContextBytes 基于 raw JSON 校验工具输出续链，避免 handler 预校验阶段全量解码大 input。 +func ValidateFunctionCallOutputContextBytes(body []byte) FunctionCallOutputValidation { + result := FunctionCallOutputValidation{} + input := gjson.GetBytes(body, "input") + if !input.IsArray() { + return result + } + + var callIDs map[string]struct{} + var referenceIDs map[string]struct{} + input.ForEach(func(_, item gjson.Result) bool { + if !item.IsObject() { + return true + } + itemType := item.Get("type").String() + switch { + case isCodexToolCallOutputItemType(itemType): + result.HasFunctionCallOutput = true + callID := strings.TrimSpace(item.Get("call_id").String()) + if callID == "" { + result.HasFunctionCallOutputMissingCallID = true + return true + } + if callIDs == nil { + callIDs = make(map[string]struct{}) + } + callIDs[callID] = struct{}{} + case isCodexToolCallContextItemType(itemType): + if strings.TrimSpace(item.Get("call_id").String()) != "" { + result.HasToolCallContext = true + } + case itemType == "item_reference": + idValue := strings.TrimSpace(item.Get("id").String()) + if idValue == "" { + return true + } + if referenceIDs == nil { + referenceIDs = make(map[string]struct{}) + } + referenceIDs[idValue] = struct{}{} + } + return !(result.HasFunctionCallOutput && result.HasToolCallContext) + }) + if !result.HasFunctionCallOutput || result.HasToolCallContext || len(callIDs) == 0 || len(referenceIDs) == 0 { + return result + } + allReferenced := true + for callID := range callIDs { + if _, ok := referenceIDs[callID]; !ok { + allReferenced = false + break + } + } + result.HasItemReferenceForAllCallIDs = allReferenced + return result +} + // ValidateFunctionCallOutputContext 为 handler 提供低开销校验结果： // 1) 无工具输出直接返回 // 2) 若已存在工具调用上下文则提前返回 diff --git a/backend/internal/service/openai_tool_continuation_test.go b/backend/internal/service/openai_tool_continuation_test.go index 0e0552f654b5..4610652b6c2c 100644 --- a/backend/internal/service/openai_tool_continuation_test.go +++ b/backend/internal/service/openai_tool_continuation_test.go @@ -1,6 +1,7 @@ package service import ( + "encoding/json" "testing" "github.com/stretchr/testify/require" @@ -118,3 +119,68 @@ func TestHasItemReferenceForCallIDs(t *testing.T) { require.True(t, HasItemReferenceForCallIDs(req, []string{"call_1", "call_2"})) require.False(t, HasItemReferenceForCallIDs(req, []string{"call_1", "call_3"})) } + +func TestValidateFunctionCallOutputContextBytesMatchesMapValidation(t *testing.T) { + // handler 预校验走 raw JSON 扫描，语义必须与 service 内部 map 校验保持一致。 + cases := []struct { + name string + body map[string]any + }{ + { + name: "no_input", + body: map[string]any{"model": "gpt-5.4"}, + }, + { + name: "missing_call_id", + body: map[string]any{"input": []any{map[string]any{"type": "function_call_output"}}}, + }, + { + name: "call_id_without_reference", + body: map[string]any{"input": []any{map[string]any{"type": "function_call_output", "call_id": "call_1"}}}, + }, + { + name: "matching_reference", + body: map[string]any{"input": []any{ + map[string]any{"type": "function_call_output", "call_id": "call_1"}, + map[string]any{"type": "item_reference", "id": "call_1"}, + }}, + }, + { + name: "partial_reference", + body: map[string]any{"input": []any{ + map[string]any{"type": "function_call_output", "call_id": "call_1"}, + map[string]any{"type": "tool_search_output", "call_id": "call_2"}, + map[string]any{"type": "item_reference", "id": "call_1"}, + }}, + }, + { + name: "tool_context", + body: map[string]any{"input": []any{ + map[string]any{"type": "function_call_output", "call_id": "call_1"}, + map[string]any{"type": "function_call", "call_id": "call_1"}, + }}, + }, + { + name: "all_codex_tool_outputs", + body: map[string]any{"input": []any{ + map[string]any{"type": "function_call_output", "call_id": "call_function"}, + map[string]any{"type": "tool_search_output", "call_id": "call_search"}, + map[string]any{"type": "custom_tool_call_output", "call_id": "call_custom"}, + map[string]any{"type": "mcp_tool_call_output", "call_id": "call_mcp"}, + map[string]any{"type": "item_reference", "id": "call_function"}, + map[string]any{"type": "item_reference", "id": "call_search"}, + map[string]any{"type": "item_reference", "id": "call_custom"}, + map[string]any{"type": "item_reference", "id": "call_mcp"}, + }}, + }, + } + + for _, tt := range cases { + t.Run(tt.name, func(t *testing.T) { + bodyBytes, err := json.Marshal(tt.body) + require.NoError(t, err) + + require.Equal(t, ValidateFunctionCallOutputContext(tt.body), ValidateFunctionCallOutputContextBytes(bodyBytes)) + }) + } +} From 09af6ebd405d42f39f0b9fcf42924e59e903ba40 Mon Sep 17 00:00:00 2001 From: name <136912576+is7Qin@users.noreply.github.com> Date: Sat, 30 May 2026 16:58:58 +0800 Subject: [PATCH 22/36] refactor(gateway): avoid extra OpenAI WS body map copy --- backend/internal/service/openai_gateway_service.go | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/backend/internal/service/openai_gateway_service.go b/backend/internal/service/openai_gateway_service.go index 62182655c722..2df2206df5c6 100644 --- a/backend/internal/service/openai_gateway_service.go +++ b/backend/internal/service/openai_gateway_service.go @@ -2793,13 +2793,8 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco // 命中 WS 时仅走 WebSocket Mode；不再自动回退 HTTP。 if wsDecision.Transport == OpenAIUpstreamTransportResponsesWebsocketV2 { + // WS 分支不会再回落 HTTP；重连恢复可直接更新 reqBody，避免额外保留一份完整顶层 map。 wsReqBody := reqBody - if len(reqBody) > 0 { - wsReqBody = make(map[string]any, len(reqBody)) - for k, v := range reqBody { - wsReqBody[k] = v - } - } _, hasPreviousResponseID := wsReqBody["previous_response_id"] logOpenAIWSModeDebug( "forward_start account_id=%d account_type=%s model=%s stream=%v has_previous_response_id=%v", From 34de99ee0e3ae69c03774ae5e354c2e6494ca67f Mon Sep 17 00:00:00 2001 From: name <136912576+is7Qin@users.noreply.github.com> Date: Sat, 30 May 2026 19:51:03 +0800 Subject: [PATCH 23/36] refactor(gateway): cap upstream error body reads --- .../service/antigravity_gateway_service.go | 30 ++- .../antigravity_gateway_service_test.go | 9 + .../gateway_forward_as_chat_completions.go | 2 +- .../service/gateway_forward_as_responses.go | 2 +- backend/internal/service/gateway_service.go | 62 +++-- .../service/gateway_service_benchmark_test.go | 214 +++++++++++++++++- .../gemini_chat_completions_compat_service.go | 4 +- .../service/gemini_messages_compat_service.go | 24 +- backend/internal/service/openai_embeddings.go | 2 +- .../openai_gateway_chat_completions.go | 2 +- .../openai_gateway_chat_completions_raw.go | 2 +- .../service/openai_gateway_messages.go | 2 +- .../openai_gateway_responses_chat_fallback.go | 2 +- .../service/openai_gateway_service.go | 34 ++- backend/internal/service/openai_images.go | 23 +- .../service/openai_images_responses.go | 2 +- .../internal/service/openai_images_test.go | 31 ++- .../service/openai_tool_continuation.go | 6 +- 18 files changed, 392 insertions(+), 61 deletions(-) diff --git a/backend/internal/service/antigravity_gateway_service.go b/backend/internal/service/antigravity_gateway_service.go index 2b849bdd8de0..de62cfab3a70 100644 --- a/backend/internal/service/antigravity_gateway_service.go +++ b/backend/internal/service/antigravity_gateway_service.go @@ -662,7 +662,7 @@ urlFallbackLoop: // 统一处理错误响应 if resp.StatusCode >= 400 { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() if overagesInjected && shouldMarkCreditsExhausted(resp, respBody, nil) { @@ -875,6 +875,22 @@ type AntigravityGatewayService struct { internal500Cache Internal500CounterCache // INTERNAL 500 渐进惩罚计数器 } +func (s *AntigravityGatewayService) upstreamErrorBodyReadLimit() int64 { + limit := gatewayUpstreamErrorBodyReadLimit + if s != nil && s.settingService != nil && s.settingService.cfg != nil && s.settingService.cfg.Gateway.LogUpstreamErrorBody && s.settingService.cfg.Gateway.LogUpstreamErrorBodyMaxBytes > int(limit) { + limit = int64(s.settingService.cfg.Gateway.LogUpstreamErrorBodyMaxBytes) + } + return limit +} + +func (s *AntigravityGatewayService) readUpstreamErrorBody(resp *http.Response) []byte { + if resp == nil || resp.Body == nil { + return nil + } + body, _ := io.ReadAll(io.LimitReader(resp.Body, s.upstreamErrorBodyReadLimit())) + return body +} + func NewAntigravityGatewayService( accountRepo AccountRepository, cache GatewayCache, @@ -1090,7 +1106,7 @@ func (s *AntigravityGatewayService) TestConnection(ctx context.Context, account } defer func() { _ = result.resp.Body.Close() }() - respBody, err := io.ReadAll(io.LimitReader(result.resp.Body, 2<<20)) + respBody, err := io.ReadAll(io.LimitReader(result.resp.Body, s.upstreamErrorBodyReadLimit())) if err != nil { return nil, fmt.Errorf("读取响应失败: %w", err) } @@ -1427,7 +1443,7 @@ func (s *AntigravityGatewayService) Forward(ctx context.Context, c *gin.Context, defer func() { _ = resp.Body.Close() }() if resp.StatusCode >= 400 { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody := s.readUpstreamErrorBody(resp) // 优先检测 thinking block 的 signature 相关错误（400）并重试一次： // Antigravity /v1internal 链路在部分场景会对 thought/thinking signature 做严格校验， @@ -1622,7 +1638,7 @@ func (s *AntigravityGatewayService) Forward(ctx context.Context, c *gin.Context, resp = retryResp respBody = nil } else { - retryBody, _ := io.ReadAll(io.LimitReader(retryResp.Body, 2<<20)) + retryBody := s.readUpstreamErrorBody(retryResp) _ = retryResp.Body.Close() respBody = retryBody resp = &http.Response{ @@ -2189,7 +2205,7 @@ func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Co // 处理错误响应 if resp.StatusCode >= 400 { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody := s.readUpstreamErrorBody(resp) contentType := resp.Header.Get("Content-Type") // 尽早关闭原始响应体，释放连接；后续逻辑仍可能需要读取 body，因此用内存副本重新包装。 _ = resp.Body.Close() @@ -2270,7 +2286,7 @@ func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Co if retryResp.StatusCode < 400 { resp = retryResp } else { - retryRespBody, _ := io.ReadAll(io.LimitReader(retryResp.Body, 2<<20)) + retryRespBody := s.readUpstreamErrorBody(retryResp) _ = retryResp.Body.Close() retryOpsBody := retryRespBody if retryUnwrapped, unwrapErr := s.unwrapV1InternalResponse(retryRespBody); unwrapErr == nil && len(retryUnwrapped) > 0 { @@ -4252,7 +4268,7 @@ func (s *AntigravityGatewayService) ForwardUpstream(ctx context.Context, c *gin. // 处理错误响应 if resp.StatusCode >= 400 { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody := s.readUpstreamErrorBody(resp) // 429 错误时标记账号限流 if resp.StatusCode == http.StatusTooManyRequests { diff --git a/backend/internal/service/antigravity_gateway_service_test.go b/backend/internal/service/antigravity_gateway_service_test.go index 2212437467cf..903c8a1a6afc 100644 --- a/backend/internal/service/antigravity_gateway_service_test.go +++ b/backend/internal/service/antigravity_gateway_service_test.go @@ -42,6 +42,15 @@ func newAntigravityTestService(cfg *config.Config) *AntigravityGatewayService { } } +func TestAntigravityUpstreamErrorBodyReadLimit_RespectsDiagnosticLimit(t *testing.T) { + svc := newAntigravityTestService(&config.Config{Gateway: config.GatewayConfig{ + LogUpstreamErrorBody: true, + LogUpstreamErrorBodyMaxBytes: int(gatewayUpstreamErrorBodyReadLimit) + 1024, + }}) + + require.Equal(t, int64(svc.settingService.cfg.Gateway.LogUpstreamErrorBodyMaxBytes), svc.upstreamErrorBodyReadLimit()) +} + func TestStripSignatureSensitiveBlocksFromClaudeRequest(t *testing.T) { req := &antigravity.ClaudeRequest{ Model: "claude-sonnet-4-5", diff --git a/backend/internal/service/gateway_forward_as_chat_completions.go b/backend/internal/service/gateway_forward_as_chat_completions.go index 729483e3be95..1df450d6b4a7 100644 --- a/backend/internal/service/gateway_forward_as_chat_completions.go +++ b/backend/internal/service/gateway_forward_as_chat_completions.go @@ -148,7 +148,7 @@ func (s *GatewayService) ForwardAsChatCompletions( // 12. Handle error response with failover if resp.StatusCode >= 400 { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody, _ := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() resp.Body = io.NopCloser(bytes.NewReader(respBody)) diff --git a/backend/internal/service/gateway_forward_as_responses.go b/backend/internal/service/gateway_forward_as_responses.go index 3baea018bd32..22951b88b1e7 100644 --- a/backend/internal/service/gateway_forward_as_responses.go +++ b/backend/internal/service/gateway_forward_as_responses.go @@ -145,7 +145,7 @@ func (s *GatewayService) ForwardAsResponses( // 12. Handle error response with failover if resp.StatusCode >= 400 { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody, _ := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() resp.Body = io.NopCloser(bytes.NewReader(respBody)) diff --git a/backend/internal/service/gateway_service.go b/backend/internal/service/gateway_service.go index 78a1f4bdabbd..812780dc4b37 100644 --- a/backend/internal/service/gateway_service.go +++ b/backend/internal/service/gateway_service.go @@ -23,6 +23,7 @@ import ( "sync/atomic" "syscall" "time" + "unsafe" "github.com/Wei-Shaw/sub2api/internal/config" "github.com/Wei-Shaw/sub2api/internal/pkg/claude" @@ -56,6 +57,8 @@ const ( defaultModelsListCacheTTL = 15 * time.Second postUsageBillingTimeout = 15 * time.Second debugGatewayBodyEnv = "SUB2API_DEBUG_GATEWAY_BODY" + // 上游错误体只需要提取错误 JSON/日志摘要，默认 512KiB 避免错误风暴叠加大请求体。 + gatewayUpstreamErrorBodyReadLimit int64 = 512 << 10 ) const ( @@ -833,8 +836,16 @@ func (s *GatewayService) extractCacheableContent(parsed *ParsedRequest) string { return systemText } +func parseRawJSONView(raw []byte) gjson.Result { + if len(raw) == 0 { + return gjson.Result{} + } + // 这里只做同步只读解析，避免 gjson.ParseBytes 为大 messages/contents 复制整段 raw。 + return gjson.Parse(*(*string)(unsafe.Pointer(&raw))) +} + func extractTextFromSystemRaw(raw []byte) string { - system := gjson.ParseBytes(raw) + system := parseRawJSONView(raw) switch system.Type { case gjson.String: return system.String() @@ -880,7 +891,7 @@ func appendMessageTextsFromRaw(builder *strings.Builder, raw []byte) { if builder == nil || len(raw) == 0 { return } - messages := gjson.ParseBytes(raw) + messages := parseRawJSONView(raw) if !messages.IsArray() { return } @@ -903,7 +914,7 @@ func appendMessageTextsFromRaw(builder *strings.Builder, raw []byte) { } func extractCacheableTextFromSystemRaw(raw []byte) string { - system := gjson.ParseBytes(raw) + system := parseRawJSONView(raw) if !system.IsArray() { return "" } @@ -920,7 +931,7 @@ func extractCacheableTextFromSystemRaw(raw []byte) string { } func extractCacheableTextFromMessagesRaw(raw []byte) string { - messages := gjson.ParseBytes(raw) + messages := parseRawJSONView(raw) if !messages.IsArray() { return "" } @@ -4676,7 +4687,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A // 优先检测thinking block签名错误（400）并重试一次 if resp.StatusCode == 400 { - respBody, readErr := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody, readErr := s.readUpstreamErrorBody(resp) if readErr == nil { _ = resp.Body.Close() @@ -4739,7 +4750,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A break } - retryRespBody, retryReadErr := io.ReadAll(io.LimitReader(retryResp.Body, 2<<20)) + retryRespBody, retryReadErr := s.readUpstreamErrorBody(retryResp) _ = retryResp.Body.Close() if retryReadErr == nil && retryResp.StatusCode == 400 && s.isSignatureErrorPattern(ctx, account, retryRespBody) { appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ @@ -4889,7 +4900,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A break } - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody, _ := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ Platform: account.Platform, @@ -4936,7 +4947,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A // 处理重试耗尽的情况 if resp.StatusCode >= 400 && s.shouldRetryUpstreamError(account, resp.StatusCode) { if s.shouldFailoverUpstreamError(resp.StatusCode) { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody, _ := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() resp.Body = io.NopCloser(bytes.NewReader(respBody)) @@ -4971,7 +4982,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A // 处理可切换账号的错误 if resp.StatusCode >= 400 && s.shouldFailoverUpstreamError(resp.StatusCode) { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody, _ := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() resp.Body = io.NopCloser(bytes.NewReader(respBody)) @@ -5003,7 +5014,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A if resp.StatusCode >= 400 { // 可选：对部分 400 触发 failover（默认关闭以保持语义） if resp.StatusCode == 400 && s.cfg != nil && s.cfg.Gateway.FailoverOn400 { - respBody, readErr := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody, readErr := s.readUpstreamErrorBody(resp) if readErr != nil { // ReadAll failed, fall back to normal error handling without consuming the stream return s.handleErrorResponse(ctx, resp, c, account, reqModel) @@ -5221,7 +5232,7 @@ func (s *GatewayService) forwardAnthropicAPIKeyPassthroughWithInput( break } - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody, _ := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ Platform: account.Platform, @@ -5259,7 +5270,7 @@ func (s *GatewayService) forwardAnthropicAPIKeyPassthroughWithInput( if resp.StatusCode >= 400 && s.shouldRetryUpstreamError(account, resp.StatusCode) { if s.shouldFailoverUpstreamError(resp.StatusCode) { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody, _ := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() resp.Body = io.NopCloser(bytes.NewReader(respBody)) @@ -5293,7 +5304,7 @@ func (s *GatewayService) forwardAnthropicAPIKeyPassthroughWithInput( } if resp.StatusCode >= 400 && s.shouldFailoverUpstreamError(resp.StatusCode) { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody, _ := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() resp.Body = io.NopCloser(bytes.NewReader(respBody)) @@ -6011,7 +6022,7 @@ func (s *GatewayService) executeBedrockUpstream( break } - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody, _ := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ Platform: account.Platform, @@ -6056,7 +6067,7 @@ func (s *GatewayService) handleBedrockUpstreamErrors( // retry exhausted + failover if s.shouldRetryUpstreamError(account, resp.StatusCode) { if s.shouldFailoverUpstreamError(resp.StatusCode) { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody, _ := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() resp.Body = io.NopCloser(bytes.NewReader(respBody)) @@ -6083,7 +6094,7 @@ func (s *GatewayService) handleBedrockUpstreamErrors( // non-retryable failover if s.shouldFailoverUpstreamError(resp.StatusCode) { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody, _ := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() resp.Body = io.NopCloser(bytes.NewReader(respBody)) @@ -7226,8 +7237,19 @@ func isCountTokensUnsupported404(statusCode int, body []byte) bool { return strings.Contains(msg, "count_tokens") && strings.Contains(msg, "not found") } +func (s *GatewayService) readUpstreamErrorBody(resp *http.Response) ([]byte, error) { + if resp == nil || resp.Body == nil { + return nil, nil + } + limit := gatewayUpstreamErrorBodyReadLimit + if s != nil && s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody && s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes > int(limit) { + limit = int64(s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes) + } + return io.ReadAll(io.LimitReader(resp.Body, limit)) +} + func (s *GatewayService) handleErrorResponse(ctx context.Context, resp *http.Response, c *gin.Context, account *Account, requestedModel ...string) (*ForwardResult, error) { - body, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + body, _ := s.readUpstreamErrorBody(resp) // 调试日志：打印上游错误响应 logger.LegacyPrintf("service.gateway", "[Forward] Upstream error (non-retryable): Account=%d(%s) Status=%d RequestID=%s Body=%s", @@ -7380,7 +7402,7 @@ func (s *GatewayService) handleErrorResponse(ctx context.Context, resp *http.Res } func (s *GatewayService) handleRetryExhaustedSideEffects(ctx context.Context, resp *http.Response, account *Account) { - body, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + body, _ := s.readUpstreamErrorBody(resp) statusCode := resp.StatusCode // OAuth/Setup Token 账号的 403：标记账号异常 @@ -7394,7 +7416,7 @@ func (s *GatewayService) handleRetryExhaustedSideEffects(ctx context.Context, re } func (s *GatewayService) handleFailoverSideEffects(ctx context.Context, resp *http.Response, account *Account, requestedModel ...string) { - body, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + body, _ := s.readUpstreamErrorBody(resp) if len(requestedModel) > 0 { s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, body, requestedModel[0]) return @@ -7407,7 +7429,7 @@ func (s *GatewayService) handleFailoverSideEffects(ctx context.Context, resp *ht // API Key 未配置错误码：仅返回错误，不标记账号 func (s *GatewayService) handleRetryExhaustedError(ctx context.Context, resp *http.Response, c *gin.Context, account *Account) (*ForwardResult, error) { // Capture upstream error body before side-effects consume the stream. - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody, _ := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() resp.Body = io.NopCloser(bytes.NewReader(respBody)) diff --git a/backend/internal/service/gateway_service_benchmark_test.go b/backend/internal/service/gateway_service_benchmark_test.go index 8b30cb249e3d..b2df143591f3 100644 --- a/backend/internal/service/gateway_service_benchmark_test.go +++ b/backend/internal/service/gateway_service_benchmark_test.go @@ -4,9 +4,14 @@ import ( "strconv" "strings" "testing" + + "github.com/Wei-Shaw/sub2api/internal/domain" ) -var benchmarkStringSink string +var ( + benchmarkStringSink string + benchmarkIntSink int +) // BenchmarkGenerateSessionHash_Metadata 关注 JSON 解析与正则匹配开销。 func BenchmarkGenerateSessionHash_Metadata(b *testing.B) { @@ -23,6 +28,121 @@ func BenchmarkGenerateSessionHash_Metadata(b *testing.B) { } } +func BenchmarkParseGatewayRequest_LargeAnthropicMessages(b *testing.B) { + for _, size := range benchmarkBodySizes() { + b.Run(size.name, func(b *testing.B) { + body := buildLargeAnthropicMessagesBody(size.bytes, false) + + b.SetBytes(int64(len(body))) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), domain.PlatformAnthropic) + if err != nil { + b.Fatalf("解析 Anthropic 请求失败: %v", err) + } + benchmarkIntSink = len(parsed.MessagesRaw()) + } + }) + } +} + +func BenchmarkParseGatewayRequest_LargeGeminiContents(b *testing.B) { + for _, size := range benchmarkBodySizes() { + b.Run(size.name, func(b *testing.B) { + body := buildLargeGeminiContentsBody(size.bytes) + + b.SetBytes(int64(len(body))) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), domain.PlatformGemini) + if err != nil { + b.Fatalf("解析 Gemini 请求失败: %v", err) + } + benchmarkIntSink = len(parsed.MessagesRaw()) + } + }) + } +} + +func BenchmarkGenerateSessionHash_LargeAnthropicMessages(b *testing.B) { + svc := &GatewayService{} + for _, size := range benchmarkBodySizes() { + b.Run(size.name, func(b *testing.B) { + body := buildLargeAnthropicMessagesBody(size.bytes, true) + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), domain.PlatformAnthropic) + if err != nil { + b.Fatalf("解析请求失败: %v", err) + } + + b.SetBytes(int64(len(body))) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + benchmarkStringSink = svc.GenerateSessionHash(parsed) + } + }) + } +} + +func BenchmarkOpenAIResponses_LargeInputMeta(b *testing.B) { + for _, size := range benchmarkBodySizes() { + b.Run(size.name, func(b *testing.B) { + body := buildLargeOpenAIResponsesBody(size.bytes) + + b.SetBytes(int64(len(body))) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + model, stream, promptCacheKey := extractOpenAIRequestMetaFromBody(body) + benchmarkStringSink = model + promptCacheKey + if stream { + benchmarkIntSink++ + } + } + }) + } +} + +func BenchmarkOpenAIResponses_LargeInputDecodeMap(b *testing.B) { + for _, size := range benchmarkBodySizes() { + b.Run(size.name, func(b *testing.B) { + body := buildLargeOpenAIResponsesBody(size.bytes) + + b.SetBytes(int64(len(body))) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + reqBody, err := getOpenAIRequestBodyMap(nil, body) + if err != nil { + b.Fatalf("解析 OpenAI 请求失败: %v", err) + } + benchmarkIntSink = len(reqBody) + } + }) + } +} + +func BenchmarkOpenAIResponses_LargeInputFunctionCallValidation(b *testing.B) { + for _, size := range benchmarkBodySizes() { + b.Run(size.name, func(b *testing.B) { + body := buildLargeOpenAIResponsesToolContinuationBody(size.bytes) + + b.SetBytes(int64(len(body))) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + validation := ValidateFunctionCallOutputContextBytes(body) + if !validation.HasFunctionCallOutput || !validation.HasItemReferenceForAllCallIDs { + b.Fatalf("工具续链校验结果异常: %+v", validation) + } + benchmarkIntSink++ + } + }) + } +} + // BenchmarkExtractCacheableContent_System 关注字符串拼接路径的性能。 func BenchmarkExtractCacheableContent_System(b *testing.B) { svc := &GatewayService{} @@ -34,6 +154,21 @@ func BenchmarkExtractCacheableContent_System(b *testing.B) { } } +func benchmarkBodySizes() []struct { + name string + bytes int +} { + return []struct { + name string + bytes int + }{ + {name: "4MB", bytes: 4 << 20}, + {name: "8MB", bytes: 8 << 20}, + {name: "16MB", bytes: 16 << 20}, + {name: "32MB", bytes: 32 << 20}, + } +} + func buildSystemCacheableRequest(parts int) *ParsedRequest { var builder strings.Builder builder.WriteString(`{"system":[`) @@ -52,3 +187,80 @@ func buildSystemCacheableRequest(parts int) *ParsedRequest { } return parsed } + +func buildLargeAnthropicMessagesBody(targetBytes int, includeCacheControl bool) []byte { + var builder strings.Builder + builder.Grow(targetBytes + 1024) + builder.WriteString(`{"model":"claude-sonnet-4-5","stream":true,"system":[{"type":"text","text":"system seed"}],"messages":[`) + for i := 0; builder.Len() < targetBytes; i++ { + if i > 0 { + builder.WriteByte(',') + } + builder.WriteString(`{"role":"user","content":[{"type":"text","text":"`) + builder.WriteString(strings.Repeat("anthropic payload ", 64)) + builder.WriteString(strconv.Itoa(i)) + builder.WriteByte('"') + if includeCacheControl && i%32 == 0 { + builder.WriteString(`,"cache_control":{"type":"ephemeral"}`) + } + builder.WriteString(`}]}`) + } + builder.WriteString(`]}`) + return []byte(builder.String()) +} + +func buildLargeGeminiContentsBody(targetBytes int) []byte { + var builder strings.Builder + builder.Grow(targetBytes + 1024) + builder.WriteString(`{"model":"gemini-2.5-pro","systemInstruction":{"parts":[{"text":"system seed"}]},"contents":[`) + for i := 0; builder.Len() < targetBytes; i++ { + if i > 0 { + builder.WriteByte(',') + } + builder.WriteString(`{"role":"user","parts":[{"text":"`) + builder.WriteString(strings.Repeat("gemini payload ", 64)) + builder.WriteString(strconv.Itoa(i)) + builder.WriteString(`"}]}`) + } + builder.WriteString(`]}`) + return []byte(builder.String()) +} + +func buildLargeOpenAIResponsesBody(targetBytes int) []byte { + var builder strings.Builder + builder.Grow(targetBytes + 1024) + builder.WriteString(`{"model":"gpt-5.4","stream":true,"prompt_cache_key":"session-benchmark","input":[`) + for i := 0; builder.Len() < targetBytes; i++ { + if i > 0 { + builder.WriteByte(',') + } + builder.WriteString(`{"type":"message","role":"user","content":[{"type":"input_text","text":"`) + builder.WriteString(strings.Repeat("openai responses payload ", 48)) + builder.WriteString(strconv.Itoa(i)) + builder.WriteString(`"}]}`) + } + builder.WriteString(`],"tools":[{"type":"function","name":"lookup","parameters":{"type":"object","properties":{"query":{"type":"string"}}}}]}`) + return []byte(builder.String()) +} + +func buildLargeOpenAIResponsesToolContinuationBody(targetBytes int) []byte { + var builder strings.Builder + builder.Grow(targetBytes + 1024) + builder.WriteString(`{"model":"gpt-5.4","stream":true,"previous_response_id":"resp_benchmark","input":[`) + for i := 0; builder.Len() < targetBytes; i++ { + if i > 0 { + builder.WriteByte(',') + } + callID := "call_" + strconv.Itoa(i) + builder.WriteString(`{"type":"item_reference","id":"`) + builder.WriteString(callID) + builder.WriteString(`"},{"type":"function_call_output","call_id":"`) + builder.WriteString(callID) + builder.WriteString(`","output":"`) + builder.WriteString(strings.Repeat("tool output payload ", 48)) + builder.WriteString(strconv.Itoa(i)) + builder.WriteString(`"}`) + } + builder.WriteString(`]}`) + return []byte(builder.String()) +} diff --git a/backend/internal/service/gemini_chat_completions_compat_service.go b/backend/internal/service/gemini_chat_completions_compat_service.go index dcc3213b1b69..ffea15950335 100644 --- a/backend/internal/service/gemini_chat_completions_compat_service.go +++ b/backend/internal/service/gemini_chat_completions_compat_service.go @@ -151,7 +151,7 @@ func (s *GeminiMessagesCompatService) forwardClaudeBodyAsChatCompletions( } if resp.StatusCode >= 400 && s.shouldRetryGeminiUpstreamError(account, resp.StatusCode) { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() if resp.StatusCode == http.StatusForbidden && isGeminiInsufficientScope(resp.Header, respBody) { resp = &http.Response{ @@ -207,7 +207,7 @@ func (s *GeminiMessagesCompatService) forwardClaudeBodyAsChatCompletions( reasoningEffort := extractCCReasoningEffortFromBody(originalChatBody) if resp.StatusCode >= 400 { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody := s.readUpstreamErrorBody(resp) s.handleGeminiUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody) evBody := unwrapIfNeeded(account.Type == AccountTypeOAuth, respBody) diff --git a/backend/internal/service/gemini_messages_compat_service.go b/backend/internal/service/gemini_messages_compat_service.go index 64f19b2ec758..86073d9ca549 100644 --- a/backend/internal/service/gemini_messages_compat_service.go +++ b/backend/internal/service/gemini_messages_compat_service.go @@ -56,6 +56,18 @@ type GeminiMessagesCompatService struct { responseHeaderFilter *responseheaders.CompiledHeaderFilter } +func (s *GeminiMessagesCompatService) readUpstreamErrorBody(resp *http.Response) []byte { + if resp == nil || resp.Body == nil { + return nil + } + limit := gatewayUpstreamErrorBodyReadLimit + if s != nil && s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody && s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes > int(limit) { + limit = int64(s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes) + } + body, _ := io.ReadAll(io.LimitReader(resp.Body, limit)) + return body +} + func NewGeminiMessagesCompatService( accountRepo AccountRepository, groupRepo GroupRepository, @@ -789,7 +801,7 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex // Special-case: signature/thought_signature validation errors are not transient, but may be fixed by // downgrading Claude thinking/tool history to plain text (conservative two-stage retry). if resp.StatusCode == http.StatusBadRequest && signatureRetryStage < 2 { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() if isGeminiSignatureRelatedError(respBody) { @@ -860,7 +872,7 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex } if resp.StatusCode >= 400 && s.shouldRetryGeminiUpstreamError(account, resp.StatusCode) { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() // Don't treat insufficient-scope as transient. if resp.StatusCode == 403 && isGeminiInsufficientScope(resp.Header, respBody) { @@ -919,7 +931,7 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex defer func() { _ = resp.Body.Close() }() if resp.StatusCode >= 400 { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody := s.readUpstreamErrorBody(resp) // 统一错误策略：自定义错误码 + 临时不可调度 if s.rateLimitService != nil { switch s.rateLimitService.CheckErrorPolicy(ctx, account, resp.StatusCode, respBody) { @@ -1329,7 +1341,7 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin. } if resp.StatusCode >= 400 && s.shouldRetryGeminiUpstreamError(account, resp.StatusCode) { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() // Don't treat insufficient-scope as transient. if resp.StatusCode == 403 && isGeminiInsufficientScope(resp.Header, respBody) { @@ -1410,7 +1422,7 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin. isOAuth := account.Type == AccountTypeOAuth if resp.StatusCode >= 400 { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody := s.readUpstreamErrorBody(resp) // 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. @@ -1619,7 +1631,7 @@ func (s *GeminiMessagesCompatService) checkErrorPolicyInLoop( if resp.StatusCode < 400 || s.rateLimitService == nil { return false, resp } - body, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + body := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() rebuilt = &http.Response{ StatusCode: resp.StatusCode, diff --git a/backend/internal/service/openai_embeddings.go b/backend/internal/service/openai_embeddings.go index 359df3bb80a2..7c710259f3fb 100644 --- a/backend/internal/service/openai_embeddings.go +++ b/backend/internal/service/openai_embeddings.go @@ -104,7 +104,7 @@ func (s *OpenAIGatewayService) ForwardEmbeddings( defer func() { _ = resp.Body.Close() }() if resp.StatusCode >= 400 { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() resp.Body = io.NopCloser(bytes.NewReader(respBody)) diff --git a/backend/internal/service/openai_gateway_chat_completions.go b/backend/internal/service/openai_gateway_chat_completions.go index 807ff43afc8a..6e91d85c83a3 100644 --- a/backend/internal/service/openai_gateway_chat_completions.go +++ b/backend/internal/service/openai_gateway_chat_completions.go @@ -241,7 +241,7 @@ func (s *OpenAIGatewayService) ForwardAsChatCompletions( // 8. Handle error response with failover if resp.StatusCode >= 400 { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() resp.Body = io.NopCloser(bytes.NewReader(respBody)) diff --git a/backend/internal/service/openai_gateway_chat_completions_raw.go b/backend/internal/service/openai_gateway_chat_completions_raw.go index e351fa75ba7a..3ff6fac4aa8e 100644 --- a/backend/internal/service/openai_gateway_chat_completions_raw.go +++ b/backend/internal/service/openai_gateway_chat_completions_raw.go @@ -183,7 +183,7 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions( // 7. Handle error response with failover if resp.StatusCode >= 400 { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() resp.Body = io.NopCloser(bytes.NewReader(respBody)) diff --git a/backend/internal/service/openai_gateway_messages.go b/backend/internal/service/openai_gateway_messages.go index 291c217e85df..4398bd27cb38 100644 --- a/backend/internal/service/openai_gateway_messages.go +++ b/backend/internal/service/openai_gateway_messages.go @@ -300,7 +300,7 @@ func (s *OpenAIGatewayService) ForwardAsAnthropic( // 8. Handle error response with failover if resp.StatusCode >= 400 { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() resp.Body = io.NopCloser(bytes.NewReader(respBody)) diff --git a/backend/internal/service/openai_gateway_responses_chat_fallback.go b/backend/internal/service/openai_gateway_responses_chat_fallback.go index cfab389a3151..205b27f7766f 100644 --- a/backend/internal/service/openai_gateway_responses_chat_fallback.go +++ b/backend/internal/service/openai_gateway_responses_chat_fallback.go @@ -163,7 +163,7 @@ func (s *OpenAIGatewayService) forwardResponsesViaRawChatCompletions( defer func() { _ = resp.Body.Close() }() if resp.StatusCode >= 400 { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() resp.Body = io.NopCloser(bytes.NewReader(respBody)) diff --git a/backend/internal/service/openai_gateway_service.go b/backend/internal/service/openai_gateway_service.go index 2df2206df5c6..d1460e5a0887 100644 --- a/backend/internal/service/openai_gateway_service.go +++ b/backend/internal/service/openai_gateway_service.go @@ -49,6 +49,8 @@ const ( // OpenAI WS Mode 失败后的重连次数上限（不含首次尝试）。 // 与 Codex 客户端保持一致：失败后最多重连 5 次。 openAIWSReconnectRetryLimit = 5 + // 上游错误体只需要提取错误 JSON/日志摘要，默认 512KiB 避免错误风暴叠加大请求体。 + openAIUpstreamErrorBodyReadLimit int64 = 512 << 10 // OpenAI WS Mode 重连退避默认值（可由配置覆盖）。 openAIWSRetryBackoffInitialDefault = 120 * time.Millisecond openAIWSRetryBackoffMaxDefault = 2 * time.Second @@ -2296,8 +2298,28 @@ func marshalOpenAIUpstreamJSON(v any) ([]byte, error) { return out, nil } +func openAIUpstreamErrorBodyReadLimitForConfig(cfg *config.Config) int64 { + limit := openAIUpstreamErrorBodyReadLimit + if cfg != nil && cfg.Gateway.LogUpstreamErrorBody && cfg.Gateway.LogUpstreamErrorBodyMaxBytes > int(limit) { + limit = int64(cfg.Gateway.LogUpstreamErrorBodyMaxBytes) + } + return limit +} + +func (s *OpenAIGatewayService) readUpstreamErrorBody(resp *http.Response) []byte { + if resp == nil || resp.Body == nil { + return nil + } + cfg := (*config.Config)(nil) + if s != nil { + cfg = s.cfg + } + body, _ := io.ReadAll(io.LimitReader(resp.Body, openAIUpstreamErrorBodyReadLimitForConfig(cfg))) + return body +} + func (s *OpenAIGatewayService) handleFailoverSideEffects(ctx context.Context, resp *http.Response, account *Account, requestedModel ...string) { - body, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + body := s.readUpstreamErrorBody(resp) if len(requestedModel) > 0 { s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, body, requestedModel[0]) return @@ -3045,7 +3067,7 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco // Handle error response if resp.StatusCode >= 400 { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() resp.Body = io.NopCloser(bytes.NewReader(respBody)) @@ -3563,7 +3585,7 @@ func (s *OpenAIGatewayService) handleFailoverErrorResponsePassthrough( account *Account, requestBody []byte, ) error { - body, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + body := s.readUpstreamErrorBody(resp) upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(body)) upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) @@ -3605,7 +3627,7 @@ func (s *OpenAIGatewayService) handleErrorResponsePassthrough( account *Account, requestBody []byte, ) error { - body, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + body := s.readUpstreamErrorBody(resp) upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(body)) upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) @@ -4298,7 +4320,7 @@ func (s *OpenAIGatewayService) handleErrorResponse( requestBody []byte, requestedModel ...string, ) (*OpenAIForwardResult, error) { - body, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + body := s.readUpstreamErrorBody(resp) upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(body)) upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) @@ -4459,7 +4481,7 @@ func (s *OpenAIGatewayService) handleCompatErrorResponse( writeError compatErrorWriter, requestedModel ...string, ) (*OpenAIForwardResult, error) { - body, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + body := s.readUpstreamErrorBody(resp) upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(body)) if upstreamMsg == "" { diff --git a/backend/internal/service/openai_images.go b/backend/internal/service/openai_images.go index 1bcd947c4495..beb34780f3bf 100644 --- a/backend/internal/service/openai_images.go +++ b/backend/internal/service/openai_images.go @@ -622,7 +622,7 @@ func (s *OpenAIGatewayService) forwardOpenAIImagesAPIKey( return nil, fmt.Errorf("upstream request failed: %s", safeErr) } if resp.StatusCode >= 400 { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() resp.Body = io.NopCloser(bytes.NewReader(respBody)) upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody)) @@ -1276,21 +1276,22 @@ func resolveOpenAIImageBytes( headers http.Header, conversationID string, pointer openAIImagePointerInfo, + errorBodyReadLimit int64, ) ([]byte, error) { if normalized := normalizeOpenAIImageBase64(pointer.B64JSON); normalized != "" { return base64.StdEncoding.DecodeString(normalized) } if downloadURL := strings.TrimSpace(pointer.DownloadURL); downloadURL != "" { - return downloadOpenAIImageBytes(ctx, client, headers, downloadURL) + return downloadOpenAIImageBytes(ctx, client, headers, downloadURL, errorBodyReadLimit) } if strings.TrimSpace(pointer.Pointer) == "" { return nil, fmt.Errorf("image asset is missing pointer, url, and base64 data") } - downloadURL, err := fetchOpenAIImageDownloadURL(ctx, client, headers, conversationID, pointer.Pointer) + downloadURL, err := fetchOpenAIImageDownloadURL(ctx, client, headers, conversationID, pointer.Pointer, errorBodyReadLimit) if err != nil { return nil, err } - return downloadOpenAIImageBytes(ctx, client, headers, downloadURL) + return downloadOpenAIImageBytes(ctx, client, headers, downloadURL, errorBodyReadLimit) } func normalizeOpenAIImageBase64(raw string) string { @@ -1395,6 +1396,7 @@ func fetchOpenAIImageDownloadURL( headers http.Header, conversationID string, pointer string, + errorBodyReadLimit int64, ) (string, error) { url := "" allowConversationRetry := false @@ -1425,7 +1427,7 @@ func fetchOpenAIImageDownloadURL( } else if resp.IsSuccessState() && strings.TrimSpace(result.DownloadURL) != "" { return strings.TrimSpace(result.DownloadURL), nil } else { - statusErr := newOpenAIImageStatusError(resp, "fetch image download url failed") + statusErr := newOpenAIImageStatusError(resp, "fetch image download url failed", errorBodyReadLimit) if !allowConversationRetry || !isOpenAIImageTransientConversationNotFoundError(statusErr) { return "", statusErr } @@ -1450,7 +1452,7 @@ func fetchOpenAIImageDownloadURL( return "", lastErr } -func downloadOpenAIImageBytes(ctx context.Context, client *req.Client, headers http.Header, downloadURL string) ([]byte, error) { +func downloadOpenAIImageBytes(ctx context.Context, client *req.Client, headers http.Header, downloadURL string, errorBodyReadLimit int64) ([]byte, error) { request := client.R(). SetContext(ctx). DisableAutoReadResponse() @@ -1478,7 +1480,7 @@ func downloadOpenAIImageBytes(ctx context.Context, client *req.Client, headers h } }() if resp.StatusCode < 200 || resp.StatusCode >= 300 { - return nil, newOpenAIImageStatusError(resp, "download image bytes failed") + return nil, newOpenAIImageStatusError(resp, "download image bytes failed", errorBodyReadLimit) } return io.ReadAll(io.LimitReader(resp.Body, openAIImageMaxDownloadBytes)) } @@ -1505,7 +1507,7 @@ func (e *openAIImageStatusError) Error() string { return "openai image backend request failed" } -func newOpenAIImageStatusError(resp *req.Response, fallback string) error { +func newOpenAIImageStatusError(resp *req.Response, fallback string, errorBodyReadLimit int64) error { if resp == nil { if strings.TrimSpace(fallback) == "" { fallback = "openai image backend request failed" @@ -1526,7 +1528,10 @@ func newOpenAIImageStatusError(resp *req.Response, fallback string) error { requestURL = resp.Request.URL.String() } if resp.Body != nil { - body, _ = io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + if errorBodyReadLimit <= 0 { + errorBodyReadLimit = openAIUpstreamErrorBodyReadLimit + } + body, _ = io.ReadAll(io.LimitReader(resp.Body, errorBodyReadLimit)) _ = resp.Body.Close() } } diff --git a/backend/internal/service/openai_images_responses.go b/backend/internal/service/openai_images_responses.go index 849ad7920c13..db9c7b1627f5 100644 --- a/backend/internal/service/openai_images_responses.go +++ b/backend/internal/service/openai_images_responses.go @@ -1172,7 +1172,7 @@ func (s *OpenAIGatewayService) forwardOpenAIImagesOAuth( return nil, fmt.Errorf("upstream request failed: %s", safeErr) } if resp.StatusCode >= 400 { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() resp.Body = io.NopCloser(bytes.NewReader(respBody)) upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody)) diff --git a/backend/internal/service/openai_images_test.go b/backend/internal/service/openai_images_test.go index a87e96c10778..c3efdc931aa8 100644 --- a/backend/internal/service/openai_images_test.go +++ b/backend/internal/service/openai_images_test.go @@ -4,6 +4,7 @@ import ( "bytes" "context" "errors" + "fmt" "io" "mime/multipart" "net/http" @@ -14,6 +15,7 @@ import ( "github.com/Wei-Shaw/sub2api/internal/config" "github.com/gin-gonic/gin" + "github.com/imroc/req/v3" "github.com/stretchr/testify/require" "github.com/tidwall/gjson" ) @@ -398,11 +400,38 @@ func TestCollectOpenAIImagePointers_RecognizesDirectAssets(t *testing.T) { func TestResolveOpenAIImageBytes_PrefersInlineBase64(t *testing.T) { data, err := resolveOpenAIImageBytes(context.Background(), nil, nil, "", openAIImagePointerInfo{ B64JSON: "data:image/png;base64,QUJD", - }) + }, openAIUpstreamErrorBodyReadLimit) require.NoError(t, err) require.Equal(t, []byte("ABC"), data) } +func TestNewOpenAIImageStatusError_UsesProvidedReadLimit(t *testing.T) { + padding := strings.Repeat("x", int(openAIUpstreamErrorBodyReadLimit)+1024) + body := fmt.Sprintf(`{"error":{"padding":"%s","message":"diagnostic-marker"}}`, padding) + resp := &req.Response{Response: &http.Response{ + StatusCode: http.StatusBadGateway, + Header: http.Header{}, + Body: io.NopCloser(strings.NewReader(body)), + }} + + err := newOpenAIImageStatusError(resp, "download image bytes failed", int64(len(body))) + require.Error(t, err) + require.Equal(t, "diagnostic-marker", err.Error()) + + var statusErr *openAIImageStatusError + require.ErrorAs(t, err, &statusErr) + require.Len(t, statusErr.ResponseBody, len(body)) +} + +func TestOpenAIUpstreamErrorBodyReadLimitForConfig_RespectsDiagnosticLimit(t *testing.T) { + cfg := &config.Config{Gateway: config.GatewayConfig{ + LogUpstreamErrorBody: true, + LogUpstreamErrorBodyMaxBytes: int(openAIUpstreamErrorBodyReadLimit) + 1024, + }} + + require.Equal(t, int64(cfg.Gateway.LogUpstreamErrorBodyMaxBytes), openAIUpstreamErrorBodyReadLimitForConfig(cfg)) +} + func TestAccountSupportsOpenAIImageCapability_OAuthSupportsNative(t *testing.T) { account := &Account{ Platform: PlatformOpenAI, diff --git a/backend/internal/service/openai_tool_continuation.go b/backend/internal/service/openai_tool_continuation.go index 36064549cd85..92bc2613bb6f 100644 --- a/backend/internal/service/openai_tool_continuation.go +++ b/backend/internal/service/openai_tool_continuation.go @@ -157,7 +157,11 @@ func AnalyzeToolContinuationSignals(reqBody map[string]any) ToolContinuationSign // ValidateFunctionCallOutputContextBytes 基于 raw JSON 校验工具输出续链，避免 handler 预校验阶段全量解码大 input。 func ValidateFunctionCallOutputContextBytes(body []byte) FunctionCallOutputValidation { result := FunctionCallOutputValidation{} - input := gjson.GetBytes(body, "input") + if len(body) == 0 { + return result + } + // handler 热路径只读扫描 input，避免 GetBytes 为大 Responses body 复制整段 JSON。 + input := parseRawJSONView(body).Get("input") if !input.IsArray() { return result } From 6a5f6b96b6d438bd785ec14bda8840ba80ff6c29 Mon Sep 17 00:00:00 2001 From: name <136912576+is7Qin@users.noreply.github.com> Date: Sun, 31 May 2026 01:01:58 +0800 Subject: [PATCH 24/36] refactor(gateway): introduce OpenAI request view Cache hot-path request scalars before full body decoding so later branches can avoid repeated map work while preserving current decode behavior. --- .../service/openai_gateway_service.go | 41 +++++++++++++++---- .../openai_gateway_service_hotpath_test.go | 20 +++++++++ 2 files changed, 53 insertions(+), 8 deletions(-) diff --git a/backend/internal/service/openai_gateway_service.go b/backend/internal/service/openai_gateway_service.go index d1460e5a0887..87ffb2cb2197 100644 --- a/backend/internal/service/openai_gateway_service.go +++ b/backend/internal/service/openai_gateway_service.go @@ -2346,7 +2346,8 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco } originalBody := body - reqModel, reqStream, promptCacheKey := extractOpenAIRequestMetaFromBody(body) + requestView := newOpenAIRequestView(body) + reqModel, reqStream, promptCacheKey := requestView.Model, requestView.Stream, requestView.PromptCacheKey originalModel := reqModel if account.Type == AccountTypeAPIKey && !openai_compat.ShouldUseResponsesAPI(account.Extra) { @@ -2396,7 +2397,7 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco return s.forwardOpenAIPassthrough(ctx, c, account, originalBody, reqModel, reasoningEffort, reqStream, startTime) } - reqBody, err := getOpenAIRequestBodyMap(c, body) + reqBody, err := requestView.Decode(c) if err != nil { return nil, err } @@ -6274,15 +6275,39 @@ func deriveOpenAIReasoningEffortFromModel(model string) string { return normalizeOpenAIReasoningEffort(parts[len(parts)-1]) } -func extractOpenAIRequestMetaFromBody(body []byte) (model string, stream bool, promptCacheKey string) { +type openAIRequestView struct { + body []byte + Model string + Stream bool + PromptCacheKey string + PreviousResponseID string + ServiceTier string + ReasoningEffort string +} + +func newOpenAIRequestView(body []byte) openAIRequestView { if len(body) == 0 { - return "", false, "" + return openAIRequestView{} + } + return openAIRequestView{ + body: body, + Model: strings.TrimSpace(gjson.GetBytes(body, "model").String()), + Stream: gjson.GetBytes(body, "stream").Bool(), + PromptCacheKey: strings.TrimSpace(gjson.GetBytes(body, "prompt_cache_key").String()), + PreviousResponseID: strings.TrimSpace(gjson.GetBytes(body, "previous_response_id").String()), + ServiceTier: strings.TrimSpace(gjson.GetBytes(body, "service_tier").String()), + ReasoningEffort: strings.TrimSpace(gjson.GetBytes(body, "reasoning.effort").String()), } +} + +// Decode 保留阶段一既有 full-map 行为；后续阶段会把调用点下沉到复杂分支。 +func (v openAIRequestView) Decode(c *gin.Context) (map[string]any, error) { + return getOpenAIRequestBodyMap(c, v.body) +} - model = strings.TrimSpace(gjson.GetBytes(body, "model").String()) - stream = gjson.GetBytes(body, "stream").Bool() - promptCacheKey = strings.TrimSpace(gjson.GetBytes(body, "prompt_cache_key").String()) - return model, stream, promptCacheKey +func extractOpenAIRequestMetaFromBody(body []byte) (model string, stream bool, promptCacheKey string) { + view := newOpenAIRequestView(body) + return view.Model, view.Stream, view.PromptCacheKey } // normalizeOpenAIPassthroughOAuthBody 将透传 OAuth 请求体收敛为旧链路关键行为： diff --git a/backend/internal/service/openai_gateway_service_hotpath_test.go b/backend/internal/service/openai_gateway_service_hotpath_test.go index af0d21c4a176..df17ed2ef023 100644 --- a/backend/internal/service/openai_gateway_service_hotpath_test.go +++ b/backend/internal/service/openai_gateway_service_hotpath_test.go @@ -9,6 +9,26 @@ import ( "github.com/stretchr/testify/require" ) +func TestOpenAIRequestView_ExtractsRawScalars(t *testing.T) { + view := newOpenAIRequestView([]byte(`{"model":" gpt-5 ","stream":true,"prompt_cache_key":" ses-1 ","previous_response_id":" resp-1 ","service_tier":" fast ","reasoning":{"effort":" medium "}}`)) + + require.Equal(t, "gpt-5", view.Model) + require.True(t, view.Stream) + require.Equal(t, "ses-1", view.PromptCacheKey) + require.Equal(t, "resp-1", view.PreviousResponseID) + require.Equal(t, "fast", view.ServiceTier) + require.Equal(t, "medium", view.ReasoningEffort) +} + +func TestOpenAIRequestView_DecodeKeepsFullMapBehavior(t *testing.T) { + view := newOpenAIRequestView([]byte(`{"model":"gpt-5","stream":true,"input":[{"type":"message","content":"hi"}]}`)) + + reqBody, err := view.Decode(nil) + require.NoError(t, err) + require.Equal(t, "gpt-5", reqBody["model"]) + require.IsType(t, []any{}, reqBody["input"]) +} + func TestExtractOpenAIRequestMetaFromBody(t *testing.T) { tests := []struct { name string From b65dde634bf7a5338f22b46ad0f9d98a197526fc Mon Sep 17 00:00:00 2001 From: wucm667 Date: Sun, 31 May 2026 08:39:37 +0800 Subject: [PATCH 25/36] =?UTF-8?q?fix(usage):=20=E4=BF=AE=E6=AD=A3=20OpenAI?= =?UTF-8?q?=205h=20=E7=94=A8=E9=87=8F=E7=99=BE=E5=88=86=E6=AF=94=E8=AF=AD?= =?UTF-8?q?=E4=B9=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../account_test_service_openai_test.go | 4 +-- .../service/account_usage_service_test.go | 29 +++++++++++++++- .../service/openai_gateway_service.go | 17 ++++++++-- ...nai_gateway_service_codex_snapshot_test.go | 34 +++++++++++++++++++ .../service/openai_gateway_service_test.go | 2 +- .../service/ratelimit_service_openai_test.go | 32 ++++++++--------- 6 files changed, 96 insertions(+), 22 deletions(-) diff --git a/backend/internal/service/account_test_service_openai_test.go b/backend/internal/service/account_test_service_openai_test.go index 910567fb25b6..970c723a9b63 100644 --- a/backend/internal/service/account_test_service_openai_test.go +++ b/backend/internal/service/account_test_service_openai_test.go @@ -132,7 +132,7 @@ func TestAccountTestService_OpenAISuccessPersistsSnapshotFromHeaders(t *testing. require.Len(t, upstream.requests, 1) require.Equal(t, HTTPUpstreamProfileOpenAI, HTTPUpstreamProfileFromContext(upstream.requests[0].Context())) require.NotEmpty(t, repo.updatedExtra) - require.Equal(t, 42.0, repo.updatedExtra["codex_5h_used_percent"]) + require.Equal(t, 58.0, repo.updatedExtra["codex_5h_used_percent"]) require.Equal(t, 88.0, repo.updatedExtra["codex_7d_used_percent"]) require.Contains(t, recorder.Body.String(), "test_complete") } @@ -170,7 +170,7 @@ func TestAccountTestService_OpenAI429PersistsSnapshotAndRateLimitState(t *testin resp.Header.Set("x-codex-primary-used-percent", "100") resp.Header.Set("x-codex-primary-reset-after-seconds", "604800") resp.Header.Set("x-codex-primary-window-minutes", "10080") - resp.Header.Set("x-codex-secondary-used-percent", "100") + resp.Header.Set("x-codex-secondary-used-percent", "0") resp.Header.Set("x-codex-secondary-reset-after-seconds", "18000") resp.Header.Set("x-codex-secondary-window-minutes", "300") diff --git a/backend/internal/service/account_usage_service_test.go b/backend/internal/service/account_usage_service_test.go index e0390c4c2a77..5f37aadbab7b 100644 --- a/backend/internal/service/account_usage_service_test.go +++ b/backend/internal/service/account_usage_service_test.go @@ -73,7 +73,7 @@ func TestExtractOpenAICodexProbeUpdatesAccepts429WithCodexHeaders(t *testing.T) headers.Set("x-codex-primary-used-percent", "100") headers.Set("x-codex-primary-reset-after-seconds", "604800") headers.Set("x-codex-primary-window-minutes", "10080") - headers.Set("x-codex-secondary-used-percent", "100") + headers.Set("x-codex-secondary-used-percent", "0") headers.Set("x-codex-secondary-reset-after-seconds", "18000") headers.Set("x-codex-secondary-window-minutes", "300") @@ -92,6 +92,33 @@ func TestExtractOpenAICodexProbeUpdatesAccepts429WithCodexHeaders(t *testing.T) } } +func TestBuildCodexUsageProgressFromExtra_UsesCanonicalUsedPercent(t *testing.T) { + t.Parallel() + now := time.Date(2026, 5, 30, 7, 4, 9, 0, time.UTC) + extra := map[string]any{ + "codex_5h_used_percent": 94.0, + "codex_5h_reset_at": now.Add(2 * time.Hour).Format(time.RFC3339), + "codex_7d_used_percent": 93.0, + "codex_7d_reset_at": now.Add(5 * 24 * time.Hour).Format(time.RFC3339), + } + + fiveHour := buildCodexUsageProgressFromExtra(extra, "5h", now) + if fiveHour == nil { + t.Fatal("expected non-nil 5h progress") + } + if fiveHour.Utilization != 94.0 { + t.Fatalf("5h Utilization = %v, want 94", fiveHour.Utilization) + } + + sevenDay := buildCodexUsageProgressFromExtra(extra, "7d", now) + if sevenDay == nil { + t.Fatal("expected non-nil 7d progress") + } + if sevenDay.Utilization != 93.0 { + t.Fatalf("7d Utilization = %v, want 93", sevenDay.Utilization) + } +} + func TestAccountUsageService_PersistOpenAICodexProbeSnapshotOnlyUpdatesExtra(t *testing.T) { t.Parallel() diff --git a/backend/internal/service/openai_gateway_service.go b/backend/internal/service/openai_gateway_service.go index cd5a40157c9e..c9e88bde6dfe 100644 --- a/backend/internal/service/openai_gateway_service.go +++ b/backend/internal/service/openai_gateway_service.go @@ -126,6 +126,19 @@ type NormalizedCodexLimits struct { Window7dMinutes *int } +func normalizeCodexFiveHourUsedPercent(raw *float64) *float64 { + if raw == nil { + return nil + } + // OpenAI's 5h Codex quota header is remaining%, despite the upstream header + // name saying "used"; the canonical codex_5h_used_percent field stores used%. + used := 100 - *raw + if used < 0 { + used = 0 + } + return &used +} + // Normalize converts primary/secondary fields to canonical 5h/7d fields. // Strategy: Compare window_minutes to determine which is 5h vs 7d. // Returns nil if snapshot is nil or has no useful data. @@ -184,7 +197,7 @@ func (s *OpenAICodexUsageSnapshot) Normalize() *NormalizedCodexLimits { // Assign values if use5hFromPrimary { - result.Used5hPercent = s.PrimaryUsedPercent + result.Used5hPercent = normalizeCodexFiveHourUsedPercent(s.PrimaryUsedPercent) result.Reset5hSeconds = s.PrimaryResetAfterSeconds result.Window5hMinutes = s.PrimaryWindowMinutes result.Used7dPercent = s.SecondaryUsedPercent @@ -194,7 +207,7 @@ func (s *OpenAICodexUsageSnapshot) Normalize() *NormalizedCodexLimits { result.Used7dPercent = s.PrimaryUsedPercent result.Reset7dSeconds = s.PrimaryResetAfterSeconds result.Window7dMinutes = s.PrimaryWindowMinutes - result.Used5hPercent = s.SecondaryUsedPercent + result.Used5hPercent = normalizeCodexFiveHourUsedPercent(s.SecondaryUsedPercent) result.Reset5hSeconds = s.SecondaryResetAfterSeconds result.Window5hMinutes = s.SecondaryWindowMinutes } diff --git a/backend/internal/service/openai_gateway_service_codex_snapshot_test.go b/backend/internal/service/openai_gateway_service_codex_snapshot_test.go index 654dd4cabe8c..22f5fa749248 100644 --- a/backend/internal/service/openai_gateway_service_codex_snapshot_test.go +++ b/backend/internal/service/openai_gateway_service_codex_snapshot_test.go @@ -104,6 +104,40 @@ func TestBuildCodexUsageExtraUpdates_UsesSnapshotUpdatedAt(t *testing.T) { } } +func TestBuildCodexUsageExtraUpdates_NormalizesFiveHourRemainingToUsedPercent(t *testing.T) { + primaryUsed := 93.0 + primaryReset := 86400 + primaryWindow := 10080 + secondaryRemaining := 6.0 + secondaryReset := 3600 + secondaryWindow := 300 + + snapshot := &OpenAICodexUsageSnapshot{ + PrimaryUsedPercent: &primaryUsed, + PrimaryResetAfterSeconds: &primaryReset, + PrimaryWindowMinutes: &primaryWindow, + SecondaryUsedPercent: &secondaryRemaining, + SecondaryResetAfterSeconds: &secondaryReset, + SecondaryWindowMinutes: &secondaryWindow, + UpdatedAt: "2026-05-30T07:04:09Z", + } + + updates := buildCodexUsageExtraUpdates(snapshot, time.Time{}) + if updates == nil { + t.Fatal("expected non-nil updates") + } + + if got := updates["codex_secondary_used_percent"]; got != 6.0 { + t.Fatalf("codex_secondary_used_percent = %v, want raw upstream value 6", got) + } + if got := updates["codex_5h_used_percent"]; got != 94.0 { + t.Fatalf("codex_5h_used_percent = %v, want 94", got) + } + if got := updates["codex_7d_used_percent"]; got != 93.0 { + t.Fatalf("codex_7d_used_percent = %v, want 93", got) + } +} + func TestBuildCodexUsageExtraUpdates_FallbackToNowWhenUpdatedAtInvalid(t *testing.T) { primaryUsed := 15.0 primaryReset := 30 diff --git a/backend/internal/service/openai_gateway_service_test.go b/backend/internal/service/openai_gateway_service_test.go index 8aad2fa6adb5..5c4e979d33a3 100644 --- a/backend/internal/service/openai_gateway_service_test.go +++ b/backend/internal/service/openai_gateway_service_test.go @@ -1774,7 +1774,7 @@ func TestOpenAIUpdateCodexUsageSnapshotFromHeaders(t *testing.T) { select { case updates := <-repo.updateExtraCalls: - require.Equal(t, 12.0, updates["codex_5h_used_percent"]) + require.Equal(t, 88.0, updates["codex_5h_used_percent"]) require.Equal(t, 34.0, updates["codex_7d_used_percent"]) require.Equal(t, 600, updates["codex_5h_reset_after_seconds"]) require.Equal(t, 86400, updates["codex_7d_reset_after_seconds"]) diff --git a/backend/internal/service/ratelimit_service_openai_test.go b/backend/internal/service/ratelimit_service_openai_test.go index aa5a070c74fc..107ac27ef414 100644 --- a/backend/internal/service/ratelimit_service_openai_test.go +++ b/backend/internal/service/ratelimit_service_openai_test.go @@ -51,7 +51,7 @@ func TestCalculateOpenAI429ResetTime_5hExhausted(t *testing.T) { headers.Set("x-codex-primary-used-percent", "50") headers.Set("x-codex-primary-reset-after-seconds", "500000") headers.Set("x-codex-primary-window-minutes", "10080") // 7 days - headers.Set("x-codex-secondary-used-percent", "100") + headers.Set("x-codex-secondary-used-percent", "0") headers.Set("x-codex-secondary-reset-after-seconds", "3600") // 1 hour headers.Set("x-codex-secondary-window-minutes", "300") // 5 hours @@ -122,7 +122,7 @@ func TestCalculateOpenAI429ResetTime_ReversedWindowOrder(t *testing.T) { // Test when OpenAI sends primary as 5h and secondary as 7d (reversed) headers := http.Header{} - headers.Set("x-codex-primary-used-percent", "100") // This is 5h + headers.Set("x-codex-primary-used-percent", "0") // This is 5h remaining% headers.Set("x-codex-primary-reset-after-seconds", "3600") // 1 hour headers.Set("x-codex-primary-window-minutes", "300") // 5 hours - smaller! headers.Set("x-codex-secondary-used-percent", "50") @@ -180,7 +180,7 @@ func TestHandle429_OpenAIPersistsCodexSnapshotImmediately(t *testing.T) { headers.Set("x-codex-primary-used-percent", "100") headers.Set("x-codex-primary-reset-after-seconds", "604800") headers.Set("x-codex-primary-window-minutes", "10080") - headers.Set("x-codex-secondary-used-percent", "100") + headers.Set("x-codex-secondary-used-percent", "0") headers.Set("x-codex-secondary-reset-after-seconds", "18000") headers.Set("x-codex-secondary-window-minutes", "300") @@ -224,7 +224,7 @@ func TestNormalizedCodexLimits(t *testing.T) { pUsed := 100.0 pReset := 384607 pWindow := 10080 - sUsed := 3.0 + sRemaining := 3.0 sReset := 17369 sWindow := 300 @@ -232,7 +232,7 @@ func TestNormalizedCodexLimits(t *testing.T) { PrimaryUsedPercent: &pUsed, PrimaryResetAfterSeconds: &pReset, PrimaryWindowMinutes: &pWindow, - SecondaryUsedPercent: &sUsed, + SecondaryUsedPercent: &sRemaining, SecondaryResetAfterSeconds: &sReset, SecondaryWindowMinutes: &sWindow, } @@ -249,8 +249,8 @@ func TestNormalizedCodexLimits(t *testing.T) { if normalized.Reset7dSeconds == nil || *normalized.Reset7dSeconds != 384607 { t.Errorf("expected Reset7dSeconds=384607, got %v", normalized.Reset7dSeconds) } - if normalized.Used5hPercent == nil || *normalized.Used5hPercent != 3.0 { - t.Errorf("expected Used5hPercent=3, got %v", normalized.Used5hPercent) + if normalized.Used5hPercent == nil || *normalized.Used5hPercent != 97.0 { + t.Errorf("expected Used5hPercent=97, got %v", normalized.Used5hPercent) } if normalized.Reset5hSeconds == nil || *normalized.Reset5hSeconds != 17369 { t.Errorf("expected Reset5hSeconds=17369, got %v", normalized.Reset5hSeconds) @@ -338,11 +338,11 @@ func TestRateLimitService_HandleUpstreamError_403FallsBackToRawBody(t *testing.T func TestNormalizedCodexLimits_OnlySecondaryData(t *testing.T) { // Test when only secondary has data, no window_minutes - sUsed := 60.0 + sRemaining := 60.0 sReset := 3000 snapshot := &OpenAICodexUsageSnapshot{ - SecondaryUsedPercent: &sUsed, + SecondaryUsedPercent: &sRemaining, SecondaryResetAfterSeconds: &sReset, // No window_minutes, no primary data } @@ -354,8 +354,8 @@ func TestNormalizedCodexLimits_OnlySecondaryData(t *testing.T) { // Legacy assumption: primary=7d, secondary=5h // So secondary goes to 5h - if normalized.Used5hPercent == nil || *normalized.Used5hPercent != 60.0 { - t.Errorf("expected Used5hPercent=60, got %v", normalized.Used5hPercent) + if normalized.Used5hPercent == nil || *normalized.Used5hPercent != 40.0 { + t.Errorf("expected Used5hPercent=40, got %v", normalized.Used5hPercent) } if normalized.Reset5hSeconds == nil || *normalized.Reset5hSeconds != 3000 { t.Errorf("expected Reset5hSeconds=3000, got %v", normalized.Reset5hSeconds) @@ -370,13 +370,13 @@ func TestNormalizedCodexLimits_BothDataNoWindowMinutes(t *testing.T) { // Test when both have data but no window_minutes pUsed := 100.0 pReset := 400000 - sUsed := 50.0 + sRemaining := 30.0 sReset := 10000 snapshot := &OpenAICodexUsageSnapshot{ PrimaryUsedPercent: &pUsed, PrimaryResetAfterSeconds: &pReset, - SecondaryUsedPercent: &sUsed, + SecondaryUsedPercent: &sRemaining, SecondaryResetAfterSeconds: &sReset, // No window_minutes } @@ -393,8 +393,8 @@ func TestNormalizedCodexLimits_BothDataNoWindowMinutes(t *testing.T) { if normalized.Reset7dSeconds == nil || *normalized.Reset7dSeconds != 400000 { t.Errorf("expected Reset7dSeconds=400000, got %v", normalized.Reset7dSeconds) } - if normalized.Used5hPercent == nil || *normalized.Used5hPercent != 50.0 { - t.Errorf("expected Used5hPercent=50, got %v", normalized.Used5hPercent) + if normalized.Used5hPercent == nil || *normalized.Used5hPercent != 70.0 { + t.Errorf("expected Used5hPercent=70, got %v", normalized.Used5hPercent) } if normalized.Reset5hSeconds == nil || *normalized.Reset5hSeconds != 10000 { t.Errorf("expected Reset5hSeconds=10000, got %v", normalized.Reset5hSeconds) @@ -425,7 +425,7 @@ func TestCalculateOpenAI429ResetTime_UserProvidedScenario(t *testing.T) { // This is the exact scenario from the user: // codex_7d_used_percent: 100 // codex_7d_reset_after_seconds: 384607 (约4.5天后重置) - // codex_5h_used_percent: 3 + // codex_5h_used_percent: 97 (from upstream 3% remaining) // codex_5h_reset_after_seconds: 17369 (约4.8小时后重置) svc := &RateLimitService{} From bf3787de1f5842d869a74bdb59af92d57288a35d Mon Sep 17 00:00:00 2001 From: wucm667 Date: Sun, 31 May 2026 08:43:20 +0800 Subject: [PATCH 26/36] fix(gateway): allow Claude Code count_tokens --- .../internal/service/claude_code_validator.go | 13 ++- .../service/claude_code_validator_test.go | 91 +++++++++++++++++++ 2 files changed, 102 insertions(+), 2 deletions(-) diff --git a/backend/internal/service/claude_code_validator.go b/backend/internal/service/claude_code_validator.go index 4e8ced67954b..2c5ded6f9b73 100644 --- a/backend/internal/service/claude_code_validator.go +++ b/backend/internal/service/claude_code_validator.go @@ -56,7 +56,7 @@ func NewClaudeCodeValidator() *ClaudeCodeValidator { // 采用与 claude-relay-service 完全一致的验证策略： // // Step 1: User-Agent 检查 (必需) - 必须是 claude-cli/x.x.x -// Step 2: 对于非 messages 路径，只要 UA 匹配就通过 +// Step 2: 对于非 messages 路径和 /messages/count_tokens，只要 UA 匹配就通过 // Step 3: 检查 max_tokens=1 + haiku 探测请求绕过（UA 已验证） // Step 4: 对于 messages 路径，进行严格验证： // - System prompt 相似度检查 @@ -71,12 +71,17 @@ func (v *ClaudeCodeValidator) Validate(r *http.Request, body map[string]any) boo return false } - // Step 2: 非 messages 路径，只要 UA 匹配就通过 + // Step 2: 非 messages 路径只要 UA 匹配就通过 path := r.URL.Path if !strings.Contains(path, "messages") { return true } + // count_tokens 是 Claude Code 官方辅助请求，通常不携带完整 messages system prompt。 + if isMessagesCountTokensPath(path) { + return true + } + // Step 3: 检查 max_tokens=1 + haiku 探测请求绕过 // 这类请求用于 Claude Code 验证 API 连通性，不携带 system prompt if isMaxTokensOneHaiku, ok := IsMaxTokensOneHaikuRequestFromContext(r.Context()); ok && isMaxTokensOneHaiku { @@ -128,6 +133,10 @@ func (v *ClaudeCodeValidator) Validate(r *http.Request, body map[string]any) boo return true } +func isMessagesCountTokensPath(path string) bool { + return strings.HasSuffix(path, "/messages/count_tokens") +} + // hasClaudeCodeSystemPrompt 检查请求是否包含 Claude Code 系统提示词 // 使用字符串相似度匹配（Dice coefficient） func (v *ClaudeCodeValidator) hasClaudeCodeSystemPrompt(body map[string]any) bool { diff --git a/backend/internal/service/claude_code_validator_test.go b/backend/internal/service/claude_code_validator_test.go index f87c56e839b7..a4b3050587fb 100644 --- a/backend/internal/service/claude_code_validator_test.go +++ b/backend/internal/service/claude_code_validator_test.go @@ -48,6 +48,97 @@ func TestClaudeCodeValidator_MessagesWithoutProbeStillNeedStrictValidation(t *te require.False(t, ok) } +func TestClaudeCodeValidator_CountTokensPathUAOnly(t *testing.T) { + validator := NewClaudeCodeValidator() + req := httptest.NewRequest(http.MethodPost, "http://example.com/v1/messages/count_tokens", nil) + req.Header.Set("User-Agent", "claude-cli/2.1.156 (Claude Code)") + + ok := validator.Validate(req, map[string]any{ + "model": "claude-opus-4-8", + }) + require.True(t, ok) +} + +func TestClaudeCodeValidator_CountTokensPathRequiresUA(t *testing.T) { + validator := NewClaudeCodeValidator() + req := httptest.NewRequest(http.MethodPost, "http://example.com/v1/messages/count_tokens", nil) + req.Header.Set("User-Agent", "curl/8.0.0") + + ok := validator.Validate(req, map[string]any{ + "model": "claude-opus-4-8", + }) + require.False(t, ok) +} + +func TestClaudeCodeValidator_MessagesPathFullValid(t *testing.T) { + validator := NewClaudeCodeValidator() + req := httptest.NewRequest(http.MethodPost, "http://example.com/v1/messages", nil) + req.Header.Set("User-Agent", "claude-cli/2.1.156 (Claude Code)") + req.Header.Set("X-App", "claude-code") + req.Header.Set("anthropic-beta", "claude-code-20250219") + req.Header.Set("anthropic-version", "2023-06-01") + + ok := validator.Validate(req, map[string]any{ + "model": "claude-opus-4-8", + "stream": true, + "system": []any{ + map[string]any{ + "type": "text", + "text": "You are Claude Code, Anthropic's official CLI for Claude.", + }, + }, + "metadata": map[string]any{ + "user_id": "user_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa_account__session_aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa", + }, + }) + require.True(t, ok) +} + +func TestClaudeCodeValidator_MessagesPathRejectsNonClaudeCodeUA(t *testing.T) { + validator := NewClaudeCodeValidator() + req := httptest.NewRequest(http.MethodPost, "http://example.com/v1/messages", nil) + req.Header.Set("User-Agent", "curl/8.0.0") + req.Header.Set("X-App", "claude-code") + req.Header.Set("anthropic-beta", "claude-code-20250219") + req.Header.Set("anthropic-version", "2023-06-01") + + ok := validator.Validate(req, map[string]any{ + "model": "claude-opus-4-8", + "stream": true, + "system": []any{ + map[string]any{ + "type": "text", + "text": "You are Claude Code, Anthropic's official CLI for Claude.", + }, + }, + "metadata": map[string]any{ + "user_id": "user_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa_account__session_aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa", + }, + }) + require.False(t, ok) +} + +func TestClaudeCodeValidator_MessagesPathWithoutSystemPromptStillRejected(t *testing.T) { + validator := NewClaudeCodeValidator() + req := httptest.NewRequest(http.MethodPost, "http://example.com/v1/messages", nil) + req.Header.Set("User-Agent", "claude-cli/2.1.156 (Claude Code)") + req.Header.Set("X-App", "claude-code") + req.Header.Set("anthropic-beta", "claude-code-20250219") + req.Header.Set("anthropic-version", "2023-06-01") + + ok := validator.Validate(req, map[string]any{ + "model": "claude-opus-4-8", + "stream": true, + "messages": []any{ + map[string]any{"role": "user", "content": "hello"}, + }, + "metadata": map[string]any{ + "user_id": "user_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa_account__session_aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa", + }, + }) + require.False(t, ok) +} + func TestClaudeCodeValidator_NonMessagesPathUAOnly(t *testing.T) { validator := NewClaudeCodeValidator() req := httptest.NewRequest(http.MethodPost, "http://example.com/v1/models", nil) From 1e2193c3d27b7d770158fc7e133772b218bbe4dd Mon Sep 17 00:00:00 2001 From: gsh Date: Sun, 31 May 2026 15:09:06 +0800 Subject: [PATCH 27/36] fix: avoid websocket usage dedup conflicts --- .../openai_gateway_record_usage_test.go | 31 +++++++++++++++++++ .../service/openai_gateway_service.go | 5 +++ 2 files changed, 36 insertions(+) diff --git a/backend/internal/service/openai_gateway_record_usage_test.go b/backend/internal/service/openai_gateway_record_usage_test.go index 9769a82ed90a..318c0861136f 100644 --- a/backend/internal/service/openai_gateway_record_usage_test.go +++ b/backend/internal/service/openai_gateway_record_usage_test.go @@ -721,6 +721,37 @@ func TestOpenAIGatewayServiceRecordUsage_PrefersClientRequestIDOverUpstreamReque require.Equal(t, "client:openai-client-stable-123", usageRepo.lastLog.RequestID) } +func TestOpenAIGatewayServiceRecordUsage_WSModePrefersUpstreamRequestIDOverClientRequestID(t *testing.T) { + usageRepo := &openAIRecordUsageLogRepoStub{} + billingRepo := &openAIRecordUsageBillingRepoStub{result: &UsageBillingApplyResult{Applied: true}} + userRepo := &openAIRecordUsageUserRepoStub{} + subRepo := &openAIRecordUsageSubRepoStub{} + svc := newOpenAIRecordUsageServiceWithBillingRepoForTest(usageRepo, billingRepo, userRepo, subRepo, nil) + + ctx := context.WithValue(context.Background(), ctxkey.ClientRequestID, "openai-ws-connection-123") + err := svc.RecordUsage(ctx, &OpenAIRecordUsageInput{ + Result: &OpenAIForwardResult{ + RequestID: "resp_openai_ws_turn_456", + OpenAIWSMode: true, + Usage: OpenAIUsage{ + InputTokens: 8, + OutputTokens: 4, + }, + Model: "gpt-5.1", + Duration: time.Second, + }, + APIKey: &APIKey{ID: 10050}, + User: &User{ID: 20050}, + Account: &Account{ID: 30050}, + }) + + require.NoError(t, err) + require.NotNil(t, billingRepo.lastCmd) + require.Equal(t, "resp_openai_ws_turn_456", billingRepo.lastCmd.RequestID) + require.NotNil(t, usageRepo.lastLog) + require.Equal(t, "resp_openai_ws_turn_456", usageRepo.lastLog.RequestID) +} + func TestOpenAIGatewayServiceRecordUsage_GeneratesRequestIDWhenAllSourcesMissing(t *testing.T) { usageRepo := &openAIRecordUsageLogRepoStub{} billingRepo := &openAIRecordUsageBillingRepoStub{result: &UsageBillingApplyResult{Applied: true}} diff --git a/backend/internal/service/openai_gateway_service.go b/backend/internal/service/openai_gateway_service.go index cd5a40157c9e..10080f31737a 100644 --- a/backend/internal/service/openai_gateway_service.go +++ b/backend/internal/service/openai_gateway_service.go @@ -5761,6 +5761,11 @@ func (s *OpenAIGatewayService) RecordUsage(ctx context.Context, input *OpenAIRec durationMs := int(result.Duration.Milliseconds()) accountRateMultiplier := account.BillingRateMultiplier() requestID := resolveUsageBillingRequestID(ctx, result.RequestID) + if result.OpenAIWSMode { + if upstreamRequestID := strings.TrimSpace(result.RequestID); upstreamRequestID != "" { + requestID = upstreamRequestID + } + } // 确定 RequestedModel（渠道映射前的原始模型） requestedModel := result.Model From 8ac2e23fb3e16323cea0a6d7dea8986467487486 Mon Sep 17 00:00:00 2001 From: name <136912576+is7Qin@users.noreply.github.com> Date: Sun, 31 May 2026 16:00:35 +0800 Subject: [PATCH 28/36] refactor(gateway): defer OpenAI request map decoding Keep the Responses hot path on raw request bytes until complex mutation or retry branches need a decoded map, reducing large-body retention under high concurrency. --- .../service/gateway_service_benchmark_test.go | 75 ++ .../service/image_generation_intent.go | 90 ++- .../service/image_generation_intent_test.go | 11 + .../service/openai_gateway_service.go | 687 ++++++++++-------- .../openai_gateway_service_hotpath_test.go | 454 ++++++++++++ .../openai_ws_forwarder_success_test.go | 85 +++ 6 files changed, 1086 insertions(+), 316 deletions(-) diff --git a/backend/internal/service/gateway_service_benchmark_test.go b/backend/internal/service/gateway_service_benchmark_test.go index b2df143591f3..f6cd14041544 100644 --- a/backend/internal/service/gateway_service_benchmark_test.go +++ b/backend/internal/service/gateway_service_benchmark_test.go @@ -124,6 +124,64 @@ func BenchmarkOpenAIResponses_LargeInputDecodeMap(b *testing.B) { } } +func BenchmarkOpenAIResponses_LargeInputRawPatch(b *testing.B) { + for _, size := range benchmarkBodySizes() { + b.Run(size.name, func(b *testing.B) { + body := buildLargeOpenAIResponsesBody(size.bytes) + + b.SetBytes(int64(len(body))) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + view := newOpenAIRequestView(body) + view.MarkPatchSet("instructions", "You are a helpful coding assistant.") + view.MarkPatchSet("reasoning.effort", "none") + patched, err := view.ApplyPatches() + if err != nil { + b.Fatalf("应用 OpenAI raw patch 失败: %v", err) + } + benchmarkIntSink = len(patched) + } + }) + } +} + +func BenchmarkOpenAIResponses_LargeInputImageBillingRaw(b *testing.B) { + for _, size := range benchmarkBodySizes() { + b.Run(size.name, func(b *testing.B) { + body := buildLargeOpenAIResponsesImageToolBody(size.bytes) + + b.SetBytes(int64(len(body))) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + cfg, err := resolveOpenAIResponsesImageBillingConfigDetailedFromBody(body, "gpt-5.4") + if err != nil { + b.Fatalf("解析 OpenAI 图片计费配置失败: %v", err) + } + benchmarkStringSink = cfg.Model + cfg.SizeTier + cfg.InputSize + } + }) + } +} + +func BenchmarkOpenAIResponses_LargeInputEmptyBase64Guard(b *testing.B) { + for _, size := range benchmarkBodySizes() { + b.Run(size.name, func(b *testing.B) { + body := buildLargeOpenAIResponsesBody(size.bytes) + + b.SetBytes(int64(len(body))) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + if openAIRequestBodyMayContainEmptyBase64InputImage(body) { + benchmarkIntSink++ + } + } + }) + } +} + func BenchmarkOpenAIResponses_LargeInputFunctionCallValidation(b *testing.B) { for _, size := range benchmarkBodySizes() { b.Run(size.name, func(b *testing.B) { @@ -264,3 +322,20 @@ func buildLargeOpenAIResponsesToolContinuationBody(targetBytes int) []byte { builder.WriteString(`]}`) return []byte(builder.String()) } + +func buildLargeOpenAIResponsesImageToolBody(targetBytes int) []byte { + var builder strings.Builder + builder.Grow(targetBytes + 1024) + builder.WriteString(`{"model":"gpt-5.4","stream":false,"tools":[{"type":"image_generation","model":"gpt-image-2","size":"2048x1152"}],"input":[`) + for i := 0; builder.Len() < targetBytes; i++ { + if i > 0 { + builder.WriteByte(',') + } + builder.WriteString(`{"type":"message","role":"user","content":[{"type":"input_text","text":"`) + builder.WriteString(strings.Repeat("openai image billing payload ", 48)) + builder.WriteString(strconv.Itoa(i)) + builder.WriteString(`"}]}`) + } + builder.WriteString(`]}`) + return []byte(builder.String()) +} diff --git a/backend/internal/service/image_generation_intent.go b/backend/internal/service/image_generation_intent.go index 4aca1239ff97..80b6c66d56a7 100644 --- a/backend/internal/service/image_generation_intent.go +++ b/backend/internal/service/image_generation_intent.go @@ -1,7 +1,6 @@ package service import ( - "encoding/json" "strings" "github.com/tidwall/gjson" @@ -91,7 +90,7 @@ func openAIJSONToolsContainImageGeneration(tools gjson.Result) bool { } found := false tools.ForEach(func(_, item gjson.Result) bool { - if strings.TrimSpace(item.Get("type").String()) == "image_generation" { + if openAIJSONString(item.Get("type")) == "image_generation" { found = true return false } @@ -100,6 +99,36 @@ func openAIJSONToolsContainImageGeneration(tools gjson.Result) bool { return found } +func openAIRequestBodyHasImageGenerationTool(body []byte) bool { + if len(body) == 0 || !gjson.ValidBytes(body) { + return false + } + return openAIJSONToolsContainImageGeneration(gjson.GetBytes(body, "tools")) +} + +func openAIRequestBodyImageGenerationToolNeedsNormalization(body []byte) bool { + if len(body) == 0 || !gjson.ValidBytes(body) { + return false + } + tools := gjson.GetBytes(body, "tools") + if !tools.IsArray() { + return false + } + needsNormalization := false + tools.ForEach(func(_, item gjson.Result) bool { + if openAIJSONString(item.Get("type")) != "image_generation" { + return true + } + // 只有旧字段需要迁移时才进入 map 修改，纯计费读取保持 raw 路径。 + if item.Get("format").Exists() || item.Get("compression").Exists() { + needsNormalization = true + return false + } + return true + }) + return needsNormalization +} + func openAIJSONToolChoiceSelectsImageGeneration(choice gjson.Result) bool { if !choice.Exists() { return false @@ -159,17 +188,6 @@ func apiKeyGroup(apiKey *APIKey) *Group { return apiKey.Group } -func cloneRequestMapForImageIntent(body []byte) map[string]any { - if len(body) == 0 { - return nil - } - var out map[string]any - if err := json.Unmarshal(body, &out); err != nil { - return nil - } - return out -} - type OpenAIResponsesImageBillingConfig struct { Model string SizeTier string @@ -225,8 +243,43 @@ func resolveOpenAIResponsesImageBillingConfigFromBody(body []byte, fallbackModel } func resolveOpenAIResponsesImageBillingConfigDetailedFromBody(body []byte, fallbackModel string) (OpenAIResponsesImageBillingConfig, error) { - reqBody := cloneRequestMapForImageIntent(body) - return resolveOpenAIResponsesImageBillingConfigDetailed(reqBody, fallbackModel) + imageModel := "" + imageSize := "" + hasImageTool := false + if len(body) > 0 && gjson.ValidBytes(body) { + tools := gjson.GetBytes(body, "tools") + if tools.IsArray() { + tools.ForEach(func(_, item gjson.Result) bool { + if openAIJSONString(item.Get("type")) != "image_generation" { + return true + } + hasImageTool = true + imageModel = openAIJSONString(item.Get("model")) + imageSize = openAIJSONString(item.Get("size")) + return false + }) + } + if imageSize == "" { + imageSize = openAIJSONString(gjson.GetBytes(body, "size")) + } + if imageModel == "" { + bodyModel := openAIJSONString(gjson.GetBytes(body, "model")) + if isOpenAIImageBillingModelAlias(bodyModel) || !hasImageTool { + imageModel = bodyModel + } + } + } + if imageModel == "" && hasImageTool { + imageModel = "gpt-image-2" + } + if imageModel == "" { + imageModel = strings.TrimSpace(fallbackModel) + } + return OpenAIResponsesImageBillingConfig{ + Model: imageModel, + SizeTier: normalizeOpenAIImageSizeTier(imageSize), + InputSize: imageSize, + }, nil } func isOpenAIImageBillingModelAlias(model string) bool { @@ -236,3 +289,10 @@ func isOpenAIImageBillingModelAlias(model string) bool { } return isOpenAIImageGenerationModel(normalized) || strings.Contains(normalized, "image") } + +func openAIJSONString(value gjson.Result) string { + if value.Type != gjson.String { + return "" + } + return strings.TrimSpace(value.String()) +} diff --git a/backend/internal/service/image_generation_intent_test.go b/backend/internal/service/image_generation_intent_test.go index 4621e9d98572..59aab39c4b4b 100644 --- a/backend/internal/service/image_generation_intent_test.go +++ b/backend/internal/service/image_generation_intent_test.go @@ -84,6 +84,17 @@ func TestResolveOpenAIResponsesImageBillingConfigToolModelWins(t *testing.T) { require.Equal(t, "2K", imageSize) } +func TestResolveOpenAIResponsesImageBillingConfigFromBodyIgnoresUnrelatedLargeInput(t *testing.T) { + cfg, err := resolveOpenAIResponsesImageBillingConfigDetailedFromBody( + []byte(`{"model":"mapped-text-model","tools":[{"type":"image_generation","model":"gpt-image-2","size":"2048x1152"}],"input":[{"type":"message","content":[{"type":"input_text","text":"hi","nonce":1e1000000}]}]}`), + "requested-model", + ) + require.NoError(t, err) + require.Equal(t, "gpt-image-2", cfg.Model) + require.Equal(t, "2K", cfg.SizeTier) + require.Equal(t, "2048x1152", cfg.InputSize) +} + func TestResolveOpenAIResponsesImageBillingConfigSupportsOfficialAndCustomSizes(t *testing.T) { tests := []struct { name string diff --git a/backend/internal/service/openai_gateway_service.go b/backend/internal/service/openai_gateway_service.go index 87ffb2cb2197..d17b85d0d26f 100644 --- a/backend/internal/service/openai_gateway_service.go +++ b/backend/internal/service/openai_gateway_service.go @@ -2397,172 +2397,83 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco return s.forwardOpenAIPassthrough(ctx, c, account, originalBody, reqModel, reasoningEffort, reqStream, startTime) } - reqBody, err := requestView.Decode(c) - if err != nil { - return nil, err - } - - if v, ok := reqBody["model"].(string); ok { - reqModel = v - originalModel = reqModel - } - if v, ok := reqBody["stream"].(bool); ok { - reqStream = v - } - if promptCacheKey == "" { - if v, ok := reqBody["prompt_cache_key"].(string); ok { - promptCacheKey = strings.TrimSpace(v) - } - } - apiKey := getAPIKeyFromContext(c) - imageGenerationAllowed := GroupAllowsImageGeneration(nil) - if apiKey != nil { - imageGenerationAllowed = GroupAllowsImageGeneration(apiKey.Group) - } - codexImageGenerationBridgeEnabled := isCodexCLI && imageGenerationAllowed && s.isCodexImageGenerationBridgeEnabled(ctx, account, apiKey) - if IsImageGenerationIntentMap(openAIResponsesEndpoint, reqModel, reqBody) && !imageGenerationAllowed { - MarkOpsClientBusinessLimited(c, OpsClientBusinessLimitedReasonLocalFeatureGate) - c.JSON(http.StatusForbidden, gin.H{ - "error": gin.H{ - "type": "permission_error", - "message": ImageGenerationPermissionMessage(), - }, - }) - return nil, errors.New("image generation disabled for group") - } - - // Track if body needs re-serialization bodyModified := false - // 单字段补丁快速路径：只要整个变更集最终可归约为同一路径的 set/delete，就避免全量 Marshal。 - patchDisabled := false - patchHasOp := false - patchDelete := false - patchPath := "" - var patchValue any - markPatchSet := func(path string, value any) { - if strings.TrimSpace(path) == "" { - patchDisabled = true - return + var reqBody map[string]any + ensureReqBody := func() (map[string]any, error) { + if requestView.HasPatches() { + patchedBody, patchErr := requestView.ApplyPatches() + if patchErr != nil { + return nil, patchErr + } + body = patchedBody + requestView = newOpenAIRequestView(body) + reqBody = nil + bodyModified = false } - if patchDisabled { - return + if reqBody != nil { + return reqBody, nil } - if !patchHasOp { - patchHasOp = true - patchDelete = false - patchPath = path - patchValue = value - return + decoded, decodeErr := requestView.Decode(c) + if decodeErr != nil { + return nil, decodeErr } - if patchDelete || patchPath != path { - patchDisabled = true + reqBody = decoded + return reqBody, nil + } + markPatchSet := func(path string, value any) { + bodyModified = true + if requestView.patchesDisabled { + if reqBody != nil { + setOpenAIRequestMapPath(reqBody, path, value) + } return } - patchValue = value + requestView.MarkPatchSet(path, value) } markPatchDelete := func(path string) { - if strings.TrimSpace(path) == "" { - patchDisabled = true - return - } - if patchDisabled { - return - } - if !patchHasOp { - patchHasOp = true - patchDelete = true - patchPath = path + bodyModified = true + if requestView.patchesDisabled { + if reqBody != nil { + deleteOpenAIRequestMapPath(reqBody, path) + } return } - if !patchDelete || patchPath != path { - patchDisabled = true - } + requestView.MarkPatchDelete(path) } disablePatch := func() { - patchDisabled = true - } - - // 非透传模式下，instructions 为空时注入默认指令。 - if isInstructionsEmpty(reqBody) && !compatMessagesBridge { - reqBody["instructions"] = "You are a helpful coding assistant." - bodyModified = true - markPatchSet("instructions", "You are a helpful coding assistant.") + requestView.DisablePatches() } - - if codexImageGenerationBridgeEnabled && ensureOpenAIResponsesImageGenerationTool(reqBody) { + markDecodedModified := func() { bodyModified = true disablePatch() - logger.LegacyPrintf("service.openai_gateway", "[OpenAI] Injected /responses image_generation tool for Codex client") } - if normalizeOpenAIResponsesImageGenerationTools(reqBody) { - bodyModified = true - disablePatch() - logger.LegacyPrintf("service.openai_gateway", "[OpenAI] Normalized /responses image_generation tool payload") + apiKey := getAPIKeyFromContext(c) + imageGenerationAllowed := GroupAllowsImageGeneration(nil) + if apiKey != nil { + imageGenerationAllowed = GroupAllowsImageGeneration(apiKey.Group) } - if codexImageGenerationBridgeEnabled && applyCodexImageGenerationBridgeInstructions(reqBody) { - bodyModified = true - disablePatch() - logger.LegacyPrintf("service.openai_gateway", "[OpenAI] Added Codex image_generation bridge instructions") + codexImageGenerationBridgeEnabled := isCodexCLI && imageGenerationAllowed && s.isCodexImageGenerationBridgeEnabled(ctx, account, apiKey) + imageIntent := IsImageGenerationIntent(openAIResponsesEndpoint, reqModel, body) + if imageIntent && !imageGenerationAllowed { + MarkOpsClientBusinessLimited(c, OpsClientBusinessLimitedReasonLocalFeatureGate) + c.JSON(http.StatusForbidden, gin.H{"error": gin.H{"type": "permission_error", "message": ImageGenerationPermissionMessage()}}) + return nil, errors.New("image generation disabled for group") + } + + instructions := gjson.GetBytes(body, "instructions") + instructionsEmpty := !instructions.Exists() || instructions.Type != gjson.String || strings.TrimSpace(instructions.String()) == "" + if instructionsEmpty && !compatMessagesBridge { + markPatchSet("instructions", "You are a helpful coding assistant.") } - // 对所有请求执行模型映射（包含 Codex CLI）。 billingModel := account.GetMappedModel(reqModel) if billingModel != reqModel { logger.LegacyPrintf("service.openai_gateway", "[OpenAI] Model mapping applied: %s -> %s (account: %s, isCodexCLI: %v)", reqModel, billingModel, account.Name, isCodexCLI) - reqBody["model"] = billingModel - bodyModified = true + reqModel = billingModel markPatchSet("model", billingModel) } upstreamModel := billingModel - if imageGenerationAllowed && normalizeOpenAIResponsesImageOnlyModel(reqBody) { - bodyModified = true - disablePatch() - if model, ok := reqBody["model"].(string); ok { - upstreamModel = strings.TrimSpace(model) - } - logger.LegacyPrintf( - "service.openai_gateway", - "[OpenAI] Normalized /responses image-only model request inbound_model=%s image_model=%s upstream_model=%s", - reqModel, - billingModel, - upstreamModel, - ) - } - if err := validateOpenAIResponsesImageModel(reqBody, upstreamModel); err != nil { - setOpsUpstreamError(c, http.StatusBadRequest, err.Error(), "") - c.JSON(http.StatusBadRequest, gin.H{ - "error": gin.H{ - "type": "invalid_request_error", - "message": err.Error(), - "param": "model", - }, - }) - return nil, err - } - if hasOpenAIImageGenerationTool(reqBody) { - logger.LegacyPrintf( - "service.openai_gateway", - "[OpenAI] /responses image_generation request inbound_model=%s mapped_model=%s account_type=%s", - reqModel, - upstreamModel, - account.Type, - ) - } - if err := validateCodexSparkInput(reqBody, upstreamModel); err != nil { - setOpsUpstreamError(c, http.StatusBadRequest, err.Error(), "") - c.JSON(http.StatusBadRequest, gin.H{ - "error": gin.H{ - "type": "invalid_request_error", - "message": err.Error(), - "param": "input", - }, - }) - return nil, err - } - - // Compact-only model 映射：仅在 /responses/compact 路径生效，且优先级高于 - // OAuth 模型规范化（避免 OAuth 规范化覆盖 compact-only 自定义模型）。 isCompactRequest := isOpenAIResponsesCompactPath(c) compactMapped := false if isCompactRequest { @@ -2570,69 +2481,99 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco if compactMappedModel != "" && compactMappedModel != billingModel { compactMapped = true upstreamModel = compactMappedModel - reqBody["model"] = compactMappedModel - bodyModified = true + reqModel = compactMappedModel markPatchSet("model", compactMappedModel) logger.LegacyPrintf("service.openai_gateway", "[OpenAI] Compact model mapping applied: %s -> %s (account: %s, isCodexCLI: %v)", billingModel, compactMappedModel, account.Name, isCodexCLI) } } - - // OpenAI OAuth 账号走 ChatGPT internal Codex endpoint，需要将模型名规范化为 - // 上游可识别的 Codex/GPT 系列。API Key 账号则应保留原始/映射后的模型名， - // 以兼容自定义 base_url 的 OpenAI-compatible 上游。 - if model, ok := reqBody["model"].(string); ok { - if !compactMapped { - upstreamModel = normalizeOpenAIModelForUpstream(account, model) - if upstreamModel != "" && upstreamModel != model { - logger.LegacyPrintf("service.openai_gateway", "[OpenAI] Upstream model resolved: %s -> %s (account: %s, type: %s, isCodexCLI: %v)", - model, upstreamModel, account.Name, account.Type, isCodexCLI) - reqBody["model"] = upstreamModel - bodyModified = true - markPatchSet("model", upstreamModel) - } + if !compactMapped { + modelForNormalize := reqModel + if modelForNormalize == "" { + modelForNormalize = requestView.Model + } + upstreamModel = normalizeOpenAIModelForUpstream(account, modelForNormalize) + if upstreamModel != "" && upstreamModel != modelForNormalize { + logger.LegacyPrintf("service.openai_gateway", "[OpenAI] Upstream model resolved: %s -> %s (account: %s, type: %s, isCodexCLI: %v)", modelForNormalize, upstreamModel, account.Name, account.Type, isCodexCLI) + reqModel = upstreamModel + markPatchSet("model", upstreamModel) } + } + if strings.TrimSpace(gjson.GetBytes(body, "reasoning.effort").String()) == "minimal" { + markPatchSet("reasoning.effort", "none") + logger.LegacyPrintf("service.openai_gateway", "[OpenAI] Normalized reasoning.effort: minimal -> none (account: %s)", account.Name) + } - // 移除 gpt-5.2-codex 以下的版本 verbosity 参数 - // 确保高版本模型向低版本模型映射不报错 - if !SupportsVerbosity(upstreamModel) { - if text, ok := reqBody["text"].(map[string]any); ok { - if _, exists := text["verbosity"]; exists { - delete(text, "verbosity") - bodyModified = true - markPatchDelete("text.verbosity") - } + imageIntent = imageIntent || IsImageGenerationIntent(openAIResponsesEndpoint, reqModel, nil) || isOpenAIImageGenerationModel(upstreamModel) + if imageIntent && !imageGenerationAllowed { + MarkOpsClientBusinessLimited(c, OpsClientBusinessLimitedReasonLocalFeatureGate) + c.JSON(http.StatusForbidden, gin.H{"error": gin.H{"type": "permission_error", "message": ImageGenerationPermissionMessage()}}) + return nil, errors.New("image generation disabled for group") + } + + if imageGenerationAllowed && (codexImageGenerationBridgeEnabled || isOpenAIImageGenerationModel(requestView.Model) || openAIRequestBodyImageGenerationToolNeedsNormalization(body) || isOpenAIImageGenerationModel(upstreamModel)) { + decoded, decodeErr := ensureReqBody() + if decodeErr != nil { + return nil, decodeErr + } + if codexImageGenerationBridgeEnabled && ensureOpenAIResponsesImageGenerationTool(decoded) { + markDecodedModified() + logger.LegacyPrintf("service.openai_gateway", "[OpenAI] Injected /responses image_generation tool for Codex client") + } + if normalizeOpenAIResponsesImageGenerationTools(decoded) { + markDecodedModified() + logger.LegacyPrintf("service.openai_gateway", "[OpenAI] Normalized /responses image_generation tool payload") + } + if normalizeOpenAIResponsesImageOnlyModel(decoded) { + markDecodedModified() + if model, ok := decoded["model"].(string); ok { + upstreamModel = strings.TrimSpace(model) } + logger.LegacyPrintf("service.openai_gateway", "[OpenAI] Normalized /responses image-only model request inbound_model=%s image_model=%s upstream_model=%s", requestView.Model, billingModel, upstreamModel) + } + if err := validateOpenAIResponsesImageModel(decoded, upstreamModel); err != nil { + setOpsUpstreamError(c, http.StatusBadRequest, err.Error(), "") + c.JSON(http.StatusBadRequest, gin.H{"error": gin.H{"type": "invalid_request_error", "message": err.Error(), "param": "model"}}) + return nil, err + } + if hasOpenAIImageGenerationTool(decoded) { + imageIntent = true + logger.LegacyPrintf("service.openai_gateway", "[OpenAI] /responses image_generation request inbound_model=%s mapped_model=%s account_type=%s", requestView.Model, upstreamModel, account.Type) + } + if codexImageGenerationBridgeEnabled && applyCodexImageGenerationBridgeInstructions(decoded) { + markDecodedModified() + logger.LegacyPrintf("service.openai_gateway", "[OpenAI] Added Codex image_generation bridge instructions") } + } else if imageGenerationAllowed && imageIntent && openAIRequestBodyHasImageGenerationTool(body) { + logger.LegacyPrintf("service.openai_gateway", "[OpenAI] /responses image_generation request inbound_model=%s mapped_model=%s account_type=%s", requestView.Model, upstreamModel, account.Type) } - // 规范化 reasoning.effort 参数（minimal -> none），与上游允许值对齐。 - if reasoning, ok := reqBody["reasoning"].(map[string]any); ok { - if effort, ok := reasoning["effort"].(string); ok && effort == "minimal" { - reasoning["effort"] = "none" - bodyModified = true - markPatchSet("reasoning.effort", "none") - logger.LegacyPrintf("service.openai_gateway", "[OpenAI] Normalized reasoning.effort: minimal -> none (account: %s)", account.Name) + if isCodexSparkModel(upstreamModel) && openAIRequestBodyMayContainImageInput(body) { + decoded, decodeErr := ensureReqBody() + if decodeErr != nil { + return nil, decodeErr + } + if err := validateCodexSparkInput(decoded, upstreamModel); err != nil { + setOpsUpstreamError(c, http.StatusBadRequest, err.Error(), "") + c.JSON(http.StatusBadRequest, gin.H{"error": gin.H{"type": "invalid_request_error", "message": err.Error(), "param": "input"}}) + return nil, err } } if account.Type == AccountTypeOAuth { + decoded, decodeErr := ensureReqBody() + if decodeErr != nil { + return nil, decodeErr + } codexResult := codexTransformResult{} if compatMessagesBridge { - codexResult = applyCodexOAuthTransformWithOptions(reqBody, codexOAuthTransformOptions{ - IsCodexCLI: isCodexCLI, - IsCompact: isCompactRequest, - SkipDefaultInstructions: true, - PreserveToolCallIDs: true, - }) - ensureCodexOAuthInstructionsField(reqBody) - bodyModified = true - disablePatch() + codexResult = applyCodexOAuthTransformWithOptions(decoded, codexOAuthTransformOptions{IsCodexCLI: isCodexCLI, IsCompact: isCompactRequest, SkipDefaultInstructions: true, PreserveToolCallIDs: true}) + ensureCodexOAuthInstructionsField(decoded) + markDecodedModified() } else { - codexResult = applyCodexOAuthTransform(reqBody, isCodexCLI, isCompactRequest) + codexResult = applyCodexOAuthTransform(decoded, isCodexCLI, isCompactRequest) } if codexResult.Modified { - bodyModified = true - disablePatch() + markDecodedModified() } if codexResult.NormalizedModel != "" { upstreamModel = codexResult.NormalizedModel @@ -2642,90 +2583,57 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco } } - // Handle max_output_tokens based on platform and account type + if !SupportsVerbosity(upstreamModel) && gjson.GetBytes(body, "text.verbosity").Exists() { + markPatchDelete("text.verbosity") + } + if !isCodexCLI { - if maxOutputTokens, hasMaxOutputTokens := reqBody["max_output_tokens"]; hasMaxOutputTokens { + maxOutputTokens := gjson.GetBytes(body, "max_output_tokens") + if maxOutputTokens.Exists() { switch account.Platform { case PlatformOpenAI: - // For OpenAI API Key, remove max_output_tokens (not supported) - // For OpenAI OAuth (Responses API), keep it (supported) if account.Type == AccountTypeAPIKey { - delete(reqBody, "max_output_tokens") - bodyModified = true markPatchDelete("max_output_tokens") } case PlatformAnthropic: - // For Anthropic (Claude), convert to max_tokens - delete(reqBody, "max_output_tokens") - markPatchDelete("max_output_tokens") - if _, hasMaxTokens := reqBody["max_tokens"]; !hasMaxTokens { - reqBody["max_tokens"] = maxOutputTokens - disablePatch() + decoded, decodeErr := ensureReqBody() + if decodeErr != nil { + return nil, decodeErr } - bodyModified = true + delete(decoded, "max_output_tokens") + if _, hasMaxTokens := decoded["max_tokens"]; !hasMaxTokens { + decoded["max_tokens"] = maxOutputTokens.Value() + } + markDecodedModified() case PlatformGemini: - // For Gemini, remove (will be handled by Gemini-specific transform) - delete(reqBody, "max_output_tokens") - bodyModified = true markPatchDelete("max_output_tokens") default: - // For unknown platforms, remove to be safe - delete(reqBody, "max_output_tokens") - bodyModified = true markPatchDelete("max_output_tokens") } } - - // Also handle max_completion_tokens (similar logic) - if _, hasMaxCompletionTokens := reqBody["max_completion_tokens"]; hasMaxCompletionTokens { - if account.Type == AccountTypeAPIKey || account.Platform != PlatformOpenAI { - delete(reqBody, "max_completion_tokens") - bodyModified = true - markPatchDelete("max_completion_tokens") - } + if gjson.GetBytes(body, "max_completion_tokens").Exists() && (account.Type == AccountTypeAPIKey || account.Platform != PlatformOpenAI) { + markPatchDelete("max_completion_tokens") } - - // Remove unsupported fields (not supported by upstream OpenAI API) - unsupportedFields := []string{"prompt_cache_retention", "safety_identifier"} - for _, unsupportedField := range unsupportedFields { - if _, has := reqBody[unsupportedField]; has { - delete(reqBody, unsupportedField) - bodyModified = true + for _, unsupportedField := range []string{"prompt_cache_retention", "safety_identifier"} { + if gjson.GetBytes(body, unsupportedField).Exists() { markPatchDelete(unsupportedField) } } } - - // 仅在 WSv2 模式保留 previous_response_id，其他模式（HTTP/WSv1）统一过滤。 - // 注意：该规则同样适用于 Codex CLI 请求，避免 WSv1 向上游透传不支持字段。 - if wsDecision.Transport != OpenAIUpstreamTransportResponsesWebsocketV2 { - if _, has := reqBody["previous_response_id"]; has { - delete(reqBody, "previous_response_id") - bodyModified = true - markPatchDelete("previous_response_id") - } + if wsDecision.Transport != OpenAIUpstreamTransportResponsesWebsocketV2 && gjson.GetBytes(body, "previous_response_id").Exists() { + markPatchDelete("previous_response_id") } - - if sanitizeEmptyBase64InputImagesInOpenAIRequestBodyMap(reqBody) { - bodyModified = true - disablePatch() + if openAIRequestBodyMayContainEmptyBase64InputImage(body) { + decoded, decodeErr := ensureReqBody() + if decodeErr != nil { + return nil, decodeErr + } + if sanitizeEmptyBase64InputImagesInOpenAIRequestBodyMap(decoded) { + markDecodedModified() + } } - // Apply OpenAI fast policy (参照 Claude BetaPolicy 的 fast-mode 过滤)： - // 针对 body 的 service_tier 字段（"priority" 即 fast，"flex"），按策略 - // 执行 filter（删除字段）或 block（拒绝请求）。对 gpt-5.5 等模型屏蔽 - // fast 时在此生效。 - // - // 注意： - // 1. 此处统一使用 upstreamModel（已经过 GetMappedModel + - // normalizeOpenAIModelForUpstream + Codex OAuth normalize），与 - // chat-completions / messages 入口保持一致，避免不同入口因为模型 - // 维度不同而出现 whitelist 命中差异。 - // 2. action=pass 时也要把 raw "fast" 归一化为 "priority" 写回 body， - // 否则 native /responses 入口透传 "fast" 给上游会被拒。chat- - // completions 入口由 normalizeResponsesBodyServiceTier 完成同一 - // 行为，这里手工实现等效逻辑。 - if rawTier, ok := reqBody["service_tier"].(string); ok { + if rawTier := requestView.ServiceTier; rawTier != "" { if normTier := normalizedOpenAIServiceTierValue(rawTier); normTier != "" { action, errMsg := s.evaluateOpenAIFastPolicy(ctx, account, upstreamModel, normTier) switch action { @@ -2738,46 +2646,51 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco writeOpenAIFastPolicyBlockedResponse(c, blocked) return nil, blocked case BetaPolicyActionFilter: - delete(reqBody, "service_tier") - bodyModified = true - disablePatch() + markPatchDelete("service_tier") default: - // pass：若客户端传的是别名 "fast"，归一化为 "priority" - // 后写回 body，确保上游收到的是其能识别的规范值。 if normTier != rawTier { - reqBody["service_tier"] = normTier - bodyModified = true markPatchSet("service_tier", normTier) } } } } - if IsImageGenerationIntentMap(openAIResponsesEndpoint, reqModel, reqBody) && !imageGenerationAllowed { - MarkOpsClientBusinessLimited(c, OpsClientBusinessLimitedReasonLocalFeatureGate) - c.JSON(http.StatusForbidden, gin.H{ - "error": gin.H{ - "type": "permission_error", - "message": ImageGenerationPermissionMessage(), - }, - }) - return nil, errors.New("image generation disabled for group") + if bodyModified { + if requestView.HasPatches() { + if patchedBody, patchErr := requestView.ApplyPatches(); patchErr == nil { + body = patchedBody + requestView = newOpenAIRequestView(body) + reqBody = nil + bodyModified = false + } + } + if bodyModified { + decoded, decodeErr := ensureReqBody() + if decodeErr != nil { + return nil, decodeErr + } + var marshalErr error + body, marshalErr = marshalOpenAIUpstreamJSON(decoded) + if marshalErr != nil { + return nil, fmt.Errorf("serialize request body: %w", marshalErr) + } + requestView = newOpenAIRequestView(body) + } } imageBillingModel := "" imageSizeTier := "" imageInputSize := "" - if IsImageGenerationIntentMap(openAIResponsesEndpoint, reqModel, reqBody) { + if imageIntent { + var imageCfg OpenAIResponsesImageBillingConfig var imageCfgErr error - imageCfg, imageCfgErr := resolveOpenAIResponsesImageBillingConfigDetailed(reqBody, billingModel) + if reqBody != nil { + imageCfg, imageCfgErr = resolveOpenAIResponsesImageBillingConfigDetailed(reqBody, billingModel) + } else { + imageCfg, imageCfgErr = resolveOpenAIResponsesImageBillingConfigDetailedFromBody(body, billingModel) + } if imageCfgErr != nil { setOpsUpstreamError(c, http.StatusBadRequest, imageCfgErr.Error(), "") - c.JSON(http.StatusBadRequest, gin.H{ - "error": gin.H{ - "type": "invalid_request_error", - "message": imageCfgErr.Error(), - "param": "size", - }, - }) + c.JSON(http.StatusBadRequest, gin.H{"error": gin.H{"type": "invalid_request_error", "message": imageCfgErr.Error(), "param": "size"}}) return nil, imageCfgErr } imageBillingModel = imageCfg.Model @@ -2785,29 +2698,6 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco imageInputSize = imageCfg.InputSize } - // Re-serialize body only if modified - if bodyModified { - serializedByPatch := false - if !patchDisabled && patchHasOp { - var patchErr error - if patchDelete { - body, patchErr = sjson.DeleteBytes(body, patchPath) - } else { - body, patchErr = sjson.SetBytes(body, patchPath, patchValue) - } - if patchErr == nil { - serializedByPatch = true - } - } - if !serializedByPatch { - var marshalErr error - body, marshalErr = marshalOpenAIUpstreamJSON(reqBody) - if marshalErr != nil { - return nil, fmt.Errorf("serialize request body: %w", marshalErr) - } - } - } - // Get access token token, _, err := s.GetAccessToken(ctx, account) if err != nil { @@ -2816,8 +2706,11 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco // 命中 WS 时仅走 WebSocket Mode；不再自动回退 HTTP。 if wsDecision.Transport == OpenAIUpstreamTransportResponsesWebsocketV2 { - // WS 分支不会再回落 HTTP；重连恢复可直接更新 reqBody，避免额外保留一份完整顶层 map。 - wsReqBody := reqBody + // WS 分支需要结构化 payload 与重连恢复，命中后再触发 full-map decode。 + wsReqBody, err := ensureReqBody() + if err != nil { + return nil, err + } _, hasPreviousResponseID := wsReqBody["previous_response_id"] logOpenAIWSModeDebug( "forward_start account_id=%d account_type=%s model=%s stream=%v has_previous_response_id=%v", @@ -3076,8 +2969,12 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) upstreamCode := extractUpstreamErrorCode(respBody) if !httpInvalidEncryptedContentRetryTried && resp.StatusCode == http.StatusBadRequest && upstreamCode == "invalid_encrypted_content" { - if trimOpenAIEncryptedReasoningItems(reqBody) { - body, err = marshalOpenAIUpstreamJSON(reqBody) + decoded, decodeErr := ensureReqBody() + if decodeErr != nil { + return nil, decodeErr + } + if trimOpenAIEncryptedReasoningItems(decoded) { + body, err = marshalOpenAIUpstreamJSON(decoded) if err != nil { return nil, fmt.Errorf("serialize invalid_encrypted_content retry body: %w", err) } @@ -3118,8 +3015,8 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco } defer func() { _ = resp.Body.Close() }() - reasoningEffort := extractOpenAIReasoningEffort(reqBody, originalModel) - serviceTier := extractOpenAIServiceTier(reqBody) + reasoningEffort := extractOpenAIReasoningEffortFromBody(body, originalModel) + serviceTier := extractOpenAIServiceTierFromBody(body) // 上游接受后只保留计费需要的标量，避免响应处理期间继续保活完整 input/tools map。 reqBody = nil @@ -6283,6 +6180,14 @@ type openAIRequestView struct { PreviousResponseID string ServiceTier string ReasoningEffort string + patches []openAIRequestPatch + patchesDisabled bool +} + +type openAIRequestPatch struct { + path string + delete bool + value any } func newOpenAIRequestView(body []byte) openAIRequestView { @@ -6305,6 +6210,110 @@ func (v openAIRequestView) Decode(c *gin.Context) (map[string]any, error) { return getOpenAIRequestBodyMap(c, v.body) } +func (v *openAIRequestView) MarkPatchSet(path string, value any) { + if v == nil || v.patchesDisabled { + return + } + path = strings.TrimSpace(path) + if path == "" { + v.DisablePatches() + return + } + v.patches = append(v.patches, openAIRequestPatch{path: path, value: value}) +} + +func (v *openAIRequestView) MarkPatchDelete(path string) { + if v == nil || v.patchesDisabled { + return + } + path = strings.TrimSpace(path) + if path == "" { + v.DisablePatches() + return + } + v.patches = append(v.patches, openAIRequestPatch{path: path, delete: true}) +} + +func (v *openAIRequestView) DisablePatches() { + if v == nil { + return + } + v.patchesDisabled = true + v.patches = nil +} + +func (v openAIRequestView) HasPatches() bool { + return !v.patchesDisabled && len(v.patches) > 0 +} + +func (v openAIRequestView) ApplyPatches() ([]byte, error) { + if v.patchesDisabled || len(v.patches) == 0 { + return nil, errors.New("openai request patches disabled") + } + body := v.body + for _, patch := range v.patches { + var err error + if patch.delete { + body, err = sjson.DeleteBytes(body, patch.path) + } else { + body, err = sjson.SetBytes(body, patch.path, patch.value) + } + if err != nil { + return nil, err + } + } + return body, nil +} + +func setOpenAIRequestMapPath(reqBody map[string]any, path string, value any) { + path = strings.TrimSpace(path) + if reqBody == nil || path == "" { + return + } + parts := strings.Split(path, ".") + current := reqBody + for _, part := range parts[:len(parts)-1] { + part = strings.TrimSpace(part) + if part == "" { + return + } + next, _ := current[part].(map[string]any) + if next == nil { + next = map[string]any{} + current[part] = next + } + current = next + } + last := strings.TrimSpace(parts[len(parts)-1]) + if last != "" { + current[last] = value + } +} + +func deleteOpenAIRequestMapPath(reqBody map[string]any, path string) { + path = strings.TrimSpace(path) + if reqBody == nil || path == "" { + return + } + parts := strings.Split(path, ".") + current := reqBody + for _, part := range parts[:len(parts)-1] { + part = strings.TrimSpace(part) + if part == "" { + return + } + next, _ := current[part].(map[string]any) + if next == nil { + return + } + current = next + } + last := strings.TrimSpace(parts[len(parts)-1]) + if last != "" { + delete(current, last) + } +} + func extractOpenAIRequestMetaFromBody(body []byte) (model string, stream bool, promptCacheKey string) { view := newOpenAIRequestView(body) return view.Model, view.Stream, view.PromptCacheKey @@ -6758,8 +6767,84 @@ func buildOpenAIFastPolicyBlockedWSEvent(err *OpenAIFastBlockedError) []byte { return payload } +func openAIRequestBodyMayContainImageInput(body []byte) bool { + if len(body) == 0 { + return false + } + input := gjson.GetBytes(body, "input") + messages := gjson.GetBytes(body, "messages") + return openAIJSONValueMayContainImageInput(input) || openAIJSONValueMayContainImageInput(messages) +} + +func openAIJSONValueMayContainImageInput(value gjson.Result) bool { + if !value.Exists() { + return false + } + if value.IsArray() { + found := false + value.ForEach(func(_, item gjson.Result) bool { + if openAIJSONValueMayContainImageInput(item) { + found = true + return false + } + return true + }) + return found + } + if value.IsObject() { + if strings.TrimSpace(value.Get("type").String()) == "input_image" || value.Get("image_url").Exists() { + return true + } + return openAIJSONValueMayContainImageInput(value.Get("content")) + } + return false +} + +func openAIRequestBodyMayContainEmptyBase64InputImage(body []byte) bool { + if len(body) == 0 || !openAIRequestBodyMayContainInputImageToken(body) { + return false + } + input := gjson.GetBytes(body, "input") + if !input.Exists() { + return false + } + return openAIJSONValueMayContainEmptyBase64InputImage(input) +} + +func openAIRequestBodyMayContainInputImageToken(body []byte) bool { + if bytes.Contains(body, []byte("input_image")) { + return true + } + // JSON 字符串任意字符都可能被 unicode escape，遇到 \u 时交给 gjson 解码后的结构扫描兜底。 + return bytes.Contains(body, []byte("\\u")) +} + +func openAIJSONValueMayContainEmptyBase64InputImage(value gjson.Result) bool { + if !value.Exists() { + return false + } + if value.IsArray() { + found := false + value.ForEach(func(_, item gjson.Result) bool { + if openAIJSONValueMayContainEmptyBase64InputImage(item) { + found = true + return false + } + return true + }) + return found + } + if value.IsObject() { + if strings.TrimSpace(value.Get("type").String()) == "input_image" && isEmptyBase64DataURI(value.Get("image_url").String()) { + return true + } + return openAIJSONValueMayContainEmptyBase64InputImage(value.Get("content")) + } + return false +} + func sanitizeEmptyBase64InputImagesInOpenAIBody(body []byte) ([]byte, bool, error) { - if len(body) == 0 || !bytes.Contains(body, []byte(`"image_url"`)) || !bytes.Contains(body, []byte(`base64,`)) { + if !openAIRequestBodyMayContainEmptyBase64InputImage(body) { return body, false, nil } diff --git a/backend/internal/service/openai_gateway_service_hotpath_test.go b/backend/internal/service/openai_gateway_service_hotpath_test.go index df17ed2ef023..d240cec84fd9 100644 --- a/backend/internal/service/openai_gateway_service_hotpath_test.go +++ b/backend/internal/service/openai_gateway_service_hotpath_test.go @@ -1,12 +1,18 @@ package service import ( + "context" "encoding/json" + "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" + "github.com/tidwall/gjson" ) func TestOpenAIRequestView_ExtractsRawScalars(t *testing.T) { @@ -29,6 +35,454 @@ func TestOpenAIRequestView_DecodeKeepsFullMapBehavior(t *testing.T) { require.IsType(t, []any{}, reqBody["input"]) } +func TestOpenAIRequestView_ApplyPatches(t *testing.T) { + view := newOpenAIRequestView([]byte(`{"model":"gpt-5","previous_response_id":"resp_1","reasoning":{"effort":"minimal"},"input":[{"type":"message","content":"hi"}]}`)) + view.MarkPatchSet("model", "gpt-5.1") + view.MarkPatchDelete("previous_response_id") + view.MarkPatchSet("reasoning.effort", "none") + + patched, err := view.ApplyPatches() + require.NoError(t, err) + require.JSONEq(t, `{"model":"gpt-5.1","reasoning":{"effort":"none"},"input":[{"type":"message","content":"hi"}]}`, string(patched)) +} + +func TestOpenAIRequestView_ApplyPatchesDisabled(t *testing.T) { + view := newOpenAIRequestView([]byte(`{"model":"gpt-5"}`)) + view.MarkPatchSet("model", "gpt-5.1") + view.DisablePatches() + + _, err := view.ApplyPatches() + require.Error(t, err) +} + +func TestOpenAIRequestView_HasPatches(t *testing.T) { + view := newOpenAIRequestView([]byte(`{"model":"gpt-5"}`)) + require.False(t, view.HasPatches()) + + view.MarkPatchSet("model", "gpt-5.1") + require.True(t, view.HasPatches()) + + view.DisablePatches() + require.False(t, view.HasPatches()) +} + +func TestOpenAIGatewayService_Forward_HTTPPatchPathKeepsLargeInputRaw(t *testing.T) { + gin.SetMode(gin.TestMode) + upstream := &httpUpstreamRecorder{ + resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader( + `{"usage":{"input_tokens":1,"output_tokens":2,"input_tokens_details":{"cached_tokens":0}}}`, + )), + }, + } + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream} + account := &Account{ + ID: 1, + Name: "openai-apikey", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "sk-test", + "base_url": "https://example.com", + }, + Extra: map[string]any{"use_responses_api": true}, + } + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil) + SetOpenAIClientTransport(c, OpenAIClientTransportHTTP) + + body := []byte(`{"model":"gpt-5","stream":false,"reasoning":{"effort":"minimal"},"input":[{"type":"message","content":[{"type":"input_text","text":"hi","nonce":9007199254740993}]}]}`) + result, err := svc.Forward(context.Background(), c, account, body) + require.NoError(t, err) + require.NotNil(t, result) + require.NotNil(t, upstream.lastReq) + require.JSONEq(t, `{"model":"gpt-5","stream":false,"reasoning":{"effort":"none"},"instructions":"You are a helpful coding assistant.","input":[{"type":"message","content":[{"type":"input_text","text":"hi","nonce":9007199254740993}]}]}`, string(upstream.lastBody)) + require.Equal(t, "9007199254740993", gjson.GetBytes(upstream.lastBody, "input.0.content.0.nonce").Raw) +} + +func TestOpenAIGatewayService_Forward_DecodedMutationKeepsLaterFieldDeletes(t *testing.T) { + gin.SetMode(gin.TestMode) + upstream := &httpUpstreamRecorder{ + resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"usage":{"input_tokens":1,"output_tokens":2}}`)), + }, + } + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream} + account := &Account{ + ID: 2, + Name: "openai-apikey", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "sk-test", + "base_url": "https://example.com", + }, + Extra: map[string]any{"use_responses_api": true}, + } + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil) + SetOpenAIClientTransport(c, OpenAIClientTransportHTTP) + + body := []byte(`{"model":"gpt-5.4","stream":false,"max_completion_tokens":12,"tools":[{"type":"image_generation","format":"png"}],"input":[{"type":"message","content":"draw"}]}`) + result, err := svc.Forward(context.Background(), c, account, body) + require.NoError(t, err) + require.NotNil(t, result) + require.False(t, gjson.GetBytes(upstream.lastBody, "max_completion_tokens").Exists()) + require.False(t, gjson.GetBytes(upstream.lastBody, "tools.0.format").Exists()) + require.Equal(t, "png", gjson.GetBytes(upstream.lastBody, "tools.0.output_format").String()) +} + +func TestOpenAIGatewayService_Forward_MappedImageModelUsesImageGate(t *testing.T) { + gin.SetMode(gin.TestMode) + upstream := &httpUpstreamRecorder{ + resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"usage":{"input_tokens":1,"output_tokens":2}}`)), + }, + } + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream} + account := &Account{ + ID: 3, + Name: "openai-apikey", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "sk-test", + "base_url": "https://example.com", + "model_mapping": map[string]any{"draw-alias": "gpt-image-2"}, + }, + Extra: map[string]any{"use_responses_api": true}, + } + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil) + c.Set("api_key", &APIKey{Group: &Group{AllowImageGeneration: false}}) + SetOpenAIClientTransport(c, OpenAIClientTransportHTTP) + + body := []byte(`{"model":"draw-alias","stream":false,"input":"draw"}`) + result, err := svc.Forward(context.Background(), c, account, body) + require.Error(t, err) + require.Nil(t, result) + require.Nil(t, upstream.lastReq) + require.Equal(t, http.StatusForbidden, rec.Code) +} + +func TestOpenAIGatewayService_Forward_TextDataImageDoesNotForceMapMarshal(t *testing.T) { + gin.SetMode(gin.TestMode) + upstream := &httpUpstreamRecorder{ + resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"usage":{"input_tokens":1,"output_tokens":2}}`)), + }, + } + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream} + account := &Account{ + ID: 4, + Name: "openai-apikey", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "sk-test", + "base_url": "https://example.com", + }, + Extra: map[string]any{"use_responses_api": true}, + } + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil) + SetOpenAIClientTransport(c, OpenAIClientTransportHTTP) + + body := []byte(`{"model":"gpt-5","stream":false,"input":[{"type":"message","content":[{"type":"input_text","text":"literal data:image/png;base64, only","nonce":1e1000000}]}]}`) + result, err := svc.Forward(context.Background(), c, account, body) + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, "1e1000000", gjson.GetBytes(upstream.lastBody, "input.0.content.0.nonce").Raw) +} + +func TestOpenAIGatewayService_Forward_ImageToolBillingDoesNotForceFullDecode(t *testing.T) { + gin.SetMode(gin.TestMode) + upstream := &httpUpstreamRecorder{ + resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader( + `{"output":[{"id":"ig_1","type":"image_generation_call","result":"final-image"}],"usage":{"input_tokens":1,"output_tokens":2}}`, + )), + }, + } + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream} + account := &Account{ + ID: 9, + Name: "openai-apikey", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "sk-test", + "base_url": "https://example.com", + }, + Extra: map[string]any{"use_responses_api": true}, + } + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil) + SetOpenAIClientTransport(c, OpenAIClientTransportHTTP) + + body := []byte(`{"model":"gpt-5","stream":false,"tools":[{"type":"image_generation","model":"gpt-image-2","size":"2048x1152"}],"input":[{"type":"message","content":[{"type":"input_text","text":"draw","nonce":1e1000000}]}]}`) + result, err := svc.Forward(context.Background(), c, account, body) + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, "1e1000000", gjson.GetBytes(upstream.lastBody, "input.0.content.0.nonce").Raw) + require.Equal(t, 1, result.ImageCount) + require.Equal(t, "2K", result.ImageSize) + require.Equal(t, "gpt-image-2", result.BillingModel) +} + +func TestOpenAIGatewayService_Forward_HTTPRetryRecoveryDoesNotDecodeBeforeError(t *testing.T) { + gin.SetMode(gin.TestMode) + upstream := &httpUpstreamRecorder{ + responses: []*http.Response{ + { + StatusCode: http.StatusBadRequest, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"error":{"code":"invalid_encrypted_content","type":"invalid_request_error","message":"bad encrypted content"}}`)), + }, + { + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"usage":{"input_tokens":1,"output_tokens":2}}`)), + }, + }, + } + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream} + account := &Account{ + ID: 10, + Name: "openai-apikey", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "sk-test", + "base_url": "https://example.com", + }, + Extra: map[string]any{"use_responses_api": true}, + } + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil) + SetOpenAIClientTransport(c, OpenAIClientTransportHTTP) + + body := []byte(`{"model":"gpt-5","stream":false,"input":[{"type":"reasoning","encrypted_content":"gAAA","summary":[{"type":"summary_text","text":"keep me"}]},{"type":"message","content":[{"type":"input_text","text":"hi","nonce":9007199254740993}]}]}`) + result, err := svc.Forward(context.Background(), c, account, body) + require.NoError(t, err) + require.NotNil(t, result) + require.Len(t, upstream.bodies, 2) + require.Equal(t, "gAAA", gjson.GetBytes(upstream.bodies[0], "input.0.encrypted_content").String()) + require.Equal(t, "9007199254740993", gjson.GetBytes(upstream.bodies[0], "input.1.content.0.nonce").Raw) + require.False(t, gjson.GetBytes(upstream.bodies[1], "input.0.encrypted_content").Exists()) + require.Equal(t, "summary_text", gjson.GetBytes(upstream.bodies[1], "input.0.summary.0.type").String()) +} + +func TestOpenAIGatewayService_Forward_CodexSparkRejectsEscapedInputImage(t *testing.T) { + gin.SetMode(gin.TestMode) + upstream := &httpUpstreamRecorder{ + resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"usage":{"input_tokens":1,"output_tokens":2}}`)), + }, + } + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream} + account := &Account{ + ID: 5, + Name: "openai-apikey", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "sk-test", + "base_url": "https://example.com", + }, + Extra: map[string]any{"use_responses_api": true}, + } + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil) + SetOpenAIClientTransport(c, OpenAIClientTransportHTTP) + + body := []byte(`{"model":"gpt-5.3-codex-spark","stream":false,"input":[{"type":"input_` + "\\u0069" + `mage","file_id":"file_1"}]}`) + result, err := svc.Forward(context.Background(), c, account, body) + require.Error(t, err) + require.Nil(t, result) + require.Nil(t, upstream.lastReq) + require.Equal(t, http.StatusBadRequest, rec.Code) +} + +func TestOpenAIGatewayService_Forward_CodexBridgeInjectionSetsImageBilling(t *testing.T) { + gin.SetMode(gin.TestMode) + upstream := &httpUpstreamRecorder{ + resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader( + `{"output":[{"id":"ig_1","type":"image_generation_call","result":"final-image","size":"1024x1024"}],"usage":{"input_tokens":1,"output_tokens":2}}`, + )), + }, + } + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + cfg.Gateway.ForceCodexCLI = true + cfg.Gateway.CodexImageGenerationBridgeEnabled = true + svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream} + account := &Account{ + ID: 7, + Name: "openai-apikey", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "sk-test", + "base_url": "https://example.com", + }, + Extra: map[string]any{"use_responses_api": true}, + } + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil) + c.Set("api_key", &APIKey{Group: &Group{AllowImageGeneration: true}}) + SetOpenAIClientTransport(c, OpenAIClientTransportHTTP) + + body := []byte(`{"model":"gpt-5","stream":false,"input":"draw if needed"}`) + result, err := svc.Forward(context.Background(), c, account, body) + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, 1, result.ImageCount) + require.Equal(t, "2K", result.ImageSize) + require.Equal(t, "gpt-image-2", result.BillingModel) +} + +func TestOpenAIGatewayService_Forward_HTTPDeletesPreviousResponseIDWhenPresent(t *testing.T) { + gin.SetMode(gin.TestMode) + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + account := &Account{ + ID: 8, + Name: "openai-apikey", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "sk-test", + "base_url": "https://example.com", + }, + Extra: map[string]any{"use_responses_api": true}, + } + + for _, body := range [][]byte{ + []byte(`{"model":"gpt-5","stream":false,"previous_response_id":"","input":"hi"}`), + []byte(`{"model":"gpt-5","stream":false,"previous_response_id":null,"input":"hi"}`), + } { + upstream := &httpUpstreamRecorder{ + resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"usage":{"input_tokens":1,"output_tokens":2}}`)), + }, + } + svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream} + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil) + SetOpenAIClientTransport(c, OpenAIClientTransportHTTP) + + result, err := svc.Forward(context.Background(), c, account, body) + require.NoError(t, err) + require.NotNil(t, result) + require.False(t, gjson.GetBytes(upstream.lastBody, "previous_response_id").Exists()) + } +} + +func TestOpenAIRequestBodyMayContainEmptyBase64InputImageSeesEscapedJSON(t *testing.T) { + body := []byte(`{"input":[{"type":"message","content":[{"type":"input_image","image_` + "\\u0075" + `rl":"data:image/png;base64` + "\\u002c" + ` "}]}]}`) + + require.True(t, openAIRequestBodyMayContainEmptyBase64InputImage(body)) +} + +func TestOpenAIRequestBodyMayContainEmptyBase64InputImageSeesEscapedImageType(t *testing.T) { + body := []byte(`{"input":[{"type":"message","content":[{"type":"input_` + "\\u0069" + `mage","image_url":"data:image/png;base64, "}]}]}`) + + require.True(t, openAIRequestBodyMayContainEmptyBase64InputImage(body)) +} + +func TestOpenAIRequestBodyMayContainEmptyBase64InputImageSeesEscapedInputPrefix(t *testing.T) { + body := []byte(`{"input":[{"type":"message","content":[{"type":"inp` + "\\u0075" + `t_image","image_url":"data:image/png;base64, "}]}]}`) + + require.True(t, openAIRequestBodyMayContainEmptyBase64InputImage(body)) +} + +func TestOpenAIGatewayService_Forward_ImageOnlyModelKeepsSupportedVerbosity(t *testing.T) { + gin.SetMode(gin.TestMode) + upstream := &httpUpstreamRecorder{ + resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"usage":{"input_tokens":1,"output_tokens":2}}`)), + }, + } + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream} + account := &Account{ + ID: 6, + Name: "openai-apikey", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "sk-test", + "base_url": "https://example.com", + }, + Extra: map[string]any{"use_responses_api": true}, + } + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil) + SetOpenAIClientTransport(c, OpenAIClientTransportHTTP) + + body := []byte(`{"model":"gpt-image-2","stream":false,"text":{"verbosity":"low"},"input":"draw"}`) + result, err := svc.Forward(context.Background(), c, account, body) + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, "low", gjson.GetBytes(upstream.lastBody, "text.verbosity").String()) + require.Equal(t, openAIImagesResponsesMainModel, gjson.GetBytes(upstream.lastBody, "model").String()) +} + func TestExtractOpenAIRequestMetaFromBody(t *testing.T) { tests := []struct { name string diff --git a/backend/internal/service/openai_ws_forwarder_success_test.go b/backend/internal/service/openai_ws_forwarder_success_test.go index e949560f82a8..bd2622079333 100644 --- a/backend/internal/service/openai_ws_forwarder_success_test.go +++ b/backend/internal/service/openai_ws_forwarder_success_test.go @@ -171,6 +171,91 @@ func TestOpenAIGatewayService_Forward_WSv2_SuccessAndBindSticky(t *testing.T) { require.Equal(t, "resp_new_1", gjson.GetBytes(responseBody, "id").String()) } +func TestOpenAIGatewayService_Forward_WSv2_UsesPatchedBodyAfterValidationDecode(t *testing.T) { + gin.SetMode(gin.TestMode) + + type receivedPayload struct { + MaxCompletionTokensExists bool + } + receivedCh := make(chan receivedPayload, 1) + + upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }} + wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + t.Errorf("upgrade websocket failed: %v", err) + return + } + defer func() { _ = conn.Close() }() + + var request map[string]any + if err := conn.ReadJSON(&request); err != nil { + t.Errorf("read ws request failed: %v", err) + return + } + requestJSON := requestToJSONString(request) + receivedCh <- receivedPayload{MaxCompletionTokensExists: gjson.Get(requestJSON, "max_completion_tokens").Exists()} + + if err := conn.WriteJSON(map[string]any{ + "type": "response.completed", + "response": map[string]any{ + "id": "resp_patched_ws_1", + "model": "gpt-5.3-codex-spark", + "usage": map[string]any{"input_tokens": 1, "output_tokens": 1}, + }, + }); err != nil { + t.Errorf("write response.completed failed: %v", err) + return + } + })) + defer wsServer.Close() + + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil) + c.Request.Header.Set("User-Agent", "unit-test-agent/1.0") + + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + cfg.Security.URLAllowlist.AllowInsecureHTTP = true + cfg.Gateway.OpenAIWS.Enabled = true + cfg.Gateway.OpenAIWS.APIKeyEnabled = true + cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true + cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3 + cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 30 + cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 10 + + svc := &OpenAIGatewayService{ + cfg: cfg, + openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), + toolCorrector: NewCodexToolCorrector(), + } + + account := &Account{ + ID: 10, + Name: "openai-ws", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "sk-test", + "base_url": wsServer.URL, + }, + Extra: map[string]any{"responses_websockets_v2_enabled": true}, + } + + body := []byte(`{"model":"gpt-5.4","stream":false,"max_completion_tokens":12,"tools":[{"type":"image_generation"}],"input":[{"type":"input_text","text":"hello"}]}`) + result, err := svc.Forward(context.Background(), c, account, body) + require.NoError(t, err) + require.NotNil(t, result) + require.True(t, result.OpenAIWSMode) + + received := <-receivedCh + require.False(t, received.MaxCompletionTokensExists) +} + func TestOpenAIGatewayService_Forward_WSv2_ImageGenerationCountsOutputs(t *testing.T) { gin.SetMode(gin.TestMode) From f10bca81559b18fa003c3b1bcc4fedf148ddbd7f Mon Sep 17 00:00:00 2001 From: visa2 Date: Sun, 31 May 2026 16:07:52 +0800 Subject: [PATCH 29/36] =?UTF-8?q?refactor(apicompat):=20redesign=20the=20C?= =?UTF-8?q?odex=20Responses=20=E2=86=94=20Chat=20Completions=20bridge?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Codex CLI speaks the OpenAI Responses protocol (streaming, store:false), while many upstreams (e.g. DeepSeek in thinking mode) only expose Chat Completions. The bridge that translates between the two had grown field by field and leaned on Go's serialization defaults, which both the Responses client (Codex) and the Chat upstream reject in ways the official OpenAI endpoints tolerate. Problems this fixes (all observed running Codex CLI against a DeepSeek upstream): - Streaming reasoning was never shown in the Codex TUI (the answer appeared with no visible thinking): reasoning deltas were emitted before the reasoning item was opened, so the strict client discarded them. - A tool-using turn could wedge the session into a "no response" state: the function_call stream was never closed (no function_call_arguments.done / output_item.done), so Codex never saw the tool call complete. - Parallel tool calls were rejected upstream (400/502): each function_call became its own assistant message, producing consecutive assistant messages with mismatched tool replies. - A tool turn was rejected with "reasoning_content in the thinking mode must be passed back": the reasoning that produced the tool call was dropped instead of being returned on the assistant message. - Items with no Chat equivalent (web_search_call, ...) and Codex's command-approval notice landed between an assistant tool_calls message and its tool reply, triggering "An assistant message with 'tool_calls' must be followed by tool messages responding to each 'tool_call_id'". - Interrupt/reconnect left an unanswered or dangling tool_call in the history, triggering the same 400. The shared root cause is reliance on serialization defaults — omitempty dropping protocol-required zero values, and unrecognized item types falling through a generic path — rather than deliberately reproducing the target protocol. The bridge is reworked into two explicit layers. Request direction (Responses input -> Chat messages): a parse -> build -> normalize pipeline. - reasoning_content is carried back on the assistant message that produced a tool call (DeepSeek thinking mode requires it to continue the same thought) - consecutive function_call items (parallel tool calls) are merged into a single assistant message's tool_calls array - item types with no Chat equivalent are skipped instead of leaking through a generic path - normalizeChatMessages is the single invariant gate: it guarantees every assistant tool_calls message is immediately followed by one tool reply per tool_call_id — reordering any intervening message (such as a command-approval notice) to after the replies, dropping unanswered tool_calls and orphan tool replies, and preserving bare passthrough tool messages. Response direction (Chat SSE -> Responses SSE): ResponsesStreamEvent.MarshalJSON constructs each streamed event explicitly so protocol-required fields are always present (output_index/content_index/summary_index at 0, message content:[], reasoning summary:[], function_call call_id/name/arguments, output_text part text/annotations/logprobs). This is a single source of truth that removes any post-hoc JSON patching. Reasoning is emitted as its own output item, opened before its deltas, and tool calls are fully closed (function_call_arguments.done + output_item.done with complete arguments). Tests cover request-direction message invariants against golden Codex request shapes (parallel calls, unknown items, intervening messages, partial/dangling calls), per-event wire completeness, and streaming lifecycle ordering. Co-Authored-By: Claude Opus 4.8 --- .../chatcompletions_responses_bridge.go | 476 +++++++++++++++--- ...tions_responses_request_invariants_test.go | 187 +++++++ ...letions_responses_stream_lifecycle_test.go | 103 ++++ .../apicompat/responses_stream_event_wire.go | 199 ++++++++ .../responses_stream_event_wire_test.go | 109 ++++ backend/internal/pkg/apicompat/types.go | 4 + 6 files changed, 1019 insertions(+), 59 deletions(-) create mode 100644 backend/internal/pkg/apicompat/chatcompletions_responses_request_invariants_test.go create mode 100644 backend/internal/pkg/apicompat/chatcompletions_responses_stream_lifecycle_test.go create mode 100644 backend/internal/pkg/apicompat/responses_stream_event_wire.go create mode 100644 backend/internal/pkg/apicompat/responses_stream_event_wire_test.go diff --git a/backend/internal/pkg/apicompat/chatcompletions_responses_bridge.go b/backend/internal/pkg/apicompat/chatcompletions_responses_bridge.go index 09b680c7c73c..cc51cba2d211 100644 --- a/backend/internal/pkg/apicompat/chatcompletions_responses_bridge.go +++ b/backend/internal/pkg/apicompat/chatcompletions_responses_bridge.go @@ -42,14 +42,24 @@ func ResponsesToChatCompletionsRequest(req *ResponsesRequest) (*ChatCompletionsR return out, nil } +// responsesInputToChatMessages converts a Responses request's instructions + +// input[] into Chat Completions messages. It is a three-stage pipeline: +// +// parse — instructions become a system message; input[] is split into items +// build — buildChatMessagesFromItems walks items, attaching reasoning to the +// assistant message that produced a tool call, merging parallel tool +// calls into one assistant message, and skipping item types that have +// no Chat equivalent +// normalize — normalizeChatMessages enforces the invariants DeepSeek requires +// +// The build + normalize split keeps every protocol rule in one place rather than +// scattered across per-item cases, and makes unknown future codex item types +// fail safe instead of leaking into the upstream request. func responsesInputToChatMessages(instructions string, inputRaw json.RawMessage) ([]ChatMessage, error) { var messages []ChatMessage if strings.TrimSpace(instructions) != "" { content, _ := json.Marshal(instructions) - messages = append(messages, ChatMessage{ - Role: "system", - Content: content, - }) + messages = append(messages, ChatMessage{Role: "system", Content: content}) } inputRaw = bytesTrimSpace(inputRaw) @@ -57,13 +67,11 @@ func responsesInputToChatMessages(instructions string, inputRaw json.RawMessage) return messages, nil } + // Bare string input is a single user turn. var inputText string if err := json.Unmarshal(inputRaw, &inputText); err == nil { content, _ := json.Marshal(inputText) - messages = append(messages, ChatMessage{ - Role: "user", - Content: content, - }) + messages = append(messages, ChatMessage{Role: "user", Content: content}) return messages, nil } @@ -72,6 +80,24 @@ func responsesInputToChatMessages(instructions string, inputRaw json.RawMessage) return nil, fmt.Errorf("parse responses input: %w", err) } + built, err := buildChatMessagesFromItems(messages, rawItems) + if err != nil { + return nil, err + } + return normalizeChatMessages(built), nil +} + +// buildChatMessagesFromItems walks the Responses input items and appends the +// corresponding Chat messages. +func buildChatMessagesFromItems(messages []ChatMessage, rawItems []json.RawMessage) ([]ChatMessage, error) { + // pendingReasoning holds the reasoning text from a reasoning item until the + // assistant message it belongs to is emitted. DeepSeek's thinking mode + // requires the reasoning_content that produced a tool call to be passed back + // on that assistant message; dropping it yields a 400. It only survives + // across an assistant message (so a following tool call in the same turn + // still receives it); any other role ends the thinking span. + var pendingReasoning string + for _, raw := range rawItems { raw = bytesTrimSpace(raw) if len(raw) == 0 || string(raw) == "null" { @@ -84,6 +110,7 @@ func responsesInputToChatMessages(instructions string, inputRaw json.RawMessage) if textErr := json.Unmarshal(raw, &text); textErr == nil { content, _ := json.Marshal(text) messages = append(messages, ChatMessage{Role: "user", Content: content}) + pendingReasoning = "" continue } return nil, fmt.Errorf("parse responses input item: %w", err) @@ -92,22 +119,40 @@ func responsesInputToChatMessages(instructions string, inputRaw json.RawMessage) role := chatCompletionsBridgeRole(rawString(item["role"])) itemType := rawString(item["type"]) switch itemType { + case "reasoning": + if txt := extractResponsesReasoningText(item); txt != "" { + pendingReasoning = txt + } + continue case "function_call": arguments := rawString(item["arguments"]) if strings.TrimSpace(arguments) == "" { arguments = "{}" } - messages = append(messages, ChatMessage{ - Role: "assistant", - ToolCalls: []ChatToolCall{{ - ID: rawString(item["call_id"]), - Type: "function", - Function: ChatFunctionCall{ - Name: rawString(item["name"]), - Arguments: arguments, - }, - }}, - }) + toolCall := ChatToolCall{ + ID: rawString(item["call_id"]), + Type: "function", + Function: ChatFunctionCall{ + Name: rawString(item["name"]), + Arguments: arguments, + }, + } + // Parallel tool calls arrive as consecutive function_call items and + // must share one assistant message; the matching tool replies then + // follow it. Merge into the immediately preceding assistant message. + if n := len(messages); n > 0 && messages[n-1].Role == "assistant" { + messages[n-1].ToolCalls = append(messages[n-1].ToolCalls, toolCall) + if messages[n-1].ReasoningContent == "" { + messages[n-1].ReasoningContent = pendingReasoning + } + } else { + messages = append(messages, ChatMessage{ + Role: "assistant", + ToolCalls: []ChatToolCall{toolCall}, + ReasoningContent: pendingReasoning, + }) + } + pendingReasoning = "" continue case "function_call_output": content, _ := json.Marshal(rawString(item["output"])) @@ -116,10 +161,12 @@ func responsesInputToChatMessages(instructions string, inputRaw json.RawMessage) ToolCallID: rawString(item["call_id"]), Content: content, }) + pendingReasoning = "" continue case "input_text", "text": content, _ := json.Marshal(rawString(item["text"])) messages = append(messages, ChatMessage{Role: "user", Content: content}) + pendingReasoning = "" continue case "input_image": content, err := chatContentFromSingleResponsesPart(itemType, item) @@ -127,6 +174,18 @@ func responsesInputToChatMessages(instructions string, inputRaw json.RawMessage) return nil, err } messages = append(messages, ChatMessage{Role: "user", Content: content}) + pendingReasoning = "" + continue + } + + // Only genuine message items become chat messages. Codex emits other + // Responses item types with no Chat equivalent (web_search_call, + // local_shell_call, custom tool calls, file_search_call, ...). Converting + // them via the generic path would insert a spurious message between an + // assistant tool_calls message and its tool reply, which DeepSeek rejects + // ("insufficient tool messages following tool_calls message"). Skip them. + if itemType != "" && itemType != "message" { + pendingReasoning = "" continue } @@ -140,15 +199,128 @@ func responsesInputToChatMessages(instructions string, inputRaw json.RawMessage) if err != nil { return nil, err } - messages = append(messages, ChatMessage{ - Role: role, - Content: chatContent, - }) + messages = append(messages, ChatMessage{Role: role, Content: chatContent}) + // Reasoning only survives across an assistant text message. + if role != "assistant" { + pendingReasoning = "" + } } return messages, nil } +// normalizeChatMessages is the single place that enforces the tool-call +// invariant the DeepSeek / OpenAI Chat Completions schema requires: an assistant +// message with tool_calls must be immediately followed by one tool message per +// tool_call_id, in order, with nothing in between. +// +// Codex histories violate this in several ways that the builder alone can't fix: +// - a non-tool message lands between an assistant tool_calls message and its +// tool replies (e.g. an "Approved command prefix saved" system notice codex +// injects mid tool-execution); +// - a parallel tool_call's sibling output never arrives, or a call is left +// dangling by a mid-execution reconnect (unanswered tool_call); +// - a tool reply has no announcing assistant tool_call (orphan). +// +// It rebuilds the sequence so each assistant's answered tool_calls are followed +// directly by their replies (in call order); unanswered tool_calls are dropped +// (and an assistant left with neither tool_calls nor content is dropped); orphan +// tool replies and intervening messages are emitted in their natural position +// but never between an assistant tool_calls message and its replies. +func normalizeChatMessages(messages []ChatMessage) []ChatMessage { + // Index every tool reply by its tool_call_id (last wins on duplicates). + replies := make(map[string]ChatMessage) + for _, m := range messages { + if m.Role == "tool" && m.ToolCallID != "" { + replies[m.ToolCallID] = m + } + } + + out := make([]ChatMessage, 0, len(messages)) + for _, m := range messages { + switch { + case m.Role == "tool": + // A bare tool message with no tool_call_id is a direct Chat + // Completions passthrough; keep it in place. A tool reply whose id is + // announced by an assistant is emitted right after that assistant + // (skip the standalone occurrence). Any other tool reply is an orphan + // and is dropped. + if m.ToolCallID == "" { + out = append(out, m) + } + continue + case len(m.ToolCalls) > 0: + kept := make([]ChatToolCall, 0, len(m.ToolCalls)) + for _, tc := range m.ToolCalls { + if tc.ID == "" { + continue + } + if _, ok := replies[tc.ID]; ok { + kept = append(kept, tc) + } + } + if len(kept) == 0 { + // No answered tool_calls left: keep as a plain message if it has + // content, otherwise drop it entirely. + if isBlankChatContent(m.Content) { + continue + } + m.ToolCalls = nil + out = append(out, m) + continue + } + m.ToolCalls = kept + out = append(out, m) + for _, tc := range kept { + out = append(out, replies[tc.ID]) + } + default: + out = append(out, m) + } + } + return out +} + +// isBlankChatContent reports whether a chat message content holds no usable text. +func isBlankChatContent(raw json.RawMessage) bool { + raw = bytesTrimSpace(raw) + if len(raw) == 0 || string(raw) == "null" || string(raw) == `""` { + return true + } + return chatMessageContentText(raw) == "" +} + +// extractResponsesReasoningText pulls the reasoning text out of a Responses +// reasoning item. The Chat→Responses bridge writes the upstream reasoning_content +// verbatim into the summary_text parts (see closeChatReasoningItem), so codex +// round-trips it there; prefer summary[].text and fall back to content. +func extractResponsesReasoningText(item map[string]json.RawMessage) string { + var parts []string + collect := func(raw json.RawMessage) { + raw = bytesTrimSpace(raw) + if len(raw) == 0 || string(raw) == "null" { + return + } + var arr []map[string]json.RawMessage + if err := json.Unmarshal(raw, &arr); err == nil { + for _, p := range arr { + if t := rawString(p["text"]); t != "" { + parts = append(parts, t) + } + } + return + } + if t := rawString(raw); t != "" { + parts = append(parts, t) + } + } + collect(item["summary"]) + if len(parts) == 0 { + collect(item["content"]) + } + return strings.Join(parts, "\n") +} + func chatCompletionsBridgeRole(role string) string { trimmed := strings.TrimSpace(role) if trimmed == "" { @@ -448,10 +620,32 @@ type ChatCompletionsToResponsesStreamState struct { CreatedSent bool CompletedSent bool + // nextOutputIndex assigns sequential output_index values to items as they + // are opened (reasoning, message, tool calls), so the streamed indices match + // the order of items in the final response.output array. + nextOutputIndex int + + // Reasoning item lifecycle. DeepSeek-style upstreams stream all + // reasoning_content before any content, so reasoning is modeled as its own + // "reasoning" output item that must be opened (output_item.added) before any + // reasoning delta and closed before the message/tool items open. + ReasoningItemID string + ReasoningIndex int + ReasoningOpen bool + ReasoningDone bool + + // Message item + output_text content-part lifecycle. MessageItemID string - Text strings.Builder - Reasoning strings.Builder - ToolCalls map[int]*ChatToolCall + MessageIndex int + TextPartOpen bool + + Text strings.Builder + Reasoning strings.Builder + + // Tool-call lifecycle, keyed by the upstream tool_call index. + ToolCalls map[int]*ChatToolCall + ToolItemIDs map[int]string + ToolOutputIndex map[int]int FinishReason string Usage *ResponsesUsage @@ -460,13 +654,21 @@ type ChatCompletionsToResponsesStreamState struct { // NewChatCompletionsToResponsesStreamState returns an initialized stream state. func NewChatCompletionsToResponsesStreamState(model string) *ChatCompletionsToResponsesStreamState { return &ChatCompletionsToResponsesStreamState{ - ResponseID: generateResponsesID(), - Model: model, - Created: time.Now().Unix(), - ToolCalls: make(map[int]*ChatToolCall), + ResponseID: generateResponsesID(), + Model: model, + Created: time.Now().Unix(), + ToolCalls: make(map[int]*ChatToolCall), + ToolItemIDs: make(map[int]string), + ToolOutputIndex: make(map[int]int), } } +func (state *ChatCompletionsToResponsesStreamState) allocOutputIndex() int { + idx := state.nextOutputIndex + state.nextOutputIndex++ + return idx +} + // ChatCompletionsChunkToResponsesEvents converts one Chat Completions stream // chunk into zero or more Responses stream events. func ChatCompletionsChunkToResponsesEvents( @@ -490,24 +692,34 @@ func ChatCompletionsChunkToResponsesEvents( events = append(events, ensureChatToResponsesCreated(state)...) for _, choice := range chunk.Choices { - if choice.Delta.Content != nil { + // Reasoning is emitted as its own output item and must be opened + // (output_item.added + reasoning_summary_part.added) before the first + // delta, otherwise a strict client discards the delta. The leading + // empty-string reasoning delta upstreams send is filtered out. + if choice.Delta.ReasoningContent != nil && *choice.Delta.ReasoningContent != "" { + events = append(events, ensureChatReasoningItem(state)...) + _, _ = state.Reasoning.WriteString(*choice.Delta.ReasoningContent) + events = append(events, chatToResponsesEvent(state, "response.reasoning_summary_text.delta", &ResponsesStreamEvent{ + OutputIndex: state.ReasoningIndex, + SummaryIndex: 0, + Delta: *choice.Delta.ReasoningContent, + ItemID: state.ReasoningItemID, + })) + } + if choice.Delta.Content != nil && *choice.Delta.Content != "" { + // First real content closes the reasoning item, then opens the + // message item and its output_text content part. + events = append(events, closeChatReasoningItem(state)...) events = append(events, ensureChatToResponsesMessageItem(state)...) + events = append(events, ensureChatToResponsesTextPart(state)...) _, _ = state.Text.WriteString(*choice.Delta.Content) events = append(events, chatToResponsesEvent(state, "response.output_text.delta", &ResponsesStreamEvent{ - OutputIndex: 0, + OutputIndex: state.MessageIndex, ContentIndex: 0, Delta: *choice.Delta.Content, ItemID: state.MessageItemID, })) } - if choice.Delta.ReasoningContent != nil { - _, _ = state.Reasoning.WriteString(*choice.Delta.ReasoningContent) - events = append(events, chatToResponsesEvent(state, "response.reasoning_summary_text.delta", &ResponsesStreamEvent{ - OutputIndex: 0, - SummaryIndex: 0, - Delta: *choice.Delta.ReasoningContent, - })) - } for _, toolCall := range choice.Delta.ToolCalls { idx := 0 if toolCall.Index != nil { @@ -515,6 +727,8 @@ func ChatCompletionsChunkToResponsesEvents( } stored, ok := state.ToolCalls[idx] if !ok { + // A tool call closes any open reasoning item first. + events = append(events, closeChatReasoningItem(state)...) copyCall := toolCall if copyCall.ID == "" { copyCall.ID = generateItemID() @@ -522,11 +736,14 @@ func ChatCompletionsChunkToResponsesEvents( copyCall.Type = "function" state.ToolCalls[idx] = ©Call stored = ©Call + itemID := generateItemID() + state.ToolItemIDs[idx] = itemID + state.ToolOutputIndex[idx] = state.allocOutputIndex() events = append(events, chatToResponsesEvent(state, "response.output_item.added", &ResponsesStreamEvent{ - OutputIndex: idx + 1, + OutputIndex: state.ToolOutputIndex[idx], Item: &ResponsesOutput{ Type: "function_call", - ID: generateItemID(), + ID: itemID, CallID: stored.ID, Name: stored.Function.Name, Status: "in_progress", @@ -543,7 +760,8 @@ func ChatCompletionsChunkToResponsesEvents( if toolCall.Function.Arguments != "" { stored.Function.Arguments += toolCall.Function.Arguments events = append(events, chatToResponsesEvent(state, "response.function_call_arguments.delta", &ResponsesStreamEvent{ - OutputIndex: idx + 1, + OutputIndex: state.ToolOutputIndex[idx], + ItemID: state.ToolItemIDs[idx], Delta: toolCall.Function.Arguments, CallID: stored.ID, Name: stored.Function.Name, @@ -565,24 +783,44 @@ func FinalizeChatCompletionsResponsesStream(state *ChatCompletionsToResponsesStr } var events []ResponsesStreamEvent events = append(events, ensureChatToResponsesCreated(state)...) + + // Close a reasoning item that never transitioned to content (reasoning-only + // or empty completion). + events = append(events, closeChatReasoningItem(state)...) + if state.MessageItemID != "" { - events = append(events, chatToResponsesEvent(state, "response.output_text.done", &ResponsesStreamEvent{ - OutputIndex: 0, - ContentIndex: 0, - Text: state.Text.String(), - ItemID: state.MessageItemID, - })) + if state.TextPartOpen { + events = append(events, chatToResponsesEvent(state, "response.output_text.done", &ResponsesStreamEvent{ + OutputIndex: state.MessageIndex, + ContentIndex: 0, + Text: state.Text.String(), + ItemID: state.MessageItemID, + })) + events = append(events, chatToResponsesEvent(state, "response.content_part.done", &ResponsesStreamEvent{ + OutputIndex: state.MessageIndex, + ContentIndex: 0, + ItemID: state.MessageItemID, + Part: &ResponsesContentPart{Type: "output_text", Text: state.Text.String()}, + })) + } events = append(events, chatToResponsesEvent(state, "response.output_item.done", &ResponsesStreamEvent{ - OutputIndex: 0, + OutputIndex: state.MessageIndex, Item: &ResponsesOutput{ - Type: "message", - ID: state.MessageItemID, - Role: "assistant", - Status: "completed", + Type: "message", + ID: state.MessageItemID, + Role: "assistant", + Content: []ResponsesContentPart{{Type: "output_text", Text: state.Text.String()}}, + Status: "completed", }, })) } + // Close every function_call item opened during the stream. Codex finalizes a + // tool call only after function_call_arguments.done + output_item.done for + // that item; without them the call never completes and the session wedges. + // Mirrors cc-switch's finalize_tools. + events = append(events, closeChatToolItems(state)...) + status := "completed" var incompleteDetails *ResponsesIncompleteDetails if state.FinishReason == "length" { @@ -621,22 +859,142 @@ func ensureChatToResponsesCreated(state *ChatCompletionsToResponsesStreamState) })} } +// ensureChatReasoningItem opens the reasoning output item (output_item.added + +// reasoning_summary_part.added) before the first reasoning delta. Codex renders +// streaming reasoning only when this summary-part lifecycle is present. +func ensureChatReasoningItem(state *ChatCompletionsToResponsesStreamState) []ResponsesStreamEvent { + if state.ReasoningOpen || state.ReasoningDone { + return nil + } + state.ReasoningOpen = true + state.ReasoningItemID = generateItemID() + state.ReasoningIndex = state.allocOutputIndex() + return []ResponsesStreamEvent{ + chatToResponsesEvent(state, "response.output_item.added", &ResponsesStreamEvent{ + OutputIndex: state.ReasoningIndex, + Item: &ResponsesOutput{Type: "reasoning", ID: state.ReasoningItemID, Status: "in_progress"}, + }), + chatToResponsesEvent(state, "response.reasoning_summary_part.added", &ResponsesStreamEvent{ + OutputIndex: state.ReasoningIndex, + SummaryIndex: 0, + ItemID: state.ReasoningItemID, + Part: &ResponsesContentPart{Type: "summary_text"}, + }), + } +} + +// closeChatReasoningItem emits the reasoning item's terminal events +// (reasoning_summary_text.done + reasoning_summary_part.done + output_item.done). +func closeChatReasoningItem(state *ChatCompletionsToResponsesStreamState) []ResponsesStreamEvent { + if !state.ReasoningOpen { + return nil + } + state.ReasoningOpen = false + state.ReasoningDone = true + reasoning := state.Reasoning.String() + return []ResponsesStreamEvent{ + chatToResponsesEvent(state, "response.reasoning_summary_text.done", &ResponsesStreamEvent{ + OutputIndex: state.ReasoningIndex, + SummaryIndex: 0, + Text: reasoning, + ItemID: state.ReasoningItemID, + }), + chatToResponsesEvent(state, "response.reasoning_summary_part.done", &ResponsesStreamEvent{ + OutputIndex: state.ReasoningIndex, + SummaryIndex: 0, + ItemID: state.ReasoningItemID, + Part: &ResponsesContentPart{Type: "summary_text", Text: reasoning}, + }), + chatToResponsesEvent(state, "response.output_item.done", &ResponsesStreamEvent{ + OutputIndex: state.ReasoningIndex, + Item: &ResponsesOutput{ + Type: "reasoning", + ID: state.ReasoningItemID, + Status: "completed", + Summary: []ResponsesSummary{{Type: "summary_text", Text: reasoning}}, + }, + }), + } +} + func ensureChatToResponsesMessageItem(state *ChatCompletionsToResponsesStreamState) []ResponsesStreamEvent { if state.MessageItemID != "" { return nil } state.MessageItemID = generateItemID() + state.MessageIndex = state.allocOutputIndex() return []ResponsesStreamEvent{chatToResponsesEvent(state, "response.output_item.added", &ResponsesStreamEvent{ - OutputIndex: 0, + OutputIndex: state.MessageIndex, Item: &ResponsesOutput{ - Type: "message", - ID: state.MessageItemID, - Role: "assistant", - Status: "in_progress", + Type: "message", + ID: state.MessageItemID, + Role: "assistant", + Status: "in_progress", + Content: []ResponsesContentPart{{Type: "output_text"}}, }, })} } +func ensureChatToResponsesTextPart(state *ChatCompletionsToResponsesStreamState) []ResponsesStreamEvent { + if state.TextPartOpen { + return nil + } + state.TextPartOpen = true + return []ResponsesStreamEvent{chatToResponsesEvent(state, "response.content_part.added", &ResponsesStreamEvent{ + OutputIndex: state.MessageIndex, + ContentIndex: 0, + ItemID: state.MessageItemID, + Part: &ResponsesContentPart{Type: "output_text", Text: ""}, + })} +} + +// closeChatToolItems emits function_call_arguments.done + output_item.done for +// every tool call opened during the stream, carrying the full call_id/name/ +// arguments so codex can deserialize and execute the call. Mirrors cc-switch's +// finalize_tools. +func closeChatToolItems(state *ChatCompletionsToResponsesStreamState) []ResponsesStreamEvent { + if len(state.ToolCalls) == 0 { + return nil + } + var events []ResponsesStreamEvent + for i := 0; i < len(state.ToolCalls); i++ { + toolCall, ok := state.ToolCalls[i] + if !ok || toolCall == nil { + continue + } + itemID, opened := state.ToolItemIDs[i] + if !opened { + continue + } + arguments := toolCall.Function.Arguments + if strings.TrimSpace(arguments) == "" { + arguments = "{}" + } + outputIndex := state.ToolOutputIndex[i] + events = append(events, + chatToResponsesEvent(state, "response.function_call_arguments.done", &ResponsesStreamEvent{ + OutputIndex: outputIndex, + ItemID: itemID, + CallID: toolCall.ID, + Name: toolCall.Function.Name, + Arguments: arguments, + }), + chatToResponsesEvent(state, "response.output_item.done", &ResponsesStreamEvent{ + OutputIndex: outputIndex, + Item: &ResponsesOutput{ + Type: "function_call", + ID: itemID, + CallID: toolCall.ID, + Name: toolCall.Function.Name, + Arguments: arguments, + Status: "completed", + }, + }), + ) + } + return events +} + func (state *ChatCompletionsToResponsesStreamState) chatOutput() []ResponsesOutput { var outputs []ResponsesOutput if state.Reasoning.Len() > 0 { diff --git a/backend/internal/pkg/apicompat/chatcompletions_responses_request_invariants_test.go b/backend/internal/pkg/apicompat/chatcompletions_responses_request_invariants_test.go new file mode 100644 index 000000000000..e54a453279d0 --- /dev/null +++ b/backend/internal/pkg/apicompat/chatcompletions_responses_request_invariants_test.go @@ -0,0 +1,187 @@ +package apicompat + +import ( + "encoding/json" + "testing" + + "github.com/stretchr/testify/require" +) + +// assertChatInvariants enforces the DeepSeek / OpenAI Chat Completions message +// invariants that, when violated, surface as upstream 400s. Used to validate the +// request-direction converter against golden codex request shapes. +func assertChatInvariants(t *testing.T, messages []ChatMessage) { + t.Helper() + for i, m := range messages { + // Every assistant tool_calls message must be immediately followed by one + // tool message per tool_call_id, in order. + if len(m.ToolCalls) > 0 { + for j, tc := range m.ToolCalls { + k := i + 1 + j + require.Lessf(t, k, len(messages), "tool_call %s has no following tool message", tc.ID) + require.Equalf(t, "tool", messages[k].Role, "tool_call %s not followed by a tool message", tc.ID) + require.Equalf(t, tc.ID, messages[k].ToolCallID, "tool reply order mismatch for %s", tc.ID) + } + } + // No two consecutive assistant messages. + if i > 0 && m.Role == "assistant" && messages[i-1].Role == "assistant" { + t.Fatalf("consecutive assistant messages at %d", i) + } + // No orphan tool replies. + if m.Role == "tool" { + require.NotEmptyf(t, m.ToolCallID, "tool message without tool_call_id at %d", i) + } + } +} + +func convertGolden(t *testing.T, input string) []ChatMessage { + t.Helper() + msgs, err := responsesInputToChatMessages("You are a helpful assistant.", json.RawMessage(input)) + require.NoError(t, err) + return msgs +} + +// Golden sample: a single tool-call turn (codex runs one shell/curl command), +// the shape that produced the original "no response" / 400. +func TestGolden_SingleToolCall(t *testing.T) { + msgs := convertGolden(t, `[ + {"type":"message","role":"user","content":[{"type":"input_text","text":"latest sha?"}]}, + {"type":"reasoning","summary":[{"type":"summary_text","text":"need to run curl"}]}, + {"type":"function_call","call_id":"call_a","name":"exec_command","arguments":"{\"cmd\":\"curl x\"}"}, + {"type":"function_call_output","call_id":"call_a","output":"deadbeef"} + ]`) + assertChatInvariants(t, msgs) + // reasoning_content must ride on the assistant tool-call message. + var asst *ChatMessage + for i := range msgs { + if len(msgs[i].ToolCalls) > 0 { + asst = &msgs[i] + } + } + require.NotNil(t, asst) + require.Equal(t, "need to run curl", asst.ReasoningContent) +} + +// Golden sample: parallel tool calls (codex runs git log + git tag at once). +func TestGolden_ParallelToolCalls(t *testing.T) { + msgs := convertGolden(t, `[ + {"type":"message","role":"user","content":[{"type":"input_text","text":"features?"}]}, + {"type":"reasoning","summary":[{"type":"summary_text","text":"inspect repo"}]}, + {"type":"function_call","call_id":"c0","name":"exec_command","arguments":"{\"cmd\":\"git log\"}"}, + {"type":"function_call","call_id":"c1","name":"exec_command","arguments":"{\"cmd\":\"git tag\"}"}, + {"type":"function_call_output","call_id":"c0","output":"log"}, + {"type":"function_call_output","call_id":"c1","output":"tags"} + ]`) + assertChatInvariants(t, msgs) + // Both parallel calls share ONE assistant message. + var toolMsgs int + for _, m := range msgs { + if len(m.ToolCalls) == 2 { + require.Equal(t, "c0", m.ToolCalls[0].ID) + require.Equal(t, "c1", m.ToolCalls[1].ID) + } + if m.Role == "tool" { + toolMsgs++ + } + } + require.Equal(t, 2, toolMsgs) +} + +// Golden sample: an unknown item type (web_search_call from a 联网查询) sitting +// between a function_call and its output must not break tool↔reply adjacency. +func TestGolden_UnknownItemBetweenToolCallAndOutput(t *testing.T) { + msgs := convertGolden(t, `[ + {"type":"message","role":"user","content":[{"type":"input_text","text":"search"}]}, + {"type":"reasoning","summary":[{"type":"summary_text","text":"let me search"}]}, + {"type":"function_call","call_id":"c0","name":"exec_command","arguments":"{}"}, + {"type":"web_search_call","id":"ws_1","status":"completed","action":{"type":"search","query":"x"}}, + {"type":"function_call_output","call_id":"c0","output":"result"} + ]`) + assertChatInvariants(t, msgs) +} + +// Sequential tool calls (a tool reply between two calls) must stay in distinct +// assistant messages. +func TestRequest_SequentialToolCallsStaySeparate(t *testing.T) { + msgs := convertGolden(t, `[ + {"type":"function_call","call_id":"c1","name":"exec","arguments":"{}"}, + {"type":"function_call_output","call_id":"c1","output":"r1"}, + {"type":"function_call","call_id":"c2","name":"exec","arguments":"{}"}, + {"type":"function_call_output","call_id":"c2","output":"r2"} + ]`) + assertChatInvariants(t, msgs) + assistants := 0 + for _, m := range msgs { + if len(m.ToolCalls) == 1 { + assistants++ + } + } + require.Equal(t, 2, assistants) +} + +// Golden sample: codex injects a message (e.g. an "Approved command prefix +// saved" notice) between a function_call and its output. The intervening message +// must be moved after the tool reply so the assistant tool_calls is immediately +// followed by its reply. +func TestGolden_MessageBetweenToolCallAndOutput(t *testing.T) { + msgs := convertGolden(t, `[ + {"type":"message","role":"user","content":[{"type":"input_text","text":"do it"}]}, + {"type":"reasoning","summary":[{"type":"summary_text","text":"run cmd"}]}, + {"type":"function_call","call_id":"A","name":"exec","arguments":"{}"}, + {"type":"message","role":"developer","content":[{"type":"input_text","text":"Approved command prefix saved"}]}, + {"type":"function_call_output","call_id":"A","output":"ok"} + ]`) + assertChatInvariants(t, msgs) + // The assistant tool_calls message is immediately followed by its tool reply. + for i, m := range msgs { + if len(m.ToolCalls) > 0 { + require.Equal(t, "tool", msgs[i+1].Role) + require.Equal(t, "A", msgs[i+1].ToolCallID) + } + } +} + +// Golden sample: a parallel tool call where one sibling's output is missing +// (codex interrupted/reconnected mid-execution). The unanswered tool_call must +// be dropped so the remaining assistant tool_calls are all answered. +func TestGolden_PartialParallelDropsUnansweredCall(t *testing.T) { + msgs := convertGolden(t, `[ + {"type":"message","role":"user","content":[{"type":"input_text","text":"q"}]}, + {"type":"reasoning","summary":[{"type":"summary_text","text":"r"}]}, + {"type":"function_call","call_id":"A","name":"exec","arguments":"{}"}, + {"type":"function_call","call_id":"B","name":"exec","arguments":"{}"}, + {"type":"function_call_output","call_id":"A","output":"oa"} + ]`) + assertChatInvariants(t, msgs) + for _, m := range msgs { + for _, tc := range m.ToolCalls { + require.NotEqual(t, "B", tc.ID, "unanswered tool_call B should have been dropped") + } + } +} + +// Golden sample: a dangling tool_call at the end of the history (no output yet). +// The assistant message holding only that call must be dropped entirely. +func TestGolden_DanglingToolCallDropped(t *testing.T) { + msgs := convertGolden(t, `[ + {"type":"message","role":"user","content":[{"type":"input_text","text":"q"}]}, + {"type":"reasoning","summary":[{"type":"summary_text","text":"r"}]}, + {"type":"function_call","call_id":"A","name":"exec","arguments":"{}"} + ]`) + assertChatInvariants(t, msgs) + for _, m := range msgs { + require.Empty(t, m.ToolCalls, "dangling unanswered tool_call should have been dropped") + } +} + +// normalizeChatMessages drops an orphan tool reply whose tool_call was never +// announced. +func TestNormalize_DropsOrphanToolReply(t *testing.T) { + msgs := convertGolden(t, `[ + {"type":"message","role":"user","content":[{"type":"input_text","text":"q"}]}, + {"type":"function_call_output","call_id":"ghost","output":"orphan"} + ]`) + for _, m := range msgs { + require.NotEqualf(t, "tool", m.Role, "orphan tool reply should have been dropped") + } +} diff --git a/backend/internal/pkg/apicompat/chatcompletions_responses_stream_lifecycle_test.go b/backend/internal/pkg/apicompat/chatcompletions_responses_stream_lifecycle_test.go new file mode 100644 index 000000000000..beb473038bc2 --- /dev/null +++ b/backend/internal/pkg/apicompat/chatcompletions_responses_stream_lifecycle_test.go @@ -0,0 +1,103 @@ +package apicompat + +import ( + "encoding/json" + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +func collectStreamEvents(t *testing.T, chunks []string) []ResponsesStreamEvent { + t.Helper() + state := NewChatCompletionsToResponsesStreamState("deepseek-v4-pro") + var events []ResponsesStreamEvent + for _, payload := range chunks { + var chunk ChatCompletionsChunk + require.NoError(t, json.Unmarshal([]byte(payload), &chunk)) + events = append(events, ChatCompletionsChunkToResponsesEvents(&chunk, state)...) + } + events = append(events, FinalizeChatCompletionsResponsesStream(state)...) + return events +} + +// TestStream_ReasoningOpensItemBeforeDelta guards the bug where a strict client +// (Codex) drops reasoning deltas that reference an item not yet opened. +func TestStream_ReasoningOpensItemBeforeDelta(t *testing.T) { + events := collectStreamEvents(t, []string{ + `{"choices":[{"index":0,"delta":{"role":"assistant","content":null,"reasoning_content":""}}]}`, + `{"choices":[{"index":0,"delta":{"reasoning_content":"think"}}]}`, + `{"choices":[{"index":0,"delta":{"content":"hello"}}]}`, + `{"choices":[{"index":0,"delta":{"content":""},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":2,"total_tokens":3}}`, + }) + + open := map[int]string{} // output_index -> item type + for _, e := range events { + switch e.Type { + case "response.output_item.added": + require.NotNil(t, e.Item) + open[e.OutputIndex] = e.Item.Type + case "response.reasoning_summary_text.delta": + require.Equalf(t, "reasoning", open[e.OutputIndex], "reasoning delta before its item was opened") + case "response.output_text.delta": + require.Equalf(t, "message", open[e.OutputIndex], "text delta before its item was opened") + } + } +} + +// TestStream_ToolCallLifecycleComplete guards that a tool call is fully closed +// (function_call_arguments.done + output_item.done with full arguments), which +// codex needs to execute the call. +func TestStream_ToolCallLifecycleComplete(t *testing.T) { + events := collectStreamEvents(t, []string{ + `{"choices":[{"index":0,"delta":{"role":"assistant","reasoning_content":"plan"}}]}`, + `{"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"call_a","type":"function","function":{"name":"exec","arguments":""}}]}}]}`, + `{"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\"cmd\":\"ls\"}"}}]}}]}`, + `{"choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":1,"completion_tokens":2,"total_tokens":3}}`, + }) + + var sawAdded, sawArgsDone, sawItemDone bool + for _, e := range events { + switch e.Type { + case "response.output_item.added": + if e.Item != nil && e.Item.Type == "function_call" { + sawAdded = true + } + case "response.function_call_arguments.done": + sawArgsDone = true + require.Equal(t, `{"cmd":"ls"}`, e.Arguments) + case "response.output_item.done": + if e.Item != nil && e.Item.Type == "function_call" { + sawItemDone = true + require.Equal(t, `{"cmd":"ls"}`, e.Item.Arguments) + require.Equal(t, "call_a", e.Item.CallID) + } + } + } + require.True(t, sawAdded, "function_call output_item.added missing") + require.True(t, sawArgsDone, "function_call_arguments.done missing") + require.True(t, sawItemDone, "function_call output_item.done missing") +} + +// TestStream_SSEWireComplete drives the full stream through SSE encoding and +// asserts the function_call events carry complete fields on the wire. +func TestStream_SSEWireComplete(t *testing.T) { + events := collectStreamEvents(t, []string{ + `{"choices":[{"index":0,"delta":{"role":"assistant","reasoning_content":"plan"}}]}`, + `{"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"call_a","type":"function","function":{"name":"exec","arguments":"{}"}}]}}]}`, + `{"choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}`, + }) + + var addedLine string + for _, e := range events { + sse, err := ResponsesEventToSSE(e) + require.NoError(t, err) + if e.Type == "response.output_item.added" && e.Item != nil && e.Item.Type == "function_call" { + addedLine = sse + } + } + require.NotEmpty(t, addedLine) + // The function_call added event must carry arguments:"" on the wire. + require.True(t, strings.Contains(addedLine, `"arguments":""`), "added line missing arguments: %s", addedLine) + require.Contains(t, addedLine, `"call_id":"call_a"`) +} diff --git a/backend/internal/pkg/apicompat/responses_stream_event_wire.go b/backend/internal/pkg/apicompat/responses_stream_event_wire.go new file mode 100644 index 000000000000..df7a82e39374 --- /dev/null +++ b/backend/internal/pkg/apicompat/responses_stream_event_wire.go @@ -0,0 +1,199 @@ +package apicompat + +import "encoding/json" + +// MarshalJSON renders a ResponsesStreamEvent into its wire form. +// +// The OpenAI Responses streaming protocol requires several fields to be present +// even when they hold a zero value: output_index/content_index/summary_index are +// meaningful at 0, a function_call item must always carry call_id/name/arguments +// (arguments may be ""), a message item must carry content:[] and an output_text +// part must carry text/annotations/logprobs. Go's `omitempty` drops exactly those +// zero values, and strict clients (Codex CLI) reject items/deltas whose required +// fields are missing. +// +// Rather than marshalling with omitempty and patching the JSON afterwards, every +// streamed event type is constructed explicitly here — the Go analogue of the +// reference gateways' (cc-switch, CCX) per-event object construction. This is the +// single source of truth for Responses SSE field presence and applies uniformly +// to every emitter (Chat→Responses bridge and Anthropic→Responses converter). +// +// Event types not listed fall back to the default struct marshalling, which +// bounds the blast radius of this method to the streamed item/part/text/tool +// events. +func (e ResponsesStreamEvent) MarshalJSON() ([]byte, error) { + switch e.Type { + case "response.output_text.delta", "response.output_text.done": + m := e.wireBase() + e.putItemID(m) + m["output_index"] = e.OutputIndex + m["content_index"] = e.ContentIndex + if e.Type == "response.output_text.done" { + m["text"] = e.Text + } else { + m["delta"] = e.Delta + } + return json.Marshal(m) + + case "response.content_part.added", "response.content_part.done": + m := e.wireBase() + e.putItemID(m) + m["output_index"] = e.OutputIndex + m["content_index"] = e.ContentIndex + m["part"] = outputTextPartWire(e.Part) + return json.Marshal(m) + + case "response.reasoning_summary_text.delta", "response.reasoning_summary_text.done": + m := e.wireBase() + e.putItemID(m) + m["output_index"] = e.OutputIndex + m["summary_index"] = e.SummaryIndex + if e.Type == "response.reasoning_summary_text.done" { + m["text"] = e.Text + } else { + m["delta"] = e.Delta + } + return json.Marshal(m) + + case "response.reasoning_summary_part.added", "response.reasoning_summary_part.done": + m := e.wireBase() + e.putItemID(m) + m["output_index"] = e.OutputIndex + m["summary_index"] = e.SummaryIndex + m["part"] = summaryTextPartWire(e.Part) + return json.Marshal(m) + + case "response.output_item.added", "response.output_item.done": + m := e.wireBase() + m["output_index"] = e.OutputIndex + m["item"] = responsesItemWire(e.Item) + return json.Marshal(m) + + case "response.function_call_arguments.delta", "response.function_call_arguments.done": + m := e.wireBase() + e.putItemID(m) + m["output_index"] = e.OutputIndex + if e.CallID != "" { + m["call_id"] = e.CallID + } + if e.Name != "" { + m["name"] = e.Name + } + if e.Type == "response.function_call_arguments.done" { + m["arguments"] = e.Arguments + } else { + m["delta"] = e.Delta + } + return json.Marshal(m) + + default: + // response.created / completed / done / failed / incomplete and any + // event type not shaped above keep the default struct marshalling. + type alias ResponsesStreamEvent + return json.Marshal(alias(e)) + } +} + +func (e ResponsesStreamEvent) wireBase() map[string]any { + m := map[string]any{ + "type": e.Type, + "sequence_number": e.SequenceNumber, + } + return m +} + +func (e ResponsesStreamEvent) putItemID(m map[string]any) { + if e.ItemID != "" { + m["item_id"] = e.ItemID + } +} + +// outputTextPartWire renders a content part for a message's output_text, always +// carrying text/annotations/logprobs (matching cc-switch's push_text_delta). +func outputTextPartWire(part *ResponsesContentPart) map[string]any { + text := "" + if part != nil { + text = part.Text + } + return map[string]any{ + "type": "output_text", + "text": text, + "annotations": []any{}, + "logprobs": []any{}, + } +} + +// summaryTextPartWire renders a reasoning summary part. +func summaryTextPartWire(part *ResponsesContentPart) map[string]any { + text := "" + if part != nil { + text = part.Text + } + return map[string]any{ + "type": "summary_text", + "text": text, + } +} + +// responsesItemWire renders an output_item with every field the item's type +// requires to be present, including the empty arrays/strings that omitempty +// would otherwise drop. Mirrors cc-switch's response_function_call_item and the +// message/reasoning item shapes codex expects. +func responsesItemWire(item *ResponsesOutput) map[string]any { + if item == nil { + return map[string]any{} + } + m := map[string]any{ + "type": item.Type, + "id": item.ID, + } + if item.Status != "" { + m["status"] = item.Status + } + switch item.Type { + case "message": + role := item.Role + if role == "" { + role = "assistant" + } + m["role"] = role + m["content"] = messageContentWire(item.Content) + case "reasoning": + m["summary"] = reasoningSummaryWire(item.Summary) + if item.EncryptedContent != "" { + m["encrypted_content"] = item.EncryptedContent + } + case "function_call": + m["call_id"] = item.CallID + m["name"] = item.Name + m["arguments"] = item.Arguments + } + return m +} + +// messageContentWire renders a message item's content array; always an array +// (never null), with each output_text part carrying its text. +func messageContentWire(parts []ResponsesContentPart) []map[string]any { + out := make([]map[string]any, 0, len(parts)) + for _, p := range parts { + typ := p.Type + if typ == "" { + typ = "output_text" + } + out = append(out, map[string]any{"type": typ, "text": p.Text}) + } + return out +} + +// reasoningSummaryWire renders a reasoning item's summary array; always an array. +func reasoningSummaryWire(summary []ResponsesSummary) []map[string]any { + out := make([]map[string]any, 0, len(summary)) + for _, s := range summary { + typ := s.Type + if typ == "" { + typ = "summary_text" + } + out = append(out, map[string]any{"type": typ, "text": s.Text}) + } + return out +} diff --git a/backend/internal/pkg/apicompat/responses_stream_event_wire_test.go b/backend/internal/pkg/apicompat/responses_stream_event_wire_test.go new file mode 100644 index 000000000000..fbef45af4544 --- /dev/null +++ b/backend/internal/pkg/apicompat/responses_stream_event_wire_test.go @@ -0,0 +1,109 @@ +package apicompat + +import ( + "encoding/json" + "testing" + + "github.com/stretchr/testify/require" +) + +// marshalEvent marshals through the custom MarshalJSON and returns the decoded +// object plus the set of top-level keys. +func marshalEvent(t *testing.T, e ResponsesStreamEvent) map[string]any { + t.Helper() + b, err := json.Marshal(e) + require.NoError(t, err) + var m map[string]any + require.NoError(t, json.Unmarshal(b, &m)) + return m +} + +// TestWire_IndexFieldsPresentAtZero guards the omitempty trap: output_index/ +// content_index/summary_index must serialize even when 0. +func TestWire_IndexFieldsPresentAtZero(t *testing.T) { + m := marshalEvent(t, ResponsesStreamEvent{ + Type: "response.output_text.delta", OutputIndex: 0, ContentIndex: 0, ItemID: "msg_1", Delta: "hi", + }) + require.Contains(t, m, "output_index") + require.Contains(t, m, "content_index") + require.EqualValues(t, 0, m["output_index"]) + + r := marshalEvent(t, ResponsesStreamEvent{ + Type: "response.reasoning_summary_text.delta", OutputIndex: 0, SummaryIndex: 0, ItemID: "rs_1", Delta: "think", + }) + require.Contains(t, r, "output_index") + require.Contains(t, r, "summary_index") +} + +// TestWire_FunctionCallItemAlwaysComplete guards that a function_call item +// always carries call_id/name/arguments, including arguments:"" on .added. +func TestWire_FunctionCallItemAlwaysComplete(t *testing.T) { + added := marshalEvent(t, ResponsesStreamEvent{ + Type: "response.output_item.added", + OutputIndex: 1, + Item: &ResponsesOutput{Type: "function_call", ID: "fc_1", CallID: "call_a", Name: "exec", Status: "in_progress"}, + }) + item := added["item"].(map[string]any) + for _, k := range []string{"call_id", "name", "arguments"} { + require.Containsf(t, item, k, "function_call item missing %q", k) + } + require.Equal(t, "", item["arguments"]) +} + +// TestWire_MessageItemContentAlwaysArray guards content:[] presence. +func TestWire_MessageItemContentAlwaysArray(t *testing.T) { + m := marshalEvent(t, ResponsesStreamEvent{ + Type: "response.output_item.added", + OutputIndex: 0, + Item: &ResponsesOutput{Type: "message", ID: "msg_1", Role: "assistant", Status: "in_progress"}, + }) + item := m["item"].(map[string]any) + require.Contains(t, item, "content") + _, ok := item["content"].([]any) + require.True(t, ok, "content must be an array") +} + +// TestWire_ReasoningItemSummaryAlwaysArray guards summary:[] presence. +func TestWire_ReasoningItemSummaryAlwaysArray(t *testing.T) { + m := marshalEvent(t, ResponsesStreamEvent{ + Type: "response.output_item.added", + OutputIndex: 0, + Item: &ResponsesOutput{Type: "reasoning", ID: "rs_1", Status: "in_progress"}, + }) + item := m["item"].(map[string]any) + require.Contains(t, item, "summary") + _, ok := item["summary"].([]any) + require.True(t, ok, "summary must be an array") +} + +// TestWire_ContentPartCarriesAnnotationsLogprobs guards the output_text part shape. +func TestWire_ContentPartCarriesAnnotationsLogprobs(t *testing.T) { + m := marshalEvent(t, ResponsesStreamEvent{ + Type: "response.content_part.added", OutputIndex: 0, ContentIndex: 0, ItemID: "msg_1", + Part: &ResponsesContentPart{Type: "output_text", Text: ""}, + }) + part := m["part"].(map[string]any) + require.Equal(t, "output_text", part["type"]) + require.Contains(t, part, "text") + require.Contains(t, part, "annotations") + require.Contains(t, part, "logprobs") +} + +// TestWire_ArgumentsDonePresentEvenEmpty guards arguments presence on done. +func TestWire_ArgumentsDonePresentEvenEmpty(t *testing.T) { + m := marshalEvent(t, ResponsesStreamEvent{ + Type: "response.function_call_arguments.done", OutputIndex: 1, ItemID: "fc_1", CallID: "call_a", Name: "exec", Arguments: "", + }) + require.Contains(t, m, "arguments") + require.Equal(t, "", m["arguments"]) +} + +// TestWire_UnknownEventFallsBackToDefault ensures non-streamed event types keep +// default marshalling (the response object is preserved). +func TestWire_UnknownEventFallsBackToDefault(t *testing.T) { + m := marshalEvent(t, ResponsesStreamEvent{ + Type: "response.completed", + Response: &ResponsesResponse{ID: "resp_1", Object: "response", Status: "completed"}, + }) + require.Contains(t, m, "response") +} diff --git a/backend/internal/pkg/apicompat/types.go b/backend/internal/pkg/apicompat/types.go index b4451f235bb3..d29378027896 100644 --- a/backend/internal/pkg/apicompat/types.go +++ b/backend/internal/pkg/apicompat/types.go @@ -406,6 +406,10 @@ type ResponsesStreamEvent struct { // Reuses Text/Delta fields above, SummaryIndex identifies which summary part SummaryIndex int `json:"summary_index,omitempty"` + // response.content_part.added / done and + // response.reasoning_summary_part.added / done + Part *ResponsesContentPart `json:"part,omitempty"` + // error event fields Code string `json:"code,omitempty"` Param string `json:"param,omitempty"` From 1afae0a4cded8d4ca81bc7ea70bf55bbbaa5c72c Mon Sep 17 00:00:00 2001 From: name <136912576+is7Qin@users.noreply.github.com> Date: Sun, 31 May 2026 16:38:42 +0800 Subject: [PATCH 30/36] fix(gateway): address OpenAI OOM review lint --- .../service/gateway_service_benchmark_test.go | 94 +++++++++---------- .../service/openai_gateway_service.go | 19 +++- .../openai_gateway_service_hotpath_test.go | 45 +++++++++ .../service/openai_tool_continuation.go | 2 +- 4 files changed, 109 insertions(+), 51 deletions(-) diff --git a/backend/internal/service/gateway_service_benchmark_test.go b/backend/internal/service/gateway_service_benchmark_test.go index f6cd14041544..42a711db9b62 100644 --- a/backend/internal/service/gateway_service_benchmark_test.go +++ b/backend/internal/service/gateway_service_benchmark_test.go @@ -229,16 +229,16 @@ func benchmarkBodySizes() []struct { func buildSystemCacheableRequest(parts int) *ParsedRequest { var builder strings.Builder - builder.WriteString(`{"system":[`) + _, _ = builder.WriteString(`{"system":[`) for i := 0; i < parts; i++ { if i > 0 { - builder.WriteByte(',') + _ = builder.WriteByte(',') } - builder.WriteString(`{"text":"system_part_`) - builder.WriteString(strconv.Itoa(i)) - builder.WriteString(`","cache_control":{"type":"ephemeral"}}`) + _, _ = builder.WriteString(`{"text":"system_part_`) + _, _ = builder.WriteString(strconv.Itoa(i)) + _, _ = builder.WriteString(`","cache_control":{"type":"ephemeral"}}`) } - builder.WriteString(`]}`) + _, _ = builder.WriteString(`]}`) parsed, err := ParseGatewayRequest(NewRequestBodyRef([]byte(builder.String())), "") if err != nil { panic(err) @@ -249,93 +249,93 @@ func buildSystemCacheableRequest(parts int) *ParsedRequest { func buildLargeAnthropicMessagesBody(targetBytes int, includeCacheControl bool) []byte { var builder strings.Builder builder.Grow(targetBytes + 1024) - builder.WriteString(`{"model":"claude-sonnet-4-5","stream":true,"system":[{"type":"text","text":"system seed"}],"messages":[`) + _, _ = builder.WriteString(`{"model":"claude-sonnet-4-5","stream":true,"system":[{"type":"text","text":"system seed"}],"messages":[`) for i := 0; builder.Len() < targetBytes; i++ { if i > 0 { - builder.WriteByte(',') + _ = builder.WriteByte(',') } - builder.WriteString(`{"role":"user","content":[{"type":"text","text":"`) - builder.WriteString(strings.Repeat("anthropic payload ", 64)) - builder.WriteString(strconv.Itoa(i)) - builder.WriteByte('"') + _, _ = builder.WriteString(`{"role":"user","content":[{"type":"text","text":"`) + _, _ = builder.WriteString(strings.Repeat("anthropic payload ", 64)) + _, _ = builder.WriteString(strconv.Itoa(i)) + _ = builder.WriteByte('"') if includeCacheControl && i%32 == 0 { - builder.WriteString(`,"cache_control":{"type":"ephemeral"}`) + _, _ = builder.WriteString(`,"cache_control":{"type":"ephemeral"}`) } - builder.WriteString(`}]}`) + _, _ = builder.WriteString(`}]}`) } - builder.WriteString(`]}`) + _, _ = builder.WriteString(`]}`) return []byte(builder.String()) } func buildLargeGeminiContentsBody(targetBytes int) []byte { var builder strings.Builder builder.Grow(targetBytes + 1024) - builder.WriteString(`{"model":"gemini-2.5-pro","systemInstruction":{"parts":[{"text":"system seed"}]},"contents":[`) + _, _ = builder.WriteString(`{"model":"gemini-2.5-pro","systemInstruction":{"parts":[{"text":"system seed"}]},"contents":[`) for i := 0; builder.Len() < targetBytes; i++ { if i > 0 { - builder.WriteByte(',') + _ = builder.WriteByte(',') } - builder.WriteString(`{"role":"user","parts":[{"text":"`) - builder.WriteString(strings.Repeat("gemini payload ", 64)) - builder.WriteString(strconv.Itoa(i)) - builder.WriteString(`"}]}`) + _, _ = builder.WriteString(`{"role":"user","parts":[{"text":"`) + _, _ = builder.WriteString(strings.Repeat("gemini payload ", 64)) + _, _ = builder.WriteString(strconv.Itoa(i)) + _, _ = builder.WriteString(`"}]}`) } - builder.WriteString(`]}`) + _, _ = builder.WriteString(`]}`) return []byte(builder.String()) } func buildLargeOpenAIResponsesBody(targetBytes int) []byte { var builder strings.Builder builder.Grow(targetBytes + 1024) - builder.WriteString(`{"model":"gpt-5.4","stream":true,"prompt_cache_key":"session-benchmark","input":[`) + _, _ = builder.WriteString(`{"model":"gpt-5.4","stream":true,"prompt_cache_key":"session-benchmark","input":[`) for i := 0; builder.Len() < targetBytes; i++ { if i > 0 { - builder.WriteByte(',') + _ = builder.WriteByte(',') } - builder.WriteString(`{"type":"message","role":"user","content":[{"type":"input_text","text":"`) - builder.WriteString(strings.Repeat("openai responses payload ", 48)) - builder.WriteString(strconv.Itoa(i)) - builder.WriteString(`"}]}`) + _, _ = builder.WriteString(`{"type":"message","role":"user","content":[{"type":"input_text","text":"`) + _, _ = builder.WriteString(strings.Repeat("openai responses payload ", 48)) + _, _ = builder.WriteString(strconv.Itoa(i)) + _, _ = builder.WriteString(`"}]}`) } - builder.WriteString(`],"tools":[{"type":"function","name":"lookup","parameters":{"type":"object","properties":{"query":{"type":"string"}}}}]}`) + _, _ = builder.WriteString(`],"tools":[{"type":"function","name":"lookup","parameters":{"type":"object","properties":{"query":{"type":"string"}}}}]}`) return []byte(builder.String()) } func buildLargeOpenAIResponsesToolContinuationBody(targetBytes int) []byte { var builder strings.Builder builder.Grow(targetBytes + 1024) - builder.WriteString(`{"model":"gpt-5.4","stream":true,"previous_response_id":"resp_benchmark","input":[`) + _, _ = builder.WriteString(`{"model":"gpt-5.4","stream":true,"previous_response_id":"resp_benchmark","input":[`) for i := 0; builder.Len() < targetBytes; i++ { if i > 0 { - builder.WriteByte(',') + _ = builder.WriteByte(',') } callID := "call_" + strconv.Itoa(i) - builder.WriteString(`{"type":"item_reference","id":"`) - builder.WriteString(callID) - builder.WriteString(`"},{"type":"function_call_output","call_id":"`) - builder.WriteString(callID) - builder.WriteString(`","output":"`) - builder.WriteString(strings.Repeat("tool output payload ", 48)) - builder.WriteString(strconv.Itoa(i)) - builder.WriteString(`"}`) + _, _ = builder.WriteString(`{"type":"item_reference","id":"`) + _, _ = builder.WriteString(callID) + _, _ = builder.WriteString(`"},{"type":"function_call_output","call_id":"`) + _, _ = builder.WriteString(callID) + _, _ = builder.WriteString(`","output":"`) + _, _ = builder.WriteString(strings.Repeat("tool output payload ", 48)) + _, _ = builder.WriteString(strconv.Itoa(i)) + _, _ = builder.WriteString(`"}`) } - builder.WriteString(`]}`) + _, _ = builder.WriteString(`]}`) return []byte(builder.String()) } func buildLargeOpenAIResponsesImageToolBody(targetBytes int) []byte { var builder strings.Builder builder.Grow(targetBytes + 1024) - builder.WriteString(`{"model":"gpt-5.4","stream":false,"tools":[{"type":"image_generation","model":"gpt-image-2","size":"2048x1152"}],"input":[`) + _, _ = builder.WriteString(`{"model":"gpt-5.4","stream":false,"tools":[{"type":"image_generation","model":"gpt-image-2","size":"2048x1152"}],"input":[`) for i := 0; builder.Len() < targetBytes; i++ { if i > 0 { - builder.WriteByte(',') + _ = builder.WriteByte(',') } - builder.WriteString(`{"type":"message","role":"user","content":[{"type":"input_text","text":"`) - builder.WriteString(strings.Repeat("openai image billing payload ", 48)) - builder.WriteString(strconv.Itoa(i)) - builder.WriteString(`"}]}`) + _, _ = builder.WriteString(`{"type":"message","role":"user","content":[{"type":"input_text","text":"`) + _, _ = builder.WriteString(strings.Repeat("openai image billing payload ", 48)) + _, _ = builder.WriteString(strconv.Itoa(i)) + _, _ = builder.WriteString(`"}]}`) } - builder.WriteString(`]}`) + _, _ = builder.WriteString(`]}`) return []byte(builder.String()) } diff --git a/backend/internal/service/openai_gateway_service.go b/backend/internal/service/openai_gateway_service.go index d17b85d0d26f..8fd4bc603349 100644 --- a/backend/internal/service/openai_gateway_service.go +++ b/backend/internal/service/openai_gateway_service.go @@ -2544,6 +2544,7 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco logger.LegacyPrintf("service.openai_gateway", "[OpenAI] Added Codex image_generation bridge instructions") } } else if imageGenerationAllowed && imageIntent && openAIRequestBodyHasImageGenerationTool(body) { + // 完整 image_generation tool 只做 raw 计费读取，校验/桥接/旧字段迁移命中时才展开大 input map。 logger.LegacyPrintf("service.openai_gateway", "[OpenAI] /responses image_generation request inbound_model=%s mapped_model=%s account_type=%s", requestView.Model, upstreamModel, account.Type) } @@ -6215,7 +6216,7 @@ func (v *openAIRequestView) MarkPatchSet(path string, value any) { return } path = strings.TrimSpace(path) - if path == "" { + if !isSimpleOpenAIRequestPatchPath(path) { v.DisablePatches() return } @@ -6227,13 +6228,25 @@ func (v *openAIRequestView) MarkPatchDelete(path string) { return } path = strings.TrimSpace(path) - if path == "" { + if !isSimpleOpenAIRequestPatchPath(path) { v.DisablePatches() return } v.patches = append(v.patches, openAIRequestPatch{path: path, delete: true}) } +func isSimpleOpenAIRequestPatchPath(path string) bool { + if path == "" || strings.ContainsRune(path, '\\') { + return false + } + for _, part := range strings.Split(path, ".") { + if strings.TrimSpace(part) == "" { + return false + } + } + return true +} + func (v *openAIRequestView) DisablePatches() { if v == nil { return @@ -6772,7 +6785,7 @@ func openAIRequestBodyMayContainImageInput(body []byte) bool { return false } input := gjson.GetBytes(body, "input") - messages := gjson.GetBytes(body, "messages") + messages := gjson.GetBytes(body, "messages.#-1") return openAIJSONValueMayContainImageInput(input) || openAIJSONValueMayContainImageInput(messages) } diff --git a/backend/internal/service/openai_gateway_service_hotpath_test.go b/backend/internal/service/openai_gateway_service_hotpath_test.go index d240cec84fd9..92a0d1ac03dc 100644 --- a/backend/internal/service/openai_gateway_service_hotpath_test.go +++ b/backend/internal/service/openai_gateway_service_hotpath_test.go @@ -46,6 +46,15 @@ func TestOpenAIRequestView_ApplyPatches(t *testing.T) { require.JSONEq(t, `{"model":"gpt-5.1","reasoning":{"effort":"none"},"input":[{"type":"message","content":"hi"}]}`, string(patched)) } +func TestOpenAIRequestView_RejectsEscapedPatchPath(t *testing.T) { + view := newOpenAIRequestView([]byte(`{"metadata":{"user.id":"old"}}`)) + view.MarkPatchSet(`metadata.user\.id`, "new") + + require.False(t, view.HasPatches()) + _, err := view.ApplyPatches() + require.Error(t, err) +} + func TestOpenAIRequestView_ApplyPatchesDisabled(t *testing.T) { view := newOpenAIRequestView([]byte(`{"model":"gpt-5"}`)) view.MarkPatchSet("model", "gpt-5.1") @@ -260,6 +269,42 @@ func TestOpenAIGatewayService_Forward_ImageToolBillingDoesNotForceFullDecode(t * require.Equal(t, "gpt-image-2", result.BillingModel) } +func TestOpenAIGatewayService_Forward_ImageToolWithImageOnlyModelIsNormalized(t *testing.T) { + gin.SetMode(gin.TestMode) + upstream := &httpUpstreamRecorder{ + resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"usage":{"input_tokens":1,"output_tokens":2}}`)), + }, + } + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream} + account := &Account{ + ID: 11, + Name: "openai-apikey", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "sk-test", + "base_url": "https://example.com", + }, + Extra: map[string]any{"use_responses_api": true}, + } + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil) + SetOpenAIClientTransport(c, OpenAIClientTransportHTTP) + + body := []byte(`{"model":"gpt-image-2","stream":false,"tools":[{"type":"image_generation","model":"gpt-image-2"}],"input":"draw"}`) + result, err := svc.Forward(context.Background(), c, account, body) + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, openAIImagesResponsesMainModel, gjson.GetBytes(upstream.lastBody, "model").String()) +} + func TestOpenAIGatewayService_Forward_HTTPRetryRecoveryDoesNotDecodeBeforeError(t *testing.T) { gin.SetMode(gin.TestMode) upstream := &httpUpstreamRecorder{ diff --git a/backend/internal/service/openai_tool_continuation.go b/backend/internal/service/openai_tool_continuation.go index 92bc2613bb6f..6515c0c4e567 100644 --- a/backend/internal/service/openai_tool_continuation.go +++ b/backend/internal/service/openai_tool_continuation.go @@ -199,7 +199,7 @@ func ValidateFunctionCallOutputContextBytes(body []byte) FunctionCallOutputValid } referenceIDs[idValue] = struct{}{} } - return !(result.HasFunctionCallOutput && result.HasToolCallContext) + return !result.HasFunctionCallOutput || !result.HasToolCallContext }) if !result.HasFunctionCallOutput || result.HasToolCallContext || len(callIDs) == 0 || len(referenceIDs) == 0 { return result From a01686c637c81c29d19d7139f154c7ca9d4a1add Mon Sep 17 00:00:00 2001 From: moonagic Date: Sun, 31 May 2026 21:44:40 +0800 Subject: [PATCH 31/36] fix antigravity gemini rate limit and account scheduling Squash of 4 commits: - Fix Gemini rate limit scheduling - fix antigravity gemini rate limit scheduling - fix antigravity gemini limited account scheduling - fix antigravity test stubs for default lint --- backend/internal/handler/gateway_handler.go | 12 +- .../internal/handler/gemini_v1beta_handler.go | 13 +- .../internal/repository/scheduler_cache.go | 1 + .../repository/scheduler_cache_unit_test.go | 26 ++++ .../antigravity_default_test_stubs_test.go | 61 +++++++++ .../service/antigravity_gateway_service.go | 126 +++++++++++++----- .../antigravity_gateway_service_test.go | 80 +++++++++++ .../service/antigravity_quota_scope.go | 12 +- .../service/antigravity_rate_limit_test.go | 95 +++++++++++++ .../antigravity_single_account_retry_test.go | 8 +- .../service/antigravity_smart_retry_test.go | 28 ++-- backend/internal/service/error_policy_test.go | 54 ++++++++ .../service/gateway_multiplatform_test.go | 100 ++++++++++++++ backend/internal/service/model_rate_limit.go | 32 ++++- .../internal/service/model_rate_limit_test.go | 62 +++++++++ .../scheduler_snapshot_hydration_test.go | 88 ++++++++++++ 16 files changed, 747 insertions(+), 51 deletions(-) create mode 100644 backend/internal/service/antigravity_default_test_stubs_test.go diff --git a/backend/internal/handler/gateway_handler.go b/backend/internal/handler/gateway_handler.go index eb5c4a42b393..50c9fe34cefe 100644 --- a/backend/internal/handler/gateway_handler.go +++ b/backend/internal/handler/gateway_handler.go @@ -440,7 +440,17 @@ func (h *GatewayHandler) Messages(c *gin.Context) { // 记录 Forward 前已写入字节数，Forward 后若增加则说明 SSE 内容已发，禁止 failover writerSizeBeforeForward := c.Writer.Size() if account.Platform == service.PlatformAntigravity { - result, err = h.antigravityGatewayService.ForwardGemini(requestCtx, c, account, reqModel, "generateContent", reqStream, body, hasBoundSession) + result, err = h.antigravityGatewayService.ForwardGemini( + requestCtx, + c, + account, + reqModel, + "generateContent", + reqStream, + body, + hasBoundSession, + service.WithForwardGeminiSession(derefGroupID(apiKey.GroupID), sessionKey), + ) } else { result, err = h.geminiCompatService.Forward(requestCtx, c, account, body) } diff --git a/backend/internal/handler/gemini_v1beta_handler.go b/backend/internal/handler/gemini_v1beta_handler.go index 0b33ca3e3efe..dfe0b2a8ddb8 100644 --- a/backend/internal/handler/gemini_v1beta_handler.go +++ b/backend/internal/handler/gemini_v1beta_handler.go @@ -477,8 +477,19 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) { if fs.SwitchCount > 0 { requestCtx = service.WithAccountSwitchCount(requestCtx, fs.SwitchCount, h.metadataBridgeEnabled()) } + sessionGroupID := derefGroupID(apiKey.GroupID) if account.Platform == service.PlatformAntigravity && account.Type != service.AccountTypeAPIKey { - result, err = h.antigravityGatewayService.ForwardGemini(requestCtx, c, account, modelName, action, stream, body, hasBoundSession) + result, err = h.antigravityGatewayService.ForwardGemini( + requestCtx, + c, + account, + modelName, + action, + stream, + body, + hasBoundSession, + service.WithForwardGeminiSession(sessionGroupID, sessionKey), + ) } else { result, err = h.geminiCompatService.ForwardNative(requestCtx, c, account, modelName, action, stream, body) } diff --git a/backend/internal/repository/scheduler_cache.go b/backend/internal/repository/scheduler_cache.go index cf19dedaa0b1..921aa081cae5 100644 --- a/backend/internal/repository/scheduler_cache.go +++ b/backend/internal/repository/scheduler_cache.go @@ -559,6 +559,7 @@ func filterSchedulerExtra(extra map[string]any) map[string]any { "auto_pause_7d_threshold", "auto_pause_5h_disabled", "auto_pause_7d_disabled", + "model_rate_limits", } filtered := make(map[string]any) for _, key := range keys { diff --git a/backend/internal/repository/scheduler_cache_unit_test.go b/backend/internal/repository/scheduler_cache_unit_test.go index a4667591722a..c14721cd7c18 100644 --- a/backend/internal/repository/scheduler_cache_unit_test.go +++ b/backend/internal/repository/scheduler_cache_unit_test.go @@ -108,3 +108,29 @@ func TestBuildSchedulerMetadataAccount_KeepsQuotaAutoPauseFields(t *testing.T) { require.Equal(t, true, got.Extra["auto_pause_5h_disabled"]) require.Equal(t, false, got.Extra["auto_pause_7d_disabled"]) } + +func TestBuildSchedulerMetadataAccount_KeepsModelRateLimits(t *testing.T) { + account := service.Account{ + ID: 90, + Platform: service.PlatformAntigravity, + Extra: map[string]any{ + "model_rate_limits": map[string]any{ + "gemini-3-flash": map[string]any{ + "rate_limit_reset_at": "2026-05-30T10:10:00Z", + }, + "antigravity:gemini": map[string]any{ + "rate_limit_reset_at": "2026-05-30T10:10:00Z", + }, + }, + "unused_large_field": "drop-me", + }, + } + + got := buildSchedulerMetadataAccount(account) + + limits, ok := got.Extra["model_rate_limits"].(map[string]any) + require.True(t, ok) + require.Contains(t, limits, "gemini-3-flash") + require.Contains(t, limits, "antigravity:gemini") + require.Nil(t, got.Extra["unused_large_field"]) +} diff --git a/backend/internal/service/antigravity_default_test_stubs_test.go b/backend/internal/service/antigravity_default_test_stubs_test.go new file mode 100644 index 000000000000..d3c2c57a272b --- /dev/null +++ b/backend/internal/service/antigravity_default_test_stubs_test.go @@ -0,0 +1,61 @@ +//go:build !unit + +package service + +import ( + "context" + "time" +) + +type defaultRateLimitCall struct { + accountID int64 + resetAt time.Time +} + +type defaultModelRateLimitCall struct { + accountID int64 + modelKey string + resetAt time.Time +} + +type defaultExtraUpdateCall struct { + accountID int64 + updates map[string]any +} + +type stubAntigravityAccountRepo struct { + AccountRepository + rateCalls []defaultRateLimitCall + modelRateLimitCalls []defaultModelRateLimitCall + extraUpdateCalls []defaultExtraUpdateCall +} + +func (s *stubAntigravityAccountRepo) SetRateLimited(_ context.Context, id int64, resetAt time.Time) error { + s.rateCalls = append(s.rateCalls, defaultRateLimitCall{accountID: id, resetAt: resetAt}) + return nil +} + +func (s *stubAntigravityAccountRepo) SetModelRateLimit(_ context.Context, id int64, modelKey string, resetAt time.Time, _ ...string) error { + s.modelRateLimitCalls = append(s.modelRateLimitCalls, defaultModelRateLimitCall{accountID: id, modelKey: modelKey, resetAt: resetAt}) + return nil +} + +func (s *stubAntigravityAccountRepo) UpdateExtra(_ context.Context, id int64, updates map[string]any) error { + s.extraUpdateCalls = append(s.extraUpdateCalls, defaultExtraUpdateCall{accountID: id, updates: updates}) + return nil +} + +type defaultDeleteSessionCall struct { + groupID int64 + sessionHash string +} + +type stubSmartRetryCache struct { + GatewayCache + deleteCalls []defaultDeleteSessionCall +} + +func (c *stubSmartRetryCache) DeleteSessionAccountID(_ context.Context, groupID int64, sessionHash string) error { + c.deleteCalls = append(c.deleteCalls, defaultDeleteSessionCall{groupID: groupID, sessionHash: sessionHash}) + return nil +} diff --git a/backend/internal/service/antigravity_gateway_service.go b/backend/internal/service/antigravity_gateway_service.go index 2b849bdd8de0..8aca2e130516 100644 --- a/backend/internal/service/antigravity_gateway_service.go +++ b/backend/internal/service/antigravity_gateway_service.go @@ -228,12 +228,11 @@ func (s *AntigravityGatewayService) handleSmartRetry(p antigravityRetryLoopParam p.prefix, resp.StatusCode, modelName, p.account.ID, rateLimitDuration, truncateForLog(respBody, 200)) resetAt := time.Now().Add(rateLimitDuration) - if !setModelRateLimitByModelName(p.ctx, p.accountRepo, p.account.ID, modelName, p.prefix, resp.StatusCode, resetAt, false) { + if !s.setAntigravityModelRateLimits(p.ctx, p.accountRepo, p.account, modelName, p.prefix, resp.StatusCode, resetAt, false) { p.handleError(p.ctx, p.prefix, p.account, resp.StatusCode, resp.Header, respBody, p.requestedModel, p.groupID, p.sessionHash, p.isStickySession) logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d rate_limited account=%d (no model mapping)", p.prefix, resp.StatusCode, p.account.ID) - } else { - s.updateAccountModelRateLimitInCache(p.ctx, p.account, modelName, resetAt) } + s.clearStickySession(p.ctx, p.groupID, p.sessionHash) // 返回账号切换信号，让上层切换账号重试 return &smartRetryResult{ @@ -392,20 +391,10 @@ func (s *AntigravityGatewayService) handleSmartRetry(p antigravityRetryLoopParam p.prefix, resp.StatusCode, maxAttempts, modelName, p.account.ID, rateLimitDuration, truncateForLog(retryBody, 200)) resetAt := time.Now().Add(rateLimitDuration) - if p.accountRepo != nil && modelName != "" { - if err := p.accountRepo.SetModelRateLimit(p.ctx, p.account.ID, modelName, resetAt); err != nil { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d model_rate_limit_failed model=%s error=%v", p.prefix, resp.StatusCode, modelName, err) - } else { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d model_rate_limited_after_smart_retry model=%s account=%d reset_in=%v", - p.prefix, resp.StatusCode, modelName, p.account.ID, rateLimitDuration) - s.updateAccountModelRateLimitInCache(p.ctx, p.account, modelName, resetAt) - } - } + s.setAntigravityModelRateLimits(p.ctx, p.accountRepo, p.account, modelName, p.prefix, resp.StatusCode, resetAt, true) // 清除粘性会话绑定，避免下次请求仍命中限流账号 - if s.cache != nil && p.sessionHash != "" { - _ = s.cache.DeleteSessionAccountID(p.ctx, p.groupID, p.sessionHash) - } + s.clearStickySession(p.ctx, p.groupID, p.sessionHash) // 返回账号切换信号，让上层切换账号重试 return &smartRetryResult{ @@ -938,8 +927,14 @@ func (s *AntigravityGatewayService) checkErrorPolicy(ctx context.Context, accoun func (s *AntigravityGatewayService) applyErrorPolicy(p antigravityRetryLoopParams, statusCode int, headers http.Header, respBody []byte) (handled bool, outStatus int, retErr error) { switch s.checkErrorPolicy(p.ctx, p.account, statusCode, respBody) { case ErrorPolicySkipped: + if s.handleAntigravityModelRateLimitBeforePolicy(p, statusCode, headers, respBody) { + return true, statusCode, nil + } return true, http.StatusInternalServerError, nil case ErrorPolicyMatched: + if s.handleAntigravityModelRateLimitBeforePolicy(p, statusCode, headers, respBody) { + return true, statusCode, nil + } _ = p.handleError(p.ctx, p.prefix, p.account, statusCode, headers, respBody, p.requestedModel, p.groupID, p.sessionHash, p.isStickySession) return true, statusCode, nil @@ -951,6 +946,31 @@ func (s *AntigravityGatewayService) applyErrorPolicy(p antigravityRetryLoopParam return false, statusCode, nil } +func (s *AntigravityGatewayService) handleAntigravityModelRateLimitBeforePolicy(p antigravityRetryLoopParams, statusCode int, headers http.Header, respBody []byte) bool { + if statusCode != http.StatusTooManyRequests && statusCode != http.StatusServiceUnavailable { + return false + } + if p.account == nil || p.account.Platform != PlatformAntigravity { + return false + } + _, shouldRateLimitModel, waitDuration, modelName, isModelCapacityExhausted := shouldTriggerAntigravitySmartRetry(p.account, respBody) + if isModelCapacityExhausted || !shouldRateLimitModel || strings.TrimSpace(modelName) == "" { + return false + } + rateLimitDuration := waitDuration + if rateLimitDuration <= 0 { + rateLimitDuration = antigravityDefaultRateLimitDuration + } + resetAt := time.Now().Add(rateLimitDuration) + if !s.setAntigravityModelRateLimits(p.ctx, p.accountRepo, p.account, modelName, p.prefix, statusCode, resetAt, false) { + return false + } + s.clearStickySession(p.ctx, p.groupID, p.sessionHash) + logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d model_rate_limited_before_error_policy model=%s account=%d reset_in=%v", + p.prefix, statusCode, modelName, p.account.ID, rateLimitDuration) + return true +} + // mapAntigravityModel 获取映射后的模型名 // 完全依赖映射配置：账户映射（通配符）→ 默认映射兜底（DefaultAntigravityModelMapping） // 注意：返回空字符串表示模型不被支持，调度时会过滤掉该账号 @@ -958,6 +978,7 @@ func mapAntigravityModel(account *Account, requestedModel string) string { if account == nil { return "" } + requestedModel = strings.TrimPrefix(requestedModel, "models/") // 获取映射表（未配置时自动使用 DefaultAntigravityModelMapping） mapping := account.GetModelMapping() @@ -2057,8 +2078,28 @@ func stripSignatureSensitiveBlocksFromClaudeRequest(req *antigravity.ClaudeReque // └─ retryDelay < 7s → 等待后重试 1 次 // ├─ 成功 → 正常返回 // └─ 失败 → 设置模型限流 + 清除粘性绑定 → 切换账号 -func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Context, account *Account, originalModel string, action string, stream bool, body []byte, isStickySession bool) (*ForwardResult, error) { +type ForwardGeminiOption func(*forwardGeminiOptions) + +type forwardGeminiOptions struct { + groupID int64 + sessionHash string +} + +func WithForwardGeminiSession(groupID int64, sessionHash string) ForwardGeminiOption { + return func(opts *forwardGeminiOptions) { + opts.groupID = groupID + opts.sessionHash = sessionHash + } +} + +func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Context, account *Account, originalModel string, action string, stream bool, body []byte, isStickySession bool, options ...ForwardGeminiOption) (*ForwardResult, error) { startTime := time.Now() + forwardOpts := forwardGeminiOptions{} + for _, apply := range options { + if apply != nil { + apply(&forwardOpts) + } + } sessionID := getSessionID(c) prefix := logPrefix(sessionID, account.Name) @@ -2163,8 +2204,8 @@ func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Co handleError: s.handleUpstreamError, requestedModel: originalModel, isStickySession: isStickySession, // ForwardGemini 由上层判断粘性会话 - groupID: 0, // ForwardGemini 方法没有 groupID，由上层处理粘性会话清除 - sessionHash: "", // ForwardGemini 方法没有 sessionHash，由上层处理粘性会话清除 + groupID: forwardOpts.groupID, + sessionHash: forwardOpts.sessionHash, }) if err != nil { // 检查是否是账号切换信号，转换为 UpstreamFailoverError 让 Handler 切换账号 @@ -2262,8 +2303,8 @@ func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Co handleError: s.handleUpstreamError, requestedModel: originalModel, isStickySession: isStickySession, - groupID: 0, - sessionHash: "", + groupID: forwardOpts.groupID, + sessionHash: forwardOpts.sessionHash, }) if retryErr == nil { retryResp := retryResult.resp @@ -2339,7 +2380,7 @@ func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Co if unwrapErr != nil || len(unwrappedForOps) == 0 { unwrappedForOps = respBody } - s.handleUpstreamError(ctx, prefix, account, resp.StatusCode, resp.Header, respBody, originalModel, 0, "", isStickySession) + s.handleUpstreamError(ctx, prefix, account, resp.StatusCode, resp.Header, respBody, originalModel, forwardOpts.groupID, forwardOpts.sessionHash, isStickySession) upstreamMsg := strings.TrimSpace(extractAntigravityErrorMessage(unwrappedForOps)) upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) upstreamDetail := s.getUpstreamErrorDetail(unwrappedForOps) @@ -2550,6 +2591,34 @@ func setModelRateLimitByModelName(ctx context.Context, repo AccountRepository, a return true } +func (s *AntigravityGatewayService) setAntigravityModelRateLimits(ctx context.Context, repo AccountRepository, account *Account, modelName, prefix string, statusCode int, resetAt time.Time, afterSmartRetry bool) bool { + if account == nil || repo == nil { + return false + } + keys := antigravityModelRateLimitKeys(modelName) + if len(keys) == 0 { + return false + } + + success := false + for _, key := range keys { + if setModelRateLimitByModelName(ctx, repo, account.ID, key, prefix, statusCode, resetAt, afterSmartRetry) { + s.updateAccountModelRateLimitInCache(ctx, account, key, resetAt) + success = true + } + } + return success +} + +func (s *AntigravityGatewayService) clearStickySession(ctx context.Context, groupID int64, sessionHash string) { + if s == nil || s.cache == nil || strings.TrimSpace(sessionHash) == "" { + return + } + if err := s.cache.DeleteSessionAccountID(ctx, groupID, sessionHash); err != nil { + logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Forward] sticky_session_clear_failed group_id=%d session=%s err=%v", groupID, shortSessionHash(sessionHash), err) + } +} + func antigravityFallbackCooldownSeconds() (time.Duration, bool) { raw := strings.TrimSpace(os.Getenv(antigravityFallbackSecondsEnv)) if raw == "" { @@ -2628,7 +2697,7 @@ func parseAntigravitySmartRetryInfo(body []byte) *antigravitySmartRetryInfo { if atType == googleRPCTypeErrorInfo { if meta, ok := dm["metadata"].(map[string]any); ok { if model, ok := meta["model"].(string); ok { - modelName = model + modelName = normalizeAntigravityModelName(model) } } // 检查 reason @@ -2802,13 +2871,7 @@ func (s *AntigravityGatewayService) setModelRateLimitAndClearSession(p *handleMo logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d model_rate_limited model=%s account=%d reset_in=%v", p.prefix, p.statusCode, info.ModelName, p.account.ID, info.RetryDelay) - // 设置模型限流状态（数据库） - if err := s.accountRepo.SetModelRateLimit(p.ctx, p.account.ID, info.ModelName, resetAt); err != nil { - logger.LegacyPrintf("service.antigravity_gateway", "%s model_rate_limit_failed model=%s error=%v", p.prefix, info.ModelName, err) - } - - // 立即更新 Redis 快照中账号的限流状态，避免并发请求重复选中 - s.updateAccountModelRateLimitInCache(p.ctx, p.account, info.ModelName, resetAt) + s.setAntigravityModelRateLimits(p.ctx, s.accountRepo, p.account, info.ModelName, p.prefix, p.statusCode, resetAt, false) // 清除粘性会话绑定 if p.cache != nil && p.sessionHash != "" { @@ -2898,12 +2961,11 @@ func (s *AntigravityGatewayService) handleUpstreamError( } if modelKey != "" { ra := s.resolveResetTime(resetAt, defaultDur) - if err := s.accountRepo.SetModelRateLimit(ctx, account.ID, modelKey, ra); err != nil { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=429 model_rate_limit_set_failed model=%s error=%v", prefix, modelKey, err) + if !s.setAntigravityModelRateLimits(ctx, s.accountRepo, account, modelKey, prefix, statusCode, ra, false) { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=429 model_rate_limit_set_failed model=%s", prefix, modelKey) } else { logger.LegacyPrintf("service.antigravity_gateway", "%s status=429 model_rate_limited model=%s account=%d reset_at=%v reset_in=%v", prefix, modelKey, account.ID, ra.Format("15:04:05"), time.Until(ra).Truncate(time.Second)) - s.updateAccountModelRateLimitInCache(ctx, account, modelKey, ra) } return nil } diff --git a/backend/internal/service/antigravity_gateway_service_test.go b/backend/internal/service/antigravity_gateway_service_test.go index 2212437467cf..1aa1ac9d989c 100644 --- a/backend/internal/service/antigravity_gateway_service_test.go +++ b/backend/internal/service/antigravity_gateway_service_test.go @@ -491,6 +491,86 @@ func TestAntigravityGatewayService_ForwardGemini_StickySessionForceCacheBilling( require.True(t, failoverErr.ForceCacheBilling, "ForceCacheBilling should be true for sticky session switch") } +func TestAntigravityGatewayService_ForwardGemini_ClearsStickySessionOnGeminiRateLimit(t *testing.T) { + gin.SetMode(gin.TestMode) + writer := httptest.NewRecorder() + c, _ := gin.CreateTestContext(writer) + + body, err := json.Marshal(map[string]any{ + "contents": []map[string]any{ + {"role": "user", "parts": []map[string]any{{"text": "hi"}}}, + }, + }) + require.NoError(t, err) + + req := httptest.NewRequest(http.MethodPost, "/v1beta/models/gemini-3-flash-preview:generateContent", bytes.NewReader(body)) + c.Request = req + + respBody := []byte(`{ + "error": { + "status": "RESOURCE_EXHAUSTED", + "details": [ + {"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-flash"}, "reason": "RATE_LIMIT_EXCEEDED"}, + {"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "15s"} + ] + } + }`) + upstream := &httpUpstreamStub{resp: &http.Response{ + StatusCode: http.StatusTooManyRequests, + Header: http.Header{}, + Body: io.NopCloser(bytes.NewReader(respBody)), + }} + repo := &stubAntigravityAccountRepo{} + cache := &stubSmartRetryCache{} + svc := &AntigravityGatewayService{ + tokenProvider: &AntigravityTokenProvider{}, + httpUpstream: upstream, + accountRepo: repo, + cache: cache, + } + + account := &Account{ + ID: 44, + Name: "acc-gemini-runtime-rate-limited", + Platform: PlatformAntigravity, + Type: AccountTypeOAuth, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Credentials: map[string]any{ + "access_token": "token", + "expires_at": time.Now().Add(time.Hour).Format(time.RFC3339), + "project_id": "proj", + }, + Extra: map[string]any{ + "mixed_scheduling": true, + }, + } + + result, err := svc.ForwardGemini( + context.Background(), + c, + account, + "gemini-3-flash-preview", + "generateContent", + false, + body, + true, + WithForwardGeminiSession(77, "gemini:sticky-runtime"), + ) + + require.Nil(t, result) + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, http.StatusServiceUnavailable, failoverErr.StatusCode) + require.Len(t, repo.modelRateLimitCalls, 2) + require.Equal(t, "gemini-3-flash", repo.modelRateLimitCalls[0].modelKey) + require.Equal(t, antigravityGeminiModelRateLimitKey, repo.modelRateLimitCalls[1].modelKey) + require.Len(t, cache.deleteCalls, 1) + require.Equal(t, int64(77), cache.deleteCalls[0].groupID) + require.Equal(t, "gemini:sticky-runtime", cache.deleteCalls[0].sessionHash) +} + // TestAntigravityGatewayService_Forward_BillsWithMappedModel // 验证：Antigravity Claude 转发返回的计费模型使用映射后的模型 func TestAntigravityGatewayService_Forward_BillsWithMappedModel(t *testing.T) { diff --git a/backend/internal/service/antigravity_quota_scope.go b/backend/internal/service/antigravity_quota_scope.go index b536d16cf893..758626331b33 100644 --- a/backend/internal/service/antigravity_quota_scope.go +++ b/backend/internal/service/antigravity_quota_scope.go @@ -8,7 +8,17 @@ import ( func normalizeAntigravityModelName(model string) string { normalized := strings.ToLower(strings.TrimSpace(model)) - normalized = strings.TrimPrefix(normalized, "models/") + if idx := strings.LastIndex(normalized, "/publishers/google/models/"); idx != -1 { + normalized = normalized[idx+len("/publishers/google/models/"):] + } else if idx := strings.LastIndex(normalized, "/publishers/anthropic/models/"); idx != -1 { + normalized = normalized[idx+len("/publishers/anthropic/models/"):] + } else if idx := strings.LastIndex(normalized, "/models/"); idx != -1 { + normalized = normalized[idx+len("/models/"):] + } else { + normalized = strings.TrimPrefix(normalized, "publishers/google/models/") + normalized = strings.TrimPrefix(normalized, "publishers/anthropic/models/") + normalized = strings.TrimPrefix(normalized, "models/") + } return normalized } diff --git a/backend/internal/service/antigravity_rate_limit_test.go b/backend/internal/service/antigravity_rate_limit_test.go index c3e49458d876..374b29f610cc 100644 --- a/backend/internal/service/antigravity_rate_limit_test.go +++ b/backend/internal/service/antigravity_rate_limit_test.go @@ -821,6 +821,51 @@ func TestSetModelRateLimitByModelName_NotConvertToScope(t *testing.T) { require.NotEqual(t, "claude_sonnet", call.modelKey, "should NOT be scope") } +func TestSetAntigravityModelRateLimits_GeminiWritesFamilyScope(t *testing.T) { + repo := &stubAntigravityAccountRepo{} + svc := &AntigravityGatewayService{} + account := &Account{ID: 789, Platform: PlatformAntigravity} + resetAt := time.Now().Add(30 * time.Second) + + success := svc.setAntigravityModelRateLimits( + context.Background(), + repo, + account, + "gemini-3-pro", + "[test]", + 429, + resetAt, + false, + ) + + require.True(t, success) + require.Len(t, repo.modelRateLimitCalls, 2) + require.Equal(t, "gemini-3-pro", repo.modelRateLimitCalls[0].modelKey) + require.Equal(t, antigravityGeminiModelRateLimitKey, repo.modelRateLimitCalls[1].modelKey) +} + +func TestSetAntigravityModelRateLimits_ClaudeDoesNotWriteGeminiScope(t *testing.T) { + repo := &stubAntigravityAccountRepo{} + svc := &AntigravityGatewayService{} + account := &Account{ID: 790, Platform: PlatformAntigravity} + resetAt := time.Now().Add(30 * time.Second) + + success := svc.setAntigravityModelRateLimits( + context.Background(), + repo, + account, + "claude-sonnet-4-5", + "[test]", + 429, + resetAt, + false, + ) + + require.True(t, success) + require.Len(t, repo.modelRateLimitCalls, 1) + require.Equal(t, "claude-sonnet-4-5", repo.modelRateLimitCalls[0].modelKey) +} + func TestAntigravityRetryLoop_PreCheck_SwitchesWhenRateLimited(t *testing.T) { upstream := &recordingOKUpstream{} account := &Account{ @@ -1124,3 +1169,53 @@ func TestSchedulerSnapshotService_UpdateAccountInCache(t *testing.T) { require.ErrorIs(t, err, expectedErr) }) } +func TestNormalizeAntigravityModelName(t *testing.T) { + tests := []struct { + name string + model string + expected string + }{ + { + name: "plain model name", + model: "gemini-1.5-pro", + expected: "gemini-1.5-pro", + }, + { + name: "models/ prefix", + model: "models/gemini-1.5-pro", + expected: "gemini-1.5-pro", + }, + { + name: "publishers/google/models/ prefix", + model: "publishers/google/models/gemini-1.5-pro", + expected: "gemini-1.5-pro", + }, + { + name: "projects/.../publishers/google/models/ path", + model: "projects/my-proj/locations/us-central1/publishers/google/models/gemini-2.5-flash", + expected: "gemini-2.5-flash", + }, + { + name: "publishers/anthropic/models/ prefix", + model: "publishers/anthropic/models/claude-sonnet-4-5", + expected: "claude-sonnet-4-5", + }, + { + name: "projects/.../publishers/anthropic/models/ path", + model: "projects/my-proj/locations/global/publishers/anthropic/models/claude-sonnet-4-5", + expected: "claude-sonnet-4-5", + }, + { + name: "mixed case and spaces", + model: " Models/Gemini-1.5-Pro ", + expected: "gemini-1.5-pro", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + actual := normalizeAntigravityModelName(tt.model) + require.Equal(t, tt.expected, actual) + }) + } +} diff --git a/backend/internal/service/antigravity_single_account_retry_test.go b/backend/internal/service/antigravity_single_account_retry_test.go index 675e9c0cbf4d..6e58ab75d2cb 100644 --- a/backend/internal/service/antigravity_single_account_retry_test.go +++ b/backend/internal/service/antigravity_single_account_retry_test.go @@ -196,8 +196,10 @@ func TestHandleSmartRetry_503_LongDelay_NoSingleAccountRetry_StillSwitches(t *te require.Nil(t, result.resp, "should not return resp when switchError is set") // 对照：多账号模式应设模型限流 - require.Len(t, repo.modelRateLimitCalls, 1, + require.Len(t, repo.modelRateLimitCalls, 2, "multi-account mode SHOULD set model rate limit") + require.Equal(t, "gemini-3-pro-high", repo.modelRateLimitCalls[0].modelKey) + require.Equal(t, antigravityGeminiModelRateLimitKey, repo.modelRateLimitCalls[1].modelKey) } // TestHandleSmartRetry_429_LongDelay_SingleAccountRetry_StillSwitches @@ -412,8 +414,10 @@ func TestHandleSmartRetry_503_ShortDelay_NoSingleAccountRetry_SetsRateLimit(t *t // 对照：多账号模式应返回 switchError require.NotNil(t, result.switchError, "multi-account mode should return switchError for 503") // 对照：多账号模式应设模型限流 - require.Len(t, repo.modelRateLimitCalls, 1, + require.Len(t, repo.modelRateLimitCalls, 2, "multi-account mode should set model rate limit") + require.Equal(t, "gemini-3-flash", repo.modelRateLimitCalls[0].modelKey) + require.Equal(t, antigravityGeminiModelRateLimitKey, repo.modelRateLimitCalls[1].modelKey) } // --------------------------------------------------------------------------- diff --git a/backend/internal/service/antigravity_smart_retry_test.go b/backend/internal/service/antigravity_smart_retry_test.go index e3b60a279027..9f06b13b0632 100644 --- a/backend/internal/service/antigravity_smart_retry_test.go +++ b/backend/internal/service/antigravity_smart_retry_test.go @@ -328,9 +328,10 @@ func TestHandleSmartRetry_ShortDelay_SmartRetryFailed_ReturnsSwitchError(t *test require.Equal(t, "gemini-3-flash", result.switchError.RateLimitedModel) require.False(t, result.switchError.IsStickySession) - // 验证模型限流已设置 - require.Len(t, repo.modelRateLimitCalls, 1) + // 验证模型限流已设置：Gemini 同时写入精确模型和家族级 scope + require.Len(t, repo.modelRateLimitCalls, 2) require.Equal(t, "gemini-3-flash", repo.modelRateLimitCalls[0].modelKey) + require.Equal(t, antigravityGeminiModelRateLimitKey, repo.modelRateLimitCalls[1].modelKey) require.Len(t, upstream.calls, 1, "should have made one retry call (max attempts)") } @@ -1104,10 +1105,9 @@ func TestHandleSmartRetry_ShortDelay_StickySession_SuccessRetry_NoDeleteSession( require.Len(t, cache.deleteCalls, 0, "should NOT call DeleteSessionAccountID on successful retry") } -// TestHandleSmartRetry_LongDelay_StickySession_NoDeleteInHandleSmartRetry -// 长延迟路径（情况1）在 handleSmartRetry 中不直接调用 DeleteSessionAccountID -// （清除由 handler 层的 shouldClearStickySession 在下次请求时处理） -func TestHandleSmartRetry_LongDelay_StickySession_NoDeleteInHandleSmartRetry(t *testing.T) { +// TestHandleSmartRetry_LongDelay_StickySession_ClearsSession +// 长延迟路径（情况1）应立即清除 sticky 绑定，避免下一次请求继续命中已限流账号。 +func TestHandleSmartRetry_LongDelay_StickySession_ClearsSession(t *testing.T) { repo := &stubAntigravityAccountRepo{} cache := &stubSmartRetryCache{} account := &Account{ @@ -1159,10 +1159,9 @@ func TestHandleSmartRetry_LongDelay_StickySession_NoDeleteInHandleSmartRetry(t * require.NotNil(t, result.switchError) require.True(t, result.switchError.IsStickySession) - // 长延迟路径不在 handleSmartRetry 中调用 DeleteSessionAccountID - // （由上游 handler 的 shouldClearStickySession 处理） - require.Len(t, cache.deleteCalls, 0, - "long delay path should NOT call DeleteSessionAccountID in handleSmartRetry (handled by handler layer)") + require.Len(t, cache.deleteCalls, 1, "long delay path should clear sticky session in handleSmartRetry") + require.Equal(t, int64(42), cache.deleteCalls[0].groupID) + require.Equal(t, "sticky-hash-long-delay", cache.deleteCalls[0].sessionHash) } // TestHandleSmartRetry_ShortDelay_NetworkError_StickySession_ClearsSession @@ -1227,6 +1226,10 @@ func TestHandleSmartRetry_ShortDelay_NetworkError_StickySession_ClearsSession(t require.Len(t, cache.deleteCalls, 1, "should call DeleteSessionAccountID after network error exhausts retry") require.Equal(t, int64(99), cache.deleteCalls[0].groupID) require.Equal(t, "sticky-net-error", cache.deleteCalls[0].sessionHash) + + require.Len(t, repo.modelRateLimitCalls, 2) + require.Equal(t, "gemini-3-flash", repo.modelRateLimitCalls[0].modelKey) + require.Equal(t, antigravityGeminiModelRateLimitKey, repo.modelRateLimitCalls[1].modelKey) } // TestHandleSmartRetry_ShortDelay_503_StickySession_FailedRetry_ClearsSession @@ -1308,9 +1311,10 @@ func TestHandleSmartRetry_ShortDelay_503_StickySession_FailedRetry_ClearsSession require.Equal(t, int64(77), cache.deleteCalls[0].groupID) require.Equal(t, "sticky-503-short", cache.deleteCalls[0].sessionHash) - // 验证模型限流已设置 - require.Len(t, repo.modelRateLimitCalls, 1) + // 验证模型限流已设置：Gemini 同时写入精确模型和家族级 scope + require.Len(t, repo.modelRateLimitCalls, 2) require.Equal(t, "gemini-3-pro", repo.modelRateLimitCalls[0].modelKey) + require.Equal(t, antigravityGeminiModelRateLimitKey, repo.modelRateLimitCalls[1].modelKey) } // TestAntigravityRetryLoop_SmartRetryFailed_StickySession_SwitchErrorPropagates diff --git a/backend/internal/service/error_policy_test.go b/backend/internal/service/error_policy_test.go index 297a954c3422..2aa7a421d1f7 100644 --- a/backend/internal/service/error_policy_test.go +++ b/backend/internal/service/error_policy_test.go @@ -389,6 +389,60 @@ func TestApplyErrorPolicy(t *testing.T) { } } +func TestApplyErrorPolicy_GeminiRateLimitBypassesCustomSkip(t *testing.T) { + repo := &stubAntigravityAccountRepo{} + cache := &stubSmartRetryCache{} + rlSvc := NewRateLimitService(repo, nil, &config.Config{}, nil, nil) + svc := &AntigravityGatewayService{ + rateLimitService: rlSvc, + accountRepo: repo, + cache: cache, + } + + account := &Account{ + ID: 31, + Type: AccountTypeAPIKey, + Platform: PlatformAntigravity, + Credentials: map[string]any{ + "custom_error_codes_enabled": true, + "custom_error_codes": []any{float64(500)}, + }, + } + body := []byte(`{ + "error": { + "status": "RESOURCE_EXHAUSTED", + "details": [ + {"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-flash"}, "reason": "RATE_LIMIT_EXCEEDED"}, + {"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "15s"} + ] + } + }`) + p := antigravityRetryLoopParams{ + ctx: context.Background(), + prefix: "[test]", + account: account, + accountRepo: repo, + groupID: 42, + sessionHash: "gemini:sticky", + handleError: func(context.Context, string, *Account, int, http.Header, []byte, string, int64, string, bool) *handleModelRateLimitResult { + t.Fatal("model rate limit should be handled before custom error fallback") + return nil + }, + } + + handled, outStatus, retErr := svc.applyErrorPolicy(p, http.StatusTooManyRequests, http.Header{}, body) + + require.True(t, handled) + require.Equal(t, http.StatusTooManyRequests, outStatus) + require.NoError(t, retErr) + require.Len(t, repo.modelRateLimitCalls, 2) + require.Equal(t, "gemini-3-flash", repo.modelRateLimitCalls[0].modelKey) + require.Equal(t, antigravityGeminiModelRateLimitKey, repo.modelRateLimitCalls[1].modelKey) + require.Len(t, cache.deleteCalls, 1) + require.Equal(t, int64(42), cache.deleteCalls[0].groupID) + require.Equal(t, "gemini:sticky", cache.deleteCalls[0].sessionHash) +} + // --------------------------------------------------------------------------- // errorPolicyRepoStub — minimal AccountRepository stub for error policy tests // --------------------------------------------------------------------------- diff --git a/backend/internal/service/gateway_multiplatform_test.go b/backend/internal/service/gateway_multiplatform_test.go index a00b7dc1d50d..7a6acaaca8a2 100644 --- a/backend/internal/service/gateway_multiplatform_test.go +++ b/backend/internal/service/gateway_multiplatform_test.go @@ -1229,6 +1229,106 @@ func TestGatewayService_selectAccountWithMixedScheduling(t *testing.T) { require.Equal(t, int64(2), acc.ID, "应选择优先级最高的账户（包含启用混合调度的antigravity）") }) + t.Run("混合调度-Gemini家族限流后跳过Antigravity账户", func(t *testing.T) { + resetAt := time.Now().Add(10 * time.Minute).Format(time.RFC3339) + repo := &mockAccountRepoForPlatform{ + accounts: []Account{ + { + ID: 1, + Platform: PlatformAntigravity, + Priority: 1, + Status: StatusActive, + Schedulable: true, + Extra: map[string]any{ + "mixed_scheduling": true, + modelRateLimitsKey: map[string]any{ + antigravityGeminiModelRateLimitKey: map[string]any{ + "rate_limit_reset_at": resetAt, + }, + }, + }, + }, + { + ID: 2, + Platform: PlatformAntigravity, + Priority: 1, + Status: StatusActive, + Schedulable: true, + Extra: map[string]any{ + "mixed_scheduling": true, + modelRateLimitsKey: map[string]any{ + antigravityGeminiModelRateLimitKey: map[string]any{ + "rate_limit_reset_at": resetAt, + }, + }, + }, + }, + { + ID: 3, + Platform: PlatformAntigravity, + Priority: 2, + Status: StatusActive, + Schedulable: true, + Extra: map[string]any{"mixed_scheduling": true}, + }, + }, + accountsByID: map[int64]*Account{}, + } + for i := range repo.accounts { + repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i] + } + + svc := &GatewayService{ + accountRepo: repo, + cache: &mockGatewayCacheForPlatform{}, + cfg: testConfig(), + } + + acc, err := svc.selectAccountWithMixedScheduling(ctx, nil, "", "gemini-3-pro-preview", nil, PlatformGemini) + require.NoError(t, err) + require.NotNil(t, acc) + require.Equal(t, int64(3), acc.ID) + }) + + t.Run("混合调度-Gemini家族限流不影响Claude调度", func(t *testing.T) { + resetAt := time.Now().Add(10 * time.Minute).Format(time.RFC3339) + repo := &mockAccountRepoForPlatform{ + accounts: []Account{ + { + ID: 1, + Platform: PlatformAntigravity, + Priority: 1, + Status: StatusActive, + Schedulable: true, + Extra: map[string]any{ + "mixed_scheduling": true, + modelRateLimitsKey: map[string]any{ + antigravityGeminiModelRateLimitKey: map[string]any{ + "rate_limit_reset_at": resetAt, + }, + }, + }, + }, + {ID: 2, Platform: PlatformAnthropic, Priority: 2, Status: StatusActive, Schedulable: true}, + }, + accountsByID: map[int64]*Account{}, + } + for i := range repo.accounts { + repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i] + } + + svc := &GatewayService{ + accountRepo: repo, + cache: &mockGatewayCacheForPlatform{}, + cfg: testConfig(), + } + + acc, err := svc.selectAccountWithMixedScheduling(ctx, nil, "", "claude-sonnet-4-5", nil, PlatformAnthropic) + require.NoError(t, err) + require.NotNil(t, acc) + require.Equal(t, int64(1), acc.ID) + }) + t.Run("混合调度-路由优先选择路由账号", func(t *testing.T) { groupID := int64(30) requestedModel := "claude-sonnet-4-5" diff --git a/backend/internal/service/model_rate_limit.go b/backend/internal/service/model_rate_limit.go index c45615cc4b41..420f14467c4e 100644 --- a/backend/internal/service/model_rate_limit.go +++ b/backend/internal/service/model_rate_limit.go @@ -6,7 +6,10 @@ import ( "time" ) -const modelRateLimitsKey = "model_rate_limits" +const ( + modelRateLimitsKey = "model_rate_limits" + antigravityGeminiModelRateLimitKey = "antigravity:gemini" +) // isRateLimitActiveForKey 检查指定 key 的限流是否生效 func (a *Account) isRateLimitActiveForKey(key string) bool { @@ -35,6 +38,9 @@ func (a *Account) isModelRateLimitedWithContext(ctx context.Context, requestedMo modelKey := a.GetMappedModel(requestedModel) if a.Platform == PlatformAntigravity { modelKey = resolveFinalAntigravityModelKey(ctx, a, requestedModel) + if isAntigravityGeminiModel(modelKey) && a.isRateLimitActiveForKey(antigravityGeminiModelRateLimitKey) { + return true + } } modelKey = strings.TrimSpace(modelKey) if modelKey == "" { @@ -62,7 +68,13 @@ func (a *Account) GetModelRateLimitRemainingTimeWithContext(ctx context.Context, if modelKey == "" { return 0 } - return a.getRateLimitRemainingForKey(modelKey) + remaining := a.getRateLimitRemainingForKey(modelKey) + if a.Platform == PlatformAntigravity && isAntigravityGeminiModel(modelKey) { + if familyRemaining := a.getRateLimitRemainingForKey(antigravityGeminiModelRateLimitKey); familyRemaining > remaining { + return familyRemaining + } + } + return remaining } func resolveFinalAntigravityModelKey(ctx context.Context, account *Account, requestedModel string) string { @@ -77,6 +89,22 @@ func resolveFinalAntigravityModelKey(ctx context.Context, account *Account, requ return modelKey } +func isAntigravityGeminiModel(model string) bool { + return strings.HasPrefix(normalizeAntigravityModelName(model), "gemini-") +} + +func antigravityModelRateLimitKeys(model string) []string { + model = strings.TrimSpace(model) + if model == "" { + return nil + } + keys := []string{model} + if isAntigravityGeminiModel(model) && model != antigravityGeminiModelRateLimitKey { + keys = append(keys, antigravityGeminiModelRateLimitKey) + } + return keys +} + func (a *Account) modelRateLimitResetAt(scope string) *time.Time { if a == nil || a.Extra == nil || scope == "" { return nil diff --git a/backend/internal/service/model_rate_limit_test.go b/backend/internal/service/model_rate_limit_test.go index b79b9688ac35..3cce645981c5 100644 --- a/backend/internal/service/model_rate_limit_test.go +++ b/backend/internal/service/model_rate_limit_test.go @@ -121,6 +121,36 @@ func TestIsModelRateLimited(t *testing.T) { requestedModel: "gemini-3-pro-preview", expected: true, }, + { + name: "antigravity platform - gemini family rate limit blocks mapped preview", + account: &Account{ + Platform: PlatformAntigravity, + Extra: map[string]any{ + modelRateLimitsKey: map[string]any{ + antigravityGeminiModelRateLimitKey: map[string]any{ + "rate_limit_reset_at": future, + }, + }, + }, + }, + requestedModel: "gemini-3-pro-preview", + expected: true, + }, + { + name: "antigravity platform - gemini family rate limit does not block claude", + account: &Account{ + Platform: PlatformAntigravity, + Extra: map[string]any{ + modelRateLimitsKey: map[string]any{ + antigravityGeminiModelRateLimitKey: map[string]any{ + "rate_limit_reset_at": future, + }, + }, + }, + }, + requestedModel: "claude-sonnet-4-5", + expected: false, + }, { name: "non-antigravity platform - gemini-3-pro-preview NOT mapped", account: &Account{ @@ -306,6 +336,38 @@ func TestGetModelRateLimitRemainingTime(t *testing.T) { minExpected: 4 * time.Minute, maxExpected: 6 * time.Minute, }, + { + name: "antigravity platform - gemini family rate limit remaining", + account: &Account{ + Platform: PlatformAntigravity, + Extra: map[string]any{ + modelRateLimitsKey: map[string]any{ + antigravityGeminiModelRateLimitKey: map[string]any{ + "rate_limit_reset_at": future10m, + }, + }, + }, + }, + requestedModel: "gemini-3-pro-preview", + minExpected: 9 * time.Minute, + maxExpected: 11 * time.Minute, + }, + { + name: "antigravity platform - gemini family remaining ignored for claude", + account: &Account{ + Platform: PlatformAntigravity, + Extra: map[string]any{ + modelRateLimitsKey: map[string]any{ + antigravityGeminiModelRateLimitKey: map[string]any{ + "rate_limit_reset_at": future10m, + }, + }, + }, + }, + requestedModel: "claude-sonnet-4-5", + minExpected: 0, + maxExpected: 0, + }, } for _, tt := range tests { diff --git a/backend/internal/service/scheduler_snapshot_hydration_test.go b/backend/internal/service/scheduler_snapshot_hydration_test.go index 778cab235165..0a1d0a0a52e7 100644 --- a/backend/internal/service/scheduler_snapshot_hydration_test.go +++ b/backend/internal/service/scheduler_snapshot_hydration_test.go @@ -6,6 +6,8 @@ import ( "context" "testing" "time" + + "github.com/Wei-Shaw/sub2api/internal/config" ) type snapshotHydrationCache struct { @@ -186,3 +188,89 @@ func TestGatewaySelectAccountWithLoadAwareness_HydratesSelectedAccountFromSchedu t.Fatalf("expected hydrated api key, got %q", got) } } + +func TestGatewaySelectAccountWithLoadAwareness_SkipsAntigravityGeminiFamilyRateLimitedSnapshot(t *testing.T) { + resetAt := time.Now().Add(10 * time.Minute).Format(time.RFC3339) + cache := &snapshotHydrationCache{ + snapshot: []*Account{ + { + ID: 1, + Platform: PlatformAntigravity, + Type: AccountTypeOAuth, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 1, + AccountGroups: []AccountGroup{ + {AccountID: 1, GroupID: 22}, + }, + GroupIDs: []int64{22}, + Extra: map[string]any{ + "mixed_scheduling": true, + modelRateLimitsKey: map[string]any{ + antigravityGeminiModelRateLimitKey: map[string]any{ + "rate_limit_reset_at": resetAt, + }, + }, + }, + }, + { + ID: 2, + Platform: PlatformAntigravity, + Type: AccountTypeOAuth, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 2, + AccountGroups: []AccountGroup{ + {AccountID: 2, GroupID: 22}, + }, + GroupIDs: []int64{22}, + Extra: map[string]any{ + "mixed_scheduling": true, + }, + }, + }, + accounts: map[int64]*Account{ + 1: {ID: 1, Platform: PlatformAntigravity, Type: AccountTypeOAuth}, + 2: {ID: 2, Platform: PlatformAntigravity, Type: AccountTypeOAuth}, + }, + } + groupID := int64(22) + svc := &GatewayService{ + schedulerSnapshot: NewSchedulerSnapshotService(cache, nil, nil, nil, nil), + groupRepo: &mockGroupRepoForGateway{ + groups: map[int64]*Group{ + groupID: { + ID: groupID, + Platform: PlatformGemini, + Status: StatusActive, + Hydrated: true, + }, + }, + }, + concurrencyService: NewConcurrencyService(&mockConcurrencyCache{}), + cfg: &config.Config{ + Gateway: config.GatewayConfig{ + Scheduling: config.GatewaySchedulingConfig{ + LoadBatchEnabled: true, + StickySessionMaxWaiting: 3, + StickySessionWaitTimeout: time.Second, + FallbackWaitTimeout: time.Second, + FallbackMaxWaiting: 10, + }, + }, + }, + } + + result, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, "", "gemini-3-flash-preview", nil, "", 0) + if err != nil { + t.Fatalf("SelectAccountWithLoadAwareness error: %v", err) + } + if result == nil || result.Account == nil { + t.Fatalf("expected selected account") + } + if result.Account.ID != 2 { + t.Fatalf("expected scheduler to skip Gemini-family limited antigravity account 1, got %d", result.Account.ID) + } +} From 0560340bd4c748c436053febd9009674301ecfb2 Mon Sep 17 00:00:00 2001 From: fatelei Date: Mon, 1 Jun 2026 08:41:35 +0800 Subject: [PATCH 32/36] fix: change balance to pointer type --- .../internal/handler/admin/user_handler.go | 16 +++--- backend/internal/service/admin_service.go | 11 +++- .../service/admin_service_create_user_test.go | 55 ++++++++++++++++++- frontend/src/api/admin/users.ts | 3 + .../components/admin/user/UserCreateModal.vue | 14 +++-- 5 files changed, 83 insertions(+), 16 deletions(-) diff --git a/backend/internal/handler/admin/user_handler.go b/backend/internal/handler/admin/user_handler.go index db35472e99d3..061bc1b2b632 100644 --- a/backend/internal/handler/admin/user_handler.go +++ b/backend/internal/handler/admin/user_handler.go @@ -34,14 +34,14 @@ func NewUserHandler(adminService service.AdminService, concurrencyService *servi // CreateUserRequest represents admin create user request type CreateUserRequest struct { - Email string `json:"email" binding:"required,email"` - Password string `json:"password" binding:"required,min=6"` - Username string `json:"username"` - Notes string `json:"notes"` - Balance float64 `json:"balance"` - Concurrency int `json:"concurrency"` - RPMLimit int `json:"rpm_limit"` - AllowedGroups []int64 `json:"allowed_groups"` + Email string `json:"email" binding:"required,email"` + Password string `json:"password" binding:"required,min=6"` + Username string `json:"username"` + Notes string `json:"notes"` + Balance *float64 `json:"balance"` + Concurrency int `json:"concurrency"` + RPMLimit int `json:"rpm_limit"` + AllowedGroups []int64 `json:"allowed_groups"` } // UpdateUserRequest represents admin update user request diff --git a/backend/internal/service/admin_service.go b/backend/internal/service/admin_service.go index eb5994d54988..09aeef867895 100644 --- a/backend/internal/service/admin_service.go +++ b/backend/internal/service/admin_service.go @@ -120,7 +120,7 @@ type CreateUserInput struct { Password string Username string Notes string - Balance float64 + Balance *float64 Concurrency int RPMLimit int AllowedGroups []int64 @@ -661,12 +661,19 @@ func (s *adminServiceImpl) GetUser(ctx context.Context, id int64) (*User, error) } func (s *adminServiceImpl) CreateUser(ctx context.Context, input *CreateUserInput) (*User, error) { + balance := 0.0 + if input.Balance != nil { + balance = *input.Balance + } else if s.settingService != nil { + balance = s.settingService.GetDefaultBalance(ctx) + } + user := &User{ Email: input.Email, Username: input.Username, Notes: input.Notes, Role: RoleUser, // Always create as regular user, never admin - Balance: input.Balance, + Balance: balance, Concurrency: input.Concurrency, RPMLimit: input.RPMLimit, Status: StatusActive, diff --git a/backend/internal/service/admin_service_create_user_test.go b/backend/internal/service/admin_service_create_user_test.go index c5b1e38d3903..5e9578abecca 100644 --- a/backend/internal/service/admin_service_create_user_test.go +++ b/backend/internal/service/admin_service_create_user_test.go @@ -14,13 +14,14 @@ import ( func TestAdminService_CreateUser_Success(t *testing.T) { repo := &userRepoStub{nextID: 10} svc := &adminServiceImpl{userRepo: repo} + balance := 12.5 input := &CreateUserInput{ Email: "user@test.com", Password: "strong-pass", Username: "tester", Notes: "note", - Balance: 12.5, + Balance: &balance, Concurrency: 7, AllowedGroups: []int64{3, 5}, } @@ -32,7 +33,7 @@ func TestAdminService_CreateUser_Success(t *testing.T) { require.Equal(t, input.Email, user.Email) require.Equal(t, input.Username, user.Username) require.Equal(t, input.Notes, user.Notes) - require.Equal(t, input.Balance, user.Balance) + require.Equal(t, balance, user.Balance) require.Equal(t, input.Concurrency, user.Concurrency) require.Equal(t, input.AllowedGroups, user.AllowedGroups) require.Equal(t, RoleUser, user.Role) @@ -42,6 +43,56 @@ func TestAdminService_CreateUser_Success(t *testing.T) { require.Equal(t, user, repo.created[0]) } +func TestAdminService_CreateUser_UsesDefaultBalanceWhenBalanceOmitted(t *testing.T) { + repo := &userRepoStub{nextID: 11} + cfg := &config.Config{ + Default: config.DefaultConfig{ + UserBalance: 0, + }, + } + settingService := NewSettingService(&settingRepoStub{values: map[string]string{ + SettingKeyDefaultBalance: "0.02", + }}, cfg) + svc := &adminServiceImpl{userRepo: repo, settingService: settingService} + + user, err := svc.CreateUser(context.Background(), &CreateUserInput{ + Email: "default-balance@test.com", + Password: "strong-pass", + }) + + require.NoError(t, err) + require.NotNil(t, user) + require.Equal(t, 0.02, user.Balance) + require.Len(t, repo.created, 1) + require.Equal(t, 0.02, repo.created[0].Balance) +} + +func TestAdminService_CreateUser_ExplicitZeroBalanceOverridesDefault(t *testing.T) { + repo := &userRepoStub{nextID: 12} + cfg := &config.Config{ + Default: config.DefaultConfig{ + UserBalance: 0, + }, + } + settingService := NewSettingService(&settingRepoStub{values: map[string]string{ + SettingKeyDefaultBalance: "0.02", + }}, cfg) + svc := &adminServiceImpl{userRepo: repo, settingService: settingService} + balance := 0.0 + + user, err := svc.CreateUser(context.Background(), &CreateUserInput{ + Email: "zero-balance@test.com", + Password: "strong-pass", + Balance: &balance, + }) + + require.NoError(t, err) + require.NotNil(t, user) + require.Equal(t, 0.0, user.Balance) + require.Len(t, repo.created, 1) + require.Equal(t, 0.0, repo.created[0].Balance) +} + func TestAdminService_CreateUser_EmailExists(t *testing.T) { repo := &userRepoStub{createErr: ErrEmailExists} svc := &adminServiceImpl{userRepo: repo} diff --git a/frontend/src/api/admin/users.ts b/frontend/src/api/admin/users.ts index fabc69bcef01..c33eaceeede7 100644 --- a/frontend/src/api/admin/users.ts +++ b/frontend/src/api/admin/users.ts @@ -115,8 +115,11 @@ export async function getById(id: number): Promise { export async function create(userData: { email: string password: string + username?: string + notes?: string balance?: number concurrency?: number + rpm_limit?: number allowed_groups?: number[] | null }): Promise { const { data } = await apiClient.post('/admin/users', userData) diff --git a/frontend/src/components/admin/user/UserCreateModal.vue b/frontend/src/components/admin/user/UserCreateModal.vue index 2966a23bee5c..a638e79cf526 100644 --- a/frontend/src/components/admin/user/UserCreateModal.vue +++ b/frontend/src/components/admin/user/UserCreateModal.vue @@ -28,7 +28,7 @@
- +
@@ -69,18 +69,24 @@ import Icon from '@/components/icons/Icon.vue' const props = defineProps<{ show: boolean }>() const emit = defineEmits(['close', 'success']); const { t } = useI18n() -const form = reactive({ email: '', password: '', username: '', notes: '', balance: 0, concurrency: 1, rpm_limit: 0 }) +const form = reactive({ email: '', password: '', username: '', notes: '', balance: '', concurrency: 1, rpm_limit: 0 }) const { loading, submit } = useForm({ form, submitFn: async (data) => { - await adminAPI.users.create(data) + const { balance: rawBalance, ...rest } = data + const balance = String(rawBalance).trim() + const payload: typeof rest & { balance?: number } = { ...rest } + if (balance !== '') { + payload.balance = Number(balance) + } + await adminAPI.users.create(payload) emit('success'); emit('close') }, successMsg: t('admin.users.userCreated') }) -watch(() => props.show, (v) => { if(v) Object.assign(form, { email: '', password: '', username: '', notes: '', balance: 0, concurrency: 1, rpm_limit: 0 }) }) +watch(() => props.show, (v) => { if(v) Object.assign(form, { email: '', password: '', username: '', notes: '', balance: '', concurrency: 1, rpm_limit: 0 }) }) const generateRandomPassword = () => { const chars = 'ABCDEFGHJKLMNPQRSTUVWXYZabcdefghjkmnpqrstuvwxyz23456789!@#$%^&*' From c8cd91e3ce1b588c374e310ddcb80bd975e7af0d Mon Sep 17 00:00:00 2001 From: wucm667 Date: Sun, 31 May 2026 08:47:13 +0800 Subject: [PATCH 33/36] =?UTF-8?q?test(openai):=20=E8=A6=86=E7=9B=96=20fail?= =?UTF-8?q?over=20=E8=AF=B7=E6=B1=82=E4=BD=93=E9=87=8D=E6=98=A0=E5=B0=84?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../openai_failover_cached_body_test.go | 134 ++++++++++++++++++ 1 file changed, 134 insertions(+) create mode 100644 backend/internal/service/openai_failover_cached_body_test.go diff --git a/backend/internal/service/openai_failover_cached_body_test.go b/backend/internal/service/openai_failover_cached_body_test.go new file mode 100644 index 000000000000..776f1b1ac56c --- /dev/null +++ b/backend/internal/service/openai_failover_cached_body_test.go @@ -0,0 +1,134 @@ +package service + +import ( + "bytes" + "context" + "errors" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" +) + +func TestOpenAIGatewayService_Forward_FailoverReparsesCachedBodyForNextAccount(t *testing.T) { + gin.SetMode(gin.TestMode) + + tests := []struct { + name string + requestModel string + firstMapping map[string]any + secondMapping map[string]any + wantFirst string + wantSecond string + }{ + { + name: "both accounts have mapping", + firstMapping: map[string]any{"alias-model": "base-model-a"}, + secondMapping: map[string]any{"alias-model": "base-model-b"}, + wantFirst: "base-model-a", + wantSecond: "base-model-b", + }, + { + name: "first account has mapping second account has none", + requestModel: "gpt-5.4-high", + firstMapping: map[string]any{"gpt-5.4-high": "gpt-5.4"}, + wantFirst: "gpt-5.4", + wantSecond: "gpt-5.4", + }, + { + name: "first account has no mapping second account has mapping", + secondMapping: map[string]any{"alias-model": "base-model-b"}, + wantFirst: "alias-model", + wantSecond: "base-model-b", + }, + { + name: "legacy context cache is ignored when mappings differ", + firstMapping: map[string]any{"alias-model": "base-model-a"}, + secondMapping: map[string]any{"alias-model": "base-model-b"}, + wantFirst: "base-model-a", + wantSecond: "base-model-b", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + requestModel := tt.requestModel + if requestModel == "" { + requestModel = "alias-model" + } + body := []byte(`{"model":"` + requestModel + `","stream":false,"instructions":"cache-test","input":"hello"}`) + + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body)) + c.Request.Header.Set("Content-Type", "application/json") + + upstream := &httpUpstreamRecorder{responses: []*http.Response{ + { + StatusCode: http.StatusTooManyRequests, + Header: http.Header{"Content-Type": []string{"application/json"}, "x-request-id": []string{"rid-failover-a"}}, + Body: io.NopCloser(strings.NewReader(`{"error":{"type":"rate_limit_error","message":"rate limited"}}`)), + }, + { + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}, "x-request-id": []string{"rid-ok-b"}}, + Body: io.NopCloser(strings.NewReader(`{"id":"resp_123","status":"completed","model":"ok","output":[],"usage":{"input_tokens":1,"output_tokens":1}}`)), + }, + }} + svc := &OpenAIGatewayService{httpUpstream: upstream} + + firstAccount := openAIFailoverCachedBodyTestAccount(1, "account-a", tt.firstMapping) + secondAccount := openAIFailoverCachedBodyTestAccount(2, "account-b", tt.secondMapping) + + _, err := svc.Forward(context.Background(), c, firstAccount, body) + require.Error(t, err) + var failoverErr *UpstreamFailoverError + require.True(t, errors.As(err, &failoverErr)) + require.Len(t, upstream.bodies, 1) + require.Equal(t, tt.wantFirst, gjson.GetBytes(upstream.bodies[0], "model").String()) + + c.Set("openai_parsed_request_body", map[string]any{"model": tt.wantFirst, "stream": true}) + result, err := svc.Forward(context.Background(), c, secondAccount, body) + require.NoError(t, err) + require.NotNil(t, result) + require.Len(t, upstream.bodies, 2) + require.Equal(t, tt.wantSecond, gjson.GetBytes(upstream.bodies[1], "model").String()) + }) + } +} + +func TestGetOpenAIRequestBodyMap_IgnoresLegacyContextCache(t *testing.T) { + gin.SetMode(gin.TestMode) + + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Set("openai_parsed_request_body", map[string]any{"model": "base-model-a", "stream": true}) + + got, err := getOpenAIRequestBodyMap(c, []byte(`{"model":"alias-model","stream":false}`)) + require.NoError(t, err) + require.Equal(t, "alias-model", got["model"]) + require.Equal(t, false, got["stream"]) +} + +func openAIFailoverCachedBodyTestAccount(id int64, name string, mapping map[string]any) *Account { + credentials := map[string]any{"access_token": "oauth-token", "chatgpt_account_id": "chatgpt-account"} + if mapping != nil { + credentials["model_mapping"] = mapping + } + return &Account{ + ID: id, + Name: name, + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Concurrency: 1, + Credentials: credentials, + Status: StatusActive, + Schedulable: true, + RateMultiplier: f64p(1), + } +} From 08e19bb15c2091885dbe214297d168c045c474c6 Mon Sep 17 00:00:00 2001 From: xlx0852 <33200995+xlx0852@users.noreply.github.com> Date: Mon, 1 Jun 2026 10:32:32 +0800 Subject: [PATCH 34/36] fix(openai): bridge oversized websocket requests --- backend/internal/config/config.go | 18 + backend/internal/config/config_test.go | 9 + .../handler/openai_gateway_handler.go | 2 +- .../service/openai_gateway_service.go | 3 + .../internal/service/openai_ws_forwarder.go | 258 ++++++++-- ...penai_ws_forwarder_ingress_session_test.go | 155 ++++++ .../internal/service/openai_ws_http_bridge.go | 387 +++++++++++++++ .../service/openai_ws_http_bridge_test.go | 461 ++++++++++++++++++ 8 files changed, 1253 insertions(+), 40 deletions(-) create mode 100644 backend/internal/service/openai_ws_http_bridge.go create mode 100644 backend/internal/service/openai_ws_http_bridge_test.go diff --git a/backend/internal/config/config.go b/backend/internal/config/config.go index df9dcefc6017..0df7d09e84a4 100644 --- a/backend/internal/config/config.go +++ b/backend/internal/config/config.go @@ -885,6 +885,12 @@ type GatewayOpenAIWSConfig struct { StoreDisabledForceNewConn bool `mapstructure:"store_disabled_force_new_conn"` // PrewarmGenerateEnabled: 是否启用 WSv2 generate=false 预热（默认 false） PrewarmGenerateEnabled bool `mapstructure:"prewarm_generate_enabled"` + // ClientReadLimitBytes: 入站客户端 WS 单帧读取上限。 + ClientReadLimitBytes int64 `mapstructure:"client_read_limit_bytes"` + // HTTPBridgeEnabled: 首包过大时，保持客户端 WS，改用 HTTP Responses 上游。 + HTTPBridgeEnabled bool `mapstructure:"http_bridge_enabled"` + // HTTPBridgeThresholdBytes: 触发 HTTP bridge 的入站 WS payload 阈值。 + HTTPBridgeThresholdBytes int64 `mapstructure:"http_bridge_threshold_bytes"` // Feature 开关：v2 优先于 v1 ResponsesWebsockets bool `mapstructure:"responses_websockets"` @@ -1806,6 +1812,9 @@ func setDefaults() { viper.SetDefault("gateway.openai_ws.store_disabled_conn_mode", "strict") viper.SetDefault("gateway.openai_ws.store_disabled_force_new_conn", true) viper.SetDefault("gateway.openai_ws.prewarm_generate_enabled", false) + viper.SetDefault("gateway.openai_ws.client_read_limit_bytes", 64*1024*1024) + viper.SetDefault("gateway.openai_ws.http_bridge_enabled", true) + viper.SetDefault("gateway.openai_ws.http_bridge_threshold_bytes", 15*1024*1024) viper.SetDefault("gateway.openai_ws.responses_websockets", false) viper.SetDefault("gateway.openai_ws.responses_websockets_v2", true) viper.SetDefault("gateway.openai_ws.max_conns_per_account", 128) @@ -2543,6 +2552,15 @@ func (c *Config) Validate() error { if c.Gateway.OpenAIWS.PrewarmCooldownMS < 0 { return fmt.Errorf("gateway.openai_ws.prewarm_cooldown_ms must be non-negative") } + if c.Gateway.OpenAIWS.ClientReadLimitBytes <= 0 { + return fmt.Errorf("gateway.openai_ws.client_read_limit_bytes must be positive") + } + if c.Gateway.OpenAIWS.HTTPBridgeThresholdBytes < 0 { + return fmt.Errorf("gateway.openai_ws.http_bridge_threshold_bytes must be non-negative") + } + if c.Gateway.OpenAIWS.HTTPBridgeEnabled && c.Gateway.OpenAIWS.HTTPBridgeThresholdBytes == 0 { + return fmt.Errorf("gateway.openai_ws.http_bridge_threshold_bytes must be positive when http_bridge_enabled is true") + } if c.Gateway.OpenAIWS.FallbackCooldownSeconds < 0 { return fmt.Errorf("gateway.openai_ws.fallback_cooldown_seconds must be non-negative") } diff --git a/backend/internal/config/config_test.go b/backend/internal/config/config_test.go index 1eae5ed95953..9478b5102dae 100644 --- a/backend/internal/config/config_test.go +++ b/backend/internal/config/config_test.go @@ -134,6 +134,15 @@ func TestLoadDefaultOpenAIWSConfig(t *testing.T) { if cfg.Gateway.OpenAIWS.PrewarmCooldownMS != 300 { t.Fatalf("Gateway.OpenAIWS.PrewarmCooldownMS = %d, want 300", cfg.Gateway.OpenAIWS.PrewarmCooldownMS) } + if cfg.Gateway.OpenAIWS.ClientReadLimitBytes != 64*1024*1024 { + t.Fatalf("Gateway.OpenAIWS.ClientReadLimitBytes = %d, want %d", cfg.Gateway.OpenAIWS.ClientReadLimitBytes, 64*1024*1024) + } + if !cfg.Gateway.OpenAIWS.HTTPBridgeEnabled { + t.Fatalf("Gateway.OpenAIWS.HTTPBridgeEnabled = false, want true") + } + if cfg.Gateway.OpenAIWS.HTTPBridgeThresholdBytes != 15*1024*1024 { + t.Fatalf("Gateway.OpenAIWS.HTTPBridgeThresholdBytes = %d, want %d", cfg.Gateway.OpenAIWS.HTTPBridgeThresholdBytes, 15*1024*1024) + } if cfg.Gateway.OpenAIWS.RetryBackoffInitialMS != 120 { t.Fatalf("Gateway.OpenAIWS.RetryBackoffInitialMS = %d, want 120", cfg.Gateway.OpenAIWS.RetryBackoffInitialMS) } diff --git a/backend/internal/handler/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index 0aa477b08cd8..f3d4caf08dbf 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -1167,7 +1167,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { defer func() { _ = wsConn.CloseNow() }() - wsConn.SetReadLimit(16 * 1024 * 1024) + wsConn.SetReadLimit(service.ResolveOpenAIWSClientReadLimitBytes(h.cfg)) ctx := c.Request.Context() readCtx, cancel := context.WithTimeout(ctx, 30*time.Second) diff --git a/backend/internal/service/openai_gateway_service.go b/backend/internal/service/openai_gateway_service.go index 89ddaa7d1c5c..b1a9559452b5 100644 --- a/backend/internal/service/openai_gateway_service.go +++ b/backend/internal/service/openai_gateway_service.go @@ -256,6 +256,9 @@ type OpenAIForwardResult struct { ImageOutputSizes []string ImageSizeSource string ImageSizeBreakdown map[string]int + + wsReplayInput []json.RawMessage + wsReplayInputExists bool } type OpenAIWSRetryMetricsSnapshot struct { diff --git a/backend/internal/service/openai_ws_forwarder.go b/backend/internal/service/openai_ws_forwarder.go index dd2f45facbd3..66c66134b0ee 100644 --- a/backend/internal/service/openai_ws_forwarder.go +++ b/backend/internal/service/openai_ws_forwarder.go @@ -1560,6 +1560,38 @@ func openAIWSRawItemsHasFunctionCallOutput(items []json.RawMessage) bool { return false } +func openAIWSRawItemsHaveToolCallContextForOutputs(items []json.RawMessage) bool { + if len(items) == 0 { + return false + } + contextCallIDs := make(map[string]struct{}) + outputCallIDs := make(map[string]struct{}) + for _, item := range items { + itemType := gjson.GetBytes(item, "type").String() + callID := strings.TrimSpace(gjson.GetBytes(item, "call_id").String()) + switch { + case isCodexToolCallContextItemType(itemType): + if callID != "" { + contextCallIDs[callID] = struct{}{} + } + case isCodexToolCallOutputItemType(itemType): + if callID == "" { + return false + } + outputCallIDs[callID] = struct{}{} + } + } + if len(outputCallIDs) == 0 || len(contextCallIDs) == 0 { + return false + } + for callID := range outputCallIDs { + if _, ok := contextCallIDs[callID]; !ok { + return false + } + } + return true +} + func openAIWSRawPayloadHasToolCallOutput(payload []byte) bool { if len(payload) == 0 { return false @@ -2664,6 +2696,27 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( }, nil } + writeClientMessage := func(message []byte) error { + writeCtx, cancel := context.WithTimeout(ctx, s.openAIWSWriteTimeout()) + defer cancel() + return clientConn.Write(writeCtx, coderws.MessageText, message) + } + + readClientMessage := func() ([]byte, error) { + msgType, payload, readErr := clientConn.Read(ctx) + if readErr != nil { + return nil, readErr + } + if msgType != coderws.MessageText && msgType != coderws.MessageBinary { + return nil, NewOpenAIWSClientCloseError( + coderws.StatusPolicyViolation, + fmt.Sprintf("unsupported websocket client message type: %s", msgType.String()), + nil, + ) + } + return payload, nil + } + firstPayload, err := parseClientPayload(firstClientMessage) if err != nil { return err @@ -2672,25 +2725,152 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( turnState := strings.TrimSpace(c.GetHeader(openAIWSTurnStateHeader)) stateStore := s.getOpenAIWSStateStore() groupID := getOpenAIGroupIDFromContext(c) - sessionHash := s.GenerateSessionHash(c, firstPayload.rawForHash) - if turnState == "" && stateStore != nil && sessionHash != "" { - if savedTurnState, ok := stateStore.GetSessionTurnState(groupID, sessionHash); ok { - turnState = savedTurnState + storeDisabledConnMode := s.openAIWSStoreDisabledConnMode() + sessionHash := "" + preferredConnID := "" + storeDisabled := false + refreshIngressRouteState := func(payload openAIWSClientPayload) { + sessionHash = s.GenerateSessionHash(c, payload.rawForHash) + if turnState == "" && stateStore != nil && sessionHash != "" { + if savedTurnState, ok := stateStore.GetSessionTurnState(groupID, sessionHash); ok { + turnState = savedTurnState + } } - } - preferredConnID := "" - if stateStore != nil && firstPayload.previousResponseID != "" { - if connID, ok := stateStore.GetResponseConn(firstPayload.previousResponseID); ok { - preferredConnID = connID + preferredConnID = "" + if stateStore != nil && payload.previousResponseID != "" { + if connID, ok := stateStore.GetResponseConn(payload.previousResponseID); ok { + preferredConnID = connID + } + } + + storeDisabled = s.isOpenAIWSStoreDisabledInRequestRaw(payload.payloadRaw, account) + if stateStore != nil && storeDisabled && payload.previousResponseID == "" && sessionHash != "" { + if connID, ok := stateStore.GetSessionConn(groupID, sessionHash); ok { + preferredConnID = connID + } } } + refreshIngressRouteState(firstPayload) - storeDisabled := s.isOpenAIWSStoreDisabledInRequestRaw(firstPayload.payloadRaw, account) - storeDisabledConnMode := s.openAIWSStoreDisabledConnMode() - if stateStore != nil && storeDisabled && firstPayload.previousResponseID == "" && sessionHash != "" { - if connID, ok := stateStore.GetSessionConn(groupID, sessionHash); ok { - preferredConnID = connID + if s.shouldBridgeOpenAIWSHTTP(firstPayload.payloadBytes, firstPayload.previousResponseID) { + logOpenAIWSModeInfo( + "ingress_ws_http_bridge_start account_id=%d account_type=%s payload_bytes=%d threshold_bytes=%d has_session_hash=%v store_disabled=%v", + account.ID, + account.Type, + firstPayload.payloadBytes, + s.openAIWSHTTPBridgeThresholdBytes(), + sessionHash != "", + storeDisabled, + ) + currentBridgePayload := firstPayload + var bridgeReplayInput []json.RawMessage + bridgeReplayInputExists := false + for turn := 1; ; turn++ { + if turn > 1 && hooks != nil && hooks.BeforeRequest != nil { + if err := hooks.BeforeRequest(turn, currentBridgePayload.payloadRaw, currentBridgePayload.originalModel); err != nil { + return err + } + } + if hooks != nil && hooks.BeforeTurn != nil { + if err := hooks.BeforeTurn(turn); err != nil { + return err + } + } + if turnState != "" && c != nil && c.Request != nil { + c.Request.Header.Set(openAIWSTurnStateHeader, turnState) + } + bridgePayloadRaw := currentBridgePayload.payloadRaw + bridgePayloadBytes := currentBridgePayload.payloadBytes + needsBridgeReplay := currentBridgePayload.previousResponseID != "" || openAIWSRawPayloadHasToolCallOutput(currentBridgePayload.payloadRaw) + turnReplayInput, turnReplayInputExists, replayInputErr := buildOpenAIWSReplayInputSequence( + bridgeReplayInput, + bridgeReplayInputExists, + currentBridgePayload.payloadRaw, + needsBridgeReplay, + ) + if replayInputErr != nil { + return fmt.Errorf("build websocket http bridge replay input: %w", replayInputErr) + } + if needsBridgeReplay && turnReplayInputExists { + updatedPayload, setInputErr := setOpenAIWSPayloadInputSequence( + currentBridgePayload.payloadRaw, + turnReplayInput, + true, + ) + if setInputErr != nil { + return fmt.Errorf("set websocket http bridge replay input: %w", setInputErr) + } + bridgePayloadRaw = updatedPayload + bridgePayloadBytes = len(updatedPayload) + logOpenAIWSModeInfo( + "ingress_ws_http_bridge_replay_input account_id=%d turn=%d input_items=%d previous_response_id_present=%v has_tool_output=%v", + account.ID, + turn, + len(turnReplayInput), + currentBridgePayload.previousResponseID != "", + openAIWSRawPayloadHasToolCallOutput(currentBridgePayload.payloadRaw), + ) + } + result, bridgeErr := s.proxyOpenAIWSHTTPBridgeTurn( + ctx, + c, + account, + token, + bridgePayloadRaw, + bridgePayloadBytes, + currentBridgePayload.originalModel, + currentBridgePayload.imageBillingModel, + currentBridgePayload.imageSizeTier, + currentBridgePayload.imageInputSize, + turn, + writeClientMessage, + ) + if hooks != nil && hooks.AfterTurn != nil { + hooks.AfterTurn(turn, result, bridgeErr) + } + if bridgeErr != nil { + return bridgeErr + } + if result == nil { + return errors.New("websocket http bridge turn result is nil") + } + bridgeReplayInput = cloneOpenAIWSRawMessages(turnReplayInput) + bridgeReplayInputExists = turnReplayInputExists + if result.wsReplayInputExists { + bridgeReplayInput = append(bridgeReplayInput, cloneOpenAIWSRawMessages(result.wsReplayInput)...) + bridgeReplayInputExists = true + } + if bridgeTurnState := strings.TrimSpace(result.ResponseHeaders.Get(openAIWSTurnStateHeader)); bridgeTurnState != "" { + turnState = bridgeTurnState + if stateStore != nil && sessionHash != "" { + stateStore.BindSessionTurnState(groupID, sessionHash, bridgeTurnState, s.openAIWSSessionStickyTTL()) + } + } + responseID := strings.TrimSpace(result.RequestID) + if responseID != "" && stateStore != nil { + ttl := s.openAIWSResponseStickyTTL() + logOpenAIWSBindResponseAccountWarn(groupID, account.ID, responseID, stateStore.BindResponseAccount(ctx, groupID, responseID, account.ID, ttl)) + } + nextClientMessage, readErr := readClientMessage() + if readErr != nil { + if isOpenAIWSClientDisconnectError(readErr) { + closeStatus, closeReason := summarizeOpenAIWSReadCloseError(readErr) + logOpenAIWSModeInfo( + "ingress_ws_http_bridge_client_closed account_id=%d close_status=%s close_reason=%s", + account.ID, + closeStatus, + truncateOpenAIWSLogValue(closeReason, openAIWSHeaderValueMaxLen), + ) + return nil + } + return fmt.Errorf("read client websocket request: %w", readErr) + } + nextPayload, parseErr := parseClientPayload(nextClientMessage) + if parseErr != nil { + return parseErr + } + currentBridgePayload = nextPayload } } @@ -2844,27 +3024,6 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( return lease, nil } - writeClientMessage := func(message []byte) error { - writeCtx, cancel := context.WithTimeout(ctx, s.openAIWSWriteTimeout()) - defer cancel() - return clientConn.Write(writeCtx, coderws.MessageText, message) - } - - readClientMessage := func() ([]byte, error) { - msgType, payload, readErr := clientConn.Read(ctx) - if readErr != nil { - return nil, readErr - } - if msgType != coderws.MessageText && msgType != coderws.MessageBinary { - return nil, NewOpenAIWSClientCloseError( - coderws.StatusPolicyViolation, - fmt.Sprintf("unsupported websocket client message type: %s", msgType.String()), - nil, - ) - } - return payload, nil - } - sendAndRelay := func(turn int, lease *openAIWSConnLease, payload []byte, payloadBytes int, originalModel string, imageBillingModel string, imageSizeTier string, imageInputSize string) (*OpenAIForwardResult, error) { if lease == nil { return nil, errors.New("upstream websocket lease is nil") @@ -2901,6 +3060,7 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( eventCount := 0 tokenEventCount := 0 terminalEventCount := 0 + replayCollector := &openAIWSToolCallReplayCollector{} firstEventType := "" lastEventType := "" needModelReplace := false @@ -3031,6 +3191,7 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( upstreamMessage = corrected } } + replayCollector.AddEvent(eventType, upstreamMessage) if err := writeClientMessage(upstreamMessage); err != nil { if isOpenAIWSClientDisconnectError(err) { clientDisconnected = true @@ -3094,6 +3255,10 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( Duration: time.Since(turnStart), FirstTokenMs: firstTokenMs, } + if replayInput := replayCollector.Items(); len(replayInput) > 0 { + result.wsReplayInput = replayInput + result.wsReplayInputExists = true + } if imageCount > 0 { result.ImageCount = imageCount result.ImageSize = imageSizeTier @@ -3487,9 +3652,12 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( if forcePreferredConn { // 携带 function_call_output 的请求不能丢弃 previous_response_id： // 上游 API 需要 response chain 来匹配 tool_result 与之前的 tool_use， - // 丢弃后会导致 "No tool call found for function call output" 400 错误。 + // 除非 replay input 已经包含与每个 tool_result 匹配的 tool_use 上下文。 hasFCOutput := hasFunctionCallOutput - if !turnPrevRecoveryTried && currentPreviousResponseID != "" && !hasFCOutput { + hasReplayToolContext := hasFCOutput && + currentTurnReplayInputExists && + openAIWSRawItemsHaveToolCallContextForOutputs(currentTurnReplayInput) + if !turnPrevRecoveryTried && currentPreviousResponseID != "" && (!hasFCOutput || hasReplayToolContext) { updatedPayload, removed, dropErr := dropPreviousResponseIDFromRawPayload(currentPayload) if dropErr != nil || !removed { reason := "not_removed" @@ -3521,11 +3689,13 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( ) } else { logOpenAIWSModeInfo( - "ingress_ws_preflight_ping_recovery account_id=%d turn=%d conn_id=%s action=drop_previous_response_id_retry previous_response_id=%s", + "ingress_ws_preflight_ping_recovery account_id=%d turn=%d conn_id=%s action=drop_previous_response_id_retry previous_response_id=%s has_function_call_output=%v has_replay_tool_context=%v", account.ID, turn, truncateOpenAIWSLogValue(sessionConnID, openAIWSIDValueMaxLen), truncateOpenAIWSLogValue(currentPreviousResponseID, openAIWSIDValueMaxLen), + hasFCOutput, + hasReplayToolContext, ) turnPrevRecoveryTried = true currentPayload = updatedWithInput @@ -3537,12 +3707,18 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( } } if hasFCOutput && currentPreviousResponseID != "" { + reason := "function_call_output_missing_replay_context" + if hasReplayToolContext { + reason = "function_call_output_replay_not_applied" + } logOpenAIWSModeInfo( - "ingress_ws_preflight_ping_recovery_skip account_id=%d turn=%d conn_id=%s reason=function_call_output action=fail_close previous_response_id=%s", + "ingress_ws_preflight_ping_recovery_skip account_id=%d turn=%d conn_id=%s reason=%s action=fail_close previous_response_id=%s has_replay_tool_context=%v", account.ID, turn, truncateOpenAIWSLogValue(sessionConnID, openAIWSIDValueMaxLen), + reason, truncateOpenAIWSLogValue(currentPreviousResponseID, openAIWSIDValueMaxLen), + hasReplayToolContext, ) } resetSessionLease(true) @@ -3622,6 +3798,10 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( lastTurnPayload = cloneOpenAIWSPayloadBytes(currentPayload) lastTurnReplayInput = cloneOpenAIWSRawMessages(currentTurnReplayInput) lastTurnReplayInputExists = currentTurnReplayInputExists + if result.wsReplayInputExists { + lastTurnReplayInput = append(lastTurnReplayInput, cloneOpenAIWSRawMessages(result.wsReplayInput)...) + lastTurnReplayInputExists = true + } nextStrictState, strictStateErr := buildOpenAIWSIngressPreviousTurnStrictState(currentPayload) if strictStateErr != nil { lastTurnStrictState = nil diff --git a/backend/internal/service/openai_ws_forwarder_ingress_session_test.go b/backend/internal/service/openai_ws_forwarder_ingress_session_test.go index b7f1bc4f78da..069a9dee1f3a 100644 --- a/backend/internal/service/openai_ws_forwarder_ingress_session_test.go +++ b/backend/internal/service/openai_ws_forwarder_ingress_session_test.go @@ -2305,6 +2305,161 @@ func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_StoreDisabledStr require.Equal(t, "world", gjson.Get(secondWrite, "input.1.text").String()) } +func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_StoreDisabledPreflightPingFailReplaysFunctionCallOutputWithContext(t *testing.T) { + gin.SetMode(gin.TestMode) + prevPreflightPingIdle := openAIWSIngressPreflightPingIdle + openAIWSIngressPreflightPingIdle = 0 + defer func() { + openAIWSIngressPreflightPingIdle = prevPreflightPingIdle + }() + + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + cfg.Security.URLAllowlist.AllowInsecureHTTP = true + cfg.Gateway.OpenAIWS.Enabled = true + cfg.Gateway.OpenAIWS.OAuthEnabled = true + cfg.Gateway.OpenAIWS.APIKeyEnabled = true + cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true + cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2 + cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0 + cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 2 + cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8 + cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3 + cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3 + cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3 + + firstConn := &openAIWSPreflightFailConn{ + events: [][]byte{ + []byte(`{"type":"response.completed","response":{"id":"resp_turn_ping_replay_ctx_1","model":"gpt-5.1","output":[{"type":"function_call","id":"fc_replay_1","call_id":"call_replay_1","name":"shell","arguments":"{}"}],"usage":{"input_tokens":1,"output_tokens":1}}}`), + }, + } + secondConn := &openAIWSCaptureConn{ + events: [][]byte{ + []byte(`{"type":"response.completed","response":{"id":"resp_turn_ping_replay_ctx_2","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`), + }, + } + dialer := &openAIWSQueueDialer{ + conns: []openAIWSClientConn{firstConn, secondConn}, + } + pool := newOpenAIWSConnPool(cfg) + pool.setClientDialerForTest(dialer) + + svc := &OpenAIGatewayService{ + cfg: cfg, + httpUpstream: &httpUpstreamRecorder{}, + cache: &stubGatewayCache{}, + openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), + toolCorrector: NewCodexToolCorrector(), + openaiWSPool: pool, + } + + account := &Account{ + ID: 128, + Name: "openai-ingress-preflight-replay-function-output-with-context", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "sk-test", + }, + Extra: map[string]any{ + "responses_websockets_v2_enabled": true, + }, + } + + serverErrCh := make(chan error, 1) + wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{ + CompressionMode: coderws.CompressionContextTakeover, + }) + if err != nil { + serverErrCh <- err + return + } + defer func() { + _ = conn.CloseNow() + }() + + rec := httptest.NewRecorder() + ginCtx, _ := gin.CreateTestContext(rec) + req := r.Clone(r.Context()) + req.Header = req.Header.Clone() + req.Header.Set("User-Agent", "unit-test-agent/1.0") + ginCtx.Request = req + + readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second) + msgType, firstMessage, readErr := conn.Read(readCtx) + cancel() + if readErr != nil { + serverErrCh <- readErr + return + } + if msgType != coderws.MessageText && msgType != coderws.MessageBinary { + serverErrCh <- errors.New("unsupported websocket client message type") + return + } + + serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, nil) + })) + defer wsServer.Close() + + dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second) + clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil) + cancelDial() + require.NoError(t, err) + defer func() { + _ = clientConn.CloseNow() + }() + + writeMessage := func(payload string) { + writeCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + require.NoError(t, clientConn.Write(writeCtx, coderws.MessageText, []byte(payload))) + } + readMessage := func() []byte { + readCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + msgType, message, readErr := clientConn.Read(readCtx) + require.NoError(t, readErr) + require.Equal(t, coderws.MessageText, msgType) + return message + } + + writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"store":false,"input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"call tool"}]}]}`) + firstTurn := readMessage() + require.Equal(t, "resp_turn_ping_replay_ctx_1", gjson.GetBytes(firstTurn, "response.id").String()) + + writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"store":false,"previous_response_id":"resp_turn_ping_replay_ctx_1","input":[{"type":"function_call_output","call_id":"call_replay_1","output":"ok"}]}`) + secondTurn := readMessage() + require.Equal(t, "resp_turn_ping_replay_ctx_2", gjson.GetBytes(secondTurn, "response.id").String()) + + require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done")) + select { + case serverErr := <-serverErrCh: + require.NoError(t, serverErr) + case <-time.After(5 * time.Second): + t.Fatal("等待 ingress websocket function_call_output 自包含重放后结束超时") + } + + require.Equal(t, 2, dialer.DialCount(), "带完整 tool 上下文的 function_call_output 应在 ping 失败后换新连接重放") + require.Equal(t, 1, firstConn.WriteCount()) + require.GreaterOrEqual(t, firstConn.PingCount(), 1) + secondConn.mu.Lock() + secondWrites := append([]map[string]any(nil), secondConn.writes...) + secondConn.mu.Unlock() + require.Len(t, secondWrites, 1) + secondWrite := requestToJSONString(secondWrites[0]) + require.False(t, gjson.Get(secondWrite, "previous_response_id").Exists()) + require.Equal(t, 3, len(gjson.Get(secondWrite, "input").Array())) + require.Equal(t, "message", gjson.Get(secondWrite, "input.0.type").String()) + require.Equal(t, "function_call", gjson.Get(secondWrite, "input.1.type").String()) + require.Equal(t, "call_replay_1", gjson.Get(secondWrite, "input.1.call_id").String()) + require.Equal(t, "function_call_output", gjson.Get(secondWrite, "input.2.type").String()) + require.Equal(t, "call_replay_1", gjson.Get(secondWrite, "input.2.call_id").String()) +} + func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_StoreDisabledPreflightPingFailClosesWhenFunctionCallOutputNeedsPreviousResponseID(t *testing.T) { gin.SetMode(gin.TestMode) prevPreflightPingIdle := openAIWSIngressPreflightPingIdle diff --git a/backend/internal/service/openai_ws_http_bridge.go b/backend/internal/service/openai_ws_http_bridge.go new file mode 100644 index 000000000000..1f0f32a05754 --- /dev/null +++ b/backend/internal/service/openai_ws_http_bridge.go @@ -0,0 +1,387 @@ +package service + +import ( + "bufio" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "strings" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/gin-gonic/gin" + "github.com/tidwall/gjson" +) + +const ( + openAIWSClientReadLimitBytesDefault int64 = 64 * 1024 * 1024 + openAIWSHTTPBridgeThresholdBytesDefault int64 = 15 * 1024 * 1024 + openAIWSHTTPBridgeErrorBodyLimitBytes = 64 * 1024 +) + +func ResolveOpenAIWSClientReadLimitBytes(cfg *config.Config) int64 { + if cfg == nil || cfg.Gateway.OpenAIWS.ClientReadLimitBytes <= 0 { + return openAIWSClientReadLimitBytesDefault + } + return cfg.Gateway.OpenAIWS.ClientReadLimitBytes +} + +func (s *OpenAIGatewayService) openAIWSHTTPBridgeEnabled() bool { + return s != nil && s.cfg != nil && s.cfg.Gateway.OpenAIWS.HTTPBridgeEnabled +} + +func (s *OpenAIGatewayService) openAIWSHTTPBridgeThresholdBytes() int64 { + if s == nil || s.cfg == nil || s.cfg.Gateway.OpenAIWS.HTTPBridgeThresholdBytes <= 0 { + return openAIWSHTTPBridgeThresholdBytesDefault + } + return s.cfg.Gateway.OpenAIWS.HTTPBridgeThresholdBytes +} + +func (s *OpenAIGatewayService) shouldBridgeOpenAIWSHTTP(payloadBytes int, previousResponseID string) bool { + if !s.openAIWSHTTPBridgeEnabled() { + return false + } + if strings.TrimSpace(previousResponseID) != "" { + return false + } + threshold := s.openAIWSHTTPBridgeThresholdBytes() + return threshold > 0 && int64(payloadBytes) >= threshold +} + +func prepareOpenAIWSHTTPBridgeBody(payload []byte) ([]byte, error) { + var body map[string]any + if err := json.Unmarshal(payload, &body); err != nil { + return nil, err + } + if body == nil { + return nil, errors.New("response.create payload must be a JSON object") + } + delete(body, "type") + delete(body, "generate") + delete(body, "previous_response_id") + body["stream"] = true + return json.Marshal(body) +} + +type openAIWSToolCallReplayCollector struct { + items []json.RawMessage + seen map[string]struct{} +} + +func (c *openAIWSToolCallReplayCollector) AddEvent(eventType string, message []byte) { + switch strings.TrimSpace(eventType) { + case "response.output_item.done": + c.addItem(gjson.GetBytes(message, "item")) + case "response.completed", "response.done": + output := gjson.GetBytes(message, "response.output") + if !output.IsArray() { + return + } + for _, item := range output.Array() { + c.addItem(item) + } + } +} + +func (c *openAIWSToolCallReplayCollector) Items() []json.RawMessage { + return cloneOpenAIWSRawMessages(c.items) +} + +func (c *openAIWSToolCallReplayCollector) addItem(item gjson.Result) { + if !item.Exists() || item.Type != gjson.JSON { + return + } + raw := strings.TrimSpace(item.Raw) + if raw == "" || !strings.HasPrefix(raw, "{") { + return + } + if !isCodexToolCallContextItemType(item.Get("type").String()) { + return + } + key := strings.TrimSpace(item.Get("id").String()) + if key == "" { + key = strings.TrimSpace(item.Get("call_id").String()) + } + if key == "" { + key = raw + } + if c.seen == nil { + c.seen = make(map[string]struct{}) + } + if _, ok := c.seen[key]; ok { + return + } + c.seen[key] = struct{}{} + c.items = append(c.items, json.RawMessage(raw)) +} + +func buildOpenAIWSHTTPBridgeErrorEvent(statusCode int, message string) []byte { + message = strings.TrimSpace(message) + if message == "" { + message = http.StatusText(statusCode) + } + if message == "" { + message = "upstream request failed" + } + event := map[string]any{ + "type": "error", + "status": statusCode, + "error": map[string]any{ + "type": "upstream_error", + "message": message, + }, + } + body, err := json.Marshal(event) + if err != nil { + return []byte(`{"type":"error","error":{"type":"upstream_error","message":"upstream request failed"}}`) + } + return body +} + +func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn( + ctx context.Context, + c *gin.Context, + account *Account, + token string, + payload []byte, + payloadBytes int, + originalModel string, + imageBillingModel string, + imageSizeTier string, + imageInputSize string, + turn int, + writeClientMessage func([]byte) error, +) (*OpenAIForwardResult, error) { + if s == nil { + return nil, errors.New("service is nil") + } + if s.httpUpstream == nil { + return nil, errors.New("openai http upstream is nil") + } + if account == nil { + return nil, errors.New("account is nil") + } + if writeClientMessage == nil { + return nil, errors.New("client websocket writer is nil") + } + + body, err := prepareOpenAIWSHTTPBridgeBody(payload) + if err != nil { + return nil, fmt.Errorf("prepare http bridge body: %w", err) + } + + upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx) + upstreamReq, err := s.buildUpstreamRequestOpenAIPassthrough(upstreamCtx, c, account, body, token) + releaseUpstreamCtx() + if err != nil { + return nil, err + } + + proxyURL := "" + if account.ProxyID != nil && account.Proxy != nil { + proxyURL = account.Proxy.URL() + } + if c != nil { + c.Set("openai_passthrough", true) + c.Set("openai_ws_http_bridge", true) + } + + turnStart := time.Now() + resp, err := s.httpUpstream.Do(upstreamReq, proxyURL, account.ID, account.Concurrency) + if err != nil { + safeErr := sanitizeUpstreamErrorMessage(err.Error()) + _ = writeClientMessage(buildOpenAIWSHTTPBridgeErrorEvent(http.StatusBadGateway, "Upstream request failed")) + return nil, fmt.Errorf("upstream http bridge request failed: %s", safeErr) + } + defer func() { _ = resp.Body.Close() }() + + if resp.StatusCode >= 400 { + respBody, _ := io.ReadAll(io.LimitReader(resp.Body, openAIWSHTTPBridgeErrorBodyLimitBytes)) + upstreamMsg := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractUpstreamErrorMessage(respBody))) + if upstreamMsg == "" { + upstreamMsg = http.StatusText(resp.StatusCode) + } + _ = writeClientMessage(buildOpenAIWSHTTPBridgeErrorEvent(resp.StatusCode, upstreamMsg)) + return nil, fmt.Errorf("upstream http bridge error: status=%d message=%s", resp.StatusCode, upstreamMsg) + } + + responseID := "" + usage := OpenAIUsage{} + imageCounter := newOpenAIImageOutputCounter() + var firstTokenMs *int + reqStream := openAIWSPayloadBoolFromRaw(body, "stream", true) + eventCount := 0 + tokenEventCount := 0 + terminalEventCount := 0 + replayCollector := &openAIWSToolCallReplayCollector{} + firstEventType := "" + lastEventType := "" + sawDone := false + wroteDownstream := false + clientDisconnected := false + mappedModel := "" + needModelReplace := false + var mappedModelBytes []byte + if originalModel != "" { + mappedModel = normalizeOpenAIModelForUpstream(account, account.GetMappedModel(originalModel)) + needModelReplace = mappedModel != "" && mappedModel != originalModel + if needModelReplace { + mappedModelBytes = []byte(mappedModel) + } + } + + resultWithUsage := func() *OpenAIForwardResult { + imageCount := imageCounter.Count() + result := &OpenAIForwardResult{ + RequestID: responseID, + Usage: usage, + Model: originalModel, + UpstreamModel: mappedModel, + ServiceTier: extractOpenAIServiceTierFromBody(body), + ReasoningEffort: extractOpenAIReasoningEffortFromBody(body, originalModel), + Stream: reqStream, + OpenAIWSMode: true, + ResponseHeaders: cloneHeader(resp.Header), + Duration: time.Since(turnStart), + FirstTokenMs: firstTokenMs, + } + if replayInput := replayCollector.Items(); len(replayInput) > 0 { + result.wsReplayInput = replayInput + result.wsReplayInputExists = true + } + if imageCount > 0 { + result.ImageCount = imageCount + result.ImageSize = imageSizeTier + result.ImageInputSize = imageInputSize + result.ImageOutputSizes = imageCounter.Sizes() + result.BillingModel = imageBillingModel + } + return result + } + + scanner := bufio.NewScanner(resp.Body) + maxLineSize := defaultMaxLineSize + if s.cfg != nil && s.cfg.Gateway.MaxLineSize > 0 { + maxLineSize = s.cfg.Gateway.MaxLineSize + } + scanBuf := getSSEScannerBuf64K() + scanner.Buffer(scanBuf[:0], maxLineSize) + defer putSSEScannerBuf64K(scanBuf) + + for scanner.Scan() { + line := scanner.Text() + data, ok := extractOpenAISSEDataLine(line) + if !ok { + continue + } + trimmedData := strings.TrimSpace(data) + if trimmedData == "" { + continue + } + if trimmedData == "[DONE]" { + sawDone = true + continue + } + + upstreamMessage := []byte(trimmedData) + eventType, eventResponseID, _ := parseOpenAIWSEventEnvelope(upstreamMessage) + if responseID == "" && eventResponseID != "" { + responseID = eventResponseID + } + if eventType != "" { + eventCount++ + if firstEventType == "" { + firstEventType = eventType + } + lastEventType = eventType + } + if isOpenAIWSTokenEvent(eventType) { + tokenEventCount++ + if firstTokenMs == nil { + ms := int(time.Since(turnStart).Milliseconds()) + firstTokenMs = &ms + } + } + if openAIWSEventShouldParseUsage(eventType) { + parseOpenAIWSResponseUsageFromCompletedEvent(upstreamMessage, &usage) + } + imageCounter.AddSSEData(upstreamMessage) + + if needModelReplace && len(mappedModelBytes) > 0 && openAIWSEventMayContainModel(eventType) && strings.Contains(trimmedData, mappedModel) { + upstreamMessage = replaceOpenAIWSMessageModel(upstreamMessage, mappedModel, originalModel) + } + if s.toolCorrector != nil && openAIWSEventMayContainToolCalls(eventType) && openAIWSMessageLikelyContainsToolCalls(upstreamMessage) { + if corrected, changed := s.toolCorrector.CorrectToolCallsInSSEBytes(upstreamMessage); changed { + upstreamMessage = corrected + } + } + replayCollector.AddEvent(eventType, upstreamMessage) + + if !clientDisconnected { + if err := writeClientMessage(upstreamMessage); err != nil { + if isOpenAIWSClientDisconnectError(err) { + clientDisconnected = true + closeStatus, closeReason := summarizeOpenAIWSReadCloseError(err) + logOpenAIWSModeInfo( + "ingress_ws_http_bridge_client_disconnected_drain account_id=%d turn=%d close_status=%s close_reason=%s", + account.ID, + turn, + closeStatus, + truncateOpenAIWSLogValue(closeReason, openAIWSHeaderValueMaxLen), + ) + } else { + return nil, wrapOpenAIWSIngressTurnError( + "write_client", + fmt.Errorf("write client websocket event: %w", err), + wroteDownstream, + ) + } + } else { + wroteDownstream = true + } + } + + if eventType == "error" { + errCodeRaw, errTypeRaw, errMsgRaw := parseOpenAIWSErrorEventFields(upstreamMessage) + s.persistOpenAIWSRateLimitSignal(ctx, account, resp.Header, upstreamMessage, errCodeRaw, errTypeRaw, errMsgRaw) + errMessage := strings.TrimSpace(errMsgRaw) + if errMessage == "" { + errMessage = "upstream error event" + } + return resultWithUsage(), errors.New(errMessage) + } + if isOpenAIWSTerminalEvent(eventType) { + terminalEventCount++ + firstTokenMsValue := -1 + if firstTokenMs != nil { + firstTokenMsValue = *firstTokenMs + } + logOpenAIWSModeInfo( + "ingress_ws_http_bridge_turn_completed account_id=%d turn=%d response_id=%s payload_bytes=%d duration_ms=%d events=%d token_events=%d terminal_events=%d first_event=%s last_event=%s first_token_ms=%d client_disconnected=%v", + account.ID, + turn, + truncateOpenAIWSLogValue(responseID, openAIWSIDValueMaxLen), + payloadBytes, + time.Since(turnStart).Milliseconds(), + eventCount, + tokenEventCount, + terminalEventCount, + truncateOpenAIWSLogValue(firstEventType, openAIWSLogValueMaxLen), + truncateOpenAIWSLogValue(lastEventType, openAIWSLogValueMaxLen), + firstTokenMsValue, + clientDisconnected, + ) + return resultWithUsage(), nil + } + } + if err := scanner.Err(); err != nil { + return resultWithUsage(), fmt.Errorf("read upstream http bridge stream: %w", err) + } + if sawDone && eventCount > 0 { + return resultWithUsage(), nil + } + return resultWithUsage(), errors.New("upstream http bridge stream ended before terminal event") +} diff --git a/backend/internal/service/openai_ws_http_bridge_test.go b/backend/internal/service/openai_ws_http_bridge_test.go new file mode 100644 index 000000000000..0a1d6b564bb3 --- /dev/null +++ b/backend/internal/service/openai_ws_http_bridge_test.go @@ -0,0 +1,461 @@ +package service + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + coderws "github.com/coder/websocket" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" +) + +func TestPrepareOpenAIWSHTTPBridgeBodyStripsWSFields(t *testing.T) { + body, err := prepareOpenAIWSHTTPBridgeBody([]byte(`{"type":"response.create","generate":true,"model":"gpt-5","stream":false,"previous_response_id":"resp_prev","input":"hi"}`)) + require.NoError(t, err) + require.False(t, gjson.GetBytes(body, "type").Exists()) + require.False(t, gjson.GetBytes(body, "generate").Exists()) + require.False(t, gjson.GetBytes(body, "previous_response_id").Exists()) + require.Equal(t, "gpt-5", gjson.GetBytes(body, "model").String()) + require.True(t, gjson.GetBytes(body, "stream").Bool()) + require.Equal(t, "hi", gjson.GetBytes(body, "input").String()) +} + +func TestOpenAIWSHTTPBridgeDecisionKeepsSmallFramesOnWS(t *testing.T) { + svc := &OpenAIGatewayService{ + cfg: &config.Config{ + Gateway: config.GatewayConfig{ + OpenAIWS: config.GatewayOpenAIWSConfig{ + HTTPBridgeEnabled: true, + HTTPBridgeThresholdBytes: 100, + }, + }, + }, + } + + require.False(t, svc.shouldBridgeOpenAIWSHTTP(99, "")) + require.True(t, svc.shouldBridgeOpenAIWSHTTP(100, "")) + require.False(t, svc.shouldBridgeOpenAIWSHTTP(1000, "resp_existing")) + + svc.cfg.Gateway.OpenAIWS.HTTPBridgeEnabled = false + require.False(t, svc.shouldBridgeOpenAIWSHTTP(1000, "")) +} + +func TestOpenAIWSHTTPBridgeRelaysSSEFramesAsWebSocketMessages(t *testing.T) { + gin.SetMode(gin.TestMode) + + sseBody := strings.Join([]string{ + `data: {"type":"response.created","response":{"id":"resp_bridge","model":"gpt-5"}}`, + "", + `data: {"type":"response.output_text.delta","response":{"id":"resp_bridge"},"delta":"ok"}`, + "", + `data: {"type":"response.completed","response":{"id":"resp_bridge","model":"gpt-5","usage":{"input_tokens":3,"output_tokens":2}}}`, + "", + }, "\n") + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{ + "Content-Type": []string{"text/event-stream"}, + "x-request-id": []string{"rid_bridge"}, + }, + Body: io.NopCloser(strings.NewReader(sseBody)), + }} + svc := &OpenAIGatewayService{ + cfg: &config.Config{ + Gateway: config.GatewayConfig{ + MaxLineSize: defaultMaxLineSize, + OpenAIWS: config.GatewayOpenAIWSConfig{ + HTTPBridgeEnabled: true, + HTTPBridgeThresholdBytes: 1, + }, + }, + }, + httpUpstream: upstream, + toolCorrector: NewCodexToolCorrector(), + } + account := &Account{ + ID: 7, + Name: "api-key", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Concurrency: 1, + Status: StatusActive, + } + payload := []byte(`{"type":"response.create","generate":true,"model":"gpt-5","stream":true,"input":"hi"}`) + + type bridgeResult struct { + result *OpenAIForwardResult + err error + } + resultCh := make(chan bridgeResult, 1) + wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{CompressionMode: coderws.CompressionContextTakeover}) + if err != nil { + resultCh <- bridgeResult{err: err} + return + } + defer func() { _ = conn.CloseNow() }() + + rec := httptest.NewRecorder() + ginCtx, _ := gin.CreateTestContext(rec) + req := r.Clone(r.Context()) + req.Header = req.Header.Clone() + ginCtx.Request = req + + writeClient := func(message []byte) error { + writeCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second) + defer cancel() + return conn.Write(writeCtx, coderws.MessageText, message) + } + result, bridgeErr := svc.proxyOpenAIWSHTTPBridgeTurn( + r.Context(), + ginCtx, + account, + "sk-test", + payload, + len(payload), + "gpt-5", + "", + "", + "", + 1, + writeClient, + ) + resultCh <- bridgeResult{result: result, err: bridgeErr} + })) + defer wsServer.Close() + + dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second) + clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil) + cancelDial() + require.NoError(t, err) + defer func() { _ = clientConn.CloseNow() }() + + readEvent := func() []byte { + readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second) + msgType, event, readErr := clientConn.Read(readCtx) + cancelRead() + require.NoError(t, readErr) + require.Equal(t, coderws.MessageText, msgType) + return event + } + + created := readEvent() + delta := readEvent() + completed := readEvent() + + require.Equal(t, "response.created", gjson.GetBytes(created, "type").String()) + require.Equal(t, "response.output_text.delta", gjson.GetBytes(delta, "type").String()) + require.Equal(t, "response.completed", gjson.GetBytes(completed, "type").String()) + + select { + case bridge := <-resultCh: + require.NoError(t, bridge.err) + require.NotNil(t, bridge.result) + require.Equal(t, "resp_bridge", bridge.result.RequestID) + require.Equal(t, 3, bridge.result.Usage.InputTokens) + require.Equal(t, 2, bridge.result.Usage.OutputTokens) + require.True(t, bridge.result.OpenAIWSMode) + case <-time.After(3 * time.Second): + t.Fatal("timed out waiting for bridge result") + } + + require.NotNil(t, upstream.lastReq) + require.Equal(t, http.MethodPost, upstream.lastReq.Method) + require.False(t, gjson.GetBytes(upstream.lastBody, "type").Exists()) + require.False(t, gjson.GetBytes(upstream.lastBody, "generate").Exists()) + require.True(t, gjson.GetBytes(upstream.lastBody, "stream").Bool()) +} + +func TestOpenAIWSHTTPBridgeAcceptsFirstFrameAboveLegacy16MiB(t *testing.T) { + gin.SetMode(gin.TestMode) + + sseBody := strings.Join([]string{ + `data: {"type":"response.created","response":{"id":"resp_large_bridge","model":"gpt-5"}}`, + "", + `data: {"type":"response.completed","response":{"id":"resp_large_bridge","model":"gpt-5","usage":{"input_tokens":9,"output_tokens":1}}}`, + "", + }, "\n") + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{ + "Content-Type": []string{"text/event-stream"}, + "x-request-id": []string{"rid_large_bridge"}, + }, + Body: io.NopCloser(strings.NewReader(sseBody)), + }} + cfg := &config.Config{ + Gateway: config.GatewayConfig{ + MaxLineSize: defaultMaxLineSize, + OpenAIWS: config.GatewayOpenAIWSConfig{ + Enabled: true, + APIKeyEnabled: true, + ResponsesWebsocketsV2: true, + ClientReadLimitBytes: 64 * 1024 * 1024, + HTTPBridgeEnabled: true, + HTTPBridgeThresholdBytes: 15 * 1024 * 1024, + }, + }, + } + svc := &OpenAIGatewayService{ + cfg: cfg, + httpUpstream: upstream, + toolCorrector: NewCodexToolCorrector(), + } + account := &Account{ + ID: 9, + Name: "api-key", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Credentials: map[string]any{"api_key": "sk-upstream"}, + Extra: map[string]any{ + "openai_apikey_responses_websockets_v2_enabled": true, + }, + Concurrency: 1, + Status: StatusActive, + } + + payload := []byte(`{"type":"response.create","generate":true,"model":"gpt-5","stream":true,"input":"` + strings.Repeat("x", 17*1024*1024) + `"}`) + require.Greater(t, len(payload), 16*1024*1024) + require.Less(t, int64(len(payload)), ResolveOpenAIWSClientReadLimitBytes(cfg)) + + errCh := make(chan error, 1) + wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{CompressionMode: coderws.CompressionContextTakeover}) + if err != nil { + errCh <- err + return + } + defer func() { _ = conn.CloseNow() }() + conn.SetReadLimit(ResolveOpenAIWSClientReadLimitBytes(cfg)) + + readCtx, cancelRead := context.WithTimeout(r.Context(), 10*time.Second) + msgType, firstMessage, err := conn.Read(readCtx) + cancelRead() + if err != nil { + errCh <- err + return + } + if msgType != coderws.MessageText && msgType != coderws.MessageBinary { + errCh <- NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, "unexpected client websocket message type", nil) + return + } + + rec := httptest.NewRecorder() + ginCtx, _ := gin.CreateTestContext(rec) + req := r.Clone(r.Context()) + req.Header = req.Header.Clone() + req.Header.Set("User-Agent", "codex_cli_rs/0.135.0") + ginCtx.Request = req + + proxyCtx, cancelProxy := context.WithTimeout(r.Context(), 20*time.Second) + defer cancelProxy() + errCh <- svc.ProxyResponsesWebSocketFromClient(proxyCtx, ginCtx, conn, account, "sk-test", firstMessage, nil) + })) + defer wsServer.Close() + + dialCtx, cancelDial := context.WithTimeout(context.Background(), 5*time.Second) + clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil) + cancelDial() + require.NoError(t, err) + defer func() { _ = clientConn.CloseNow() }() + + writeCtx, cancelWrite := context.WithTimeout(context.Background(), 20*time.Second) + err = clientConn.Write(writeCtx, coderws.MessageText, payload) + cancelWrite() + require.NoError(t, err) + + var eventTypes []string + for { + readCtx, cancelRead := context.WithTimeout(context.Background(), 10*time.Second) + msgType, event, readErr := clientConn.Read(readCtx) + cancelRead() + require.NoError(t, readErr) + require.Equal(t, coderws.MessageText, msgType) + + eventType := gjson.GetBytes(event, "type").String() + eventTypes = append(eventTypes, eventType) + if eventType == "response.completed" { + break + } + } + require.Contains(t, eventTypes, "response.created") + require.Contains(t, eventTypes, "response.completed") + + require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done")) + select { + case proxyErr := <-errCh: + require.NoError(t, proxyErr) + case <-time.After(10 * time.Second): + t.Fatal("timed out waiting for websocket bridge proxy to finish") + } + + require.NotNil(t, upstream.lastReq) + require.Equal(t, http.MethodPost, upstream.lastReq.Method) + require.Greater(t, len(upstream.lastBody), 16*1024*1024) + require.False(t, gjson.GetBytes(upstream.lastBody, "type").Exists()) + require.False(t, gjson.GetBytes(upstream.lastBody, "generate").Exists()) + require.True(t, gjson.GetBytes(upstream.lastBody, "stream").Bool()) + require.Equal(t, "gpt-5", gjson.GetBytes(upstream.lastBody, "model").String()) +} + +func TestOpenAIWSHTTPBridgeKeepsContinuationFramesOnHTTPWithoutPreviousResponseID(t *testing.T) { + gin.SetMode(gin.TestMode) + + firstSSEBody := strings.Join([]string{ + `data: {"type":"response.completed","response":{"id":"resp_bridge_first","model":"gpt-5.1","output":[{"type":"function_call","id":"fc_bridge_1","call_id":"call_bridge_1","name":"shell","arguments":"{}"}],"usage":{"input_tokens":9,"output_tokens":1}}}`, + "", + }, "\n") + secondSSEBody := strings.Join([]string{ + `data: {"type":"response.completed","response":{"id":"resp_bridge_second","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`, + "", + }, "\n") + upstream := &httpUpstreamRecorder{responses: []*http.Response{ + { + StatusCode: http.StatusOK, + Header: http.Header{ + "Content-Type": []string{"text/event-stream"}, + }, + Body: io.NopCloser(strings.NewReader(firstSSEBody)), + }, + { + StatusCode: http.StatusOK, + Header: http.Header{ + "Content-Type": []string{"text/event-stream"}, + }, + Body: io.NopCloser(strings.NewReader(secondSSEBody)), + }, + }} + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + cfg.Security.URLAllowlist.AllowInsecureHTTP = true + cfg.Gateway.OpenAIWS.Enabled = true + cfg.Gateway.OpenAIWS.OAuthEnabled = true + cfg.Gateway.OpenAIWS.APIKeyEnabled = true + cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true + cfg.Gateway.OpenAIWS.HTTPBridgeEnabled = true + cfg.Gateway.OpenAIWS.HTTPBridgeThresholdBytes = 1 + cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1 + cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0 + cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1 + cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8 + cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3 + cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3 + cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3 + + captureConn := &openAIWSCaptureConn{} + captureDialer := &openAIWSCaptureDialer{conn: captureConn} + pool := newOpenAIWSConnPool(cfg) + pool.setClientDialerForTest(captureDialer) + + svc := &OpenAIGatewayService{ + cfg: cfg, + httpUpstream: upstream, + cache: &stubGatewayCache{}, + openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), + toolCorrector: NewCodexToolCorrector(), + openaiWSPool: pool, + } + account := &Account{ + ID: 19, + Name: "api-key-bridge-handoff", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Credentials: map[string]any{"api_key": "sk-upstream"}, + Extra: map[string]any{ + "responses_websockets_v2_enabled": true, + }, + Concurrency: 1, + Status: StatusActive, + Schedulable: true, + } + + errCh := make(chan error, 1) + wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{CompressionMode: coderws.CompressionContextTakeover}) + if err != nil { + errCh <- err + return + } + defer func() { _ = conn.CloseNow() }() + + readCtx, cancelRead := context.WithTimeout(r.Context(), 3*time.Second) + msgType, firstMessage, err := conn.Read(readCtx) + cancelRead() + if err != nil { + errCh <- err + return + } + if msgType != coderws.MessageText && msgType != coderws.MessageBinary { + errCh <- NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, "unexpected client websocket message type", nil) + return + } + + rec := httptest.NewRecorder() + ginCtx, _ := gin.CreateTestContext(rec) + req := r.Clone(r.Context()) + req.Header = req.Header.Clone() + req.Header.Set("User-Agent", "codex_cli_rs/0.135.0") + ginCtx.Request = req + + errCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, nil) + })) + defer wsServer.Close() + + dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second) + clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil) + cancelDial() + require.NoError(t, err) + defer func() { _ = clientConn.CloseNow() }() + + writeMessage := func(payload string) { + writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second) + defer cancelWrite() + require.NoError(t, clientConn.Write(writeCtx, coderws.MessageText, []byte(payload))) + } + readMessage := func() []byte { + readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second) + defer cancelRead() + msgType, event, readErr := clientConn.Read(readCtx) + require.NoError(t, readErr) + require.Equal(t, coderws.MessageText, msgType) + return event + } + + writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":true,"input":"first"}`) + firstTurnEvent := readMessage() + require.Equal(t, "response.completed", gjson.GetBytes(firstTurnEvent, "type").String()) + require.Equal(t, "resp_bridge_first", gjson.GetBytes(firstTurnEvent, "response.id").String()) + + writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"previous_response_id":"resp_bridge_first","input":[{"type":"function_call_output","call_id":"call_bridge_1","output":"ok"}]}`) + secondTurnEvent := readMessage() + require.Equal(t, "response.completed", gjson.GetBytes(secondTurnEvent, "type").String()) + require.Equal(t, "resp_bridge_second", gjson.GetBytes(secondTurnEvent, "response.id").String()) + + require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done")) + select { + case proxyErr := <-errCh: + require.NoError(t, proxyErr) + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for websocket bridge proxy to finish") + } + + require.Len(t, upstream.bodies, 2, "进入 HTTP bridge 后同一客户端 WS 连接内应保持 HTTP/SSE bridge") + require.False(t, gjson.GetBytes(upstream.bodies[0], "previous_response_id").Exists()) + require.False(t, gjson.GetBytes(upstream.bodies[1], "previous_response_id").Exists()) + secondInput := gjson.GetBytes(upstream.bodies[1], "input").Array() + require.Len(t, secondInput, 3) + require.Equal(t, "first", secondInput[0].String()) + require.Equal(t, "function_call", secondInput[1].Get("type").String()) + require.Equal(t, "call_bridge_1", secondInput[1].Get("call_id").String()) + require.Equal(t, "function_call_output", secondInput[2].Get("type").String()) + require.Equal(t, "call_bridge_1", secondInput[2].Get("call_id").String()) + require.Equal(t, 0, captureDialer.DialCount()) + require.Empty(t, captureConn.writes) +} From 2a075a85bd6746bd3d2093f56cd713d041d16663 Mon Sep 17 00:00:00 2001 From: Kanshan03 <92105202+Kanshan03@users.noreply.github.com> Date: Mon, 1 Jun 2026 11:25:28 +0800 Subject: [PATCH 35/36] =?UTF-8?q?fix(openai):=20=E6=B3=A8=E5=85=A5=20WS=20?= =?UTF-8?q?Codex=20=E7=94=9F=E5=9B=BE=E6=A1=A5=E6=8E=A5=E5=B7=A5=E5=85=B7?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../internal/service/openai_ws_forwarder.go | 32 ++++- ...penai_ws_forwarder_ingress_session_test.go | 136 ++++++++++++++++++ 2 files changed, 166 insertions(+), 2 deletions(-) diff --git a/backend/internal/service/openai_ws_forwarder.go b/backend/internal/service/openai_ws_forwarder.go index dd2f45facbd3..22ed214b256e 100644 --- a/backend/internal/service/openai_ws_forwarder.go +++ b/backend/internal/service/openai_ws_forwarder.go @@ -2477,6 +2477,7 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( wsPath = normalizeOpenAIWSLogValue(parsedURL.Path) } debugEnabled := isOpenAIWSModeDebugEnabled() + isCodexCLI := openai.IsCodexOfficialClientByHeaders(c.GetHeader("User-Agent"), c.GetHeader("originator")) || (s.cfg != nil && s.cfg.Gateway.ForceCodexCLI) type openAIWSClientPayload struct { payloadRaw []byte @@ -2586,6 +2587,34 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( } normalized = next } + apiKey := getAPIKeyFromContext(c) + imageGenerationAllowed := GroupAllowsImageGeneration(apiKeyGroup(apiKey)) + codexBridgeEnabled := isCodexCLI && imageGenerationAllowed && s.isCodexImageGenerationBridgeEnabled(ctx, account, apiKey) + if codexBridgeEnabled { + payloadMap := make(map[string]any) + if err := json.Unmarshal(normalized, &payloadMap); err != nil { + return openAIWSClientPayload{}, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, "invalid websocket request payload", err) + } + bridgeModified := false + if ensureOpenAIResponsesImageGenerationTool(payloadMap) { + bridgeModified = true + logOpenAIWSModeInfo("ingress_ws_codex_image_tool_injected account_id=%d", account.ID) + } + if normalizeOpenAIResponsesImageGenerationTools(payloadMap) { + bridgeModified = true + } + if applyCodexImageGenerationBridgeInstructions(payloadMap) { + bridgeModified = true + logOpenAIWSModeInfo("ingress_ws_codex_image_bridge_instructions_added account_id=%d", account.ID) + } + if bridgeModified { + rebuilt, marshalErr := json.Marshal(payloadMap) + if marshalErr != nil { + return openAIWSClientPayload{}, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, "invalid websocket request payload", marshalErr) + } + normalized = rebuilt + } + } upstreamModel := normalizeOpenAIModelForUpstream(account, account.GetMappedModel(originalModel)) if modelMissing || upstreamModel != originalModel { next, setErr := applyPayloadMutation(normalized, "model", upstreamModel) @@ -2595,7 +2624,7 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( normalized = next } imageIntent := IsImageGenerationIntent(openAIResponsesEndpoint, originalModel, normalized) - if imageIntent && !GroupAllowsImageGeneration(apiKeyGroup(getAPIKeyFromContext(c))) { + if imageIntent && !imageGenerationAllowed { return openAIWSClientPayload{}, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, ImageGenerationPermissionMessage(), nil) } imageBillingModel := "" @@ -2694,7 +2723,6 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( } } - isCodexCLI := openai.IsCodexOfficialClientByHeaders(c.GetHeader("User-Agent"), c.GetHeader("originator")) || (s.cfg != nil && s.cfg.Gateway.ForceCodexCLI) wsHeaders, _ := s.buildOpenAIWSHeaders(c, account, token, wsDecision, isCodexCLI, turnState, strings.TrimSpace(c.GetHeader(openAIWSTurnMetadataHeader)), firstPayload.promptCacheKey) baseAcquireReq := openAIWSAcquireRequest{ Account: account, diff --git a/backend/internal/service/openai_ws_forwarder_ingress_session_test.go b/backend/internal/service/openai_ws_forwarder_ingress_session_test.go index b7f1bc4f78da..5e4b70c2a8af 100644 --- a/backend/internal/service/openai_ws_forwarder_ingress_session_test.go +++ b/backend/internal/service/openai_ws_forwarder_ingress_session_test.go @@ -298,6 +298,142 @@ func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_FollowupCreateCa require.Equal(t, "resp_omit_model_1", gjson.Get(requestToJSONString(captureConn.writes[1]), "previous_response_id").String()) } +func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_InjectsCodexImageBridge(t *testing.T) { + gin.SetMode(gin.TestMode) + + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + cfg.Security.URLAllowlist.AllowInsecureHTTP = true + cfg.Gateway.OpenAIWS.Enabled = true + cfg.Gateway.OpenAIWS.OAuthEnabled = true + cfg.Gateway.OpenAIWS.APIKeyEnabled = true + cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true + cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1 + cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0 + cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1 + cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8 + cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3 + cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3 + cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3 + + captureConn := &openAIWSCaptureConn{ + events: [][]byte{ + []byte(`{"type":"response.completed","response":{"id":"resp_codex_image_bridge","model":"gpt-5.5","usage":{"input_tokens":1,"output_tokens":1}}}`), + }, + } + captureDialer := &openAIWSCaptureDialer{conn: captureConn} + pool := newOpenAIWSConnPool(cfg) + pool.setClientDialerForTest(captureDialer) + + svc := &OpenAIGatewayService{ + cfg: cfg, + httpUpstream: &httpUpstreamRecorder{}, + cache: &stubGatewayCache{}, + openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), + toolCorrector: NewCodexToolCorrector(), + openaiWSPool: pool, + } + + groupID := int64(3) + apiKey := &APIKey{ + ID: 1, + UserID: 1, + GroupID: &groupID, + Group: &Group{ + ID: groupID, + AllowImageGeneration: true, + }, + } + account := &Account{ + ID: 31, + Name: "openai-codex-image-ws", + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Credentials: map[string]any{ + "access_token": "test-token", + }, + Extra: map[string]any{ + "openai_oauth_responses_websockets_v2_enabled": true, + "codex_image_generation_bridge": true, + }, + } + + serverErrCh := make(chan error, 1) + wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{ + CompressionMode: coderws.CompressionContextTakeover, + }) + if err != nil { + serverErrCh <- err + return + } + defer func() { + _ = conn.CloseNow() + }() + + rec := httptest.NewRecorder() + ginCtx, _ := gin.CreateTestContext(rec) + req := r.Clone(r.Context()) + req.Header = req.Header.Clone() + req.Header.Set("User-Agent", "codex_cli_rs/0.98.0") + ginCtx.Request = req + ginCtx.Set("api_key", apiKey) + + readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second) + msgType, firstMessage, readErr := conn.Read(readCtx) + cancel() + if readErr != nil { + serverErrCh <- readErr + return + } + if msgType != coderws.MessageText && msgType != coderws.MessageBinary { + serverErrCh <- errors.New("unsupported websocket client message type") + return + } + + serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "test-token", firstMessage, nil) + })) + defer wsServer.Close() + + dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second) + clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil) + cancelDial() + require.NoError(t, err) + defer func() { + _ = clientConn.CloseNow() + }() + + writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second) + err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"gpt-5.5","stream":false,"input":"draw a cat"}`)) + cancelWrite() + require.NoError(t, err) + + readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second) + msgType, message, err := clientConn.Read(readCtx) + cancelRead() + require.NoError(t, err) + require.Equal(t, coderws.MessageText, msgType) + require.Equal(t, "resp_codex_image_bridge", gjson.GetBytes(message, "response.id").String()) + + _ = clientConn.Close(coderws.StatusNormalClosure, "done") + + select { + case serverErr := <-serverErrCh: + require.NoError(t, serverErr) + case <-time.After(5 * time.Second): + t.Fatal("等待 ingress websocket 结束超时") + } + + require.Len(t, captureConn.writes, 1) + upstreamPayload := requestToJSONString(captureConn.writes[0]) + require.True(t, gjson.Get(upstreamPayload, `tools.#(type=="image_generation")`).Exists()) + require.Equal(t, "png", gjson.Get(upstreamPayload, `tools.#(type=="image_generation").output_format`).String()) + require.Contains(t, gjson.Get(upstreamPayload, "instructions").String(), "image_generation") +} + func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_DedicatedModeDoesNotReuseConnAcrossSessions(t *testing.T) { gin.SetMode(gin.TestMode) From 003b2786dacfd9c1c5342c499ae57590979fef8d Mon Sep 17 00:00:00 2001 From: visa2 Date: Mon, 1 Jun 2026 12:03:15 +0800 Subject: [PATCH 36/36] test(apicompat): check type assertions in responses stream wire tests errcheck (check-type-assertions) flagged unchecked single-value type assertions; switch to the comma-ok form so golangci-lint passes. Co-Authored-By: Claude Opus 4.8 --- .../responses_stream_event_wire_test.go | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/backend/internal/pkg/apicompat/responses_stream_event_wire_test.go b/backend/internal/pkg/apicompat/responses_stream_event_wire_test.go index fbef45af4544..b4f6871d5f4f 100644 --- a/backend/internal/pkg/apicompat/responses_stream_event_wire_test.go +++ b/backend/internal/pkg/apicompat/responses_stream_event_wire_test.go @@ -43,7 +43,8 @@ func TestWire_FunctionCallItemAlwaysComplete(t *testing.T) { OutputIndex: 1, Item: &ResponsesOutput{Type: "function_call", ID: "fc_1", CallID: "call_a", Name: "exec", Status: "in_progress"}, }) - item := added["item"].(map[string]any) + item, ok := added["item"].(map[string]any) + require.True(t, ok, "item must be an object") for _, k := range []string{"call_id", "name", "arguments"} { require.Containsf(t, item, k, "function_call item missing %q", k) } @@ -57,9 +58,10 @@ func TestWire_MessageItemContentAlwaysArray(t *testing.T) { OutputIndex: 0, Item: &ResponsesOutput{Type: "message", ID: "msg_1", Role: "assistant", Status: "in_progress"}, }) - item := m["item"].(map[string]any) + item, ok := m["item"].(map[string]any) + require.True(t, ok, "item must be an object") require.Contains(t, item, "content") - _, ok := item["content"].([]any) + _, ok = item["content"].([]any) require.True(t, ok, "content must be an array") } @@ -70,9 +72,10 @@ func TestWire_ReasoningItemSummaryAlwaysArray(t *testing.T) { OutputIndex: 0, Item: &ResponsesOutput{Type: "reasoning", ID: "rs_1", Status: "in_progress"}, }) - item := m["item"].(map[string]any) + item, ok := m["item"].(map[string]any) + require.True(t, ok, "item must be an object") require.Contains(t, item, "summary") - _, ok := item["summary"].([]any) + _, ok = item["summary"].([]any) require.True(t, ok, "summary must be an array") } @@ -82,7 +85,8 @@ func TestWire_ContentPartCarriesAnnotationsLogprobs(t *testing.T) { Type: "response.content_part.added", OutputIndex: 0, ContentIndex: 0, ItemID: "msg_1", Part: &ResponsesContentPart{Type: "output_text", Text: ""}, }) - part := m["part"].(map[string]any) + part, ok := m["part"].(map[string]any) + require.True(t, ok, "part must be an object") require.Equal(t, "output_text", part["type"]) require.Contains(t, part, "text") require.Contains(t, part, "annotations")