diff --git a/backend/internal/pkg/xai/usage.go b/backend/internal/pkg/xai/usage.go new file mode 100644 index 0000000000..77aa382900 --- /dev/null +++ b/backend/internal/pkg/xai/usage.go @@ -0,0 +1,25 @@ +package xai + +// IncludeIndependentReasoningTokens adds reasoning tokens to billed output +// only when total_tokens proves they are not already folded into output. +// +// Official xAI Chat Completions example: prompt=32, completion=9, +// reasoning=94, total=135. Responses example: input=32, output=9, +// reasoning=110, total=151. OpenAI's canonical completion_tokens already +// includes reasoning and total equals input+output. +func IncludeIndependentReasoningTokens(input, output, total, reasoning int64) int64 { + if input < 0 || output < 0 || reasoning <= 0 || total <= 0 { + return output + } + if total == input+output { + return output + } + gap := total - input - output + if gap <= 0 { + return output + } + if reasoning < gap { + gap = reasoning + } + return output + gap +} diff --git a/backend/internal/pkg/xai/usage_test.go b/backend/internal/pkg/xai/usage_test.go new file mode 100644 index 0000000000..79b8f8639c --- /dev/null +++ b/backend/internal/pkg/xai/usage_test.go @@ -0,0 +1,30 @@ +package xai + +import "testing" + +func TestIncludeIndependentReasoningTokens(t *testing.T) { + t.Parallel() + tests := []struct { + name string + input, output, total, reason int64 + want int64 + }{ + {name: "xAI chat example", input: 32, output: 9, total: 135, reason: 94, want: 103}, + {name: "xAI responses example", input: 32, output: 9, total: 151, reason: 110, want: 119}, + {name: "OpenAI inclusive output", input: 32, output: 103, total: 135, reason: 94, want: 103}, + {name: "gap smaller than reasoning", input: 10, output: 20, total: 33, reason: 5, want: 23}, + {name: "absent total", input: 10, output: 20, total: 0, reason: 5, want: 20}, + {name: "total already complete", input: 10, output: 20, total: 30, reason: 5, want: 20}, + {name: "inconsistent undershoot", input: 10, output: 20, total: 25, reason: 5, want: 20}, + {name: "zero reasoning", input: 10, output: 20, total: 35, reason: 0, want: 20}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + got := IncludeIndependentReasoningTokens(tt.input, tt.output, tt.total, tt.reason) + if got != tt.want { + t.Fatalf("got %d, want %d", got, tt.want) + } + }) + } +} diff --git a/backend/internal/service/grok_media.go b/backend/internal/service/grok_media.go index 3d7ad1fe31..85dd1b14ec 100644 --- a/backend/internal/service/grok_media.go +++ b/backend/internal/service/grok_media.go @@ -59,6 +59,8 @@ type GrokMediaRequestInfo struct { N int Size string SizeTier string + AspectRatio string + ImageResolution string Resolution string DurationSeconds int InputImageURLs []string @@ -127,6 +129,8 @@ func ParseGrokMediaRequest(contentType string, body []byte) GrokMediaRequestInfo info.Prompt = strings.TrimSpace(info.Prompt) info.Size = strings.TrimSpace(info.Size) info.SizeTier = NormalizeImageBillingTierOrDefault(info.Size) + info.AspectRatio = strings.TrimSpace(info.AspectRatio) + info.ImageResolution = grokImagineImageResolution(info.ImageResolution) info.Resolution = NormalizeVideoBillingResolutionOrDefault(info.Resolution) info.DurationSeconds = NormalizeVideoBillingDurationSecondsOrDefault(info.DurationSeconds) if info.N <= 0 { @@ -142,7 +146,8 @@ func parseGrokMediaJSONRequest(body []byte, info *GrokMediaRequestInfo) { info.Model = strings.TrimSpace(gjson.GetBytes(body, "model").String()) info.Prompt = strings.TrimSpace(gjson.GetBytes(body, "prompt").String()) info.Size = strings.TrimSpace(gjson.GetBytes(body, "size").String()) - info.Resolution = strings.TrimSpace(gjson.GetBytes(body, "resolution").String()) + info.AspectRatio = strings.TrimSpace(gjson.GetBytes(body, "aspect_ratio").String()) + assignGrokMediaResolution(strings.TrimSpace(gjson.GetBytes(body, "resolution").String()), info) if duration := gjson.GetBytes(body, "duration"); duration.Exists() && duration.Type == gjson.Number { info.DurationSeconds = int(duration.Int()) } @@ -255,8 +260,10 @@ func parseGrokMediaMultipartRequest(contentType string, body []byte, info *GrokM info.Prompt = value case "size": info.Size = value + case "aspect_ratio": + info.AspectRatio = value case "resolution": - info.Resolution = value + assignGrokMediaResolution(value, info) case "duration": if duration, err := strconv.Atoi(value); err == nil { info.DurationSeconds = duration @@ -934,6 +941,12 @@ func prepareGrokMediaForwardBody(endpoint GrokMediaEndpoint, body []byte, conten if info.Size != "" { payload["size"] = info.Size } + if info.ImageResolution != "" { + payload["resolution"] = info.ImageResolution + } + if info.AspectRatio != "" { + payload["aspect_ratio"] = info.AspectRatio + } images := make([]map[string]string, 0, len(info.InputImageURLs)+len(info.Uploads)) for _, imageURL := range info.InputImageURLs { @@ -1106,10 +1119,7 @@ func sanitizeGrokMediaForwardBody(endpoint GrokMediaEndpoint, body []byte, conte } switch endpoint { case GrokMediaEndpointImagesGenerations, GrokMediaEndpointImagesEdits: - if !gjson.GetBytes(body, "size").Exists() { - return body, contentType, nil - } - out, err := sjson.DeleteBytes(body, "size") + out, err := applyGrokImagineImageGeometry(body) if err != nil { return nil, "", fmt.Errorf("sanitize grok media size: %w", err) } diff --git a/backend/internal/service/grok_media_image_geometry.go b/backend/internal/service/grok_media_image_geometry.go new file mode 100644 index 0000000000..51e2a6fd69 --- /dev/null +++ b/backend/internal/service/grok_media_image_geometry.go @@ -0,0 +1,149 @@ +package service + +import ( + "math" + "strconv" + "strings" + + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" +) + +// Official Imagine image geometry: https://docs.x.ai/developers/model-capabilities/images/generation +var grokImagineAspectRatioValues = []struct { + label string + ratio float64 +}{ + {"1:1", 1}, + {"16:9", 16.0 / 9.0}, + {"9:16", 9.0 / 16.0}, + {"4:3", 4.0 / 3.0}, + {"3:4", 3.0 / 4.0}, + {"3:2", 1.5}, + {"2:3", 2.0 / 3.0}, + {"2:1", 2}, + {"1:2", 0.5}, + {"19.5:9", 19.5 / 9.0}, + {"9:19.5", 9.0 / 19.5}, + {"20:9", 20.0 / 9.0}, + {"9:20", 9.0 / 20.0}, +} + +func applyGrokImagineImageGeometry(body []byte) ([]byte, error) { + size := strings.TrimSpace(gjson.GetBytes(body, "size").String()) + resolution := grokImagineImageResolution(gjson.GetBytes(body, "resolution").String()) + aspect := strings.TrimSpace(gjson.GetBytes(body, "aspect_ratio").String()) + out := append([]byte(nil), body...) + + if resolution == "" { + if derived := grokImagineImageResolutionFromSize(size); derived != "" { + next, err := sjson.SetBytes(out, "resolution", derived) + if err != nil { + return nil, err + } + out = next + } + } else if gjson.GetBytes(body, "resolution").String() != resolution { + next, err := sjson.SetBytes(out, "resolution", resolution) + if err != nil { + return nil, err + } + out = next + } + + if aspect == "" { + if derived := grokImagineAspectRatioFromSize(size); derived != "" { + next, err := sjson.SetBytes(out, "aspect_ratio", derived) + if err != nil { + return nil, err + } + out = next + } + } + + if !gjson.GetBytes(out, "size").Exists() { + return out, nil + } + return sjson.DeleteBytes(out, "size") +} + +func assignGrokMediaResolution(value string, info *GrokMediaRequestInfo) { + if info == nil { + return + } + value = strings.TrimSpace(value) + if value == "" { + return + } + if img := grokImagineImageResolution(value); img != "" { + info.ImageResolution = img + return + } + info.Resolution = value +} + +func grokImagineImageResolution(value string) string { + switch strings.ToLower(strings.TrimSpace(value)) { + case "1k": + return "1k" + case "2k": + return "2k" + default: + return "" + } +} + +func grokImagineImageResolutionFromSize(size string) string { + if explicit := grokImagineImageResolution(size); explicit != "" { + return explicit + } + tier, ok := ClassifyImageBillingTier(size) + if !ok { + return "" + } + if tier == ImageBillingSize1K { + return "1k" + } + return "2k" +} + +func grokImagineAspectRatioFromSize(size string) string { + width, height, ok := parseImageBillingDimensions(strings.TrimSpace(size)) + if !ok || width <= 0 || height <= 0 { + return "" + } + div := grokImagineGCD(width, height) + exact := strconv.Itoa(width/div) + ":" + strconv.Itoa(height/div) + for _, candidate := range grokImagineAspectRatioValues { + if candidate.label == exact { + return exact + } + } + ratio := float64(width) / float64(height) + bestLabel := "" + bestDelta := math.MaxFloat64 + for _, candidate := range grokImagineAspectRatioValues { + delta := math.Abs(ratio - candidate.ratio) + if delta < bestDelta { + bestDelta = delta + bestLabel = candidate.label + } + } + return bestLabel +} + +func grokImagineGCD(a, b int) int { + if a < 0 { + a = -a + } + if b < 0 { + b = -b + } + for b != 0 { + a, b = b, a%b + } + if a == 0 { + return 1 + } + return a +} diff --git a/backend/internal/service/grok_media_image_geometry_test.go b/backend/internal/service/grok_media_image_geometry_test.go new file mode 100644 index 0000000000..cac26c3bd1 --- /dev/null +++ b/backend/internal/service/grok_media_image_geometry_test.go @@ -0,0 +1,60 @@ +package service + +import ( + "testing" + + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" +) + +func TestApplyGrokImagineImageGeometryMapsOpenAISize(t *testing.T) { + t.Parallel() + + out, err := applyGrokImagineImageGeometry([]byte(`{"model":"grok-imagine-image-2.0","prompt":"hi","size":"1152x1536"}`)) + require.NoError(t, err) + require.False(t, gjson.GetBytes(out, "size").Exists()) + require.Equal(t, "2k", gjson.GetBytes(out, "resolution").String()) + require.Equal(t, "3:4", gjson.GetBytes(out, "aspect_ratio").String()) +} + +func TestApplyGrokImagineImageGeometryKeepsClientGeometry(t *testing.T) { + t.Parallel() + + out, err := applyGrokImagineImageGeometry([]byte(`{"size":"1024x1024","resolution":"2K","aspect_ratio":"16:9"}`)) + require.NoError(t, err) + require.False(t, gjson.GetBytes(out, "size").Exists()) + require.Equal(t, "2k", gjson.GetBytes(out, "resolution").String()) + require.Equal(t, "16:9", gjson.GetBytes(out, "aspect_ratio").String()) +} + +func TestSanitizeGrokMediaForwardBodyConvertsImageSize(t *testing.T) { + t.Parallel() + + out, contentType, err := sanitizeGrokMediaForwardBody( + GrokMediaEndpointImagesGenerations, + []byte(`{"model":"grok-imagine-image","prompt":"hi","size":"1024x1024"}`), + "application/json", + ) + require.NoError(t, err) + require.Equal(t, "application/json", contentType) + require.False(t, gjson.GetBytes(out, "size").Exists()) + require.Equal(t, "1k", gjson.GetBytes(out, "resolution").String()) + require.Equal(t, "1:1", gjson.GetBytes(out, "aspect_ratio").String()) +} + +func TestParseGrokMediaRequestKeepsImageResolutionOutOfVideoNormalize(t *testing.T) { + t.Parallel() + + info := ParseGrokMediaRequest("application/json", []byte(`{"model":"grok-imagine-image-2.0","resolution":"2K","aspect_ratio":"16:9"}`)) + require.Equal(t, "2k", info.ImageResolution) + require.Equal(t, "16:9", info.AspectRatio) + require.Equal(t, VideoBillingResolution480P, info.Resolution) +} + +func TestGrokImagineAspectRatioFromSize(t *testing.T) { + t.Parallel() + require.Equal(t, "1:1", grokImagineAspectRatioFromSize("1024x1024")) + require.Equal(t, "3:4", grokImagineAspectRatioFromSize("1152x1536")) + require.Equal(t, "4:3", grokImagineAspectRatioFromSize("1536x1152")) + require.Equal(t, "16:9", grokImagineAspectRatioFromSize("1792x1024")) +} diff --git a/backend/internal/service/grok_upstream_errors.go b/backend/internal/service/grok_upstream_errors.go index 1d1bddc6d6..db204fc70b 100644 --- a/backend/internal/service/grok_upstream_errors.go +++ b/backend/internal/service/grok_upstream_errors.go @@ -114,8 +114,9 @@ func isGrokAccountAccessCode(value string) bool { "subscription_required", "entitlement_required", "not_entitled", - "plan_required", - "permission_denied": + "plan_required": + // permission-denied is omitted: xAI reuses it for both entitlement + // refusals and request-scoped safety blocks, so the message decides. return true default: return false @@ -171,6 +172,7 @@ func grokContentPolicyMessage(value string) bool { "prompt violates policy", "input violates content policy", "input violates policy", + "violates usage guidelines", } { if strings.Contains(lower, phrase) { return true diff --git a/backend/internal/service/grok_upstream_errors_test.go b/backend/internal/service/grok_upstream_errors_test.go index 0dee756b72..3bd5ad3e73 100644 --- a/backend/internal/service/grok_upstream_errors_test.go +++ b/backend/internal/service/grok_upstream_errors_test.go @@ -83,6 +83,24 @@ func TestIsGrokContentPolicyRejection(t *testing.T) { body: `{"error":{"code":"policy_violation","message":"request blocked by policy"}}`, want: true, }, + { + name: "permission-denied usage guidelines is request scoped", + status: http.StatusForbidden, + body: `{"code":"permission-denied","error":"Content violates usage guidelines. "}`, + want: true, + }, + { + name: "permission-denied entitlement stays on the account path", + status: http.StatusForbidden, + body: `{"code":"permission-denied","error":"Access to the chat endpoint is denied"}`, + want: false, + }, + { + name: "structured account code overrides usage guidelines phrase", + status: http.StatusForbidden, + body: `{"error":{"code":"account_suspended","message":"Content violates usage guidelines."}}`, + want: false, + }, { name: "wrong status", status: http.StatusBadRequest, @@ -279,6 +297,21 @@ func TestGrokContentPolicySSEErrorDoesNotMutateOrFailover(t *testing.T) { require.False(t, svc.isOpenAIAccountRuntimeBlocked(account)) } +func TestGrokPermissionDeniedContentRefusalDoesNotMutateOrFailover(t *testing.T) { + repo := &grokQuotaAccountRepo{} + svc := &OpenAIGatewayService{accountRepo: repo} + account := &Account{ID: 4785, Platform: PlatformGrok, Type: AccountTypeOAuth} + body := []byte(`{"code":"permission-denied","error":"Content violates usage guidelines. "}`) + + svc.handleGrokAccountUpstreamError(context.Background(), account, http.StatusForbidden, nil, body) + + require.Zero(t, repo.tempUnschedCalls) + require.Zero(t, repo.rateLimitedCalls) + require.Zero(t, repo.updateCalls) + require.False(t, svc.isOpenAIAccountRuntimeBlocked(account)) + require.False(t, svc.shouldFailoverGrokUpstreamError(http.StatusForbidden, body)) +} + func TestHandleGrokAccountUpstreamErrorEntitlement403KeepsDefaultCooldown(t *testing.T) { repo := &grokQuotaAccountRepo{} svc := &OpenAIGatewayService{accountRepo: repo} diff --git a/backend/internal/service/openai_gateway_grok.go b/backend/internal/service/openai_gateway_grok.go index 9a16c6c7a7..cca3b28ab3 100644 --- a/backend/internal/service/openai_gateway_grok.go +++ b/backend/internal/service/openai_gateway_grok.go @@ -856,33 +856,98 @@ func grokResponsesToolDedupKey(tool gjson.Result) string { return "json:" + normalizeCompatSeedJSON(json.RawMessage(tool.Raw)) } -// sanitizeGrokReasoningNullContent 删除 reasoning 项中的 "content": null。 -// xAI 的 untagged enum 反序列化器拒收该字段，返回 422。 +// sanitizeGrokReasoningNullContent drops explicit JSON nulls from Responses +// input items. xAI's untagged ModelInput decoder 422s on those fields. +// Compaction items stay unmodified per the compact contract. func sanitizeGrokReasoningNullContent(body []byte) ([]byte, error) { input := gjson.GetBytes(body, "input") - if !input.Exists() || !input.IsArray() { + if !input.Exists() || (!input.IsArray() && !input.IsObject()) { return body, nil } - items := input.Array() + var decoded map[string]any + decoder := json.NewDecoder(bytes.NewReader(body)) + decoder.UseNumber() + if err := decoder.Decode(&decoded); err != nil { + return body, nil + } + rawInput, ok := decoded["input"] + if !ok { + return body, nil + } + cleaned, changed := stripExplicitNullsFromGrokInput(rawInput) + if !changed { + return body, nil + } + decoded["input"] = cleaned + out, err := marshalOpenAIUpstreamJSON(decoded) + if err != nil { + return nil, err + } + return out, nil +} + +func stripExplicitNullsFromGrokInput(value any) (any, bool) { + switch node := value.(type) { + case []any: + changed := false + for i, item := range node { + itemMap, ok := item.(map[string]any) + if !ok { + next, childChanged := stripExplicitNullsFromGrokInput(item) + if childChanged { + node[i] = next + changed = true + } + continue + } + if isOpenAICompactionType(stringValue(itemMap["type"])) { + continue + } + next, childChanged := stripExplicitNullsFromJSONObject(itemMap) + if childChanged { + node[i] = next + changed = true + } + } + return node, changed + case map[string]any: + if isOpenAICompactionType(stringValue(node["type"])) { + return node, false + } + return stripExplicitNullsFromJSONObject(node) + default: + return value, false + } +} + +func stripExplicitNullsFromJSONObject(node map[string]any) (map[string]any, bool) { + if node == nil { + return node, false + } changed := false - for i := len(items) - 1; i >= 0; i-- { - item := items[i] - if strings.TrimSpace(item.Get("type").String()) != "reasoning" { + for key, child := range node { + if child == nil { + delete(node, key) + changed = true continue } - contentResult := item.Get("content") - if contentResult.Exists() && contentResult.Type == gjson.Null { - var err error - body, err = sjson.DeleteBytes(body, fmt.Sprintf("input.%d.content", i)) - if err != nil { - return nil, err + switch typed := child.(type) { + case map[string]any: + next, childChanged := stripExplicitNullsFromJSONObject(typed) + if childChanged { + node[key] = next + changed = true + } + case []any: + next, childChanged := stripExplicitNullsFromGrokInput(typed) + if childChanged { + node[key] = next + changed = true } - changed = true } } - _ = changed - return body, nil + return node, changed } var grokResponsesSupportedToolTypes = map[string]struct{}{ diff --git a/backend/internal/service/openai_gateway_grok_test.go b/backend/internal/service/openai_gateway_grok_test.go index 92f9901142..146b310176 100644 --- a/backend/internal/service/openai_gateway_grok_test.go +++ b/backend/internal/service/openai_gateway_grok_test.go @@ -1177,7 +1177,7 @@ func TestForwardGrokMediaAppliesAccountModelMappingAfterEndpointNormalization(t modelMapping: map[string]any{"grok-imagine-image-quality": "vendor-image-model"}, wantRequestModel: "grok-imagine-image-quality", wantUpstream: "vendor-image-model", - wantBody: `{"model":"vendor-image-model","prompt":"draw"}`, + wantBody: `{"model":"vendor-image-model","prompt":"draw","resolution":"1k","aspect_ratio":"1:1"}`, responseBody: `{"data":[{"url":"https://images.test/mapped.png"}]}`, }, { @@ -1300,7 +1300,8 @@ func TestForwardGrokMediaImagesGenerationStripsUnsupportedSize(t *testing.T) { result, err := svc.ForwardGrokMedia(context.Background(), c, account, GrokMediaEndpointImagesGenerations, "", body, "application/json") require.NoError(t, err) - require.JSONEq(t, `{"model":"grok-imagine-image","prompt":"draw a cat"}`, string(upstream.lastBody)) + require.JSONEq(t, `{"model":"grok-imagine-image","prompt":"draw a cat","resolution":"1k","aspect_ratio":"1:1"}`, string(upstream.lastBody)) + require.False(t, gjson.GetBytes(upstream.lastBody, "size").Exists()) require.Equal(t, ImageBillingSize1K, result.ImageSize) require.Equal(t, "1024x1024", result.ImageInputSize) } @@ -1361,6 +1362,58 @@ func TestForwardGrokMediaImagesEditMultipartConvertsToJSON(t *testing.T) { require.Equal(t, "vendor-image-edit", result.UpstreamModel) } +func TestForwardGrokMediaImagesEditMultipartPreservesExplicitGeometry(t *testing.T) { + t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true") + gin.SetMode(gin.TestMode) + + var buf bytes.Buffer + writer := multipart.NewWriter(&buf) + require.NoError(t, writer.WriteField("model", "grok-imagine-edit")) + require.NoError(t, writer.WriteField("prompt", "edit this private image")) + require.NoError(t, writer.WriteField("size", "1024x1024")) + require.NoError(t, writer.WriteField("resolution", "2k")) + require.NoError(t, writer.WriteField("aspect_ratio", "16:9")) + partHeader := textproto.MIMEHeader{} + partHeader.Set("Content-Disposition", `form-data; name="image"; filename="input.png"`) + partHeader.Set("Content-Type", "image/png") + part, err := writer.CreatePart(partHeader) + require.NoError(t, err) + _, err = part.Write([]byte{0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a}) + require.NoError(t, err) + require.NoError(t, writer.Close()) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/edits", bytes.NewReader(buf.Bytes())) + c.Request.Header.Set("Content-Type", writer.FormDataContentType()) + + account := &Account{ + ID: 67, + Name: "grok", + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "api-key", + "base_url": "https://xai.test/v1", + }, + } + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"data":[{"url":"https://images.test/edited.png"}]}`)), + }} + svc := &OpenAIGatewayService{httpUpstream: upstream} + + result, err := svc.ForwardGrokMedia(context.Background(), c, account, GrokMediaEndpointImagesEdits, "", buf.Bytes(), writer.FormDataContentType()) + require.NoError(t, err) + require.False(t, gjson.GetBytes(upstream.lastBody, "size").Exists()) + require.Equal(t, "2k", gjson.GetBytes(upstream.lastBody, "resolution").String()) + require.Equal(t, "16:9", gjson.GetBytes(upstream.lastBody, "aspect_ratio").String()) + require.Equal(t, ImageBillingSize1K, result.ImageSize) + require.Equal(t, "1024x1024", result.ImageInputSize) +} + func TestForwardGrokMediaVideoGenerationReturnsUsageAndResponseID(t *testing.T) { t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true") gin.SetMode(gin.TestMode) @@ -3446,6 +3499,43 @@ func TestPatchGrokResponsesBody_StripsReasoningContentNull(t *testing.T) { require.False(t, reasoning.Get("content").Exists(), "content: null should be stripped") } +func TestPatchGrokResponsesBody_StripsNestedInputNulls(t *testing.T) { + t.Parallel() + + body := []byte(`{ + "model": "grok-4.6", + "input": [ + {"type":"message","role":"user","content":[{"type":"input_text","text":"hi","logprobs":null}]} + ] + }`) + patched, err := patchGrokResponsesBody(body, "grok-4.6") + require.NoError(t, err) + require.False(t, gjson.GetBytes(patched, "input.0.content.0.logprobs").Exists()) + require.Equal(t, "hi", gjson.GetBytes(patched, "input.0.content.0.text").String()) +} + +func TestStripExplicitNullsFromGrokInputSkipsCompaction(t *testing.T) { + t.Parallel() + + input := []any{ + map[string]any{"type": "message", "role": "user", "content": []any{map[string]any{"type": "input_text", "text": "hi", "logprobs": nil}}}, + map[string]any{"type": "compaction", "id": "cmp_1", "encrypted_content": "blob", "status": nil}, + } + out, changed := stripExplicitNullsFromGrokInput(input) + require.True(t, changed) + items, ok := out.([]any) + require.True(t, ok) + require.Len(t, items, 2) + msg := items[0].(map[string]any) + content := msg["content"].([]any)[0].(map[string]any) + _, hasLogprobs := content["logprobs"] + require.False(t, hasLogprobs) + compaction := items[1].(map[string]any) + require.Equal(t, "blob", compaction["encrypted_content"]) + _, hasStatus := compaction["status"] + require.True(t, hasStatus, "compaction items must stay unmodified") +} + func TestPatchGrokResponsesBody_KeepsReasoningContentNonNull(t *testing.T) { t.Parallel() diff --git a/backend/internal/service/openai_gateway_response_handling.go b/backend/internal/service/openai_gateway_response_handling.go index 45a8e06de6..3c5a64e5fe 100644 --- a/backend/internal/service/openai_gateway_response_handling.go +++ b/backend/internal/service/openai_gateway_response_handling.go @@ -15,6 +15,7 @@ import ( "github.com/Wei-Shaw/sub2api/internal/pkg/apicompat" "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/gin-gonic/gin" "github.com/tidwall/gjson" @@ -1185,18 +1186,16 @@ func openAIUsageFromGJSON(value gjson.Result) (OpenAIUsage, bool) { if outputTokens == 0 { outputTokens = value.Get("completion_tokens").Int() } - // xAI may report reasoning_tokens separately from visible output tokens. - // Only add it when total_tokens proves it is an independent component; - // OpenAI's canonical completion_tokens already includes reasoning details. + // xAI reports visible output separately from reasoning_tokens; OpenAI + // folds reasoning into completion/output. Use total_tokens to tell them apart. reasoningTokens := max(int(firstPositiveGJSONInt( value.Get("completion_tokens_details.reasoning_tokens"), value.Get("output_tokens_details.reasoning_tokens"), )), 0) if reasoningTokens > 0 { - totalTokens := value.Get("total_tokens").Int() - if totalTokens > 0 && totalTokens == inputTokens+outputTokens+int64(reasoningTokens) { - outputTokens += int64(reasoningTokens) - } + outputTokens = xai.IncludeIndependentReasoningTokens( + inputTokens, outputTokens, value.Get("total_tokens").Int(), int64(reasoningTokens), + ) } cacheReadTokens := openAICacheReadTokensFromUsage(value) cacheCreationTokens := openAICacheCreationTokensFromUsage(value) diff --git a/backend/internal/service/openai_ws_v2/passthrough_relay.go b/backend/internal/service/openai_ws_v2/passthrough_relay.go index abc94e8f1c..2cad27d88f 100644 --- a/backend/internal/service/openai_ws_v2/passthrough_relay.go +++ b/backend/internal/service/openai_ws_v2/passthrough_relay.go @@ -13,6 +13,8 @@ import ( coderws "github.com/coder/websocket" "github.com/tidwall/gjson" + + "github.com/Wei-Shaw/sub2api/internal/pkg/xai" ) type FrameConn interface { @@ -948,6 +950,15 @@ func parseUsageAndAccumulate( // 解析失败时不做部分字段累加，避免计费 usage 出现“半有效”状态。 return Usage{} } + reasoningTokens := usageResult.Get("output_tokens_details.reasoning_tokens").Int() + if reasoningTokens == 0 { + reasoningTokens = usageResult.Get("completion_tokens_details.reasoning_tokens").Int() + } + if reasoningTokens > 0 { + outputTokens = int(xai.IncludeIndependentReasoningTokens( + int64(inputTokens), int64(outputTokens), usageResult.Get("total_tokens").Int(), reasoningTokens, + )) + } parsedUsage := Usage{ InputTokens: inputTokens, OutputTokens: outputTokens, diff --git a/backend/internal/service/openai_ws_v2/passthrough_relay_internal_test.go b/backend/internal/service/openai_ws_v2/passthrough_relay_internal_test.go index ebb2394a12..51f1afb6ee 100644 --- a/backend/internal/service/openai_ws_v2/passthrough_relay_internal_test.go +++ b/backend/internal/service/openai_ws_v2/passthrough_relay_internal_test.go @@ -324,6 +324,29 @@ func TestParseUsageAndEnrichCoverage(t *testing.T) { enrichResult(nil, state, 0) } +func TestParseUsageAndAccumulateIncludesIndependentReasoningTokens(t *testing.T) { + t.Parallel() + + state := &relayState{} + got := parseUsageAndAccumulate( + state, + []byte(`{"type":"response.completed","response":{"usage":{"input_tokens":32,"output_tokens":9,"total_tokens":151,"output_tokens_details":{"reasoning_tokens":110}}}}`), + "response.completed", + nil, + ) + require.Equal(t, 32, got.InputTokens) + require.Equal(t, 119, got.OutputTokens) + + state = &relayState{} + got = parseUsageAndAccumulate( + state, + []byte(`{"type":"response.completed","response":{"usage":{"input_tokens":32,"output_tokens":119,"total_tokens":151,"output_tokens_details":{"reasoning_tokens":110}}}}`), + "response.completed", + nil, + ) + require.Equal(t, 119, got.OutputTokens, "inclusive Responses output must not double-count reasoning") +} + func TestParseUsageAndAccumulateAcceptsChatUsageAliases(t *testing.T) { t.Parallel()