From c4128580f287240ff9ed40764ecfab77c2f7ddd4 Mon Sep 17 00:00:00 2001 From: JRBaggins <45261448+JRBaggins@users.noreply.github.com> Date: Thu, 2 Jul 2026 11:42:10 +0800 Subject: [PATCH] fix: handle OpenAI OAuth count_tokens scope errors --- backend/go.mod | 2 + backend/go.sum | 4 + .../service/openai_gateway_count_tokens.go | 356 ++++++++++++++++-- .../openai_gateway_count_tokens_test.go | 290 ++++++++++++-- 4 files changed, 595 insertions(+), 57 deletions(-) diff --git a/backend/go.mod b/backend/go.mod index 62be56c86e..d4e0af6c13 100644 --- a/backend/go.mod +++ b/backend/go.mod @@ -86,6 +86,7 @@ require ( github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect github.com/distribution/reference v0.6.0 // indirect + github.com/dlclark/regexp2/v2 v2.1.0 // indirect github.com/docker/docker v28.5.1+incompatible // indirect github.com/docker/go-connections v0.6.0 // indirect github.com/docker/go-units v0.5.0 // indirect @@ -155,6 +156,7 @@ require ( github.com/testcontainers/testcontainers-go v0.40.0 // indirect github.com/tidwall/match v1.1.1 // indirect github.com/tidwall/pretty v1.2.0 // indirect + github.com/tiktoken-go/tokenizer v0.8.0 // indirect github.com/tklauser/go-sysconf v0.3.12 // indirect github.com/tklauser/numcpus v0.6.1 // indirect github.com/twitchyliquid64/golang-asm v0.15.1 // indirect diff --git a/backend/go.sum b/backend/go.sum index fbc04494ce..1574798015 100644 --- a/backend/go.sum +++ b/backend/go.sum @@ -108,6 +108,8 @@ github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/r github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc= github.com/distribution/reference v0.6.0 h1:0IXCQ5g4/QMHHkarYzh5l+u8T3t73zM5QvfrDyIgxBk= github.com/distribution/reference v0.6.0/go.mod h1:BbU0aIcezP1/5jX/8MP0YiH4SdvB5Y4f/wlDRiLyi3E= +github.com/dlclark/regexp2/v2 v2.1.0 h1:jHXRmHRZGbuQzDZjMlCAXOvQb75iv3HyLDzXGj5H1AY= +github.com/dlclark/regexp2/v2 v2.1.0/go.mod h1:Bz5TMy5d8fPK0ximH0Yi9KvsRHNnvXqUx9XG6a4wB+I= github.com/docker/docker v28.5.1+incompatible h1:Bm8DchhSD2J6PsFzxC35TZo4TLGR2PdW/E69rU45NhM= github.com/docker/docker v28.5.1+incompatible/go.mod h1:eEKB0N0r5NX/I1kEveEz05bcu8tLC/8azJZsviup8Sk= github.com/docker/go-connections v0.6.0 h1:LlMG9azAe1TqfR7sO+NJttz1gy6KO7VJBh+pMmjSD94= @@ -360,6 +362,8 @@ github.com/tidwall/pretty v1.2.0 h1:RWIZEg2iJ8/g6fDDYzMpobmaoGh5OLl4AXtGUGPcqCs= github.com/tidwall/pretty v1.2.0/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU= github.com/tidwall/sjson v1.2.5 h1:kLy8mja+1c9jlljvWTlSazM7cKDRfJuR/bOJhcY5NcY= github.com/tidwall/sjson v1.2.5/go.mod h1:Fvgq9kS/6ociJEDnK0Fk1cpYF4FIW6ZF7LAe+6jwd28= +github.com/tiktoken-go/tokenizer v0.8.0 h1:drHWno2Zx3eAm/hk/LmvBKXPpSImB7BRyh/ru4+3Q7Y= +github.com/tiktoken-go/tokenizer v0.8.0/go.mod h1:pTmPz4r14MV3JkUGAmAcdLdYhSxN68MCjrP+EoxBdx0= github.com/tklauser/go-sysconf v0.3.12 h1:0QaGUFOdQaIVdPgfITYzaTegZvdCjmYO52cSFAEVmqU= github.com/tklauser/go-sysconf v0.3.12/go.mod h1:Ho14jnntGE1fpdOqQEEaiKRpvIavV0hSfmBq8nJbHYI= github.com/tklauser/numcpus v0.6.1 h1:ng9scYS7az0Bk4OZLvrNXNSAO2Pxr1XXRAPyjhIx+Fk= diff --git a/backend/internal/service/openai_gateway_count_tokens.go b/backend/internal/service/openai_gateway_count_tokens.go index dda1c8e648..4a01b143e9 100644 --- a/backend/internal/service/openai_gateway_count_tokens.go +++ b/backend/internal/service/openai_gateway_count_tokens.go @@ -13,9 +13,16 @@ import ( "github.com/Wei-Shaw/sub2api/internal/pkg/logger" "github.com/gin-gonic/gin" "github.com/tidwall/gjson" + "github.com/tiktoken-go/tokenizer" "go.uber.org/zap" ) +const ( + openAIResponsesInputItemTokenOverhead = 3 + openAIResponsesContentPartOverhead = 1 + openAIInputTokensFallbackMinimum = 1 +) + type openAIInputTokensCountRequest struct { Model string `json:"model"` Instructions string `json:"instructions,omitempty"` @@ -24,6 +31,14 @@ type openAIInputTokensCountRequest struct { ToolChoice json.RawMessage `json:"tool_choice,omitempty"` } +type openAIInputTokensCountPrepared struct { + Request openAIInputTokensCountRequest + OriginalModel string + NormalizedModel string + BillingModel string + UpstreamModel string +} + // ForwardCountTokensAsAnthropic bridges Anthropic /v1/messages/count_tokens to // OpenAI POST /v1/responses/input_tokens and returns Anthropic-compatible output. func (s *OpenAIGatewayService) ForwardCountTokensAsAnthropic( @@ -38,31 +53,13 @@ func (s *OpenAIGatewayService) ForwardCountTokensAsAnthropic( return fmt.Errorf("count_tokens: missing account") } - var anthropicReq apicompat.AnthropicRequest - if err := json.Unmarshal(body, &anthropicReq); err != nil { - writeAnthropicCountTokensError(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body") - return fmt.Errorf("parse anthropic count_tokens request: %w", err) - } - - originalModel := anthropicReq.Model - applyOpenAICompatModelNormalization(&anthropicReq) - normalizedModel := anthropicReq.Model - billingModel := resolveOpenAIForwardModel(account, normalizedModel, strings.TrimSpace(defaultMappedModel)) - upstreamModel := normalizeOpenAIModelForUpstream(account, billingModel) - - responsesReq, err := apicompat.AnthropicToResponses(&anthropicReq) + prepared, err := prepareOpenAIInputTokensCountRequest(body, account, defaultMappedModel) if err != nil { - writeAnthropicCountTokensError(c, http.StatusBadRequest, "invalid_request_error", "Failed to convert request body") - return fmt.Errorf("convert anthropic request to responses: %w", err) + writeAnthropicCountTokensError(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body") + return err } - upstreamBody, err := marshalOpenAIUpstreamJSON(openAIInputTokensCountRequest{ - Model: upstreamModel, - Instructions: responsesReq.Instructions, - Input: responsesReq.Input, - Tools: responsesReq.Tools, - ToolChoice: responsesReq.ToolChoice, - }) + upstreamBody, err := marshalOpenAIUpstreamJSON(prepared.Request) if err != nil { writeAnthropicCountTokensError(c, http.StatusInternalServerError, "api_error", "Failed to build request") return fmt.Errorf("marshal openai input_tokens body: %w", err) @@ -70,10 +67,10 @@ func (s *OpenAIGatewayService) ForwardCountTokensAsAnthropic( logger.L().Debug("openai count_tokens: model mapping applied", zap.Int64("account_id", account.ID), - zap.String("original_model", originalModel), - zap.String("normalized_model", normalizedModel), - zap.String("billing_model", billingModel), - zap.String("upstream_model", upstreamModel), + zap.String("original_model", prepared.OriginalModel), + zap.String("normalized_model", prepared.NormalizedModel), + zap.String("billing_model", prepared.BillingModel), + zap.String("upstream_model", prepared.UpstreamModel), ) token, _, err := s.GetAccessToken(ctx, account) @@ -108,15 +105,16 @@ func (s *OpenAIGatewayService) ForwardCountTokensAsAnthropic( } if resp.StatusCode >= 400 { + upstreamMsg := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractUpstreamErrorMessage(respBody))) + if account.Type == AccountTypeOAuth && isOpenAIOAuthInputTokensUnsupported(resp.StatusCode, respBody) { + writeOpenAIOAuthInputTokensFallback(c, account, prepared, resp.StatusCode) + return nil + } + if s.rateLimitService != nil { s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody) } - upstreamMsg := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractUpstreamErrorMessage(respBody))) - if account.Type == AccountTypeOAuth && isOpenAIOAuthInputTokensUnsupported(resp.StatusCode) { - writeAnthropicCountTokensError(c, http.StatusNotFound, "not_found_error", "Token counting is not supported for this OpenAI account type") - return nil - } if isOpenAIInputTokensUnsupported(resp.StatusCode, respBody) { writeAnthropicCountTokensError(c, http.StatusNotFound, "not_found_error", "Token counting is not supported by upstream") return nil @@ -158,6 +156,42 @@ func (s *OpenAIGatewayService) ForwardCountTokensAsAnthropic( return nil } +func prepareOpenAIInputTokensCountRequest( + body []byte, + account *Account, + defaultMappedModel string, +) (*openAIInputTokensCountPrepared, error) { + var anthropicReq apicompat.AnthropicRequest + if err := json.Unmarshal(body, &anthropicReq); err != nil { + return nil, fmt.Errorf("parse anthropic count_tokens request: %w", err) + } + + originalModel := anthropicReq.Model + applyOpenAICompatModelNormalization(&anthropicReq) + normalizedModel := anthropicReq.Model + billingModel := resolveOpenAIForwardModel(account, normalizedModel, strings.TrimSpace(defaultMappedModel)) + upstreamModel := normalizeOpenAIModelForUpstream(account, billingModel) + + responsesReq, err := apicompat.AnthropicToResponses(&anthropicReq) + if err != nil { + return nil, fmt.Errorf("convert anthropic request to responses: %w", err) + } + + return &openAIInputTokensCountPrepared{ + Request: openAIInputTokensCountRequest{ + Model: upstreamModel, + Instructions: responsesReq.Instructions, + Input: responsesReq.Input, + Tools: responsesReq.Tools, + ToolChoice: responsesReq.ToolChoice, + }, + OriginalModel: originalModel, + NormalizedModel: normalizedModel, + BillingModel: billingModel, + UpstreamModel: upstreamModel, + }, nil +} + func (s *OpenAIGatewayService) buildInputTokensUpstreamRequest( ctx context.Context, c *gin.Context, @@ -218,11 +252,267 @@ func isOpenAIInputTokensUnsupported(statusCode int, body []byte) bool { return strings.Contains(msg, "input_tokens") && strings.Contains(msg, "not found") } -func isOpenAIOAuthInputTokensUnsupported(statusCode int) bool { +func writeOpenAIOAuthInputTokensFallback(c *gin.Context, account *Account, prepared *openAIInputTokensCountPrepared, statusCode int) { + estimated := openAIInputTokensFallbackMinimum + if got, err := estimateOpenAIInputTokens(prepared.Request); err == nil { + if got > 0 { + estimated = got + } + logger.L().Info("openai count_tokens: oauth fallback to local tiktoken estimate", + zap.Int64("account_id", account.ID), + zap.Int("upstream_status", statusCode), + zap.Int("estimated_input_tokens", estimated), + zap.String("upstream_model", prepared.UpstreamModel), + ) + } else { + logger.L().Warn("openai count_tokens: oauth local tiktoken fallback failed, using minimum estimate", + zap.Int64("account_id", account.ID), + zap.Int("upstream_status", statusCode), + zap.Int("estimated_input_tokens", estimated), + zap.String("upstream_model", prepared.UpstreamModel), + zap.Error(err), + ) + } + + c.JSON(http.StatusOK, gin.H{ + "input_tokens": estimated, + }) +} + +func isOpenAIOAuthInputTokensUnsupported(statusCode int, body []byte) bool { switch statusCode { case http.StatusUnauthorized, http.StatusForbidden, http.StatusNotFound: - return true default: return false } + + bodyLower := strings.ToLower(string(body)) + msg := strings.ToLower(strings.TrimSpace(extractUpstreamErrorMessage(body))) + code := strings.ToLower(strings.TrimSpace(extractUpstreamErrorCode(body))) + + if code == "missing_scope" || + strings.Contains(bodyLower, "api.responses.write") || + strings.Contains(bodyLower, "missing scopes") || + strings.Contains(bodyLower, "insufficient_scope") { + return true + } + + if statusCode == http.StatusNotFound && isOpenAIInputTokensUnsupported(statusCode, body) { + return true + } + + return strings.Contains(msg, "input_tokens") && + (strings.Contains(msg, "not found") || + strings.Contains(msg, "not supported") || + strings.Contains(msg, "unsupported")) +} + +func estimateOpenAIInputTokens(req openAIInputTokensCountRequest) (int, error) { + codec, err := openAIInputTokensCodecForModel(req.Model) + if err != nil { + return 0, err + } + + total := 0 + addCount := func(text string) error { + text = strings.TrimSpace(text) + if text == "" { + return nil + } + n, err := codec.Count(text) + if err != nil { + return err + } + total += n + return nil + } + + if err := addCount(req.Instructions); err != nil { + return 0, err + } + inputTokens, err := estimateOpenAIInputTokensForInput(codec, req.Input) + if err != nil { + return 0, err + } + total += inputTokens + + for _, tool := range req.Tools { + raw, err := marshalOpenAIUpstreamJSON(tool) + if err != nil { + return 0, err + } + if err := addCount(string(raw)); err != nil { + return 0, err + } + } + if len(req.ToolChoice) > 0 { + compacted, err := compactOpenAIInputTokensJSON(req.ToolChoice) + if err != nil { + return 0, err + } + if err := addCount(compacted); err != nil { + return 0, err + } + } + + if total < 0 { + return 0, nil + } + return total, nil +} + +func estimateOpenAIInputTokensForInput(codec tokenizer.Codec, raw json.RawMessage) (int, error) { + if len(bytes.TrimSpace(raw)) == 0 { + return 0, nil + } + + var plainText string + if err := json.Unmarshal(raw, &plainText); err == nil { + return codec.Count(plainText) + } + + var items []apicompat.ResponsesInputItem + if err := json.Unmarshal(raw, &items); err == nil { + return estimateOpenAIInputTokensForInputItems(codec, items) + } + + compacted, err := compactOpenAIInputTokensJSON(raw) + if err != nil { + return 0, err + } + return codec.Count(compacted) +} + +func estimateOpenAIInputTokensForInputItems(codec tokenizer.Codec, items []apicompat.ResponsesInputItem) (int, error) { + total := 0 + countText := func(text string) error { + text = strings.TrimSpace(text) + if text == "" { + return nil + } + n, err := codec.Count(text) + if err != nil { + return err + } + total += n + return nil + } + + for _, item := range items { + total += openAIResponsesInputItemTokenOverhead + if err := countText(item.Role); err != nil { + return 0, err + } + if item.Type != "" && item.Type != "message" { + if err := countText(item.Type); err != nil { + return 0, err + } + } + if err := countText(item.Name); err != nil { + return 0, err + } + if err := countText(item.Arguments); err != nil { + return 0, err + } + if err := countText(item.Output); err != nil { + return 0, err + } + if err := countText(item.CallID); err != nil { + return 0, err + } + if err := countText(item.ID); err != nil { + return 0, err + } + + if len(bytes.TrimSpace(item.Content)) == 0 { + continue + } + + var contentText string + if err := json.Unmarshal(item.Content, &contentText); err == nil { + if err := countText(contentText); err != nil { + return 0, err + } + continue + } + + var parts []apicompat.ResponsesContentPart + if err := json.Unmarshal(item.Content, &parts); err == nil { + for _, part := range parts { + total += openAIResponsesContentPartOverhead + switch part.Type { + case "input_text", "output_text", "text": + if err := countText(part.Text); err != nil { + return 0, err + } + case "input_image": + if err := countText(estimateOpenAIInputImageText(part.ImageURL)); err != nil { + return 0, err + } + default: + if err := countText(part.Type); err != nil { + return 0, err + } + } + } + continue + } + + compacted, err := compactOpenAIInputTokensJSON(item.Content) + if err != nil { + return 0, err + } + if err := countText(compacted); err != nil { + return 0, err + } + } + + return total, nil +} + +func estimateOpenAIInputImageText(imageURL string) string { + trimmed := strings.TrimSpace(imageURL) + if trimmed == "" { + return "" + } + if strings.HasPrefix(strings.ToLower(trimmed), "data:") { + if comma := strings.Index(trimmed, ","); comma > 0 { + return trimmed[:comma] + } + } + return trimmed +} + +func compactOpenAIInputTokensJSON(raw json.RawMessage) (string, error) { + if len(bytes.TrimSpace(raw)) == 0 { + return "", nil + } + var buf bytes.Buffer + if err := json.Compact(&buf, raw); err != nil { + return "", err + } + return buf.String(), nil +} + +func openAIInputTokensCodecForModel(model string) (tokenizer.Codec, error) { + switch openAIInputTokensEncodingForModel(model) { + case tokenizer.Cl100kBase: + return tokenizer.Get(tokenizer.Cl100kBase) + default: + return tokenizer.Get(tokenizer.O200kBase) + } +} + +func openAIInputTokensEncodingForModel(model string) tokenizer.Encoding { + normalized := strings.ToLower(strings.TrimSpace(model)) + switch { + case strings.HasPrefix(normalized, "gpt-3.5"), + (strings.HasPrefix(normalized, "gpt-4") && + !strings.HasPrefix(normalized, "gpt-4o") && + !strings.HasPrefix(normalized, "gpt-4.1")), + strings.HasPrefix(normalized, "text-embedding-"): + return tokenizer.Cl100kBase + default: + return tokenizer.O200kBase + } } diff --git a/backend/internal/service/openai_gateway_count_tokens_test.go b/backend/internal/service/openai_gateway_count_tokens_test.go index dac8291e66..3ab0178a0d 100644 --- a/backend/internal/service/openai_gateway_count_tokens_test.go +++ b/backend/internal/service/openai_gateway_count_tokens_test.go @@ -3,18 +3,40 @@ package service import ( "bytes" "context" + "encoding/json" + "fmt" "io" "net/http" "net/http/httptest" + "os" + "strconv" "strings" "testing" + "time" "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/apicompat" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" "github.com/tidwall/gjson" ) +type countTokensRuntimeStateRepo struct { + AccountRepository + tempUnschedCalls int + setErrorCalls int +} + +func (r *countTokensRuntimeStateRepo) SetTempUnschedulable(_ context.Context, _ int64, _ time.Time, _ string) error { + r.tempUnschedCalls++ + return nil +} + +func (r *countTokensRuntimeStateRepo) SetError(_ context.Context, _ int64, _ string) error { + r.setErrorCalls++ + return nil +} + func TestOpenAIGatewayService_ForwardCountTokensAsAnthropic_APIKeyUsesResponsesInputTokens(t *testing.T) { gin.SetMode(gin.TestMode) @@ -66,23 +88,7 @@ func TestOpenAIGatewayService_ForwardCountTokensAsAnthropic_APIKeyUsesResponsesI func TestOpenAIGatewayService_ForwardCountTokensAsAnthropic_OAuthFallsBackWhenPlatformEndpointUnsupported(t *testing.T) { gin.SetMode(gin.TestMode) - rec := httptest.NewRecorder() - c, _ := gin.CreateTestContext(rec) body := []byte(`{"model":"claude-opus-4-1","messages":[{"role":"user","content":"hello"}]}`) - c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages/count_tokens", bytes.NewReader(body)) - c.Request.Header.Set("Content-Type", "application/json") - c.Request.Header.Set("User-Agent", "Claude-Code/1.0") - - upstream := &httpUpstreamRecorder{resp: &http.Response{ - StatusCode: http.StatusUnauthorized, - Header: http.Header{"Content-Type": []string{"application/json"}}, - Body: io.NopCloser(strings.NewReader(`{"error":{"type":"invalid_request_error","message":"unauthorized"}}`)), - }} - - svc := &OpenAIGatewayService{ - cfg: &config.Config{}, - httpUpstream: upstream, - } account := &Account{ ID: 202, Name: "openai-oauth", @@ -90,18 +96,254 @@ func TestOpenAIGatewayService_ForwardCountTokensAsAnthropic_OAuthFallsBackWhenPl Type: AccountTypeOAuth, Concurrency: 1, Credentials: map[string]any{ - "access_token": "oauth-token", + "access_token": "oauth-token", + "refresh_token": "oauth-refresh-token", }, Status: StatusActive, Schedulable: true, } - err := svc.ForwardCountTokensAsAnthropic(context.Background(), c, account, body, "gpt-5.4") + prepared, err := prepareOpenAIInputTokensCountRequest(body, account, "gpt-5.4") require.NoError(t, err) - require.Equal(t, http.StatusNotFound, rec.Code) - require.Contains(t, rec.Body.String(), "Token counting is not supported for this OpenAI account type") - require.NotNil(t, upstream.lastReq) - require.Equal(t, "https://api.openai.com/v1/responses/input_tokens", upstream.lastReq.URL.String()) - require.Equal(t, "Bearer oauth-token", upstream.lastReq.Header.Get("authorization")) - require.Empty(t, upstream.lastReq.Header.Get("Chatgpt-Account-Id")) + expectedEstimate, err := estimateOpenAIInputTokens(prepared.Request) + require.NoError(t, err) + + cases := []struct { + name string + statusCode int + body string + }{ + { + name: "401_missing_responses_write_scope", + statusCode: http.StatusUnauthorized, + body: `{"error":{"type":"invalid_request_error","code":"missing_scope","message":"You have insufficient permissions for this operation. Missing scopes: api.responses.write."}}`, + }, + { + name: "403_missing_responses_write_scope", + statusCode: http.StatusForbidden, + body: `{"error":{"type":"invalid_request_error","code":"missing_scope","message":"Missing scopes: api.responses.write"}}`, + }, + { + name: "404_input_tokens_unsupported", + statusCode: http.StatusNotFound, + body: `{"error":{"type":"invalid_request_error","message":"The /v1/responses/input_tokens endpoint was not found"}}`, + }, + } + + for _, tt := range cases { + t.Run(tt.name, func(t *testing.T) { + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages/count_tokens", bytes.NewReader(body)) + c.Request.Header.Set("Content-Type", "application/json") + c.Request.Header.Set("User-Agent", "Claude-Code/1.0") + + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: tt.statusCode, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(tt.body)), + }} + repo := &countTokensRuntimeStateRepo{} + svc := &OpenAIGatewayService{ + cfg: &config.Config{}, + httpUpstream: upstream, + rateLimitService: &RateLimitService{accountRepo: repo, cfg: &config.Config{}}, + } + + err := svc.ForwardCountTokensAsAnthropic(context.Background(), c, account, body, "gpt-5.4") + require.NoError(t, err) + require.Equal(t, http.StatusOK, rec.Code) + require.JSONEq(t, `{"input_tokens":`+strconv.Itoa(expectedEstimate)+`}`, rec.Body.String()) + require.NotNil(t, upstream.lastReq) + require.Equal(t, "https://api.openai.com/v1/responses/input_tokens", upstream.lastReq.URL.String()) + require.Equal(t, "Bearer oauth-token", upstream.lastReq.Header.Get("authorization")) + require.Empty(t, upstream.lastReq.Header.Get("Chatgpt-Account-Id")) + require.Zero(t, repo.tempUnschedCalls, "OAuth input_tokens unsupported errors must not temp-unschedule the account") + require.Zero(t, repo.setErrorCalls, "OAuth input_tokens unsupported errors must not mark the account error") + }) + } +} + +func TestOpenAIGatewayService_OpenAIOAuthInputTokensFallbackUsesMinimumWhenEstimateFails(t *testing.T) { + gin.SetMode(gin.TestMode) + + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + prepared := &openAIInputTokensCountPrepared{ + Request: openAIInputTokensCountRequest{ + Model: "gpt-5", + Input: json.RawMessage(`[`), + }, + UpstreamModel: "gpt-5", + } + + writeOpenAIOAuthInputTokensFallback(c, &Account{ID: 303}, prepared, http.StatusUnauthorized) + + require.Equal(t, http.StatusOK, rec.Code) + require.JSONEq(t, `{"input_tokens":1}`, rec.Body.String()) +} + +func TestEstimateOpenAIInputTokens_RequestSamples(t *testing.T) { + cases := []struct { + name string + req openAIInputTokensCountRequest + want int + }{ + { + name: "simple text input", + req: openAIInputTokensCountRequest{ + Model: "gpt-5", + Input: json.RawMessage(`[{"role":"user","content":"hello world"}]`), + }, + want: 6, + }, + { + name: "instructions plus tool schema", + req: openAIInputTokensCountRequest{ + Model: "gpt-5", + Instructions: "You are helpful.", + Input: json.RawMessage(`[{"role":"user","content":"lookup weather in shanghai"}]`), + Tools: []apicompat.ResponsesTool{ + { + Type: "function", + Name: "lookup_weather", + Description: "Look up current weather", + Parameters: json.RawMessage(`{"type":"object","properties":{"city":{"type":"string"}},"required":["city"]}`), + }, + }, + }, + want: 50, + }, + { + name: "input parts and tool output", + req: openAIInputTokensCountRequest{ + Model: "gpt-4.1", + Input: json.RawMessage(`[ + {"role":"user","content":[{"type":"input_text","text":"first line"},{"type":"input_text","text":"second line"}]}, + {"type":"function_call_output","call_id":"call_123","output":"{\"ok\":true}"} + ]`), + }, + want: 24, + }, + } + + for _, tt := range cases { + t.Run(tt.name, func(t *testing.T) { + got, err := estimateOpenAIInputTokens(tt.req) + require.NoError(t, err) + require.Equal(t, tt.want, got) + }) + } +} + +func TestOpenAIInputTokensEncodingForModel(t *testing.T) { + cases := []struct { + model string + want string + }{ + {model: "gpt-5", want: "o200k_base"}, + {model: "gpt-5.3-codex", want: "o200k_base"}, + {model: "gpt-4o-mini", want: "o200k_base"}, + {model: "gpt-4.1", want: "o200k_base"}, + {model: "gpt-4-turbo", want: "cl100k_base"}, + {model: "gpt-3.5-turbo", want: "cl100k_base"}, + } + + for _, tt := range cases { + t.Run(tt.model, func(t *testing.T) { + require.Equal(t, tt.want, string(openAIInputTokensEncodingForModel(tt.model))) + }) + } +} + +func TestEstimateOpenAIInputTokens_CompareWithOpenAIAPI(t *testing.T) { + apiKey := strings.TrimSpace(os.Getenv("OPENAI_API_KEY")) + if apiKey == "" { + t.Skip("OPENAI_API_KEY not set") + } + + client := &http.Client{Timeout: 30 * time.Second} + cases := []struct { + name string + anthropicBody []byte + defaultOpenAIModel string + }{ + { + name: "simple user text", + defaultOpenAIModel: "gpt-5", + anthropicBody: []byte(`{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"hello world from sub2api"}]}`), + }, + { + name: "system plus tool", + defaultOpenAIModel: "gpt-5", + anthropicBody: []byte(`{"model":"claude-sonnet-4-5","system":"You are helpful.","messages":[{"role":"user","content":"find weather in shanghai"}],"tools":[{"name":"lookup_weather","description":"Look up current weather","input_schema":{"type":"object","properties":{"city":{"type":"string"}},"required":["city"]}}]}`), + }, + { + name: "multi turn text", + defaultOpenAIModel: "gpt-4.1", + anthropicBody: []byte(`{"model":"claude-opus-4-1","messages":[{"role":"user","content":"summarize this repo"},{"role":"assistant","content":"which repo?"},{"role":"user","content":"sub2api"}]}`), + }, + } + + account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey} + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + prepared, err := prepareOpenAIInputTokensCountRequest(tc.anthropicBody, account, tc.defaultOpenAIModel) + require.NoError(t, err) + + estimated, err := estimateOpenAIInputTokens(prepared.Request) + require.NoError(t, err) + + actual, err := callOpenAIInputTokensAPIForTest(client, apiKey, prepared.Request) + require.NoError(t, err) + + diff := estimated - actual + if diff < 0 { + diff = -diff + } + t.Logf("model=%s estimated=%d actual=%d diff=%d", prepared.Request.Model, estimated, actual, diff) + require.LessOrEqual(t, diff, maxLocalInt(24, actual/4)) + }) + } +} + +func callOpenAIInputTokensAPIForTest(client *http.Client, apiKey string, reqBody openAIInputTokensCountRequest) (int, error) { + body, err := marshalOpenAIUpstreamJSON(reqBody) + if err != nil { + return 0, err + } + req, err := http.NewRequest(http.MethodPost, openaiPlatformAPIInputTokensURL, bytes.NewReader(body)) + if err != nil { + return 0, err + } + req.Header.Set("authorization", "Bearer "+apiKey) + req.Header.Set("content-type", "application/json") + req.Header.Set("accept", "application/json") + + resp, err := client.Do(req) + if err != nil { + return 0, err + } + defer func() { _ = resp.Body.Close() }() + + respBody, err := io.ReadAll(resp.Body) + if err != nil { + return 0, err + } + if resp.StatusCode >= 400 { + return 0, fmt.Errorf("openai input_tokens api error: status=%d body=%s", resp.StatusCode, string(respBody)) + } + + value := gjson.GetBytes(respBody, "input_tokens") + if !value.Exists() { + return 0, fmt.Errorf("openai input_tokens api missing input_tokens: %s", string(respBody)) + } + return int(value.Int()), nil +} + +func maxLocalInt(a, b int) int { + if a > b { + return a + } + return b }