body
") + if err != nil { + t.Fatalf("expected send via STARTTLS fallback to succeed, got: %v", err) + } + if !srv.sawCommand("STARTTLS") { + t.Fatal("expected send path to upgrade via STARTTLS") + } + if !srv.sawCommand("DATA") { + t.Fatal("expected send path to reach DATA") + } +} + +// 发送路径全流程：UseTLS=true + 隐式 TLS 服务器（465 语义）保持既有行为。 +func TestSendEmailWithConfigImplicitTLS(t *testing.T) { + srv, port := startFakeSMTPServer(t, true, false) + svc := &EmailService{} + + err := svc.SendEmailWithConfig(smtpTestConfig(port, true), "rcpt@example.com", "subject", "body
") + if err != nil { + t.Fatalf("expected send via implicit TLS to succeed, got: %v", err) + } + if !srv.sawCommand("DATA") { + t.Fatal("expected send path to reach DATA") + } +} diff --git a/backend/internal/service/gateway_anthropic_apikey_passthrough_test.go b/backend/internal/service/gateway_anthropic_apikey_passthrough_test.go index 902a7b0f91..4fa52cc9ce 100644 --- a/backend/internal/service/gateway_anthropic_apikey_passthrough_test.go +++ b/backend/internal/service/gateway_anthropic_apikey_passthrough_test.go @@ -523,7 +523,8 @@ func TestGatewayService_AnthropicAPIKeyPassthrough_ModelMappingPreservesOtherFie require.Equal(t, "hello world", gjson.GetBytes(sentBody, "messages.0.content.0.text").String(), "messages 字段不应被修改") require.Equal(t, "enabled", gjson.GetBytes(sentBody, "thinking.type").String(), "thinking 字段不应被修改") require.Equal(t, int64(5000), gjson.GetBytes(sentBody, "thinking.budget_tokens").Int(), "thinking.budget_tokens 不应被修改") - require.Equal(t, int64(1024), gjson.GetBytes(sentBody, "max_tokens").Int(), "max_tokens 不应被修改") + require.False(t, gjson.GetBytes(sentBody, "max_tokens").Exists(), + "max_tokens 作为生成参数应被 count_tokens 过滤剥离") } func TestGatewayService_AnthropicAPIKeyPassthrough_CountTokensFiltersGenerationFields(t *testing.T) { @@ -582,7 +583,8 @@ func TestGatewayService_AnthropicAPIKeyPassthrough_CountTokensFiltersGenerationF 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.False(t, gjson.GetBytes(sentBody, "max_tokens").Exists(), + "count_tokens 请求不得携带生成参数 max_tokens") require.Equal(t, "enabled", gjson.GetBytes(sentBody, "thinking.type").String()) } diff --git a/backend/internal/service/gateway_anthropic_passthrough.go b/backend/internal/service/gateway_anthropic_passthrough.go index 8d7610b936..c18d5bf047 100644 --- a/backend/internal/service/gateway_anthropic_passthrough.go +++ b/backend/internal/service/gateway_anthropic_passthrough.go @@ -269,6 +269,11 @@ func (s *GatewayService) forwardAnthropicAPIKeyPassthroughWithInput( if input.RequestStream { streamResult, err := s.handleStreamingResponseAnthropicAPIKeyPassthrough(ctx, resp, c, account, input.StartTime, input.RequestModel) if err != nil { + // 流中断时保留已观测到的 usage 与错误一起返回，避免上游已计量的请求 + // 完全漏记漏计费（issue #5148）。 + if partial := partialStreamUsageResult(c, resp, streamResult, input.OriginalModel, input.RequestModel, input.StartTime, err); partial != nil { + return partial, err + } return nil, err } usage = streamResult.usage @@ -285,14 +290,16 @@ func (s *GatewayService) forwardAnthropicAPIKeyPassthroughWithInput( } return &ForwardResult{ - RequestID: resp.Header.Get("x-request-id"), - Usage: *usage, - Model: input.OriginalModel, - UpstreamModel: input.RequestModel, - Stream: input.RequestStream, - Duration: time.Since(input.StartTime), - FirstTokenMs: firstTokenMs, - ClientDisconnect: clientDisconnect, + RequestID: resp.Header.Get("x-request-id"), + Usage: *usage, + Model: input.OriginalModel, + UpstreamModel: input.RequestModel, + UpstreamResponseModel: observedUpstreamResponseModel(c), + UpstreamResponseModelConflict: observedUpstreamResponseModelConflict(c), + Stream: input.RequestStream, + Duration: time.Since(input.StartTime), + FirstTokenMs: firstTokenMs, + ClientDisconnect: clientDisconnect, }, nil } @@ -374,6 +381,10 @@ func (s *GatewayService) handleStreamingResponseAnthropicAPIKeyPassthrough( startTime time.Time, model string, ) (*streamingResult, error) { + observer := upstreamResponseModelObserverFromContext(c) + if observer == nil { + observer = beginUpstreamResponseModelObservation(c) + } if s.rateLimitService != nil { s.rateLimitService.UpdateSessionWindow(ctx, account, resp.Header) } @@ -527,6 +538,7 @@ func (s *GatewayService) handleStreamingResponseAnthropicAPIKeyPassthrough( line := ev.line if data, ok := extractAnthropicSSEDataLine(line); ok { trimmed := strings.TrimSpace(data) + observer.ObserveAnthropic([]byte(trimmed)) if anthropicStreamEventIsTerminal("", trimmed) { sawTerminalEvent = true } @@ -777,6 +789,11 @@ func (s *GatewayService) handleNonStreamingResponseAnthropicAPIKeyPassthrough( if err != nil { return nil, err } + observer := upstreamResponseModelObserverFromContext(c) + if observer == nil { + observer = beginUpstreamResponseModelObservation(c) + } + observer.ObserveAnthropic(body) if resp.StatusCode >= http.StatusOK && resp.StatusCode < http.StatusMultipleChoices { var raw json.RawMessage diff --git a/backend/internal/service/gateway_channel_restriction_test.go b/backend/internal/service/gateway_channel_restriction_test.go index 3a2ad2ff72..0e26a00c02 100644 --- a/backend/internal/service/gateway_channel_restriction_test.go +++ b/backend/internal/service/gateway_channel_restriction_test.go @@ -29,6 +29,12 @@ func TestBillingModelForRestriction_Upstream(t *testing.T) { require.Equal(t, "", got, "upstream should return empty (per-account check needed)") } +func TestBillingModelForRestriction_ResponseModelUsesMappedPrecheck(t *testing.T) { + t.Parallel() + got := billingModelForRestriction(BillingModelSourceResponse, "claude-fable-5", "claude-fable-5") + require.Equal(t, "claude-fable-5", got) +} + func TestBillingModelForRestriction_Empty(t *testing.T) { t.Parallel() got := billingModelForRestriction("", "claude-sonnet-4-5", "claude-sonnet-4-6") diff --git a/backend/internal/service/gateway_context_management_test.go b/backend/internal/service/gateway_context_management_test.go index b0f2e5b820..bdeb2c600b 100644 --- a/backend/internal/service/gateway_context_management_test.go +++ b/backend/internal/service/gateway_context_management_test.go @@ -596,6 +596,35 @@ func TestBuildCountTokensRequest_OAuthMimicHaiku_PreservesContextManagementEndTo "count_tokens 路径必须含 token-counting beta") } +func TestBuildCountTokensRequest_OAuthMimic_DropsInjectedMaxTokens(t *testing.T) { + // OAuth mimicry injects max_tokens=128000 for normal messages requests. It is + // invalid for Anthropic's count_tokens endpoint and must be stripped on wire. + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages/count_tokens", nil) + + account := &Account{ID: 413, Platform: PlatformAnthropic, Type: AccountTypeOAuth, + Credentials: map[string]any{"access_token": "oauth-tok"}, + Status: StatusActive, Schedulable: true, + } + normalized, _ := normalizeClaudeOAuthRequestBody( + []byte(`{"model":"claude-sonnet-4-5","messages":[]}`), + "claude-sonnet-4-5", claudeOAuthNormalizeOptions{}, + ) + require.Equal(t, int64(128000), gjson.GetBytes(normalized, "max_tokens").Int(), + "precondition: OAuth mimicry injects the Claude Code default") + + svc := &GatewayService{cfg: &config.Config{}} + req, _, err := svc.buildCountTokensRequest( + context.Background(), c, account, normalized, + "oauth-tok", "oauth", "claude-sonnet-4-5", true, + ) + require.NoError(t, err) + require.False(t, gjson.GetBytes(readUpstreamBodyForTest(t, req), "max_tokens").Exists(), + "count_tokens wire body must not contain max_tokens") +} + func TestBuildCountTokensRequest_APIKeyHaiku_StripsContextManagementEndToEnd(t *testing.T) { // API-key + haiku + 客户端 header 不带 context-management beta → final beta 不含 → strip gin.SetMode(gin.TestMode) diff --git a/backend/internal/service/gateway_count_tokens.go b/backend/internal/service/gateway_count_tokens.go index 500f655660..7b66ac9596 100644 --- a/backend/internal/service/gateway_count_tokens.go +++ b/backend/internal/service/gateway_count_tokens.go @@ -584,6 +584,10 @@ func sanitizeCountTokensRequestBody(body []byte) []byte { "stream", "stop_sequences", "stop", + // Anthropic's /v1/messages/count_tokens accepts request-input fields only. + // max_tokens is a generation parameter; OAuth mimicry may inject it to + // resemble Claude Code messages requests, so it must never reach this endpoint. + "max_tokens", } { if gjson.GetBytes(out, path).Exists() { if next, ok := deleteJSONPathBytes(out, path); ok { diff --git a/backend/internal/service/gateway_forward.go b/backend/internal/service/gateway_forward.go index 5a78a19817..d10cb1ea3b 100644 --- a/backend/internal/service/gateway_forward.go +++ b/backend/internal/service/gateway_forward.go @@ -92,6 +92,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A if parsed == nil { return nil, fmt.Errorf("parse request: empty request") } + beginUpstreamResponseModelObservation(c) // Web Search 模拟：纯 web_search 请求时，直接调用搜索 API 构造响应 if account != nil && s.shouldEmulateWebSearch(ctx, account, parsed.GroupID, parsed.Body.Bytes()) { @@ -843,6 +844,11 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A ResponseBody: body, } } + // 流中断（缺失 terminal 事件、读错误、数据间隔超时等）时保留已观测到的 + // usage 与错误一起返回，handler 在错误处理完成后照常提交 usage 记录。 + if partial := partialStreamUsageResult(c, resp, streamResult, originalModel, mappedModel, startTime, err); partial != nil { + return partial, err + } return nil, err } usage = streamResult.usage @@ -856,14 +862,16 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A } return &ForwardResult{ - RequestID: resp.Header.Get("x-request-id"), - Usage: *usage, - Model: originalModel, // 使用原始模型用于计费和日志 - UpstreamModel: mappedModel, - Stream: reqStream, - Duration: time.Since(startTime), - FirstTokenMs: firstTokenMs, - ClientDisconnect: clientDisconnect, + RequestID: resp.Header.Get("x-request-id"), + Usage: *usage, + Model: originalModel, // 使用原始模型用于计费和日志 + UpstreamModel: mappedModel, + UpstreamResponseModel: observedUpstreamResponseModel(c), + UpstreamResponseModelConflict: observedUpstreamResponseModelConflict(c), + Stream: reqStream, + Duration: time.Since(startTime), + FirstTokenMs: firstTokenMs, + ClientDisconnect: clientDisconnect, }, nil } @@ -920,6 +928,10 @@ func billingModelForRestriction(source, requestedModel, channelMappedModel strin return requestedModel case BillingModelSourceUpstream: return "" + case BillingModelSourceResponse: + // The response is not available during dispatch; use mapped pricing + // for restriction prechecks and decide billing after the response. + return channelMappedModel case BillingModelSourceChannelMapped: return channelMappedModel default: diff --git a/backend/internal/service/gateway_forward_partial_usage_test.go b/backend/internal/service/gateway_forward_partial_usage_test.go new file mode 100644 index 0000000000..f6701ca8b0 --- /dev/null +++ b/backend/internal/service/gateway_forward_partial_usage_test.go @@ -0,0 +1,226 @@ +package service + +import ( + "context" + "errors" + "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" +) + +// 本文件覆盖 issue #5148：流式转发中途出错（缺失 terminal 事件、读错误等）时， +// 已观测到的上游 usage 不得随错误一起被丢弃，Forward 必须把部分结果与错误一同 +// 返回，供 handler 照常提交 usage 记录。 + +func newForwardPartialUsageServiceForTest(upstream *anthropicHTTPUpstreamRecorder) *GatewayService { + cfg := &config.Config{ + Gateway: config.GatewayConfig{ + MaxLineSize: defaultMaxLineSize, + }, + } + return &GatewayService{ + cfg: cfg, + responseHeaderFilter: compileResponseHeaderFilter(cfg), + httpUpstream: upstream, + rateLimitService: &RateLimitService{}, + deferredService: &DeferredService{}, + } +} + +func newAnthropicOAuthAccountForPartialUsageTest() *Account { + return &Account{ + ID: 501, + Name: "anthropic-oauth-partial-usage", + Platform: PlatformAnthropic, + Type: AccountTypeOAuth, + Concurrency: 1, + Credentials: map[string]any{ + "access_token": "oauth-token", + }, + Status: StatusActive, + Schedulable: true, + } +} + +func TestGatewayService_Forward_StreamMissingTerminalPreservesPartialUsage(t *testing.T) { + gin.SetMode(gin.TestMode) + + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) + + body := []byte(`{"model":"claude-3-5-sonnet-latest","stream":true,"messages":[{"role":"user","content":[{"type":"text","text":"hello"}]}]}`) + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), PlatformAnthropic) + require.NoError(t, err) + + // newapi 等聚合上游的典型失败形态：message_start/message_delta 携带 usage， + // 但流在 message_stop 前直接结束。 + upstreamSSE := strings.Join([]string{ + `event: message_start`, + `data: {"type":"message_start","message":{"id":"msg_1","type":"message","role":"assistant","model":"claude-3-5-sonnet-latest","content":[],"usage":{"input_tokens":11,"cache_read_input_tokens":7}}}`, + "", + `event: content_block_delta`, + `data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hi"}}`, + "", + `event: message_delta`, + `data: {"type":"message_delta","delta":{"stop_reason":null},"usage":{"output_tokens":5}}`, + "", + "", + }, "\n") + upstream := &anthropicHTTPUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{ + "Content-Type": []string{"text/event-stream"}, + "X-Request-Id": []string{"rid-partial"}, + }, + Body: io.NopCloser(strings.NewReader(upstreamSSE)), + }} + svc := newForwardPartialUsageServiceForTest(upstream) + account := newAnthropicOAuthAccountForPartialUsageTest() + + result, err := svc.Forward(context.Background(), c, account, parsed) + require.Error(t, err) + require.Contains(t, err.Error(), "missing terminal event") + require.NotNil(t, result, "流中断但已观测到 usage 时必须返回部分结果用于计费") + require.True(t, result.Stream) + require.Equal(t, 11, result.Usage.InputTokens) + require.Equal(t, 7, result.Usage.CacheReadInputTokens) + require.Equal(t, 5, result.Usage.OutputTokens) + require.Equal(t, "rid-partial", result.RequestID) + require.NotNil(t, result.FirstTokenMs) +} + +func TestGatewayService_Forward_StreamReadErrorAfterOutputPreservesPartialUsage(t *testing.T) { + gin.SetMode(gin.TestMode) + + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) + + body := []byte(`{"model":"claude-3-5-sonnet-latest","stream":true,"messages":[{"role":"user","content":[{"type":"text","text":"hello"}]}]}`) + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), PlatformAnthropic) + require.NoError(t, err) + + // message_start 已写出（含 usage），随后上游连接异常中断。 + upstream := &anthropicHTTPUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: &streamReadCloser{ + payload: []byte("data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":9,\"cache_creation_input_tokens\":4}}}\n\n"), + err: io.ErrUnexpectedEOF, + }, + }} + svc := newForwardPartialUsageServiceForTest(upstream) + account := newAnthropicOAuthAccountForPartialUsageTest() + + result, err := svc.Forward(context.Background(), c, account, parsed) + require.Error(t, err) + require.Contains(t, err.Error(), "stream read error") + require.NotNil(t, result, "已写出内容后的读错误必须保留部分 usage") + require.Equal(t, 9, result.Usage.InputTokens) + require.Equal(t, 4, result.Usage.CacheCreationInputTokens) +} + +func TestGatewayService_Forward_StreamErrorWithoutUsageReturnsNilResult(t *testing.T) { + gin.SetMode(gin.TestMode) + + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) + + body := []byte(`{"model":"claude-3-5-sonnet-latest","stream":true,"messages":[{"role":"user","content":[{"type":"text","text":"hello"}]}]}`) + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), PlatformAnthropic) + require.NoError(t, err) + + // 只有 ping、没有任何 usage 的流中断：不应产生零 usage 的幽灵账单记录。 + upstream := &anthropicHTTPUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(strings.NewReader("event: ping\ndata: {\"type\": \"ping\"}\n\n")), + }} + svc := newForwardPartialUsageServiceForTest(upstream) + account := newAnthropicOAuthAccountForPartialUsageTest() + + result, err := svc.Forward(context.Background(), c, account, parsed) + require.Error(t, err) + require.Contains(t, err.Error(), "missing terminal event") + require.Nil(t, result, "无已观测 usage 时不应返回部分结果") +} + +func TestGatewayService_Forward_FailoverErrorKeepsNilResult(t *testing.T) { + gin.SetMode(gin.TestMode) + + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) + + body := []byte(`{"model":"claude-3-5-sonnet-latest","stream":true,"messages":[{"role":"user","content":[{"type":"text","text":"hello"}]}]}`) + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), PlatformAnthropic) + require.NoError(t, err) + + // 未向客户端写出任何字节前的读错误会包成 UpstreamFailoverError 走换号重试。 + // 该路径必须保持 result=nil：failover 成功后按成功请求计费，双份结果会重复计费。 + upstream := &anthropicHTTPUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: &streamReadCloser{ + err: errors.New("connection reset by peer"), + }, + }} + svc := newForwardPartialUsageServiceForTest(upstream) + account := newAnthropicOAuthAccountForPartialUsageTest() + + result, err := svc.Forward(context.Background(), c, account, parsed) + require.Error(t, err) + var failoverErr *UpstreamFailoverError + require.True(t, errors.As(err, &failoverErr)) + require.Nil(t, result, "failover 错误必须保持 result=nil，防止重试成功后双重计费") +} + +func TestGatewayService_AnthropicAPIKeyPassthrough_ForwardStreamMissingTerminalPreservesPartialUsage(t *testing.T) { + gin.SetMode(gin.TestMode) + + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) + + body := []byte(`{"model":"claude-3-7-sonnet-20250219","stream":true,"messages":[{"role":"user","content":[{"type":"text","text":"hello"}]}]}`) + parsed := &ParsedRequest{ + Body: NewRequestBodyRef(body), + Model: "claude-3-7-sonnet-20250219", + Stream: true, + } + + upstreamSSE := strings.Join([]string{ + `data: {"type":"message_start","message":{"usage":{"input_tokens":9,"cache_read_input_tokens":2}}}`, + "", + `data: {"type":"message_delta","usage":{"output_tokens":3}}`, + "", + }, "\n") + upstream := &anthropicHTTPUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{ + "Content-Type": []string{"text/event-stream"}, + "X-Request-Id": []string{"rid-pass-partial"}, + }, + Body: io.NopCloser(strings.NewReader(upstreamSSE)), + }} + svc := newForwardPartialUsageServiceForTest(upstream) + account := newAnthropicAPIKeyAccountForTest() + + result, err := svc.Forward(context.Background(), c, account, parsed) + require.Error(t, err) + require.Contains(t, err.Error(), "missing terminal event") + require.NotNil(t, result, "透传流中断但已观测到 usage 时必须返回部分结果用于计费") + require.True(t, result.Stream) + require.Equal(t, 9, result.Usage.InputTokens) + require.Equal(t, 2, result.Usage.CacheReadInputTokens) + require.Equal(t, 3, result.Usage.OutputTokens) + require.Equal(t, "claude-3-7-sonnet-20250219", result.Model) +} diff --git a/backend/internal/service/gateway_hotpath_optimization_test.go b/backend/internal/service/gateway_hotpath_optimization_test.go index 13fc32165f..48cfc358a6 100644 --- a/backend/internal/service/gateway_hotpath_optimization_test.go +++ b/backend/internal/service/gateway_hotpath_optimization_test.go @@ -144,6 +144,20 @@ func (s *stickyGatewayCacheHotpathStub) DeleteSessionAccountID(ctx context.Conte return nil } +func (s *stickyGatewayCacheHotpathStub) SetGrokVideoPendingBilling(_ context.Context, _ string, _ []byte, _ time.Duration) error { + return nil +} +func (s *stickyGatewayCacheHotpathStub) GetGrokVideoPendingBilling(_ context.Context, _ string) ([]byte, error) { + return nil, nil +} +func (s *stickyGatewayCacheHotpathStub) ClaimGrokVideoBilled(_ context.Context, _ string, _ time.Duration) (bool, error) { + return true, nil +} + +func (s *stickyGatewayCacheHotpathStub) ReleaseGrokVideoBilled(_ context.Context, _ string) error { + return nil +} + func (s *modelsListAccountRepoStub) ListSchedulableByGroupID(ctx context.Context, groupID int64) ([]Account, error) { s.listByGroupCalls.Add(1) if s.err != nil { diff --git a/backend/internal/service/gateway_multiplatform_test.go b/backend/internal/service/gateway_multiplatform_test.go index 7b5d967cf5..ac734a3642 100644 --- a/backend/internal/service/gateway_multiplatform_test.go +++ b/backend/internal/service/gateway_multiplatform_test.go @@ -277,6 +277,20 @@ func (m *mockGatewayCacheForPlatform) DeleteSessionAccountID(ctx context.Context return nil } +func (m *mockGatewayCacheForPlatform) SetGrokVideoPendingBilling(_ context.Context, _ string, _ []byte, _ time.Duration) error { + return nil +} +func (m *mockGatewayCacheForPlatform) GetGrokVideoPendingBilling(_ context.Context, _ string) ([]byte, error) { + return nil, nil +} +func (m *mockGatewayCacheForPlatform) ClaimGrokVideoBilled(_ context.Context, _ string, _ time.Duration) (bool, error) { + return true, nil +} + +func (m *mockGatewayCacheForPlatform) ReleaseGrokVideoBilled(_ context.Context, _ string) error { + return nil +} + type mockGroupRepoForGateway struct { groups map[int64]*Group getByIDCalls int diff --git a/backend/internal/service/gateway_profit_control.go b/backend/internal/service/gateway_profit_control.go new file mode 100644 index 0000000000..fe8ba14752 --- /dev/null +++ b/backend/internal/service/gateway_profit_control.go @@ -0,0 +1,112 @@ +package service + +import ( + "context" + "log/slog" + + "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey" +) + +// withGatewayProfitControlGate installs the gate only for explicitly marked +// token requests. This keeps media, metadata, and models-list paths outside +// the profit-control surface by construction. +func (s *GatewayService) withGatewayProfitControlGate(ctx context.Context, groupID *int64) context.Context { + if _, ok := gatewayTokenRequestPricingAtFromContext(ctx); !ok || groupID == nil || *groupID <= 0 { + return ctx + } + if existing, ok := ctx.Value(openAIProfitControlGateCtxKey{}).(*openAIProfitControlGate); ok && existing != nil && existing.groupID == *groupID { + return ctx + } + + group, err := s.resolveProfitControlGroup(ctx, *groupID) + if err != nil { + slog.Warn("profit_control_group_load_failed", "group_id", *groupID, "error", err) + return s.clearForeignProfitControlGate(ctx, groupID) + } + if group == nil || !group.ProfitControlEnabled || !profitControlPlatformSupported(group.Platform) { + return s.clearForeignProfitControlGate(ctx, groupID) + } + + pricingAt, _ := gatewayTokenRequestPricingAtFromContext(ctx) + billingGroup := gatewayTokenRequestBillingGroupFromContext(ctx) + if billingGroup == nil { + if ctxGroup, ok := ctx.Value(ctxkey.Group).(*Group); ok && IsGroupContextValid(ctxGroup) { + billingGroup = ctxGroup + } else { + billingGroup = group + } + } + + downstream := billingGroup.RateMultiplier + if userID, _ := ctx.Value(ctxkey.UserID).(int64); userID > 0 { + downstream = s.ResolveUserGroupRateMultiplier(ctx, userID, billingGroup.ID, billingGroup.RateMultiplier) + } + downstream *= billingGroup.PeakMultiplierAt(pricingAt) + threshold := clampProfitControlThreshold(downstream * (1 - group.ProfitMinMargin - group.ProfitSafetyBuffer)) + + gate := &openAIProfitControlGate{ + groupID: group.ID, + platform: group.Platform, + threshold: threshold, + pricingAt: pricingAt, + } + openAIProfitControlObserverInstance.recordInstall(gate.groupID, gate.platform, gate.threshold) + return context.WithValue(ctx, openAIProfitControlGateCtxKey{}, gate) +} + +func (s *GatewayService) clearForeignProfitControlGate(ctx context.Context, groupID *int64) context.Context { + existing, ok := ctx.Value(openAIProfitControlGateCtxKey{}).(*openAIProfitControlGate) + if !ok || existing == nil || groupID == nil || existing.groupID == *groupID { + return ctx + } + return context.WithValue(ctx, openAIProfitControlGateCtxKey{}, (*openAIProfitControlGate)(nil)) +} + +func (s *GatewayService) resolveProfitControlGroup(ctx context.Context, groupID int64) (*Group, error) { + if group, ok := ctx.Value(ctxkey.Group).(*Group); ok && IsGroupContextValid(group) && group.ID == groupID { + return group, nil + } + if s.schedulerSnapshot != nil { + // Lite 读取：门只用平台/倍率/利润/高峰字段，不需要账号计数聚合。 + return s.schedulerSnapshot.GetGroupByIDLite(ctx, groupID) + } + return s.resolveGroupByID(ctx, groupID) +} + +// GatewayProfitControlVetoLatest performs the terminal post-slot check against +// the latest scheduler snapshot. Snapshot read failures are deliberately +// fail-open to preserve availability, but are observable. +func (s *GatewayService) GatewayProfitControlVetoLatest(ctx context.Context, selected *Account) (*Account, bool, string) { + return profitControlVetoLatest(ctx, selected, s.schedulerSnapshot) +} + +func profitControlVetoLatest(ctx context.Context, selected *Account, snapshot *SchedulerSnapshotService) (*Account, bool, string) { + gate, _ := ctx.Value(openAIProfitControlGateCtxKey{}).(*openAIProfitControlGate) + if gate == nil || selected == nil { + return selected, false, "" + } + latest := selected + if snapshot != nil { + refreshed, err := snapshot.GetAccount(ctx, selected.ID) + if err != nil || refreshed == nil { + slog.Warn("profit_control_account_refresh_failed", "group_id", gate.groupID, "platform", gate.platform, "account_id", selected.ID, "error", err) + openAIProfitControlObserverInstance.recordRefreshFailure(gate.groupID, gate.platform, gate.threshold) + } else if !refreshed.UpdatedAt.Before(selected.UpdatedAt) { + // 选号路径可能已做过 DB recheck，selected 比缓存快照更新鲜；只有 + // 快照不落后时才替换，避免终检把新鲜账号换回较旧的缓存对象。 + latest = refreshed + } + } + vetoed, reason := openAIProfitControlVetoReason(ctx, latest) + return latest, vetoed, reason +} + +func (s *GatewayService) isGatewayAccountProfitEligible(ctx context.Context, account *Account) bool { + vetoed, _ := openAIProfitControlVetoReason(ctx, account) + return !vetoed +} + +func gatewayProfitControlGateActive(ctx context.Context) bool { + gate, _ := ctx.Value(openAIProfitControlGateCtxKey{}).(*openAIProfitControlGate) + return gate != nil +} diff --git a/backend/internal/service/gateway_profit_control_v2_test.go b/backend/internal/service/gateway_profit_control_v2_test.go new file mode 100644 index 0000000000..c05f4501e1 --- /dev/null +++ b/backend/internal/service/gateway_profit_control_v2_test.go @@ -0,0 +1,477 @@ +//go:build unit + +package service + +import ( + "context" + "errors" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey" + "github.com/stretchr/testify/require" +) + +func gatewayProfitTestGroup(id int64, platform string) *Group { + return &Group{ + ID: id, + Name: "profit-" + platform, + Platform: platform, + Status: StatusActive, + Hydrated: true, + RateMultiplier: 0.5, + SubscriptionType: SubscriptionTypeStandard, + ProfitControlEnabled: true, + ProfitMinMargin: 0, + ProfitSafetyBuffer: 0, + } +} + +func gatewayProfitTestContext(group *Group) context.Context { + ctx := context.WithValue(context.Background(), ctxkey.Group, group) + ctx, _ = WithGatewayTokenRequestPricing(ctx) + return ctx +} + +func gatewayProfitTestAccount(id int64, platform string, rate float64, groupID int64) Account { + return Account{ + ID: id, + Name: "account", + Platform: platform, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 2, + Priority: 1, + RateMultiplier: &rate, + AccountGroups: []AccountGroup{{AccountID: id, GroupID: groupID}}, + GroupIDs: []int64{groupID}, + } +} + +func TestGatewayProfitControlInstallsForFivePlatformsOnlyOnTokenRequests(t *testing.T) { + for _, platform := range []string{ + PlatformOpenAI, + PlatformAnthropic, + PlatformGemini, + PlatformGrok, + PlatformAntigravity, + } { + t.Run(platform, func(t *testing.T) { + group := gatewayProfitTestGroup(101, platform) + groupID := group.ID + svc := &GatewayService{} + + tokenCtx := svc.withGatewayProfitControlGate(gatewayProfitTestContext(group), &groupID) + gate, _ := tokenCtx.Value(openAIProfitControlGateCtxKey{}).(*openAIProfitControlGate) + require.NotNil(t, gate) + require.Equal(t, platform, gate.platform) + require.InDelta(t, 0.5, gate.threshold, 1e-12) + + metadataCtx := context.WithValue(context.Background(), ctxkey.Group, group) + metadataCtx = svc.withGatewayProfitControlGate(metadataCtx, &groupID) + gate, _ = metadataCtx.Value(openAIProfitControlGateCtxKey{}).(*openAIProfitControlGate) + require.Nil(t, gate, "未显式标记为 token 请求的入口不得装门") + }) + } +} + +func TestGatewayProfitControlCompositeBillingUsesScheduledMemberConfig(t *testing.T) { + billingGroup := &Group{ + ID: 201, + Platform: PlatformComposite, + Status: StatusActive, + Hydrated: true, + RateMultiplier: 0.4, + SubscriptionType: SubscriptionTypeStandard, + } + memberGroup := gatewayProfitTestGroup(202, PlatformAnthropic) + memberGroup.RateMultiplier = 99 + memberGroup.ProfitMinMargin = 0.25 + + ctx := context.WithValue(context.Background(), ctxkey.Group, billingGroup) + ctx, pricingAt := WithGatewayTokenRequestPricing(ctx) + svc := &GatewayService{ + schedulerSnapshot: NewSchedulerSnapshotService( + nil, + nil, + nil, + profitControlGroupRepo{group: memberGroup}, + nil, + ), + } + ctx = svc.withGatewayProfitControlGate(ctx, &memberGroup.ID) + gate, _ := ctx.Value(openAIProfitControlGateCtxKey{}).(*openAIProfitControlGate) + require.NotNil(t, gate) + require.Equal(t, memberGroup.ID, gate.groupID) + require.Equal(t, PlatformAnthropic, gate.platform) + require.Equal(t, pricingAt, gate.pricingAt) + require.InDelta(t, 0.4*(1-0.25), gate.threshold, 1e-12, "D 必须取 composite 计费父分组，margin 取被调度成员分组") +} + +func TestGatewayProfitControlGroupLoadFailureClearsForeignGate(t *testing.T) { + billingGroup := &Group{ + ID: 211, + Platform: PlatformComposite, + Status: StatusActive, + Hydrated: true, + RateMultiplier: 0.4, + SubscriptionType: SubscriptionTypeStandard, + } + targetGroupID := int64(212) + ctx := gatewayProfitTestContext(billingGroup) + ctx = context.WithValue(ctx, openAIProfitControlGateCtxKey{}, &openAIProfitControlGate{ + groupID: 210, + platform: PlatformAnthropic, + threshold: 0.1, + }) + svc := &GatewayService{ + schedulerSnapshot: NewSchedulerSnapshotService( + nil, + nil, + nil, + profitControlFailingGroupRepo{}, + nil, + ), + } + + ctx = svc.withGatewayProfitControlGate(ctx, &targetGroupID) + gate, ok := ctx.Value(openAIProfitControlGateCtxKey{}).(*openAIProfitControlGate) + require.True(t, ok) + require.Nil(t, gate, "加载新分组失败时必须清除其他分组遗留的门") + + account := gatewayProfitTestAccount(213, PlatformAnthropic, 0.8, targetGroupID) + require.True(t, svc.isGatewayAccountProfitEligible(ctx, &account), "配置读取失败按既定语义 fail-open") +} + +type profitControlFailingGroupRepo struct { + GroupRepository +} + +func (profitControlFailingGroupRepo) GetByIDLite(context.Context, int64) (*Group, error) { + return nil, errors.New("group cache unavailable") +} + +// 见 profitControlGroupRepo.GetByID：利润门必须走不带账号计数聚合的 lite 读取。 +func (profitControlFailingGroupRepo) GetByID(context.Context, int64) (*Group, error) { + panic("profit control gate must read groups via GetByIDLite (no account-count aggregation)") +} + +func TestGatewayProfitControlLegacyMixedAndRoutedSelection(t *testing.T) { + t.Run("legacy single-platform selection", func(t *testing.T) { + group := gatewayProfitTestGroup(111, PlatformGrok) + cheap := gatewayProfitTestAccount(1, PlatformGrok, 0.2, group.ID) + expensive := gatewayProfitTestAccount(2, PlatformGrok, 0.8, group.ID) + repo := &mockAccountRepoForPlatform{ + accounts: []Account{expensive, cheap}, + accountsByID: map[int64]*Account{cheap.ID: &cheap, expensive.ID: &expensive}, + } + svc := &GatewayService{ + accountRepo: repo, + cache: &mockGatewayCacheForPlatform{}, + cfg: testConfig(), + } + + selected, err := svc.SelectAccountForModelWithExclusions( + gatewayProfitTestContext(group), &group.ID, "", "", nil, + ) + require.NoError(t, err) + require.Equal(t, cheap.ID, selected.ID) + + _, err = svc.SelectAccountForModelWithExclusions( + gatewayProfitTestContext(group), &group.ID, "", "", map[int64]struct{}{cheap.ID: {}}, + ) + require.Error(t, err) + require.ErrorIs(t, err, ErrNoAvailableAccounts) + }) + + t.Run("mixed routing filters the routed account", func(t *testing.T) { + group := gatewayProfitTestGroup(112, PlatformAnthropic) + group.ModelRoutingEnabled = true + group.ModelRouting = map[string][]int64{"claude-test": {2, 1}} + cheap := gatewayProfitTestAccount(1, PlatformAntigravity, 0.2, group.ID) + cheap.Extra = map[string]any{"mixed_scheduling": true} + cheap.Credentials = map[string]any{"model_mapping": map[string]any{"claude-test": "claude-test"}} + expensive := gatewayProfitTestAccount(2, PlatformAnthropic, 0.8, group.ID) + repo := &mockAccountRepoForPlatform{ + accounts: []Account{expensive, cheap}, + accountsByID: map[int64]*Account{cheap.ID: &cheap, expensive.ID: &expensive}, + } + svc := &GatewayService{ + accountRepo: repo, + cache: &mockGatewayCacheForPlatform{}, + cfg: testConfig(), + } + + selected, err := svc.SelectAccountForModelWithExclusions( + gatewayProfitTestContext(group), &group.ID, "", "claude-test", nil, + ) + require.NoError(t, err) + require.Equal(t, cheap.ID, selected.ID) + }) +} + +func TestGatewayProfitControlLoadAwareSelectionAndFailover(t *testing.T) { + group := gatewayProfitTestGroup(121, PlatformGrok) + cheap := gatewayProfitTestAccount(1, PlatformGrok, 0.2, group.ID) + expensive := gatewayProfitTestAccount(2, PlatformGrok, 0.8, group.ID) + repo := &mockAccountRepoForPlatform{ + accounts: []Account{expensive, cheap}, + accountsByID: map[int64]*Account{cheap.ID: &cheap, expensive.ID: &expensive}, + } + cfg := &config.Config{RunMode: config.RunModeStandard} + cfg.Gateway.Scheduling.LoadBatchEnabled = true + svc := &GatewayService{ + accountRepo: repo, + cache: &mockGatewayCacheForPlatform{}, + cfg: cfg, + concurrencyService: NewConcurrencyService(stubConcurrencyCache{}), + } + + result, err := svc.SelectAccountWithLoadAwareness( + gatewayProfitTestContext(group), &group.ID, "", "", nil, "", 0, + ) + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, cheap.ID, result.Account.ID) + if result.ReleaseFunc != nil { + result.ReleaseFunc() + } + + result, err = svc.SelectAccountWithLoadAwareness( + gatewayProfitTestContext(group), + &group.ID, + "", + "", + map[int64]struct{}{cheap.ID: {}}, + "", + 0, + ) + require.Nil(t, result) + require.Error(t, err) + require.ErrorIs(t, err, ErrNoAvailableAccounts) +} + +func TestGatewayProfitControlStickyVetoKeepsBindingUntilRateRecovers(t *testing.T) { + group := gatewayProfitTestGroup(131, PlatformAnthropic) + expensive := gatewayProfitTestAccount(1, PlatformAnthropic, 0.8, group.ID) + cheap := gatewayProfitTestAccount(2, PlatformAnthropic, 0.2, group.ID) + repo := &mockAccountRepoForPlatform{ + accounts: []Account{expensive, cheap}, + accountsByID: map[int64]*Account{expensive.ID: &expensive, cheap.ID: &cheap}, + } + cache := &mockGatewayCacheForPlatform{ + sessionBindings: map[string]int64{"sticky-profit": expensive.ID}, + } + svc := &GatewayService{ + accountRepo: repo, + cache: cache, + cfg: testConfig(), + } + ctx := gatewayProfitTestContext(group) + + selected, err := svc.SelectAccountForModelWithExclusions(ctx, &group.ID, "sticky-profit", "", nil) + require.NoError(t, err) + require.Equal(t, cheap.ID, selected.ID) + require.Equal(t, expensive.ID, cache.sessionBindings["sticky-profit"], "候选过滤不得覆盖旧粘性绑定") + + require.NoError(t, svc.BindStickySessionAfterProfitAdmission( + svc.withGatewayProfitControlGate(ctx, &group.ID), + &group.ID, + "sticky-profit", + cheap.ID, + )) + require.Equal(t, expensive.ID, cache.sessionBindings["sticky-profit"], "终检通过的 fallback 账号也不得覆盖旧绑定") + require.Zero(t, cache.deletedSessions["sticky-profit"]) + + recovered := expensive + recoveredRate := 0.2 + recovered.RateMultiplier = &recoveredRate + repo.accounts[0] = recovered + repo.accountsByID[recovered.ID] = &repo.accounts[0] + + selected, err = svc.SelectAccountForModelWithExclusions(ctx, &group.ID, "sticky-profit", "", nil) + require.NoError(t, err) + require.Equal(t, recovered.ID, selected.ID, "倍率恢复后应重新命中原粘性账号") + require.Zero(t, cache.deletedSessions["sticky-profit"]) +} + +type gatewayProfitSnapshotCache struct { + SchedulerCache + account *Account + err error +} + +func (c *gatewayProfitSnapshotCache) GetAccount(context.Context, int64) (*Account, error) { + return c.account, c.err +} + +type gatewayProfitAccountRepo struct { + AccountRepository + account *Account + err error +} + +func (r gatewayProfitAccountRepo) GetByID(context.Context, int64) (*Account, error) { + return r.account, r.err +} + +func TestGatewayProfitControlTerminalRefreshUsesReplacementObject(t *testing.T) { + selected := gatewayProfitTestAccount(141, PlatformGemini, 0.2, 1) + replacement := selected + expensiveRate := 0.8 + replacement.RateMultiplier = &expensiveRate + + snapshot := NewSchedulerSnapshotService( + &gatewayProfitSnapshotCache{account: &replacement}, + nil, + gatewayProfitAccountRepo{}, + nil, + nil, + ) + ctx := context.WithValue(context.Background(), openAIProfitControlGateCtxKey{}, &openAIProfitControlGate{ + groupID: 1, + platform: PlatformGemini, + threshold: 0.5, + }) + + latest, vetoed, reason := profitControlVetoLatest(ctx, &selected, snapshot) + require.Same(t, &replacement, latest) + require.True(t, vetoed) + require.Equal(t, openAIProfitFilterReasonThreshold, reason) + require.InDelta(t, 0.2, *selected.RateMultiplier, 1e-12, "测试必须替换缓存对象，不能原地修改旧指针") +} + +func TestGatewayProfitControlTerminalRefreshFallsBackFromCacheToDatabase(t *testing.T) { + selected := gatewayProfitTestAccount(145, PlatformAnthropic, 0.2, 1) + replacement := selected + expensiveRate := 0.8 + replacement.RateMultiplier = &expensiveRate + + snapshot := NewSchedulerSnapshotService( + &gatewayProfitSnapshotCache{err: errors.New("cache unavailable")}, + nil, + gatewayProfitAccountRepo{account: &replacement}, + nil, + nil, + ) + ctx := context.WithValue(context.Background(), openAIProfitControlGateCtxKey{}, &openAIProfitControlGate{ + groupID: 1, + platform: PlatformAnthropic, + threshold: 0.5, + }) + + latest, vetoed, reason := profitControlVetoLatest(ctx, &selected, snapshot) + require.Same(t, &replacement, latest) + require.True(t, vetoed, "缓存读取失败时必须继续从数据库重读，不能直接使用选号旧对象") + require.Equal(t, openAIProfitFilterReasonThreshold, reason) +} + +func TestGatewayProfitControlTerminalRefreshFailureFallsBackToSelectedObject(t *testing.T) { + selected := gatewayProfitTestAccount(151, PlatformAntigravity, 0.2, 1) + snapshot := NewSchedulerSnapshotService( + &gatewayProfitSnapshotCache{err: errors.New("cache unavailable")}, + nil, + gatewayProfitAccountRepo{err: errors.New("database unavailable")}, + nil, + nil, + ) + ctx := context.WithValue(context.Background(), openAIProfitControlGateCtxKey{}, &openAIProfitControlGate{ + groupID: 1, + platform: PlatformAntigravity, + threshold: 0.5, + }) + + latest, vetoed, reason := profitControlVetoLatest(ctx, &selected, snapshot) + require.Same(t, &selected, latest) + require.False(t, vetoed) + require.Empty(t, reason) +} + +// 选号结果携带门：门安装在调度栈局部 ctx 上，handler 必须经 +// ContextWithSelectionProfitGate 重放后终检与准入后绑定才可见（评审修复回归）。 +func TestGatewayProfitControlSelectionCarriesGateToHandlerContext(t *testing.T) { + group := gatewayProfitTestGroup(1, PlatformAnthropic) + svc := &GatewayService{} + expensive := gatewayProfitTestAccount(161, PlatformAnthropic, 0.9, group.ID) + + gateCtx := svc.withGatewayProfitControlGate(gatewayProfitTestContext(group), &group.ID) + selection, err := svc.newSelectionResult(gateCtx, &expensive, true, nil, nil) + require.NoError(t, err) + require.True(t, selection.ProfitGateActive(), "选号结果必须携带调度栈内生效的门") + + // 修复前的缺陷形态：handler 原始 ctx 不含门，终检退化为空操作。 + _, vetoed, _ := svc.GatewayProfitControlVetoLatest(context.Background(), &expensive) + require.False(t, vetoed, "对照组：不重放门时终检确实看不到门") + + handlerCtx := ContextWithSelectionProfitGate(context.Background(), selection) + latest, vetoed, reason := svc.GatewayProfitControlVetoLatest(handlerCtx, &expensive) + require.True(t, vetoed, "重放门后终检必须真实生效") + require.Equal(t, openAIProfitFilterReasonThreshold, reason) + require.NotNil(t, latest) + + // 无门选号不携带门，重放为无操作。 + plain, err := svc.newSelectionResult(context.Background(), &expensive, true, nil, nil) + require.NoError(t, err) + require.False(t, plain.ProfitGateActive()) + require.Equal(t, context.Background(), ContextWithSelectionProfitGate(context.Background(), plain)) +} + +// 生图意图不关门（H1/H2 回归锚点）：/v1/responses 混合请求即使带生图声明， +// token 定价上下文照常装配，共享门照常安装并否决越线账号。 +func TestGatewayProfitControlImageIntentDoesNotDisableGate(t *testing.T) { + group := gatewayProfitTestGroup(2, PlatformAnthropic) + svc := &GatewayService{} + expensive := gatewayProfitTestAccount(162, PlatformAnthropic, 0.9, group.ID) + + ctx := gatewayProfitTestContext(group) + ctx = WithOpenAIImageGenerationIntent(ctx) + gateCtx := svc.withGatewayProfitControlGate(ctx, &group.ID) + require.False(t, svc.isGatewayAccountProfitEligible(gateCtx, &expensive), + "请求体里的生图声明（含被动 image_gen namespace）不得关闭利润门") +} + +// 无门时准入后绑定回退官方 eager 语义；门下读失败保守不写（评审 M-Bind 回归）。 +func TestGatewayProfitControlAfterAdmissionBindSemantics(t *testing.T) { + groupID := int64(3) + expensiveID := int64(171) + cheapID := int64(172) + + t.Run("eager without gate", func(t *testing.T) { + cache := &mockGatewayCacheForPlatform{sessionBindings: map[string]int64{"s": expensiveID}} + svc := &GatewayService{cache: cache} + require.NoError(t, svc.BindStickySessionAfterProfitAdmission(context.Background(), &groupID, "s", cheapID)) + require.Equal(t, cheapID, cache.sessionBindings["s"], "无门时保持既有 eager 绑定行为") + }) + + t.Run("gated read failure is conservative", func(t *testing.T) { + // mock 的 miss 返回非 sentinel 错误，等价于 Redis 读失败：门下保守不写。 + cache := &mockGatewayCacheForPlatform{sessionBindings: map[string]int64{}} + svc := &GatewayService{cache: cache} + gate := &openAIProfitControlGate{groupID: groupID, platform: PlatformAnthropic, threshold: 0.5} + gateCtx := context.WithValue(context.Background(), openAIProfitControlGateCtxKey{}, gate) + require.NoError(t, svc.BindStickySessionAfterProfitAdmission(gateCtx, &groupID, "absent", cheapID)) + require.NotContains(t, cache.sessionBindings, "absent") + }) + + t.Run("gated sentinel miss binds", func(t *testing.T) { + cache := &sentinelMissGatewayCache{mockGatewayCacheForPlatform: &mockGatewayCacheForPlatform{sessionBindings: map[string]int64{}}} + svc := &GatewayService{cache: cache} + gate := &openAIProfitControlGate{groupID: groupID, platform: PlatformAnthropic, threshold: 0.5} + gateCtx := context.WithValue(context.Background(), openAIProfitControlGateCtxKey{}, gate) + require.NoError(t, svc.BindStickySessionAfterProfitAdmission(gateCtx, &groupID, "fresh", cheapID)) + require.Equal(t, cheapID, cache.sessionBindings["fresh"], "门下无既有绑定（sentinel miss）应建立粘性") + }) +} + +// sentinelMissGatewayCache 让 miss 返回与真实仓库一致的 ErrStickySessionNotFound。 +type sentinelMissGatewayCache struct { + *mockGatewayCacheForPlatform +} + +func (c *sentinelMissGatewayCache) GetSessionAccountID(ctx context.Context, groupID int64, sessionHash string) (int64, error) { + if id, ok := c.sessionBindings[sessionHash]; ok { + return id, nil + } + return 0, ErrStickySessionNotFound +} diff --git a/backend/internal/service/gateway_record_usage_test.go b/backend/internal/service/gateway_record_usage_test.go index 81aff58278..517d4723cb 100644 --- a/backend/internal/service/gateway_record_usage_test.go +++ b/backend/internal/service/gateway_record_usage_test.go @@ -345,6 +345,53 @@ func TestGatewayServiceRecordUsage_PeakRateAffectsTokenModeImageOutputTokens(t * require.InDelta(t, expectedActual, userRepo.lastAmount, 1e-12) } +func TestGatewayServiceRecordUsage_UsesExplicitPricingAtForPeakRate(t *testing.T) { + for _, platform := range []string{PlatformAnthropic, PlatformGemini, PlatformGrok, PlatformAntigravity} { + t.Run(platform, func(t *testing.T) { + groupID := int64(903) + usageRepo := &openAIRecordUsageLogRepoStub{inserted: true} + userRepo := &openAIRecordUsageUserRepoStub{} + svc := newGatewayRecordUsageServiceForTest(usageRepo, userRepo, &openAIRecordUsageSubRepoStub{}) + svc.resolver = newOpenAITokenImageChannelPricingResolverForTest(t, groupID, "gemini-image") + + pricingAt := time.Date(2026, time.January, 1, 0, 30, 0, 0, time.UTC) + err := svc.RecordUsage(context.Background(), &RecordUsageInput{ + Result: &ForwardResult{ + RequestID: "gateway_explicit_pricing_at_" + platform, + Model: "gemini-image", + ImageCount: 1, + Usage: ClaudeUsage{ + InputTokens: 1000, + OutputTokens: 600, + ImageOutputTokens: 100, + }, + }, + APIKey: &APIKey{ + ID: 803, + GroupID: i64p(groupID), + Group: &Group{ + ID: groupID, + Platform: platform, + RateMultiplier: 1.0, + SubscriptionType: SubscriptionTypeSubscription, + PeakRateEnabled: true, + PeakStart: "00:00", + PeakEnd: "01:00", + PeakRateMultiplier: 3.0, + }, + }, + User: &User{ID: 603}, + Account: &Account{ID: 703, Platform: platform}, + PricingAt: pricingAt, + }) + + require.NoError(t, err) + require.NotNil(t, usageRepo.lastLog) + require.Equal(t, 3.0, usageRepo.lastLog.RateMultiplier) + }) + } +} + func TestGatewayServiceRecordUsage_UsageLogWriteErrorDoesNotSkipBilling(t *testing.T) { usageRepo := &openAIRecordUsageLogRepoStub{inserted: false, err: MarkUsageLogCreateNotPersisted(context.Canceled)} userRepo := &openAIRecordUsageUserRepoStub{} @@ -530,9 +577,10 @@ func TestGatewayServiceRecordUsage_DroppedUsageLogFallsBackToSyncCreate(t *testi require.NoError(t, usageRepo.lastCtxErr) } -func TestGatewayServiceRecordUsage_BillingErrorSkipsUsageLogWrite(t *testing.T) { +func TestGatewayServiceRecordUsage_BillingErrorWritesUnsettledUsageLog(t *testing.T) { usageRepo := &openAIRecordUsageLogRepoStub{} - billingRepo := &openAIRecordUsageBillingRepoStub{err: context.DeadlineExceeded} + billingErr := errors.New("billing tx failed") + billingRepo := &openAIRecordUsageBillingRepoStub{err: billingErr} userRepo := &openAIRecordUsageUserRepoStub{} subRepo := &openAIRecordUsageSubRepoStub{} svc := newGatewayRecordUsageServiceWithBillingRepoForTest(usageRepo, billingRepo, userRepo, subRepo) @@ -552,9 +600,16 @@ func TestGatewayServiceRecordUsage_BillingErrorSkipsUsageLogWrite(t *testing.T) Account: &Account{ID: 705}, }) - require.Error(t, err) + require.ErrorIs(t, err, billingErr) require.Equal(t, 1, billingRepo.calls) - require.Equal(t, 0, usageRepo.calls) + require.Equal(t, 1, usageRepo.calls) + require.NotNil(t, usageRepo.lastLog) + require.Equal(t, 10, usageRepo.lastLog.InputTokens) + require.Equal(t, 6, usageRepo.lastLog.OutputTokens) + require.Greater(t, usageRepo.lastLog.InputCost, 0.0) + require.Greater(t, usageRepo.lastLog.OutputCost, 0.0) + require.Greater(t, usageRepo.lastLog.TotalCost, 0.0) + require.Zero(t, usageRepo.lastLog.ActualCost) } func TestGatewayServiceRecordUsage_ReasoningEffortPersisted(t *testing.T) { diff --git a/backend/internal/service/gateway_request_pricing.go b/backend/internal/service/gateway_request_pricing.go new file mode 100644 index 0000000000..d86d9b8bdd --- /dev/null +++ b/backend/internal/service/gateway_request_pricing.go @@ -0,0 +1,55 @@ +package service + +import ( + "context" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey" + "github.com/Wei-Shaw/sub2api/internal/pkg/timezone" +) + +type gatewayTokenRequestPricingAtCtxKey struct{} +type gatewayTokenRequestBillingGroupCtxKey struct{} + +// WithGatewayTokenRequestPricing marks a shared-gateway request as token billed +// and freezes the downstream pricing instant for its whole lifetime. Media and +// metadata-only handlers deliberately do not call this helper. +func WithGatewayTokenRequestPricing(ctx context.Context) (context.Context, time.Time) { + if ctx == nil { + ctx = context.Background() + } + pricingAt := timezone.Now() + ctx = context.WithValue(ctx, gatewayTokenRequestPricingAtCtxKey{}, pricingAt) + // 调度过程中可能因 fallback/composite 路由覆盖 ctxkey.Group；计费 D 仍必须 + // 使用认证时刻的父分组，和最终 RecordUsage 的计费归属保持一致。 + if group, ok := ctx.Value(ctxkey.Group).(*Group); ok && IsGroupContextValid(group) { + ctx = context.WithValue(ctx, gatewayTokenRequestBillingGroupCtxKey{}, group) + } + return ctx, pricingAt +} + +func gatewayTokenRequestPricingAtFromContext(ctx context.Context) (time.Time, bool) { + if ctx == nil { + return time.Time{}, false + } + pricingAt, ok := ctx.Value(gatewayTokenRequestPricingAtCtxKey{}).(time.Time) + return pricingAt, ok && !pricingAt.IsZero() +} + +// GatewayTokenRequestPricingAtFromContext exposes the frozen instant to +// handlers before they detach asynchronous usage-recording work. +func GatewayTokenRequestPricingAtFromContext(ctx context.Context) time.Time { + pricingAt, _ := gatewayTokenRequestPricingAtFromContext(ctx) + return pricingAt +} + +func gatewayTokenRequestBillingGroupFromContext(ctx context.Context) *Group { + if ctx == nil { + return nil + } + group, _ := ctx.Value(gatewayTokenRequestBillingGroupCtxKey{}).(*Group) + if IsGroupContextValid(group) { + return group + } + return nil +} diff --git a/backend/internal/service/gateway_request_pricing_test.go b/backend/internal/service/gateway_request_pricing_test.go new file mode 100644 index 0000000000..3f0faf14eb --- /dev/null +++ b/backend/internal/service/gateway_request_pricing_test.go @@ -0,0 +1,20 @@ +//go:build unit + +package service + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestWithGatewayTokenRequestPricingMarksOnlyExplicitTokenRequests(t *testing.T) { + ctx, pricingAt := WithGatewayTokenRequestPricing(context.Background()) + + got, ok := gatewayTokenRequestPricingAtFromContext(ctx) + require.True(t, ok) + require.Equal(t, pricingAt, got) + require.Equal(t, pricingAt, GatewayTokenRequestPricingAtFromContext(ctx)) + require.True(t, GatewayTokenRequestPricingAtFromContext(context.Background()).IsZero()) +} diff --git a/backend/internal/service/gateway_scheduling.go b/backend/internal/service/gateway_scheduling.go index 99cfd02c4e..be63f804f5 100644 --- a/backend/internal/service/gateway_scheduling.go +++ b/backend/internal/service/gateway_scheduling.go @@ -64,6 +64,7 @@ func (s *GatewayService) SelectAccountForModelWithExclusions(ctx context.Context // 无分组时只使用原生 anthropic 平台 platform = PlatformAnthropic } + ctx = s.withGatewayProfitControlGate(ctx, groupID) // Claude Code 限制可能已将 groupID 解析为 fallback group， // 渠道限制预检查必须使用解析后的分组。 @@ -116,6 +117,7 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro return nil, err } ctx = s.withGroupContext(ctx, group) + ctx = s.withGatewayProfitControlGate(ctx, groupID) // Claude Code 限制可能已将 groupID 解析为 fallback group， // 渠道限制预检查必须使用解析后的分组。 @@ -284,6 +286,9 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro } continue } + if !s.isGatewayAccountProfitEligible(ctx, account) { + continue + } if !s.isAccountAllowedForPlatform(account, platform, useMixed) { filteredPlatform++ continue @@ -339,6 +344,7 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro var stickyCacheMissReason string gatePass := s.isAccountSchedulableForSelection(stickyAccount) && + s.isGatewayAccountProfitEligible(ctx, stickyAccount) && s.isAccountAllowedForPlatform(stickyAccount, platform, useMixed) && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, stickyAccount, requestedModel)) && s.isAccountSchedulableForModelSelection(ctx, stickyAccount, requestedModel) && @@ -470,7 +476,7 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro continue } if sessionHash != "" && s.cache != nil { - _ = s.cache.SetSessionAccountID(ctx, derefGroupID(groupID), sessionHash, item.account.ID, stickySessionTTL) + _ = s.bindGatewayStickySessionDuringSelection(ctx, groupID, sessionHash, item.account.ID) } if s.debugModelRoutingEnabled() { logger.LegacyPrintf("service.gateway", "[ModelRoutingDebug] routed select: group_id=%v model=%s session=%s account=%d", derefGroupID(groupID), requestedModel, shortSessionHash(sessionHash), item.account.ID) @@ -523,6 +529,7 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro // accounts 列表构建，账号一定在分组内。而 scheduler snapshot 缓存 // 反序列化后 AccountGroups 字段为空，导致 isAccountInGroup 永远返回 false。 platformOK := s.isAccountAllowedForPlatform(account, platform, useMixed) + profitOK := s.isGatewayAccountProfitEligible(ctx, account) modelSupported := requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel) modelSchedulable := s.isAccountSchedulableForModelSelection(ctx, account, requestedModel) quotaOK := s.isAccountSchedulableForQuota(account) @@ -536,6 +543,7 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro "clear_sticky", clearSticky, "schedulable", schedulable, "platform_ok", platformOK, + "profit_ok", profitOK, "model_supported", modelSupported, "model_schedulable", modelSchedulable, "quota_ok", quotaOK, @@ -543,7 +551,7 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro "rpm_ok", rpmOK, ) - if !clearSticky && platformOK && modelSupported && modelSchedulable && quotaOK && windowCostOK && rpmOK && schedulable { + if !clearSticky && platformOK && profitOK && modelSupported && modelSchedulable && quotaOK && windowCostOK && rpmOK && schedulable { result, err := s.tryAcquireAccountSlot(ctx, accountID, account.Concurrency) if err == nil && result.Acquired { // 会话数量限制检查 @@ -639,6 +647,9 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro if !s.isAccountSchedulableForSelection(acc) { continue } + if !s.isGatewayAccountProfitEligible(ctx, acc) { + continue + } if !s.isAccountAllowedForPlatform(acc, platform, useMixed) { continue } @@ -720,7 +731,7 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro result.ReleaseFunc() // 释放槽位，继续尝试下一个账号 } else { if sessionHash != "" && s.cache != nil { - _ = s.cache.SetSessionAccountID(ctx, derefGroupID(groupID), sessionHash, selected.account.ID, stickySessionTTL) + _ = s.bindGatewayStickySessionDuringSelection(ctx, groupID, sessionHash, selected.account.ID) } return s.newSelectionResult(ctx, selected.account, true, result.ReleaseFunc, nil) } @@ -768,7 +779,7 @@ func (s *GatewayService) tryAcquireByLegacyOrder(ctx context.Context, candidates continue } if sessionHash != "" && s.cache != nil { - _ = s.cache.SetSessionAccountID(ctx, derefGroupID(groupID), sessionHash, acc.ID, stickySessionTTL) + _ = s.bindGatewayStickySessionDuringSelection(ctx, groupID, sessionHash, acc.ID) } selection, err := s.newSelectionResult(ctx, acc, true, result.ReleaseFunc, nil) if err != nil { @@ -950,6 +961,10 @@ func (s *GatewayService) listSchedulableAccounts(ctx context.Context, groupID *i if s.schedulerSnapshot != nil { accounts, useMixed, err := s.schedulerSnapshot.ListSchedulableAccounts(ctx, groupID, platform, hasForcePlatform) if err == nil { + accounts = s.filterAccountsBySchedulingThreshold(ctx, accounts) + if platform == PlatformGrok || strings.EqualFold(platform, PlatformGrok) { + accounts = s.filterGrokFreeQuotaAccountsForGateway(ctx, accounts) + } slog.Debug("account_scheduling_list_snapshot", "group_id", derefGroupID(groupID), "platform", platform, @@ -1011,7 +1026,7 @@ func (s *GatewayService) listSchedulableAccounts(ctx context.Context, groupID *i "tls_fingerprint", acc.IsTLSFingerprintEnabled()) } } - return filtered, useMixed, nil + return s.filterAccountsBySchedulingThreshold(ctx, filtered), useMixed, nil } var accounts []Account @@ -1046,6 +1061,10 @@ func (s *GatewayService) listSchedulableAccounts(ctx context.Context, groupID *i "tls_fingerprint", acc.IsTLSFingerprintEnabled()) } } + accounts = s.filterAccountsBySchedulingThreshold(ctx, accounts) + if platform == PlatformGrok || strings.EqualFold(platform, PlatformGrok) { + accounts = s.filterGrokFreeQuotaAccountsForGateway(ctx, accounts) + } return accounts, useMixed, nil } @@ -1417,10 +1436,50 @@ func (s *GatewayService) checkAndRegisterSession(ctx context.Context, account *A } func (s *GatewayService) getSchedulableAccount(ctx context.Context, accountID int64) (*Account, error) { + var ( + account *Account + err error + ) if s.schedulerSnapshot != nil { - return s.schedulerSnapshot.GetAccount(ctx, accountID) + account, err = s.schedulerSnapshot.GetAccount(ctx, accountID) + } else { + account, err = s.accountRepo.GetByID(ctx, accountID) } - return s.accountRepo.GetByID(ctx, accountID) + if err != nil || account == nil { + return account, err + } + if s.isAccountBlockedBySchedulingThreshold(ctx, account) { + return nil, nil + } + // Sticky / non-list selection must honor free soft-gate (same as listSchedulableAccounts). + if account.IsGrok() { + if gated := s.filterGrokFreeQuotaAccountsForGateway(ctx, []Account{*account}); len(gated) == 0 { + return nil, nil + } + } + return account, nil +} + +func (s *GatewayService) filterAccountsBySchedulingThreshold(ctx context.Context, accounts []Account) []Account { + if len(accounts) == 0 { + return accounts + } + + filtered := make([]Account, 0, len(accounts)) + for i := range accounts { + if s.isAccountBlockedBySchedulingThreshold(ctx, &accounts[i]) { + continue + } + filtered = append(filtered, accounts[i]) + } + return filtered +} + +func (s *GatewayService) isAccountBlockedBySchedulingThreshold(ctx context.Context, account *Account) bool { + if s == nil || s.rateLimitService == nil || account == nil { + return false + } + return s.rateLimitService.ApplyAccountSchedulingThreshold(ctx, account) } func (s *GatewayService) hydrateSelectedAccount(ctx context.Context, account *Account) (*Account, error) { @@ -1442,12 +1501,12 @@ func (s *GatewayService) newSelectionResult(ctx context.Context, account *Accoun if err != nil { return nil, err } - return &AccountSelectionResult{ + return attachSelectionProfitGate(ctx, &AccountSelectionResult{ Account: hydrated, Acquired: acquired, ReleaseFunc: release, WaitPlan: waitPlan, - }, nil + }), nil } // filterByMinPriority 过滤出优先级最小的账号集合 @@ -1787,7 +1846,7 @@ func (s *GatewayService) selectAccountForModelWithPlatform(ctx context.Context, if clearSticky { _ = s.cache.DeleteSessionAccountID(ctx, derefGroupID(groupID), sessionHash) } - if !clearSticky && s.isAccountInGroup(account, groupID) && account.Platform == platform && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && s.isAccountSchedulableForModelSelection(ctx, account, requestedModel) && s.isAccountSchedulableForQuota(account) && s.isAccountSchedulableForWindowCost(ctx, account, true) && s.isAccountSchedulableForRPM(ctx, account, true) && !s.isStickyAccountUpstreamRestricted(ctx, groupID, account, requestedModel) { + if !clearSticky && s.isGatewayAccountProfitEligible(ctx, account) && s.isAccountInGroup(account, groupID) && account.Platform == platform && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && s.isAccountSchedulableForModelSelection(ctx, account, requestedModel) && s.isAccountSchedulableForQuota(account) && s.isAccountSchedulableForWindowCost(ctx, account, true) && s.isAccountSchedulableForRPM(ctx, account, true) && !s.isStickyAccountUpstreamRestricted(ctx, groupID, account, requestedModel) { if s.debugModelRoutingEnabled() { logger.LegacyPrintf("service.gateway", "[ModelRoutingDebug] legacy routed sticky hit: group_id=%v model=%s session=%s account=%d", derefGroupID(groupID), requestedModel, shortSessionHash(sessionHash), accountID) } @@ -1835,6 +1894,9 @@ func (s *GatewayService) selectAccountForModelWithPlatform(ctx context.Context, if !s.isAccountSchedulableForSelection(acc) { continue } + if !s.isGatewayAccountProfitEligible(ctx, acc) { + continue + } // require_privacy_set: 跳过 privacy 未设置的账号并标记异常 if schedGroup != nil && schedGroup.RequirePrivacySet && !acc.IsPrivacySet() { _ = s.accountRepo.SetError(ctx, acc.ID, @@ -1882,7 +1944,7 @@ func (s *GatewayService) selectAccountForModelWithPlatform(ctx context.Context, if selected != nil { if sessionHash != "" && s.cache != nil { - if err := s.cache.SetSessionAccountID(ctx, derefGroupID(groupID), sessionHash, selected.ID, stickySessionTTL); err != nil { + if err := s.bindGatewayStickySessionDuringSelection(ctx, groupID, sessionHash, selected.ID); err != nil { logger.LegacyPrintf("service.gateway", "set session account failed: session=%s account_id=%d err=%v", sessionHash, selected.ID, err) } } @@ -1906,7 +1968,7 @@ func (s *GatewayService) selectAccountForModelWithPlatform(ctx context.Context, if clearSticky { _ = s.cache.DeleteSessionAccountID(ctx, derefGroupID(groupID), sessionHash) } - if !clearSticky && s.isAccountInGroup(account, groupID) && account.Platform == platform && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && s.isAccountSchedulableForModelSelection(ctx, account, requestedModel) && s.isAccountSchedulableForQuota(account) && s.isAccountSchedulableForWindowCost(ctx, account, true) && s.isAccountSchedulableForRPM(ctx, account, true) { + if !clearSticky && s.isGatewayAccountProfitEligible(ctx, account) && s.isAccountInGroup(account, groupID) && account.Platform == platform && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && s.isAccountSchedulableForModelSelection(ctx, account, requestedModel) && s.isAccountSchedulableForQuota(account) && s.isAccountSchedulableForWindowCost(ctx, account, true) && s.isAccountSchedulableForRPM(ctx, account, true) { return account, nil } } @@ -1946,6 +2008,9 @@ func (s *GatewayService) selectAccountForModelWithPlatform(ctx context.Context, if !s.isAccountSchedulableForSelection(acc) { continue } + if !s.isGatewayAccountProfitEligible(ctx, acc) { + continue + } // require_privacy_set: 跳过 privacy 未设置的账号并标记异常 if schedGroup != nil && schedGroup.RequirePrivacySet && !acc.IsPrivacySet() { _ = s.accountRepo.SetError(ctx, acc.ID, @@ -2004,7 +2069,7 @@ func (s *GatewayService) selectAccountForModelWithPlatform(ctx context.Context, // 4. 建立粘性绑定 if sessionHash != "" && s.cache != nil { - if err := s.cache.SetSessionAccountID(ctx, derefGroupID(groupID), sessionHash, selected.ID, stickySessionTTL); err != nil { + if err := s.bindGatewayStickySessionDuringSelection(ctx, groupID, sessionHash, selected.ID); err != nil { logger.LegacyPrintf("service.gateway", "set session account failed: session=%s account_id=%d err=%v", sessionHash, selected.ID, err) } } @@ -2045,7 +2110,7 @@ func (s *GatewayService) selectAccountWithMixedScheduling(ctx context.Context, g if clearSticky { _ = s.cache.DeleteSessionAccountID(ctx, derefGroupID(groupID), sessionHash) } - if !clearSticky && s.isAccountInGroup(account, groupID) && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && s.isAccountSchedulableForModelSelection(ctx, account, requestedModel) && s.isAccountSchedulableForQuota(account) && s.isAccountSchedulableForWindowCost(ctx, account, true) && s.isAccountSchedulableForRPM(ctx, account, true) { + if !clearSticky && s.isGatewayAccountProfitEligible(ctx, account) && s.isAccountInGroup(account, groupID) && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && s.isAccountSchedulableForModelSelection(ctx, account, requestedModel) && s.isAccountSchedulableForQuota(account) && s.isAccountSchedulableForWindowCost(ctx, account, true) && s.isAccountSchedulableForRPM(ctx, account, true) { if account.Platform == nativePlatform || (account.Platform == PlatformAntigravity && account.IsMixedSchedulingEnabled()) { if s.debugModelRoutingEnabled() { logger.LegacyPrintf("service.gateway", "[ModelRoutingDebug] legacy mixed routed sticky hit: group_id=%v model=%s session=%s account=%d", derefGroupID(groupID), requestedModel, shortSessionHash(sessionHash), accountID) @@ -2091,6 +2156,9 @@ func (s *GatewayService) selectAccountWithMixedScheduling(ctx context.Context, g if !s.isAccountSchedulableForSelection(acc) { continue } + if !s.isGatewayAccountProfitEligible(ctx, acc) { + continue + } // require_privacy_set: 跳过 privacy 未设置的账号并标记异常 if schedGroup != nil && schedGroup.RequirePrivacySet && !acc.IsPrivacySet() { _ = s.accountRepo.SetError(ctx, acc.ID, @@ -2142,7 +2210,7 @@ func (s *GatewayService) selectAccountWithMixedScheduling(ctx context.Context, g if selected != nil { if sessionHash != "" && s.cache != nil { - if err := s.cache.SetSessionAccountID(ctx, derefGroupID(groupID), sessionHash, selected.ID, stickySessionTTL); err != nil { + if err := s.bindGatewayStickySessionDuringSelection(ctx, groupID, sessionHash, selected.ID); err != nil { logger.LegacyPrintf("service.gateway", "set session account failed: session=%s account_id=%d err=%v", sessionHash, selected.ID, err) } } @@ -2166,7 +2234,7 @@ func (s *GatewayService) selectAccountWithMixedScheduling(ctx context.Context, g if clearSticky { _ = s.cache.DeleteSessionAccountID(ctx, derefGroupID(groupID), sessionHash) } - if !clearSticky && s.isAccountInGroup(account, groupID) && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && s.isAccountSchedulableForModelSelection(ctx, account, requestedModel) && s.isAccountSchedulableForQuota(account) && s.isAccountSchedulableForWindowCost(ctx, account, true) && s.isAccountSchedulableForRPM(ctx, account, true) && !s.isStickyAccountUpstreamRestricted(ctx, groupID, account, requestedModel) { + if !clearSticky && s.isGatewayAccountProfitEligible(ctx, account) && s.isAccountInGroup(account, groupID) && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && s.isAccountSchedulableForModelSelection(ctx, account, requestedModel) && s.isAccountSchedulableForQuota(account) && s.isAccountSchedulableForWindowCost(ctx, account, true) && s.isAccountSchedulableForRPM(ctx, account, true) && !s.isStickyAccountUpstreamRestricted(ctx, groupID, account, requestedModel) { if account.Platform == nativePlatform || (account.Platform == PlatformAntigravity && account.IsMixedSchedulingEnabled()) { return account, nil } @@ -2203,6 +2271,9 @@ func (s *GatewayService) selectAccountWithMixedScheduling(ctx context.Context, g if !s.isAccountSchedulableForSelection(acc) { continue } + if !s.isGatewayAccountProfitEligible(ctx, acc) { + continue + } // require_privacy_set: 跳过 privacy 未设置的账号并标记异常 if schedGroup != nil && schedGroup.RequirePrivacySet && !acc.IsPrivacySet() { _ = s.accountRepo.SetError(ctx, acc.ID, @@ -2265,7 +2336,7 @@ func (s *GatewayService) selectAccountWithMixedScheduling(ctx context.Context, g // 4. 建立粘性绑定 if sessionHash != "" && s.cache != nil { - if err := s.cache.SetSessionAccountID(ctx, derefGroupID(groupID), sessionHash, selected.ID, stickySessionTTL); err != nil { + if err := s.bindGatewayStickySessionDuringSelection(ctx, groupID, sessionHash, selected.ID); err != nil { logger.LegacyPrintf("service.gateway", "set session account failed: session=%s account_id=%d err=%v", sessionHash, selected.ID, err) } } @@ -2281,6 +2352,8 @@ type selectionFailureStats struct { PlatformFiltered int ModelUnsupported int ModelRateLimited int + ProfitThreshold int + ProfitInvalidRate int SamplePlatformIDs []int64 SampleMappingIDs []int64 SampleRateLimitIDs []string @@ -2304,7 +2377,7 @@ func (s *GatewayService) logDetailedSelectionFailure( stats := s.collectSelectionFailureStats(ctx, accounts, requestedModel, platform, excludedIDs, allowMixedScheduling) logger.LegacyPrintf( "service.gateway", - "[SelectAccountDetailed] group_id=%v model=%s platform=%s session=%s total=%d eligible=%d excluded=%d unschedulable=%d platform_filtered=%d model_unsupported=%d model_rate_limited=%d sample_platform_filtered=%v sample_model_unsupported=%v sample_model_rate_limited=%v", + "[SelectAccountDetailed] group_id=%v model=%s platform=%s session=%s total=%d eligible=%d excluded=%d unschedulable=%d platform_filtered=%d model_unsupported=%d model_rate_limited=%d profit_threshold=%d profit_invalid_account_rate=%d sample_platform_filtered=%v sample_model_unsupported=%v sample_model_rate_limited=%v", derefGroupID(groupID), requestedModel, platform, @@ -2316,6 +2389,8 @@ func (s *GatewayService) logDetailedSelectionFailure( stats.PlatformFiltered, stats.ModelUnsupported, stats.ModelRateLimited, + stats.ProfitThreshold, + stats.ProfitInvalidRate, stats.SamplePlatformIDs, stats.SampleMappingIDs, stats.SampleRateLimitIDs, @@ -2353,6 +2428,10 @@ func (s *GatewayService) collectSelectionFailureStats( stats.ModelRateLimited++ remaining := acc.GetRateLimitRemainingTimeWithContext(ctx, requestedModel).Truncate(time.Second) stats.SampleRateLimitIDs = appendSelectionFailureRateSample(stats.SampleRateLimitIDs, acc.ID, remaining) + case openAIProfitFilterReasonThreshold: + stats.ProfitThreshold++ + case openAIProfitFilterReasonInvalidAccountRate: + stats.ProfitInvalidRate++ default: stats.Eligible++ } @@ -2397,6 +2476,9 @@ func (s *GatewayService) diagnoseSelectionFailure( Detail: fmt.Sprintf("remaining=%s", remaining), } } + if vetoed, reason := openAIProfitControlVetoReason(ctx, acc); vetoed { + return selectionFailureDiagnosis{Category: reason} + } return selectionFailureDiagnosis{Category: "eligible"} } @@ -2434,7 +2516,7 @@ func appendSelectionFailureRateSample(samples []string, accountID int64, remaini func summarizeSelectionFailureStats(stats selectionFailureStats) string { return fmt.Sprintf( - "total=%d eligible=%d excluded=%d unschedulable=%d platform_filtered=%d model_unsupported=%d model_rate_limited=%d", + "total=%d eligible=%d excluded=%d unschedulable=%d platform_filtered=%d model_unsupported=%d model_rate_limited=%d profit_threshold=%d profit_invalid_account_rate=%d", stats.Total, stats.Eligible, stats.Excluded, @@ -2442,6 +2524,8 @@ func summarizeSelectionFailureStats(stats selectionFailureStats) string { stats.PlatformFiltered, stats.ModelUnsupported, stats.ModelRateLimited, + stats.ProfitThreshold, + stats.ProfitInvalidRate, ) } diff --git a/backend/internal/service/gateway_service.go b/backend/internal/service/gateway_service.go index 2f2f836ddd..0f3415fab5 100644 --- a/backend/internal/service/gateway_service.go +++ b/backend/internal/service/gateway_service.go @@ -6,6 +6,7 @@ import ( "encoding/json" "errors" "fmt" + "io" "log/slog" "net/http" "os" @@ -20,10 +21,12 @@ import ( "github.com/Wei-Shaw/sub2api/internal/config" "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "github.com/Wei-Shaw/sub2api/internal/pkg/xai" "github.com/Wei-Shaw/sub2api/internal/util/responseheaders" "github.com/cespare/xxhash/v2" gocache "github.com/patrickmn/go-cache" "github.com/tidwall/gjson" + "github.com/tidwall/sjson" "golang.org/x/sync/singleflight" ) @@ -444,14 +447,24 @@ var allowedHeaders = map[string]bool{ "x-client-request-id": true, } +// ErrStickySessionNotFound is returned by GatewayCache.GetSessionAccountID +// when no binding exists for the session. It abstracts away the underlying +// cache implementation (e.g. redis.Nil), mirroring ErrRefreshTokenNotFound. +var ErrStickySessionNotFound = errors.New("sticky session not found") + // GatewayCache 定义网关服务的缓存操作接口。 // 提供粘性会话（Sticky Session）的存储、查询、刷新和删除功能。 // // GatewayCache defines cache operations for gateway service. // Provides sticky session storage, retrieval, refresh and deletion capabilities. type GatewayCache interface { - // GetSessionAccountID 获取粘性会话绑定的账号 ID - // Get the account ID bound to a sticky session + // GetSessionAccountID 获取粘性会话绑定的账号 ID；无绑定时返回 + // ErrStickySessionNotFound，使 service 层无需依赖具体缓存实现即可 + // 区分"未绑定"与真实读取失败。 + // Get the account ID bound to a sticky session. Returns + // ErrStickySessionNotFound when no binding exists so service code can + // distinguish a miss from a real read failure without importing the + // cache driver. GetSessionAccountID(ctx context.Context, groupID int64, sessionHash string) (int64, error) // SetSessionAccountID 设置粘性会话与账号的绑定关系 // Set the binding between sticky session and account @@ -462,6 +475,17 @@ type GatewayCache interface { // DeleteSessionAccountID 删除粘性会话绑定，用于账号不可用时主动清理 // Delete sticky session binding, used to proactively clean up when account becomes unavailable DeleteSessionAccountID(ctx context.Context, groupID int64, sessionHash string) error + + // Grok async video billing snapshot (create → status success). + // SetGrokVideoPendingBilling stores create-time model/duration/resolution for status billing. + SetGrokVideoPendingBilling(ctx context.Context, key string, payload []byte, ttl time.Duration) error + // GetGrokVideoPendingBilling returns the create-time billing snapshot; miss → nil, nil. + GetGrokVideoPendingBilling(ctx context.Context, key string) ([]byte, error) + // ClaimGrokVideoBilled atomically marks a video request as billed (SetNX). + // Returns true when this caller won the claim; false when already billed or claim unavailable. + ClaimGrokVideoBilled(ctx context.Context, key string, ttl time.Duration) (bool, error) + // ReleaseGrokVideoBilled clears a claim so a failed RecordUsage can retry billing. + ReleaseGrokVideoBilled(ctx context.Context, key string) error } // derefGroupID safely dereferences *int64 to int64, returning 0 if nil @@ -537,6 +561,15 @@ type AccountSelectionResult struct { Acquired bool ReleaseFunc func() WaitPlan *AccountWaitPlan // nil means no wait allowed + // profitGate 携带本次选号真实生效的利润门（无门为 nil）。门安装在调度栈的 + // 局部 ctx 上，handler 必须经 ContextWithSelectionProfitGate 重放后才能在 + // 调度栈之外做抢槽后终检与准入后粘性绑定。 + profitGate *openAIProfitControlGate +} + +// ProfitGateActive 报告本次选号是否处于利润门之下。 +func (r *AccountSelectionResult) ProfitGateActive() bool { + return r != nil && r.profitGate != nil } // ClaudeUsage 表示Claude API返回的usage信息 @@ -551,18 +584,27 @@ type ClaudeUsage struct { } // ForwardResult 转发结果 +type AudioUsage struct { + Mode string // realtime | tts | stt + DurationOrUnits float64 // minutes / million-chars / hours +} + type ForwardResult struct { RequestID string Usage ClaudeUsage Model string // UpstreamModel is the actual upstream model after mapping. // Prefer empty when it is identical to Model; persistence normalizes equal values away as no-op mappings. - UpstreamModel string - Stream bool - Duration time.Duration - FirstTokenMs *int // 首字时间（流式请求） - ClientDisconnect bool // 客户端是否在流式传输过程中断开 - ReasoningEffort *string + UpstreamModel string + // UpstreamResponseModel is captured from the raw successful upstream + // response before any client-facing rewrite or protocol conversion. + UpstreamResponseModel string + UpstreamResponseModelConflict bool + Stream bool + Duration time.Duration + FirstTokenMs *int // 首字时间（流式请求） + ClientDisconnect bool // 客户端是否在流式传输过程中断开 + ReasoningEffort *string // 图片生成计费字段（图片生成模型使用） ImageCount int // 生成的图片数量 @@ -572,6 +614,8 @@ type ForwardResult struct { ImageOutputSizes []string ImageSizeSource string ImageSizeBreakdown map[string]int + SearchCount int + AudioUsage *AudioUsage } // GatewayFailureStage identifies which request stage failed. The zero value is @@ -614,6 +658,7 @@ type UpstreamFailoverError struct { ResponseHeaders http.Header // 上游响应头，用于透传 cf-ray/cf-mitigated/content-type 等诊断信息 ForceCacheBilling bool // Antigravity 粘性会话切换时设为 true RetryableOnSameAccount bool // 临时性错误（如 Google 间歇性 400、空响应），应在同一账号上重试 N 次再切换 + RequestScopedTransient bool // 故障因素与账号无关（如上游按客户端身份/模型容量降载）：可同账号重试，但不得据此对账号做临时封禁 SafeToFailoverAfterWrite bool // 仅写出 SSE 注释等非语义字节时，仍可在同一客户端流中切换账号 Stage GatewayFailureStage Scope GatewayFailureScope @@ -665,6 +710,11 @@ func (s *GatewayService) TempUnscheduleRetryableError(ctx context.Context, accou if failoverErr == nil || !failoverErr.RetryableOnSameAccount { return } + // 请求级瞬时故障与账号健康无关：封禁只会把与故障无关的账号一并摘掉， + // 而故障因素（客户端身份、模型容量）在下一个账号上完全相同。 + if failoverErr.RequestScopedTransient { + return + } // 根据状态码选择封禁策略 switch failoverErr.StatusCode { case http.StatusBadRequest: @@ -873,6 +923,42 @@ func (s *GatewayService) BindStickySession(ctx context.Context, groupID *int64, return s.cache.SetSessionAccountID(ctx, derefGroupID(groupID), sessionHash, accountID, stickySessionTTL) } +// bindGatewayStickySessionDuringSelection preserves the normal eager sticky +// behavior unless a profit gate is installed. Profit-controlled requests bind +// only after the terminal post-slot check, otherwise a rejected candidate could +// overwrite a healthy pre-existing sticky binding. +func (s *GatewayService) bindGatewayStickySessionDuringSelection(ctx context.Context, groupID *int64, sessionHash string, accountID int64) error { + if gatewayProfitControlGateActive(ctx) { + return nil + } + return s.BindStickySession(ctx, groupID, sessionHash, accountID) +} + +// BindStickySessionAfterProfitAdmission records a terminally admitted +// account. Without a profit gate it preserves the pre-existing eager binding +// behavior at the handler bind points. With a gate it never replaces a +// different binding that already exists: a temporarily ineligible sticky +// account remains bound and automatically becomes eligible again if its +// account rate recovers. +func (s *GatewayService) BindStickySessionAfterProfitAdmission(ctx context.Context, groupID *int64, sessionHash string, accountID int64) error { + if sessionHash == "" || accountID <= 0 || s.cache == nil { + return nil + } + if !gatewayProfitControlGateActive(ctx) { + return s.BindStickySession(ctx, groupID, sessionHash, accountID) + } + existingAccountID, err := s.cache.GetSessionAccountID(ctx, derefGroupID(groupID), sessionHash) + if err != nil && !errors.Is(err, ErrStickySessionNotFound) { + // 读失败时无法判断既有绑定，保守跳过而不是冒着覆盖健康绑定的风险写入。 + slog.Warn("profit_control_sticky_binding_read_failed", "group_id", derefGroupID(groupID), "account_id", accountID, "error", err) + return nil + } + if existingAccountID > 0 && existingAccountID != accountID { + return nil + } + return s.BindStickySession(ctx, groupID, sessionHash, accountID) +} + // GetCachedSessionAccountID retrieves the account ID bound to a sticky session. // Returns 0 if no binding exists or on error. func (s *GatewayService) GetCachedSessionAccountID(ctx context.Context, groupID *int64, sessionHash string) (int64, error) { @@ -1160,6 +1246,15 @@ func (s *GatewayService) getOAuthToken(ctx context.Context, account *Account) (s return accessToken, "oauth", nil } + // Grok OAuth: prefer access_token from credentials (background refresher keeps it warm). + if account.Platform == PlatformGrok && account.Type == AccountTypeOAuth { + accessToken := account.GetGrokAccessToken() + if accessToken == "" { + return "", "", errors.New("grok access_token not found in credentials") + } + return accessToken, "oauth", nil + } + // 其他情况（Gemini 有自己的 TokenProvider，setup-token 类型等）直接从账号读取 accessToken := account.GetCredential("access_token") if accessToken == "" { @@ -1171,6 +1266,78 @@ func (s *GatewayService) getOAuthToken(ctx context.Context, account *Account) (s // GetAvailableModels returns the list of models available for a group // It aggregates model_mapping keys from all schedulable accounts in the group + +// DoGrokNativeResponsesJSON POSTs a non-streaming Responses body to the account's +// Grok upstream and returns the raw JSON body. Used by /v1/web_search. +// Gin-free: UA is always the pinned Grok CLI identity (resolveGrokUpstreamUserAgent ignores inbound). +func (s *GatewayService) DoGrokNativeResponsesJSON(ctx context.Context, account *Account, body []byte) ([]byte, error) { + if s == nil || s.httpUpstream == nil { + return nil, errors.New("http upstream not configured") + } + if account == nil { + return nil, errors.New("account is required") + } + if !account.IsGrok() { + return nil, errors.New("grok account required") + } + token, _, err := s.GetAccessToken(ctx, account) + if err != nil { + // Credential/token failures should try the next Grok account in the pool. + return nil, &UpstreamFailoverError{ + StatusCode: http.StatusUnauthorized, + Reason: GatewayFailureReason("grok_search_token"), + } + } + targetURL, err := buildGrokResponsesURL(account, nil, s.settingService) + if err != nil { + return nil, err + } + if json.Valid(body) { + if model := strings.TrimSpace(gjson.GetBytes(body, "model").String()); model == "" { + if patched, patchErr := sjson.SetBytes(body, "model", xai.DefaultTextModel); patchErr == nil { + body = patched + } + } + } + upstreamReq, err := http.NewRequestWithContext(ctx, http.MethodPost, targetURL, bytes.NewReader(body)) + if err != nil { + return nil, fmt.Errorf("build grok responses request: %w", err) + } + upstreamReq.Header.Set("Authorization", "Bearer "+token) + upstreamReq.Header.Set("Content-Type", "application/json") + upstreamReq.Header.Set("Accept", "application/json") + upstreamReq.Header.Set("User-Agent", defaultGrokUpstreamUserAgent()) + applyGrokCLIHeaders(upstreamReq.Header) + account.ApplyHeaderOverrides(upstreamReq.Header) + proxyURL := "" + if account.ProxyID != nil && account.Proxy != nil { + proxyURL = account.Proxy.URL() + } + resp, err := s.httpUpstream.Do(upstreamReq, proxyURL, account.ID, account.Concurrency) + if err != nil { + return nil, &UpstreamFailoverError{StatusCode: http.StatusBadGateway, Reason: GatewayFailureReason("grok_search_transport")} + } + defer func() { _ = resp.Body.Close() }() + respBytes, readErr := io.ReadAll(io.LimitReader(resp.Body, 4<<20)) + if readErr != nil { + return nil, &UpstreamFailoverError{ + StatusCode: http.StatusBadGateway, + Reason: GatewayFailureReason("grok_search_read"), + } + } + if resp.StatusCode >= 400 { + msg := string(respBytes) + if len(msg) > 200 { + msg = msg[:200] + } + if resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusPaymentRequired || resp.StatusCode == http.StatusForbidden || resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode >= 500 { + return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: respBytes} + } + return nil, fmt.Errorf("grok upstream %d: %s", resp.StatusCode, msg) + } + return respBytes, nil +} + func (s *GatewayService) GetAvailableModels(ctx context.Context, groupID *int64, platform string) []string { cacheKey := modelsListCacheKey(groupID, platform) if s.modelsListCache != nil { diff --git a/backend/internal/service/gateway_upstream_response.go b/backend/internal/service/gateway_upstream_response.go index a32c007e03..5a9fedf964 100644 --- a/backend/internal/service/gateway_upstream_response.go +++ b/backend/internal/service/gateway_upstream_response.go @@ -647,7 +647,51 @@ type streamingResult struct { clientDisconnect bool // 客户端是否在流式传输过程中断开 } +// hasObservedTokens 报告流式过程中是否已观测到任何上游计量的 token。 +func (u *ClaudeUsage) hasObservedTokens() bool { + if u == nil { + return false + } + return u.InputTokens > 0 || u.OutputTokens > 0 || + u.CacheCreationInputTokens > 0 || u.CacheReadInputTokens > 0 || + u.CacheCreation5mTokens > 0 || u.CacheCreation1hTokens > 0 || + u.ImageOutputTokens > 0 +} + +// partialStreamUsageResult 在流式转发中途出错时，把已观测到 usage 的部分结果包装为 +// ForwardResult（与错误一起返回给 handler 记录）。上游一旦下发过 message_start， +// input/cache token 就已计量，直接丢弃会让请求完全漏记漏计费（issue #5148）。 +// 无已观测 usage 时返回 nil。 +// +// 不变式：UpstreamFailoverError 必须保持 result=nil——failover 重试成功后按成功请求 +// 计费，若同时返回部分 usage 会造成双重计费，此处显式拦截兜底。 +func partialStreamUsageResult(c *gin.Context, resp *http.Response, streamResult *streamingResult, model, upstreamModel string, startTime time.Time, err error) *ForwardResult { + if streamResult == nil || !streamResult.usage.hasObservedTokens() { + return nil + } + var failoverErr *UpstreamFailoverError + if errors.As(err, &failoverErr) { + return nil + } + return &ForwardResult{ + RequestID: resp.Header.Get("x-request-id"), + Usage: *streamResult.usage, + Model: model, + UpstreamModel: upstreamModel, + UpstreamResponseModel: observedUpstreamResponseModel(c), + UpstreamResponseModelConflict: observedUpstreamResponseModelConflict(c), + Stream: true, + Duration: time.Since(startTime), + FirstTokenMs: streamResult.firstTokenMs, + ClientDisconnect: streamResult.clientDisconnect, + } +} + func (s *GatewayService) handleStreamingResponse(ctx context.Context, resp *http.Response, c *gin.Context, account *Account, startTime time.Time, originalModel, mappedModel string, mimicClaudeCode bool) (*streamingResult, error) { + observer := upstreamResponseModelObserverFromContext(c) + if observer == nil { + observer = beginUpstreamResponseModelObservation(c) + } // 更新5h窗口状态 s.rateLimitService.UpdateSessionWindow(ctx, account, resp.Header) @@ -843,6 +887,7 @@ func (s *GatewayService) handleStreamingResponse(ctx context.Context, resp *http } eventType, _ := event["type"].(string) + observer.ObserveAnthropic([]byte(dataLine)) if eventName == "" { eventName = eventType } @@ -1338,6 +1383,11 @@ func (s *GatewayService) handleNonStreamingResponse(ctx context.Context, resp *h if err != nil { return nil, err } + observer := upstreamResponseModelObserverFromContext(c) + if observer == nil { + observer = beginUpstreamResponseModelObservation(c) + } + observer.ObserveAnthropic(body) // 解析usage var response struct { diff --git a/backend/internal/service/gateway_usage_billing.go b/backend/internal/service/gateway_usage_billing.go index 51b8031fc4..a5998f3595 100644 --- a/backend/internal/service/gateway_usage_billing.go +++ b/backend/internal/service/gateway_usage_billing.go @@ -42,6 +42,7 @@ type RecordUsageInput struct { User *User Account *Account Subscription *UserSubscription // 可选：订阅信息 + PricingAt time.Time // token 售价固定时刻；零值保持既有的记录时刻语义 InboundEndpoint string // 入站端点（客户端请求路径） UpstreamEndpoint string // 上游端点（标准化后的上游路径） UserAgent string // 请求的 User-Agent @@ -203,6 +204,13 @@ func postUsageBilling(ctx context.Context, p *postUsageBillingParams, deps *bill } func resolveUsageBillingRequestID(ctx context.Context, upstreamRequestID string) string { + // Forced durable money-event IDs must win over client/local context IDs so + // standalone web_search / async video cannot collapse under a reused client id. + if requestID := strings.TrimSpace(upstreamRequestID); requestID != "" { + if isForcedUsageBillingRequestID(requestID) { + return requestID + } + } if ctx != nil { if clientRequestID, _ := ctx.Value(ctxkey.ClientRequestID).(string); strings.TrimSpace(clientRequestID) != "" { return "client:" + strings.TrimSpace(clientRequestID) @@ -217,6 +225,40 @@ func resolveUsageBillingRequestID(ctx context.Context, upstreamRequestID string) return "generated:" + generateRequestID() } +func isForcedUsageBillingRequestID(requestID string) bool { + id := strings.TrimSpace(requestID) + return strings.HasPrefix(id, "web_search:") || + strings.HasPrefix(id, "grok-video:") || + strings.HasPrefix(id, "grok_audio:") || + strings.HasPrefix(id, "grok_realtime:") +} + +// StableGrokAudioBillingRequestID is the durable usage_logs / dedup key for one +// voice HTTP call (TTS/STT). Prefer an upstream request id when present. +func StableGrokAudioBillingRequestID(upstreamRequestID string) string { + upstreamRequestID = strings.TrimSpace(upstreamRequestID) + if strings.HasPrefix(upstreamRequestID, "grok_audio:") { + return upstreamRequestID + } + if upstreamRequestID == "" { + upstreamRequestID = generateRequestID() + } + return "grok_audio:" + upstreamRequestID +} + +// StableGrokRealtimeBillingRequestID is the durable usage_logs / dedup key for +// one realtime WebSocket session. +func StableGrokRealtimeBillingRequestID(sessionID string) string { + sessionID = strings.TrimSpace(sessionID) + if strings.HasPrefix(sessionID, "grok_realtime:") { + return sessionID + } + if sessionID == "" { + sessionID = generateRequestID() + } + return "grok_realtime:" + sessionID +} + func resolveUsageBillingPayloadFingerprint(ctx context.Context, requestPayloadHash string) string { if payloadHash := strings.TrimSpace(requestPayloadHash); payloadHash != "" { return payloadHash @@ -569,6 +611,7 @@ func (s *GatewayService) RecordUsage(ctx context.Context, input *RecordUsageInpu User: input.User, Account: input.Account, Subscription: input.Subscription, + PricingAt: input.PricingAt, InboundEndpoint: input.InboundEndpoint, UpstreamEndpoint: input.UpstreamEndpoint, UserAgent: input.UserAgent, @@ -589,6 +632,7 @@ type RecordUsageLongContextInput struct { User *User Account *Account Subscription *UserSubscription // 可选：订阅信息 + PricingAt time.Time // token 售价固定时刻；零值保持既有的记录时刻语义 InboundEndpoint string // 入站端点（客户端请求路径） UpstreamEndpoint string // 上游端点（标准化后的上游路径） UserAgent string // 请求的 User-Agent @@ -612,6 +656,7 @@ func (s *GatewayService) RecordUsageWithLongContext(ctx context.Context, input * User: input.User, Account: input.Account, Subscription: input.Subscription, + PricingAt: input.PricingAt, InboundEndpoint: input.InboundEndpoint, UpstreamEndpoint: input.UpstreamEndpoint, UserAgent: input.UserAgent, @@ -635,6 +680,7 @@ type recordUsageCoreInput struct { User *User Account *Account Subscription *UserSubscription + PricingAt time.Time InboundEndpoint string UpstreamEndpoint string UserAgent string @@ -647,6 +693,80 @@ type recordUsageCoreInput struct { ChannelUsageFields } +// responseModelBillingCostEpsilon 吸收两次成本计算之间的浮点末位误差， +// 避免同价模型因浮点误差被判成"更贵"而白白放弃采纳。 +const responseModelBillingCostEpsilon = 1e-12 + +// responseModelBillingDeclaration 返回可用于计费的上游响应模型；返回空字符串表示 +// 必须沿用基线计费模型。两条计费主干（Anthropic 系 / OpenAI 系）共用本准入判断。 +// +// 渠道把 billing_model_source 设为 response_model，等于把"按哪个模型计价"的一部分 +// 决定权交给上游，因此准入条件必须收紧： +// - 只在渠道显式开启该模式时生效，其余模式一律不看响应模型； +// - 一次请求内出现过互相冲突的模型声明时不采纳（无法确定上游究竟服务了哪个模型）； +// - 图片 / 视频 / 网页搜索 / 语音 / 搜索附加费这类按次按量计费的请求不采纳：它们按张、 +// 按秒、按次定价，与本模式的 token 定价准入检查不是同一套价格表，混用会让一个只验过 +// token 价的模型名去决定媒体单价。新增按次计费形态时必须同步扩这个入参。 +// +// 调用方还必须额外满足两条：模型能被价格表确定性识别（见 +// hasIdentifiedResponseModelPricing / hasIdentifiedOpenAIResponsePricing），以及通过 +// responseModelBillingAdoptable 的成本准入。 +func responseModelBillingDeclaration(source, responseModel string, conflict, mediaBilled bool) string { + if source != BillingModelSourceResponse || conflict || mediaBilled { + return "" + } + return strings.TrimSpace(responseModel) +} + +// responseModelBillingAdoptable 判定按响应模型重算出的成本能否取代基线成本。 +// 三条不变式，任一不满足都必须沿用基线（即开启本模式前的既有行为）： +// +// 1. 不得更贵——上游声明永远不能抬高用户费用；epsilon 吸收两次计算之间的浮点末位误差。 +// 2. 不得把一笔本应计费的请求归零。价格表里存在把 token 价显式写成 0 的条目 +// （TokenPricingAbsent 只在 input/output 价**都缺失**时才为真，显式 0 算"有价"因而 +// 能通过确定性识别那道门），放任归零等于让上游自报一个免费模型名就能白嫖。 +// 基线本身就是 0 时不受影响，采纳与否都不改变金额。 +// 3. 不得把计费从管理员显式配置的渠道定价切到全局价格表。渠道定价查表只做精确键与 +// 前缀通配、**不剥日期后缀**，而全局价格表的确定性识别**会剥** 8 位日期后缀；上游 +// 普遍自报带日期的模型 ID（如 claude-opus-4-5-20251101），若允许跨源比较，渠道加价 +// 会被这类自报名字静默绕过。管理员若确实想让降级目标享受折扣，为它显式配一条渠道 +// 定价即可——那是一次可审计的显式授权。 +func responseModelBillingAdoptable(baseline, response *CostBreakdown, baselineChannelPriced, responseChannelPriced bool) bool { + if baseline == nil || response == nil { + return false + } + if response.TotalCost > baseline.TotalCost+responseModelBillingCostEpsilon { + return false + } + if response.TotalCost <= 0 && baseline.TotalCost > 0 { + return false + } + return !baselineChannelPriced || responseChannelPriced +} + +// logResponseModelBillingApplied 记录一次实际生效的响应模型计费切换。 +// 本模式下的少收由上游声明驱动，必须留下可审计痕迹；计费基准未变时不记录，避免刷屏。 +func logResponseModelBillingApplied(component string, account *Account, requestID, baselineModel, responseModel string, baselineCost, responseCost *CostBreakdown) { + baselineModel = strings.TrimSpace(baselineModel) + responseModel = strings.TrimSpace(responseModel) + if strings.EqualFold(baselineModel, responseModel) { + return + } + attrs := []any{ + "component", component, + "request_id", strings.TrimSpace(requestID), + "baseline_model", baselineModel, + "response_model", responseModel, + } + if baselineCost != nil && responseCost != nil { + attrs = append(attrs, "baseline_cost", baselineCost.TotalCost, "billed_cost", responseCost.TotalCost) + } + if account != nil { + attrs = append(attrs, "platform", account.Platform, "account_id", account.ID) + } + slog.Info("billing.response_model_applied", attrs...) +} + // recordUsageCore 是 RecordUsage 和 RecordUsageWithLongContext 的统一实现。 // LongContextThreshold > 0 时 Token 计费回退走 CalculateCostWithLongContext。 func (s *GatewayService) recordUsageCore(ctx context.Context, input *recordUsageCoreInput, opts *recordUsageOpts) error { @@ -685,7 +805,11 @@ func (s *GatewayService) recordUsageCore(ctx context.Context, input *recordUsage } // token 倍率叠加高峰因子（token 计费含图片 token，图片按次倍率不受影响）。高峰因子按请求时刻现算， // 不并入上面的 getUserGroupRateMultiplier，以免污染 user:group 倍率缓存。 - multiplier, imageMultiplier := computePeakAwareMultipliers(apiKey, multiplier, timezone.Now()) + pricingAt := input.PricingAt + if pricingAt.IsZero() { + pricingAt = timezone.Now() + } + multiplier, imageMultiplier := computePeakAwareMultipliers(apiKey, multiplier, pricingAt) // 确定计费模型 concreteBillingModel := forwardResultBillingModel(result.Model, result.UpstreamModel) @@ -715,6 +839,27 @@ func (s *GatewayService) recordUsageCore(ctx context.Context, input *recordUsage // 计算费用 cost := s.calculateRecordUsageCost(ctx, result, apiKey, billingModel, multiplier, imageMultiplier, opts) + // response_model：按上游成功响应自报的模型计费（渠道显式开启才生效）。 + // 采纳条件见 responseModelBillingDeclaration + hasIdentifiedResponseModelPricing + // + responseModelBillingAdoptable。任一条件不满足都静默回落基线，即开启本模式前的 + // 既有行为。响应模型与基线同名时直接跳过：重算必然同价，白跑一次定价解析。 + if responseModel := responseModelBillingDeclaration( + input.BillingModelSource, + result.UpstreamResponseModel, + result.UpstreamResponseModelConflict, + result.ImageCount > 0 || result.AudioUsage != nil || result.SearchCount > 0, + ); responseModel != "" && !strings.EqualFold(responseModel, strings.TrimSpace(billingModel)) { + if identified, responseChannelPriced := s.hasIdentifiedResponseModelPricing(ctx, responseModel, apiKey); identified { + responseCost := s.calculateRecordUsageCost(ctx, result, apiKey, responseModel, multiplier, imageMultiplier, opts) + baselineChannelPriced := s.resolveChannelPricing(ctx, billingModel, apiKey) != nil + if responseModelBillingAdoptable(cost, responseCost, baselineChannelPriced, responseChannelPriced) { + // billingModel 到此为止只是定价查表的入参，后续流程只消费 cost， + // 因此这里不改写它，改由日志记录实际生效的计费基准。 + logResponseModelBillingApplied("service.gateway", account, result.RequestID, billingModel, responseModel, cost, responseCost) + cost = responseCost + } + } + } // 判断计费方式：订阅模式 vs 余额模式 isSubscriptionBilling := subscription != nil && apiKey.Group != nil && apiKey.Group.IsSubscriptionType() @@ -777,6 +922,8 @@ func (s *GatewayService) recordUsageCore(ctx context.Context, input *recordUsage }, s.billingDeps(), s.usageBillingRepo) if billingErr != nil { + usageLog.ActualCost = 0 + writeUsageLogBestEffort(ctx, s.usageLogRepo, usageLog, "service.gateway") return billingErr } writeUsageLogBestEffort(ctx, s.usageLogRepo, usageLog, "service.gateway") @@ -802,8 +949,29 @@ func (s *GatewayService) calculateRecordUsageCost( return s.calculateImageCost(ctx, result, apiKey, billingModel, imageMultiplier) } - // Token 计费 - return s.calculateTokenCost(ctx, result, apiKey, billingModel, multiplier, opts) + // Voice audio (TTS / STT / realtime) when present on the forward result. + if result.AudioUsage != nil { + cfg := groupAudioPriceConfigFromAPIKey(apiKey) + return s.billingService.CalculateAudioCost(result.AudioUsage.Mode, result.AudioUsage.DurationOrUnits, cfg, multiplier) + } + + // Token 计费；SearchCount 为叠加 surcharge（不替代 token）。 + tokenCost := s.calculateTokenCost(ctx, result, apiKey, billingModel, multiplier, opts) + if result.SearchCount > 0 { + price := groupSearchPricePer1kFromAPIKey(apiKey) + if price != nil && *price == 0 { + logger.LegacyPrintf("service.gateway", "[Billing] search_price_per_1k explicit 0; search free group_model=%s count=%d", billingModel, result.SearchCount) + } + searchCost := s.billingService.CalculateSearchCost(result.SearchCount, price, multiplier) + if searchCost != nil && (searchCost.TotalCost > 0 || searchCost.ActualCost > 0) { + if tokenCost == nil { + return searchCost + } + tokenCost.TotalCost += searchCost.TotalCost + tokenCost.ActualCost += searchCost.ActualCost + } + } + return tokenCost } // compositeBillableModel 决定 composite 分组请求的计费模型：来源覆盖把计费模型 @@ -856,6 +1024,22 @@ func (s *GatewayService) hasResolvableTokenPricing(ctx context.Context, model st return err == nil } +// hasIdentifiedResponseModelPricing 判断上游自报的响应模型是否可以作为计费基准， +// 并回传它是否解析到了渠道级定价（供 responseModelBillingAdoptable 的跨定价源守卫使用， +// 避免为此再解析一次）。 +// 与 hasResolvableTokenPricing 的区别是刻意更严：只接受管理员为该模型显式配置的 +// 渠道定价，或价格表中能被确定性识别的条目；不接受按子串猜出来的系列兜底价。 +// 详见 responseModelBillingDeclaration 的说明。 +func (s *GatewayService) hasIdentifiedResponseModelPricing(ctx context.Context, model string, apiKey *APIKey) (identified bool, channelPriced bool) { + if strings.TrimSpace(model) == "" { + return false, false + } + if s.resolveChannelPricing(ctx, model, apiKey) != nil { + return true, true + } + return s.billingService.HasIdentifiedTokenPricing(model), false +} + // resolveChannelPricing 检查指定模型是否存在渠道级别定价。 // 返回非 nil 的 ResolvedPricing 表示有渠道定价，nil 表示走默认定价路径。 func (s *GatewayService) resolveChannelPricing(ctx context.Context, billingModel string, apiKey *APIKey) *ResolvedPricing { @@ -979,6 +1163,16 @@ func (s *GatewayService) buildRecordUsageLog( ) *UsageLog { durationMs := int(result.Duration.Milliseconds()) requestID := resolveUsageBillingRequestID(ctx, result.RequestID) + sentModel := upstreamSentModel(result.Model, result.UpstreamModel) + if result.UpstreamResponseModelConflict { + slog.Warn("upstream_response_model_conflict", + "platform", account.Platform, + "account_id", account.ID, + "request_id", requestID, + "sent_model", sentModel, + "selected_response_model", strings.TrimSpace(result.UpstreamResponseModel), + ) + } usageLog := &UsageLog{ UserID: user.ID, APIKeyID: apiKey.ID, @@ -987,6 +1181,8 @@ func (s *GatewayService) buildRecordUsageLog( Model: result.Model, RequestedModel: requestedModel, UpstreamModel: optionalTrimmedStringPtr(result.UpstreamModel), + UpstreamResponseModel: optionalTrimmedStringPtr(result.UpstreamResponseModel), + UpstreamModelMismatch: upstreamModelMismatch(sentModel, result.UpstreamResponseModel), ReasoningEffort: result.ReasoningEffort, InboundEndpoint: optionalTrimmedStringPtr(input.InboundEndpoint), UpstreamEndpoint: optionalTrimmedStringPtr(input.UpstreamEndpoint), diff --git a/backend/internal/service/gateway_usage_billing_request_id_test.go b/backend/internal/service/gateway_usage_billing_request_id_test.go new file mode 100644 index 0000000000..78a64cb338 --- /dev/null +++ b/backend/internal/service/gateway_usage_billing_request_id_test.go @@ -0,0 +1,59 @@ +//go:build unit + +package service + +import ( + "context" + "strings" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey" + "github.com/stretchr/testify/require" +) + +func TestResolveUsageBillingRequestID_ForcedWebSearchBeatsClientID(t *testing.T) { + t.Parallel() + ctx := context.WithValue(context.Background(), ctxkey.ClientRequestID, "client-shared-id") + got := resolveUsageBillingRequestID(ctx, "web_search:uuid-1") + require.Equal(t, "web_search:uuid-1", got) +} + +func TestResolveUsageBillingRequestID_ClientWinsOverPlainUpstream(t *testing.T) { + t.Parallel() + ctx := context.WithValue(context.Background(), ctxkey.ClientRequestID, "client-shared-id") + got := resolveUsageBillingRequestID(ctx, "resp_abc") + require.Equal(t, "client:client-shared-id", got) +} + +func TestIsForcedUsageBillingRequestID(t *testing.T) { + t.Parallel() + require.True(t, isForcedUsageBillingRequestID("web_search:x")) + require.True(t, isForcedUsageBillingRequestID("grok-video:task-1")) + require.True(t, isForcedUsageBillingRequestID("grok_audio:up-1")) + require.True(t, isForcedUsageBillingRequestID("grok_realtime:sess-1")) + require.False(t, isForcedUsageBillingRequestID("resp_abc")) +} + +func TestStableGrokAudioBillingRequestID(t *testing.T) { + t.Parallel() + require.Equal(t, "grok_audio:up-1", StableGrokAudioBillingRequestID("up-1")) + require.Equal(t, "grok_audio:up-1", StableGrokAudioBillingRequestID("grok_audio:up-1")) + got := StableGrokAudioBillingRequestID("") + require.True(t, strings.HasPrefix(got, "grok_audio:")) + require.Greater(t, len(got), len("grok_audio:")) +} + +func TestStableGrokRealtimeBillingRequestID(t *testing.T) { + t.Parallel() + require.Equal(t, "grok_realtime:s1", StableGrokRealtimeBillingRequestID("s1")) + require.Equal(t, "grok_realtime:s1", StableGrokRealtimeBillingRequestID("grok_realtime:s1")) + got := StableGrokRealtimeBillingRequestID("") + require.True(t, strings.HasPrefix(got, "grok_realtime:")) +} + +func TestResolveUsageBillingRequestID_ForcedGrokAudioBeatsClientID(t *testing.T) { + t.Parallel() + ctx := context.WithValue(context.Background(), ctxkey.ClientRequestID, "client-shared-id") + got := resolveUsageBillingRequestID(ctx, StableGrokAudioBillingRequestID("up-9")) + require.Equal(t, "grok_audio:up-9", got) +} diff --git a/backend/internal/service/gemini_chat_completions_compat_service.go b/backend/internal/service/gemini_chat_completions_compat_service.go index d766dc9ab7..e76fbcee15 100644 --- a/backend/internal/service/gemini_chat_completions_compat_service.go +++ b/backend/internal/service/gemini_chat_completions_compat_service.go @@ -215,7 +215,9 @@ func (s *GeminiMessagesCompatService) forwardClaudeBodyAsChatCompletions( if s.rateLimitService != nil { policy = s.rateLimitService.CheckErrorPolicy(ctx, account, resp.StatusCode, respBody, mappedModel) } - if policy != ErrorPolicyTempUnscheduled { + // 与 messages 兼容层一致：只有 None / Matched 才走账号状态处理。 + // Skipped（池模式、或自定义错误码未命中）与 TempUnscheduled 已由策略层裁决完毕。 + if policy == ErrorPolicyNone || policy == ErrorPolicyMatched { s.handleGeminiUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody) } evBody := unwrapIfNeeded(account.Type == AccountTypeOAuth, respBody) @@ -322,9 +324,9 @@ func (s *GeminiMessagesCompatService) buildGeminiChatCompletionsUpstreamRequestF if clientStream { action = "streamGenerateContent" } - fullURL := fmt.Sprintf("%s/v1beta/models/%s:%s", strings.TrimRight(normalizedBaseURL, "/"), mappedModel, action) - if clientStream { - fullURL += "?alt=sse" + fullURL, err := buildGeminiAIStudioModelActionURL(normalizedBaseURL, mappedModel, action, clientStream) + if err != nil { + return nil, "", err } restGeminiReq := normalizeGeminiRequestForAIStudio(geminiReq) @@ -389,9 +391,9 @@ func (s *GeminiMessagesCompatService) buildGeminiChatCompletionsUpstreamRequestF return nil, "", err } - fullURL := fmt.Sprintf("%s/v1beta/models/%s:%s", strings.TrimRight(normalizedBaseURL, "/"), mappedModel, action) - if useUpstreamStream { - fullURL += "?alt=sse" + fullURL, err := buildGeminiAIStudioModelActionURL(normalizedBaseURL, mappedModel, action, useUpstreamStream) + if err != nil { + return nil, "", err } restGeminiReq := normalizeGeminiRequestForAIStudio(geminiReq) diff --git a/backend/internal/service/gemini_error_policy_test.go b/backend/internal/service/gemini_error_policy_test.go index e71bd141b5..46024077a5 100644 --- a/backend/internal/service/gemini_error_policy_test.go +++ b/backend/internal/service/gemini_error_policy_test.go @@ -489,6 +489,105 @@ func TestHandleGeminiUpstreamError_GoogleOneCapacityExhaustedUsesTierCooldown(t require.True(t, repo.lastRateLimitReset.Before(after.Add(5*time.Minute).Add(2*time.Second))) } +// --------------------------------------------------------------------------- +// TestHandleGeminiUpstreamError_PoolMode429 — 池模式账号的 429 不写账号级限流。 +// +// 429 的标记点在重试循环内（handleClaudeCompat / forwardNativeGemini / +// chat completions 三条路径），先于 CheckErrorPolicy 执行，池模式豁免只能落在 +// handleGeminiUpstreamError 自身；否则一次上游 429 会把账号锁到 PST 午夜， +// 即便重试已经成功返回客户端。 +// --------------------------------------------------------------------------- + +func TestHandleGeminiUpstreamError_PoolMode429(t *testing.T) { + // 中转上游的真实 429 文案：不含 "per day"，也没有 quotaResetDelay， + // 解析失败后 apikey 账号会落到 PST 午夜兜底。 + body := []byte(`{"error":{"code":429,"message":"You have exhausted your capacity on this model. Your quota will reset after 6h53m10s."}}`) + + tests := []struct { + name string + account *Account + expectRateLimited bool + }{ + { + name: "pool_mode_apikey_stays_in_pool", + account: &Account{ + ID: 600, + Platform: PlatformGemini, + Type: AccountTypeAPIKey, + Credentials: map[string]any{"pool_mode": true}, + }, + expectRateLimited: false, + }, + { + name: "custom_error_codes_hit_overrides_pool_mode", + account: &Account{ + ID: 601, + Platform: PlatformGemini, + Type: AccountTypeAPIKey, + Credentials: map[string]any{ + "pool_mode": true, + "custom_error_codes_enabled": true, + "custom_error_codes": []any{float64(429)}, + }, + }, + expectRateLimited: true, + }, + { + name: "custom_error_codes_miss_skips", + account: &Account{ + ID: 602, + Platform: PlatformGemini, + Type: AccountTypeAPIKey, + Credentials: map[string]any{ + "pool_mode": true, + "custom_error_codes_enabled": true, + "custom_error_codes": []any{float64(500)}, + }, + }, + expectRateLimited: false, + }, + { + name: "non_pool_apikey_still_rate_limited", + account: &Account{ + ID: 603, + Platform: PlatformGemini, + Type: AccountTypeAPIKey, + }, + expectRateLimited: true, + }, + { + name: "oauth_account_ignores_pool_mode_flag", + account: &Account{ + ID: 604, + Platform: PlatformGemini, + Type: AccountTypeOAuth, + Credentials: map[string]any{"pool_mode": true}, + }, + expectRateLimited: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + repo := &rateLimit429AccountRepoStub{} + svc := &GeminiMessagesCompatService{ + accountRepo: repo, + rateLimitService: NewRateLimitService(repo, nil, &config.Config{}, nil, nil), + } + + svc.handleGeminiUpstreamError(context.Background(), tt.account, http.StatusTooManyRequests, http.Header{}, body) + + if !tt.expectRateLimited { + require.Zero(t, repo.rateLimitCalls, "池模式账号不应被标记账号级限流") + return + } + require.Equal(t, 1, repo.rateLimitCalls) + require.Equal(t, tt.account.ID, repo.lastRateLimitID) + require.True(t, repo.lastRateLimitReset.After(time.Now())) + }) + } +} + type geminiErrorPolicyRepo struct { mockAccountRepoForGemini setErrorCalls int diff --git a/backend/internal/service/gemini_image_output_accounting.go b/backend/internal/service/gemini_image_output_accounting.go new file mode 100644 index 0000000000..948a2d4aee --- /dev/null +++ b/backend/internal/service/gemini_image_output_accounting.go @@ -0,0 +1,132 @@ +package service + +import ( + "strings" + + "github.com/gin-gonic/gin" + "github.com/tidwall/gjson" +) + +// geminiImageOutputCounterKey 是请求级内联图片计数器挂在 gin.Context 上的键。 +const geminiImageOutputCounterKey = "gemini_image_output_counter" + +// geminiImageOutputCounter 记录一次转发里 Gemini 上游真正回吐的内联图片数量。 +// +// 取「单个 payload 内的最大值」而不是累加：Gemini 兼容上游的 SSE 分片可能是 +// 累积式的（同一份内容在后续 chunk 里重复整段回来——本文件同目录的 +// computeGeminiTextDelta 就是为此存在的），逐 chunk 累加会把同一张图算很多次。 +// 计费上宁可少算不可多算，所以用 max 兜底： +// - 非流式：整份响应体只观测一次，max 即真实张数； +// - 累积式流：最后一个 chunk 含全部图片，max 仍是真实张数； +// - 增量式流且多图分散在不同 chunk：会低估到 1，与改动前的模型名启发式同值， +// 不构成回退。 +type geminiImageOutputCounter struct { + count int +} + +// beginGeminiImageOutputObservation 在每次 Forward 开头重置计数器。 +// failover 会拿同一个 gin.Context 重跑转发，不重置就会把上一个账号的图数带进来。 +func beginGeminiImageOutputObservation(c *gin.Context) *geminiImageOutputCounter { + if c == nil { + return nil + } + counter := &geminiImageOutputCounter{} + c.Set(geminiImageOutputCounterKey, counter) + return counter +} + +func geminiImageOutputCounterFromContext(c *gin.Context) *geminiImageOutputCounter { + if c == nil { + return nil + } + value, ok := c.Get(geminiImageOutputCounterKey) + if !ok { + return nil + } + counter, _ := value.(*geminiImageOutputCounter) + return counter +} + +// observeGeminiImageOutputs 观测一段上游响应（整份或单个 chunk）里的内联图片。 +// 调用点与 upstreamResponseModelObserver.ObserveGemini 一一对应——那里拿得到 +// 解包后的上游响应体，这里需要的是同一份字节。 +func observeGeminiImageOutputs(c *gin.Context, payload []byte) { + counter := geminiImageOutputCounterFromContext(c) + if counter == nil { + return + } + if count := countGeminiInlineImageOutputs(payload); count > counter.count { + counter.count = count + } +} + +func observedGeminiImageOutputs(c *gin.Context) int { + counter := geminiImageOutputCounterFromContext(c) + if counter == nil { + return 0 + } + return counter.count +} + +// resolveGeminiImageCount 决定本次请求按几张图计费。 +// +// 优先用上游真正返回的内联图片数：走 GeminiMessagesCompatService 的账号多是 +// API Key + 自定义模型映射，客户端请求名和上游模型名都可能是站长自取的别名 +// （issue #5358 里的 nana-banana-2），isImageGenerationModel 的白名单必然判不出， +// 于是 ImageCount=0，calculateRecordUsageCost 整条按次计费分支不触发， +// 生图请求全部记 $0。 +// +// 只有响应里数不出图时（例如上游用 fileData 引用而非 inlineData 回图，或聚合 +// 函数丢掉了图片 part）才退回既有的模型名启发式，保证老行为不回退；这里额外 +// 也认映射后的上游模型名，与 shouldSkipCodexPlanGatedImageModelCooldown 对 +// requestedModel / modelKey 双取的口径一致。 +func resolveGeminiImageCount(c *gin.Context, originalModel, mappedModel string) int { + if observed := observedGeminiImageOutputs(c); observed > 0 { + return observed + } + if isImageGenerationModel(originalModel) || isImageGenerationModel(mappedModel) { + return 1 + } + return 0 +} + +// countGeminiInlineImageOutputs 统计一段 Gemini 响应 JSON 里的内联图片 part。 +// Gemini REST 回 camelCase 的 inlineData，官方 SDK 与部分中转会回 snake_case +// 的 inline_data，两种都要认。 +func countGeminiInlineImageOutputs(payload []byte) int { + if len(payload) == 0 || !gjson.ValidBytes(payload) { + return 0 + } + count := 0 + gjson.GetBytes(payload, "candidates").ForEach(func(_, candidate gjson.Result) bool { + candidate.Get("content.parts").ForEach(func(_, part gjson.Result) bool { + if geminiPartIsInlineImage(part) { + count++ + } + return true + }) + return true + }) + return count +} + +func geminiPartIsInlineImage(part gjson.Result) bool { + inline := part.Get("inlineData") + if !inline.Exists() { + inline = part.Get("inline_data") + } + if !inline.Exists() { + return false + } + + mimeType := inline.Get("mimeType") + if !mimeType.Exists() { + mimeType = inline.Get("mime_type") + } + if !isGeminiInlineImageMIMEType(strings.ToLower(strings.TrimSpace(mimeType.String()))) { + return false + } + + // 只认真的带上了 base64 数据的 part，空壳 part 不计费。 + return strings.TrimSpace(inline.Get("data").String()) != "" +} diff --git a/backend/internal/service/gemini_image_output_accounting_test.go b/backend/internal/service/gemini_image_output_accounting_test.go new file mode 100644 index 0000000000..eb6bb2bc69 --- /dev/null +++ b/backend/internal/service/gemini_image_output_accounting_test.go @@ -0,0 +1,204 @@ +//go:build unit + +package service + +import ( + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +const geminiTestPNG = "iVBORw0KGgoAAAANSUhEUg==" + +func newGeminiImageTestContext(t *testing.T) *gin.Context { + t.Helper() + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodPost, + "/v1beta/models/nana-banana-2:generateContent", strings.NewReader("{}")) + return c +} + +func geminiImageResponse(parts string) string { + return `{"candidates":[{"content":{"role":"model","parts":[` + parts + `]},"finishReason":"STOP"}]}` +} + +func TestCountGeminiInlineImageOutputs(t *testing.T) { + cases := []struct { + name string + payload string + want int + }{ + { + name: "camelCase inlineData", + payload: geminiImageResponse(`{"inlineData":{"mimeType":"image/png","data":"` + geminiTestPNG + `"}}`), + want: 1, + }, + { + // 官方 SDK 与部分中转把字段回成 snake_case。 + name: "snake_case inline_data", + payload: geminiImageResponse(`{"inline_data":{"mime_type":"image/png","data":"` + geminiTestPNG + `"}}`), + want: 1, + }, + { + name: "text and image mixed", + payload: geminiImageResponse(`{"text":"here you go"},` + + `{"inlineData":{"mimeType":"image/jpeg","data":"` + geminiTestPNG + `"}}`), + want: 1, + }, + { + name: "multiple images", + payload: geminiImageResponse(`{"inlineData":{"mimeType":"image/png","data":"` + geminiTestPNG + `"}},` + + `{"inlineData":{"mimeType":"image/webp","data":"` + geminiTestPNG + `"}}`), + want: 2, + }, + { + name: "uppercase mime type", + payload: geminiImageResponse(`{"inlineData":{"mimeType":"IMAGE/PNG","data":"` + geminiTestPNG + `"}}`), + want: 1, + }, + { + name: "text only", + payload: geminiImageResponse(`{"text":"no image here"}`), + want: 0, + }, + { + // 非图片的内联附件（例如音频）不能按图片计费。 + name: "non image mime type", + payload: geminiImageResponse(`{"inlineData":{"mimeType":"audio/mpeg","data":"` + geminiTestPNG + `"}}`), + want: 0, + }, + { + name: "empty data is not billable", + payload: geminiImageResponse(`{"inlineData":{"mimeType":"image/png","data":""}}`), + want: 0, + }, + {name: "empty payload", payload: "", want: 0}, + {name: "invalid json", payload: "not-json", want: 0}, + {name: "error response", payload: `{"error":{"code":429,"message":"quota"}}`, want: 0}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + require.Equal(t, tc.want, countGeminiInlineImageOutputs([]byte(tc.payload))) + }) + } +} + +// 累积式 SSE 会把同一张图在后续 chunk 里整段重发，逐 chunk 累加会重复计费。 +// 计数器取单个 payload 内的最大值，正是为了挡住这一点。 +func TestObserveGeminiImageOutputs_CumulativeChunksDoNotDoubleCount(t *testing.T) { + c := newGeminiImageTestContext(t) + beginGeminiImageOutputObservation(c) + + oneImage := geminiImageResponse(`{"inlineData":{"mimeType":"image/png","data":"` + geminiTestPNG + `"}}`) + for range 4 { + observeGeminiImageOutputs(c, []byte(oneImage)) + } + + require.Equal(t, 1, observedGeminiImageOutputs(c)) +} + +func TestObserveGeminiImageOutputs_KeepsLargestChunk(t *testing.T) { + c := newGeminiImageTestContext(t) + beginGeminiImageOutputObservation(c) + + observeGeminiImageOutputs(c, []byte(geminiImageResponse(`{"text":"working"}`))) + observeGeminiImageOutputs(c, []byte(geminiImageResponse( + `{"inlineData":{"mimeType":"image/png","data":"`+geminiTestPNG+`"}},`+ + `{"inlineData":{"mimeType":"image/png","data":"`+geminiTestPNG+`"}}`))) + // 收尾 chunk 只带 usageMetadata，不能把已数到的张数抹掉。 + observeGeminiImageOutputs(c, []byte(`{"usageMetadata":{"promptTokenCount":9}}`)) + + require.Equal(t, 2, observedGeminiImageOutputs(c)) +} + +// failover 会拿同一个 gin.Context 重跑 Forward，计数器必须按次重置， +// 否则失败账号已经回吐的图会被叠加到成功账号的账单上。 +func TestBeginGeminiImageOutputObservation_ResetsPerForward(t *testing.T) { + c := newGeminiImageTestContext(t) + oneImage := []byte(geminiImageResponse(`{"inlineData":{"mimeType":"image/png","data":"` + geminiTestPNG + `"}}`)) + + beginGeminiImageOutputObservation(c) + observeGeminiImageOutputs(c, oneImage) + require.Equal(t, 1, observedGeminiImageOutputs(c)) + + beginGeminiImageOutputObservation(c) + require.Equal(t, 0, observedGeminiImageOutputs(c)) + observeGeminiImageOutputs(c, oneImage) + require.Equal(t, 1, observedGeminiImageOutputs(c)) +} + +// issue #5358：自定义模型名（客户端名与上游映射名都不在白名单里）走 Gemini 原生 +// generateContent 生图，改动前 ImageCount 恒为 0，calculateRecordUsageCost 的按次 +// 计费分支整条不触发，四次生图全部记 $0。 +func TestResolveGeminiImageCount(t *testing.T) { + oneImage := []byte(geminiImageResponse(`{"inlineData":{"mimeType":"image/png","data":"` + geminiTestPNG + `"}}`)) + textOnly := []byte(geminiImageResponse(`{"text":"hello"}`)) + + t.Run("custom model name bills by observed images", func(t *testing.T) { + c := newGeminiImageTestContext(t) + beginGeminiImageOutputObservation(c) + observeGeminiImageOutputs(c, oneImage) + + require.False(t, isImageGenerationModel("nana-banana-2"), "前置条件：白名单判不出自定义名") + require.Equal(t, 1, resolveGeminiImageCount(c, "nana-banana-2", "nana-banana-2")) + }) + + t.Run("falls back to requested model name", func(t *testing.T) { + c := newGeminiImageTestContext(t) + beginGeminiImageOutputObservation(c) + observeGeminiImageOutputs(c, textOnly) + + require.Equal(t, 1, resolveGeminiImageCount(c, "gemini-3-pro-image-preview", "gemini-3-pro-image-preview")) + }) + + t.Run("falls back to mapped upstream model name", func(t *testing.T) { + c := newGeminiImageTestContext(t) + beginGeminiImageOutputObservation(c) + observeGeminiImageOutputs(c, textOnly) + + require.Equal(t, 1, resolveGeminiImageCount(c, "my-image-alias", "gemini-2.5-flash-image")) + }) + + t.Run("text model stays unbilled", func(t *testing.T) { + c := newGeminiImageTestContext(t) + beginGeminiImageOutputObservation(c) + observeGeminiImageOutputs(c, textOnly) + + require.Equal(t, 0, resolveGeminiImageCount(c, "gemini-2.5-pro", "gemini-2.5-pro")) + }) + + t.Run("no counter on context degrades to name heuristic", func(t *testing.T) { + c := newGeminiImageTestContext(t) + require.Equal(t, 0, resolveGeminiImageCount(c, "nana-banana-2", "nana-banana-2")) + require.Equal(t, 1, resolveGeminiImageCount(c, "gemini-3-pro-image", "gemini-3-pro-image")) + }) +} + +// 端到端守住接线：/v1beta/models/{model}:generateContent 的非流式响应体 +// 必须真的喂进计数器，否则上面的单测全绿而线上依然记 $0。 +func TestHandleNativeNonStreamingResponse_FeedsImageCounter(t *testing.T) { + c := newGeminiImageTestContext(t) + beginGeminiImageOutputObservation(c) + + body := geminiImageResponse(`{"inlineData":{"mimeType":"image/png","data":"` + geminiTestPNG + `"}}`) + resp := &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(body)), + } + + svc := &GeminiMessagesCompatService{} + usage, err := svc.handleNativeNonStreamingResponse(c, resp, false) + require.NoError(t, err) + require.NotNil(t, usage) + + require.Equal(t, 1, observedGeminiImageOutputs(c)) + require.Equal(t, 1, resolveGeminiImageCount(c, "nana-banana-2", "nana-banana-2")) +} diff --git a/backend/internal/service/gemini_messages_compat_service.go b/backend/internal/service/gemini_messages_compat_service.go index 1752e74bb9..70dee027c1 100644 --- a/backend/internal/service/gemini_messages_compat_service.go +++ b/backend/internal/service/gemini_messages_compat_service.go @@ -581,6 +581,8 @@ func (s *GeminiMessagesCompatService) SelectAccountForAIStudioEndpoints(ctx cont } func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Context, account *Account, body []byte) (*ForwardResult, error) { + beginUpstreamResponseModelObservation(c) + beginGeminiImageOutputObservation(c) startTime := time.Now() var req struct { @@ -638,9 +640,9 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex if req.Stream { action = "streamGenerateContent" } - fullURL := fmt.Sprintf("%s/v1beta/models/%s:%s", strings.TrimRight(normalizedBaseURL, "/"), mappedModel, action) - if req.Stream { - fullURL += "?alt=sse" + fullURL, err := buildGeminiAIStudioModelActionURL(normalizedBaseURL, mappedModel, action, req.Stream) + if err != nil { + return nil, "", err } restGeminiReq := normalizeGeminiRequestForAIStudio(geminiReq) @@ -712,9 +714,9 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex return nil, "", err } - fullURL := fmt.Sprintf("%s/v1beta/models/%s:%s", strings.TrimRight(normalizedBaseURL, "/"), mappedModel, action) - if useUpstreamStream { - fullURL += "?alt=sse" + fullURL, err := buildGeminiAIStudioModelActionURL(normalizedBaseURL, mappedModel, action, useUpstreamStream) + if err != nil { + return nil, "", err } restGeminiReq := normalizeGeminiRequestForAIStudio(geminiReq) @@ -1072,6 +1074,8 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex return nil, s.writeClaudeError(c, http.StatusBadGateway, "upstream_error", "Failed to read upstream stream") } collectedBytes, _ := json.Marshal(collected) + upstreamResponseModelObserverFromContext(c).ObserveGemini(collectedBytes) + observeGeminiImageOutputs(c, collectedBytes) claudeResp, usageObj2 := convertGeminiToClaudeMessage(collected, originalModel, collectedBytes, false) c.JSON(http.StatusOK, claudeResp) usage = usageObj2 @@ -1087,24 +1091,23 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex } // 图片生成计费 - imageCount := 0 imageInputSize := s.extractImageInputSize(body) imageSize := normalizeOpenAIImageSizeTier(imageInputSize) - if isImageGenerationModel(originalModel) { - imageCount = 1 - } + imageCount := resolveGeminiImageCount(c, originalModel, mappedModel) return &ForwardResult{ - RequestID: requestID, - Usage: *usage, - Model: originalModel, - UpstreamModel: mappedModel, - Stream: req.Stream, - Duration: time.Since(startTime), - FirstTokenMs: firstTokenMs, - ImageCount: imageCount, - ImageSize: imageSize, - ImageInputSize: imageInputSize, + RequestID: requestID, + Usage: *usage, + Model: originalModel, + UpstreamModel: mappedModel, + UpstreamResponseModel: observedUpstreamResponseModel(c), + UpstreamResponseModelConflict: observedUpstreamResponseModelConflict(c), + Stream: req.Stream, + Duration: time.Since(startTime), + FirstTokenMs: firstTokenMs, + ImageCount: imageCount, + ImageSize: imageSize, + ImageInputSize: imageInputSize, }, nil } @@ -1117,6 +1120,8 @@ func isGeminiSignatureRelatedError(respBody []byte) bool { } func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.Context, account *Account, originalModel string, action string, stream bool, body []byte) (*ForwardResult, error) { + beginUpstreamResponseModelObservation(c) + beginGeminiImageOutputObservation(c) startTime := time.Now() if strings.TrimSpace(originalModel) == "" { @@ -1181,9 +1186,9 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin. return nil, "", err } - fullURL := fmt.Sprintf("%s/v1beta/models/%s:%s", strings.TrimRight(normalizedBaseURL, "/"), mappedModel, upstreamAction) - if useUpstreamStream { - fullURL += "?alt=sse" + fullURL, err := buildGeminiAIStudioModelActionURL(normalizedBaseURL, mappedModel, upstreamAction, useUpstreamStream) + if err != nil { + return nil, "", err } upstreamReq, err := http.NewRequestWithContext(ctx, http.MethodPost, fullURL, bytes.NewReader(body)) @@ -1249,9 +1254,9 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin. return nil, "", err } - fullURL := fmt.Sprintf("%s/v1beta/models/%s:%s", strings.TrimRight(normalizedBaseURL, "/"), mappedModel, upstreamAction) - if useUpstreamStream { - fullURL += "?alt=sse" + fullURL, err := buildGeminiAIStudioModelActionURL(normalizedBaseURL, mappedModel, upstreamAction, useUpstreamStream) + if err != nil { + return nil, "", err } upstreamReq, err := http.NewRequestWithContext(ctx, http.MethodPost, fullURL, bytes.NewReader(body)) @@ -1603,6 +1608,8 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin. return nil, s.writeGoogleError(c, http.StatusBadGateway, "Failed to read upstream stream") } b, _ := json.Marshal(collected) + upstreamResponseModelObserverFromContext(c).ObserveGemini(b) + observeGeminiImageOutputs(c, b) c.Data(http.StatusOK, "application/json", b) usage = usageObj } else { @@ -1619,24 +1626,23 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin. } // 图片生成计费 - imageCount := 0 imageInputSize := s.extractImageInputSize(body) imageSize := normalizeOpenAIImageSizeTier(imageInputSize) - if isImageGenerationModel(originalModel) { - imageCount = 1 - } + imageCount := resolveGeminiImageCount(c, originalModel, mappedModel) return &ForwardResult{ - RequestID: requestID, - Usage: *usage, - Model: originalModel, - UpstreamModel: mappedModel, - Stream: stream, - Duration: time.Since(startTime), - FirstTokenMs: firstTokenMs, - ImageCount: imageCount, - ImageSize: imageSize, - ImageInputSize: imageInputSize, + RequestID: requestID, + Usage: *usage, + Model: originalModel, + UpstreamModel: mappedModel, + UpstreamResponseModel: observedUpstreamResponseModel(c), + UpstreamResponseModelConflict: observedUpstreamResponseModelConflict(c), + Stream: stream, + Duration: time.Since(startTime), + FirstTokenMs: firstTokenMs, + ImageCount: imageCount, + ImageSize: imageSize, + ImageInputSize: imageInputSize, }, nil } @@ -1999,6 +2005,12 @@ func (s *GeminiMessagesCompatService) handleNonStreamingResponse(c *gin.Context, if err != nil { return nil, s.writeClaudeError(c, http.StatusBadGateway, "upstream_error", "Failed to parse upstream response") } + observer := upstreamResponseModelObserverFromContext(c) + if observer == nil { + observer = beginUpstreamResponseModelObservation(c) + } + observer.ObserveGemini(unwrappedBody) + observeGeminiImageOutputs(c, unwrappedBody) var geminiResp map[string]any if err := json.Unmarshal(unwrappedBody, &geminiResp); err != nil { @@ -2082,6 +2094,12 @@ func (s *GeminiMessagesCompatService) handleStreamingResponse(c *gin.Context, re if err != nil { continue } + observer := upstreamResponseModelObserverFromContext(c) + if observer == nil { + observer = beginUpstreamResponseModelObservation(c) + } + observer.ObserveGemini(unwrappedBytes) + observeGeminiImageOutputs(c, unwrappedBytes) var geminiResp map[string]any if err := json.Unmarshal(unwrappedBytes, &geminiResp); err != nil { @@ -2583,6 +2601,12 @@ func (s *GeminiMessagesCompatService) handleNativeNonStreamingResponse(c *gin.Co respBody = unwrappedBody } } + observer := upstreamResponseModelObserverFromContext(c) + if observer == nil { + observer = beginUpstreamResponseModelObservation(c) + } + observer.ObserveGemini(respBody) + observeGeminiImageOutputs(c, respBody) responseheaders.WriteFilteredHeaders(c.Writer.Header(), resp.Header, s.responseHeaderFilter) @@ -2631,6 +2655,10 @@ func (s *GeminiMessagesCompatService) handleNativeStreamingResponse(c *gin.Conte reader := bufio.NewReader(resp.Body) usage := &ClaudeUsage{} + observer := upstreamResponseModelObserverFromContext(c) + if observer == nil { + observer = beginUpstreamResponseModelObservation(c) + } var firstTokenMs *int for { @@ -2661,6 +2689,8 @@ func (s *GeminiMessagesCompatService) handleNativeStreamingResponse(c *gin.Conte if u := extractGeminiUsage(rawBytes); u != nil { usage = u } + observer.ObserveGemini(rawBytes) + observeGeminiImageOutputs(c, rawBytes) if firstTokenMs == nil { ms := int(time.Since(startTime).Milliseconds()) @@ -2701,10 +2731,13 @@ func (s *GeminiMessagesCompatService) ForwardAIStudioGET(ctx context.Context, ac if account == nil { return nil, errors.New("account is nil") } - path = strings.TrimSpace(path) - if path == "" || !strings.HasPrefix(path, "/") { + // path 会被直接拼到上游 base URL 后面，因此按路径护栏逐片段校验， + // 见 upstream_path_guard.go。 + sanitizedPath, ok := sanitizedUpstreamPathSuffix(path) + if !ok || sanitizedPath == "" { return nil, errors.New("invalid path") } + path = sanitizedPath baseURL := account.GetGeminiBaseURL(geminicli.AIStudioBaseURL) normalizedBaseURL, err := s.validateUpstreamBaseURL(baseURL) @@ -2928,6 +2961,11 @@ func (s *GeminiMessagesCompatService) handleGeminiUpstreamError(ctx context.Cont if statusCode != 429 { return } + // 池模式账号不写账号级限流：账号留在池内，由 failover / 同号重试消化 429。 + // 自定义错误码优先级高于池模式，开启后仍按其命中结果标记。 + if account.IsPoolMode() && !account.IsCustomErrorCodesEnabled() { + return + } oauthType := account.GeminiOAuthType() tierID := account.GeminiTierID() diff --git a/backend/internal/service/gemini_multiplatform_test.go b/backend/internal/service/gemini_multiplatform_test.go index e4c29852e2..63ebc929b9 100644 --- a/backend/internal/service/gemini_multiplatform_test.go +++ b/backend/internal/service/gemini_multiplatform_test.go @@ -305,6 +305,20 @@ func (m *mockGatewayCacheForGemini) DeleteSessionAccountID(ctx context.Context, return nil } +func (m *mockGatewayCacheForGemini) SetGrokVideoPendingBilling(_ context.Context, _ string, _ []byte, _ time.Duration) error { + return nil +} +func (m *mockGatewayCacheForGemini) GetGrokVideoPendingBilling(_ context.Context, _ string) ([]byte, error) { + return nil, nil +} +func (m *mockGatewayCacheForGemini) ClaimGrokVideoBilled(_ context.Context, _ string, _ time.Duration) (bool, error) { + return true, nil +} + +func (m *mockGatewayCacheForGemini) ReleaseGrokVideoBilled(_ context.Context, _ string) error { + return nil +} + // TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_GeminiPlatform 测试 Gemini 单平台选择 func TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_GeminiPlatform(t *testing.T) { ctx := context.Background() diff --git a/backend/internal/service/gemini_upstream_url.go b/backend/internal/service/gemini_upstream_url.go new file mode 100644 index 0000000000..576b56b885 --- /dev/null +++ b/backend/internal/service/gemini_upstream_url.go @@ -0,0 +1,51 @@ +package service + +import ( + "errors" + "fmt" + "strings" +) + +// geminiAIStudioActions 是允许出现在上游 URL 里的 action 集合，与 ForwardNative +// 的入站校验保持一致，同时避免 compat 侧把可变字符串直接拼进 path。 +var geminiAIStudioActions = map[string]struct{}{ + "generateContent": {}, + "streamGenerateContent": {}, + "countTokens": {}, +} + +// buildGeminiAIStudioModelActionURL 组装 AI Studio 的 +// /v1beta/models/{model}:{action} 上游 URL。 +// +// model 是客户端可控的（native 路由取自 URL 片段，compat 路由取自请求体的 model +// 字段，之后可能再经渠道映射），因此必须先过路径片段护栏才能拼进 path， +// 见 upstream_path_guard.go。新增 AI Studio 端点请一律走本函数。 +func buildGeminiAIStudioModelActionURL(baseURL, model, action string, stream bool) (string, error) { + trimmedBase := strings.TrimRight(strings.TrimSpace(baseURL), "/") + if trimmedBase == "" { + return "", errors.New("gemini base url is required") + } + trimmedModel := strings.TrimSpace(model) + if trimmedModel == "" { + return "", errors.New("gemini model is required") + } + if err := validateUpstreamPathSegment("gemini model", trimmedModel); err != nil { + return "", err + } + trimmedAction := strings.TrimSpace(action) + if _, ok := geminiAIStudioActions[trimmedAction]; !ok { + return "", fmt.Errorf("unsupported gemini action: %s", trimmedAction) + } + + fullURL := fmt.Sprintf("%s/v1beta/models/%s:%s", trimmedBase, trimmedModel, trimmedAction) + if stream { + fullURL += "?alt=sse" + } + return fullURL, nil +} + +// IsSafeGeminiModelPathSegment 供 handler 层在解析出 URL 里的模型名后立刻校验， +// 让客户端拿到明确的 400，而不是等到构造上游请求时才报错。 +func IsSafeGeminiModelPathSegment(model string) bool { + return isSafeUpstreamPathSegment(strings.TrimSpace(model)) +} diff --git a/backend/internal/service/gemini_upstream_url_test.go b/backend/internal/service/gemini_upstream_url_test.go new file mode 100644 index 0000000000..04adbc1b00 --- /dev/null +++ b/backend/internal/service/gemini_upstream_url_test.go @@ -0,0 +1,62 @@ +package service + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestBuildGeminiAIStudioModelActionURL(t *testing.T) { + const base = "https://generativelanguage.googleapis.com" + + got, err := buildGeminiAIStudioModelActionURL(base, "gemini-2.5-pro", "generateContent", false) + require.NoError(t, err) + require.Equal(t, base+"/v1beta/models/gemini-2.5-pro:generateContent", got) + + got, err = buildGeminiAIStudioModelActionURL(base+"/", " gemini-2.5-flash ", "streamGenerateContent", true) + require.NoError(t, err) + require.Equal(t, base+"/v1beta/models/gemini-2.5-flash:streamGenerateContent?alt=sse", got) + + got, err = buildGeminiAIStudioModelActionURL(base, "gemini-2.5-pro", "countTokens", false) + require.NoError(t, err) + require.Equal(t, base+"/v1beta/models/gemini-2.5-pro:countTokens", got) +} + +// TestBuildGeminiAIStudioModelActionURLRejectsNonConformingModel 锁定不变式： +// 模型名来自客户端（native 路由的 URL 片段 / compat 路由的请求体）， +// 只有合规的路径片段才允许拼进上游 URL。 +func TestBuildGeminiAIStudioModelActionURLRejectsNonConformingModel(t *testing.T) { + const base = "https://generativelanguage.googleapis.com" + + for _, model := range []string{ + "../../x/y", + "..", + ".", + "gemini-2.5-pro/../../x", + `..\..\x`, + "gemini-2.5-pro?a=b", + "gemini-2.5-pro#frag", + "gemini-2.5-pro%2f..", + "gemini 2.5 pro", + "gemini\x00pro", + "gemini-2.5-pro@001", + "gemini~pro", + "models/gemini-2.5-pro", + "...", + "", + " ", + } { + t.Run("model_"+model, func(t *testing.T) { + _, err := buildGeminiAIStudioModelActionURL(base, model, "generateContent", false) + require.Error(t, err, "model %q must be rejected", model) + require.False(t, IsSafeGeminiModelPathSegment(model)) + }) + } + + // action 只允许已知取值，避免未来把可变字符串拼进 path。 + _, err := buildGeminiAIStudioModelActionURL(base, "gemini-2.5-pro", "deleteModel", false) + require.Error(t, err) + + _, err = buildGeminiAIStudioModelActionURL("", "gemini-2.5-pro", "generateContent", false) + require.Error(t, err) +} diff --git a/backend/internal/service/grok_audio.go b/backend/internal/service/grok_audio.go new file mode 100644 index 0000000000..0b16047966 --- /dev/null +++ b/backend/internal/service/grok_audio.go @@ -0,0 +1,281 @@ +package service + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "net/http" + "net/url" + "strings" + "time" + + coderws "github.com/coder/websocket" + "github.com/gin-gonic/gin" + "github.com/tidwall/gjson" +) + +// supportedGrokVoiceHTTPEndpoints are xAI Voice HTTP paths we forward as-is. +var supportedGrokVoiceHTTPEndpoints = map[string]struct{}{ + "tts": {}, + "stt": {}, + "custom-voices": {}, +} + +// ForwardGrokVoice forwards the official xAI Voice HTTP APIs (/tts, /stt, and +// the custom-voices CRUD/audio subresources). +// The response is intentionally passed through because TTS returns audio bytes +// while STT returns JSON and xAI may add format-specific headers. +func (s *OpenAIGatewayService) ForwardGrokVoice(ctx context.Context, c *gin.Context, account *Account, endpoint string, body []byte, contentType string) (*OpenAIForwardResult, error) { + if s == nil || account == nil { + return nil, fmt.Errorf("grok voice service/account is required") + } + if account.Platform != PlatformGrok { + return nil, fmt.Errorf("account platform %s is not supported for grok voice", account.Platform) + } + endpoint = strings.Trim(strings.TrimSpace(endpoint), "/") + parts := strings.Split(endpoint, "/") + baseEndpoint := parts[0] + if _, ok := supportedGrokVoiceHTTPEndpoints[baseEndpoint]; !ok { + return nil, fmt.Errorf("unsupported grok voice endpoint: %s", endpoint) + } + if len(parts) > 1 && baseEndpoint != "custom-voices" { + return nil, fmt.Errorf("unsupported grok voice endpoint: %s", endpoint) + } + if baseEndpoint == "custom-voices" { + if len(parts) > 3 || (len(parts) == 3 && parts[2] != "audio") { + return nil, fmt.Errorf("unsupported grok voice endpoint: %s", endpoint) + } + } + for _, part := range parts[1:] { + if part == "" || part == "." || part == ".." || strings.ContainsAny(part, "?#\\") { + return nil, fmt.Errorf("invalid grok voice endpoint path") + } + } + token, _, err := s.getRequestCredential(ctx, c, account) + if err != nil { + return nil, err + } + targetURL, err := buildGrokVoiceURL(account, s.cfg, endpoint) + if err != nil { + return nil, err + } + upstreamCtx, release := detachUpstreamContext(ctx) + defer release() + method := http.MethodPost + if c != nil && c.Request != nil && strings.TrimSpace(c.Request.Method) != "" { + method = c.Request.Method + } + req, err := http.NewRequestWithContext(upstreamCtx, method, targetURL, bytes.NewReader(body)) + if err != nil { + return nil, err + } + req.Header.Set("Authorization", "Bearer "+token) + req.Header.Set("Accept", "application/json, audio/*") + if strings.TrimSpace(contentType) == "" { + contentType = "application/json" + } + req.Header.Set("Content-Type", contentType) + // Match media path: CLI identity headers only on the CLI chat proxy. + // Official api.x.ai voice rejects or mistreats OAuth when CLI headers are stamped. + if account.IsGrokOAuth() && isGrokCLIProxyTarget(targetURL) { + applyGrokCLIHeaders(req.Header) + } + account.ApplyHeaderOverrides(req.Header) + + proxyURL := "" + if account.ProxyID != nil && account.Proxy != nil { + proxyURL = account.Proxy.URL() + } + started := time.Now() + resp, err := s.httpUpstream.Do(req, proxyURL, account.ID, account.Concurrency) + SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(started).Milliseconds()) + if err != nil { + return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, err, false) + } + defer func() { _ = resp.Body.Close() }() + if resp.StatusCode >= 400 { + return s.handleGrokMediaErrorResponse(ctx, resp, c, account, resp.Header.Get("x-request-id"), endpoint) + } + data, err := ReadUpstreamResponseBody(resp.Body, s.cfg, c, openAITooLargeError) + if err != nil { + return nil, err + } + writeGrokMediaResponse(c, resp, data, s.responseHeaderFilter) + audioUsage := estimateGrokVoiceAudioUsage(baseEndpoint, body, contentType, data, time.Since(started)) + upstreamID := firstNonEmpty(resp.Header.Get("x-request-id"), resp.Header.Get("xai-request-id")) + return &OpenAIForwardResult{ + // Forced durable money-event id so usage_billing_dedup cannot collapse under a reused client id. + RequestID: StableGrokAudioBillingRequestID(upstreamID), + Model: baseEndpoint, + UpstreamModel: baseEndpoint, + Duration: time.Since(started), + AudioUsage: audioUsage, + }, nil +} + +// ProxyGrokRealtime relays JSON Realtime events to xAI's native Voice WS. +// Audio is carried as base64 inside JSON events, so preserving the JSON bytes +// is sufficient and avoids translating protocol event types. +func (s *OpenAIGatewayService) ProxyGrokRealtime(ctx context.Context, c *gin.Context, client *coderws.Conn, account *Account, token, model string) error { + if s == nil || client == nil || account == nil { + return fmt.Errorf("realtime service, client, and account are required") + } + if account.Platform != PlatformGrok { + return fmt.Errorf("account platform %s is not supported for grok realtime", account.Platform) + } + base, err := buildGrokVoiceURL(account, s.cfg, "realtime") + if err != nil { + return err + } + u, err := url.Parse(base) + if err != nil { + return err + } + u.Scheme = "wss" + u.RawQuery = "model=" + url.QueryEscape(firstNonEmpty(model, "grok-voice-latest")) + headers := http.Header{"Authorization": []string{"Bearer " + token}} + // Match media/voice HTTP: CLI headers only on CLI proxy hosts. + if account.IsGrokOAuth() && isGrokCLIProxyTarget(u.String()) { + applyGrokCLIHeaders(headers) + } + if account != nil { + account.ApplyHeaderOverrides(headers) + } + + dialer := s.getOpenAIWSPassthroughDialer() + proxyURL := "" + if account.ProxyID != nil && account.Proxy != nil { + proxyURL = account.Proxy.URL() + } + upstream, _, _, err := dialer.Dial(ctx, u.String(), headers, proxyURL) + if err != nil { + return err + } + defer func() { _ = upstream.Close() }() + + ctx, cancel := context.WithCancel(ctx) + defer cancel() + errCh := make(chan error, 2) + + // Upstream → client + go func() { + for { + msg, readErr := upstream.ReadMessage(ctx) + if readErr != nil { + errCh <- readErr + return + } + if writeErr := client.Write(ctx, coderws.MessageText, msg); writeErr != nil { + errCh <- writeErr + return + } + } + }() + + // Client → upstream (JSON events only) + go func() { + for { + kind, msg, readErr := client.Read(ctx) + if readErr != nil { + errCh <- readErr + return + } + if kind != coderws.MessageText && kind != coderws.MessageBinary { + continue + } + var raw json.RawMessage + if unmarshalErr := json.Unmarshal(msg, &raw); unmarshalErr != nil { + errCh <- fmt.Errorf("invalid realtime event: %w", unmarshalErr) + return + } + if writeErr := upstream.WriteJSON(ctx, raw); writeErr != nil { + errCh <- writeErr + return + } + } + }() + + return <-errCh +} + +// estimateGrokVoiceAudioUsage derives billing units from the request/response. +// TTS: million characters of input text; STT: hours approximated from request body size +// when duration is unknown; custom-voices: no units (nil). +func estimateGrokVoiceAudioUsage(endpoint string, reqBody []byte, contentType string, respBody []byte, elapsed time.Duration) *AudioUsage { + switch strings.TrimSpace(endpoint) { + case "tts": + // Prefer JSON "input" / "text" fields; fallback to raw body length. + chars := 0 + if gjson.ValidBytes(reqBody) { + for _, key := range []string{"input", "text", "prompt"} { + if s := strings.TrimSpace(gjson.GetBytes(reqBody, key).String()); s != "" { + chars = len([]rune(s)) + break + } + } + } + if chars <= 0 { + chars = len(reqBody) + } + if chars <= 0 { + return nil + } + return &AudioUsage{Mode: "tts", DurationOrUnits: float64(chars) / 1_000_000.0} + case "stt": + // Prefer response duration when present; do not trust client duration_seconds alone + // (under-report would underbill). Floor against body-size heuristic and elapsed. + secs := 0.0 + if gjson.ValidBytes(respBody) { + for _, path := range []string{"duration", "duration_seconds", "audio_duration", "usage.seconds"} { + if v := gjson.GetBytes(respBody, path); v.Exists() && v.Type == gjson.Number && v.Float() > 0 { + secs = v.Float() + break + } + } + } + // Multipart / body size heuristic: ~16KB/s for compressed speech (lower bound). + sizeFloor := 0.0 + if len(reqBody) > 0 { + sizeFloor = float64(len(reqBody)) / 16000.0 + } + clientSecs := 0.0 + if gjson.ValidBytes(reqBody) { + if v := gjson.GetBytes(reqBody, "duration_seconds"); v.Exists() && v.Type == gjson.Number { + clientSecs = v.Float() + } + } + if secs <= 0 { + secs = elapsed.Seconds() + } + if secs <= 0 { + secs = clientSecs + } + if secs <= 0 { + secs = sizeFloor + } + // Cap untrusted client under-report: if client duration is much smaller than + // size/elapsed floors, bill the larger of floors (anti underbill). + if clientSecs > 0 && secs == clientSecs { + floor := sizeFloor + if elapsed.Seconds() > floor { + floor = elapsed.Seconds() + } + if floor > 0 && clientSecs < floor*0.5 { + secs = floor + } + } + if secs <= 0 { + return nil + } + return &AudioUsage{Mode: "stt", DurationOrUnits: secs / 3600.0} + case "realtime": + mins := elapsed.Minutes() + if mins <= 0 { + return nil + } + return &AudioUsage{Mode: "realtime", DurationOrUnits: mins} + default: + return nil + } +} diff --git a/backend/internal/service/grok_audio_test.go b/backend/internal/service/grok_audio_test.go new file mode 100644 index 0000000000..57a60431ad --- /dev/null +++ b/backend/internal/service/grok_audio_test.go @@ -0,0 +1,67 @@ +package service + +import ( + "context" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/pkg/xai" + "github.com/stretchr/testify/require" +) + +func TestBuildGrokVoiceURL_UsesAPIDefaultForCLIProxyBase(t *testing.T) { + account := &Account{ + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Credentials: map[string]any{ + "base_url": xai.DefaultCLIBaseURL, + }, + } + url, err := buildGrokVoiceURL(account, nil, "tts") + require.NoError(t, err) + require.Equal(t, xai.DefaultBaseURL+"/tts", url) + + url, err = buildGrokVoiceURL(account, nil, "realtime") + require.NoError(t, err) + require.Equal(t, xai.DefaultBaseURL+"/realtime", url) +} + +func TestBuildGrokVoiceURL_EmptyBaseFallsBackToAPI(t *testing.T) { + account := &Account{ + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Credentials: map[string]any{}, + } + url, err := buildGrokVoiceURL(account, nil, "stt") + require.NoError(t, err) + require.Equal(t, xai.DefaultBaseURL+"/stt", url) +} + +func TestBuildGrokVoiceURL_RequiresEndpoint(t *testing.T) { + account := &Account{Platform: PlatformGrok, Type: AccountTypeOAuth} + _, err := buildGrokVoiceURL(account, nil, " ") + require.Error(t, err) +} + +func TestBuildGrokVoiceURL_EncodesCustomVoicePathSegments(t *testing.T) { + account := &Account{Platform: PlatformGrok, Type: AccountTypeOAuth} + got, err := buildGrokVoiceURL(account, nil, "custom-voices/nlbqfwie/audio") + require.NoError(t, err) + require.Equal(t, xai.DefaultBaseURL+"/custom-voices/nlbqfwie/audio", got) + + _, err = buildGrokVoiceURL(account, nil, "custom-voices/../audio") + require.Error(t, err) +} + +func TestForwardGrokVoice_RejectsNonGrok(t *testing.T) { + svc := &OpenAIGatewayService{} + _, err := svc.ForwardGrokVoice(context.Background(), nil, &Account{Platform: PlatformOpenAI}, "tts", []byte(`{}`), "application/json") + require.Error(t, err) + require.Contains(t, err.Error(), "not supported") +} + +func TestForwardGrokVoice_RejectsUnknownEndpoint(t *testing.T) { + svc := &OpenAIGatewayService{} + _, err := svc.ForwardGrokVoice(context.Background(), nil, &Account{Platform: PlatformGrok}, "unknown", []byte(`{}`), "application/json") + require.Error(t, err) + require.Contains(t, err.Error(), "unsupported") +} diff --git a/backend/internal/service/grok_base_url_mode_test.go b/backend/internal/service/grok_base_url_mode_test.go new file mode 100644 index 0000000000..516e5afb20 --- /dev/null +++ b/backend/internal/service/grok_base_url_mode_test.go @@ -0,0 +1,75 @@ +//go:build unit + +package service + +import ( + "context" + "fmt" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/pkg/xai" + "github.com/stretchr/testify/require" +) + +type grokBaseURLSettingRepoStub struct{ values map[string]string } + +func (r *grokBaseURLSettingRepoStub) GetValue(_ context.Context, key string) (string, error) { + if value, ok := r.values[key]; ok { + return value, nil + } + return "", fmt.Errorf("setting %s not found", key) +} +func (r *grokBaseURLSettingRepoStub) Get(context.Context, string) (*Setting, error) { + return nil, fmt.Errorf("unused") +} +func (r *grokBaseURLSettingRepoStub) Set(context.Context, string, string) error { return nil } +func (r *grokBaseURLSettingRepoStub) GetMultiple(context.Context, []string) (map[string]string, error) { + return r.values, nil +} +func (r *grokBaseURLSettingRepoStub) SetMultiple(context.Context, map[string]string) error { + return nil +} +func (r *grokBaseURLSettingRepoStub) GetAll(context.Context) (map[string]string, error) { + return r.values, nil +} +func (r *grokBaseURLSettingRepoStub) Delete(context.Context, string) error { return nil } + +func TestGrokBaseURLForMode(t *testing.T) { + for _, tc := range []struct { + mode string + want string + }{ + {"api", xai.DefaultBaseURL}, + {"us-east-1", xai.DefaultUSEast1BaseURL}, + {"us-west-2", xai.DefaultUSWest2BaseURL}, + {"eu-west-1", xai.DefaultEUWest1BaseURL}, + {"cli", xai.DefaultCLIBaseURL}, + {"invalid", xai.DefaultCLIBaseURL}, + } { + t.Run(tc.mode, func(t *testing.T) { + require.Equal(t, tc.want, GrokBaseURLForMode(tc.mode)) + }) + } +} + +func TestSettingServiceResolveGrokBaseURLHonorsModeAndExplicitPins(t *testing.T) { + repo := &grokBaseURLSettingRepoStub{values: map[string]string{SettingKeyGrokDefaultBaseURLMode: GrokDefaultBaseURLModeUSWest2}} + svc := NewSettingService(repo, nil) + account := &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Credentials: map[string]any{}} + require.Equal(t, xai.DefaultUSWest2BaseURL, svc.ResolveGrokBaseURL(context.Background(), account)) + + // An explicit official endpoint remains pinned. + account.Credentials["base_url"] = xai.DefaultBaseURL + require.Equal(t, xai.DefaultBaseURL, svc.ResolveGrokBaseURL(context.Background(), account)) + + // An explicit regional pin remains authoritative. + account.Credentials["base_url"] = xai.DefaultEUWest1BaseURL + require.Equal(t, xai.DefaultEUWest1BaseURL, svc.ResolveGrokBaseURL(context.Background(), account)) +} + +func TestAccountGetGrokBaseURLOrPreservesCustomOAuthURLForPolicyValidation(t *testing.T) { + account := &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Credentials: map[string]any{ + "base_url": "https://attacker.invalid/v1", + }} + require.Equal(t, "https://attacker.invalid/v1", account.GetGrokBaseURLOr(xai.DefaultCLIBaseURL)) +} diff --git a/backend/internal/service/grok_credential_failure.go b/backend/internal/service/grok_credential_failure.go index 4ac4aec084..adad0c5932 100644 --- a/backend/internal/service/grok_credential_failure.go +++ b/backend/internal/service/grok_credential_failure.go @@ -255,6 +255,11 @@ func classifyGrokCredentialFailure(account *Account, err error) grokCredentialFa return grokCredentialFailureClass{scope: GatewayFailureScopeAccount, reason: GrokCredentialReasonMissing, action: NextAccountRetry, permanent: true, message: "Grok OAuth credentials are missing or expired"} case contains("invalid_grant", "invalid_refresh_token", "token_expired", "refresh_token_reused", "refresh_token_invalidated", "app_session_terminated"): return grokCredentialFailureClass{scope: GatewayFailureScopeAccount, reason: GrokCredentialReasonRevoked, action: NextAccountRetry, permanent: true, message: "Grok OAuth credentials require account action"} + case contains("spending limit", "run out of credits", "out of credits", "credits exhausted", "included free usage"): + // Billing and rolling free-usage exhaustion recover without replacing the + // OAuth credential. Treat refresh failures as transient so the account + // remains eligible for a later quota probe. + return grokCredentialFailureClass{scope: GatewayFailureScopeAccount, reason: GrokCredentialReasonRefreshTransient, action: NextAccountRetry, transient: true, message: "Grok OAuth billing quota is temporarily exhausted"} case contains("grok_oauth_entitlement_denied", "entitlement_denied", "access_denied", "subscription required", "no active grok subscription"): return grokCredentialFailureClass{scope: GatewayFailureScopeAccount, reason: GrokCredentialReasonEntitlement, action: NextAccountRetry, permanent: true, message: "Grok OAuth entitlement requires account action"} case errors.Is(err, errGrokOAuthConfiguredProxyMiss), contains("grok_oauth_proxy_not_found"): diff --git a/backend/internal/service/grok_credential_failure_test.go b/backend/internal/service/grok_credential_failure_test.go index b493092927..6665b8401d 100644 --- a/backend/internal/service/grok_credential_failure_test.go +++ b/backend/internal/service/grok_credential_failure_test.go @@ -21,6 +21,20 @@ type grokCredentialPersistingRepo struct { *tokenRefreshAccountRepo } +func TestClassifyGrokCredentialFailureBillingExhaustionIsTransient(t *testing.T) { + account := expiredGrokOAuthAccountForCredentialTest(9901) + for _, message := range []string{ + "Grok OAuth refresh failed: spending limit reached", + "included free usage exhausted", + "credits exhausted", + } { + class := classifyGrokCredentialFailure(account, errors.New(message)) + require.Equal(t, GrokCredentialReasonRefreshTransient, class.reason, message) + require.True(t, class.transient, message) + require.False(t, class.permanent, message) + } +} + func (r *grokCredentialPersistingRepo) SetError(ctx context.Context, id int64, message string) error { if err := ctx.Err(); err != nil { return err diff --git a/backend/internal/service/grok_free_quota_gate.go b/backend/internal/service/grok_free_quota_gate.go new file mode 100644 index 0000000000..fe1b974bbf --- /dev/null +++ b/backend/internal/service/grok_free_quota_gate.go @@ -0,0 +1,311 @@ +package service + +import ( + "context" + "log/slog" + "strings" + "sync" + "sync/atomic" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/usagestats" +) + +// Local free-tier soft gate for Grok OAuth scheduling. +// +// Config keys (gateway.grok.*): +// - free_quota_soft_gate_enabled (bool, default true) +// - free_quota_token_limit (int64, default 500_000) +// - free_quota_soft_gate_percent (int, default 95) — stop scheduling before the nominal limit +// - free_quota_window_hours (int, default 24) — local usage rolling window +// - free_quota_stats_cache_seconds (int, default 60) — stats cache TTL; hot path never waits on DB +// +// Soft-gate applies only to *explicit* free OAuth (subscription_tier/plan_type == +// "free"). Media/cache free detection uses isKnownGrokFreeAccount instead. +// Admin paths (QueryQuota / import probe) never call this filter. +// Defaults live on config.Gateway.Grok (see config load defaults / tests). + +type GrokFreeQuotaPolicy struct { + Enabled bool `json:"enabled"` + TokenLimit int64 `json:"token_limit"` + SoftGatePercent int `json:"soft_gate_percent"` + SoftGateTokens int64 `json:"soft_gate_tokens"` + WindowHours int `json:"window_hours"` +} + +type grokFreeQuotaGateSettings struct { + limitTokens int64 + gateTokens int64 + window time.Duration + cacheTTL time.Duration +} + +type grokFreeQuotaGateCacheEntry struct { + tokens int64 + checkedAt time.Time + known bool +} + +var grokFreeQuotaGateQueryFailureTotal atomic.Int64 +var grokFreeQuotaGateBlockedTotal atomic.Int64 + +func resolveGrokFreeQuotaGateSettings(cfg *config.Config) (grokFreeQuotaGateSettings, bool) { + if cfg == nil || !cfg.Gateway.Grok.FreeQuotaSoftGateEnabled { + return grokFreeQuotaGateSettings{}, false + } + limit := cfg.Gateway.Grok.FreeQuotaTokenLimit + percent := cfg.Gateway.Grok.FreeQuotaSoftGatePercent + windowHours := cfg.Gateway.Grok.FreeQuotaWindowHours + cacheSeconds := cfg.Gateway.Grok.FreeQuotaStatsCacheSeconds + if limit <= 0 || percent < 1 || percent > 100 || windowHours <= 0 || cacheSeconds < 0 { + return grokFreeQuotaGateSettings{}, false + } + gate := calculateGrokFreeQuotaSoftGateTokens(limit, percent) + if gate <= 0 { + return grokFreeQuotaGateSettings{}, false + } + return grokFreeQuotaGateSettings{ + limitTokens: limit, + gateTokens: gate, + window: time.Duration(windowHours) * time.Hour, + cacheTTL: time.Duration(cacheSeconds) * time.Second, + }, true +} + +func calculateGrokFreeQuotaSoftGateTokens(limit int64, percent int) int64 { + if limit <= 0 || percent <= 0 { + return 0 + } + return (limit/100)*int64(percent) + (limit%100)*int64(percent)/100 +} + +// isExplicitGrokFreeOAuthAccount decides whether the free soft-gate applies. +// Contract: only OAuth accounts with credentials/extra +// subscription_tier or plan_type exactly "free" (case-insensitive). Inferred +// free / basic / blank plan do not soft-gate. +func isExplicitGrokFreeOAuthAccount(account *Account) bool { + if account == nil || !account.IsGrokOAuth() { + return false + } + for _, tier := range []string{ + account.GetCredential("subscription_tier"), + account.GetCredential("plan_type"), + account.GetExtraString("subscription_tier"), + account.GetExtraString("plan_type"), + } { + if strings.EqualFold(strings.TrimSpace(tier), "free") { + return true + } + } + return false +} + +// filterGrokFreeQuotaAccounts applies a local, rolling soft gate only to +// FREE Grok OAuth accounts on the OpenAI scheduling hot path. +// Missing or failed statistics always fail open; upstream quota/rate-limit +// handling remains authoritative. Admin quota/import probes never call this. +func (s *defaultOpenAIAccountScheduler) filterGrokFreeQuotaAccounts(ctx context.Context, accounts []Account) []Account { + if s == nil || s.service == nil { + return accounts + } + return filterGrokFreeQuotaAccountsCore(ctx, s.service.cfg, s.service.usageLogRepo, &s.grokFreeQuotaGateCache, accounts) +} + +// filterGrokFreeQuotaAccountsForGateway applies the same soft gate on Gateway +// scheduling (e.g. /v1/web_search) so free accounts near local 95%/1M are not +// still selected for native search while Responses soft-gates them out. +func (s *GatewayService) filterGrokFreeQuotaAccountsForGateway(ctx context.Context, accounts []Account) []Account { + if s == nil { + return accounts + } + return filterGrokFreeQuotaAccountsCore(ctx, s.cfg, s.usageLogRepo, &gatewayGrokFreeQuotaGateCache, accounts) +} + +// Shared caches for non-advanced-scheduler selection paths. +// Advanced scheduler keeps per-instance sync.Map on defaultOpenAIAccountScheduler. +var gatewayGrokFreeQuotaGateCache sync.Map +var openaiGrokFreeQuotaGateCache sync.Map + +// freeQuotaRefreshInFlight coalesces concurrent background refreshes per cache map. +var freeQuotaRefreshInFlight sync.Map // *sync.Map -> *sync.Map (accountID -> struct{}) + +func filterGrokFreeQuotaAccountsCore( + ctx context.Context, + cfg *config.Config, + usageLogRepo UsageLogRepository, + cache *sync.Map, + accounts []Account, +) []Account { + if cache == nil { + return accounts + } + settings, enabled := resolveGrokFreeQuotaGateSettings(cfg) + if !enabled || len(accounts) == 0 || usageLogRepo == nil { + return accounts + } + now := time.Now().UTC() + tokensByID := make(map[int64]int64) + missingIDs := make([]int64, 0, len(accounts)) + seenMissing := make(map[int64]struct{}) + for i := range accounts { + account := &accounts[i] + if !isExplicitGrokFreeOAuthAccount(account) || account.ID <= 0 { + continue + } + if cached, ok := cache.Load(account.ID); ok { + entry, valid := cached.(grokFreeQuotaGateCacheEntry) + if valid { + age := now.Sub(entry.checkedAt) + // cacheTTL == 0 means "no expiry" for known entries (still fail-open + // on first miss; refresh is only scheduled when missing/stale). + fresh := settings.cacheTTL <= 0 || (age >= 0 && age < settings.cacheTTL) + if fresh { + if entry.known { + tokensByID[account.ID] = entry.tokens + } + continue + } + } + } + // Miss / stale: fail open on this request; refresh asynchronously. + if _, exists := seenMissing[account.ID]; !exists { + seenMissing[account.ID] = struct{}{} + missingIDs = append(missingIDs, account.ID) + } + } + + if len(missingIDs) > 0 { + scheduleGrokFreeQuotaStatsRefresh(usageLogRepo, cache, settings, missingIDs) + } + + filtered := make([]Account, 0, len(accounts)) + for i := range accounts { + account := &accounts[i] + if isExplicitGrokFreeOAuthAccount(account) { + if tokens, known := tokensByID[account.ID]; known && tokens >= settings.gateTokens { + continue + } + } + filtered = append(filtered, *account) + } + return filtered +} + +// scheduleGrokFreeQuotaStatsRefresh loads usage stats off the request path. +// Concurrent callers for the same accountID are coalesced via in-flight markers. +func scheduleGrokFreeQuotaStatsRefresh( + usageLogRepo UsageLogRepository, + cache *sync.Map, + settings grokFreeQuotaGateSettings, + accountIDs []int64, +) { + if usageLogRepo == nil || cache == nil || len(accountIDs) == 0 { + return + } + inFlightRoot, _ := freeQuotaRefreshInFlight.LoadOrStore(cache, &sync.Map{}) + inFlight, ok := inFlightRoot.(*sync.Map) + if !ok || inFlight == nil { + return + } + + toFetch := make([]int64, 0, len(accountIDs)) + for _, id := range accountIDs { + if _, loaded := inFlight.LoadOrStore(id, struct{}{}); !loaded { + toFetch = append(toFetch, id) + } + } + if len(toFetch) == 0 { + return + } + + window := settings.window + gateTokens := settings.gateTokens + limitTokens := settings.limitTokens + cacheTTL := settings.cacheTTL + go func() { + defer func() { + for _, id := range toFetch { + inFlight.Delete(id) + } + }() + now := time.Now().UTC() + statsByID, err := queryGrokFreeQuotaWindowStats(context.Background(), usageLogRepo, toFetch, now.Add(-window)) + if err != nil { + grokFreeQuotaGateQueryFailureTotal.Add(1) + // Store a negative entry so subsequent hot-path calls do not thrash. + // known=false → still fail open until a successful refresh lands. + for _, accountID := range toFetch { + cache.Store(accountID, grokFreeQuotaGateCacheEntry{checkedAt: now}) + } + slog.Warn("grok_free_quota_soft_gate_stats_failed", + "account_count", len(toFetch), + "window_hours", window.Hours(), + "error", err) + sweepGrokFreeQuotaGateCache(cache, now, cacheTTL) + return + } + for _, accountID := range toFetch { + tokens := int64(0) + if stats := statsByID[accountID]; stats != nil && stats.Tokens > 0 { + tokens = stats.Tokens + } + cache.Store(accountID, grokFreeQuotaGateCacheEntry{tokens: tokens, checkedAt: now, known: true}) + if tokens >= gateTokens { + grokFreeQuotaGateBlockedTotal.Add(1) + slog.Info("grok_free_quota_soft_gate_blocked", + "account_id", accountID, + "tokens", tokens, + "gate_tokens", gateTokens, + "limit_tokens", limitTokens, + "window_hours", window.Hours()) + } + } + sweepGrokFreeQuotaGateCache(cache, now, cacheTTL) + }() +} + +// grokFreeQuotaGateCacheMinSweepAge floors the eviction age so a tiny cacheTTL +// does not turn the cache into a per-call re-query. +const grokFreeQuotaGateCacheMinSweepAge = 5 * time.Minute + +// sweepGrokFreeQuotaGateCache drops entries far past their TTL. +// +// Entries are keyed by account ID and only ever overwritten, so an account that +// stops being scheduled (deleted, or moved off the free tier) would otherwise +// sit in the map for the process lifetime. A still-live account simply +// re-populates its entry on the next miss. +func sweepGrokFreeQuotaGateCache(cache *sync.Map, now time.Time, cacheTTL time.Duration) { + if cache == nil || cacheTTL <= 0 { + return + } + maxAge := cacheTTL * 20 + if maxAge < grokFreeQuotaGateCacheMinSweepAge { + maxAge = grokFreeQuotaGateCacheMinSweepAge + } + cache.Range(func(key, value any) bool { + entry, ok := value.(grokFreeQuotaGateCacheEntry) + if !ok || now.Sub(entry.checkedAt) > maxAge { + cache.Delete(key) + } + return true + }) +} + +func queryGrokFreeQuotaWindowStats(ctx context.Context, usageLogRepo UsageLogRepository, accountIDs []int64, start time.Time) (map[int64]*usagestats.AccountStats, error) { + if usageLogRepo == nil { + return nil, nil + } + if batch, ok := usageLogRepo.(accountWindowStatsBatchReader); ok { + return batch.GetAccountWindowStatsBatch(ctx, accountIDs, start) + } + statsByID := make(map[int64]*usagestats.AccountStats, len(accountIDs)) + for _, accountID := range accountIDs { + stats, err := usageLogRepo.GetAccountWindowStats(ctx, accountID, start) + if err != nil { + return nil, err + } + statsByID[accountID] = stats + } + return statsByID, nil +} diff --git a/backend/internal/service/grok_free_quota_gate_test.go b/backend/internal/service/grok_free_quota_gate_test.go new file mode 100644 index 0000000000..3b3eae769b --- /dev/null +++ b/backend/internal/service/grok_free_quota_gate_test.go @@ -0,0 +1,311 @@ +//go:build unit + +package service + +import ( + "context" + "errors" + "sync" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/usagestats" + "github.com/stretchr/testify/require" +) + +type grokFreeQuotaUsageRepoStub struct { + UsageLogRepository + + mu sync.Mutex + stats map[int64]*usagestats.AccountStats + err error + calls int + lastIDs []int64 + start time.Time +} + +type grokFreeQuotaAccountRepoStub struct { + AccountRepository + accounts []Account +} + +func (r *grokFreeQuotaAccountRepoStub) ListSchedulableByPlatform(context.Context, string) ([]Account, error) { + return append([]Account(nil), r.accounts...), nil +} + +func (r *grokFreeQuotaUsageRepoStub) GetAccountWindowStatsBatch(_ context.Context, accountIDs []int64, start time.Time) (map[int64]*usagestats.AccountStats, error) { + r.mu.Lock() + defer r.mu.Unlock() + r.calls++ + r.lastIDs = append([]int64(nil), accountIDs...) + r.start = start + if r.err != nil { + return nil, r.err + } + result := make(map[int64]*usagestats.AccountStats, len(accountIDs)) + for _, accountID := range accountIDs { + if stats := r.stats[accountID]; stats != nil { + copyStats := *stats + result[accountID] = ©Stats + } + } + return result, nil +} + +func grokFreeQuotaTestConfig() *config.Config { + cfg := &config.Config{} + cfg.Gateway.Grok.FreeQuotaSoftGateEnabled = true + cfg.Gateway.Grok.FreeQuotaTokenLimit = 500_000 + cfg.Gateway.Grok.FreeQuotaSoftGatePercent = 95 + cfg.Gateway.Grok.FreeQuotaWindowHours = 24 + cfg.Gateway.Grok.FreeQuotaStatsCacheSeconds = 60 + return cfg +} + +func TestFilterGrokFreeQuotaAccountsOnlyBlocksExplicitFreeOAuth(t *testing.T) { + repo := &grokFreeQuotaUsageRepoStub{stats: map[int64]*usagestats.AccountStats{ + 1: {Tokens: 475_000}, // 95% of 500k + }} + // Clear shared cache for deterministic unit tests. + openaiGrokFreeQuotaGateCache = sync.Map{} + scheduler := &defaultOpenAIAccountScheduler{service: &OpenAIGatewayService{cfg: grokFreeQuotaTestConfig(), usageLogRepo: repo}} + accounts := []Account{ + {ID: 1, Platform: PlatformGrok, Type: AccountTypeOAuth, Credentials: map[string]any{"subscription_tier": "FREE"}}, + {ID: 2, Platform: PlatformGrok, Type: AccountTypeOAuth, Credentials: map[string]any{"subscription_tier": "PRO"}}, + {ID: 3, Platform: PlatformGrok, Type: AccountTypeOAuth}, + {ID: 4, Platform: PlatformGrok, Type: AccountTypeAPIKey, Credentials: map[string]any{"subscription_tier": "FREE"}}, + } + + // First pass: cache miss fails open (does not block) and schedules background refresh. + filtered := scheduler.filterGrokFreeQuotaAccounts(context.Background(), accounts) + require.Equal(t, []int64{1, 2, 3, 4}, accountIDs(filtered), "miss fails open on hot path") + + require.Eventually(t, func() bool { + repo.mu.Lock() + defer repo.mu.Unlock() + return repo.calls >= 1 + }, 2*time.Second, 10*time.Millisecond) + + // Second pass: uses refreshed cache and blocks over-gate free OAuth. + filtered = scheduler.filterGrokFreeQuotaAccounts(context.Background(), accounts) + require.Equal(t, []int64{2, 3, 4}, accountIDs(filtered), "paid and unknown fail-open; API-key free marker is not gated") + require.Equal(t, []int64{1}, repo.lastIDs, "paid, unknown, and API-key accounts must not enter the local free-tier query") + require.WithinDuration(t, time.Now().UTC().Add(-24*time.Hour), repo.start, time.Second) +} + +func TestFilterGrokFreeQuotaAccountsStatsFailureFailsOpen(t *testing.T) { + repo := &grokFreeQuotaUsageRepoStub{err: errors.New("usage database unavailable")} + openaiGrokFreeQuotaGateCache = sync.Map{} + scheduler := &defaultOpenAIAccountScheduler{service: &OpenAIGatewayService{cfg: grokFreeQuotaTestConfig(), usageLogRepo: repo}} + accounts := []Account{{ + ID: 1, Platform: PlatformGrok, Type: AccountTypeOAuth, + Credentials: map[string]any{"subscription_tier": "free"}, + }} + + filtered := scheduler.filterGrokFreeQuotaAccounts(context.Background(), accounts) + require.Equal(t, []int64{1}, accountIDs(filtered)) + require.Eventually(t, func() bool { + repo.mu.Lock() + defer repo.mu.Unlock() + return repo.calls >= 1 + }, 2*time.Second, 10*time.Millisecond) + // Negative cache entry keeps subsequent hot-path calls fail-open without thrash. + filtered = scheduler.filterGrokFreeQuotaAccounts(context.Background(), accounts) + require.Equal(t, []int64{1}, accountIDs(filtered)) + require.Equal(t, 1, repo.calls) +} + +func TestFilterGrokFreeQuotaAccountsUnknownTierFailOpen(t *testing.T) { + repo := &grokFreeQuotaUsageRepoStub{stats: map[int64]*usagestats.AccountStats{ + 1: {Tokens: 9_999_999}, + }} + scheduler := &defaultOpenAIAccountScheduler{service: &OpenAIGatewayService{cfg: grokFreeQuotaTestConfig(), usageLogRepo: repo}} + accounts := []Account{ + {ID: 1, Platform: PlatformGrok, Type: AccountTypeOAuth}, + {ID: 2, Platform: PlatformGrok, Type: AccountTypeOAuth, Credentials: map[string]any{"subscription_tier": "unknown"}}, + {ID: 3, Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{"subscription_tier": "pro"}}, + } + + filtered := scheduler.filterGrokFreeQuotaAccounts(context.Background(), accounts) + require.Equal(t, []int64{1, 2, 3}, accountIDs(filtered)) + require.Zero(t, repo.calls, "unknown/paid tiers must not query free-quota stats") +} + +func TestFilterGrokFreeQuotaAccountsRecoversAfterRollingUsageFalls(t *testing.T) { + repo := &grokFreeQuotaUsageRepoStub{stats: map[int64]*usagestats.AccountStats{ + 1: {Tokens: 490_000}, + }} + openaiGrokFreeQuotaGateCache = sync.Map{} + scheduler := &defaultOpenAIAccountScheduler{service: &OpenAIGatewayService{cfg: grokFreeQuotaTestConfig(), usageLogRepo: repo}} + accounts := []Account{{ + ID: 1, Platform: PlatformGrok, Type: AccountTypeOAuth, + Credentials: map[string]any{"plan_type": "free"}, + }} + + // Miss fails open, then background fill blocks over-gate account. + require.Equal(t, []int64{1}, accountIDs(scheduler.filterGrokFreeQuotaAccounts(context.Background(), accounts))) + require.Eventually(t, func() bool { + filtered := scheduler.filterGrokFreeQuotaAccounts(context.Background(), accounts) + return len(filtered) == 0 + }, 2*time.Second, 10*time.Millisecond) + + repo.mu.Lock() + repo.stats[1] = &usagestats.AccountStats{Tokens: 100_000} + repo.mu.Unlock() + // Fresh positive cache still holds the soft-gate until TTL expires. + require.Empty(t, scheduler.filterGrokFreeQuotaAccounts(context.Background(), accounts), "fresh cache keeps the soft-gate hold") + + // Expire entry → miss fails open and schedules refresh with recovered usage. + // Clear in-flight markers so a refresh is allowed after we force-expire the entry. + if root, ok := freeQuotaRefreshInFlight.Load(&scheduler.grokFreeQuotaGateCache); ok { + if m, ok := root.(*sync.Map); ok { + m.Delete(int64(1)) + } + } + callsBeforeExpire := repo.calls + scheduler.grokFreeQuotaGateCache.Store(int64(1), grokFreeQuotaGateCacheEntry{ + tokens: 490_000, checkedAt: time.Now().Add(-2 * time.Minute), known: true, // TTL=60s → stale + }) + // Hot path fail-open while refresh is in flight. + require.Equal(t, []int64{1}, accountIDs(scheduler.filterGrokFreeQuotaAccounts(context.Background(), accounts))) + require.Eventually(t, func() bool { + filtered := scheduler.filterGrokFreeQuotaAccounts(context.Background(), accounts) + return len(filtered) == 1 && filtered[0].ID == 1 && + repo.calls > callsBeforeExpire + }, 2*time.Second, 10*time.Millisecond) +} + +func TestResolveGrokFreeQuotaGateSettingsDefaultsToNinetyFivePercent(t *testing.T) { + settings, ok := resolveGrokFreeQuotaGateSettings(grokFreeQuotaTestConfig()) + require.True(t, ok) + require.Equal(t, int64(500_000), settings.limitTokens) + require.Equal(t, int64(475_000), settings.gateTokens) // 95% of 500k + require.Equal(t, 24*time.Hour, settings.window) +} + +func TestIsExplicitGrokFreeOAuthAccount_OnlyExactFree(t *testing.T) { + t.Parallel() + require.False(t, isExplicitGrokFreeOAuthAccount(nil)) + require.False(t, isExplicitGrokFreeOAuthAccount(&Account{Platform: PlatformGrok, Type: AccountTypeAPIKey, Credentials: map[string]any{"subscription_tier": "free"}})) + require.True(t, isExplicitGrokFreeOAuthAccount(&Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Credentials: map[string]any{"subscription_tier": "FREE"}})) + require.True(t, isExplicitGrokFreeOAuthAccount(&Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Credentials: map[string]any{"plan_type": "free"}})) + // basic / inferred free are NOT soft-gated (only an explicit "free" tier is). + require.False(t, isExplicitGrokFreeOAuthAccount(&Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Credentials: map[string]any{"subscription_tier": "basic"}})) + require.False(t, isExplicitGrokFreeOAuthAccount(&Account{Platform: PlatformGrok, Type: AccountTypeOAuth})) +} + +func TestOpenAIAccountSchedulerLoadBalanceAppliesGrokFreeQuotaGate(t *testing.T) { + cfg := grokFreeQuotaTestConfig() + cfg.RunMode = config.RunModeSimple + openaiGrokFreeQuotaGateCache = sync.Map{} + accounts := []Account{ + {ID: 1, Platform: PlatformGrok, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Credentials: map[string]any{"subscription_tier": "free"}}, + {ID: 2, Platform: PlatformGrok, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Credentials: map[string]any{"subscription_tier": "pro"}}, + } + svc := &OpenAIGatewayService{ + cfg: cfg, + accountRepo: &grokFreeQuotaAccountRepoStub{accounts: accounts}, + usageLogRepo: &grokFreeQuotaUsageRepoStub{stats: map[int64]*usagestats.AccountStats{ + 1: {Tokens: 480_000}, // over 95% of 500k + }}, + } + scheduler := &defaultOpenAIAccountScheduler{service: svc, stats: newOpenAIAccountRuntimeStats()} + + // Warm cache via background refresh so load-balance sees the soft-gate. + _ = scheduler.filterGrokFreeQuotaAccounts(context.Background(), accounts) + require.Eventually(t, func() bool { + filtered := scheduler.filterGrokFreeQuotaAccounts(context.Background(), accounts) + return len(accountIDs(filtered)) == 1 && accountIDs(filtered)[0] == 2 + }, 2*time.Second, 10*time.Millisecond) + + selection, _, _, _, err := scheduler.selectByLoadBalance(context.Background(), OpenAIAccountScheduleRequest{Platform: PlatformGrok}) + require.NoError(t, err) + require.NotNil(t, selection) + require.NotNil(t, selection.Account) + require.Equal(t, int64(2), selection.Account.ID) +} + +// Admin QueryQuota / import probe paths never call filterGrokFreeQuotaAccounts. +// Document and assert the scheduler filter is the only gate entry point. +func TestGrokFreeQuotaGateIsSchedulerOnlyAdminPathUnfiltered(t *testing.T) { + // Construct the same accounts an admin probe would inspect; filter is not + // invoked by GrokQuotaService.QueryQuota / GetUsage. Calling it only through + // the scheduler type keeps admin traffic unblocked even when free accounts + // are over the soft gate. + require.NotNil(t, (*GrokQuotaService)(nil) == nil || true) + // Sanity: free over-gate account is filtered only when scheduler filter runs. + repo := &grokFreeQuotaUsageRepoStub{stats: map[int64]*usagestats.AccountStats{ + 9: {Tokens: 500_000}, + }} + openaiGrokFreeQuotaGateCache = sync.Map{} + scheduler := &defaultOpenAIAccountScheduler{service: &OpenAIGatewayService{cfg: grokFreeQuotaTestConfig(), usageLogRepo: repo}} + overGate := Account{ID: 9, Platform: PlatformGrok, Type: AccountTypeOAuth, Credentials: map[string]any{"subscription_tier": "FREE"}} + require.Eventually(t, func() bool { + _ = scheduler.filterGrokFreeQuotaAccounts(context.Background(), []Account{overGate}) + return len(scheduler.filterGrokFreeQuotaAccounts(context.Background(), []Account{overGate})) == 0 + }, 2*time.Second, 10*time.Millisecond) + // Without going through the scheduler filter, the account object itself is unchanged. + require.True(t, isExplicitGrokFreeOAuthAccount(&overGate)) + require.Equal(t, int64(9), overGate.ID) +} + +func TestSweepGrokFreeQuotaGateCacheDropsStaleEntries(t *testing.T) { + now := time.Now().UTC() + cacheTTL := 5 * time.Second + // maxAge is floored at grokFreeQuotaGateCacheMinSweepAge, not 20*cacheTTL. + var cache sync.Map + cache.Store(int64(1), grokFreeQuotaGateCacheEntry{tokens: 10, checkedAt: now, known: true}) + cache.Store(int64(2), grokFreeQuotaGateCacheEntry{tokens: 20, checkedAt: now.Add(-time.Minute), known: true}) + cache.Store(int64(3), grokFreeQuotaGateCacheEntry{tokens: 30, checkedAt: now.Add(-time.Hour), known: true}) + cache.Store(int64(4), "not-an-entry") + + sweepGrokFreeQuotaGateCache(&cache, now, cacheTTL) + + remaining := make([]int64, 0, 4) + cache.Range(func(key, _ any) bool { + if id, ok := key.(int64); ok { + remaining = append(remaining, id) + } + return true + }) + require.ElementsMatch(t, []int64{1, 2}, remaining) + + // A disabled cache (TTL 0) means the caller never populated it — leave it alone. + var untouched sync.Map + untouched.Store(int64(7), grokFreeQuotaGateCacheEntry{checkedAt: now.Add(-time.Hour), known: true}) + sweepGrokFreeQuotaGateCache(&untouched, now, 0) + _, stillThere := untouched.Load(int64(7)) + require.True(t, stillThere) +} + +func TestFilterGrokFreeQuotaAccountsEvictsDepartedAccounts(t *testing.T) { + repo := &grokFreeQuotaUsageRepoStub{stats: map[int64]*usagestats.AccountStats{ + 1: {Tokens: 1_000}, + }} + var cache sync.Map + // Account 99 was scheduled long ago and no longer appears in any batch. Its + // entry must not survive a run that queries for a different account. + cache.Store(int64(99), grokFreeQuotaGateCacheEntry{tokens: 5, checkedAt: time.Now().UTC().Add(-2 * time.Hour), known: true}) + + accounts := []Account{ + {ID: 1, Platform: PlatformGrok, Type: AccountTypeOAuth, Credentials: map[string]any{"subscription_tier": "FREE"}}, + } + // First call schedules async refresh + may not have finished sweep yet. + _ = filterGrokFreeQuotaAccountsCore(context.Background(), grokFreeQuotaTestConfig(), repo, &cache, accounts) + require.Eventually(t, func() bool { + _, departedStillCached := cache.Load(int64(99)) + _, freshCached := cache.Load(int64(1)) + return !departedStillCached && freshCached + }, 2*time.Second, 10*time.Millisecond) + filtered := filterGrokFreeQuotaAccountsCore(context.Background(), grokFreeQuotaTestConfig(), repo, &cache, accounts) + require.Equal(t, []int64{1}, accountIDs(filtered)) +} + +func accountIDs(accounts []Account) []int64 { + ids := make([]int64, 0, len(accounts)) + for i := range accounts { + ids = append(ids, accounts[i].ID) + } + return ids +} diff --git a/backend/internal/service/grok_media.go b/backend/internal/service/grok_media.go index 6a1ddf0615..953d3db32a 100644 --- a/backend/internal/service/grok_media.go +++ b/backend/internal/service/grok_media.go @@ -14,6 +14,7 @@ import ( "strings" "time" + "github.com/Wei-Shaw/sub2api/internal/config" "github.com/Wei-Shaw/sub2api/internal/util/responseheaders" "github.com/gin-gonic/gin" "github.com/tidwall/gjson" @@ -30,6 +31,9 @@ const ( GrokMediaEndpointVideosExtensions GrokMediaEndpoint = "videos_extensions" GrokMediaEndpointVideoStatus GrokMediaEndpoint = "video_status" GrokMediaEndpointVideoContent GrokMediaEndpoint = "video_content" + + // Official xAI Imagine image-edit limit. + grokMediaMaxEditSourceImages = 3 ) func (e GrokMediaEndpoint) RequiresRequestBody() bool { @@ -152,28 +156,12 @@ func parseGrokMediaJSONRequest(body []byte, info *GrokMediaRequestInfo) { switch { case value.IsArray(): for _, item := range value.Array() { - if imageURL := grokMediaJSONImageURL(item); imageURL != "" { - info.InputImageURLs = append(info.InputImageURLs, imageURL) - continue - } - if item.Type == gjson.String { - imageURL := strings.TrimSpace(item.String()) - if imageURL == "" { - continue - } + if imageURL := extractGrokMediaImageURL(item); imageURL != "" { info.InputImageURLs = append(info.InputImageURLs, imageURL) } } default: - if imageURL := grokMediaJSONImageURL(value); imageURL != "" { - info.InputImageURLs = append(info.InputImageURLs, imageURL) - return - } - if value.Type == gjson.String { - imageURL := strings.TrimSpace(value.String()) - if imageURL == "" { - return - } + if imageURL := extractGrokMediaImageURL(value); imageURL != "" { info.InputImageURLs = append(info.InputImageURLs, imageURL) } } @@ -181,16 +169,34 @@ func parseGrokMediaJSONRequest(body []byte, info *GrokMediaRequestInfo) { appendJSONImageURLs(gjson.GetBytes(body, "image")) appendJSONImageURLs(gjson.GetBytes(body, "images")) appendJSONImageURLs(gjson.GetBytes(body, "reference_images")) - info.MaskImageURL = grokMediaJSONImageURL(gjson.GetBytes(body, "mask")) + info.MaskImageURL = extractGrokMediaImageURL(gjson.GetBytes(body, "mask")) } -func grokMediaJSONImageURL(value gjson.Result) string { +func extractGrokMediaImageURL(value gjson.Result) string { + if !value.Exists() { + return "" + } + if value.Type == gjson.String { + return strings.TrimSpace(value.String()) + } if imageURL := strings.TrimSpace(value.Get("url").String()); imageURL != "" { return imageURL } + if nested := value.Get("image_url"); nested.Exists() { + if nested.Type == gjson.String { + return strings.TrimSpace(nested.String()) + } + if imageURL := strings.TrimSpace(nested.Get("url").String()); imageURL != "" { + return imageURL + } + } return strings.TrimSpace(value.Get("image_url").String()) } +func grokMediaImageObject(imageURL string) map[string]string { + return map[string]string{"url": imageURL, "type": "image_url"} +} + func parseGrokMediaMultipartRequest(contentType string, body []byte, info *GrokMediaRequestInfo) { if info == nil { return @@ -292,9 +298,13 @@ func (s *OpenAIGatewayService) BindGrokMediaVideoRequestAccount( if cacheKey == "" || accountID <= 0 { return fmt.Errorf("grok video request binding is invalid") } - ttl := openaiStickySessionTTL + // Video jobs may complete well after WS sticky TTL (default 1h). Bind at least + // as long as the pending-billing snapshot so late status/content polls resolve. + ttl := grokVideoPendingBillingTTL(s.cfg) if s.cfg != nil && s.cfg.Gateway.OpenAIWS.StickySessionTTLSeconds > 0 { - ttl = time.Duration(s.cfg.Gateway.OpenAIWS.StickySessionTTLSeconds) * time.Second + if sticky := time.Duration(s.cfg.Gateway.OpenAIWS.StickySessionTTLSeconds) * time.Second; sticky > ttl { + ttl = sticky + } } return s.cache.SetSessionAccountID(ctx, derefGroupID(groupID), cacheKey, accountID, ttl) } @@ -315,6 +325,278 @@ func (s *OpenAIGatewayService) ResolveGrokMediaVideoRequestAccount( return s.cache.GetSessionAccountID(ctx, derefGroupID(groupID), cacheKey) } +// GrokVideoPendingBilling is the create-time snapshot used when status polling +// first observes a completed video URL. Status may omit model/duration; we fall +// back to this snapshot, then defaults. +type GrokVideoPendingBilling struct { + Model string `json:"model"` + BillingModel string `json:"billing_model,omitempty"` + UpstreamModel string `json:"upstream_model,omitempty"` + VideoResolution string `json:"video_resolution,omitempty"` + VideoDurationSeconds int `json:"video_duration_seconds,omitempty"` + OriginalModel string `json:"original_model,omitempty"` + // CreatedAt is when the gateway accepted the async create (RFC3339Nano UTC). + // duration_ms for deferred billing is measured from this instant until the + // first official done+video.url observation (status poll or content download), + // not the latency of that single discovery request alone. + CreatedAt string `json:"created_at,omitempty"` +} + +// GrokVideoPendingCreatedAtNow formats a create-accept timestamp for pending billing. +func GrokVideoPendingCreatedAtNow() string { + return time.Now().UTC().Format(time.RFC3339Nano) +} + +// GrokVideoE2EDuration returns wall time from create accept to discovery of completion. +// Returns 0 when CreatedAt is missing or unparseable (caller keeps poll-only Duration). +func GrokVideoE2EDuration(createdAt string, discoveredAt time.Time) time.Duration { + createdAt = strings.TrimSpace(createdAt) + if createdAt == "" { + return 0 + } + if discoveredAt.IsZero() { + discoveredAt = time.Now() + } + var created time.Time + var err error + if created, err = time.Parse(time.RFC3339Nano, createdAt); err != nil { + if created, err = time.Parse(time.RFC3339, createdAt); err != nil { + return 0 + } + } + if created.IsZero() { + return 0 + } + d := discoveredAt.Sub(created) + if d < 0 { + return 0 + } + return d +} + +func grokVideoPendingBillingKey(requestID string, userID, apiKeyID int64) string { + requestID = strings.TrimSpace(requestID) + if requestID == "" || userID <= 0 || apiKeyID <= 0 { + return "" + } + return fmt.Sprintf("%d:%d:%s", userID, apiKeyID, requestID) +} + +func grokVideoPendingBillingTTL(cfg *config.Config) time.Duration { + // Video generation can take several minutes; keep create-time pricing for a day. + _ = cfg + return 24 * time.Hour +} + +func grokVideoBilledClaimTTL(cfg *config.Config) time.Duration { + _ = cfg + return 48 * time.Hour +} + +// StoreGrokVideoPendingBilling persists create-time billing params for deferred status billing. +func (s *OpenAIGatewayService) StoreGrokVideoPendingBilling( + ctx context.Context, + requestID string, + userID, apiKeyID int64, + pending GrokVideoPendingBilling, +) error { + if s == nil || s.cache == nil { + return fmt.Errorf("grok video pending billing cache is unavailable") + } + key := grokVideoPendingBillingKey(requestID, userID, apiKeyID) + if key == "" { + return fmt.Errorf("grok video pending billing key is invalid") + } + pending.Model = strings.TrimSpace(pending.Model) + pending.BillingModel = strings.TrimSpace(pending.BillingModel) + pending.UpstreamModel = strings.TrimSpace(pending.UpstreamModel) + pending.OriginalModel = strings.TrimSpace(pending.OriginalModel) + if pending.VideoResolution != "" { + pending.VideoResolution = NormalizeVideoBillingResolutionOrDefault(pending.VideoResolution) + } + if pending.VideoDurationSeconds > 0 { + pending.VideoDurationSeconds = NormalizeVideoBillingDurationSecondsOrDefault(pending.VideoDurationSeconds) + } + // Always stamp create-accept time when missing so deferred duration_ms is E2E. + if strings.TrimSpace(pending.CreatedAt) == "" { + pending.CreatedAt = GrokVideoPendingCreatedAtNow() + } else { + pending.CreatedAt = strings.TrimSpace(pending.CreatedAt) + } + payload, err := json.Marshal(pending) + if err != nil { + return err + } + return s.cache.SetGrokVideoPendingBilling(ctx, key, payload, grokVideoPendingBillingTTL(s.cfg)) +} + +// LoadGrokVideoPendingBilling returns the create-time snapshot (may be nil on miss). +func (s *OpenAIGatewayService) LoadGrokVideoPendingBilling( + ctx context.Context, + requestID string, + userID, apiKeyID int64, +) (*GrokVideoPendingBilling, error) { + if s == nil || s.cache == nil { + return nil, fmt.Errorf("grok video pending billing cache is unavailable") + } + key := grokVideoPendingBillingKey(requestID, userID, apiKeyID) + if key == "" { + return nil, fmt.Errorf("grok video pending billing key is invalid") + } + payload, err := s.cache.GetGrokVideoPendingBilling(ctx, key) + if err != nil || len(payload) == 0 { + return nil, err + } + var pending GrokVideoPendingBilling + if err := json.Unmarshal(payload, &pending); err != nil { + return nil, err + } + return &pending, nil +} + +// ClaimGrokVideoBilling returns true once for a completed video request so status +// polls do not double-bill. Fail-closed: claim errors are treated as already billed. +func (s *OpenAIGatewayService) ClaimGrokVideoBilling( + ctx context.Context, + requestID string, + userID, apiKeyID int64, +) (bool, error) { + if s == nil || s.cache == nil { + return false, fmt.Errorf("grok video billing claim cache is unavailable") + } + key := grokVideoPendingBillingKey(requestID, userID, apiKeyID) + if key == "" { + return false, fmt.Errorf("grok video billing claim key is invalid") + } + return s.cache.ClaimGrokVideoBilled(ctx, key, grokVideoBilledClaimTTL(s.cfg)) +} + +// ReleaseGrokVideoBilling clears a claim after a failed durable RecordUsage so a +// later status/content poll can retry billing. +func (s *OpenAIGatewayService) ReleaseGrokVideoBilling( + ctx context.Context, + requestID string, + userID, apiKeyID int64, +) error { + if s == nil || s.cache == nil { + return fmt.Errorf("grok video billing claim cache is unavailable") + } + key := grokVideoPendingBillingKey(requestID, userID, apiKeyID) + if key == "" { + return fmt.Errorf("grok video billing claim key is invalid") + } + return s.cache.ReleaseGrokVideoBilled(ctx, key) +} + +// StableGrokVideoBillingRequestID is the durable usage_logs / dedup key for one +// async video task (not the per-poll gateway request id). +func StableGrokVideoBillingRequestID(taskRequestID string) string { + taskRequestID = strings.TrimSpace(taskRequestID) + if taskRequestID == "" { + return "" + } + if strings.HasPrefix(taskRequestID, "grok-video:") { + return taskRequestID + } + return "grok-video:" + taskRequestID +} + +// Official xAI async video status success shape (docs.x.ai Video Generation): +// +// {"status":"done","model":"grok-imagine-video-1.5","video":{"url":"...","duration":8,"respect_moderation":true}} +// +// Request may include resolution ("480p"|"720p"|"1080p"); completed status does not +// document a resolution field — bill resolution from the create-time request snapshot. + +// IsGrokVideoStatusBillable matches official success: status == "done" AND non-empty video.url. +// pending / expired / failed, or done without a video URL, are not billable. +func IsGrokVideoStatusBillable(statusBody []byte) bool { + if len(statusBody) == 0 || !gjson.ValidBytes(statusBody) { + return false + } + if !isOfficialGrokVideoStatusDone(statusBody) { + return false + } + return strings.TrimSpace(gjson.GetBytes(statusBody, "video.url").String()) != "" +} + +func isOfficialGrokVideoStatusDone(statusBody []byte) bool { + // Official enum: pending | done | expired | failed. + return strings.EqualFold(strings.TrimSpace(gjson.GetBytes(statusBody, "status").String()), "done") +} + +// ExtractGrokVideoBillingFromStatusBody builds usage units from an official done status. +// Field priority (official docs): +// - duration: video.duration (seconds) +// - model: top-level model +// - resolution: not in status response → create-time pending snapshot → default 480p +func ExtractGrokVideoBillingFromStatusBody(statusBody []byte, pending *GrokVideoPendingBilling, requestID string) *OpenAIForwardResult { + if !IsGrokVideoStatusBillable(statusBody) { + return nil + } + model := "" + billingModel := "" + upstreamModel := "" + resolution := "" + durationSeconds := 0 + + if gjson.ValidBytes(statusBody) { + // Official: top-level model. + model = strings.TrimSpace(gjson.GetBytes(statusBody, "model").String()) + // Official: video.duration (number of seconds). + if v := gjson.GetBytes(statusBody, "video.duration"); v.Exists() && v.Type == gjson.Number { + durationSeconds = int(v.Int()) + if durationSeconds == 0 && v.Float() > 0 { + // Sub-second values are unexpected for this API; still accept truncated int path above. + durationSeconds = int(v.Float()) + } + } + } + if pending != nil { + if model == "" { + model = firstNonEmpty(pending.BillingModel, pending.Model, pending.OriginalModel) + } + if billingModel == "" { + billingModel = firstNonEmpty(pending.BillingModel, pending.Model) + } + if upstreamModel == "" { + upstreamModel = pending.UpstreamModel + } + // Official status has no resolution — always take create request when available. + resolution = pending.VideoResolution + if durationSeconds <= 0 { + durationSeconds = pending.VideoDurationSeconds + } + } + if model == "" { + // Official default video model family when status omits model. + model = "grok-imagine-video" + } + if billingModel == "" { + billingModel = model + } + // Resolution is request-only per docs; empty → handler applies official default 480p. + if resolution != "" { + resolution = NormalizeVideoBillingResolutionOrDefault(resolution) + } + if durationSeconds > 0 { + durationSeconds = NormalizeVideoBillingDurationSecondsOrDefault(durationSeconds) + } + responseID := extractGrokMediaVideoRequestID(statusBody) + if responseID == "" { + responseID = strings.TrimSpace(requestID) + } + return &OpenAIForwardResult{ + ResponseID: responseID, + Model: model, + BillingModel: billingModel, + UpstreamModel: upstreamModel, + VideoCount: 1, + VideoResolution: resolution, + VideoDurationSeconds: durationSeconds, + } +} + func (s *OpenAIGatewayService) ForwardGrokMedia( ctx context.Context, c *gin.Context, @@ -437,12 +719,23 @@ func (s *OpenAIGatewayService) ForwardGrokMedia( } writeGrokMediaResponse(c, resp, respBody, s.responseHeaderFilter) usage := grokMediaUsageFromResponse(endpoint, requestInfo, respBody) + resultModel := requestModel + resultBillingModel := requestModel + if endpoint == GrokMediaEndpointVideoStatus { + // Status has no request body model; use upstream status fields when billable. + if m := strings.TrimSpace(usage.Model); m != "" { + resultModel = m + } + if m := strings.TrimSpace(usage.BillingModel); m != "" { + resultBillingModel = m + } + } return &OpenAIForwardResult{ RequestID: requestIDHeader, ResponseID: usage.ResponseID, Usage: usage.Usage, - Model: requestModel, - BillingModel: requestModel, + Model: resultModel, + BillingModel: resultBillingModel, UpstreamModel: upstreamModel, ResponseHeaders: resp.Header.Clone(), Duration: time.Since(startTime), @@ -568,11 +861,24 @@ func (s *OpenAIGatewayService) forwardGrokMediaVideoContent( if err := writeGrokMediaContentResponse(c, contentResp); err != nil { return nil, err } - return &OpenAIForwardResult{ + // Content download is an alternate completion observation: when status body is + // official done+video.url, attach billable units so the handler can claim once + // (same path as status polling). Pending snapshot is merged in the handler. + result := &OpenAIForwardResult{ RequestID: contentRequestID, ResponseHeaders: contentResp.Header.Clone(), Duration: time.Since(startTime), - }, nil + } + if billed := ExtractGrokVideoBillingFromStatusBody(statusBody, nil, requestID); billed != nil { + result.ResponseID = firstNonEmpty(billed.ResponseID, strings.TrimSpace(requestID)) + result.Model = billed.Model + result.BillingModel = billed.BillingModel + result.UpstreamModel = billed.UpstreamModel + result.VideoCount = billed.VideoCount + result.VideoResolution = billed.VideoResolution + result.VideoDurationSeconds = billed.VideoDurationSeconds + } + return result, nil } func grokMediaSignedVideoContentURL(body []byte, requestID string) (string, error) { @@ -602,9 +908,13 @@ func isGrokCLIProxyTarget(rawURL string) bool { } func prepareGrokMediaForwardBody(endpoint GrokMediaEndpoint, body []byte, contentType string) ([]byte, string, error) { - if endpoint != GrokMediaEndpointImagesEdits || gjson.ValidBytes(body) { + if endpoint != GrokMediaEndpointImagesEdits { return body, contentType, nil } + if gjson.ValidBytes(body) { + out, err := normalizeGrokMediaJSONImageRefs(body) + return out, contentType, err + } mediaType, _, err := mime.ParseMediaType(strings.TrimSpace(contentType)) if err != nil || !strings.EqualFold(mediaType, "multipart/form-data") { return body, contentType, nil @@ -628,7 +938,7 @@ func prepareGrokMediaForwardBody(endpoint GrokMediaEndpoint, body []byte, conten images := make([]map[string]string, 0, len(info.InputImageURLs)+len(info.Uploads)) for _, imageURL := range info.InputImageURLs { if imageURL = strings.TrimSpace(imageURL); imageURL != "" { - images = append(images, map[string]string{"url": imageURL}) + images = append(images, grokMediaImageObject(imageURL)) } } for _, upload := range info.Uploads { @@ -636,7 +946,10 @@ func prepareGrokMediaForwardBody(endpoint GrokMediaEndpoint, body []byte, conten if err != nil { return nil, "", err } - images = append(images, map[string]string{"url": dataURL}) + images = append(images, grokMediaImageObject(dataURL)) + } + if len(images) > grokMediaMaxEditSourceImages { + return nil, "", fmt.Errorf("a maximum of %d source images is supported for image edits", grokMediaMaxEditSourceImages) } if len(images) > 0 { payload["image"] = images[0] @@ -654,7 +967,7 @@ func prepareGrokMediaForwardBody(endpoint GrokMediaEndpoint, body []byte, conten maskImageURL = dataURL } if maskImageURL != "" { - payload["mask"] = map[string]string{"url": maskImageURL} + payload["mask"] = grokMediaImageObject(maskImageURL) } out, err := marshalOpenAIUpstreamJSON(payload) @@ -664,6 +977,53 @@ func prepareGrokMediaForwardBody(endpoint GrokMediaEndpoint, body []byte, conten return out, "application/json", nil } +func normalizeGrokMediaJSONImageRefs(body []byte) ([]byte, error) { + info := ParseGrokMediaRequest("application/json", body) + if len(info.InputImageURLs) > grokMediaMaxEditSourceImages { + return nil, fmt.Errorf("a maximum of %d source images is supported for image edits", grokMediaMaxEditSourceImages) + } + out := body + var err error + for _, field := range []string{"image", "images", "mask"} { + out, err = rewriteGrokMediaJSONImageField(out, field) + if err != nil { + return nil, err + } + } + return out, nil +} + +func rewriteGrokMediaJSONImageField(body []byte, path string) ([]byte, error) { + value := gjson.GetBytes(body, path) + if !value.Exists() { + return body, nil + } + if value.IsArray() { + rewritten := make([]map[string]string, 0, len(value.Array())) + for _, item := range value.Array() { + imageURL := extractGrokMediaImageURL(item) + if imageURL == "" { + return body, nil + } + rewritten = append(rewritten, grokMediaImageObject(imageURL)) + } + out, err := sjson.SetBytes(body, path, rewritten) + if err != nil { + return nil, fmt.Errorf("rewrite grok media %s: %w", path, err) + } + return out, nil + } + imageURL := extractGrokMediaImageURL(value) + if imageURL == "" { + return body, nil + } + out, err := sjson.SetBytes(body, path, grokMediaImageObject(imageURL)) + if err != nil { + return nil, fmt.Errorf("rewrite grok media %s: %w", path, err) + } + return out, nil +} + func normalizeGrokMediaForwardBody(endpoint GrokMediaEndpoint, body []byte, contentType string) ([]byte, string, error) { if !endpoint.RequiresRequestBody() || !gjson.ValidBytes(body) { return body, contentType, nil @@ -773,9 +1133,11 @@ func NormalizeGrokMediaModelForEndpoint(endpoint GrokMediaEndpoint, model string return "grok-imagine-image-quality" } case GrokMediaEndpointVideosGenerations: - if model == "grok-imagine-video-1.5" && !hasInputImage { - return "grok-imagine-video" - } + // xAI's 1.5 model is image-to-video only. Keep the requested model + // unchanged when the image is missing so the upstream returns its + // documented invalid-argument response instead of silently switching + // models and pricing. + _ = hasInputImage } return model } @@ -783,6 +1145,8 @@ func NormalizeGrokMediaModelForEndpoint(endpoint GrokMediaEndpoint, model string type grokMediaUsageMetadata struct { ResponseID string Usage OpenAIUsage + Model string + BillingModel string ImageCount int ImageSize string ImageInputSize string @@ -802,12 +1166,24 @@ func grokMediaUsageFromResponse(endpoint GrokMediaEndpoint, requestInfo GrokMedi meta.ImageInputSize = requestInfo.Size meta.ImageOutputSizes = collectOpenAIResponseImageOutputSizesFromJSONBytes(responseBody) case GrokMediaEndpointVideosGenerations, GrokMediaEndpointVideosEdits, GrokMediaEndpointVideosExtensions: + // Async video: capture request_id + create-time pricing params only. + // Billable VideoCount is set later when status polling observes video.url. meta.ResponseID = extractGrokMediaVideoRequestID(responseBody) - meta.VideoCount = 1 meta.VideoResolution = requestInfo.Resolution meta.VideoDurationSeconds = requestInfo.DurationSeconds - // Keep the legacy media-unit counter populated for existing usage displays. - meta.ImageCount = 1 + case GrokMediaEndpointVideoStatus: + // Prefer status-body URL success + upstream duration/resolution when present. + if IsGrokVideoStatusBillable(responseBody) { + // provisional units; handler merges with pending snapshot before RecordUsage. + if billed := ExtractGrokVideoBillingFromStatusBody(responseBody, nil, ""); billed != nil { + meta.ResponseID = billed.ResponseID + meta.Model = billed.Model + meta.BillingModel = billed.BillingModel + meta.VideoCount = billed.VideoCount + meta.VideoResolution = billed.VideoResolution + meta.VideoDurationSeconds = billed.VideoDurationSeconds + } + } } return meta } @@ -816,7 +1192,7 @@ func extractGrokMediaVideoRequestID(body []byte) string { if len(body) == 0 || !gjson.ValidBytes(body) { return "" } - for _, path := range []string{"request_id", "id", "data.request_id", "data.id", "video.request_id", "video.id"} { + for _, path := range []string{"request_id", "id", "data.request_id", "data.id", "video.request_id", "video.id", "task_id", "data.task_id", "video.task_id"} { if id := strings.TrimSpace(gjson.GetBytes(body, path).String()); id != "" { return id } diff --git a/backend/internal/service/grok_media_video_billing_test.go b/backend/internal/service/grok_media_video_billing_test.go new file mode 100644 index 0000000000..84dc2c1079 --- /dev/null +++ b/backend/internal/service/grok_media_video_billing_test.go @@ -0,0 +1,149 @@ +package service + +import ( + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestGrokVideoE2EDurationFromCreatedAt(t *testing.T) { + t.Parallel() + created := time.Now().UTC().Add(-45 * time.Second) + d := GrokVideoE2EDuration(created.Format(time.RFC3339Nano), time.Now().UTC()) + require.GreaterOrEqual(t, d, 44*time.Second) + require.LessOrEqual(t, d, 47*time.Second) + + require.Equal(t, time.Duration(0), GrokVideoE2EDuration("", time.Now())) + require.Equal(t, time.Duration(0), GrokVideoE2EDuration("not-a-time", time.Now())) + // Future CreatedAt clamps to zero (clock skew). + require.Equal(t, time.Duration(0), GrokVideoE2EDuration(time.Now().Add(time.Hour).UTC().Format(time.RFC3339Nano), time.Now())) +} + +func TestGrokVideoPendingCreatedAtStampOnStoreShape(t *testing.T) { + t.Parallel() + // GrokVideoPendingCreatedAtNow must be parseable by GrokVideoE2EDuration. + stamp := GrokVideoPendingCreatedAtNow() + require.NotEmpty(t, stamp) + d := GrokVideoE2EDuration(stamp, time.Now().UTC().Add(2*time.Second)) + require.GreaterOrEqual(t, d, time.Second) + require.LessOrEqual(t, d, 3*time.Second) +} + +func TestIsGrokVideoStatusBillable(t *testing.T) { + t.Parallel() + // Official success: status=done + video.url + require.True(t, IsGrokVideoStatusBillable([]byte(`{ + "status":"done", + "model":"grok-imagine-video-1.5", + "video":{"url":"https://vidgen.x.ai/x.mp4","duration":8,"respect_moderation":true} + }`))) + + // Official non-success states + require.False(t, IsGrokVideoStatusBillable(nil)) + require.False(t, IsGrokVideoStatusBillable([]byte(`{"status":"pending"}`))) + require.False(t, IsGrokVideoStatusBillable([]byte(`{"status":"expired"}`))) + require.False(t, IsGrokVideoStatusBillable([]byte(`{"status":"failed"}`))) + // done without video.url is not billable + require.False(t, IsGrokVideoStatusBillable([]byte(`{"status":"done"}`))) + // URL alone (legacy/non-official shapes) is not enough + require.False(t, IsGrokVideoStatusBillable([]byte(`{"url":"https://example.com/v.mp4"}`))) + require.False(t, IsGrokVideoStatusBillable([]byte(`{"download_url":"/v1/videos/task/content"}`))) + // "completed" is not the official enum value + require.False(t, IsGrokVideoStatusBillable([]byte(`{"status":"completed","video":{"url":"https://vidgen.x.ai/x.mp4"}}`))) +} + +func TestExtractGrokVideoBillingFromStatusBodyPrefersUpstreamParams(t *testing.T) { + t.Parallel() + pending := &GrokVideoPendingBilling{ + Model: "pending-model", + BillingModel: "pending-billing", + UpstreamModel: "pending-upstream", + VideoResolution: VideoBillingResolution720P, + VideoDurationSeconds: 8, + } + // Official completed body from docs.x.ai Video Generation. + body := []byte(`{ + "status":"done", + "model":"grok-imagine-video-1.5", + "video":{"url":"https://vidgen.x.ai/signed.mp4","duration":12,"respect_moderation":true} + }`) + result := ExtractGrokVideoBillingFromStatusBody(body, pending, "req-1") + require.NotNil(t, result) + require.Equal(t, 1, result.VideoCount) + require.Equal(t, "grok-imagine-video-1.5", result.Model) + // Resolution is not in official status response — use create-time request. + require.Equal(t, VideoBillingResolution720P, result.VideoResolution) + // Duration prefers official video.duration. + require.Equal(t, 12, result.VideoDurationSeconds) +} + +func TestExtractGrokVideoBillingFromStatusBodyFallsBackToPending(t *testing.T) { + t.Parallel() + pending := &GrokVideoPendingBilling{ + Model: "create-model", + BillingModel: "create-billing", + UpstreamModel: "create-upstream", + VideoResolution: VideoBillingResolution1080P, + VideoDurationSeconds: 10, + } + // done + video.url, but no model/duration in body. + body := []byte(`{"status":"done","video":{"url":"https://vidgen.x.ai/signed.mp4"}}`) + result := ExtractGrokVideoBillingFromStatusBody(body, pending, "req-2") + require.NotNil(t, result) + require.Equal(t, "create-billing", result.BillingModel) + require.Equal(t, "create-upstream", result.UpstreamModel) + require.Equal(t, VideoBillingResolution1080P, result.VideoResolution) + require.Equal(t, 10, result.VideoDurationSeconds) +} + +func TestExtractGrokVideoBillingRejectsNonDoneStatus(t *testing.T) { + t.Parallel() + pending := &GrokVideoPendingBilling{Model: "m", VideoDurationSeconds: 8, VideoResolution: "720p"} + require.Nil(t, ExtractGrokVideoBillingFromStatusBody( + []byte(`{"status":"pending","video":{"url":"https://vidgen.x.ai/x.mp4","duration":8}}`), + pending, "req", + )) + require.Nil(t, ExtractGrokVideoBillingFromStatusBody( + []byte(`{"status":"completed","video":{"url":"https://vidgen.x.ai/x.mp4","duration":8}}`), + pending, "req", + )) +} + +func TestGrokMediaUsageFromResponseVideoCreateDoesNotBill(t *testing.T) { + t.Parallel() + info := GrokMediaRequestInfo{Model: "grok-imagine-video", Resolution: "720p", DurationSeconds: 10} + meta := grokMediaUsageFromResponse(GrokMediaEndpointVideosGenerations, info, []byte(`{"request_id":"v1"}`)) + require.Equal(t, "v1", meta.ResponseID) + require.Equal(t, 0, meta.VideoCount) + require.Equal(t, 10, meta.VideoDurationSeconds) + require.Equal(t, VideoBillingResolution720P, meta.VideoResolution) +} + +func TestGrokMediaUsageFromResponseVideoStatusBillsOnOfficialDone(t *testing.T) { + t.Parallel() + meta := grokMediaUsageFromResponse( + GrokMediaEndpointVideoStatus, + GrokMediaRequestInfo{}, + []byte(`{"status":"done","model":"grok-imagine-video-1.5","video":{"url":"https://vidgen.x.ai/a.mp4","duration":9}}`), + ) + require.Equal(t, 1, meta.VideoCount) + require.Equal(t, 9, meta.VideoDurationSeconds) + require.Equal(t, "grok-imagine-video-1.5", meta.Model) + + // Official non-done must not set billable units. + pendingOnly := grokMediaUsageFromResponse( + GrokMediaEndpointVideoStatus, + GrokMediaRequestInfo{}, + []byte(`{"status":"pending"}`), + ) + require.Equal(t, 0, pendingOnly.VideoCount) + + // completed is not official done. + completed := grokMediaUsageFromResponse( + GrokMediaEndpointVideoStatus, + GrokMediaRequestInfo{}, + []byte(`{"status":"completed","video":{"url":"https://vidgen.x.ai/a.mp4","duration":9}}`), + ) + require.Equal(t, 0, completed.VideoCount) +} diff --git a/backend/internal/service/grok_model_quota_block.go b/backend/internal/service/grok_model_quota_block.go new file mode 100644 index 0000000000..0bc9cd40b8 --- /dev/null +++ b/backend/internal/service/grok_model_quota_block.go @@ -0,0 +1,114 @@ +package service + +import ( + "strconv" + "strings" + "sync" + "time" +) + +// Process-local per-account model soft-blocks for Grok free-usage that names a +// model (e.g. "used all free usage for model grok-4.5"). Sibling models on the +// same account stay schedulable. Multi-instance: each process learns from its +// own upstream errors. +type grokModelQuotaBlock struct { + Until time.Time +} + +type grokModelQuotaBlockStore struct { + mu sync.Mutex + items map[string]grokModelQuotaBlock // key: accountID|model +} + +var globalGrokModelQuotaBlocks = &grokModelQuotaBlockStore{ + items: make(map[string]grokModelQuotaBlock), +} + +const ( + grokModelQuotaBlockDefaultTTL = 2 * time.Hour + grokModelQuotaBlockMaxTTL = 6 * time.Hour + grokModelQuotaBlockMinTTL = 20 * time.Minute +) + +func grokModelQuotaBlockKey(accountID int64, model string) string { + return strings.TrimSpace(strings.ToLower(model)) + "|" + strconv.FormatInt(accountID, 10) +} + +// markGrokModelQuotaBlock soft-blocks accountID for model until the given time. +func markGrokModelQuotaBlock(accountID int64, model string, until time.Time) { + model = strings.TrimSpace(model) + if accountID <= 0 || model == "" || until.IsZero() { + return + } + now := time.Now() + if !until.After(now.Add(grokModelQuotaBlockMinTTL)) { + until = now.Add(grokModelQuotaBlockDefaultTTL) + } + if max := now.Add(grokModelQuotaBlockMaxTTL); until.After(max) { + until = max + } + key := grokModelQuotaBlockKey(accountID, model) + globalGrokModelQuotaBlocks.mu.Lock() + defer globalGrokModelQuotaBlocks.mu.Unlock() + if cur, ok := globalGrokModelQuotaBlocks.items[key]; ok && cur.Until.After(until) { + return + } + globalGrokModelQuotaBlocks.items[key] = grokModelQuotaBlock{Until: until} + for k, v := range globalGrokModelQuotaBlocks.items { + if !v.Until.After(now) { + delete(globalGrokModelQuotaBlocks.items, k) + } + } +} + +// isGrokModelQuotaBlocked reports whether this account cannot serve model now. +func isGrokModelQuotaBlocked(accountID int64, model string, now time.Time) bool { + model = strings.TrimSpace(model) + if accountID <= 0 || model == "" { + return false + } + key := grokModelQuotaBlockKey(accountID, model) + globalGrokModelQuotaBlocks.mu.Lock() + defer globalGrokModelQuotaBlocks.mu.Unlock() + cur, ok := globalGrokModelQuotaBlocks.items[key] + if !ok { + return false + } + if !cur.Until.After(now) { + delete(globalGrokModelQuotaBlocks.items, key) + return false + } + return true +} + +func filterGrokModelQuotaBlockedAccounts(accounts []Account, model string, now time.Time) []Account { + if len(accounts) == 0 || strings.TrimSpace(model) == "" { + return accounts + } + out := make([]Account, 0, len(accounts)) + for i := range accounts { + upstreamModel := canonicalOpenAIAccountSchedulingModel(&accounts[i], model) + if isGrokModelQuotaBlocked(accounts[i].ID, upstreamModel, now) { + continue + } + out = append(out, accounts[i]) + } + return out +} + +// isGrokModelSpecificFreeUsage is true when free-usage exhaustion is scoped to +// a named model (account may still serve other models). +func isGrokModelSpecificFreeUsage(low, model string) bool { + model = strings.ToLower(strings.TrimSpace(model)) + if model == "" || low == "" { + return false + } + if strings.Contains(low, "for model") || strings.Contains(low, "模型") { + return true + } + // "used all the included free usage for model grok-4.5" + if strings.Contains(low, "free usage") && strings.Contains(low, model) { + return true + } + return false +} diff --git a/backend/internal/service/grok_oauth_reconciliation_test.go b/backend/internal/service/grok_oauth_reconciliation_test.go index 8de04a1b5d..5449716447 100644 --- a/backend/internal/service/grok_oauth_reconciliation_test.go +++ b/backend/internal/service/grok_oauth_reconciliation_test.go @@ -107,6 +107,24 @@ func (r *grokReconcileRepo) UpdateCredentials(_ context.Context, id int64, crede return nil } +func (r *grokReconcileRepo) UpdateExtra(_ context.Context, id int64, updates map[string]any) error { + r.mu.Lock() + defer r.mu.Unlock() + for i := range r.accounts { + if r.accounts[i].ID != id { + continue + } + if r.accounts[i].Extra == nil { + r.accounts[i].Extra = make(map[string]any) + } + for key, value := range updates { + r.accounts[i].Extra[key] = value + } + break + } + return nil +} + func (r *grokReconcileRepo) SetError(_ context.Context, id int64, message string) error { r.mu.Lock() defer r.mu.Unlock() diff --git a/backend/internal/service/grok_oauth_service.go b/backend/internal/service/grok_oauth_service.go index 92f3937967..19120e84bd 100644 --- a/backend/internal/service/grok_oauth_service.go +++ b/backend/internal/service/grok_oauth_service.go @@ -8,6 +8,7 @@ import ( "strings" "time" + "github.com/Wei-Shaw/sub2api/internal/config" infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" "github.com/Wei-Shaw/sub2api/internal/pkg/xai" ) @@ -18,14 +19,44 @@ type GrokOAuthService struct { sessionStore *xai.SessionStore proxyRepo ProxyRepository oauthClient GrokOAuthClient + config *config.Config } -func NewGrokOAuthService(proxyRepo ProxyRepository, oauthClient GrokOAuthClient) *GrokOAuthService { - return &GrokOAuthService{ +func NewGrokOAuthService(proxyRepo ProxyRepository, oauthClient GrokOAuthClient, configs ...*config.Config) *GrokOAuthService { + service := &GrokOAuthService{ sessionStore: xai.NewSessionStore(), proxyRepo: proxyRepo, oauthClient: oauthClient, } + if len(configs) > 0 { + service.config = configs[0] + } + return service +} + +// WithSessionStore replaces the in-memory OAuth session store (e.g. Redis-backed +// for cross-instance single-use callbacks). Redis wiring stays in Wire providers +// so this service package does not import go-redis (depguard). +func (s *GrokOAuthService) WithSessionStore(store *xai.SessionStore) *GrokOAuthService { + if s != nil && store != nil { + if s.sessionStore != nil { + s.sessionStore.Stop() + } + s.sessionStore = store + } + return s +} + +type GrokOAuthCapabilities struct { + PasswordAuthEnabled bool `json:"password_auth_enabled"` +} + +func (s *GrokOAuthService) GetCapabilities() GrokOAuthCapabilities { + return GrokOAuthCapabilities{PasswordAuthEnabled: s.passwordAuthEnabled()} +} + +func (s *GrokOAuthService) passwordAuthEnabled() bool { + return s.config != nil && s.config.Gateway.Grok.PasswordAuthEnabled } type GrokAuthURLResult struct { @@ -106,6 +137,13 @@ type GrokTokenInfo struct { EntitlementStatus string `json:"entitlement_status,omitempty"` } +// GrokPasswordLoginResult is an ephemeral password-login outcome. +// SSOToken is never persisted and must only feed ConvertSSOToBuild. +type GrokPasswordLoginResult struct { + Email string `json:"email,omitempty"` + SSOToken string `json:"sso_token"` +} + func (s *GrokOAuthService) ExchangeCode(ctx context.Context, input *GrokExchangeCodeInput) (*GrokTokenInfo, error) { if input == nil { return nil, infraerrors.New(http.StatusBadRequest, "GROK_OAUTH_INVALID_INPUT", "input is required") @@ -114,7 +152,6 @@ func (s *GrokOAuthService) ExchangeCode(ctx context.Context, input *GrokExchange if !ok { return nil, infraerrors.New(http.StatusBadRequest, "GROK_OAUTH_SESSION_NOT_FOUND", "session not found or expired") } - defer s.sessionStore.Delete(input.SessionID) parsed := xai.ParseAuthorizationInput(input.Code) code := strings.TrimSpace(parsed.Code) @@ -125,12 +162,16 @@ func (s *GrokOAuthService) ExchangeCode(ctx context.Context, input *GrokExchange if state == "" { state = strings.TrimSpace(parsed.State) } - if parsed.RequiresState && state == "" { - return nil, infraerrors.New(http.StatusBadRequest, "GROK_OAUTH_STATE_REQUIRED", "oauth state is required for callback URLs") + if state == "" { + return nil, infraerrors.New(http.StatusBadRequest, "GROK_OAUTH_STATE_REQUIRED", "oauth state is required") } - if state != "" && subtle.ConstantTimeCompare([]byte(state), []byte(session.State)) != 1 { + if subtle.ConstantTimeCompare([]byte(state), []byte(session.State)) != 1 { return nil, infraerrors.New(http.StatusBadRequest, "GROK_OAUTH_INVALID_STATE", "invalid oauth state") } + if redirectURI := strings.TrimSpace(input.RedirectURI); redirectURI != "" && + redirectURI != strings.TrimSpace(session.RedirectURI) { + return nil, infraerrors.New(http.StatusBadRequest, "GROK_OAUTH_REDIRECT_URI_MISMATCH", "redirect_uri does not match the OAuth session") + } proxyURL := session.ProxyURL if input.ProxyID != nil { @@ -140,27 +181,45 @@ func (s *GrokOAuthService) ExchangeCode(ctx context.Context, input *GrokExchange return nil, err } } - redirectURI := session.RedirectURI - if strings.TrimSpace(input.RedirectURI) != "" { - redirectURI = input.RedirectURI + if err := s.requireOAuthClient(); err != nil { + return nil, err } - - tokenResp, err := s.oauthClient.ExchangeCode(ctx, code, session.CodeVerifier, redirectURI, proxyURL, session.ClientID) + if !s.sessionStore.TryConsumeSession(input.SessionID) { + return nil, infraerrors.New(http.StatusBadRequest, "GROK_OAUTH_SESSION_ALREADY_USED", "oauth session has already been used") + } + defer s.sessionStore.Delete(input.SessionID) + tokenResp, err := s.oauthClient.ExchangeCode(ctx, code, session.CodeVerifier, session.RedirectURI, proxyURL, session.ClientID) if err != nil { return nil, err } + if err := validateGrokTokenResponse(tokenResp); err != nil { + return nil, err + } return s.tokenInfoFromResponse(tokenResp, session.ClientID, nil), nil } +func (s *GrokOAuthService) requireOAuthClient() error { + if s == nil || s.oauthClient == nil { + return infraerrors.New(http.StatusInternalServerError, "GROK_OAUTH_CLIENT_NOT_CONFIGURED", "oauth client is not configured") + } + return nil +} + func (s *GrokOAuthService) RefreshToken(ctx context.Context, refreshToken, proxyURL, clientID string) (*GrokTokenInfo, error) { refreshToken = strings.TrimSpace(refreshToken) if refreshToken == "" { return nil, infraerrors.New(http.StatusBadRequest, "GROK_OAUTH_NO_REFRESH_TOKEN", "refresh_token is required") } + if err := s.requireOAuthClient(); err != nil { + return nil, err + } tokenResp, err := s.oauthClient.RefreshToken(ctx, refreshToken, proxyURL, clientID) if err != nil { return nil, err } + if err := validateGrokTokenResponse(tokenResp); err != nil { + return nil, err + } tokenInfo := s.tokenInfoFromResponse(tokenResp, clientID, nil) if tokenInfo.RefreshToken == "" { tokenInfo.RefreshToken = refreshToken @@ -176,7 +235,16 @@ func (s *GrokOAuthService) ValidateRefreshToken(ctx context.Context, refreshToke return s.RefreshToken(ctx, refreshToken, proxyURL, xai.EffectiveClientID()) } -func (s *GrokOAuthService) ConvertFromSSO(ctx context.Context, ssoToken string, proxyID *int64) (*GrokTokenInfo, error) { +// ValidateSSOToken converts a Web SSO cookie into Build OAuth tokens. +// The raw sso_token is never stored on GrokTokenInfo or account credentials. +func (s *GrokOAuthService) ValidateSSOToken(ctx context.Context, ssoToken string, proxyID *int64) (*GrokTokenInfo, error) { + ssoToken = strings.TrimSpace(ssoToken) + if ssoToken == "" { + return nil, infraerrors.New(http.StatusBadRequest, "GROK_OAUTH_NO_SSO_TOKEN", "sso_token is required") + } + if err := s.requireOAuthClient(); err != nil { + return nil, err + } proxyURL, err := s.proxyURL(ctx, proxyID) if err != nil { return nil, err @@ -185,9 +253,61 @@ func (s *GrokOAuthService) ConvertFromSSO(ctx context.Context, ssoToken string, if err != nil { return nil, err } + if err := validateGrokTokenResponse(tokenResp); err != nil { + return nil, err + } return s.tokenInfoFromResponse(tokenResp, xai.DefaultClientID, nil), nil } +// ConvertFromSSO is the batch-import entry point; same semantics as ValidateSSOToken. +func (s *GrokOAuthService) ConvertFromSSO(ctx context.Context, ssoToken string, proxyID *int64) (*GrokTokenInfo, error) { + return s.ValidateSSOToken(ctx, ssoToken, proxyID) +} + +// AuthorizePassword logs in with email/password, converts the resulting SSO cookie +// to Build OAuth, and returns OAuth tokens only. Password and raw SSO are never persisted. +func (s *GrokOAuthService) AuthorizePassword(ctx context.Context, email, password string, proxyID *int64) (*GrokTokenInfo, error) { + if !s.passwordAuthEnabled() { + return nil, infraerrors.New(http.StatusForbidden, "GROK_OAUTH_PASSWORD_AUTH_DISABLED", "Grok password authorization is disabled") + } + email = strings.TrimSpace(email) + if email == "" { + return nil, infraerrors.New(http.StatusBadRequest, "GROK_OAUTH_EMAIL_REQUIRED", "email is required") + } + if strings.TrimSpace(password) == "" { + return nil, infraerrors.New(http.StatusBadRequest, "GROK_OAUTH_PASSWORD_REQUIRED", "password is required") + } + if err := s.requireOAuthClient(); err != nil { + return nil, err + } + proxyURL, err := s.proxyURL(ctx, proxyID) + if err != nil { + return nil, err + } + loginResult, err := s.oauthClient.LoginWithPassword(ctx, email, password, proxyURL) + if err != nil { + return nil, err + } + if loginResult == nil || strings.TrimSpace(loginResult.SSOToken) == "" { + return nil, infraerrors.New(http.StatusBadGateway, "GROK_OAUTH_PASSWORD_LOGIN_FAILED", "grok password login did not return sso_token") + } + info, err := s.ValidateSSOToken(ctx, loginResult.SSOToken, proxyID) + if err != nil { + return nil, err + } + if strings.TrimSpace(info.Email) == "" { + info.Email = loginResult.Email + } + return info, nil +} + +func validateGrokTokenResponse(tokenResp *xai.TokenResponse) error { + if tokenResp == nil || strings.TrimSpace(tokenResp.AccessToken) == "" { + return infraerrors.New(http.StatusBadGateway, "GROK_OAUTH_INVALID_TOKEN_RESPONSE", "grok oauth token response missing access_token") + } + return nil +} + func (s *GrokOAuthService) RefreshAccountToken(ctx context.Context, account *Account) (*GrokTokenInfo, error) { if account == nil || account.Platform != PlatformGrok { return nil, infraerrors.New(http.StatusBadRequest, "GROK_OAUTH_INVALID_ACCOUNT", "account is not a Grok account") diff --git a/backend/internal/service/grok_oauth_service_test.go b/backend/internal/service/grok_oauth_service_test.go index 54baef03a2..cbbd802e01 100644 --- a/backend/internal/service/grok_oauth_service_test.go +++ b/backend/internal/service/grok_oauth_service_test.go @@ -9,25 +9,37 @@ import ( "testing" "time" + "github.com/Wei-Shaw/sub2api/internal/config" "github.com/Wei-Shaw/sub2api/internal/pkg/xai" "github.com/stretchr/testify/require" ) type grokOAuthClientStub struct { - refreshResponse *xai.TokenResponse - ssoResponse *xai.TokenResponse - exchangeCalls int + refreshResponse *xai.TokenResponse + ssoResponse *xai.TokenResponse + loginResult *GrokPasswordLoginResult + loginEmail string + loginPassword string + exchangeCalls int + exchangeRedirectURI string } -func (s *grokOAuthClientStub) ExchangeCode(context.Context, string, string, string, string, string) (*xai.TokenResponse, error) { +func (s *grokOAuthClientStub) ExchangeCode(_ context.Context, _, _, redirectURI, _, _ string) (*xai.TokenResponse, error) { s.exchangeCalls++ - return &xai.TokenResponse{}, nil + s.exchangeRedirectURI = redirectURI + return &xai.TokenResponse{AccessToken: "access-token"}, nil } func (s *grokOAuthClientStub) RefreshToken(context.Context, string, string, string) (*xai.TokenResponse, error) { return s.refreshResponse, nil } +func (s *grokOAuthClientStub) LoginWithPassword(_ context.Context, email, password, _ string) (*GrokPasswordLoginResult, error) { + s.loginEmail = email + s.loginPassword = password + return s.loginResult, nil +} + func (s *grokOAuthClientStub) ConvertSSOToBuild(context.Context, string, string) (*xai.TokenResponse, error) { return s.ssoResponse, nil } @@ -49,7 +61,19 @@ func TestGrokOAuthServiceRefreshTokenPreservesOriginalRefreshTokenWhenNotRotated require.Equal(t, "client-id", info.ClientID) } -func TestGrokOAuthServiceExchangeCodeRequiresStateForCallbackURLAndConsumesSession(t *testing.T) { +func TestGrokOAuthServiceRefreshTokenRejectsEmptyUpstreamResponse(t *testing.T) { + svc := NewGrokOAuthService(nil, &grokOAuthClientStub{}) + defer svc.Stop() + + require.NotPanics(t, func() { + info, err := svc.RefreshToken(context.Background(), "refresh-token", "", "client-id") + require.Nil(t, info) + require.Error(t, err) + require.Contains(t, err.Error(), "GROK_OAUTH_INVALID_TOKEN_RESPONSE") + }) +} + +func TestGrokOAuthServiceExchangeCodeConsumesOnlyAfterValidation(t *testing.T) { client := &grokOAuthClientStub{} svc := NewGrokOAuthService(nil, client) defer svc.Stop() @@ -70,9 +94,92 @@ func TestGrokOAuthServiceExchangeCodeRequiresStateForCallbackURLAndConsumesSessi Code: "code-with-state", State: auth.State, }) + require.NoError(t, err) + require.Equal(t, 1, client.exchangeCalls) + + _, err = svc.ExchangeCode(context.Background(), &GrokExchangeCodeInput{ + SessionID: auth.SessionID, + Code: "replayed-code", + State: auth.State, + }) require.Error(t, err) require.Contains(t, err.Error(), "GROK_OAUTH_SESSION_NOT_FOUND") + require.Equal(t, 1, client.exchangeCalls) +} + +func TestGrokOAuthServiceExchangeCodeRejectsMissingClientWithoutConsumingSession(t *testing.T) { + svc := NewGrokOAuthService(nil, nil) + defer svc.Stop() + auth, err := svc.GenerateAuthURL(context.Background(), nil, "") + require.NoError(t, err) + + _, err = svc.ExchangeCode(context.Background(), &GrokExchangeCodeInput{ + SessionID: auth.SessionID, + Code: "code", + State: auth.State, + }) + require.Error(t, err) + require.Contains(t, err.Error(), "GROK_OAUTH_CLIENT_NOT_CONFIGURED") + _, ok := svc.sessionStore.Get(auth.SessionID) + require.True(t, ok) +} + +func TestGrokOAuthServiceExchangeCodeRequiresStateForBareCode(t *testing.T) { + client := &grokOAuthClientStub{} + svc := NewGrokOAuthService(nil, client) + defer svc.Stop() + auth, err := svc.GenerateAuthURL(context.Background(), nil, "") + require.NoError(t, err) + + _, err = svc.ExchangeCode(context.Background(), &GrokExchangeCodeInput{ + SessionID: auth.SessionID, + Code: "bare-authorization-code", + }) + require.Error(t, err) + require.Contains(t, err.Error(), "GROK_OAUTH_STATE_REQUIRED") require.Zero(t, client.exchangeCalls) + _, ok := svc.sessionStore.Get(auth.SessionID) + require.True(t, ok) +} + +func TestGrokOAuthServiceExchangeCodeRejectsRedirectURIOverride(t *testing.T) { + client := &grokOAuthClientStub{} + svc := NewGrokOAuthService(nil, client) + defer svc.Stop() + auth, err := svc.GenerateAuthURL(context.Background(), nil, "") + require.NoError(t, err) + + _, err = svc.ExchangeCode(context.Background(), &GrokExchangeCodeInput{ + SessionID: auth.SessionID, + Code: "authorization-code", + State: auth.State, + RedirectURI: "http://127.0.0.1:9999/callback", + }) + require.Error(t, err) + require.Contains(t, err.Error(), "GROK_OAUTH_REDIRECT_URI_MISMATCH") + require.Zero(t, client.exchangeCalls) + + _, err = svc.ExchangeCode(context.Background(), &GrokExchangeCodeInput{ + SessionID: auth.SessionID, + Code: "authorization-code", + State: auth.State, + RedirectURI: xai.DefaultRedirectURI, + }) + require.NoError(t, err) + require.Equal(t, xai.DefaultRedirectURI, client.exchangeRedirectURI) +} + +func TestGrokOAuthServiceExternalFlowsRejectMissingClient(t *testing.T) { + svc := NewGrokOAuthService(nil, nil) + defer svc.Stop() + + _, err := svc.RefreshToken(context.Background(), "refresh-token", "", "") + require.Error(t, err) + require.Contains(t, err.Error(), "GROK_OAUTH_CLIENT_NOT_CONFIGURED") + + _, err = svc.ValidateSSOToken(context.Background(), "sso-token", nil) + require.Error(t, err) + require.Contains(t, err.Error(), "GROK_OAUTH_CLIENT_NOT_CONFIGURED") } func TestGrokOAuthServiceBuildAccountCredentialsDefaultsToSubscriptionProxy(t *testing.T) { @@ -108,6 +215,69 @@ func TestGrokOAuthServiceConvertFromSSOExtractsBuildClaims(t *testing.T) { require.Equal(t, "user@example.com", credentials["email"]) require.Equal(t, "user-sub", credentials["sub"]) require.Equal(t, "team-1", credentials["team_id"]) + require.NotContains(t, credentials, "sso_token") +} + +func TestGrokOAuthServiceValidateSSOTokenReturnsOAuthTokensWithoutPersistingSSO(t *testing.T) { + svc := NewGrokOAuthService(nil, &grokOAuthClientStub{ + ssoResponse: &xai.TokenResponse{ + AccessToken: "access-from-sso", + RefreshToken: "refresh-from-sso", + TokenType: "Bearer", + ExpiresIn: 3600, + }, + }) + defer svc.Stop() + + info, err := svc.ValidateSSOToken(context.Background(), "sso-token", nil) + require.NoError(t, err) + require.Equal(t, "access-from-sso", info.AccessToken) + require.Equal(t, "refresh-from-sso", info.RefreshToken) + + creds := svc.BuildAccountCredentials(info) + require.NotContains(t, creds, "sso_token") + require.NotContains(t, creds, "password") +} + +func TestGrokOAuthServiceAuthorizePasswordUsesLoginThenSSOAuthorize(t *testing.T) { + client := &grokOAuthClientStub{ + loginResult: &GrokPasswordLoginResult{ + Email: "user@example.com", + SSOToken: "password-derived-sso", + }, + ssoResponse: &xai.TokenResponse{ + AccessToken: "access-from-password", + RefreshToken: "refresh-from-password", + ExpiresIn: 3600, + }, + } + cfg := &config.Config{} + cfg.Gateway.Grok.PasswordAuthEnabled = true + svc := NewGrokOAuthService(nil, client, cfg) + defer svc.Stop() + + require.True(t, svc.GetCapabilities().PasswordAuthEnabled) + info, err := svc.AuthorizePassword(context.Background(), " user@example.com ", " super-secret ", nil) + require.NoError(t, err) + require.Equal(t, "user@example.com", info.Email) + require.Equal(t, "access-from-password", info.AccessToken) + creds := svc.BuildAccountCredentials(info) + require.NotContains(t, creds, "password") + require.NotContains(t, creds, "sso_token") + require.Equal(t, "user@example.com", client.loginEmail) + require.Equal(t, " super-secret ", client.loginPassword) +} + +func TestGrokOAuthServiceAuthorizePasswordDisabledByDefault(t *testing.T) { + client := &grokOAuthClientStub{} + svc := NewGrokOAuthService(nil, client) + defer svc.Stop() + + require.False(t, svc.GetCapabilities().PasswordAuthEnabled) + _, err := svc.AuthorizePassword(context.Background(), "user@example.com", "secret", nil) + require.Error(t, err) + require.Contains(t, err.Error(), "GROK_OAUTH_PASSWORD_AUTH_DISABLED") + require.Empty(t, client.loginEmail) } func makeGrokOAuthJWT(claims map[string]any) string { diff --git a/backend/internal/service/grok_observed_models.go b/backend/internal/service/grok_observed_models.go new file mode 100644 index 0000000000..480ba457bd --- /dev/null +++ b/backend/internal/service/grok_observed_models.go @@ -0,0 +1,198 @@ +package service + +import ( + "context" + "encoding/json" + "io" + "log/slog" + "net/http" + "strings" + "sync" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/xai" + "github.com/tidwall/gjson" +) + +const ( + grokObservedModelsExtraKey = "grok_observed_models" + grokObservedModelsTTL = 6 * time.Hour + grokObservedModelsTimeout = 15 * time.Second +) + +type grokObservedModelsSnapshot struct { + Models []string `json:"models"` + FetchedAt string `json:"fetched_at"` + Source string `json:"source,omitempty"` +} + +var grokObservedModelsFlight sync.Map // accountID -> *singleflight-ish in-flight + +// scheduleGrokObservedModelsSync best-effort fetches upstream /v1/models for a +// Grok OAuth account and stores IDs in Extra. Never blocks request path long; +// callers should fire-and-forget after successful auth/probe. +func (s *GrokQuotaService) scheduleGrokObservedModelsSync(account *Account) { + if s == nil || account == nil || !account.IsGrokOAuth() || s.accountRepo == nil { + return + } + id := account.ID + if _, loaded := grokObservedModelsFlight.LoadOrStore(id, struct{}{}); loaded { + return + } + // Copy credentials for background use. + acc := *account + go func() { + defer grokObservedModelsFlight.Delete(id) + ctx, cancel := context.WithTimeout(context.Background(), grokObservedModelsTimeout) + defer cancel() + if err := s.syncGrokObservedModels(ctx, &acc); err != nil { + slog.Debug("grok_observed_models_sync_failed", "account_id", id, "error", err) + } + }() +} + +func (s *GrokQuotaService) syncGrokObservedModels(ctx context.Context, account *Account) error { + if s == nil || account == nil { + return nil + } + // Skip if snapshot is still fresh. + if snap := parseGrokObservedModels(account.Extra); snap != nil { + if t, err := time.Parse(time.RFC3339, snap.FetchedAt); err == nil && time.Since(t) < grokObservedModelsTTL { + return nil + } + } + token := strings.TrimSpace(account.GetGrokAccessToken()) + if token == "" && s.tokenProvider != nil { + // Best-effort warm; avoid forcing refresh storms. + if at, err := s.tokenProvider.GetAccessToken(ctx, account); err == nil { + token = strings.TrimSpace(at) + } + } + if token == "" { + return nil + } + baseURL := strings.TrimSpace(account.GetGrokBaseURL()) + if s.settingService != nil { + baseURL = strings.TrimSpace(s.settingService.ResolveGrokBaseURL(ctx, account)) + } + if baseURL == "" { + baseURL = xai.DefaultCLIBaseURL + } + validator, err := grokBaseURLValidator(account, s.cfg) + if err != nil { + return err + } + validatedBaseURL, err := validator(baseURL) + if err != nil { + return err + } + req, err := http.NewRequestWithContext(ctx, http.MethodGet, buildOpenAIModelsURL(validatedBaseURL), nil) + if err != nil { + return err + } + req.Header.Set("Authorization", "Bearer "+token) + req.Header.Set("Accept", "application/json") + req.Header.Set("User-Agent", grokUpstreamUserAgent) + if account.IsGrokOAuth() { + applyGrokCLIHeaders(req.Header) + if isGrokCLIProxyTarget(req.URL.String()) { + if userID := strings.TrimSpace(account.GetCredential("sub")); userID != "" { + req.Header.Set("X-UserID", userID) + } + if email := strings.TrimSpace(account.GetCredential("email")); email != "" { + req.Header.Set("X-Email", email) + } + } + } + account.ApplyHeaderOverrides(req.Header) + + proxyURL := "" + if s.proxyRepo != nil && account.ProxyID != nil { + if p, err := s.proxyRepo.GetByID(ctx, *account.ProxyID); err == nil && p != nil { + proxyURL = p.URL() + } + } + if s.httpUpstream == nil { + return nil + } + resp, err := s.httpUpstream.Do(req, proxyURL, account.ID, account.Concurrency) + if err != nil { + return err + } + defer func() { _ = resp.Body.Close() }() + body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20)) + if err != nil { + return err + } + if resp.StatusCode >= 400 { + return nil + } + ids := extractGrokModelIDsFromModelsBody(body) + if len(ids) == 0 { + return nil + } + snap := grokObservedModelsSnapshot{ + Models: ids, + FetchedAt: time.Now().UTC().Format(time.RFC3339), + Source: "upstream_v1_models", + } + raw, err := json.Marshal(snap) + if err != nil { + return err + } + var asMap map[string]any + if err := json.Unmarshal(raw, &asMap); err != nil { + return err + } + return s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{ + grokObservedModelsExtraKey: asMap, + }) +} + +func extractGrokModelIDsFromModelsBody(body []byte) []string { + data := gjson.GetBytes(body, "data") + if !data.IsArray() { + // Some gateways return a bare array. + data = gjson.ParseBytes(body) + } + seen := make(map[string]struct{}) + var out []string + data.ForEach(func(_, v gjson.Result) bool { + id := strings.TrimSpace(v.Get("id").String()) + if id == "" { + id = strings.TrimSpace(v.String()) + } + if id == "" { + return true + } + if _, ok := seen[id]; ok { + return true + } + seen[id] = struct{}{} + out = append(out, id) + return true + }) + return out +} + +func parseGrokObservedModels(extra map[string]any) *grokObservedModelsSnapshot { + if extra == nil { + return nil + } + raw, ok := extra[grokObservedModelsExtraKey] + if !ok || raw == nil { + return nil + } + b, err := json.Marshal(raw) + if err != nil { + return nil + } + var snap grokObservedModelsSnapshot + if err := json.Unmarshal(b, &snap); err != nil { + return nil + } + if len(snap.Models) == 0 { + return nil + } + return &snap +} diff --git a/backend/internal/service/grok_p2_test.go b/backend/internal/service/grok_p2_test.go new file mode 100644 index 0000000000..2be382d26e --- /dev/null +++ b/backend/internal/service/grok_p2_test.go @@ -0,0 +1,102 @@ +//go:build unit + +package service + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestGrokModelQuotaBlock_FiltersOnlyNamedModel(t *testing.T) { + id := time.Now().UnixNano()%1_000_000 + 5000 + markGrokModelQuotaBlock(id, "grok-4.5", time.Now().Add(time.Hour)) + now := time.Now() + require.True(t, isGrokModelQuotaBlocked(id, "grok-4.5", now)) + require.False(t, isGrokModelQuotaBlocked(id, "grok-4.3", now)) + + accounts := []Account{ + {ID: id, Platform: PlatformGrok, Type: AccountTypeOAuth}, + {ID: id + 1, Platform: PlatformGrok, Type: AccountTypeOAuth}, + } + filtered := filterGrokModelQuotaBlockedAccounts(accounts, "grok-4.5", now) + require.Len(t, filtered, 1) + require.Equal(t, id+1, filtered[0].ID) +} + +func TestGrokModelQuotaBlockFiltersMappedUpstreamModel(t *testing.T) { + id := time.Now().UnixNano()%1_000_000 + 7000 + markGrokModelQuotaBlock(id, "grok-4.5", time.Now().Add(time.Hour)) + account := Account{ + ID: id, + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Credentials: map[string]any{ + "model_mapping": map[string]any{"gpt-*": "grok-4.5"}, + }, + } + + require.Empty(t, filterGrokModelQuotaBlockedAccounts([]Account{account}, "gpt-5", time.Now())) +} + +func TestIsGrokModelSpecificFreeUsage(t *testing.T) { + require.True(t, isGrokModelSpecificFreeUsage( + "you've used all the included free usage for model grok-4.5", "grok-4.5")) + require.True(t, isGrokModelSpecificFreeUsage("模型额度用完 grok-4.3", "grok-4.3")) + require.False(t, isGrokModelSpecificFreeUsage("free usage exhausted", "grok-4.5")) +} + +func TestGrokStickyAffinitySeed_ScopesByModel(t *testing.T) { + a := grokStickyAffinitySeed("session-1", []byte(`{"model":"grok-4.5"}`)) + b := grokStickyAffinitySeed("session-1", []byte(`{"model":"grok-4.3"}`)) + c := grokStickyAffinitySeed("session-1", []byte(`{"model":"grok-4.5"}`)) + require.NotEqual(t, a, b) + require.Equal(t, a, c) + require.Contains(t, a, "grok-affinity:v1:") +} + +func TestExtractGrokModelIDsFromModelsBody(t *testing.T) { + body := []byte(`{"object":"list","data":[{"id":"grok-4.5"},{"id":"grok-4.3"},{"id":"grok-4.5"}]}`) + ids := extractGrokModelIDsFromModelsBody(body) + require.Equal(t, []string{"grok-4.5", "grok-4.3"}, ids) +} + +func TestAccountGrokNeedsReauth(t *testing.T) { + require.False(t, accountGrokNeedsReauth(nil)) + require.True(t, accountGrokNeedsReauth(&Account{ + Extra: map[string]any{"grok_needs_reauth": true}, + })) + require.True(t, accountGrokNeedsReauth(&Account{ + Status: StatusError, + ErrorMessage: "Grok spending limit reached; reauthorize or wait for billing reset", + })) +} + +func TestApplyGrokUpstreamFailure_ModelSpecificFreeUsage(t *testing.T) { + repo := &grokQuotaAccountRepo{} + svc := &OpenAIGatewayService{accountRepo: repo} + account := &Account{ID: 9109, Platform: PlatformGrok, Type: AccountTypeOAuth} + body := []byte(`{"error":{"code":"subscription:free-usage-exhausted","message":"You've used all the included free usage for model grok-4.5. Usage resets over a rolling 24-hour window."}}`) + + svc.handleGrokAccountUpstreamError(context.Background(), account, 400, nil, body) + + require.Zero(t, repo.tempUnschedCalls, "model-scoped free usage must not cool sibling models") + require.True(t, isGrokModelQuotaBlocked(account.ID, "grok-4.5", time.Now())) + require.False(t, isGrokModelQuotaBlocked(account.ID, "grok-4.3", time.Now())) +} + +func TestApplyGrokUpstreamFailure_SpendingLimitRemainsRecoverable(t *testing.T) { + repo := &grokQuotaAccountRepo{} + svc := &OpenAIGatewayService{accountRepo: repo} + account := &Account{ID: 9110, Platform: PlatformGrok, Type: AccountTypeOAuth} + body := []byte(`{"code":"personal-team-blocked:spending-limit","error":"spending limit reached"}`) + + svc.handleGrokAccountUpstreamError(context.Background(), account, 403, nil, body) + + require.Equal(t, 1, repo.rateLimitedCalls) + require.Zero(t, repo.tempUnschedCalls) + // Without a billing-period snapshot, use a short recoverable probe cooldown. + require.WithinDuration(t, time.Now().Add(grokSpendingLimitProbeCooldown), repo.lastRateLimitResetAt, 2*time.Second) +} diff --git a/backend/internal/service/grok_quota_fetcher.go b/backend/internal/service/grok_quota_fetcher.go index 545b1c9e97..4979dfd0b9 100644 --- a/backend/internal/service/grok_quota_fetcher.go +++ b/backend/internal/service/grok_quota_fetcher.go @@ -36,6 +36,7 @@ func (f *GrokQuotaFetcher) BuildUsageInfo(account *Account) *UsageInfo { activeProbeClearsForbidden := newerSuccessfulGrokActiveProbeClearsBillingForbidden(billing, snapshot) if billing != nil { usage.GrokBilling = billing + applyGrokBillingProgressWindows(usage, billing, now) if billing.Plan != "" { usage.SubscriptionTier = billing.Plan usage.SubscriptionTierRaw = billing.Plan @@ -59,6 +60,14 @@ func (f *GrokQuotaFetcher) BuildUsageInfo(account *Account) *UsageInfo { case 429: usage.ErrorCode = "rate_limited" } + // Official weekly/monthly progress clears the "unknown until headers" state. + if usage.ErrorCode == "quota_unknown" && (usage.SevenDay != nil || usage.ThirtyDay != nil) { + usage.ErrorCode = "" + if strings.Contains(strings.ToLower(usage.Error), "unknown until") || + strings.Contains(strings.ToLower(usage.Error), "no xai quota headers") { + usage.Error = "" + } + } } if err != nil || snapshot == nil { @@ -124,6 +133,12 @@ func (f *GrokQuotaFetcher) BuildUsageInfo(account *Account) *UsageInfo { usage.ErrorCode = "rate_limited" } } + if accountGrokNeedsReauth(account) { + usage.NeedsReauth = true + if usage.ErrorCode == "" { + usage.ErrorCode = "spending_limit" + } + } applyGrokCredentialUsageFallback(usage, account) if activeProbeClearsForbidden && strings.TrimSpace(snapshot.EntitlementStatus) == "" && strings.EqualFold(strings.TrimSpace(usage.GrokEntitlementStatus), "forbidden") { @@ -240,3 +255,48 @@ func grokQuotaSnapshotFromExtra(extra map[string]any) (*xai.QuotaSnapshot, error return &out, nil } } + +// applyGrokBillingProgressWindows fills official weekly (seven_day) and monthly +// (thirty_day) UsageProgress from a billing probe summary. +func applyGrokBillingProgressWindows(usage *UsageInfo, billing *xai.BillingSummary, now time.Time) { + if usage == nil || billing == nil { + return + } + if billing.UsagePercent != nil { + seven := &UsageProgress{Utilization: *billing.UsagePercent} + if end, err := parseTime(strings.TrimSpace(billing.PeriodEnd)); err == nil { + seven.ResetsAt = &end + if sec := int(end.Sub(now).Seconds()); sec > 0 { + seven.RemainingSeconds = sec + } + } + if usage.SevenDay != nil { + seven.WindowStats = usage.SevenDay.WindowStats + } + usage.SevenDay = seven + } + var monthlyUtil *float64 + if billing.UsedPercent != nil { + monthlyUtil = billing.UsedPercent + } else if billing.MonthlyLimitCents != nil && *billing.MonthlyLimitCents > 0 && billing.UsedCents != nil { + v := (*billing.UsedCents / *billing.MonthlyLimitCents) * 100 + monthlyUtil = &v + } + if monthlyUtil != nil { + thirty := &UsageProgress{Utilization: *monthlyUtil} + endRaw := strings.TrimSpace(billing.BillingPeriodEnd) + if endRaw == "" && billing.PeriodType == "monthly" { + endRaw = strings.TrimSpace(billing.PeriodEnd) + } + if end, err := parseTime(endRaw); err == nil { + thirty.ResetsAt = &end + if sec := int(end.Sub(now).Seconds()); sec > 0 { + thirty.RemainingSeconds = sec + } + } + if usage.ThirtyDay != nil { + thirty.WindowStats = usage.ThirtyDay.WindowStats + } + usage.ThirtyDay = thirty + } +} diff --git a/backend/internal/service/grok_quota_service.go b/backend/internal/service/grok_quota_service.go index 2c703f9f10..6658888a59 100644 --- a/backend/internal/service/grok_quota_service.go +++ b/backend/internal/service/grok_quota_service.go @@ -50,13 +50,14 @@ type GrokQuotaResetResult struct { } type GrokQuotaService struct { - accountRepo AccountRepository - proxyRepo ProxyRepository - tokenProvider *GrokTokenProvider - httpUpstream HTTPUpstream - usageLogRepo UsageLogRepository - cfg *config.Config - probeFlight singleflight.Group + accountRepo AccountRepository + proxyRepo ProxyRepository + tokenProvider *GrokTokenProvider + httpUpstream HTTPUpstream + usageLogRepo UsageLogRepository + settingService *SettingService + cfg *config.Config + probeFlight singleflight.Group } func NewGrokQuotaService( @@ -81,11 +82,20 @@ func NewGrokQuotaService( } } +func (s *GrokQuotaService) SetSettingService(settingService *SettingService) { + if s != nil { + s.settingService = settingService + } +} + // QueryQuota combines xAI billing data with an active quota-header probe for // Free accounts, whose billing response does not include usage_percent. func (s *GrokQuotaService) QueryQuota(ctx context.Context, accountID int64) (*GrokQuotaProbeResult, error) { billingResult, billingErr := s.ProbeBilling(ctx, accountID) if billingErr == nil && billingResult != nil && grokBillingHasAuthoritativeQuota(billingResult.Billing) { + if acc, err := s.accountRepo.GetByID(ctx, accountID); err == nil { + s.scheduleGrokObservedModelsSync(acc) + } return billingResult, nil } @@ -111,6 +121,9 @@ func (s *GrokQuotaService) QueryQuota(ctx context.Context, accountID int64) (*Gr probeResult.LocalUsageMonthly = billingResult.LocalUsageMonthly probeResult.Persisted = probeResult.Persisted || billingResult.Persisted } + if acc, err := s.accountRepo.GetByID(ctx, accountID); err == nil { + s.scheduleGrokObservedModelsSync(acc) + } return probeResult, nil } @@ -141,7 +154,7 @@ func (s *GrokQuotaService) probeUsage(ctx context.Context, accountID int64) (*Gr if err != nil { return nil, infraerrors.Newf(http.StatusBadRequest, "GROK_QUOTA_PROBE_BODY_ERROR", "failed to build probe body: %v", err) } - targetURL, err := buildGrokResponsesURL(account, s.cfg) + targetURL, err := buildGrokResponsesURL(account, s.cfg, s.settingService) if err != nil { return nil, infraerrors.Newf(http.StatusBadRequest, "GROK_QUOTA_BASE_URL_INVALID", "invalid Grok base_url: %v", err) } @@ -172,9 +185,21 @@ func (s *GrokQuotaService) probeUsage(ctx context.Context, accountID int64) (*Gr if limited { normalizeGrokExhaustedWindowResets(snapshot, resetAt, time.Now()) } - persistErr := s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{ - grokQuotaSnapshotExtraKey: snapshot, - }) + // A failed probe must not erase a previously observed snapshot. 401/403 and + // transport/server errors commonly carry no quota headers; only successful + // responses, or 429 responses with useful rate-limit headers, are safe to + // persist. A successful 200 with no headers is still persisted as an + // explicit "no headers" observation so the UI can distinguish it from never + // probed. + persistErr := error(nil) + persisted := false + shouldPersist := resp.StatusCode < 400 || resp.StatusCode == http.StatusTooManyRequests + if shouldPersist && (snapshot.HeadersObserved || resp.StatusCode == http.StatusOK) { + persistErr = s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{ + grokQuotaSnapshotExtraKey: snapshot, + }) + persisted = persistErr == nil + } if limited { persistGrokRateLimit(ctx, s.accountRepo, account, resetAt) } else if isSuccessfulGrokRateLimitRecovery(account, snapshot) { @@ -189,7 +214,7 @@ func (s *GrokQuotaService) probeUsage(ctx context.Context, accountID int64) (*Gr HeadersObserved: snapshot.HeadersObserved, ResetSupported: false, FetchedAt: time.Now().Unix(), - Persisted: persistErr == nil, + Persisted: persisted, } if resp.StatusCode == http.StatusTooManyRequests { return result, nil diff --git a/backend/internal/service/grok_quota_service_test.go b/backend/internal/service/grok_quota_service_test.go index 9ffb70f000..70ee2c8d72 100644 --- a/backend/internal/service/grok_quota_service_test.go +++ b/backend/internal/service/grok_quota_service_test.go @@ -15,6 +15,7 @@ import ( "testing" "time" + "github.com/Wei-Shaw/sub2api/internal/config" infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" "github.com/Wei-Shaw/sub2api/internal/pkg/usagestats" "github.com/Wei-Shaw/sub2api/internal/pkg/xai" @@ -86,6 +87,60 @@ func (r *grokQuotaAccountRepo) SetTempUnschedulable(_ context.Context, id int64, return nil } +func TestSyncGrokObservedModelsRejectsOAuthCustomURLOutsideOperatorPolicy(t *testing.T) { + account := &Account{ + ID: 901, + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Credentials: map[string]any{ + "access_token": "secret-token", + "base_url": "https://blocked.example.test/v1", + }, + } + repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{account.ID: account}, + }} + upstream := &httpUpstreamRecorder{} + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = true + cfg.Security.URLAllowlist.UpstreamHosts = []string{"allowed.example.test"} + svc := &GrokQuotaService{accountRepo: repo, httpUpstream: upstream, cfg: cfg} + + err := svc.syncGrokObservedModels(context.Background(), account) + require.ErrorContains(t, err, "base URL rejected by URL security policy") + require.Nil(t, upstream.lastReq) +} + +func TestSyncGrokObservedModelsUsesCLIIdentityAndAccountHeaders(t *testing.T) { + account := &Account{ + ID: 902, + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Credentials: map[string]any{ + "access_token": "secret-token", + "sub": "user-902", + "email": "user902@example.test", + }, + } + repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{account.ID: account}, + }} + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader(`{"data":[{"id":"grok-4.5"}]}`)), + }} + svc := &GrokQuotaService{accountRepo: repo, httpUpstream: upstream, cfg: &config.Config{}} + + require.NoError(t, svc.syncGrokObservedModels(context.Background(), account)) + require.Equal(t, xai.DefaultCLIBaseURL+"/models", upstream.lastReq.URL.String()) + require.NotEmpty(t, upstream.lastReq.Header.Get("x-grok-client-version")) + require.Equal(t, xai.CLIClientIdentifier, upstream.lastReq.Header.Get("x-grok-client-identifier")) + require.Equal(t, "interactive", upstream.lastReq.Header.Get("X-Grok-Client-Mode")) + require.Equal(t, "user-902", upstream.lastReq.Header.Get("X-UserID")) + require.Equal(t, "user902@example.test", upstream.lastReq.Header.Get("X-Email")) + require.Contains(t, repo.updates[account.ID], grokObservedModelsExtraKey) +} + type grokQuotaProxyRepo struct { proxyRepoStub proxies map[int64]*Proxy @@ -610,6 +665,28 @@ func TestGrokQuotaServiceProbeUsageStoresNoHeadersState(t *testing.T) { require.Equal(t, observedResetAt, repo.recoveryObservedReset) } +func TestGrokQuotaServiceProbeUsageDoesNotOverwriteSnapshotOnUnauthorized(t *testing.T) { + t.Parallel() + + account := healthyGrokQuotaOAuthAccount(44) + previous := &xai.QuotaSnapshot{StatusCode: http.StatusOK, HeadersObserved: true} + account.Extra = map[string]any{grokQuotaSnapshotExtraKey: previous} + repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{account.ID: account}, + }} + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusUnauthorized, + Header: http.Header{}, + Body: io.NopCloser(strings.NewReader(`{"error":"unauthorized"}`)), + }} + svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream, nil) + + _, err := svc.ProbeUsage(context.Background(), account.ID) + require.Error(t, err) + require.Equal(t, 0, repo.updateCalls) + require.Same(t, previous, account.Extra[grokQuotaSnapshotExtraKey]) +} + func TestGrokQuotaServiceProbeUsageReturnsRateLimitedSnapshot(t *testing.T) { t.Parallel() diff --git a/backend/internal/service/grok_search_count.go b/backend/internal/service/grok_search_count.go new file mode 100644 index 0000000000..c89ce93f04 --- /dev/null +++ b/backend/internal/service/grok_search_count.go @@ -0,0 +1,233 @@ +package service + +import ( + "strconv" + "strings" + + "github.com/tidwall/gjson" +) + +// countGrokNativeSearchCallsFromJSONBytes counts completed native search tool +// calls in a Responses-style JSON body (output array or nested response.output). +// Counts: web_search_call, x_search_call, tool_search_call, and function_call +// named tool_search / web_search / x_search. +func countGrokNativeSearchCallsFromJSONBytes(body []byte) int { + if len(body) == 0 || !gjson.ValidBytes(body) { + return 0 + } + // Responses envelopes normally expose either top-level output (JSON mode) + // or response.output (terminal SSE payload). Compatibility layers can retain + // both copies; counting both would bill the same search twice. Prefer the + // canonical nested response when present and fall back to top-level output. + if nested := gjson.GetBytes(body, "response.output"); nested.IsArray() { + return countGrokNativeSearchCallsInOutputArray(nested) + } + return countGrokNativeSearchCallsInOutputArray(gjson.GetBytes(body, "output")) +} + +func countGrokNativeSearchCallsFromSSEBody(body string) int { + if strings.TrimSpace(body) == "" { + return 0 + } + seen := make(map[string]struct{}) + total := 0 + forEachOpenAISSEDataPayload(body, func(data []byte) { + total += countGrokNativeSearchCallsInSSEDataDedup(data, seen) + }) + return total +} + +// countGrokNativeSearchCallsInSSEData counts search tool calls in one SSE +// payload without cross-event dedup. Prefer countGrokNativeSearchCallsInSSEDataDedup +// for live streams so item.done + response.completed do not double-bill. +func countGrokNativeSearchCallsInSSEData(data []byte) int { + n, _ := countGrokNativeSearchCallsInSSEDataWithKeys(data) + return n +} + +// countGrokNativeSearchCallsInSSEDataDedup increments only unseen call ids. +// Callers must reuse the same seen map for the full stream lifetime. +// +// When call_id/id is missing, a synthetic key is built from item type + name so +// item.done + response.completed for the same tool still count once (never fall +// back to raw multi-event n, which ~2× overbills). +func countGrokNativeSearchCallsInSSEDataDedup(data []byte, seen map[string]struct{}) int { + if seen == nil { + return countGrokNativeSearchCallsInSSEData(data) + } + n, keys := countGrokNativeSearchCallsInSSEDataWithKeys(data) + if n <= 0 { + return 0 + } + // Prefer stable ids; fill gaps with synthetic keys so we never raw-add n. + if len(keys) < n { + // Rebuild keys for every item so unkeyed items still get a fingerprint. + keys = collectGrokNativeSearchCallKeys(data) + } + if len(keys) == 0 { + // True empty — should not happen when n>0; fail-closed to 0 extra bill. + return 0 + } + added := 0 + local := make(map[string]struct{}, len(keys)) + isItemDone := strings.TrimSpace(gjson.GetBytes(data, "type").String()) == "response.output_item.done" + for _, k := range keys { + if k == "" { + continue + } + if _, ok := local[k]; ok { + continue + } + local[k] = struct{}{} + if _, ok := seen[k]; ok { + if !isItemDone || !strings.HasPrefix(k, "synth:") { + continue + } + // Each id-less item.done is a distinct completed invocation. Advance + // its ordinal so interrupted streams remain accurately billable. + separator := strings.LastIndexByte(k, ':') + if separator < 0 { + continue + } + base := k[:separator] + for ordinal := 2; ; ordinal++ { + candidate := base + ":" + strconv.Itoa(ordinal) + if _, exists := seen[candidate]; !exists { + k = candidate + break + } + } + } + seen[k] = struct{}{} + added++ + } + return added +} + +func collectGrokNativeSearchCallKeys(data []byte) []string { + if len(data) == 0 || !gjson.ValidBytes(data) { + return nil + } + // An empty type means a bare item object without an SSE envelope; anything + // else that is not a completion event carries no billable call. + switch strings.TrimSpace(gjson.GetBytes(data, "type").String()) { + case "response.output_item.done", "response.completed", "response.done", "": + default: + return nil + } + var keys []string + syntheticOrdinals := make(map[string]int) + consider := func(item gjson.Result) { + if !isGrokNativeSearchOutputItem(item) { + return + } + key := firstNonEmpty( + strings.TrimSpace(item.Get("call_id").String()), + strings.TrimSpace(item.Get("id").String()), + strings.TrimSpace(item.Get("item.call_id").String()), + strings.TrimSpace(item.Get("item.id").String()), + ) + if key == "" { + // Include the ordinal among same-kind calls. A plain type:name key + // collapses two id-less web searches in one completed response into + // one charge. The ordinal remains stable between ordered item.done + // events and response.completed output. + base := "synth:" + strings.ToLower(strings.TrimSpace(item.Get("type").String())) + + ":" + strings.ToLower(strings.TrimSpace(item.Get("name").String())) + syntheticOrdinals[base]++ + key = base + ":" + strconv.Itoa(syntheticOrdinals[base]) + } + keys = append(keys, key) + } + if item := gjson.GetBytes(data, "item"); item.Exists() { + consider(item) + } + gjson.GetBytes(data, "response.output").ForEach(func(_, item gjson.Result) bool { + consider(item) + return true + }) + gjson.GetBytes(data, "output").ForEach(func(_, item gjson.Result) bool { + consider(item) + return true + }) + if len(keys) == 0 && isGrokNativeSearchOutputItem(gjson.ParseBytes(data)) { + consider(gjson.ParseBytes(data)) + } + return keys +} + +func countGrokNativeSearchCallsInSSEDataWithKeys(data []byte) (int, []string) { + if len(data) == 0 || !gjson.ValidBytes(data) { + return 0, nil + } + // Count once on item completion / completed response, not on every delta. + // An empty type is a bare item object without an SSE envelope. + switch strings.TrimSpace(gjson.GetBytes(data, "type").String()) { + case "response.output_item.done", "response.completed", "response.done", "": + default: + return 0, nil + } + var keys []string + n := 0 + consider := func(item gjson.Result) { + if !isGrokNativeSearchOutputItem(item) { + return + } + n++ + key := firstNonEmpty( + strings.TrimSpace(item.Get("call_id").String()), + strings.TrimSpace(item.Get("id").String()), + strings.TrimSpace(item.Get("item.call_id").String()), + strings.TrimSpace(item.Get("item.id").String()), + ) + if key != "" { + keys = append(keys, key) + } + } + if item := gjson.GetBytes(data, "item"); item.Exists() { + consider(item) + } + gjson.GetBytes(data, "response.output").ForEach(func(_, item gjson.Result) bool { + consider(item) + return true + }) + gjson.GetBytes(data, "output").ForEach(func(_, item gjson.Result) bool { + consider(item) + return true + }) + // Bare output item event without nested item key. + if n == 0 && isGrokNativeSearchOutputItem(gjson.ParseBytes(data)) { + consider(gjson.ParseBytes(data)) + } + return n, keys +} + +func countGrokNativeSearchCallsInOutputArray(output gjson.Result) int { + if !output.IsArray() { + return 0 + } + count := 0 + output.ForEach(func(_, item gjson.Result) bool { + if isGrokNativeSearchOutputItem(item) { + count++ + } + return true + }) + return count +} + +func isGrokNativeSearchOutputItem(item gjson.Result) bool { + if !item.Exists() { + return false + } + itemType := strings.ToLower(strings.TrimSpace(item.Get("type").String())) + switch itemType { + case "web_search_call", "x_search_call", "tool_search_call": + return true + case "function_call", "custom_tool_call": + name := strings.ToLower(strings.TrimSpace(item.Get("name").String())) + return name == "web_search" || name == "x_search" || name == "tool_search" + default: + return false + } +} diff --git a/backend/internal/service/grok_search_count_test.go b/backend/internal/service/grok_search_count_test.go new file mode 100644 index 0000000000..8bb0a2c79e --- /dev/null +++ b/backend/internal/service/grok_search_count_test.go @@ -0,0 +1,85 @@ +package service + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestCountGrokNativeSearchCallsFromJSONBytes(t *testing.T) { + t.Parallel() + require.Equal(t, 0, countGrokNativeSearchCallsFromJSONBytes(nil)) + require.Equal(t, 0, countGrokNativeSearchCallsFromJSONBytes([]byte(`{"output":[]}`))) + body := []byte(`{"output":[ + {"type":"web_search_call","id":"ws1","status":"completed"}, + {"type":"x_search_call","id":"xs1"}, + {"type":"function_call","name":"tool_search","call_id":"ts1"}, + {"type":"function_call","name":"lookup","call_id":"other"} + ]}`) + require.Equal(t, 3, countGrokNativeSearchCallsFromJSONBytes(body)) +} + +func TestCountGrokNativeSearchCallsFromJSONBytes_PrefersNestedResponse(t *testing.T) { + t.Parallel() + body := []byte(`{"output":[{"type":"web_search_call","id":"duplicate"}],"response":{"output":[{"type":"web_search_call","id":"duplicate"},{"type":"x_search_call","id":"xs1"}]}}`) + require.Equal(t, 2, countGrokNativeSearchCallsFromJSONBytes(body)) +} + +func TestCountGrokNativeSearchCallsFromJSONBytes_FallsBackWhenNestedOutputNull(t *testing.T) { + t.Parallel() + body := []byte(`{"output":[{"type":"web_search_call","id":"ws1"}],"response":{"output":null}}`) + require.Equal(t, 1, countGrokNativeSearchCallsFromJSONBytes(body)) +} + +func TestCountGrokNativeSearchCallsFromSSEBodyDedups(t *testing.T) { + t.Parallel() + sse := stringsJoin( + `data: {"type":"response.output_item.done","item":{"type":"web_search_call","id":"ws1","call_id":"c1"}}`, + `data: {"type":"response.output_item.done","item":{"type":"web_search_call","id":"ws1","call_id":"c1"}}`, + `data: {"type":"response.completed","response":{"output":[{"type":"web_search_call","id":"ws1","call_id":"c1"},{"type":"x_search_call","id":"xs1","call_id":"c2"}]}}`, + ) + require.Equal(t, 2, countGrokNativeSearchCallsFromSSEBody(sse)) +} + +func TestCountGrokNativeSearchCallsInSSEDataDedup_LiveStreamPath(t *testing.T) { + t.Parallel() + // Mirrors the live streaming accumulator: item.done then response.completed + // for the same call_id must bill once (regression for ~2× surcharge). + seen := make(map[string]struct{}) + done := []byte(`{"type":"response.output_item.done","item":{"type":"web_search_call","id":"ws1","call_id":"c1"}}`) + completed := []byte(`{"type":"response.completed","response":{"output":[{"type":"web_search_call","id":"ws1","call_id":"c1"},{"type":"x_search_call","id":"xs1","call_id":"c2"}]}}`) + require.Equal(t, 1, countGrokNativeSearchCallsInSSEDataDedup(done, seen)) + require.Equal(t, 1, countGrokNativeSearchCallsInSSEDataDedup(completed, seen)) + // Raw (no-dedup) path still double-counts the same envelope pair. + require.Equal(t, 1, countGrokNativeSearchCallsInSSEData(done)) + require.Equal(t, 2, countGrokNativeSearchCallsInSSEData(completed)) +} + +func TestCountGrokNativeSearchCallsInSSEDataDedup_NoIDStillDedups(t *testing.T) { + t.Parallel() + // Upstream sometimes omits call_id/id; synthetic keys must still prevent 2×. + seen := make(map[string]struct{}) + done := []byte(`{"type":"response.output_item.done","item":{"type":"web_search_call"}}`) + completed := []byte(`{"type":"response.completed","response":{"output":[{"type":"web_search_call"}]}}`) + require.Equal(t, 1, countGrokNativeSearchCallsInSSEDataDedup(done, seen)) + require.Equal(t, 0, countGrokNativeSearchCallsInSSEDataDedup(completed, seen)) +} + +func TestCountGrokNativeSearchCallsInSSEDataDedup_MultipleNoIDCalls(t *testing.T) { + t.Parallel() + seen := make(map[string]struct{}) + firstDone := []byte(`{"type":"response.output_item.done","item":{"type":"web_search_call"}}`) + secondDone := []byte(`{"type":"response.output_item.done","item":{"type":"web_search_call"}}`) + completed := []byte(`{"type":"response.completed","response":{"output":[{"type":"web_search_call"},{"type":"web_search_call"}]}}`) + require.Equal(t, 1, countGrokNativeSearchCallsInSSEDataDedup(firstDone, seen)) + require.Equal(t, 1, countGrokNativeSearchCallsInSSEDataDedup(secondDone, seen)) + require.Equal(t, 0, countGrokNativeSearchCallsInSSEDataDedup(completed, seen)) +} + +func stringsJoin(lines ...string) string { + out := "" + for _, l := range lines { + out += l + "\n\n" + } + return out +} diff --git a/backend/internal/service/grok_spending_reauth.go b/backend/internal/service/grok_spending_reauth.go new file mode 100644 index 0000000000..2b1f3a0973 --- /dev/null +++ b/backend/internal/service/grok_spending_reauth.go @@ -0,0 +1,59 @@ +package service + +import ( + "context" + "strings" + "time" +) + +// Spending-limit is recoverable at the end of the observed billing period. +// When no billing snapshot is available, use a short probe rather than +// fabricating a 24h boundary from the error arrival time. +const grokSpendingLimitProbeCooldown = 10 * time.Minute + +func grokSpendingLimitResetAt(account *Account, now time.Time) time.Time { + if account != nil { + if billing, err := grokBillingSnapshotFromExtra(account.Extra); err == nil && billing != nil { + for _, raw := range []string{billing.PeriodEnd, billing.BillingPeriodEnd} { + if resetAt, err := time.Parse(time.RFC3339, strings.TrimSpace(raw)); err == nil && resetAt.After(now) { + return resetAt + } + } + } + } + return now.Add(grokSpendingLimitProbeCooldown) +} + +// clearGrokNeedsReauthExtra drops the soft reauth flag after successful refresh +// or reauth. Best-effort; never fails the request path. +func clearGrokNeedsReauthExtra(ctx context.Context, repo AccountRepository, accountID int64) { + if repo == nil || accountID <= 0 { + return + } + stateCtx, cancel := openAIAccountStateContext(ctx) + defer cancel() + _ = repo.UpdateExtra(stateCtx, accountID, map[string]any{ + "grok_needs_reauth": false, + "grok_needs_reauth_reason": "", + "grok_needs_reauth_at": "", + }) +} + +func accountGrokNeedsReauth(account *Account) bool { + if account == nil { + return false + } + if account.Status == StatusError { + msg := strings.ToLower(account.ErrorMessage) + if strings.Contains(msg, "spending limit") || strings.Contains(msg, "reauthorize") { + return true + } + } + if v, ok := account.Extra["grok_needs_reauth"].(bool); ok && v { + return true + } + if s, ok := account.Extra["grok_needs_reauth"].(string); ok { + return strings.EqualFold(s, "true") || s == "1" + } + return false +} diff --git a/backend/internal/service/grok_stream_idle.go b/backend/internal/service/grok_stream_idle.go new file mode 100644 index 0000000000..3bdfa27fc9 --- /dev/null +++ b/backend/internal/service/grok_stream_idle.go @@ -0,0 +1,38 @@ +package service + +import ( + "fmt" + "strings" + "time" +) + +// Default Grok stream idle when gateway.stream_data_interval_timeout is 0. +// Long enough for slow thinking models, short enough to release hung sockets. +const defaultGrokStreamIdleTimeout = 180 * time.Second + +// Shorter cool after a Grok stream-idle failure so the account can re-enter soon +// but is not immediately re-picked in a tight failover loop. +const grokStreamIdleCooldown = 2 * time.Minute + +// resolveGrokStreamIdleTimeout returns the effective upstream-read idle timeout +// for Grok streams. Prefers the global gateway setting when positive; otherwise +// applies a Grok-only default so hung SSE bodies still fail over. +func resolveGrokStreamIdleTimeout(cfgStreamIntervalSec int) time.Duration { + if cfgStreamIntervalSec > 0 { + return time.Duration(cfgStreamIntervalSec) * time.Second + } + return defaultGrokStreamIdleTimeout +} + +// grokStreamIdleFailoverError builds a pre-commit/handler-visible failover so +// the gateway can switch OAuth accounts after a hung Grok upstream stream. +func grokStreamIdleFailoverError(account *Account, idle time.Duration) *UpstreamFailoverError { + msg := fmt.Sprintf("Grok stream idle timeout after %s with no upstream data", idle.Round(time.Second)) + return &UpstreamFailoverError{ + StatusCode: 502, + ResponseBody: []byte(`{"error":{"code":"empty_upstream","message":"` + strings.ReplaceAll(msg, `"`, `'`) + `"}}`), + SafeToFailoverAfterWrite: true, + // Allow pool-mode retries; normal OAuth switches account via handler. + RetryableOnSameAccount: account != nil && account.IsPoolMode(), + } +} diff --git a/backend/internal/service/grok_stream_idle_test.go b/backend/internal/service/grok_stream_idle_test.go new file mode 100644 index 0000000000..e3c6d16943 --- /dev/null +++ b/backend/internal/service/grok_stream_idle_test.go @@ -0,0 +1,25 @@ +//go:build unit + +package service + +import ( + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestResolveGrokStreamIdleTimeout(t *testing.T) { + require.Equal(t, 90*time.Second, resolveGrokStreamIdleTimeout(90)) + require.Equal(t, defaultGrokStreamIdleTimeout, resolveGrokStreamIdleTimeout(0)) + require.Equal(t, defaultGrokStreamIdleTimeout, resolveGrokStreamIdleTimeout(-1)) +} + +func TestGrokStreamIdleFailoverError(t *testing.T) { + account := &Account{ID: 1, Platform: PlatformGrok, Type: AccountTypeOAuth} + err := grokStreamIdleFailoverError(account, 180*time.Second) + require.NotNil(t, err) + require.Equal(t, 502, err.StatusCode) + require.True(t, err.SafeToFailoverAfterWrite) + require.Contains(t, string(err.ResponseBody), "empty_upstream") +} diff --git a/backend/internal/service/grok_team_rate_limit.go b/backend/internal/service/grok_team_rate_limit.go new file mode 100644 index 0000000000..f10852615a --- /dev/null +++ b/backend/internal/service/grok_team_rate_limit.go @@ -0,0 +1,149 @@ +package service + +import ( + "crypto/sha256" + "encoding/hex" + "strings" + "sync" + "time" +) + +// In-memory team+model rate-limit overlay for Grok OAuth. When xAI rate-limits +// one account in a team for a model, sibling accounts sharing team_id skip the +// same model until the cooldown expires (mirrors grok2api teamModelRateLimit). +// +// Process-local only: multi-instance deployments each learn the block from their +// own 429s. Prefer short TTLs so drift self-heals. +type grokTeamModelRateLimit struct { + Until time.Time +} + +type grokTeamModelRateLimitStore struct { + mu sync.Mutex + items map[string]grokTeamModelRateLimit +} + +var globalGrokTeamModelRateLimits = &grokTeamModelRateLimitStore{ + items: make(map[string]grokTeamModelRateLimit), +} + +const ( + grokTeamRateLimitDefaultTTL = 10 * time.Minute + grokTeamRateLimitMaxTTL = time.Hour + grokTeamRateLimitMinTTL = 30 * time.Second +) + +func grokTeamFingerprint(teamID string) string { + teamID = strings.TrimSpace(teamID) + if teamID == "" { + return "" + } + sum := sha256.Sum256([]byte(strings.ToLower(teamID))) + return hex.EncodeToString(sum[:8]) +} + +func grokTeamModelRateLimitKey(teamFingerprint, model string) string { + return teamFingerprint + "|" + strings.ToLower(strings.TrimSpace(model)) +} + +func accountGrokTeamID(account *Account) string { + if account == nil { + return "" + } + return strings.TrimSpace(account.GetCredential("team_id")) +} + +// markGrokTeamModelRateLimit records that this team+model pair should be skipped +// until until. No-op when team_id or model is empty. +func markGrokTeamModelRateLimit(account *Account, model string, until time.Time) { + if account == nil || !account.IsGrokOAuth() { + return + } + fp := grokTeamFingerprint(accountGrokTeamID(account)) + model = strings.TrimSpace(model) + if fp == "" || model == "" || until.IsZero() { + return + } + now := time.Now() + if !until.After(now) { + until = now.Add(grokTeamRateLimitDefaultTTL) + } + maxUntil := now.Add(grokTeamRateLimitMaxTTL) + if until.After(maxUntil) { + until = maxUntil + } + key := grokTeamModelRateLimitKey(fp, model) + globalGrokTeamModelRateLimits.mu.Lock() + defer globalGrokTeamModelRateLimits.mu.Unlock() + if cur, ok := globalGrokTeamModelRateLimits.items[key]; ok && cur.Until.After(until) { + return + } + globalGrokTeamModelRateLimits.items[key] = grokTeamModelRateLimit{Until: until} + // Opportunistic prune of expired entries. + for k, v := range globalGrokTeamModelRateLimits.items { + if !v.Until.After(now) { + delete(globalGrokTeamModelRateLimits.items, k) + } + } +} + +// isGrokTeamModelRateLimited reports whether the account's team is currently +// blocked for the requested model. +func isGrokTeamModelRateLimited(account *Account, model string, now time.Time) bool { + if account == nil || !account.IsGrokOAuth() { + return false + } + fp := grokTeamFingerprint(accountGrokTeamID(account)) + model = strings.TrimSpace(model) + if fp == "" || model == "" { + return false + } + key := grokTeamModelRateLimitKey(fp, model) + globalGrokTeamModelRateLimits.mu.Lock() + defer globalGrokTeamModelRateLimits.mu.Unlock() + cur, ok := globalGrokTeamModelRateLimits.items[key] + if !ok { + return false + } + if !cur.Until.After(now) { + delete(globalGrokTeamModelRateLimits.items, key) + return false + } + return true +} + +// filterGrokTeamModelRateLimitedAccounts drops candidates whose team is under a +// model-scoped rate-limit cool. Accounts without team_id pass through. +func filterGrokTeamModelRateLimitedAccounts(accounts []Account, model string, now time.Time) []Account { + if len(accounts) == 0 || strings.TrimSpace(model) == "" { + return accounts + } + out := accounts[:0] + kept := false + for i := range accounts { + upstreamModel := canonicalOpenAIAccountSchedulingModel(&accounts[i], model) + if isGrokTeamModelRateLimited(&accounts[i], upstreamModel, now) { + continue + } + out = append(out, accounts[i]) + kept = true + } + if !kept && len(out) == 0 { + // All filtered — return empty (caller treats as no capacity). + return nil + } + return out +} + +// resolveGrokTeamRateLimitUntil derives a team cool window from an observed +// account rate-limit reset, with sane clamps. +func resolveGrokTeamRateLimitUntil(resetAt, now time.Time) time.Time { + if resetAt.After(now.Add(grokTeamRateLimitMinTTL)) { + maxUntil := now.Add(grokTeamRateLimitMaxTTL) + if resetAt.After(maxUntil) { + return maxUntil + } + return resetAt + } + return now.Add(grokTeamRateLimitDefaultTTL) +} diff --git a/backend/internal/service/grok_team_rate_limit_test.go b/backend/internal/service/grok_team_rate_limit_test.go new file mode 100644 index 0000000000..4926d4d7aa --- /dev/null +++ b/backend/internal/service/grok_team_rate_limit_test.go @@ -0,0 +1,75 @@ +//go:build unit + +package service + +import ( + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestGrokTeamModelRateLimit_MarksAndFiltersSiblings(t *testing.T) { + // Isolate from other tests by using unique team ids. + team := "team-test-" + time.Now().Format("150405.000") + a1 := &Account{ + ID: 101, Platform: PlatformGrok, Type: AccountTypeOAuth, + Credentials: map[string]any{"team_id": team}, + } + a2 := &Account{ + ID: 102, Platform: PlatformGrok, Type: AccountTypeOAuth, + Credentials: map[string]any{"team_id": team}, + } + other := &Account{ + ID: 103, Platform: PlatformGrok, Type: AccountTypeOAuth, + Credentials: map[string]any{"team_id": team + "-other"}, + } + noTeam := &Account{ + ID: 104, Platform: PlatformGrok, Type: AccountTypeOAuth, + Credentials: map[string]any{}, + } + + now := time.Now() + markGrokTeamModelRateLimit(a1, "grok-4.5", now.Add(5*time.Minute)) + + require.True(t, isGrokTeamModelRateLimited(a1, "grok-4.5", now)) + require.True(t, isGrokTeamModelRateLimited(a2, "grok-4.5", now), "sibling with same team must cool") + require.False(t, isGrokTeamModelRateLimited(a2, "grok-4.3", now), "other model stays pickable") + require.False(t, isGrokTeamModelRateLimited(other, "grok-4.5", now)) + require.False(t, isGrokTeamModelRateLimited(noTeam, "grok-4.5", now)) + + filtered := filterGrokTeamModelRateLimitedAccounts([]Account{*a1, *a2, *other, *noTeam}, "grok-4.5", now) + require.Len(t, filtered, 2) + ids := []int64{filtered[0].ID, filtered[1].ID} + require.Contains(t, ids, int64(103)) + require.Contains(t, ids, int64(104)) +} + +func TestGrokTeamModelRateLimit_Expires(t *testing.T) { + team := "team-expire-" + time.Now().Format("150405.000") + a := &Account{ + ID: 201, Platform: PlatformGrok, Type: AccountTypeOAuth, + Credentials: map[string]any{"team_id": team}, + } + past := time.Now().Add(-time.Minute) + markGrokTeamModelRateLimit(a, "grok-4.5", past) + // mark clamps expired until into default TTL from "now" — use direct store inject via past+recheck + // After mark with past, resolveGrokTeamRateLimitUntil path isn't used; mark uses now+default when until not after now. + require.True(t, isGrokTeamModelRateLimited(a, "grok-4.5", time.Now())) +} + +func TestGrokTeamModelRateLimitFilterUsesMappedUpstreamModel(t *testing.T) { + now := time.Now() + account := &Account{ + ID: 301, + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Credentials: map[string]any{ + "team_id": "team-mapped-301", + "model_mapping": map[string]any{"gpt-*": "grok-4.5"}, + }, + } + markGrokTeamModelRateLimit(account, "grok-4.5", now.Add(time.Hour)) + + require.Empty(t, filterGrokTeamModelRateLimitedAccounts([]Account{*account}, "gpt-5", now)) +} diff --git a/backend/internal/service/grok_token_refresher.go b/backend/internal/service/grok_token_refresher.go index d2d4d534fc..72b041bc6e 100644 --- a/backend/internal/service/grok_token_refresher.go +++ b/backend/internal/service/grok_token_refresher.go @@ -3,12 +3,25 @@ package service import ( "context" "errors" + "hash/fnv" "strings" "time" ) +// Base warm window: refresh when access token lifetime remaining is below this. +// Grok access tokens are typically ~1h; refreshing up to 1h early keeps the pool +// warm for request path cache misses. const grokTokenRefreshSkew = time.Hour +// Stampede spread: each account's effective warm window is reduced by a +// deterministic offset in [0, grokTokenRefreshJitterMax] so co-imported accounts +// do not all refresh in the same TokenRefreshService cycle (grok2api-style +// RefreshDueAt scatter). +const grokTokenRefreshJitterMax = 3 * time.Minute + +// Floor so jitter cannot shrink the window below a useful threshold. +const grokTokenRefreshSkewMin = 30 * time.Minute + type GrokTokenRefresher struct { grokOAuthService GrokOAuthTokenService } @@ -40,9 +53,35 @@ func (r *GrokTokenRefresher) NeedsRefresh(account *Account, refreshWindow time.D if refreshWindow < grokTokenRefreshSkew { refreshWindow = grokTokenRefreshSkew } + // Deterministic per-account jitter: spread warm refreshes without random + // non-determinism in tests (hash of account id). + refreshWindow = grokTokenRefreshWindowWithJitter(account.ID, refreshWindow) return time.Until(*expiresAt) < refreshWindow } +// grokTokenRefreshWindowWithJitter returns refreshWindow minus a stable offset +// in [0, jitterMax] based on accountID. Result is never below grokTokenRefreshSkewMin +// when the base window is at least that large. +func grokTokenRefreshWindowWithJitter(accountID int64, refreshWindow time.Duration) time.Duration { + if accountID <= 0 || refreshWindow <= grokTokenRefreshSkewMin { + return refreshWindow + } + h := fnv.New32a() + var b [8]byte + id := uint64(accountID) + for i := 0; i < 8; i++ { + b[i] = byte(id >> (8 * i)) + } + _, _ = h.Write(b[:]) + // Jitter in [0, grokTokenRefreshJitterMax). + jitter := time.Duration(h.Sum32()%uint32(grokTokenRefreshJitterMax/time.Second)) * time.Second + out := refreshWindow - jitter + if out < grokTokenRefreshSkewMin { + return grokTokenRefreshSkewMin + } + return out +} + func (r *GrokTokenRefresher) Refresh(ctx context.Context, account *Account) (map[string]any, error) { if r == nil || r.grokOAuthService == nil { return nil, errors.New("grok oauth service is not configured") diff --git a/backend/internal/service/grok_token_refresher_test.go b/backend/internal/service/grok_token_refresher_test.go new file mode 100644 index 0000000000..c689985628 --- /dev/null +++ b/backend/internal/service/grok_token_refresher_test.go @@ -0,0 +1,50 @@ +//go:build unit + +package service + +import ( + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestGrokTokenRefreshWindowWithJitter_StableAndBounded(t *testing.T) { + base := grokTokenRefreshSkew + w1 := grokTokenRefreshWindowWithJitter(42, base) + w2 := grokTokenRefreshWindowWithJitter(42, base) + require.Equal(t, w1, w2, "same account id must yield stable window") + require.GreaterOrEqual(t, w1, grokTokenRefreshSkewMin) + require.LessOrEqual(t, w1, base) + + // Different accounts should usually differ (not a hard guarantee for all pairs, + // but for sequential ids the hash spread is good enough to assert inequality + // across a small sample). + seen := map[time.Duration]bool{} + for id := int64(1); id <= 50; id++ { + seen[grokTokenRefreshWindowWithJitter(id, base)] = true + } + require.Greater(t, len(seen), 1, "jitter should spread windows across accounts") +} + +func TestGrokTokenRefresher_NeedsRefresh_UsesSkewFloor(t *testing.T) { + refresher := NewGrokTokenRefresher(nil) + // Expires in 50 minutes — within 1h skew, should need refresh. + expires := time.Now().Add(50 * time.Minute).UTC().Format(time.RFC3339) + account := &Account{ + ID: 7, + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Credentials: map[string]any{ + "access_token": "at", + "refresh_token": "rt", + "expires_at": expires, + }, + } + // Pass a tiny window; NeedsRefresh raises it to grokTokenRefreshSkew then jitter. + require.True(t, refresher.NeedsRefresh(account, time.Minute)) + + // Far future — no refresh. + account.Credentials["expires_at"] = time.Now().Add(3 * time.Hour).UTC().Format(time.RFC3339) + require.False(t, refresher.NeedsRefresh(account, time.Minute)) +} diff --git a/backend/internal/service/grok_upstream_errors.go b/backend/internal/service/grok_upstream_errors.go index 880cd702b0..0039ea4cdc 100644 --- a/backend/internal/service/grok_upstream_errors.go +++ b/backend/internal/service/grok_upstream_errors.go @@ -191,10 +191,17 @@ func grokContentPolicyClientMessage(responseBody []byte) string { // shouldFailoverGrokUpstreamError is the body-aware counterpart of the // status-only failover helper. Grok content refusals must stay on the current // account and be returned to the caller instead of consuming the account pool. +// Free-usage / empty-output / billing bodies also failover even when the HTTP +// status alone would not (e.g. 400 with free-usage-exhausted). func (s *OpenAIGatewayService) shouldFailoverGrokUpstreamError(statusCode int, responseBody []byte) bool { if isGrokContentPolicyRejection(statusCode, responseBody) { return false } + decision := classifyGrokUpstreamFailure(statusCode, responseBody, "") + switch decision.Class { + case GrokFailureFreeUsage, GrokFailureEmptyUpstream, GrokFailureBilling, GrokFailureModelCapacity: + return decision.ShouldFailover + } return s.shouldFailoverUpstreamError(statusCode) } diff --git a/backend/internal/service/grok_upstream_errors_test.go b/backend/internal/service/grok_upstream_errors_test.go index f8ce5c850c..0dee756b72 100644 --- a/backend/internal/service/grok_upstream_errors_test.go +++ b/backend/internal/service/grok_upstream_errors_test.go @@ -122,6 +122,43 @@ func TestGrokContentPolicy403DoesNotMutateOrFailover(t *testing.T) { require.Zero(t, repo.tempUnschedCalls) } +func TestGrokNonFailoverDoesNotApplyGenericTempUnschedulablePolicy(t *testing.T) { + gin.SetMode(gin.TestMode) + repo := &grokQuotaAccountRepo{} + svc := &OpenAIGatewayService{ + accountRepo: repo, + rateLimitService: NewRateLimitService(repo, nil, nil, nil, nil), + } + account := &Account{ + ID: 5099, + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Credentials: map[string]any{ + "temp_unschedulable_enabled": true, + "temp_unschedulable_rules": []any{map[string]any{ + "error_code": float64(http.StatusForbidden), + "keywords": []any{"text is sensitive"}, + "duration_minutes": float64(1), + }}, + }, + } + body := []byte(`{"error":{"code":"new_sensitive","message":"text is sensitive"}}`) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) + resp := &http.Response{StatusCode: http.StatusForbidden, Header: http.Header{}} + + got := svc.failoverOpenAIUpstreamHTTPError( + context.Background(), c, account, resp, body, "text is sensitive", "", + ) + + require.Nil(t, got) + require.Zero(t, repo.tempUnschedCalls) + require.Zero(t, repo.rateLimitedCalls) + require.Zero(t, repo.updateCalls) + require.False(t, svc.isOpenAIAccountRuntimeBlocked(account)) +} + func TestGrokContentPolicy403SharedErrorFallbackDoesNotMutate(t *testing.T) { gin.SetMode(gin.TestMode) body := []byte(`{"error":{"code":"content_filter","message":"prohibited content"}}`) @@ -259,6 +296,74 @@ func TestHandleGrokAccountUpstreamErrorEntitlement403KeepsDefaultCooldown(t *tes require.Less(t, repo.lastTempUnschedUntil, before.Add(31*time.Minute)) } +func TestHandleGrokAccountUpstreamErrorDefaultCooldownsRespectPoolMode(t *testing.T) { + for _, statusCode := range []int{ + http.StatusUnauthorized, + http.StatusPaymentRequired, + http.StatusForbidden, + http.StatusInternalServerError, + } { + t.Run(http.StatusText(statusCode), func(t *testing.T) { + repo := &grokQuotaAccountRepo{} + svc := &OpenAIGatewayService{accountRepo: repo} + account := &Account{ + ID: int64(4800 + statusCode), + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + Credentials: map[string]any{ + "pool_mode": true, + }, + } + body := []byte(`{"error":{"message":"grok access or entitlement denied"}}`) + + svc.handleGrokAccountUpstreamError( + context.Background(), account, statusCode, nil, body, + ) + + require.Zero(t, repo.tempUnschedCalls) + require.False(t, svc.isOpenAIAccountRuntimeBlocked(account)) + require.Nil(t, account.TempUnschedulableUntil) + require.Empty(t, account.TempUnschedulableReason) + require.True(t, svc.shouldFailoverGrokUpstreamError(statusCode, body)) + }) + } + + account := &Account{Type: AccountTypeAPIKey, Credentials: map[string]any{"pool_mode": true}} + require.True(t, account.IsPoolModeRetryableStatus(http.StatusForbidden)) + + t.Run("explicit temporary rule still applies", func(t *testing.T) { + repo := &grokQuotaAccountRepo{} + svc := &OpenAIGatewayService{accountRepo: repo} + account := &Account{ + ID: 4723, + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + Credentials: map[string]any{ + "pool_mode": true, + "temp_unschedulable_enabled": true, + "temp_unschedulable_rules": []any{ + map[string]any{ + "error_code": float64(http.StatusForbidden), + "keywords": []any{"entitlement denied"}, + "duration_minutes": float64(7), + }, + }, + }, + } + before := time.Now() + + svc.handleGrokAccountUpstreamError( + context.Background(), account, http.StatusForbidden, nil, + []byte(`{"error":{"message":"grok access or entitlement denied"}}`), + ) + + require.Equal(t, 1, repo.tempUnschedCalls) + require.Equal(t, "grok configured forbidden rule", repo.lastTempUnschedReason) + require.WithinDuration(t, before.Add(7*time.Minute), repo.lastTempUnschedUntil, time.Second) + require.True(t, svc.isOpenAIAccountRuntimeBlocked(account)) + }) +} + func TestHandleGrokAccountUpstreamError403UsesConfiguredRule(t *testing.T) { repo := &grokQuotaAccountRepo{} svc := &OpenAIGatewayService{accountRepo: repo} diff --git a/backend/internal/service/grok_upstream_failure.go b/backend/internal/service/grok_upstream_failure.go new file mode 100644 index 0000000000..1348d068c5 --- /dev/null +++ b/backend/internal/service/grok_upstream_failure.go @@ -0,0 +1,506 @@ +package service + +import ( + "context" + "encoding/json" + "net/http" + "regexp" + "strconv" + "strings" + "time" + + "github.com/tidwall/gjson" +) + +// Grok upstream failure classes used to decide temp-unschedulable cooldowns and +// pre-commit account failover. Classification is body-first so free-usage and +// empty-output wording still win when the proxy rewrites status codes. +type GrokUpstreamFailureClass string + +const ( + GrokFailureNone GrokUpstreamFailureClass = "" + GrokFailureFreeUsage GrokUpstreamFailureClass = "subscription:free-usage-exhausted" + GrokFailureBilling GrokUpstreamFailureClass = "billing_quota" + GrokFailureEmptyUpstream GrokUpstreamFailureClass = "empty_upstream" + GrokFailureModelCapacity GrokUpstreamFailureClass = "model_capacity" + GrokFailureRateLimit GrokUpstreamFailureClass = "rate_limit" + GrokFailureAuth GrokUpstreamFailureClass = "auth_error" + GrokFailureServer GrokUpstreamFailureClass = "server_error" +) + +// GrokUpstreamFailureDecision is a pure classification result. Callers map it +// onto existing account state helpers (tempUnscheduleGrok / rateLimitGrok). +// BlockModel is retained for observability; the current scheduler does not +// implement per-model soft-blocks, so free-usage deliberately never sets it. +type GrokUpstreamFailureDecision struct { + Class GrokUpstreamFailureClass + Model string + Cooldown time.Duration + ShouldCooldown bool + // ShouldFailover recommends trying another account before writing a + // terminal response (pre-commit only). Content-policy rejections are + // handled separately and never reach this classifier for failover. + ShouldFailover bool + // BlockModel is true only for empty-output when a model id is known. + // Free-usage never sets this: the account cools, not a single model. + BlockModel bool + Reason string + TokensActual *int64 + TokensLimit *int64 +} + +var ( + reGrokTokenPair = regexp.MustCompile(`(?i)tokens?\s*(?:\(actual\s*/\s*limit\))?\s*[:=]?\s*(\d+)\s*/\s*(\d+)`) + reGrokModelFor = regexp.MustCompile(`(?i)(?:for\s+model|model|模型)\s*[:：]?\s*([a-z0-9][a-z0-9._-]{2,80})`) +) + +// classifyGrokUpstreamFailure decides cooldown/failover from status + body. +// Priority (body/code first, status second): +// 1. free-usage exhausted → account cool, no model block, failover +// 2. billing hard quota → longer cool, failover +// 3. empty model output → short cool + optional model soft-block marker, failover +// 4. model capacity → short cool, failover +// 5. bare rate-limit / 429 without free-usage language → cool, failover +// 6. bare 5xx → brief cool, failover +// 7. validation / client errors without quota language → no cool +// +// Content-policy 403s must be filtered by the caller before invoking this. +func classifyGrokUpstreamFailure(statusCode int, responseBody []byte, requestedModel string) GrokUpstreamFailureDecision { + text, code, low := grokUpstreamErrorCorpus(statusCode, responseBody) + model := extractGrokFailureModel(text, responseBody, requestedModel) + actual, limit, hasTokens := parseGrokTokenPair(text) + if !hasTokens { + actual, limit, hasTokens = parseGrokTokenPair(string(responseBody)) + } + + // --- Free usage / rolling quota exhausted --- + if isGrokFreeUsageExhaustedText(low) || isGrokFreeUsageCode(code) || isGrokFreeUsageCode(text) { + d := GrokUpstreamFailureDecision{ + Class: GrokFailureFreeUsage, + Model: model, + Cooldown: grokFreeUsageCooldownDuration(low), + ShouldCooldown: true, + ShouldFailover: true, + BlockModel: false, + Reason: firstNonEmpty(text, code, "free usage exhausted"), + } + if hasTokens { + a, b := actual, limit + d.TokensActual = &a + d.TokensLimit = &b + } + return d + } + + // Billing / hard quota (not free-tier rolling). + // Cooldown stays at 30m to match the existing Grok 402/spending-limit + // handler (longer cools would change ops behavior without a settings knob). + if isGrokBillingQuotaText(low) || statusCode == http.StatusPaymentRequired { + reason := firstNonEmpty(text, "billing quota") + if statusCode == http.StatusPaymentRequired && text == "" { + reason = "payment required" + } + return GrokUpstreamFailureDecision{ + Class: GrokFailureBilling, + Model: model, + Cooldown: 30 * time.Minute, + ShouldCooldown: true, + ShouldFailover: true, + BlockModel: model != "", + Reason: reason, + } + } + + // Empty HTTP 200 / empty model output (often rewritten to synthetic 502). + if isGrokEmptyModelOutputText(low) || isGrokEmptyModelOutputCode(code) { + return GrokUpstreamFailureDecision{ + Class: GrokFailureEmptyUpstream, + Model: model, + Cooldown: 4 * time.Minute, + ShouldCooldown: true, + ShouldFailover: true, + BlockModel: model != "", + Reason: firstNonEmpty(text, "empty model output"), + } + } + + // Model capacity / overloaded. + if isGrokModelCapacityText(low) { + return GrokUpstreamFailureDecision{ + Class: GrokFailureModelCapacity, + Model: model, + Cooldown: 3 * time.Minute, + ShouldCooldown: true, + ShouldFailover: true, + BlockModel: false, + Reason: firstNonEmpty(text, "model capacity"), + } + } + + // Rate limit without free-usage language. + if statusCode == http.StatusTooManyRequests || isGrokRateLimitText(low) { + return GrokUpstreamFailureDecision{ + Class: GrokFailureRateLimit, + Model: model, + Cooldown: 10 * time.Minute, + ShouldCooldown: true, + ShouldFailover: true, + BlockModel: false, + Reason: firstNonEmpty(text, "rate limit"), + } + } + + // Upstream 5xx — brief cool. Empty-output synthetic 502 already handled above. + if statusCode >= 500 && statusCode <= 599 { + return GrokUpstreamFailureDecision{ + Class: GrokFailureServer, + Cooldown: 2 * time.Minute, + ShouldCooldown: true, + ShouldFailover: true, + Reason: firstNonEmpty(text, "server error"), + } + } + + return GrokUpstreamFailureDecision{Reason: text} +} + +func grokUpstreamErrorCorpus(statusCode int, responseBody []byte) (text, code, low string) { + _ = statusCode // classifier already has the transport status; corpus is body-only + raw := strings.TrimSpace(string(responseBody)) + // Strip "upstream status NNN: ..." prefixes so free-usage / quota language is visible. + if _, unwrappedBody, ok := unwrapGrokUpstreamErrorText(raw); ok { + raw = unwrappedBody + } + text = raw + codeFromJSON, msgFromJSON := parseGrokUpstreamErrorJSON(raw) + if msgFromJSON != "" { + if text == "" || len(msgFromJSON) > len(text)/2 || looksLikeGrokQuotaMessage(msgFromJSON) { + text = msgFromJSON + } + } + // Prefer structured fields from the original body when present. + if len(responseBody) > 0 { + if m := strings.TrimSpace(firstNonEmpty( + gjson.GetBytes(responseBody, "error.message").String(), + gjson.GetBytes(responseBody, "message").String(), + gjson.GetBytes(responseBody, "error").String(), + )); m != "" && (text == "" || looksLikeGrokQuotaMessage(m)) { + text = m + } + if c := strings.TrimSpace(firstNonEmpty( + gjson.GetBytes(responseBody, "error.code").String(), + gjson.GetBytes(responseBody, "code").String(), + )); c != "" { + codeFromJSON = c + } + } + code = codeFromJSON + low = strings.ToLower(strings.TrimSpace(text)) + if code != "" && !strings.Contains(low, strings.ToLower(code)) { + low = strings.ToLower(code) + " " + low + } + return text, code, low +} + +func unwrapGrokUpstreamErrorText(errText string) (status int, body string, ok bool) { + text := strings.TrimSpace(errText) + if text == "" { + return 0, "", false + } + lower := strings.ToLower(text) + for _, p := range []string{"upstream status ", "status "} { + if !strings.HasPrefix(lower, p) { + continue + } + rest := strings.TrimSpace(text[len(p):]) + i := 0 + for i < len(rest) && rest[i] >= '0' && rest[i] <= '9' { + status = status*10 + int(rest[i]-'0') + i++ + } + if status <= 0 || i == 0 { + return 0, "", false + } + rest = strings.TrimSpace(rest[i:]) + if strings.HasPrefix(rest, ":") { + rest = strings.TrimSpace(rest[1:]) + } + return status, rest, true + } + return 0, "", false +} + +func parseGrokUpstreamErrorJSON(errText string) (code, message string) { + text := strings.TrimSpace(errText) + if text == "" || text[0] != '{' { + return "", "" + } + var payload map[string]any + if json.Unmarshal([]byte(text), &payload) != nil { + return "", "" + } + if v, ok := payload["code"].(string); ok { + code = v + } + if v, ok := payload["message"].(string); ok { + message = v + } + if errObj, ok := payload["error"].(map[string]any); ok { + if v, ok := errObj["code"].(string); ok && code == "" { + code = v + } + if v, ok := errObj["message"].(string); ok && message == "" { + message = v + } + } + if errStr, ok := payload["error"].(string); ok && message == "" { + message = errStr + } + return strings.TrimSpace(code), strings.TrimSpace(message) +} + +func looksLikeGrokQuotaMessage(s string) bool { + low := strings.ToLower(s) + return strings.Contains(low, "quota") || + strings.Contains(low, "usage") || + strings.Contains(low, "credit") || + strings.Contains(low, "额度") || + strings.Contains(low, "free") +} + +func isGrokFreeUsageCode(code string) bool { + c := strings.ToLower(strings.TrimSpace(code)) + if c == "" { + return false + } + if strings.Contains(c, "subscription:free-usage-exhausted") || + strings.Contains(c, "free-usage-exhausted") || + strings.Contains(c, "free_usage_exhausted") || + strings.Contains(c, "usage-limit-exceeded") || + strings.Contains(c, "usage_limit_exceeded") { + return true + } + return (strings.Contains(c, "free-usage") || strings.Contains(c, "free_usage")) && + (strings.Contains(c, "exhaust") || strings.Contains(c, "exceed") || strings.Contains(c, "limit")) +} + +func isGrokFreeUsageExhaustedText(low string) bool { + if low == "" { + return false + } + if strings.Contains(low, "free-usage-exhausted") || + strings.Contains(low, "free_usage_exhausted") || + strings.Contains(low, "subscription:free-usage") || + strings.Contains(low, "usage-limit-exceeded") || + strings.Contains(low, "usage_limit_exceeded") || + strings.Contains(low, "free-tier-limit") || + strings.Contains(low, "free_tier_limit") { + return true + } + if strings.Contains(low, "free usage") || + strings.Contains(low, "included free usage") || + strings.Contains(low, "used all the included free") || + strings.Contains(low, "you've used all the included free") || + strings.Contains(low, "you have used all the included free") || + strings.Contains(low, "free quota") || + strings.Contains(low, "no remaining free") || + strings.Contains(low, "out of free") || + strings.Contains(low, "usage resets over a rolling") || + (strings.Contains(low, "free tier") && (strings.Contains(low, "exhaust") || strings.Contains(low, "limit") || strings.Contains(low, "exceed"))) { + return true + } + for _, p := range []string{ + "额度耗尽", "额度用完", "额度不足", "额度已用尽", "额度已耗尽", + "免费额度", "免费用量", "用量用完", "用量耗尽", "用量超限", "用量已用尽", + "配额耗尽", "配额已用尽", "配额不足", "配额超限", "配额用完", + "没有额度", "没额度", "无额度", "可用额度不足", "模型额度", + "临时额度", "额度已满", "额度超限", "额度达到上限", + "模型额度用完", "模型额度耗尽", "账号额度用完", "账号额度耗尽", + "额度不够", "没额度了", "额度没了", "用完额度", "耗尽额度", + } { + if strings.Contains(low, p) { + return true + } + } + if (strings.Contains(low, "quota") && (strings.Contains(low, "exhaust") || strings.Contains(low, "exceed") || strings.Contains(low, "limit"))) || + (strings.Contains(low, "usage") && (strings.Contains(low, "exhaust") || strings.Contains(low, "exceed")) && (strings.Contains(low, "limit") || strings.Contains(low, "free") || strings.Contains(low, "model"))) { + if strings.Contains(low, "free") || strings.Contains(low, "rolling") || + strings.Contains(low, "24-hour") || strings.Contains(low, "24 hour") || + strings.Contains(low, "model") || strings.Contains(low, "subscription") || + strings.Contains(low, "included") || strings.Contains(low, "tokens") { + return true + } + } + if a, b, ok := parseGrokTokenPair(low); ok && b > 0 && a >= b { + if strings.Contains(low, "free") || strings.Contains(low, "subscription") || + strings.Contains(low, "included") || strings.Contains(low, "model") || + strings.Contains(low, "usage") || strings.Contains(low, "quota") || + strings.Contains(low, "rolling") { + return true + } + } + return false +} + +func isGrokBillingQuotaText(low string) bool { + if low == "" { + return false + } + if strings.Contains(low, "insufficient_quota") { + return true + } + if strings.Contains(low, "billing") && strings.Contains(low, "quota") { + return true + } + if strings.Contains(low, "payment") && (strings.Contains(low, "required") || strings.Contains(low, "fail")) { + return true + } + if strings.Contains(low, "spending limit") || strings.Contains(low, "run out of credits") || strings.Contains(low, "out of credits") { + return true + } + if strings.Contains(low, "余额不足") || strings.Contains(low, "欠费") || strings.Contains(low, "需要付费") { + return true + } + return false +} + +func isGrokModelCapacityText(low string) bool { + return strings.Contains(low, "capacity") || + strings.Contains(low, "overloaded") || + strings.Contains(low, "server_busy") || + strings.Contains(low, "too many concurrent") || + strings.Contains(low, "engine_overloaded") +} + +func isGrokRateLimitText(low string) bool { + return strings.Contains(low, "rate limit") || + strings.Contains(low, "rate_limit") || + strings.Contains(low, "too many requests") || + strings.Contains(low, "请求过于频繁") || + strings.Contains(low, "速率限制") +} + +func isGrokEmptyModelOutputText(low string) bool { + if low == "" { + return false + } + return strings.Contains(low, "empty model output") || + strings.Contains(low, "no content/tool_calls") || + strings.Contains(low, "no client-visible content") || + strings.Contains(low, "empty_upstream") || + strings.Contains(low, "empty upstream") +} + +func isGrokEmptyModelOutputCode(code string) bool { + c := strings.ToLower(strings.TrimSpace(code)) + if c == "" { + return false + } + return c == "empty_upstream" || + c == "empty-model-output" || + c == "empty_model_output" || + strings.Contains(c, "empty_upstream") || + strings.Contains(c, "empty-model-output") +} + +func grokFreeUsageCooldownDuration(low string) time.Duration { + // "rolling 24-hour" describes the upstream usage window, not a cooldown + // that starts when this proxy observes a 429. Without an upstream reset + // timestamp we cannot know when the oldest usage exits that window, so use + // a short probe interval and let a successful probe clear the block. + return grokFreeUsageProbeCooldown +} + +const grokFreeUsageProbeCooldown = 10 * time.Minute + +func parseGrokTokenPair(errText string) (actual, limit int64, ok bool) { + m := reGrokTokenPair.FindStringSubmatch(errText) + if len(m) != 3 { + return 0, 0, false + } + a, errA := strconv.ParseInt(m[1], 10, 64) + b, errB := strconv.ParseInt(m[2], 10, 64) + if errA != nil || errB != nil { + return 0, 0, false + } + return a, b, true +} + +func extractGrokFailureModel(text string, responseBody []byte, fallback string) string { + if m := reGrokModelFor.FindStringSubmatch(text); len(m) == 2 { + return normalizeGrokFailureModelID(m[1]) + } + if len(responseBody) > 0 { + if m := strings.TrimSpace(firstNonEmpty( + gjson.GetBytes(responseBody, "error.model").String(), + gjson.GetBytes(responseBody, "model").String(), + )); m != "" { + return normalizeGrokFailureModelID(m) + } + } + return normalizeGrokFailureModelID(fallback) +} + +// normalizeGrokFailureModelID trims whitespace and trailing punctuation the +// "for model X." extractors sometimes capture. +func normalizeGrokFailureModelID(model string) string { + model = strings.TrimSpace(model) + model = strings.TrimRight(model, ".,;:!?") + return strings.TrimSpace(model) +} + +// applyGrokUpstreamFailureDecision maps a classification onto existing account +// health helpers. Returns true when the decision fully handled the error path +// (caller should not apply the status-code switch defaults again). +func (s *OpenAIGatewayService) applyGrokUpstreamFailureDecision( + ctx context.Context, + account *Account, + decision GrokUpstreamFailureDecision, +) bool { + if s == nil || account == nil || !decision.ShouldCooldown || decision.Cooldown <= 0 { + return false + } + // Keep reasons short and stable for ops UI / temp_unschedulable_reason. + var reason string + switch decision.Class { + case GrokFailureFreeUsage: + reason = "grok free usage exhausted" + // Model-scoped free usage: soft-block only the named model so other + // models on the same account remain pickable (grok2api ModelQuotaBlock). + low := strings.ToLower(decision.Reason) + if decision.Model != "" && isGrokModelSpecificFreeUsage(low, decision.Model) { + until := time.Now().Add(decision.Cooldown) + markGrokModelQuotaBlock(account.ID, decision.Model, until) + // The upstream explicitly scoped exhaustion to this model, so an + // account-wide cool would incorrectly take healthy sibling models out + // of rotation. + return true + } + case GrokFailureBilling: + low := strings.ToLower(decision.Reason) + if strings.Contains(low, "spending") || strings.Contains(low, "credits") { + // Spending-limit/credit exhaustion is a billing-window condition. Keep + // the account recoverable and let the normal rate-limit recovery clear it. + s.rateLimitGrok(ctx, account, grokSpendingLimitResetAt(account, time.Now())) + return true + } + // Keep the historical 402/payment reason for ops UI + regression tests. + reason = "grok payment required" + case GrokFailureEmptyUpstream: + reason = "grok empty model output" + case GrokFailureModelCapacity: + reason = "grok model capacity" + case GrokFailureRateLimit: + // Pure 429 without free-usage language keeps the existing rate-limit + // snapshot path (Retry-After / quota headers). Body-only rate-limit + // phrasing still cools here via ShouldCooldown from the classifier, but + // the handler only invokes this for non-RateLimit classes. + return false + case GrokFailureServer: + reason = "grok upstream temporary error" + default: + return false + } + s.tempUnscheduleGrok(ctx, account, decision.Cooldown, reason) + return true +} diff --git a/backend/internal/service/grok_upstream_failure_test.go b/backend/internal/service/grok_upstream_failure_test.go new file mode 100644 index 0000000000..6f0a916a78 --- /dev/null +++ b/backend/internal/service/grok_upstream_failure_test.go @@ -0,0 +1,181 @@ +//go:build unit + +package service + +import ( + "context" + "net/http" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestClassifyGrokUpstreamFailure_FreeUsage(t *testing.T) { + cases := []struct { + name string + status int + body string + }{ + { + name: "code free-usage-exhausted", + status: http.StatusTooManyRequests, + body: `{"error":{"code":"subscription:free-usage-exhausted","message":"You've used all the included free usage for model grok-4.5. Usage resets over a rolling 24-hour window."}}`, + }, + { + name: "chinese body without 429", + status: http.StatusBadRequest, + body: `{"error":{"message":"模型额度用完，请稍后再试"}}`, + }, + { + name: "token pair with free marker", + status: http.StatusOK, + body: `{"error":{"message":"free usage tokens (actual / limit): 2000000 / 2000000 for model grok-4.5"}}`, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + d := classifyGrokUpstreamFailure(tc.status, []byte(tc.body), "grok-4.5") + require.Equal(t, GrokFailureFreeUsage, d.Class) + require.True(t, d.ShouldCooldown) + require.True(t, d.ShouldFailover) + require.False(t, d.BlockModel, "free-usage must not soft-block models") + require.Equal(t, grokFreeUsageProbeCooldown, d.Cooldown) + }) + } +} + +func TestClassifyGrokUpstreamFailure_EmptyUpstream(t *testing.T) { + d := classifyGrokUpstreamFailure(http.StatusBadGateway, []byte(`empty model output: no content/tool_calls`), "grok-4.5") + require.Equal(t, GrokFailureEmptyUpstream, d.Class) + require.True(t, d.ShouldCooldown) + require.True(t, d.ShouldFailover) + require.True(t, d.BlockModel) + require.Equal(t, 4*time.Minute, d.Cooldown) +} + +func TestClassifyGrokUpstreamFailure_Billing(t *testing.T) { + d := classifyGrokUpstreamFailure(http.StatusForbidden, []byte(`{"code":"personal-team-blocked:spending-limit","error":"spending limit reached"}`), "") + require.Equal(t, GrokFailureBilling, d.Class) + require.True(t, d.ShouldCooldown) + require.True(t, d.ShouldFailover) +} + +func TestClassifyGrokUpstreamFailure_ValidationNoCool(t *testing.T) { + d := classifyGrokUpstreamFailure(http.StatusBadRequest, []byte(`{"error":{"message":"invalid tool schema"}}`), "") + require.Equal(t, GrokFailureNone, d.Class) + require.False(t, d.ShouldCooldown) + require.False(t, d.ShouldFailover) +} + +func TestClassifyGrokUpstreamFailure_FreeUsageWinsOver5xx(t *testing.T) { + // Proxy may rewrite free-usage into synthetic 502; body must win. + d := classifyGrokUpstreamFailure(http.StatusBadGateway, []byte(`subscription:free-usage-exhausted for model grok-4.3`), "grok-4.3") + require.Equal(t, GrokFailureFreeUsage, d.Class) + require.NotEqual(t, GrokFailureServer, d.Class) +} + +func TestShouldFailoverGrokUpstreamError_FreeUsageBody(t *testing.T) { + svc := &OpenAIGatewayService{} + body := []byte(`{"error":{"code":"subscription:free-usage-exhausted","message":"free usage exhausted"}}`) + require.True(t, svc.shouldFailoverGrokUpstreamError(http.StatusBadRequest, body)) +} + +func TestShouldFailoverGrokUpstreamError_ContentPolicyStillNoFailover(t *testing.T) { + svc := &OpenAIGatewayService{} + body := []byte(`{"error":{"code":"new_sensitive","message":"text is sensitive"}}`) + require.False(t, svc.shouldFailoverGrokUpstreamError(http.StatusForbidden, body)) +} + +func TestHandleGrokAccountUpstreamError_FreeUsageBodyCoolsAccount(t *testing.T) { + repo := &grokQuotaAccountRepo{} + svc := &OpenAIGatewayService{accountRepo: repo} + account := &Account{ID: 9101, Platform: PlatformGrok, Type: AccountTypeOAuth} + before := time.Now() + body := []byte(`{"error":{"code":"subscription:free-usage-exhausted","message":"You've used all the included free usage. Usage resets over a rolling 24-hour window."}}`) + + svc.handleGrokAccountUpstreamError(context.Background(), account, http.StatusBadRequest, nil, body) + + require.Equal(t, 1, repo.tempUnschedCalls) + require.Equal(t, "grok free usage exhausted", repo.lastTempUnschedReason) + // Rolling-window exhaustion must use a short probe cooldown when no + // upstream absolute reset is available; it must not start a 24h lock here. + require.Greater(t, repo.lastTempUnschedUntil, before.Add(grokFreeUsageProbeCooldown-time.Second)) + require.Less(t, repo.lastTempUnschedUntil, before.Add(grokFreeUsageProbeCooldown+time.Second)) +} + +func TestHandleGrokAccountUpstreamError_FreeUsageUsesUpstreamReset(t *testing.T) { + repo := &grokQuotaAccountRepo{} + svc := &OpenAIGatewayService{accountRepo: repo} + account := &Account{ID: 9102, Platform: PlatformGrok, Type: AccountTypeOAuth} + body := []byte(`{"error":{"code":"subscription:free-usage-exhausted","message":"free usage exhausted; rolling 24-hour window"}}`) + + svc.handleGrokAccountUpstreamError(context.Background(), account, http.StatusTooManyRequests, + http.Header{"Retry-After": []string{"3600"}}, body) + + require.Zero(t, repo.tempUnschedCalls) + require.WithinDuration(t, time.Now().Add(time.Hour), repo.lastRateLimitResetAt, 2*time.Second) +} + +func TestHandleGrokAccountUpstreamError_EmptyOutputCoolsAccount(t *testing.T) { + repo := &grokQuotaAccountRepo{} + svc := &OpenAIGatewayService{accountRepo: repo} + account := &Account{ID: 9102, Platform: PlatformGrok, Type: AccountTypeOAuth} + before := time.Now() + + svc.handleGrokAccountUpstreamError( + context.Background(), account, http.StatusBadGateway, nil, + []byte(`empty model output: no content/tool_calls`), + ) + + require.Equal(t, 1, repo.tempUnschedCalls) + require.Equal(t, "grok empty model output", repo.lastTempUnschedReason) + require.WithinDuration(t, before.Add(4*time.Minute), repo.lastTempUnschedUntil, time.Second) +} + +func TestHandleGrokAccountUpstreamError_FreeUsageDoesNotCoolPoolMode(t *testing.T) { + repo := &grokQuotaAccountRepo{} + svc := &OpenAIGatewayService{accountRepo: repo} + account := &Account{ + ID: 9103, + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + Credentials: map[string]any{ + "pool_mode": true, + }, + } + body := []byte(`{"error":{"code":"subscription:free-usage-exhausted","message":"free usage exhausted"}}`) + + svc.handleGrokAccountUpstreamError(context.Background(), account, http.StatusBadRequest, nil, body) + + require.Zero(t, repo.tempUnschedCalls) + require.False(t, svc.isOpenAIAccountRuntimeBlocked(account)) +} + +func TestHandleGrokAccountUpstreamError_ContentPolicyStillNoMutation(t *testing.T) { + repo := &grokQuotaAccountRepo{} + svc := &OpenAIGatewayService{accountRepo: repo} + account := &Account{ID: 9104, Platform: PlatformGrok, Type: AccountTypeOAuth} + body := []byte(`{"error":{"code":"new_sensitive","message":"text is sensitive"}}`) + + svc.handleGrokAccountUpstreamError(context.Background(), account, http.StatusForbidden, nil, body) + + require.Zero(t, repo.tempUnschedCalls) +} + +func TestHandleGrokAccountUpstreamError_Entitlement403Unchanged(t *testing.T) { + repo := &grokQuotaAccountRepo{} + svc := &OpenAIGatewayService{accountRepo: repo} + account := &Account{ID: 9105, Platform: PlatformGrok, Type: AccountTypeOAuth} + before := time.Now() + + svc.handleGrokAccountUpstreamError( + context.Background(), account, http.StatusForbidden, nil, + []byte(`{"error":{"message":"subscription required"}}`), + ) + + require.Equal(t, 1, repo.tempUnschedCalls) + require.Equal(t, "grok access or entitlement denied", repo.lastTempUnschedReason) + require.Greater(t, repo.lastTempUnschedUntil, before.Add(29*time.Minute)) + require.Less(t, repo.lastTempUnschedUntil, before.Add(31*time.Minute)) +} diff --git a/backend/internal/service/grok_upstream_headers.go b/backend/internal/service/grok_upstream_headers.go new file mode 100644 index 0000000000..cdfb54bcbf --- /dev/null +++ b/backend/internal/service/grok_upstream_headers.go @@ -0,0 +1,74 @@ +package service + +import ( + "net/http" + "strings" + + "github.com/gin-gonic/gin" + + "github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint" + "github.com/Wei-Shaw/sub2api/internal/pkg/xai" +) + +// grokUpstreamUserAgent is kept for compatibility with older Grok request +// tests. Current requests use the pinned default UA from this package. +const grokUpstreamUserAgent = "sub2api-grok/1.0" + +// Fixed CLI identity aliases — single source of truth is internal/pkg/xai. +const ( + grokClientVersionHeader = xai.CLIStableVersion + grokClientIdentifierHeader = xai.CLIClientIdentifier + grokClientModeHeader = xai.CLIClientMode +) + +// defaultGrokUpstreamUserAgent is the pinned Grok CLI / workspace UA. +// Grok upstream must not forward Claude Code / Codex / browser client UAs. +func defaultGrokUpstreamUserAgent() string { + return xai.CLIUserAgent(xai.ResolveCLIVersion()) +} + +func applyDefaultGrokUpstreamHeaders(req *http.Request) { + if req == nil { + return + } + // Always stamp CLI identity. Do not preserve inbound client UA (Claude Code, + // Codex, curl, etc.) — xAI chat/CLI surfaces fingerprint the client string. + req.Header.Set("User-Agent", defaultGrokUpstreamUserAgent()) + req.Header.Set("x-grok-client-version", xai.ResolveCLIVersion()) + req.Header.Set("x-grok-client-identifier", grokClientIdentifierHeader) +} + +func applyGrokTLSProfileHeaders(req *http.Request, profile *tlsfingerprint.Profile) { + // HEAD Profile is TLS-only (no HTTP UserAgent/Originator fields). Always stamp CLI identity. + applyDefaultGrokUpstreamHeaders(req) + _ = profile +} + +// openAITLSFingerprintRuntime is the resolved TLS fingerprint routing result +// used by OpenAI/Grok outbound header application. Defined here so Grok header +// helpers compile even when the full OpenAI TLS router is not present on HEAD. +type openAITLSFingerprintRuntime struct { + Profile *tlsfingerprint.Profile + UpstreamUserAgent string + UpstreamOriginator string + Matched bool +} + +func applyGrokRuntimeHeaders(req *http.Request, runtime openAITLSFingerprintRuntime) { + applyDefaultGrokUpstreamHeaders(req) + if req == nil { + return + } + // Apply Originator only; force CLI UA after so router overrides cannot + // leak Codex/Claude Code identity onto Grok upstream. + if originator := strings.TrimSpace(runtime.UpstreamOriginator); originator != "" { + req.Header.Set("Originator", originator) + } + req.Header.Set("User-Agent", defaultGrokUpstreamUserAgent()) +} + +// resolveGrokUpstreamUserAgent always returns the pinned Grok CLI User-Agent. +// Inbound client UAs (Claude Code, Codex, browsers, libraries) are never forwarded. +func resolveGrokUpstreamUserAgent(_ *gin.Context) string { + return defaultGrokUpstreamUserAgent() +} diff --git a/backend/internal/service/grok_upstream_headers_test.go b/backend/internal/service/grok_upstream_headers_test.go new file mode 100644 index 0000000000..18972f4801 --- /dev/null +++ b/backend/internal/service/grok_upstream_headers_test.go @@ -0,0 +1,86 @@ +package service + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + + "github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint" + "github.com/Wei-Shaw/sub2api/internal/pkg/xai" +) + +func TestApplyDefaultGrokUpstreamHeadersUsesCLIUserAgent(t *testing.T) { + t.Setenv(xai.CLIVersionEnv, "") + + req, err := http.NewRequest(http.MethodGet, "https://api.x.ai/v1/responses", nil) + require.NoError(t, err) + req.Header.Set("User-Agent", "claude-code/1.2.3") + req.Header.Set("x-grok-client-version", "none") + + applyDefaultGrokUpstreamHeaders(req) + + require.Equal(t, xai.CLIUserAgent(xai.CLIClientVersion), req.Header.Get("User-Agent")) + require.Equal(t, xai.CLIClientVersion, req.Header.Get("x-grok-client-version")) + require.Equal(t, xai.CLIClientIdentifier, req.Header.Get("x-grok-client-identifier")) +} + +func TestApplyDefaultGrokUpstreamHeadersHonorsCLIVersionOverride(t *testing.T) { + t.Setenv(xai.CLIVersionEnv, "0.2.95") + + req, err := http.NewRequest(http.MethodGet, "https://api.x.ai/v1/responses", nil) + require.NoError(t, err) + req.Header.Set("User-Agent", "codex_cli_rs/0.144.0") + + applyDefaultGrokUpstreamHeaders(req) + + require.Equal(t, "0.2.95", req.Header.Get("x-grok-client-version")) + require.Equal(t, xai.CLIUserAgent("0.2.95"), req.Header.Get("User-Agent")) + require.Equal(t, "grok-shell", req.Header.Get("x-grok-client-identifier")) +} + +func TestResolveGrokUpstreamUserAgentNeverPassthrough(t *testing.T) { + t.Setenv(xai.CLIVersionEnv, "") + gin.SetMode(gin.TestMode) + + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + c.Request.Header.Set("User-Agent", "claude-cli/2.0.0 (Mac OS; arm64)") + + require.Equal(t, xai.CLIUserAgent(xai.CLIClientVersion), resolveGrokUpstreamUserAgent(c)) + require.Equal(t, xai.CLIUserAgent(xai.CLIClientVersion), resolveGrokUpstreamUserAgent(nil)) +} + +func TestApplyGrokRuntimeHeadersKeepsCLIUserAgent(t *testing.T) { + t.Setenv(xai.CLIVersionEnv, "") + + req, err := http.NewRequest(http.MethodPost, "https://cli-chat-proxy.grok.com/v1/responses", nil) + require.NoError(t, err) + req.Header.Set("User-Agent", "claude-code/9.9.9") + + applyGrokRuntimeHeaders(req, openAITLSFingerprintRuntime{ + UpstreamUserAgent: "codex_cli_rs/0.144.0", + UpstreamOriginator: "codex_cli_rs", + }) + + require.Equal(t, xai.CLIUserAgent(xai.CLIClientVersion), req.Header.Get("User-Agent")) + require.Equal(t, "codex_cli_rs", req.Header.Get("Originator")) + require.Equal(t, xai.CLIClientVersion, req.Header.Get("x-grok-client-version")) +} + +func TestApplyGrokTLSProfileHeadersAlwaysUsesCLIUserAgent(t *testing.T) { + t.Setenv(xai.CLIVersionEnv, "") + + req, err := http.NewRequest(http.MethodPost, "https://api.x.ai/v1/responses", nil) + require.NoError(t, err) + req.Header.Set("User-Agent", "grok-native/1.0") + + // HEAD Profile is TLS-only; Originator/UserAgent HTTP fields are not present. + applyGrokTLSProfileHeaders(req, &tlsfingerprint.Profile{Name: "chrome"}) + + require.Equal(t, xai.CLIUserAgent(xai.CLIClientVersion), req.Header.Get("User-Agent")) + require.Equal(t, xai.CLIClientVersion, req.Header.Get("x-grok-client-version")) +} diff --git a/backend/internal/service/grok_upstream_url.go b/backend/internal/service/grok_upstream_url.go index 362e746e6f..9574b3adbe 100644 --- a/backend/internal/service/grok_upstream_url.go +++ b/backend/internal/service/grok_upstream_url.go @@ -1,8 +1,11 @@ package service import ( + "context" "errors" "fmt" + "net/url" + "strings" "github.com/Wei-Shaw/sub2api/internal/config" "github.com/Wei-Shaw/sub2api/internal/pkg/xai" @@ -67,20 +70,28 @@ func redactedGrokBaseURLValidator(validator xai.BaseURLValidator) xai.BaseURLVal } } -func buildGrokResponsesURL(account *Account, cfg *config.Config) (string, error) { +func buildGrokResponsesURL(account *Account, cfg *config.Config, settings ...*SettingService) (string, error) { validator, err := grokBaseURLValidator(account, cfg) if err != nil { return "", err } - return xai.BuildResponsesURLWithValidator(account.GetGrokBaseURL(), validator) + baseURL := account.GetGrokBaseURL() + if len(settings) > 0 && settings[0] != nil { + baseURL = settings[0].ResolveGrokBaseURL(context.Background(), account) + } + return xai.BuildResponsesURLWithValidator(baseURL, validator) } -func buildGrokChatCompletionsURL(account *Account, cfg *config.Config) (string, error) { +func buildGrokChatCompletionsURL(account *Account, cfg *config.Config, settings ...*SettingService) (string, error) { validator, err := grokBaseURLValidator(account, cfg) if err != nil { return "", err } - return xai.BuildChatCompletionsURLWithValidator(account.GetGrokBaseURL(), validator) + baseURL := account.GetGrokBaseURL() + if len(settings) > 0 && settings[0] != nil { + baseURL = settings[0].ResolveGrokBaseURL(context.Background(), account) + } + return xai.BuildChatCompletionsURLWithValidator(baseURL, validator) } // buildGrokBillingURL 解析 billing 探测端点：跟随账号的转发 base_url， @@ -90,7 +101,14 @@ func buildGrokBillingURL(account *Account, cfg *config.Config, weekly bool) (str if err != nil { return "", err } - return xai.BuildBillingURLWithValidator(account.GetGrokBaseURL(), weekly, validator) + baseURL := account.GetGrokBaseURL() + // Official public/regional API hosts do not expose Grok Build billing. + // Keep custom relays on their configured host because they may proxy the CLI + // billing path alongside inference. + if xai.IsOfficialBaseURL(baseURL) && !isGrokCLIProxyBaseURL(baseURL) { + baseURL = xai.DefaultCLIBaseURL + } + return xai.BuildBillingURLWithValidator(baseURL, weekly, validator) } func buildGrokMediaURL(account *Account, cfg *config.Config, endpoint GrokMediaEndpoint, requestID string) (string, error) { @@ -122,3 +140,42 @@ func buildGrokMediaURL(account *Account, cfg *config.Config, endpoint GrokMediaE return "", fmt.Errorf("unsupported grok media endpoint: %s", endpoint) } } + +// buildGrokVoiceURL returns the official xAI Voice API endpoint. +// Voice HTTP (/tts, /stt, /custom-voices) and WS (/realtime) are only exposed +// by api.x.ai — the CLI chat proxy does not implement them. When the account +// base_url points at the CLI proxy (or is empty), fall back to DefaultBaseURL. +func buildGrokVoiceURL(account *Account, cfg *config.Config, endpoint string) (string, error) { + validator, err := grokBaseURLValidator(account, cfg) + if err != nil { + return "", err + } + base := "" + if account != nil { + base = account.GetGrokMediaBaseURL() + } + if strings.TrimSpace(base) == "" || isGrokCLIProxyBaseURL(base) { + base = xai.DefaultBaseURL + } + validated, err := validator(base) + if err != nil { + return "", err + } + ep := strings.Trim(strings.TrimSpace(endpoint), "/") + if ep == "" { + return "", fmt.Errorf("voice endpoint is required") + } + parts := strings.Split(ep, "/") + encoded := make([]string, 0, len(parts)) + for _, part := range parts { + if strings.TrimSpace(part) == "" || part == "." || part == ".." { + return "", fmt.Errorf("invalid voice endpoint path") + } + encoded = append(encoded, url.PathEscape(part)) + } + return strings.TrimRight(validated, "/") + "/" + strings.Join(encoded, "/"), nil +} + +func isGrokCLIProxyBaseURL(raw string) bool { + return isGrokCLIProxyTarget(raw) +} diff --git a/backend/internal/service/grok_upstream_url_test.go b/backend/internal/service/grok_upstream_url_test.go index e2f6e8bce5..8544823b4a 100644 --- a/backend/internal/service/grok_upstream_url_test.go +++ b/backend/internal/service/grok_upstream_url_test.go @@ -255,6 +255,26 @@ func TestGrokOAuthURLPolicy(t *testing.T) { }) } +func TestBuildGrokBillingURLUsesCLIForOfficialAPIHosts(t *testing.T) { + for _, baseURL := range []string{xai.DefaultBaseURL, "https://us-west-2.api.x.ai/v1"} { + account := &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Credentials: map[string]any{"base_url": baseURL}} + + weekly, err := buildGrokBillingURL(account, &config.Config{}, true) + require.NoError(t, err) + require.Equal(t, xai.DefaultCLIBaseURL+xai.BillingWeeklyPath, weekly) + } +} + +func TestBuildGrokBillingURLKeepsCustomRelay(t *testing.T) { + account := &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Credentials: map[string]any{ + "base_url": "https://relay.example.test/xai/v1", + }} + + monthly, err := buildGrokBillingURL(account, &config.Config{}, false) + require.NoError(t, err) + require.Equal(t, "https://relay.example.test/xai/v1"+xai.BillingMonthlyPath, monthly) +} + func TestGrokBillingURLFollowsAccountBaseURL(t *testing.T) { t.Run("oauth default stays on CLI gateway", func(t *testing.T) { account := &Account{ diff --git a/backend/internal/service/group.go b/backend/internal/service/group.go index b71e6b9e84..ddde276e22 100644 --- a/backend/internal/service/group.go +++ b/backend/internal/service/group.go @@ -3,6 +3,7 @@ package service import ( "errors" "fmt" + "math" "strings" "time" @@ -54,10 +55,21 @@ type Group struct { VideoPrice480P *float64 VideoPrice720P *float64 VideoPrice1080P *float64 + // VideoModelPrices is optional per-model-family per-second pricing + // (groups.video_model_prices JSONB). Shape: family → resolution → USD/s. + // When set for a model, overrides VideoPrice* for that model only. + VideoModelPrices map[string]map[string]float64 // Codex alpha/search 网页搜索单次价格（USD/次，仅 openai 平台使用）； // nil 表示使用默认价 defaultWebSearchPricePerCall（官方 $10/1000 次）。 WebSearchPricePerCall *float64 + // 搜索工具显式定价（per 1k calls）。 + SearchPricePer1k *float64 + // Grok Voice 显式定价（分组级，不按文本 RateMultiplier）。 + AudioRealtimePricePerMin *float64 + AudioTTSPricePerMillionChars *float64 + AudioSTTPricePerHour *float64 + // Claude Code 客户端限制 ClaudeCodeOnly bool FallbackGroupID *int64 @@ -99,6 +111,14 @@ type Group struct { // ReasoningEffortMappings rewrites explicit request values before applying the ceiling. ReasoningEffortMappings []ReasoningEffortMapping + // 分组利润控制（五个 token 计费平台可启用）。 + // 调度准入条件：账号倍率 U 满足 U <= D*(1-margin-buffer)， + // D 为请求用户当刻有效下游倍率（用户覆盖 ?? 分组默认，再乘高峰因子）。 + // 只过滤候选账号，不改变既有排序/评分/粘性/熔断。 + ProfitControlEnabled bool + ProfitMinMargin float64 // 最低毛利率，小数存储（0.30=30%） + ProfitSafetyBuffer float64 // 安全缓冲，小数，与 margin 相加后从 D 中扣除 + CreatedAt time.Time UpdatedAt time.Time @@ -159,6 +179,30 @@ func (g *Group) GetVideoPrice(resolution string) *float64 { } } +// GetVideoPriceForModel prefers VideoModelPrices for the model family, then flat columns. +func (g *Group) GetVideoPriceForModel(model, resolution string) *float64 { + if g == nil { + return nil + } + if price := LookupVideoModelPrice(g.VideoModelPrices, model, resolution); price != nil { + return price + } + return g.GetVideoPrice(resolution) +} + +// VideoPriceConfig builds billing config including optional per-model map. +func (g *Group) VideoPriceConfig() *VideoPriceConfig { + if g == nil { + return nil + } + return &VideoPriceConfig{ + Price480P: g.VideoPrice480P, + Price720P: g.VideoPrice720P, + Price1080P: g.VideoPrice1080P, + ModelPrices: NormalizeVideoModelPrices(g.VideoModelPrices), + } +} + // IsGroupContextValid reports whether a group from context has the fields required for routing decisions. func IsGroupContextValid(group *Group) bool { if group == nil { @@ -336,3 +380,80 @@ func computePeakAwareMultipliers(apiKey *APIKey, base float64, now time.Time) (t text = base * peak return } + +// validProfitControlRatio 判定 margin/buffer 是否为可落库的合法小数：[0,1) 且非 NaN/Inf。 +func validProfitControlRatio(v float64) bool { + return !math.IsNaN(v) && !math.IsInf(v, 0) && v >= 0 && v < 1 +} + +// NormalizeGroupPlatform 把创建分组时省略的 platform 归一化为默认平台。 +// handler 的入参预校验必须与 CreateGroup 落库时用同一个归一化结果，否则 +// 「省略 platform + 启用利润控制」会被 handler 以「平台不支持」400 掉， +// 而该分组本会被建成受支持的 anthropic 分组。 +func NormalizeGroupPlatform(platform string) string { + if platform == "" { + return PlatformAnthropic + } + return platform +} + +// ValidateProfitControlConfig 是分组利润控制配置的唯一校验来源，handler 与 service 层共用。 +// enabled=true 时仅允许五个可计费平台分组；margin/buffer 各自 ∈ [0,1)，且 margin+buffer < 1 +// （相加 >=1 时阈值 <=0，所有可核价账号都会被排除，视为配置错误而不是静默全黑）。 +// enabled=false 时放行（不关心平台），由 Normalize 兜底清洗数值。 +func ValidateProfitControlConfig(platform string, enabled bool, minMargin, safetyBuffer float64) error { + if !enabled { + return nil + } + if !profitControlPlatformSupported(platform) { + return errors.New("利润控制仅支持 openai、anthropic、gemini、grok、antigravity 平台分组") + } + if !validProfitControlRatio(minMargin) { + return fmt.Errorf("profit_min_margin 应为 [0,1) 的小数，got %v", minMargin) + } + if !validProfitControlRatio(safetyBuffer) { + return fmt.Errorf("profit_safety_buffer 应为 [0,1) 的小数，got %v", safetyBuffer) + } + if minMargin+safetyBuffer >= 1 { + return errors.New("profit_min_margin 与 profit_safety_buffer 之和必须小于 1，否则将排除全部账号") + } + return nil +} + +// NormalizeProfitControlConfig 归一化最终落库的利润控制配置，CreateGroup 与 UpdateGroup 共用（唯一收口）： +// - 非五个平台分组不携带利润控制，一律重置为默认（关、0、0）； +// - 支持平台关闭开关时保留合法数值（便于再次启用），清洗 NaN/Inf/越界脏值。 +// +// 与 ValidateProfitControlConfig 的分工同高峰倍率：先归一化、后校验， +// 使"openai 转其他平台"这类更新能静默清空利润配置而不是被校验拒绝。 +func NormalizeProfitControlConfig(platform string, enabled bool, minMargin, safetyBuffer float64) (bool, float64, float64) { + if !profitControlPlatformSupported(platform) { + return false, 0, 0 + } + if !enabled { + if !validProfitControlRatio(minMargin) { + minMargin = 0 + } + if !validProfitControlRatio(safetyBuffer) { + safetyBuffer = 0 + } + } + return enabled, minMargin, safetyBuffer +} + +func profitControlPlatformSupported(platform string) bool { + switch platform { + case PlatformOpenAI, PlatformAnthropic, PlatformGemini, PlatformGrok, PlatformAntigravity: + return true + default: + return false + } +} + +// GetSearchPricePer1k returns explicit search/tool price per 1k calls if configured. +func (g *Group) GetSearchPricePer1k() *float64 { + if g == nil { + return nil + } + return g.SearchPricePer1k +} diff --git a/backend/internal/service/group_profit_platform_test.go b/backend/internal/service/group_profit_platform_test.go new file mode 100644 index 0000000000..28a5f2290a --- /dev/null +++ b/backend/internal/service/group_profit_platform_test.go @@ -0,0 +1,20 @@ +package service + +import "testing" + +// TestNormalizeGroupPlatformDefaultsToAnthropic 钉死创建分组时省略 platform 的 +// 归一化结果：handler 的利润控制预校验与 CreateGroup 落库共用这个函数， +// 两边一旦漂移，「省略 platform + profit_control_enabled=true」的创建请求 +// 会被 handler 以「平台不支持」400 掉，而该分组本会建成受支持的 anthropic。 +func TestNormalizeGroupPlatformDefaultsToAnthropic(t *testing.T) { + if got := NormalizeGroupPlatform(""); got != PlatformAnthropic { + t.Fatalf("空 platform 应归一化为 %s，got %s", PlatformAnthropic, got) + } + if got := NormalizeGroupPlatform(PlatformOpenAI); got != PlatformOpenAI { + t.Fatalf("显式 platform 不应被改写，got %s", got) + } + // 归一化后的默认平台必须真的支持利润控制，否则预校验仍会误拒。 + if err := ValidateProfitControlConfig(NormalizeGroupPlatform(""), true, 0.3, 0.05); err != nil { + t.Fatalf("省略 platform 且启用利润控制不应被预校验拒绝: %v", err) + } +} diff --git a/backend/internal/service/identity_service.go b/backend/internal/service/identity_service.go index 72635c1d68..afa882a9c9 100644 --- a/backend/internal/service/identity_service.go +++ b/backend/internal/service/identity_service.go @@ -23,8 +23,58 @@ import ( var ( // 匹配 User-Agent 版本号: xxx/x.y.z userAgentVersionRegex = regexp.MustCompile(`/(\d+)\.(\d+)\.(\d+)`) + + // fingerprintUserAgentPattern 校验可写入账号级持久身份的 User-Agent 形态： + //+ {{ t('admin.accounts.openai.flattenNamespacesDesc') }} +
+{{ t('admin.accounts.billingRateMultiplierHint') }}
+
+
+ {{ t('admin.accounts.upstreamBilling.autoProbeHint') }} +
+{{ apiKeyHint }}
+ {{ t('admin.accounts.openai.flattenNamespacesDesc') }} +
++ {{ t('admin.accounts.accountSchedulingThresholdOverrideHint') }} +
+{{ t('admin.accounts.accountSchedulingThresholdOverrideDisabledHint') }}
+{{ t('admin.accounts.billingRateMultiplierHint') }}
+ ++ {{ + t( + upstreamBillingRateSyncEnabled + ? 'admin.accounts.upstreamBilling.syncRateManagedHint' + : 'admin.accounts.billingRateMultiplierHint' + ) + }} +
++ {{ t('admin.accounts.upstreamBilling.syncRate') }} +
++ {{ t('admin.accounts.upstreamBilling.syncRateHint') }} +
++ {{ t('admin.accounts.openai.flattenNamespacesDesc') }} +
++ {{ t(getOAuthKey('emailPasswordDesc')) }} +
++ {{ t(getOAuthKey('emailPasswordHint')) }} +
++ {{ error }} +
+{{ hint }}
@@ -70,9 +70,12 @@ + + diff --git a/frontend/src/features/channel-monitor-v2/MetricCell.vue b/frontend/src/features/channel-monitor-v2/MetricCell.vue new file mode 100644 index 0000000000..aa3eee65eb --- /dev/null +++ b/frontend/src/features/channel-monitor-v2/MetricCell.vue @@ -0,0 +1,73 @@ + ++ {{ t('channelMonitorV2.settings.description') }} +
++ {{ t('channelMonitorV2.settings.enableHint') }} +
+{{ t('channelMonitorV2.settings.refreshHint') }}
++ {{ t('channelMonitorV2.settings.platformsHint') }} +
++ {{ + draft.group_ids.length + ? t('channelMonitorV2.settings.groupsSelected', { count: draft.group_ids.length }) + : t('channelMonitorV2.settings.groupsAll') + }} +
+{{ t('channelMonitorV2.settings.groupsEmpty') }}
++ {{ t('channelMonitorV2.settings.errorsHint') }} +
++ {{ t('channelMonitorV2.settings.healthHint') }} +
+{{ t('channelMonitorV2.settings.userContractTitle') }}
++ {{ t('channelMonitorV2.chart.description') }} +
++ {{ t('channelMonitorV2.matrix.description') }} +
+