//go:build unit
package service
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"mime/multipart"
"net/http"
"net/http/httptest"
"net/textproto"
"strings"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
"github.com/tidwall/gjson"
)
func TestPatchGrokResponsesBodySetsMappedModelAndDropsUnsupportedFields(t *testing.T) {
t.Parallel()
body := []byte(`{
"model": "grok",
"input": "hello",
"prompt_cache_retention": "24h",
"safety_identifier": "user-1",
"reasoning": {"effort": "high"}
}`)
patched, err := patchGrokResponsesBody(body, "grok-4.3")
require.NoError(t, err)
require.True(t, json.Valid(patched))
require.Equal(t, "grok-4.3", gjson.GetBytes(patched, "model").String())
require.False(t, gjson.GetBytes(patched, "prompt_cache_retention").Exists())
require.False(t, gjson.GetBytes(patched, "safety_identifier").Exists())
require.Equal(t, "high", gjson.GetBytes(patched, "reasoning.effort").String())
}
func TestPatchGrokResponsesBodyDropsRedundantViewImageForCurrentInlineImage(t *testing.T) {
t.Parallel()
tests := []struct {
name string
body string
}{
{
name: "top-level tools",
body: `{
"model":"grok-4.6",
"input":[{"type":"message","role":"user","content":[
{"type":"input_text","text":"What text is in this image?"},
{"type":"input_image","image_url":"data:image/png;base64,AA=="}
]}],
"tools":[
{"type":"function","name":"view_image","parameters":{"type":"object"}},
{"type":"function","name":"shell_command","parameters":{"type":"object"}}
]
}`,
},
{
name: "Responses Lite additional tools",
body: `{
"model":"grok-4.6",
"input":[
{"type":"additional_tools","role":"developer","tools":[
{"type":"function","name":"view_image","parameters":{"type":"object"}},
{"type":"function","name":"shell_command","parameters":{"type":"object"}}
]},
{"type":"message","role":"user","content":[
{"type":"input_text","text":"What text is in this image?"},
{"type":"input_image","image_url":"data:image/png;base64,AA=="}
]}
]
}`,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
patched, err := patchGrokResponsesBody([]byte(tt.body), "grok-4.6")
require.NoError(t, err)
require.False(t, gjson.GetBytes(patched, `tools.#(name=="view_image")`).Exists())
require.Equal(t, "shell_command", gjson.GetBytes(patched, "tools.0.name").String())
})
}
}
func TestPatchGrokResponsesBodyKeepsNonRedundantViewImage(t *testing.T) {
t.Parallel()
tests := []struct {
name string
body string
}{
{
name: "current turn has no inline image",
body: `{"input":[{"role":"user","content":[{"type":"input_text","text":"Inspect a local image"}]}],"tools":[{"type":"function","name":"view_image"}]}`,
},
{
name: "inline image is only historical",
body: `{"input":[{"role":"user","content":[{"type":"input_image","image_url":"data:image/png;base64,AA=="}]},{"role":"assistant","content":[{"type":"output_text","text":"Done"}]},{"role":"user","content":[{"type":"input_text","text":"Inspect another local image"}]}],"tools":[{"type":"function","name":"view_image"}]}`,
},
{
name: "view image is explicitly selected",
body: `{"input":[{"role":"user","content":[{"type":"input_image","image_url":"data:image/png;base64,AA=="}]}],"tools":[{"type":"function","name":"view_image"}],"tool_choice":{"type":"function","name":"view_image"}}`,
},
{
name: "required with view image as the only tool",
body: `{"input":[{"role":"user","content":[{"type":"input_image","image_url":"data:image/png;base64,AA=="}]}],"tools":[{"type":"function","name":"view_image"}],"tool_choice":"required"}`,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
patched, err := patchGrokResponsesBody([]byte(tt.body), "grok-4.6")
require.NoError(t, err)
require.Equal(t, "view_image", gjson.GetBytes(patched, "tools.0.name").String())
})
}
}
func TestPatchGrokResponsesBodyDropsViewImageOnlyToolMetadata(t *testing.T) {
t.Parallel()
body := []byte(`{
"input":[{"role":"user","content":[{"type":"input_image","image_url":"data:image/png;base64,AA=="}]}],
"tools":[{"type":"function","name":"view_image"}],
"tool_choice":"auto",
"parallel_tool_calls":true
}`)
patched, err := patchGrokResponsesBody(body, "grok-4.6")
require.NoError(t, err)
require.False(t, gjson.GetBytes(patched, "tools").Exists())
require.False(t, gjson.GetBytes(patched, "tool_choice").Exists())
require.False(t, gjson.GetBytes(patched, "parallel_tool_calls").Exists())
}
func TestPatchGrokResponsesBodySanitizesComposerReasoningParameters(t *testing.T) {
t.Parallel()
tests := []struct {
name string
upstreamModel string
wantReasoning bool
}{
{name: "composer fast", upstreamModel: "grok-composer-2.5-fast"},
{name: "composer shorthand", upstreamModel: "grok-composer"},
{name: "composer legacy alias", upstreamModel: "composer-2.5"},
{name: "provider-prefixed composer", upstreamModel: "xai/grok-composer-2.5-fast"},
{name: "grok 4.5", upstreamModel: "grok-4.5", wantReasoning: true},
{name: "grok 4.6", upstreamModel: "grok-4.6", wantReasoning: true},
{name: "grok 4.6 latest", upstreamModel: "grok-4.6-latest", wantReasoning: true},
}
bodyTemplate := []byte(`{
"model": "grok",
"input": "hello",
"reasoning": {"effort": "medium", "summary": "auto"},
"reasoning_effort": "medium",
"reasoningEffort": "medium"
}`)
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
patched, err := patchGrokResponsesBody(append([]byte(nil), bodyTemplate...), tt.upstreamModel)
require.NoError(t, err)
require.True(t, json.Valid(patched))
require.Equal(t, tt.upstreamModel, gjson.GetBytes(patched, "model").String())
if tt.wantReasoning {
require.Equal(t, "medium", gjson.GetBytes(patched, "reasoning.effort").String())
require.Equal(t, "medium", gjson.GetBytes(patched, "reasoning_effort").String())
require.False(t, gjson.GetBytes(patched, "reasoningEffort").Exists())
return
}
require.False(t, gjson.GetBytes(patched, "reasoning").Exists())
require.False(t, gjson.GetBytes(patched, "reasoning_effort").Exists())
require.False(t, gjson.GetBytes(patched, "reasoningEffort").Exists())
})
}
}
func TestExtractGrokResponsesReasoningEffortSupportsOpenAICompatibleField(t *testing.T) {
t.Parallel()
effort := extractOpenAIReasoningEffortFromBody(
[]byte(`{"model":"grok-4.3","reasoning_effort":"high"}`),
"grok-4.3",
)
require.NotNil(t, effort)
require.Equal(t, "high", *effort)
}
func TestPatchGrokResponsesBodyDropsGrok45ReasoningUnsupportedFields(t *testing.T) {
t.Parallel()
body := []byte(`{
"model": "grok-latest",
"input": "hello",
"presence_penalty": 0.1,
"presencePenalty": 0.2,
"frequency_penalty": 0.3,
"frequencyPenalty": 0.4,
"stop": ["done"]
}`)
patched, err := patchGrokResponsesBody(body, "grok-4.5")
require.NoError(t, err)
require.True(t, json.Valid(patched))
require.Equal(t, "grok-4.5", gjson.GetBytes(patched, "model").String())
require.False(t, gjson.GetBytes(patched, "presence_penalty").Exists())
require.False(t, gjson.GetBytes(patched, "presencePenalty").Exists())
require.False(t, gjson.GetBytes(patched, "frequency_penalty").Exists())
require.False(t, gjson.GetBytes(patched, "frequencyPenalty").Exists())
require.False(t, gjson.GetBytes(patched, "stop").Exists())
}
func TestPatchGrokResponsesBodyKeepsPenaltyAndStopFieldsForNon45Models(t *testing.T) {
t.Parallel()
body := []byte(`{
"model": "grok-4.3",
"input": "hello",
"presence_penalty": 0.1,
"frequency_penalty": 0.2,
"stop": ["done"]
}`)
patched, err := patchGrokResponsesBody(body, "grok-4.3")
require.NoError(t, err)
require.True(t, json.Valid(patched))
require.Equal(t, "grok-4.3", gjson.GetBytes(patched, "model").String())
require.Equal(t, 0.1, gjson.GetBytes(patched, "presence_penalty").Float())
require.Equal(t, 0.2, gjson.GetBytes(patched, "frequency_penalty").Float())
require.Len(t, gjson.GetBytes(patched, "stop").Array(), 1)
}
func TestPatchGrokResponsesBodyDropsLogprobsForGrok420Family(t *testing.T) {
t.Parallel()
body := []byte(`{"model":"grok-4.20-0309-reasoning","input":"hello","logprobs":true,"top_logprobs":5}`)
patched, err := patchGrokResponsesBody(body, "grok-4.20-0309-reasoning")
require.NoError(t, err)
require.False(t, gjson.GetBytes(patched, "logprobs").Exists())
require.False(t, gjson.GetBytes(patched, "top_logprobs").Exists())
}
func TestPatchGrokResponsesBodyNormalizesReasoningEffortAliases(t *testing.T) {
t.Parallel()
tests := []struct {
name string
body string
upstreamModel string
path string
want string
}{
{name: "minimal nested", body: `{"input":"hi","reasoning":{"effort":"minimal"}}`, upstreamModel: "grok-4.5", path: "reasoning.effort", want: "low"},
{name: "xhigh stays high for 4.5", body: `{"input":"hi","reasoning_effort":"xhigh"}`, upstreamModel: "grok-4.5", path: "reasoning_effort", want: "high"},
{name: "xhigh nested for 4.6", body: `{"input":"hi","reasoning":{"effort":"xhigh"}}`, upstreamModel: "grok-4.6", path: "reasoning.effort", want: "xhigh"},
{name: "xhigh snake for 4.6 latest", body: `{"input":"hi","reasoning_effort":"xhigh"}`, upstreamModel: "grok-4.6-latest", path: "reasoning_effort", want: "xhigh"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
patched, err := patchGrokResponsesBody([]byte(tt.body), tt.upstreamModel)
require.NoError(t, err)
require.Equal(t, tt.want, gjson.GetBytes(patched, tt.path).String(), string(patched))
require.False(t, gjson.GetBytes(patched, "reasoningEffort").Exists())
})
}
}
func TestPatchGrokResponsesBodyAddsDefaultFunctionParameters(t *testing.T) {
patched, err := patchGrokResponsesBody(
[]byte(`{"input":"hi","tools":[{"type":"function","name":"lookup"},{"type":"function","name":"wait","parameters":null}]}`),
"grok-4.5",
)
require.NoError(t, err)
for _, tool := range gjson.GetBytes(patched, "tools").Array() {
require.Equal(t, "object", tool.Get("parameters.type").String(), string(patched))
require.True(t, tool.Get("parameters.properties").IsObject(), string(patched))
}
}
func TestNormalizeGrokChatReasoningEffort(t *testing.T) {
patched, err := normalizeGrokChatReasoningEffort([]byte(`{"reasoningEffort":"ultra"}`), "grok-4.3")
require.NoError(t, err)
require.Equal(t, "high", gjson.GetBytes(patched, "reasoning_effort").String())
require.False(t, gjson.GetBytes(patched, "reasoningEffort").Exists())
patched, err = normalizeGrokChatReasoningEffort([]byte(`{"reasoning_effort":"xhigh"}`), "grok-4.6")
require.NoError(t, err)
require.Equal(t, "xhigh", gjson.GetBytes(patched, "reasoning_effort").String())
patched, err = normalizeGrokChatReasoningEffort([]byte(`{"reasoning_effort":"high"}`), "grok-composer-2.5-fast")
require.NoError(t, err)
require.False(t, gjson.GetBytes(patched, "reasoning_effort").Exists())
}
func TestPatchGrokResponsesBodyDropsNestedUnsupportedFields(t *testing.T) {
t.Parallel()
body := []byte(`{
"model": "grok",
"input": "hello",
"external_web_access": true,
"tools": [
{"type": "function", "name": "kept_fn", "external_web_access": true, "parameters": {"type": "object", "properties": {"q": {"type": "string", "external_web_access": true}}}}
],
"metadata": {"external_web_access": false}
}`)
patched, err := patchGrokResponsesBody(body, "grok-4.3")
require.NoError(t, err)
require.True(t, json.Valid(patched))
require.False(t, strings.Contains(string(patched), "external_web_access"))
require.Equal(t, "kept_fn", gjson.GetBytes(patched, "tools.0.name").String())
}
func TestPatchGrokResponsesBodyFlattensNamespaceTools(t *testing.T) {
t.Parallel()
body := []byte(`{
"model": "grok",
"input": "hello",
"tools": [
{"type": "namespace", "name": "functions", "tools": [{"type": "function", "name": "inner"}]},
{"type": "function", "name": "kept_fn", "parameters": {"type": "object"}},
{"type": "shell", "name": "kept_shell"}
],
"tool_choice": {"type": "function", "namespace": "functions", "name": "inner"}
}`)
patched, _, err := patchGrokResponsesBodyWithClientTools(body, "grok-4.3")
require.NoError(t, err)
require.True(t, json.Valid(patched))
require.Equal(t, "grok-4.3", gjson.GetBytes(patched, "model").String())
require.Len(t, gjson.GetBytes(patched, "tools").Array(), 3)
require.False(t, gjson.GetBytes(patched, `tools.#(type=="namespace")`).Exists())
require.True(t, gjson.GetBytes(patched, `tools.#(type=="function")`).Exists())
require.True(t, gjson.GetBytes(patched, `tools.#(type=="shell")`).Exists())
require.Equal(t, "functions__inner", gjson.GetBytes(patched, "tools.0.name").String())
require.Equal(t, "functions__inner", gjson.GetBytes(patched, "tool_choice.name").String())
require.False(t, gjson.GetBytes(patched, "tool_choice.namespace").Exists())
}
func TestPatchGrokResponsesBodyDropsToolChoiceWhenNoSupportedToolsRemain(t *testing.T) {
t.Parallel()
body := []byte(`{
"model": "grok",
"input": "hello",
"tools": [
{"type": "namespace", "namespace": "functions"},
{"type": "image_generation", "model": "gpt-image-2"}
],
"tool_choice": {"type": "namespace", "namespace": "functions"}
}`)
patched, err := patchGrokResponsesBody(body, "grok-4.3")
require.NoError(t, err)
require.True(t, json.Valid(patched))
require.False(t, gjson.GetBytes(patched, "tools").Exists())
require.False(t, gjson.GetBytes(patched, "tool_choice").Exists())
}
func TestSanitizeGrokResponsesToolsKeepsToolChoiceOnlyWithSupportedTools(t *testing.T) {
t.Parallel()
tests := []struct {
name string
body string
wantTools bool
wantToolChoice bool
}{
{
name: "missing tools with string tool choice",
body: `{"input":"hello","tool_choice":"auto"}`,
},
{
name: "missing tools with object tool choice",
body: `{"input":"hello","tool_choice":{"type":"function","name":"lookup"}}`,
},
{
name: "empty tools",
body: `{"input":"hello","tools":[],"tool_choice":"auto"}`,
wantTools: true,
},
{
name: "all tools unsupported",
body: `{"input":"hello","tools":[{"type":"namespace","name":"client_tools"}],"tool_choice":"auto"}`,
},
{
name: "supported tool",
body: `{"input":"hello","tools":[{"type":"function","name":"lookup"}],"tool_choice":"auto"}`,
wantTools: true,
wantToolChoice: true,
},
{
name: "malformed non-array tools remain untouched",
body: `{"input":"hello","tools":{"type":"function","name":"lookup"},"tool_choice":"auto"}`,
wantTools: true,
wantToolChoice: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
patched, err := sanitizeGrokResponsesTools([]byte(tt.body))
require.NoError(t, err)
require.True(t, json.Valid(patched))
require.Equal(t, tt.wantTools, gjson.GetBytes(patched, "tools").Exists())
require.Equal(t, tt.wantToolChoice, gjson.GetBytes(patched, "tool_choice").Exists())
if tt.wantToolChoice {
require.Equal(t, "auto", gjson.GetBytes(patched, "tool_choice").String())
}
})
}
}
func TestPatchGrokResponsesBodyPromotesCodexAdditionalTools(t *testing.T) {
t.Parallel()
body := []byte(`{
"model": "grok",
"tools": [
{"type": "function", "name": "existing", "description": "top-level wins"},
{"type": "web_search"}
],
"tool_choice": "auto",
"input": [
{
"type": "additional_tools",
"role": "developer",
"tools": [
{"type": "function", "name": "existing", "description": "duplicate carrier definition"},
{"type": "function", "name": "wait"},
{"type": "web_search"},
{"type": "shell"},
{"type": "custom", "name": "apply_patch"},
{"type": "namespace", "name": "collaboration"}
]
},
{
"type": "message",
"role": "developer",
"content": [{"type": "input_text", "text": "system prompt"}]
},
{
"type": "message",
"role": "user",
"content": [{"type": "input_text", "text": "hello"}]
}
]
}`)
patched, _, err := patchGrokResponsesBodyWithClientTools(body, "grok-4.5")
require.NoError(t, err)
require.True(t, json.Valid(patched))
require.Equal(t, "grok-4.5", gjson.GetBytes(patched, "model").String())
require.Equal(t, 2, len(gjson.GetBytes(patched, "input").Array()))
require.False(t, gjson.GetBytes(patched, `input.#(type=="additional_tools")`).Exists())
tools := gjson.GetBytes(patched, "tools").Array()
require.Len(t, tools, 5)
require.Equal(t, "existing", tools[0].Get("name").String())
require.Equal(t, "top-level wins", tools[0].Get("description").String())
require.Equal(t, "web_search", tools[1].Get("type").String())
require.Equal(t, "wait", tools[2].Get("name").String())
require.Equal(t, "shell", tools[3].Get("type").String())
require.Equal(t, "function", tools[4].Get("type").String())
require.Equal(t, "apply_patch", tools[4].Get("name").String())
require.Equal(t, "string", tools[4].Get("parameters.properties.input.type").String())
require.False(t, gjson.GetBytes(patched, `tools.#(type=="custom")`).Exists())
require.False(t, gjson.GetBytes(patched, `tools.#(type=="namespace")`).Exists())
require.Equal(t, "auto", gjson.GetBytes(patched, "tool_choice").String())
require.Equal(t, "developer", gjson.GetBytes(patched, "input.0.role").String())
require.Equal(t, "system prompt", gjson.GetBytes(patched, "input.0.content.0.text").String())
require.Equal(t, "user", gjson.GetBytes(patched, "input.1.role").String())
require.Equal(t, "hello", gjson.GetBytes(patched, "input.1.content.0.text").String())
}
func TestForwardGrokResponsesCodexAdditionalToolsUsesMixedCacheIntent(t *testing.T) {
gin.SetMode(gin.TestMode)
body := []byte(`{
"model":"grok",
"stream":false,
"prompt_cache_key":"codex-session",
"input":[
{"type":"additional_tools","role":"developer","tools":[
{"type":"function","name":"lookup","description":"look up a key","parameters":{"type":"object"}},
{"type":"function","name":"web_search","description":"search","parameters":{"type":"object"}},
{"type":"custom","name":"apply_patch"},
{"type":"namespace","name":"collaboration"}
]},
{"type":"message","role":"user","content":[{"type":"input_text","text":"hello"}]}
]
}`)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
c.Request.Header.Set(grokClientToolCacheOptInHeader, "prefer-cache")
c.Set("api_key", &APIKey{ID: 4501})
account := healthyGrokOAuthGatewayTestAccount(4501, "access-token")
account.Credentials["subscription_tier"] = "free"
repo := &grokQuotaAccountRepo{
mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
accountsByID: map[int64]*Account{account.ID: account},
},
}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`{
"id":"resp_codex_lite","object":"response","model":"grok-4.5","status":"completed",
"output":[],"usage":{"input_tokens":10,"output_tokens":1}
}`)),
}}
svc := &OpenAIGatewayService{
httpUpstream: upstream,
grokTokenProvider: NewGrokTokenProvider(repo, nil),
accountRepo: repo,
}
result, err := svc.forwardGrokResponses(context.Background(), c, account, body, "grok", false, time.Now())
require.NoError(t, err)
require.NotNil(t, result)
require.Equal(t, "resp_codex_lite", result.ResponseID)
require.False(t, gjson.GetBytes(upstream.lastBody, `input.#(type=="additional_tools")`).Exists())
tools := gjson.GetBytes(upstream.lastBody, "tools").Array()
require.Len(t, tools, 4)
require.Equal(t, "function", tools[0].Get("type").String())
require.Equal(t, "lookup", tools[0].Get("name").String())
require.Equal(t, "web_search", tools[1].Get("type").String())
require.Equal(t, "function", tools[2].Get("type").String())
require.Equal(t, "apply_patch", tools[2].Get("name").String())
require.Equal(t, "x_search", tools[3].Get("type").String())
require.False(t, gjson.GetBytes(upstream.lastBody, "tool_choice").Exists())
require.False(t, gjson.GetBytes(upstream.lastBody, `tools.#(type=="custom")`).Exists())
require.False(t, gjson.GetBytes(upstream.lastBody, `tools.#(type=="namespace")`).Exists())
identity := gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String()
require.NotEmpty(t, identity)
require.Equal(t, identity, upstream.lastReq.Header.Get(grokConversationIDHeader))
require.Empty(t, upstream.lastReq.Header.Get(grokClientToolCacheOptInHeader))
}
func TestForwardGrokResponsesClaudeDesktopClientToolsUseCacheRoute(t *testing.T) {
gin.SetMode(gin.TestMode)
firstBody := []byte(`{
"model":"grok","stream":false,"instructions":"You are Claude Desktop.",
"tools":[
{"type":"function","name":"Read","parameters":{"type":"object"}},
{"type":"function","name":"Edit","parameters":{"type":"object"}},
{"type":"function","name":"WebSearch","parameters":{"type":"object"}},
{"type":"function","name":"mcp__workspace__bash","parameters":{"type":"object"}}
],
"input":[{"role":"user","content":[{"type":"input_text","text":"first turn"}]}]
}`)
secondBody := []byte(`{
"model":"grok","stream":false,"instructions":"You are Claude Desktop.",
"tools":[
{"type":"function","name":"Read","parameters":{"type":"object"}},
{"type":"function","name":"Edit","parameters":{"type":"object"}},
{"type":"function","name":"WebSearch","parameters":{"type":"object"}},
{"type":"function","name":"mcp__workspace__bash","parameters":{"type":"object"}}
],
"input":[
{"role":"user","content":[{"type":"input_text","text":"first turn"}]},
{"role":"assistant","content":[{"type":"output_text","text":"first answer"}]},
{"role":"user","content":[{"type":"input_text","text":"second turn"}]}
]
}`)
account := healthyGrokOAuthGatewayTestAccount(4504, "access-token")
account.Credentials["subscription_tier"] = "free"
repo := &grokQuotaAccountRepo{
mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
accountsByID: map[int64]*Account{account.ID: account},
},
}
upstream := &httpUpstreamRecorder{responses: []*http.Response{
{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`{
"id":"resp_claude_desktop_1","object":"response","model":"grok-4.5","status":"completed",
"output":[],"usage":{"input_tokens":30000,"output_tokens":10,"input_tokens_details":{"cached_tokens":0}}
}`)),
},
{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`{
"id":"resp_claude_desktop_2","object":"response","model":"grok-4.5","status":"completed",
"output":[],"usage":{"input_tokens":30100,"output_tokens":12,"input_tokens_details":{"cached_tokens":28672}}
}`)),
},
}}
svc := &OpenAIGatewayService{
httpUpstream: upstream,
grokTokenProvider: NewGrokTokenProvider(repo, nil),
accountRepo: repo,
}
newContext := func(body []byte) *gin.Context {
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
c.Request.Header.Set("User-Agent", "claude-cli/2.1.215 (external, claude-desktop-3p, agent-sdk/0.3.215)")
c.Request.Header.Set("X-App", "cli")
c.Request.Header.Set("anthropic-client-platform", "desktop_app")
c.Request.Header.Set("X-Claude-Code-Session-Id", "claude-desktop-session")
c.Set("api_key", &APIKey{ID: 4504})
return c
}
first, err := svc.forwardGrokResponses(context.Background(), newContext(firstBody), account, firstBody, "grok", false, time.Now())
require.NoError(t, err)
second, err := svc.forwardGrokResponses(context.Background(), newContext(secondBody), account, secondBody, "grok", false, time.Now())
require.NoError(t, err)
require.Equal(t, 0, first.Usage.CacheReadInputTokens)
require.Equal(t, 28672, second.Usage.CacheReadInputTokens)
require.Len(t, upstream.bodies, 2)
require.Len(t, upstream.requests, 2)
for i := range upstream.bodies {
tools := gjson.GetBytes(upstream.bodies[i], "tools").Array()
require.Len(t, tools, 6)
require.Equal(t, "Read", tools[0].Get("name").String())
require.Equal(t, "Edit", tools[1].Get("name").String())
require.Equal(t, "WebSearch", tools[2].Get("name").String())
require.Equal(t, "mcp__workspace__bash", tools[3].Get("name").String())
require.Equal(t, "web_search", tools[4].Get("type").String())
require.Equal(t, "x_search", tools[5].Get("type").String())
require.False(t, gjson.GetBytes(upstream.bodies[i], "tool_choice").Exists())
require.Empty(t, upstream.requests[i].Header.Get("X-App"))
require.Empty(t, upstream.requests[i].Header.Get("anthropic-client-platform"))
require.Empty(t, upstream.requests[i].Header.Get("X-Claude-Code-Session-Id"))
}
firstIdentity := gjson.GetBytes(upstream.bodies[0], "prompt_cache_key").String()
secondIdentity := gjson.GetBytes(upstream.bodies[1], "prompt_cache_key").String()
require.NotEmpty(t, firstIdentity)
require.Equal(t, firstIdentity, secondIdentity)
require.Equal(t, firstIdentity, upstream.requests[0].Header.Get(grokConversationIDHeader))
require.Equal(t, secondIdentity, upstream.requests[1].Header.Get(grokConversationIDHeader))
}
func TestGrokResponsesCacheIdentityIncludesPromotedCodexTools(t *testing.T) {
c := newGrokCacheTestContext(4503)
lookupBody := []byte(`{"model":"grok","input":[{"type":"additional_tools","tools":[{"type":"function","name":"lookup","parameters":{"type":"object"}}]},{"type":"message","role":"user","content":"same prompt"}]}`)
readBody := []byte(`{"model":"grok","input":[{"type":"additional_tools","tools":[{"type":"function","name":"read_file","parameters":{"type":"object"}}]},{"type":"message","role":"user","content":"same prompt"}]}`)
patchedLookup, err := patchGrokResponsesBody(lookupBody, "grok-4.5")
require.NoError(t, err)
patchedRead, err := patchGrokResponsesBody(readBody, "grok-4.5")
require.NoError(t, err)
lookupIdentity := resolveGrokCacheIdentity(c, patchedLookup, "", "grok-4.5")
readIdentity := resolveGrokCacheIdentity(c, patchedRead, "", "grok-4.5")
require.NotEmpty(t, lookupIdentity)
require.NotEmpty(t, readIdentity)
require.NotEqual(t, lookupIdentity, readIdentity)
}
func TestCodexUnsupportedAdditionalToolsDoNotBecomeToolFreeCacheIntent(t *testing.T) {
body := []byte(`{
"model":"grok","tool_choice":"auto",
"input":[
{"type":"additional_tools","role":"developer","tools":[
{"type":"custom","name":"apply_patch"},
{"type":"namespace","name":"collaboration"}
]},
{"type":"message","role":"user","content":"hello"}
]
}`)
patched, err := patchGrokResponsesBody(body, "grok-4.5")
require.NoError(t, err)
require.False(t, gjson.GetBytes(patched, "tools").Exists())
require.False(t, gjson.GetBytes(patched, "tool_choice").Exists())
mixedCacheIntent := patched
patched, err = applyGrokResponsesCacheIdentity(patched, body, "isolated-id", true)
require.NoError(t, err)
account := healthyGrokOAuthGatewayTestAccount(4502, "access-token")
account.Credentials["subscription_tier"] = "free"
patched, err = applyGrokFreeRequestToolCacheRoute(nil, patched, mixedCacheIntent, account, "isolated-id")
require.NoError(t, err)
require.False(t, gjson.GetBytes(patched, "tools").Exists())
require.False(t, gjson.GetBytes(patched, "tool_choice").Exists())
require.Equal(t, "isolated-id", gjson.GetBytes(patched, "prompt_cache_key").String())
}
func TestBuildGrokResponsesRequestUsesAccountBaseURLAndBearerToken(t *testing.T) {
t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
account := &Account{
Platform: PlatformGrok,
Type: AccountTypeOAuth,
Credentials: map[string]any{
"base_url": "https://xai.test/v1/",
},
}
req, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "access-token", "isolated-cache-id", nil)
require.NoError(t, err)
require.Equal(t, http.MethodPost, req.Method)
require.Equal(t, "https://xai.test/v1/responses", req.URL.String())
require.Equal(t, "Bearer access-token", req.Header.Get("Authorization"))
require.Equal(t, "application/json", req.Header.Get("Content-Type"))
require.Contains(t, req.Header.Get("Accept"), "text/event-stream")
require.Equal(t, grokCLIVersion, req.Header.Get("X-Grok-Client-Version"))
require.Equal(t, "isolated-cache-id", req.Header.Get(grokConversationIDHeader))
data, err := io.ReadAll(req.Body)
require.NoError(t, err)
require.Equal(t, `{"model":"grok-4.3"}`, strings.TrimSpace(string(data)))
}
func TestBuildGrokCompactRequestBodyUsesResponsesCompactionTurn(t *testing.T) {
body := []byte(`{"model":"grok-4.5","input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"hello"}]}],"tools":[{"type":"function","name":"shell"}],"stream":true}`)
patched, err := buildGrokCompactRequestBody(body)
require.NoError(t, err)
require.False(t, gjson.GetBytes(patched, "stream").Bool())
require.False(t, gjson.GetBytes(patched, "store").Bool())
require.Equal(t, "none", gjson.GetBytes(patched, "tool_choice").String())
require.Equal(t, "reasoning.encrypted_content", gjson.GetBytes(patched, "include.0").String())
require.Equal(t, "hello", gjson.GetBytes(patched, "input.0.content.0.text").String())
prompt := gjson.GetBytes(patched, "input.1.content.0.text").String()
require.Contains(t, prompt, "1. Primary Request and Intent")
require.Contains(t, prompt, "9. Optional Next Step")
require.Contains(t, prompt, "Respond with ONLY the ... block")
require.NotContains(t, prompt, "")
}
func TestConvertGrokResponseToOpenAICompact(t *testing.T) {
body := []byte(`{
"id":"resp_grok_1",
"object":"response",
"status":"completed",
"model":"grok-4.5",
"output":[
{"id":"rs_1","type":"reasoning","summary":[],"encrypted_content":"grok-encrypted-state"},
{"id":"msg_1","type":"message","role":"assistant","content":[{"type":"output_text","text":"summary text"}]}
],
"usage":{"input_tokens":10,"output_tokens":4,"total_tokens":14}
}`)
converted, err := convertGrokResponseToOpenAICompact(body)
require.NoError(t, err)
require.Equal(t, "resp_grok_1", gjson.GetBytes(converted, "id").String())
require.Len(t, gjson.GetBytes(converted, "output").Array(), 1)
require.Equal(t, "compaction", gjson.GetBytes(converted, "output.0.type").String())
require.Equal(t, "grok-encrypted-state", gjson.GetBytes(converted, "output.0.encrypted_content").String())
require.Equal(t, "summary text", gjson.GetBytes(converted, "output.0.summary.0.text").String())
require.Equal(t, int64(14), gjson.GetBytes(converted, "usage.total_tokens").Int())
}
func TestPatchGrokResponsesBodyRestoresCompactInput(t *testing.T) {
body := []byte(`{
"model":"grok-4.5",
"input":[
{"id":"cmp_1","type":"compaction","status":"completed","encrypted_content":"grok-encrypted-state","summary":[{"type":"summary_text","text":"summary text"}]},
{"type":"message","role":"user","content":[{"type":"input_text","text":"continue"}]}
]
}`)
patched, err := patchGrokResponsesBody(body, "grok-4.5")
require.NoError(t, err)
require.Equal(t, "reasoning", gjson.GetBytes(patched, "input.0.type").String())
require.Equal(t, "grok-encrypted-state", gjson.GetBytes(patched, "input.0.encrypted_content").String())
require.Equal(t, "message", gjson.GetBytes(patched, "input.1.type").String())
require.Contains(t, gjson.GetBytes(patched, "input.1.content.0.text").String(), "summary text")
require.Equal(t, "continue", gjson.GetBytes(patched, "input.2.content.0.text").String())
}
func TestConvertGrokResponseToOpenAICompactRequiresEncryptedContent(t *testing.T) {
_, err := convertGrokResponseToOpenAICompact([]byte(`{"output":[{"type":"message","content":[{"type":"output_text","text":"summary"}]}]}`))
require.ErrorContains(t, err, "reasoning.encrypted_content")
}
func TestBuildGrokResponsesRequestAllowsPublicAPIKeyBaseURLByDefault(t *testing.T) {
account := &Account{
Platform: PlatformGrok,
Type: AccountTypeAPIKey,
Credentials: map[string]any{
"base_url": "https://grok.example.test/v1/",
},
}
req, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "api-key", "", nil)
require.NoError(t, err)
require.Equal(t, "https://grok.example.test/v1/responses", req.URL.String())
require.Equal(t, "Bearer api-key", req.Header.Get("Authorization"))
require.Empty(t, req.Header.Get("X-Grok-Client-Version"))
require.NotEqual(t, grokUpstreamUserAgent, req.Header.Get("User-Agent"))
}
func TestBuildGrokResponsesRequestHonorsOAuthOfficialEndpointSwitch(t *testing.T) {
t.Parallel()
account := &Account{
Platform: PlatformGrok,
Type: AccountTypeOAuth,
Credentials: map[string]any{
"base_url": xai.DefaultBaseURL,
},
}
req, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "access-token", "", nil)
require.NoError(t, err)
require.Equal(t, xai.DefaultBaseURL+"/responses", req.URL.String())
}
func TestBuildGrokResponsesRequestAppliesHeaderOverridesLast(t *testing.T) {
t.Parallel()
account := &Account{
Platform: PlatformGrok,
Type: AccountTypeOAuth,
Credentials: map[string]any{
"base_url": "https://relay.example.test/v1",
"header_override_enabled": true,
"header_overrides": map[string]any{
"User-Agent": "relay-client/2.0",
"X-Grok-Client-Version": "9.9.9",
"X-Relay-Token": "relay-secret",
},
},
}
req, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "access-token", "conv-1", nil)
require.NoError(t, err)
require.Equal(t, "https://relay.example.test/v1/responses", req.URL.String())
// 覆写值优先于内置 CLI 身份头。名字不在 wire casing 映射中的覆写头
// 以小写键直写（HTTP/2 线上语义），需按写入形态断言。
require.Equal(t, "relay-client/2.0", req.Header.Get("User-Agent"))
require.Equal(t, []string{"9.9.9"}, req.Header["x-grok-client-version"])
require.Empty(t, req.Header.Get("X-Grok-Client-Version"))
require.Equal(t, []string{"relay-secret"}, req.Header["x-relay-token"])
// 会话路由头与认证头不受覆写影响。
require.Equal(t, "conv-1", req.Header.Get(grokConversationIDHeader))
require.Equal(t, "Bearer access-token", req.Header.Get("Authorization"))
}
func TestBuildGrokResponsesRequestIgnoresBlockedHeaderOverrides(t *testing.T) {
t.Parallel()
account := &Account{
Platform: PlatformGrok,
Type: AccountTypeAPIKey,
Credentials: map[string]any{
"header_override_enabled": true,
"header_overrides": map[string]any{
"Authorization": "Bearer stolen",
"x-grok-conv-id": "pinned-conversation",
},
},
}
req, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "api-key", "conv-2", nil)
require.NoError(t, err)
require.Equal(t, "Bearer api-key", req.Header.Get("Authorization"))
require.Equal(t, "conv-2", req.Header.Get(grokConversationIDHeader))
}
func TestGrokMediaGenerationGateCoversImagesAndVideo(t *testing.T) {
tests := []struct {
name string
endpoint GrokMediaEndpoint
want bool
}{
{name: "image generation", endpoint: GrokMediaEndpointImagesGenerations, want: true},
{name: "image edit", endpoint: GrokMediaEndpointImagesEdits, want: true},
{name: "video generation", endpoint: GrokMediaEndpointVideosGenerations, want: true},
{name: "video status", endpoint: GrokMediaEndpointVideoStatus, want: false},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
require.Equal(t, tt.want, tt.endpoint.IsGenerationRequest())
})
}
}
func TestExtractGrokMediaModelSupportsJSONAndMultipart(t *testing.T) {
require.Equal(t, "grok-imagine", ExtractGrokMediaModel("application/json", []byte(`{"model":"grok-imagine"}`)))
var buf bytes.Buffer
writer := multipart.NewWriter(&buf)
require.NoError(t, writer.WriteField("prompt", "draw a cat"))
require.NoError(t, writer.WriteField("model", "grok-imagine-edit"))
require.NoError(t, writer.Close())
require.Equal(t, "grok-imagine-edit", ExtractGrokMediaModel(writer.FormDataContentType(), buf.Bytes()))
}
func TestParseGrokMediaRequestBuildsMultipartModerationBody(t *testing.T) {
var buf bytes.Buffer
writer := multipart.NewWriter(&buf)
require.NoError(t, writer.WriteField("prompt", "edit this private image"))
require.NoError(t, writer.WriteField("model", "grok-imagine-edit"))
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())
info := ParseGrokMediaRequest(writer.FormDataContentType(), buf.Bytes())
require.Equal(t, "grok-imagine-edit", info.Model)
require.Equal(t, "edit this private image", info.Prompt)
moderationBody := info.ModerationBody()
require.NotEmpty(t, moderationBody)
require.Equal(t, "edit this private image", gjson.GetBytes(moderationBody, "prompt").String())
require.True(t, strings.HasPrefix(gjson.GetBytes(moderationBody, "images.0.image_url").String(), "data:image/"))
}
func TestParseGrokMediaVideoRequestResolution(t *testing.T) {
info := ParseGrokMediaRequest("application/json", []byte(`{"model":"grok-imagine-video","prompt":"waves","resolution":"720p"}`))
require.Equal(t, "grok-imagine-video", info.Model)
require.Equal(t, "720p", info.Resolution)
}
func TestParseGrokMediaRequestAcceptsOfficialImageURLFields(t *testing.T) {
body := []byte(`{
"model":"grok-imagine-video-1.5",
"image":{"url":"https://example.com/source.png"},
"reference_images":[{"url":"https://example.com/reference.png"}]
}`)
info := ParseGrokMediaRequest("application/json", body)
require.Equal(t, []string{
"https://example.com/source.png",
"https://example.com/reference.png",
}, info.InputImageURLs)
require.True(t, info.HasInputImage())
}
func TestNormalizeGrokMediaForwardBodyCanonicalizesImageURLAlias(t *testing.T) {
body := []byte(`{
"model":"grok-imagine-video-1.5",
"prompt":"animate",
"image":{"image_url":"https://example.com/source.png"},
"duration":8
}`)
out, contentType, err := normalizeGrokMediaForwardBody(GrokMediaEndpointVideosGenerations, body, "application/json")
require.NoError(t, err)
require.Equal(t, "application/json", contentType)
require.Equal(t, "grok-imagine-video-1.5", gjson.GetBytes(out, "model").String())
require.Equal(t, "https://example.com/source.png", gjson.GetBytes(out, "image.url").String())
require.False(t, gjson.GetBytes(out, "image.image_url").Exists())
}
func TestNormalizeGrokMediaForwardBodyPreservesImageToVideoModelForOfficialURL(t *testing.T) {
body := []byte(`{
"model":"grok-imagine-video-1.5",
"prompt":"animate",
"image":{"url":"https://example.com/source.png"}
}`)
out, _, err := normalizeGrokMediaForwardBody(GrokMediaEndpointVideosGenerations, body, "application/json")
require.NoError(t, err)
require.Equal(t, "grok-imagine-video-1.5", gjson.GetBytes(out, "model").String())
require.Equal(t, "https://example.com/source.png", gjson.GetBytes(out, "image.url").String())
}
func TestCanonicalizeGrokMediaImageURLFieldsPreservesOfficialURL(t *testing.T) {
body := []byte(`{
"image":{"url":"https://example.com/official.png","image_url":"https://example.com/legacy.png"},
"images":[
{"image_url":"https://example.com/first.png"},
{"url":"https://example.com/second.png"}
],
"reference_images":[{"image_url":"https://example.com/reference.png"}],
"mask":{"image_url":"https://example.com/mask.png"}
}`)
out, err := canonicalizeGrokMediaImageURLFields(body, "image", "images", "reference_images", "mask")
require.NoError(t, err)
require.Equal(t, "https://example.com/official.png", gjson.GetBytes(out, "image.url").String())
require.False(t, gjson.GetBytes(out, "image.image_url").Exists())
require.Equal(t, "https://example.com/first.png", gjson.GetBytes(out, "images.0.url").String())
require.False(t, gjson.GetBytes(out, "images.0.image_url").Exists())
require.Equal(t, "https://example.com/second.png", gjson.GetBytes(out, "images.1.url").String())
require.Equal(t, "https://example.com/reference.png", gjson.GetBytes(out, "reference_images.0.url").String())
require.False(t, gjson.GetBytes(out, "reference_images.0.image_url").Exists())
require.Equal(t, "https://example.com/mask.png", gjson.GetBytes(out, "mask.url").String())
require.False(t, gjson.GetBytes(out, "mask.image_url").Exists())
}
func TestCanonicalizeGrokMediaImageURLFieldsReplacesEmptyOfficialURL(t *testing.T) {
body := []byte(`{"image":{"url":" ","image_url":"https://example.com/legacy.png"}}`)
out, err := canonicalizeGrokMediaImageURLFields(body, "image")
require.NoError(t, err)
require.Equal(t, "https://example.com/legacy.png", gjson.GetBytes(out, "image.url").String())
require.False(t, gjson.GetBytes(out, "image.image_url").Exists())
}
func TestPrepareGrokImageEditNormalizesOfficialImageObjects(t *testing.T) {
body := []byte(`{
"model":"grok-imagine-image-quality",
"image":{"image_url":{"url":"https://example.com/first.png"}},
"images":["https://example.com/second.png"],
"mask":{"image_url":"https://example.com/mask.png"}
}`)
out, contentType, err := prepareGrokMediaForwardBody(GrokMediaEndpointImagesEdits, body, "application/json")
require.NoError(t, err)
require.Equal(t, "application/json", contentType)
for _, path := range []string{"image", "images.0", "mask"} {
require.Equal(t, "image_url", gjson.GetBytes(out, path+".type").String())
require.NotEmpty(t, gjson.GetBytes(out, path+".url").String())
require.False(t, gjson.GetBytes(out, path+".image_url").Exists())
}
}
func TestPrepareGrokImageEditRejectsMoreThanThreeSources(t *testing.T) {
body := []byte(`{"images":["https://example.com/1.png","https://example.com/2.png","https://example.com/3.png","https://example.com/4.png"]}`)
out, _, err := prepareGrokMediaForwardBody(GrokMediaEndpointImagesEdits, body, "application/json")
require.Error(t, err)
require.Nil(t, out)
require.Contains(t, err.Error(), "maximum of 3 source images")
}
func TestNormalizeGrokMediaModelForEndpoint(t *testing.T) {
tests := []struct {
name string
endpoint GrokMediaEndpoint
model string
hasInputImage bool
want string
}{
{name: "image generation alias", endpoint: GrokMediaEndpointImagesGenerations, model: "grok-imagine", want: "grok-imagine-image-quality"},
{name: "image edit alias", endpoint: GrokMediaEndpointImagesEdits, model: "grok-imagine", want: "grok-imagine-image-quality"},
{name: "image quality passthrough", endpoint: GrokMediaEndpointImagesGenerations, model: "grok-imagine-image-quality", want: "grok-imagine-image-quality"},
{name: "image fast passthrough", endpoint: GrokMediaEndpointImagesGenerations, model: "grok-imagine-image", want: "grok-imagine-image"},
{name: "video passthrough", endpoint: GrokMediaEndpointVideosGenerations, model: "grok-imagine-video", want: "grok-imagine-video"},
{name: "video 1.5 text-only remains explicit", endpoint: GrokMediaEndpointVideosGenerations, model: "grok-imagine-video-1.5", want: "grok-imagine-video-1.5"},
{name: "video 1.5 image-to-video passthrough", endpoint: GrokMediaEndpointVideosGenerations, model: "grok-imagine-video-1.5", hasInputImage: true, want: "grok-imagine-video-1.5"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
require.Equal(t, tt.want, NormalizeGrokMediaModelForEndpoint(tt.endpoint, tt.model, tt.hasInputImage))
})
}
}
func TestForwardGrokMediaImagesGenerationNormalizesImagineAlias(t *testing.T) {
t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
body := []byte(`{"model":"grok-imagine","prompt":"draw a cat"}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body))
c.Request.Header.Set("Content-Type", "application/json")
account := &Account{
ID: 61,
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"},
"Xai-Request-Id": []string{"xai-image-req"},
},
Body: io.NopCloser(strings.NewReader(`{"data":[{"url":"https://images.test/cat.png"}]}`)),
}}
svc := &OpenAIGatewayService{httpUpstream: upstream}
result, err := svc.ForwardGrokMedia(context.Background(), c, account, GrokMediaEndpointImagesGenerations, "", body, "application/json")
require.NoError(t, err)
require.Equal(t, "https://xai.test/v1/images/generations", upstream.lastReq.URL.String())
require.Equal(t, http.MethodPost, upstream.lastReq.Method)
require.Equal(t, "Bearer api-key", upstream.lastReq.Header.Get("Authorization"))
require.Equal(t, "application/json", upstream.lastReq.Header.Get("Content-Type"))
require.Empty(t, upstream.lastReq.Header.Get("X-Grok-Client-Version"))
require.NotEqual(t, grokUpstreamUserAgent, upstream.lastReq.Header.Get("User-Agent"))
require.JSONEq(t, `{"model":"grok-imagine-image-quality","prompt":"draw a cat"}`, string(upstream.lastBody))
require.Equal(t, http.StatusOK, recorder.Code)
require.JSONEq(t, `{"data":[{"url":"https://images.test/cat.png"}]}`, recorder.Body.String())
require.Equal(t, "xai-image-req", result.RequestID)
require.Equal(t, "grok-imagine-image-quality", result.Model)
require.Equal(t, "grok-imagine-image-quality", result.BillingModel)
require.Equal(t, 1, result.ImageCount)
require.Equal(t, ImageBillingSize2K, result.ImageSize)
}
func TestForwardGrokMediaAppliesAccountModelMappingAfterEndpointNormalization(t *testing.T) {
t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
gin.SetMode(gin.TestMode)
tests := []struct {
name string
endpoint GrokMediaEndpoint
path string
body string
modelMapping map[string]any
wantRequestModel string
wantUpstream string
wantBody string
responseBody string
}{
{
name: "image generation maps normalized image alias",
endpoint: GrokMediaEndpointImagesGenerations,
path: "/v1/images/generations",
body: `{"model":"grok-imagine","prompt":"draw a cat"}`,
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 a cat"}`,
responseBody: `{"data":[{"url":"https://images.test/mapped.png"}]}`,
},
{
name: "video generation maps text-only fallback model",
endpoint: GrokMediaEndpointVideosGenerations,
path: "/v1/videos/generations",
body: `{"model":"grok-imagine-video-1.5","prompt":"waves"}`,
modelMapping: map[string]any{"grok-imagine-video": "grok-image-video"},
wantRequestModel: "grok-imagine-video-1.5",
wantUpstream: "grok-imagine-video-1.5",
wantBody: `{"model":"grok-imagine-video-1.5","prompt":"waves"}`,
responseBody: `{"request_id":"video-request-mapped"}`,
},
{
name: "image-to-video preserves then maps the requested model",
endpoint: GrokMediaEndpointVideosGenerations,
path: "/v1/videos/generations",
body: `{"model":"grok-imagine-video-1.5","prompt":"animate","image":{"url":"https://example.com/input.png"}}`,
modelMapping: map[string]any{"grok-imagine-video-1.5": "vendor-image-video"},
wantRequestModel: "grok-imagine-video-1.5",
wantUpstream: "vendor-image-video",
wantBody: `{"model":"vendor-image-video","prompt":"animate","image":{"url":"https://example.com/input.png"}}`,
responseBody: `{"request_id":"image-video-request-mapped"}`,
},
{
name: "mapping and image sanitization compose",
endpoint: GrokMediaEndpointImagesGenerations,
path: "/v1/images/generations",
body: `{"model":"grok-imagine","prompt":"draw","size":"1024x1024"}`,
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"}`,
responseBody: `{"data":[{"url":"https://images.test/mapped.png"}]}`,
},
{
name: "whitespace mapping target safely preserves normalized model",
endpoint: GrokMediaEndpointImagesGenerations,
path: "/v1/images/generations",
body: `{"model":"grok-imagine","prompt":"draw"}`,
modelMapping: map[string]any{"grok-imagine-image-quality": " "},
wantRequestModel: "grok-imagine-image-quality",
wantUpstream: "grok-imagine-image-quality",
wantBody: `{"model":"grok-imagine-image-quality","prompt":"draw"}`,
responseBody: `{"data":[{"url":"https://images.test/mapped.png"}]}`,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodPost, tt.path, strings.NewReader(tt.body))
c.Request.Header.Set("Content-Type", "application/json")
account := &Account{
ID: 66,
Name: "grok-mapped",
Platform: PlatformGrok,
Type: AccountTypeAPIKey,
Concurrency: 1,
Credentials: map[string]any{
"api_key": "api-key",
"base_url": "https://xai.test/v1",
"model_mapping": tt.modelMapping,
},
}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(tt.responseBody)),
}}
svc := &OpenAIGatewayService{httpUpstream: upstream}
result, err := svc.ForwardGrokMedia(context.Background(), c, account, tt.endpoint, "", []byte(tt.body), "application/json")
require.NoError(t, err)
require.JSONEq(t, tt.wantBody, string(upstream.lastBody))
require.Equal(t, tt.wantRequestModel, result.Model)
require.Equal(t, tt.wantRequestModel, result.BillingModel)
require.Equal(t, tt.wantUpstream, result.UpstreamModel)
})
}
}
func TestForwardGrokMediaImagesGenerationRejectsEmptySuccessfulResponse(t *testing.T) {
t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
body := []byte(`{"model":"grok-imagine-image","prompt":"draw a cat"}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body))
c.Request.Header.Set("Content-Type", "application/json")
account := &Account{
ID: 66,
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":[]}`)),
}}
svc := &OpenAIGatewayService{httpUpstream: upstream}
result, err := svc.ForwardGrokMedia(context.Background(), c, account, GrokMediaEndpointImagesGenerations, "", body, "application/json")
require.Nil(t, result)
var failoverErr *UpstreamFailoverError
require.ErrorAs(t, err, &failoverErr)
require.Equal(t, http.StatusBadGateway, failoverErr.StatusCode)
require.JSONEq(t, `{"data":[]}`, string(failoverErr.ResponseBody))
require.Empty(t, recorder.Body.String())
}
func TestForwardGrokMediaImagesGenerationStripsUnsupportedSize(t *testing.T) {
t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
body := []byte(`{"model":"grok-imagine-image","prompt":"draw a cat","size":"1024x1024"}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body))
c.Request.Header.Set("Content-Type", "application/json")
account := &Account{
ID: 65,
Name: "grok",
Platform: PlatformGrok,
Type: AccountTypeAPIKey,
Concurrency: 1,
Credentials: map[string]any{
"api_key": "api-key",
"base_url": "https://xai.test/v1",
"model_mapping": map[string]any{"grok-imagine-edit": "vendor-image-edit"},
},
}
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/cat.png"}]}`)),
}}
svc := &OpenAIGatewayService{httpUpstream: upstream}
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.Equal(t, ImageBillingSize1K, result.ImageSize)
require.Equal(t, "1024x1024", result.ImageInputSize)
}
func TestForwardGrokMediaImagesEditMultipartConvertsToJSON(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"))
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: 62,
Name: "grok",
Platform: PlatformGrok,
Type: AccountTypeAPIKey,
Concurrency: 1,
Credentials: map[string]any{
"api_key": "api-key",
"base_url": "https://xai.test/v1",
"model_mapping": map[string]any{"grok-imagine-edit": "vendor-image-edit"},
},
}
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.Equal(t, "https://xai.test/v1/images/edits", upstream.lastReq.URL.String())
require.Equal(t, "application/json", upstream.lastReq.Header.Get("Content-Type"))
require.True(t, json.Valid(upstream.lastBody))
require.Equal(t, "vendor-image-edit", gjson.GetBytes(upstream.lastBody, "model").String())
require.Equal(t, "edit this private image", gjson.GetBytes(upstream.lastBody, "prompt").String())
require.True(t, strings.HasPrefix(gjson.GetBytes(upstream.lastBody, "image.url").String(), "data:image/png;base64,"))
require.False(t, gjson.GetBytes(upstream.lastBody, "image.image_url").Exists())
require.Equal(t, "grok-imagine-edit", result.BillingModel)
require.Equal(t, "vendor-image-edit", result.UpstreamModel)
}
func TestForwardGrokMediaVideoGenerationReturnsUsageAndResponseID(t *testing.T) {
t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
body := []byte(`{"model":"grok-imagine-video-1.5","prompt":"waves","resolution":"720p","duration":10}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/videos/generations", bytes.NewReader(body))
c.Request.Header.Set("Content-Type", "application/json")
account := &Account{
ID: 63,
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"},
"Xai-Request-Id": []string{"xai-video-generate-req"},
},
Body: io.NopCloser(strings.NewReader(`{"request_id":"video-request-123","usage":{"prompt_tokens":3,"completion_tokens":4}}`)),
}}
svc := &OpenAIGatewayService{httpUpstream: upstream}
result, err := svc.ForwardGrokMedia(context.Background(), c, account, GrokMediaEndpointVideosGenerations, "", body, "application/json")
require.NoError(t, err)
require.Equal(t, "https://xai.test/v1/videos/generations", upstream.lastReq.URL.String())
require.JSONEq(t, `{"model":"grok-imagine-video-1.5","prompt":"waves","resolution":"720p","duration":10}`, string(upstream.lastBody))
require.Equal(t, "video-request-123", result.ResponseID)
require.Equal(t, "grok-imagine-video-1.5", result.BillingModel)
require.Equal(t, 3, result.Usage.InputTokens)
require.Equal(t, 4, result.Usage.OutputTokens)
// Create accepts the job only — VideoCount stays 0 until status returns video.url.
require.Equal(t, 0, result.ImageCount)
require.Empty(t, result.ImageSize)
require.Equal(t, 0, result.VideoCount)
require.Equal(t, VideoBillingResolution720P, result.VideoResolution)
require.Equal(t, 10, result.VideoDurationSeconds)
}
func TestForwardGrokMediaVideoGenerationReturnsTaskIDAsResponseID(t *testing.T) {
t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
body := []byte(`{"model":"grok-imagine-video","prompt":"waves"}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/videos/generations", bytes.NewReader(body))
c.Request.Header.Set("Content-Type", "application/json")
account := &Account{
ID: 63,
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(`{"task_id":"video-task-123"}`)),
}}
svc := &OpenAIGatewayService{httpUpstream: upstream}
result, err := svc.ForwardGrokMedia(context.Background(), c, account, GrokMediaEndpointVideosGenerations, "", body, "application/json")
require.NoError(t, err)
require.Equal(t, "video-task-123", result.ResponseID)
}
func TestExtractGrokMediaVideoRequestIDPreservesExistingPrecedence(t *testing.T) {
body := []byte(`{
"request_id":"request-id",
"id":"id",
"task_id":"task-id",
"data":{"request_id":"data-request-id","id":"data-id","task_id":"data-task-id"},
"video":{"request_id":"video-request-id","id":"video-id","task_id":"video-task-id"}
}`)
require.Equal(t, "request-id", extractGrokMediaVideoRequestID(body))
}
func TestForwardGrokMediaVideoGenerationPreservesImageToVideoModel(t *testing.T) {
t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
body := []byte(`{"model":"grok-imagine-video-1.5","prompt":"animate","image":{"image_url":"data:image/png;base64,aW1n"}}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/videos/generations", bytes.NewReader(body))
c.Request.Header.Set("Content-Type", "application/json")
account := &Account{
ID: 63,
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(`{"request_id":"video-request-456"}`)),
}}
svc := &OpenAIGatewayService{httpUpstream: upstream}
result, err := svc.ForwardGrokMedia(context.Background(), c, account, GrokMediaEndpointVideosGenerations, "", body, "application/json")
require.NoError(t, err)
require.Equal(t, "https://xai.test/v1/videos/generations", upstream.lastReq.URL.String())
require.JSONEq(t, `{"model":"grok-imagine-video-1.5","prompt":"animate","image":{"url":"data:image/png;base64,aW1n"}}`, string(upstream.lastBody))
require.Equal(t, "video-request-456", result.ResponseID)
require.Equal(t, "grok-imagine-video-1.5", result.BillingModel)
// 未指定 duration 时按上游默认 8 秒计费。
require.Equal(t, VideoBillingDefaultDurationSeconds, result.VideoDurationSeconds)
}
func TestForwardGrokMediaOAuthImageToVideoUsesOfficialAPIForLargeBody(t *testing.T) {
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
imageData := strings.Repeat("A", 2*1024*1024)
body := []byte(`{"model":"grok-imagine-video-1.5","prompt":"animate","image":{"image_url":"data:image/png;base64,` + imageData + `"}}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/videos/generations", bytes.NewReader(body))
c.Request.Header.Set("Content-Type", "application/json")
account := &Account{
ID: 66,
Name: "grok-oauth",
Platform: PlatformGrok,
Type: AccountTypeOAuth,
Status: StatusActive,
Schedulable: true,
Concurrency: 1,
Credentials: map[string]any{
"access_token": "oauth-access-token",
"refresh_token": "oauth-refresh-token",
"expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
"base_url": xai.DefaultCLIBaseURL,
},
}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{
"Content-Type": []string{"application/json"},
},
Body: io.NopCloser(strings.NewReader(`{"request_id":"video-request-oauth"}`)),
}}
svc := &OpenAIGatewayService{httpUpstream: upstream, grokTokenProvider: NewGrokTokenProvider(nil, nil)}
_, err := svc.ForwardGrokMedia(context.Background(), c, account, GrokMediaEndpointVideosGenerations, "", body, "application/json")
require.NoError(t, err)
require.Equal(t, xai.DefaultBaseURL+"/videos/generations", upstream.lastReq.URL.String())
require.Empty(t, upstream.lastReq.Header.Get("X-XAI-Token-Auth"))
require.Empty(t, upstream.lastReq.Header.Get("x-grok-client-version"))
require.Equal(t, "data:image/png;base64,"+imageData, gjson.GetBytes(upstream.lastBody, "image.url").String())
require.False(t, gjson.GetBytes(upstream.lastBody, "image.image_url").Exists())
}
func TestForwardGrokMediaVideoStatusUsesGETWithoutBody(t *testing.T) {
t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodGet, "/v1/videos/request-123", nil)
account := &Account{
ID: 62,
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"},
"Xai-Request-Id": []string{"xai-video-req"},
},
Body: io.NopCloser(strings.NewReader(`{"id":"request-123","status":"completed"}`)),
}}
svc := &OpenAIGatewayService{httpUpstream: upstream}
result, err := svc.ForwardGrokMedia(context.Background(), c, account, GrokMediaEndpointVideoStatus, "request-123", nil, "")
require.NoError(t, err)
require.Equal(t, "https://xai.test/v1/videos/request-123", upstream.lastReq.URL.String())
require.Equal(t, http.MethodGet, upstream.lastReq.Method)
require.Equal(t, "Bearer api-key", upstream.lastReq.Header.Get("Authorization"))
require.Empty(t, upstream.lastReq.Header.Get("X-Grok-Client-Version"))
require.NotEqual(t, grokUpstreamUserAgent, upstream.lastReq.Header.Get("User-Agent"))
require.Empty(t, upstream.lastReq.Header.Get("Content-Type"))
require.Empty(t, upstream.lastBody)
require.Equal(t, http.StatusOK, recorder.Code)
require.JSONEq(t, `{"id":"request-123","status":"completed"}`, recorder.Body.String())
require.Equal(t, "xai-video-req", result.RequestID)
}
func TestForwardGrokMediaVideoMutationEndpoints(t *testing.T) {
t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
gin.SetMode(gin.TestMode)
tests := []struct {
name string
endpoint GrokMediaEndpoint
path string
}{
{name: "edit", endpoint: GrokMediaEndpointVideosEdits, path: "/videos/edits"},
{name: "extension", endpoint: GrokMediaEndpointVideosExtensions, path: "/videos/extensions"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
body := []byte(`{"model":"grok-imagine-video","prompt":"continue","video":{"url":"https://example.com/in.mp4"},"duration":6}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1"+tt.path, bytes.NewReader(body))
c.Request.Header.Set("Content-Type", "application/json")
account := &Account{
ID: 71, Name: "grok", Platform: PlatformGrok, Type: AccountTypeAPIKey, Concurrency: 1,
Credentials: map[string]any{
"api_key": "api-key",
"base_url": "https://xai.test/v1",
"model_mapping": map[string]any{"grok-imagine-video": "vendor-video-mutation"},
},
}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`{"request_id":"video-mutation-123"}`)),
}}
svc := &OpenAIGatewayService{httpUpstream: upstream}
result, err := svc.ForwardGrokMedia(context.Background(), c, account, tt.endpoint, "", body, "application/json")
require.NoError(t, err)
require.Equal(t, "https://xai.test/v1"+tt.path, upstream.lastReq.URL.String())
require.Equal(t, http.MethodPost, upstream.lastReq.Method)
require.JSONEq(t, `{"model":"vendor-video-mutation","prompt":"continue","video":{"url":"https://example.com/in.mp4"},"duration":6}`, string(upstream.lastBody))
require.Equal(t, "video-mutation-123", result.ResponseID)
require.Equal(t, 0, result.VideoCount)
require.Equal(t, 6, result.VideoDurationSeconds)
require.Equal(t, "grok-imagine-video", result.BillingModel)
require.Equal(t, "vendor-video-mutation", result.UpstreamModel)
})
}
}
func TestGrokMediaVideoRequestBindingIsScopedToUserAndAPIKey(t *testing.T) {
gin.SetMode(gin.TestMode)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodGet, "/v1/videos/video-request-123", nil)
c.Request.Header.Set("session_id", "shared-client-session")
groupID := int64(7)
cache := &stubGatewayCache{}
svc := &OpenAIGatewayService{cache: cache}
const userID int64 = 41
const apiKeyID int64 = 51
require.NotEmpty(t, svc.GenerateExplicitSessionHash(c, nil))
ctx := c.Request.Context()
hash := GrokMediaVideoRequestSessionHash("video-request-123", userID, apiKeyID)
require.NotEmpty(t, hash)
require.NoError(t, svc.BindGrokMediaVideoRequestAccount(ctx, &groupID, "video-request-123", userID, apiKeyID, 63))
accountID, err := svc.ResolveGrokMediaVideoRequestAccount(ctx, &groupID, "video-request-123", userID, apiKeyID)
require.NoError(t, err)
require.Equal(t, int64(63), accountID)
accountID, err = svc.ResolveGrokMediaVideoRequestAccount(ctx, &groupID, "video-request-123", userID+1, apiKeyID)
require.Error(t, err)
require.Zero(t, accountID)
accountID, err = svc.ResolveGrokMediaVideoRequestAccount(ctx, &groupID, "video-request-123", userID, apiKeyID+1)
require.Error(t, err)
require.Zero(t, accountID)
}
func TestForwardGrokMedia429ReconcilesRateLimitBeforeCustomErrorBypass(t *testing.T) {
t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
body := []byte(`{"model":"grok-imagine","prompt":"draw a cat"}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body))
c.Request.Header.Set("Content-Type", "application/json")
account := &Account{
ID: 64,
Name: "grok",
Platform: PlatformGrok,
Type: AccountTypeAPIKey,
Concurrency: 1,
Credentials: map[string]any{
"api_key": "api-key",
"base_url": "https://xai.test/v1",
"custom_error_codes_enabled": true,
"custom_error_codes": []any{float64(http.StatusBadRequest)},
},
}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusTooManyRequests,
Header: http.Header{
"Content-Type": []string{"application/json"},
"Xai-Request-Id": []string{"xai-error-req"},
"Retry-After": []string{"45"},
},
Body: io.NopCloser(strings.NewReader(`{"error":{"message":"do not expose this upstream detail"}}`)),
}}
repo := &grokQuotaAccountRepo{}
svc := &OpenAIGatewayService{httpUpstream: upstream, accountRepo: repo}
result, err := svc.ForwardGrokMedia(context.Background(), c, account, GrokMediaEndpointImagesGenerations, "", body, "application/json")
require.Error(t, err)
require.Nil(t, result)
require.Equal(t, http.StatusInternalServerError, recorder.Code)
require.Contains(t, recorder.Body.String(), "Upstream gateway error")
require.NotContains(t, recorder.Body.String(), "do not expose")
require.Equal(t, 1, repo.rateLimitedCalls)
require.Zero(t, repo.tempUnschedCalls)
require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
}
func TestGrokMedia429FailoverPreservesRetryAfter(t *testing.T) {
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil)
account := &Account{
ID: 641, Name: "grok-oauth", Platform: PlatformGrok, Type: AccountTypeOAuth,
Status: StatusActive, Schedulable: true,
Credentials: map[string]any{
"custom_error_codes_enabled": true,
"custom_error_codes": []any{float64(http.StatusTooManyRequests)},
},
}
repo := &grokQuotaAccountRepo{}
svc := &OpenAIGatewayService{accountRepo: repo}
resp := &http.Response{
StatusCode: http.StatusTooManyRequests,
Header: http.Header{"Retry-After": []string{"45"}},
Body: io.NopCloser(strings.NewReader(`{"error":{"message":"rate limited"}}`)),
}
result, err := svc.handleGrokMediaErrorResponse(context.Background(), resp, c, account, "request-id", "grok-imagine")
require.Nil(t, result)
var failoverErr *UpstreamFailoverError
require.ErrorAs(t, err, &failoverErr)
require.Equal(t, http.StatusTooManyRequests, failoverErr.StatusCode)
require.Equal(t, "45", failoverErr.ResponseHeaders.Get("Retry-After"))
}
func healthyGrokOAuthGatewayTestAccount(id int64, token string) *Account {
return &Account{
ID: id,
Name: "grok",
Platform: PlatformGrok,
Type: AccountTypeOAuth,
Status: StatusActive,
Schedulable: true,
Concurrency: 1,
Credentials: map[string]any{
"access_token": token,
"refresh_token": "refresh-token",
"expires_at": time.Now().Add(2 * grokTokenRefreshSkew).UTC().Format(time.RFC3339),
"base_url": xai.DefaultCLIBaseURL,
},
}
}
func TestForwardAsChatCompletionsForGrokStopFallsBackToXAIChatCompletions(t *testing.T) {
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
body := []byte(`{"model":"grok","messages":[{"role":"user","content":"hi"}],"stream":false,"stop":"done","prompt_cache_key":"raw-client-cache-key"}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
c.Set("api_key", &APIKey{ID: 5101})
account := healthyGrokOAuthGatewayTestAccount(51, "access-token")
repo := &grokQuotaAccountRepo{
mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
accountsByID: map[int64]*Account{51: account},
},
}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{
"Content-Type": []string{"application/json"},
"Xai-Request-Id": []string{"xai-req"},
"X-Ratelimit-Limit-Requests": []string{"10"},
"X-Ratelimit-Remaining-Requests": []string{"9"},
"X-Ratelimit-Limit-Tokens": []string{"1000"},
"X-Ratelimit-Remaining-Tokens": []string{"990"},
},
Body: io.NopCloser(strings.NewReader(`{"id":"chatcmpl","object":"chat.completion","model":"grok-4.3","choices":[],"usage":{"prompt_tokens":1,"completion_tokens":2,"prompt_tokens_details":{"cached_tokens":1}}}`)),
}}
svc := &OpenAIGatewayService{
httpUpstream: upstream,
grokTokenProvider: NewGrokTokenProvider(repo, nil),
accountRepo: repo,
}
result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "")
require.NoError(t, err)
require.Equal(t, xai.DefaultCLIBaseURL+"/chat/completions", upstream.lastReq.URL.String())
require.Equal(t, "Bearer access-token", upstream.lastReq.Header.Get("Authorization"))
require.NotEmpty(t, upstream.lastReq.Header.Get(grokConversationIDHeader))
require.NotEqual(t, "raw-client-cache-key", upstream.lastReq.Header.Get(grokConversationIDHeader))
require.Equal(t, "grok-4.5", gjson.GetBytes(upstream.lastBody, "model").String())
require.False(t, gjson.GetBytes(upstream.lastBody, "prompt_cache_key").Exists())
require.Equal(t, "grok", result.Model)
require.Equal(t, "grok-4.5", result.UpstreamModel)
require.Equal(t, 1, result.Usage.InputTokens)
require.Equal(t, 2, result.Usage.OutputTokens)
require.Equal(t, 1, result.Usage.CacheReadInputTokens)
require.NotNil(t, repo.updates[51][grokQuotaSnapshotExtraKey])
require.Equal(t, http.StatusOK, recorder.Code)
}
func TestForwardGrokResponsesStreamingDefaultsEmptyModelTo45AndSnapshots(t *testing.T) {
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
body := []byte(`{"input":"hi","stream":true,"reasoning_effort":"high"}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
c.Request.Header.Set("Content-Type", "application/json")
c.Request.Header.Set("OpenAI-Beta", "responses=experimental")
c.Set("api_key", &APIKey{ID: 5201})
account := healthyGrokOAuthGatewayTestAccount(52, "access-token")
repo := &grokQuotaAccountRepo{
mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
accountsByID: map[int64]*Account{52: account},
},
}
upstreamBody := strings.Join([]string{
`data: {"type":"response.output_text.delta","sequence_number":0,"delta":"ok"}`,
"",
`data: {"type":"response.completed","sequence_number":1,"response":{"id":"resp_grok","model":"grok-4.3","usage":{"input_tokens":5,"output_tokens":3,"input_tokens_details":{"cached_tokens":2}}}}`,
"",
}, "\n")
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{
"Content-Type": []string{"text/event-stream"},
"Xai-Request-Id": []string{"xai-stream-req"},
"X-Ratelimit-Limit-Requests": []string{"10"},
"X-Ratelimit-Remaining-Requests": []string{"8"},
"X-Ratelimit-Limit-Tokens": []string{"1000"},
"X-Ratelimit-Remaining-Tokens": []string{"990"},
},
Body: io.NopCloser(strings.NewReader(upstreamBody)),
}}
svc := &OpenAIGatewayService{
httpUpstream: upstream,
grokTokenProvider: NewGrokTokenProvider(repo, nil),
accountRepo: repo,
}
result, err := svc.forwardGrokResponses(context.Background(), c, account, body, "", true, time.Now())
require.NoError(t, err)
require.Equal(t, xai.DefaultCLIBaseURL+"/responses", upstream.lastReq.URL.String())
require.Equal(t, "Bearer access-token", upstream.lastReq.Header.Get("Authorization"))
require.Equal(t, "responses=experimental", upstream.lastReq.Header.Get("OpenAI-Beta"))
require.Equal(t, "grok-4.5", gjson.GetBytes(upstream.lastBody, "model").String())
require.NotEmpty(t, gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String())
require.Equal(t, gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String(), upstream.lastReq.Header.Get(grokConversationIDHeader))
require.Equal(t, "web_search", gjson.GetBytes(upstream.lastBody, "tools.0.type").String())
require.Equal(t, "x_search", gjson.GetBytes(upstream.lastBody, "tools.1.type").String())
require.Equal(t, "none", gjson.GetBytes(upstream.lastBody, "tool_choice").String())
require.Equal(t, "high", gjson.GetBytes(upstream.lastBody, "reasoning_effort").String())
require.True(t, gjson.GetBytes(upstream.lastBody, "stream").Bool())
require.True(t, result.Stream)
require.Equal(t, "resp_grok", result.ResponseID)
require.Equal(t, "xai-stream-req", result.RequestID)
require.Equal(t, 5, result.Usage.InputTokens)
require.Equal(t, 3, result.Usage.OutputTokens)
require.Equal(t, 2, result.Usage.CacheReadInputTokens)
require.NotNil(t, result.ReasoningEffort)
require.Equal(t, "high", *result.ReasoningEffort)
require.Contains(t, recorder.Header().Get("Content-Type"), "text/event-stream")
require.Contains(t, recorder.Body.String(), "response.output_text.delta")
require.NotNil(t, repo.updates[52][grokQuotaSnapshotExtraKey])
}
func TestForwardGrokResponsesAPIKeyUsesXAIResponses(t *testing.T) {
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
body := []byte(`{"model":"grok","input":"hi","stream":true}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
c.Request.Header.Set("Content-Type", "application/json")
account := &Account{
ID: 53,
Name: "grok-api-key",
Platform: PlatformGrok,
Type: AccountTypeAPIKey,
Concurrency: 2,
Credentials: map[string]any{
"api_key": "xai-test-key",
"base_url": "https://api.x.ai/v1",
},
}
upstreamBody := strings.Join([]string{
`data: {"type":"response.output_text.delta","sequence_number":0,"delta":"ok"}`,
"",
`data: {"type":"response.completed","sequence_number":1,"response":{"id":"resp_grok_api_key","model":"grok-4.5","usage":{"input_tokens":2,"output_tokens":1}}}`,
"",
}, "\n")
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
Body: io.NopCloser(strings.NewReader(upstreamBody)),
}}
svc := &OpenAIGatewayService{httpUpstream: upstream}
result, err := svc.forwardGrokResponses(context.Background(), c, account, body, "grok", true, time.Now())
require.NoError(t, err)
require.Equal(t, "https://api.x.ai/v1/responses", upstream.lastReq.URL.String())
require.Equal(t, "Bearer xai-test-key", upstream.lastReq.Header.Get("Authorization"))
require.Empty(t, upstream.lastReq.Header.Get("X-Grok-Client-Version"))
require.NotEqual(t, grokUpstreamUserAgent, upstream.lastReq.Header.Get("User-Agent"))
require.Equal(t, "grok-4.5", gjson.GetBytes(upstream.lastBody, "model").String())
require.Equal(t, "resp_grok_api_key", result.ResponseID)
require.Equal(t, 2, result.Usage.InputTokens)
require.Equal(t, 1, result.Usage.OutputTokens)
}
func TestForwardGrokResponsesRetriesInvalidEncryptedContentOnce(t *testing.T) {
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
body := []byte(`{
"model":"grok",
"input":[
{"type":"reasoning","summary":[{"type":"summary_text","text":"keep this summary"}],"encrypted_content":"encrypted-reasoning"},
{"type":"message","role":"user","content":[{"type":"input_text","text":"hi"}]}
],
"metadata":{"large_id":9007199254740993},
"stream":false
}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
c.Set("api_key", &APIKey{ID: 4535})
account := &Account{
ID: 4535,
Name: "grok-api-key",
Platform: PlatformGrok,
Type: AccountTypeAPIKey,
Status: StatusActive,
Schedulable: true,
Concurrency: 2,
Credentials: map[string]any{
"api_key": "same-token",
"base_url": "https://api.x.ai/v1",
},
}
upstream := &httpUpstreamRecorder{responses: []*http.Response{
{
StatusCode: http.StatusBadRequest,
Header: http.Header{
"Content-Type": []string{"application/json"},
"Xai-Request-Id": []string{"recoverable-first"},
},
Body: io.NopCloser(strings.NewReader(`{"code":"invalid-argument","error":"Could not decrypt the provided encrypted_content. Ensure the value is unmodified."}`)),
},
{
StatusCode: http.StatusOK,
Header: http.Header{
"Content-Type": []string{"application/json"},
"Xai-Request-Id": []string{"recovered-second"},
},
Body: io.NopCloser(strings.NewReader(`{"id":"resp_recovered","object":"response","model":"grok-4.5","status":"completed","output":[],"usage":{"input_tokens":2,"output_tokens":1}}`)),
},
}}
svc := &OpenAIGatewayService{httpUpstream: upstream}
result, err := svc.forwardGrokResponses(context.Background(), c, account, body, "grok", false, time.Now())
require.NoError(t, err)
require.NotNil(t, result)
require.Equal(t, "resp_recovered", result.ResponseID)
require.Equal(t, "recovered-second", result.RequestID)
require.Len(t, upstream.requests, 2)
require.Len(t, upstream.bodies, 2)
require.Equal(t, "reasoning", gjson.GetBytes(upstream.bodies[0], "input.0.type").String())
require.Equal(t, "encrypted-reasoning", gjson.GetBytes(upstream.bodies[0], "input.0.encrypted_content").String())
require.Equal(t, "reasoning", gjson.GetBytes(upstream.bodies[1], "input.0.type").String())
require.False(t, gjson.GetBytes(upstream.bodies[1], "input.0.encrypted_content").Exists())
require.Equal(t, "keep this summary", gjson.GetBytes(upstream.bodies[1], "input.0.summary.0.text").String())
require.Equal(t, "message", gjson.GetBytes(upstream.bodies[1], "input.1.type").String())
require.Equal(t, "9007199254740993", gjson.GetBytes(upstream.bodies[0], "metadata.large_id").Raw)
require.Equal(t, "9007199254740993", gjson.GetBytes(upstream.bodies[1], "metadata.large_id").Raw)
firstIdentity := gjson.GetBytes(upstream.bodies[0], "prompt_cache_key").String()
secondIdentity := gjson.GetBytes(upstream.bodies[1], "prompt_cache_key").String()
require.NotEmpty(t, firstIdentity)
require.Equal(t, firstIdentity, secondIdentity)
for _, req := range upstream.requests {
require.Equal(t, "Bearer same-token", req.Header.Get("Authorization"))
require.Equal(t, firstIdentity, req.Header.Get(grokConversationIDHeader))
}
require.Equal(t, StatusActive, account.Status)
_, hasUpstreamErrors := c.Get(OpsUpstreamErrorsKey)
require.False(t, hasUpstreamErrors)
_, hasTerminalStatus := c.Get(OpsUpstreamStatusCodeKey)
require.False(t, hasTerminalStatus)
}
func TestForwardGrokResponsesInvalidEncryptedContentRecoveryDoesNotOvermatch(t *testing.T) {
gin.SetMode(gin.TestMode)
matchingError := `{"code":"invalid-argument","error":"Could not decrypt the provided encrypted_content."}`
tests := []struct {
name string
requestBody string
responseBody string
}{
{
name: "different top-level code",
requestBody: `{"model":"grok","input":[{"type":"reasoning","encrypted_content":"cipher"}],"stream":false}`,
responseBody: `{"code":"bad-request","error":"Could not decrypt the provided encrypted_content."}`,
},
{
name: "message does not mention decryption",
requestBody: `{"model":"grok","input":[{"type":"reasoning","encrypted_content":"cipher"}],"stream":false}`,
responseBody: `{"code":"invalid-argument","error":"The provided encrypted_content is invalid."}`,
},
{
name: "request has no encrypted reasoning",
requestBody: `{"model":"grok","input":[{"type":"message","role":"user","content":"hi"}],"stream":false}`,
responseBody: matchingError,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
body := []byte(tt.requestBody)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
account := &Account{
ID: 4536,
Name: "grok-api-key",
Platform: PlatformGrok,
Type: AccountTypeAPIKey,
Concurrency: 1,
Credentials: map[string]any{"api_key": "token", "base_url": "https://api.x.ai/v1"},
}
upstream := &httpUpstreamRecorder{responses: []*http.Response{{
StatusCode: http.StatusBadRequest,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(tt.responseBody)),
}}}
svc := &OpenAIGatewayService{httpUpstream: upstream}
result, err := svc.forwardGrokResponses(context.Background(), c, account, body, "grok", false, time.Now())
require.Nil(t, result)
require.Error(t, err)
require.Len(t, upstream.requests, 1)
require.Len(t, upstream.bodies, 1)
})
}
}
func TestForwardGrokResponsesInvalidEncryptedContentRecoveryNestedErrorShape(t *testing.T) {
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
body := []byte(`{"model":"grok","input":[{"type":"reasoning","encrypted_content":"cipher"},{"type":"message","role":"user","content":"hi"}],"stream":false}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
account := &Account{
ID: 4538,
Name: "grok-api-key",
Platform: PlatformGrok,
Type: AccountTypeAPIKey,
Concurrency: 1,
Credentials: map[string]any{"api_key": "token", "base_url": "https://api.x.ai/v1"},
}
upstream := &httpUpstreamRecorder{responses: []*http.Response{
{
StatusCode: http.StatusBadRequest,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`{"code":"invalid-argument","error":{"message":"Could not decrypt the provided encrypted_content."}}`)),
},
{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`{"id":"resp_ok","output":[],"usage":{"input_tokens":1,"output_tokens":1}}`)),
},
}}
svc := &OpenAIGatewayService{httpUpstream: upstream}
result, err := svc.forwardGrokResponses(context.Background(), c, account, body, "grok", false, time.Now())
require.NoError(t, err)
require.NotNil(t, result)
require.Len(t, upstream.requests, 2)
require.True(t, gjson.GetBytes(upstream.bodies[0], "input.0.encrypted_content").Exists())
require.False(t, gjson.GetBytes(upstream.bodies[1], "input.0.encrypted_content").Exists())
}
func TestForwardGrokResponsesInvalidEncryptedContentRetryFailureIsTerminal(t *testing.T) {
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
body := []byte(`{"model":"grok","input":[{"type":"reasoning","encrypted_content":"cipher"},{"type":"message","role":"user","content":"hi"}],"stream":false}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
account := &Account{
ID: 4537,
Name: "grok-api-key",
Platform: PlatformGrok,
Type: AccountTypeAPIKey,
Concurrency: 1,
Credentials: map[string]any{"api_key": "same-token", "base_url": "https://api.x.ai/v1"},
}
newInvalidEncryptedResponse := func(requestID string) *http.Response {
return &http.Response{
StatusCode: http.StatusBadRequest,
Header: http.Header{
"Content-Type": []string{"application/json"},
"Xai-Request-Id": []string{requestID},
},
Body: io.NopCloser(strings.NewReader(`{"code":"invalid-argument","error":"Could not decrypt the provided encrypted_content."}`)),
}
}
upstream := &httpUpstreamRecorder{responses: []*http.Response{
newInvalidEncryptedResponse("recoverable-first"),
newInvalidEncryptedResponse("terminal-second"),
}}
svc := &OpenAIGatewayService{httpUpstream: upstream}
result, err := svc.forwardGrokResponses(context.Background(), c, account, body, "grok", false, time.Now())
require.Nil(t, result)
require.Error(t, err)
require.Len(t, upstream.requests, 2)
require.Len(t, upstream.bodies, 2)
require.True(t, gjson.GetBytes(upstream.bodies[0], "input.0.encrypted_content").Exists())
require.False(t, gjson.GetBytes(upstream.bodies[1], `input.#(type=="reasoning")`).Exists())
rawEvents, ok := c.Get(OpsUpstreamErrorsKey)
require.True(t, ok)
events, ok := rawEvents.([]*OpsUpstreamErrorEvent)
require.True(t, ok)
require.NotEmpty(t, events)
for _, event := range events {
require.NotEqual(t, "recoverable-first", event.UpstreamRequestID)
}
require.Equal(t, http.StatusBadRequest, c.GetInt(OpsUpstreamStatusCodeKey))
}
func TestForwardAsChatCompletionsForGrokAPIKeyUsesConfiguredRawEndpointWithoutOAuthIdentity(t *testing.T) {
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
body := []byte(`{"model":"grok","messages":[{"role":"user","content":"hi"}],"stream":false}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
c.Request.Header.Set("Content-Type", "application/json")
account := &Account{
ID: 706,
Platform: PlatformGrok,
Type: AccountTypeAPIKey,
Concurrency: 1,
Credentials: map[string]any{
"api_key": "third-party-key",
"base_url": "https://grok.example.test/v1",
},
}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`{"id":"chatcmpl","object":"chat.completion","model":"grok-4.5","choices":[],"usage":{"prompt_tokens":1,"completion_tokens":1}}`)),
}}
svc := &OpenAIGatewayService{cfg: rawChatCompletionsTestConfig(), httpUpstream: upstream}
_, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "")
require.NoError(t, err)
require.Equal(t, "https://grok.example.test/v1/chat/completions", upstream.lastReq.URL.String())
require.Equal(t, "Bearer third-party-key", upstream.lastReq.Header.Get("Authorization"))
require.Empty(t, upstream.lastReq.Header.Get("X-Grok-Client-Version"))
require.NotEqual(t, grokUpstreamUserAgent, upstream.lastReq.Header.Get("User-Agent"))
}
func TestForwardAsChatCompletionsForGrokAPIKeyRejectsNonStreamingResponseWithoutUsage(t *testing.T) {
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
body := []byte(`{"model":"grok","messages":[{"role":"user","content":"hi"}],"stream":false}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
c.Request.Header.Set("Content-Type", "application/json")
account := &Account{
ID: 707,
Platform: PlatformGrok,
Type: AccountTypeAPIKey,
Concurrency: 1,
Credentials: map[string]any{
"api_key": "third-party-key",
"base_url": "https://grok.example.test/v1",
},
}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(
`{"id":"resp_missing_usage","object":"chat.completion","model":"grok-4.5","choices":[{"index":0,"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}]}`,
)),
}}
svc := &OpenAIGatewayService{cfg: rawChatCompletionsTestConfig(), httpUpstream: upstream}
result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "")
require.Nil(t, result)
var failoverErr *UpstreamFailoverError
require.ErrorAs(t, err, &failoverErr)
require.Equal(t, http.StatusBadGateway, failoverErr.StatusCode)
require.Equal(t, grokMissingUsageErrorCode, gjson.GetBytes(failoverErr.ResponseBody, "error.code").String())
require.False(t, c.Writer.Written(), "an unbillable response must not be committed to the client")
require.Empty(t, recorder.Body.String())
}
func TestAccountTestServiceGrokAPIKeyUsesXAIResponses(t *testing.T) {
gin.SetMode(gin.TestMode)
account := &Account{
ID: 54,
Name: "grok-api-key",
Platform: PlatformGrok,
Type: AccountTypeAPIKey,
Concurrency: 2,
Credentials: map[string]any{
"api_key": "xai-test-key",
"base_url": "https://api.x.ai/v1",
},
}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
Body: io.NopCloser(strings.NewReader(
"data: {\"type\":\"response.output_text.delta\",\"delta\":\"ok\"}\n\n" +
"data: {\"type\":\"response.completed\"}\n\n",
)),
}}
svc := &AccountTestService{httpUpstream: upstream}
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/54/test", nil)
err := svc.testGrokAccountConnection(c, account, "grok", "", AccountTestModeDefault, AccountTestOptions{})
require.NoError(t, err)
require.Equal(t, "https://api.x.ai/v1/responses", upstream.lastReq.URL.String())
require.Equal(t, "Bearer xai-test-key", upstream.lastReq.Header.Get("Authorization"))
require.Contains(t, recorder.Body.String(), `"type":"test_complete"`)
}
func TestAccountTestServiceGrokAPIKeyAllowsConfiguredHTTPWhenGlobalPolicyDoes(t *testing.T) {
gin.SetMode(gin.TestMode)
account := &Account{
ID: 55,
Name: "grok-api-key-http",
Platform: PlatformGrok,
Type: AccountTypeAPIKey,
Concurrency: 1,
Credentials: map[string]any{
"api_key": "third-party-key",
"base_url": "http://grok.example.test/v1",
},
}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
Body: io.NopCloser(strings.NewReader(
"data: {\"type\":\"response.output_text.delta\",\"delta\":\"ok\"}\n\n" +
"data: {\"type\":\"response.completed\"}\n\n",
)),
}}
svc := &AccountTestService{cfg: rawChatCompletionsTestConfig(), httpUpstream: upstream}
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/55/test", nil)
err := svc.testGrokAccountConnection(c, account, "grok", "", AccountTestModeDefault, AccountTestOptions{})
require.NoError(t, err)
require.Equal(t, "http://grok.example.test/v1/responses", upstream.lastReq.URL.String())
require.Equal(t, "Bearer third-party-key", upstream.lastReq.Header.Get("Authorization"))
require.Empty(t, upstream.lastReq.Header.Get("X-Grok-Client-Version"))
require.Contains(t, recorder.Body.String(), `"type":"test_complete"`)
}
func TestAccountTestServiceGrokOAuthPaymentRequiredTemporarilyUnschedulesAccount(t *testing.T) {
gin.SetMode(gin.TestMode)
account := healthyGrokOAuthGatewayTestAccount(56, "access-token")
repo := &grokQuotaAccountRepo{}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusPaymentRequired,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`{"code":"personal-team-blocked:spending-limit"}`)),
}}
svc := &AccountTestService{
accountRepo: repo,
grokTokenProvider: NewGrokTokenProvider(repo, nil),
httpUpstream: upstream,
}
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/56/test", nil)
before := time.Now()
err := svc.testGrokAccountConnection(c, account, "grok", "", AccountTestModeDefault, AccountTestOptions{})
require.Error(t, err)
require.Zero(t, repo.tempUnschedCalls)
require.Equal(t, 1, repo.rateLimitedCalls)
require.Equal(t, account.ID, repo.lastRateLimitedID)
require.WithinDuration(t, before.Add(grokSpendingLimitProbeCooldown), repo.lastRateLimitResetAt, time.Second)
require.Contains(t, recorder.Body.String(), `"type":"error"`)
require.Contains(t, recorder.Body.String(), "Grok Responses API returned 402")
}
func TestForwardAsChatCompletionsForGrokStreamingUsesRawXAIChatCompletions(t *testing.T) {
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
body := []byte(`{"model":"grok","messages":[{"role":"user","content":"hi"}],"stream":true}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
c.Request.Header.Set("Content-Type", "application/json")
account := healthyGrokOAuthGatewayTestAccount(53, "access-token")
repo := &grokQuotaAccountRepo{
mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
accountsByID: map[int64]*Account{53: account},
},
}
upstreamBody := strings.Join([]string{
`data: {"id":"chatcmpl_grok","object":"chat.completion.chunk","model":"grok-4.3","choices":[{"index":0,"delta":{"content":"ok"}}]}`,
"",
`data: {"id":"chatcmpl_grok","object":"chat.completion.chunk","model":"grok-4.3","choices":[],"usage":{"prompt_tokens":6,"completion_tokens":4,"total_tokens":10,"prompt_tokens_details":{"cached_tokens":1}}}`,
"",
"data: [DONE]",
"",
}, "\n")
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{
"Content-Type": []string{"text/event-stream"},
"X-Request-Id": []string{"chat-stream-req"},
"X-Ratelimit-Limit-Requests": []string{"10"},
"X-Ratelimit-Remaining-Requests": []string{"7"},
},
Body: io.NopCloser(strings.NewReader(upstreamBody)),
}}
svc := &OpenAIGatewayService{
cfg: rawChatCompletionsTestConfig(),
httpUpstream: upstream,
grokTokenProvider: NewGrokTokenProvider(repo, nil),
accountRepo: repo,
}
result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "")
require.NoError(t, err)
require.Equal(t, xai.DefaultCLIBaseURL+"/chat/completions", upstream.lastReq.URL.String())
require.Equal(t, "Bearer access-token", upstream.lastReq.Header.Get("Authorization"))
require.Equal(t, "text/event-stream", upstream.lastReq.Header.Get("Accept"))
require.Equal(t, xai.CLIUserAgent(xai.CLIClientVersion), upstream.lastReq.Header.Get("User-Agent"))
require.Equal(t, "grok-4.5", gjson.GetBytes(upstream.lastBody, "model").String())
require.True(t, gjson.GetBytes(upstream.lastBody, "stream_options.include_usage").Bool())
require.True(t, result.Stream)
require.Equal(t, 6, result.Usage.InputTokens)
require.Equal(t, 4, result.Usage.OutputTokens)
require.Equal(t, 1, result.Usage.CacheReadInputTokens)
require.Contains(t, recorder.Body.String(), "data: [DONE]")
require.NotNil(t, repo.updates[53][grokQuotaSnapshotExtraKey])
}
func TestForwardGrokResponsesNonStreamingUsesCacheIdentityAndCachedUsage(t *testing.T) {
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
body := []byte(`{"model":"grok","input":"hi","stream":false,"tools":[{"type":"namespace","name":"client_tools"}],"tool_choice":{"type":"namespace","name":"client_tools"}}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
c.Request.Header.Set("Content-Type", "application/json")
c.Set("api_key", &APIKey{ID: 5202})
account := healthyGrokOAuthGatewayTestAccount(56, "access-token")
observedResetAt := time.Now().Add(-time.Second).UTC().Truncate(time.Second)
observedLimitedAt := observedResetAt.Add(-grokRateLimitRepeatCooldown)
account.RateLimitedAt = &observedLimitedAt
account.RateLimitResetAt = &observedResetAt
repo := &grokQuotaAccountRepo{
mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
accountsByID: map[int64]*Account{56: account},
},
recoveryClearResult: true,
}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{
"Content-Type": []string{"application/json"},
"Xai-Request-Id": []string{"xai-non-stream-req"},
},
Body: io.NopCloser(strings.NewReader(`{"id":"resp_grok_non_stream","object":"response","model":"grok-4.3","status":"completed","output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"ok"}]}],"usage":{"input_tokens":7,"output_tokens":2,"total_tokens":9,"input_tokens_details":{"cached_tokens":4}}}`)),
}}
svc := &OpenAIGatewayService{
httpUpstream: upstream,
grokTokenProvider: NewGrokTokenProvider(repo, nil),
accountRepo: repo,
}
result, err := svc.forwardGrokResponses(context.Background(), c, account, body, "grok", false, time.Now())
require.NoError(t, err)
require.NotNil(t, result)
require.False(t, result.Stream)
require.Equal(t, "resp_grok_non_stream", result.ResponseID)
require.Equal(t, 7, result.Usage.InputTokens)
require.Equal(t, 2, result.Usage.OutputTokens)
require.Equal(t, 4, result.Usage.CacheReadInputTokens)
require.Equal(t, "grok-4.5", gjson.GetBytes(upstream.lastBody, "model").String())
identity := gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String()
require.NotEmpty(t, identity)
require.Equal(t, identity, upstream.lastReq.Header.Get(grokConversationIDHeader))
// The sanitizer drops this unsupported client tool, but its explicit intent
// must still prevent native cache-routing tools from being injected.
require.False(t, gjson.GetBytes(upstream.lastBody, "tools").Exists())
require.False(t, gjson.GetBytes(upstream.lastBody, "tool_choice").Exists())
require.Equal(t, "resp_grok_non_stream", gjson.Get(recorder.Body.String(), "id").String())
require.Equal(t, 1, repo.recoveryClearCalls)
require.Equal(t, observedLimitedAt, repo.recoveryObservedAt)
require.Equal(t, observedResetAt, repo.recoveryObservedReset)
}
func TestForwardGrokResponsesFailoverKeepsCacheIdentityAcrossAccounts(t *testing.T) {
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
body := []byte(`{"model":"grok","input":[{"role":"user","content":"stable prefix"}],"stream":false}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
c.Set("api_key", &APIKey{ID: 5203})
newAccount := func(id int64, token string) *Account {
account := healthyGrokOAuthGatewayTestAccount(id, token)
account.Name = fmt.Sprintf("grok-%d", id)
return account
}
firstAccount := newAccount(58, "access-token-a")
secondAccount := newAccount(59, "access-token-b")
repo := &grokQuotaAccountRepo{
mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
accountsByID: map[int64]*Account{58: firstAccount, 59: secondAccount},
},
}
upstream := &httpUpstreamRecorder{responses: []*http.Response{
{
StatusCode: http.StatusServiceUnavailable,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`{"error":{"message":"temporary"}}`)),
},
{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`{"id":"resp_after_failover","object":"response","model":"grok-4.3","status":"completed","output":[],"usage":{"input_tokens":5,"output_tokens":1}}`)),
},
}}
svc := &OpenAIGatewayService{
httpUpstream: upstream,
grokTokenProvider: NewGrokTokenProvider(repo, nil),
accountRepo: repo,
}
_, err := svc.forwardGrokResponses(context.Background(), c, firstAccount, body, "grok", false, time.Now())
var failoverErr *UpstreamFailoverError
require.ErrorAs(t, err, &failoverErr)
result, err := svc.forwardGrokResponses(context.Background(), c, secondAccount, body, "grok", false, time.Now())
require.NoError(t, err)
require.NotNil(t, result)
require.Len(t, upstream.requests, 2)
require.Len(t, upstream.bodies, 2)
firstIdentity := gjson.GetBytes(upstream.bodies[0], "prompt_cache_key").String()
secondIdentity := gjson.GetBytes(upstream.bodies[1], "prompt_cache_key").String()
require.NotEmpty(t, firstIdentity)
require.Equal(t, firstIdentity, secondIdentity)
require.Equal(t, firstIdentity, upstream.requests[0].Header.Get(grokConversationIDHeader))
require.Equal(t, secondIdentity, upstream.requests[1].Header.Get(grokConversationIDHeader))
require.Equal(t, "Bearer access-token-a", upstream.requests[0].Header.Get("Authorization"))
require.Equal(t, "Bearer access-token-b", upstream.requests[1].Header.Get("Authorization"))
}
func TestForwardAsChatCompletionsForGrokStreamingStopFallsBackToRawXAIChatCompletions(t *testing.T) {
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
body := []byte(`{"model":"grok","messages":[{"role":"user","content":"hi"}],"stream":true,"stop":"done"}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
c.Request.Header.Set("Content-Type", "application/json")
c.Request.Header.Set(grokConversationIDHeader, "native-client-conversation")
c.Set("api_key", &APIKey{ID: 5301})
account := healthyGrokOAuthGatewayTestAccount(53, "access-token")
repo := &grokQuotaAccountRepo{
mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
accountsByID: map[int64]*Account{53: account},
},
}
upstreamBody := strings.Join([]string{
`data: {"id":"chatcmpl_grok","object":"chat.completion.chunk","model":"grok-4.3","choices":[{"index":0,"delta":{"content":"ok"}}]}`,
"",
`data: {"id":"chatcmpl_grok","object":"chat.completion.chunk","model":"grok-4.3","choices":[],"usage":{"prompt_tokens":6,"completion_tokens":4,"total_tokens":10,"prompt_tokens_details":{"cached_tokens":1}}}`,
"",
"data: [DONE]",
"",
}, "\n")
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{
"Content-Type": []string{"text/event-stream"},
"X-Request-Id": []string{"chat-stream-req"},
"X-Ratelimit-Limit-Requests": []string{"10"},
"X-Ratelimit-Remaining-Requests": []string{"7"},
},
Body: io.NopCloser(strings.NewReader(upstreamBody)),
}}
svc := &OpenAIGatewayService{
cfg: rawChatCompletionsTestConfig(),
httpUpstream: upstream,
grokTokenProvider: NewGrokTokenProvider(repo, nil),
accountRepo: repo,
}
result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "")
require.NoError(t, err)
require.Equal(t, xai.DefaultCLIBaseURL+"/chat/completions", upstream.lastReq.URL.String())
require.Equal(t, "Bearer access-token", upstream.lastReq.Header.Get("Authorization"))
require.Equal(t, "text/event-stream", upstream.lastReq.Header.Get("Accept"))
require.Equal(t, xai.CLIUserAgent(xai.CLIClientVersion), upstream.lastReq.Header.Get("User-Agent"))
require.Equal(t, grokCLIVersion, upstream.lastReq.Header.Get("X-Grok-Client-Version"))
require.NotEmpty(t, upstream.lastReq.Header.Get(grokConversationIDHeader))
require.NotEqual(t, "native-client-conversation", upstream.lastReq.Header.Get(grokConversationIDHeader))
require.Equal(t, "grok-4.5", gjson.GetBytes(upstream.lastBody, "model").String())
require.True(t, gjson.GetBytes(upstream.lastBody, "stream_options.include_usage").Bool())
require.True(t, result.Stream)
require.Equal(t, 6, result.Usage.InputTokens)
require.Equal(t, 4, result.Usage.OutputTokens)
require.Equal(t, 1, result.Usage.CacheReadInputTokens)
require.Contains(t, recorder.Body.String(), "data: [DONE]")
require.NotNil(t, repo.updates[53][grokQuotaSnapshotExtraKey])
}
func TestForwardAsChatCompletionsForGrokComposerBridgesImageInput(t *testing.T) {
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
body := []byte(`{"model":"grok-composer-2.5-fast","messages":[{"role":"system","content":"You are concise."},{"role":"user","content":[{"type":"text","text":"What is shown?"},{"type":"image_url","image_url":{"url":"data:image/png;base64,QUJD"}}]}],"stream":false}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
c.Request.Header.Set("Content-Type", "application/json")
c.Set("api_key", &APIKey{ID: 5501})
account := healthyGrokOAuthGatewayTestAccount(55, "access-token")
repo := &grokQuotaAccountRepo{
mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
accountsByID: map[int64]*Account{55: account},
},
}
upstream := &httpUpstreamRecorder{responses: []*http.Response{
{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}, "xai-request-id": []string{"vision-req"}},
Body: io.NopCloser(strings.NewReader(`{"id":"resp_vision","object":"response","model":"grok-build-0.1","output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"A small diagram with ABC letters."}]}],"usage":{"input_tokens":11,"output_tokens":7,"total_tokens":18}}`)),
},
{
StatusCode: http.StatusOK,
Header: http.Header{
"Content-Type": []string{"application/json"},
"X-Request-Id": []string{"composer-req"},
"X-Ratelimit-Limit-Requests": []string{"10"},
"X-Ratelimit-Remaining-Requests": []string{"9"},
"X-Ratelimit-Limit-Tokens": []string{"1000"},
"X-Ratelimit-Remaining-Tokens": []string{"980"},
},
Body: io.NopCloser(strings.NewReader(`{"id":"chatcmpl_composer","object":"chat.completion","model":"grok-composer-2.5-fast","choices":[{"index":0,"message":{"role":"assistant","content":"It shows ABC."},"finish_reason":"stop"}],"usage":{"prompt_tokens":3,"completion_tokens":5,"total_tokens":8}}`)),
},
}}
svc := &OpenAIGatewayService{
cfg: rawChatCompletionsTestConfig(),
httpUpstream: upstream,
grokTokenProvider: NewGrokTokenProvider(repo, nil),
accountRepo: repo,
}
result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "")
require.NoError(t, err)
require.NotNil(t, result)
require.Len(t, upstream.requests, 2)
require.Equal(t, xai.DefaultCLIBaseURL+"/responses", upstream.requests[0].URL.String())
require.Empty(t, upstream.requests[0].Header.Get(grokConversationIDHeader))
require.Equal(t, "grok-build-0.1", gjson.GetBytes(upstream.bodies[0], "model").String())
require.Equal(t, "input_image", gjson.GetBytes(upstream.bodies[0], "input.0.content.1.type").String())
require.Equal(t, xai.DefaultCLIBaseURL+"/chat/completions", upstream.requests[1].URL.String())
require.NotEmpty(t, upstream.requests[1].Header.Get(grokConversationIDHeader))
require.Equal(t, "grok-composer-2.5-fast", gjson.GetBytes(upstream.bodies[1], "model").String())
require.False(t, strings.Contains(string(upstream.bodies[1]), "image_url"))
require.Contains(t, gjson.GetBytes(upstream.bodies[1], "messages.1.content").String(), "Image 1 description")
require.Contains(t, gjson.GetBytes(upstream.bodies[1], "messages.1.content").String(), "A small diagram with ABC letters.")
require.Equal(t, 14, result.Usage.InputTokens)
require.Equal(t, 12, result.Usage.OutputTokens)
require.Equal(t, "It shows ABC.", gjson.Get(recorder.Body.String(), "choices.0.message.content").String())
require.NotNil(t, repo.updates[55][grokQuotaSnapshotExtraKey])
}
func TestForwardAsAnthropicForGrokUsesXAIResponses(t *testing.T) {
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
body := []byte(`{"model":"grok","max_tokens":32,"stream":false,"messages":[{"role":"user","content":"hi"}]}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(body))
c.Set("api_key", &APIKey{ID: 5401})
c.Request.Header.Set("OpenAI-Beta", "grok-experimental")
c.Request.Header.Set("originator", "opencode")
account := healthyGrokOAuthGatewayTestAccount(54, "access-token")
repo := &grokQuotaAccountRepo{
mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
accountsByID: map[int64]*Account{54: account},
},
}
upstream := &httpUpstreamRecorder{resp: grokMessagesSSECompletedResponse("resp_grok_messages", 3)}
svc := &OpenAIGatewayService{
httpUpstream: upstream,
grokTokenProvider: NewGrokTokenProvider(repo, nil),
accountRepo: repo,
}
result, err := svc.ForwardAsAnthropic(context.Background(), c, account, body, "", "")
require.NoError(t, err)
require.Equal(t, xai.DefaultCLIBaseURL+"/responses", upstream.lastReq.URL.String())
require.Equal(t, "Bearer access-token", upstream.lastReq.Header.Get("Authorization"))
require.Equal(t, xai.CLIUserAgent(xai.CLIClientVersion), upstream.lastReq.Header.Get("User-Agent"))
require.Equal(t, grokCLIVersion, upstream.lastReq.Header.Get("X-Grok-Client-Version"))
require.Equal(t, "grok-experimental", upstream.lastReq.Header.Get("OpenAI-Beta"))
require.Empty(t, upstream.lastReq.Header.Get("originator"))
require.Empty(t, upstream.lastReq.Header.Get("version"))
require.Equal(t, "grok-4.5", gjson.GetBytes(upstream.lastBody, "model").String())
require.NotEmpty(t, gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String())
require.Equal(t, gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String(), upstream.lastReq.Header.Get(grokConversationIDHeader))
require.Equal(t, "web_search", gjson.GetBytes(upstream.lastBody, "tools.0.type").String())
require.Equal(t, "x_search", gjson.GetBytes(upstream.lastBody, "tools.1.type").String())
require.Equal(t, "none", gjson.GetBytes(upstream.lastBody, "tool_choice").String())
require.Empty(t, upstream.lastReq.Header.Get("session_id"))
require.True(t, gjson.GetBytes(upstream.lastBody, "stream").Bool())
require.NotContains(t, string(upstream.lastBody), "chatgpt.com")
require.Equal(t, "grok", result.Model)
require.Equal(t, "grok-4.5", result.UpstreamModel)
require.Equal(t, 5, result.Usage.InputTokens)
require.Equal(t, 2, result.Usage.OutputTokens)
require.Equal(t, 3, result.Usage.CacheReadInputTokens)
require.Contains(t, recorder.Body.String(), `"type":"message"`)
require.Equal(t, int64(3), gjson.Get(recorder.Body.String(), "usage.cache_read_input_tokens").Int())
require.Contains(t, recorder.Body.String(), "ok")
}
func TestForwardAsAnthropicForGrokFunctionToolUsesCacheCapableMixedRoute(t *testing.T) {
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
body := []byte(`{
"model":"grok","max_tokens":32,"stream":false,
"messages":[{"role":"user","content":"look up alpha"}],
"tools":[{"name":"lookup","description":"look up a key","input_schema":{"type":"object","properties":{"key":{"type":"string"}},"required":["key"]}},{"name":"web_search","description":"search the web","input_schema":{"type":"object","properties":{"query":{"type":"string"}},"required":["query"]}}],
"tool_choice":{"type":"auto"}
}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(body))
c.Set("api_key", &APIKey{ID: 5403})
account := healthyGrokOAuthGatewayTestAccount(58, "access-token")
account.Extra = map[string]any{grokBillingExtraKey: map[string]any{
"status_code": http.StatusOK,
"source": "billing_probe",
"monthly_updated_at": "2026-07-15T05:00:00Z",
}}
repo := &grokQuotaAccountRepo{
mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
accountsByID: map[int64]*Account{58: account},
},
}
responseBody := strings.Join([]string{
`data: {"type":"response.completed","response":{"id":"resp_grok_function","object":"response","model":"grok-4.5","status":"completed","output":[{"type":"function_call","id":"fc_lookup","call_id":"call_lookup","name":"lookup","arguments":"{\"key\":\"alpha\"}","status":"completed"}],"usage":{"input_tokens":7000,"output_tokens":2,"total_tokens":7002,"input_tokens_details":{"cached_tokens":6144}}}}`,
"",
"data: [DONE]",
"",
}, "\n")
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
Body: io.NopCloser(strings.NewReader(responseBody)),
}}
svc := &OpenAIGatewayService{
httpUpstream: upstream,
grokTokenProvider: NewGrokTokenProvider(repo, nil),
accountRepo: repo,
}
result, err := svc.ForwardAsAnthropic(context.Background(), c, account, body, "", "")
require.NoError(t, err)
require.NotNil(t, result)
require.Equal(t, xai.DefaultCLIBaseURL+"/responses", upstream.lastReq.URL.String())
identity := gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String()
require.NotEmpty(t, identity)
require.Equal(t, identity, upstream.lastReq.Header.Get(grokConversationIDHeader))
tools := gjson.GetBytes(upstream.lastBody, "tools").Array()
require.Len(t, tools, 3)
require.Equal(t, "function", tools[0].Get("type").String())
require.Equal(t, "lookup", tools[0].Get("name").String())
require.Equal(t, "object", tools[0].Get("parameters.type").String())
require.Equal(t, "web_search", tools[1].Get("type").String())
require.Equal(t, "x_search", tools[2].Get("type").String())
require.Equal(t, "auto", gjson.GetBytes(upstream.lastBody, "tool_choice").String())
require.Equal(t, 7000, result.Usage.InputTokens)
require.Equal(t, 6144, result.Usage.CacheReadInputTokens)
clientBody := recorder.Body.String()
require.Equal(t, "tool_use", gjson.Get(clientBody, "content.0.type").String())
require.Equal(t, "call_lookup", gjson.Get(clientBody, "content.0.id").String())
require.Equal(t, "lookup", gjson.Get(clientBody, "content.0.name").String())
require.Equal(t, "alpha", gjson.Get(clientBody, "content.0.input.key").String())
require.Equal(t, "tool_use", gjson.Get(clientBody, "stop_reason").String())
require.Equal(t, int64(856), gjson.Get(clientBody, "usage.input_tokens").Int())
require.Equal(t, int64(6144), gjson.Get(clientBody, "usage.cache_read_input_tokens").Int())
}
func TestForwardAsAnthropicForGrokStreamingPreservesCacheUsage(t *testing.T) {
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
body := []byte(`{"model":"grok","max_tokens":32,"stream":true,"messages":[{"role":"user","content":"hi"}]}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(body))
c.Set("api_key", &APIKey{ID: 5402})
account := healthyGrokOAuthGatewayTestAccount(57, "access-token")
repo := &grokQuotaAccountRepo{
mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
accountsByID: map[int64]*Account{57: account},
},
}
upstream := &httpUpstreamRecorder{resp: grokMessagesSSECompletedResponse("resp_grok_messages_stream", 2)}
svc := &OpenAIGatewayService{
httpUpstream: upstream,
grokTokenProvider: NewGrokTokenProvider(repo, nil),
accountRepo: repo,
}
result, err := svc.ForwardAsAnthropic(context.Background(), c, account, body, "", "")
require.NoError(t, err)
require.NotNil(t, result)
require.Equal(t, 2, result.Usage.CacheReadInputTokens)
identity := gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String()
require.NotEmpty(t, identity)
require.Equal(t, identity, upstream.lastReq.Header.Get(grokConversationIDHeader))
require.Contains(t, recorder.Header().Get("Content-Type"), "text/event-stream")
require.Contains(t, recorder.Body.String(), `"cache_read_input_tokens":2`)
}
func grokMessagesSSECompletedResponse(responseID string, cachedTokens int) *http.Response {
body := strings.Join([]string{
fmt.Sprintf(`data: {"type":"response.completed","response":{"id":%q,"object":"response","model":"grok-4.3","status":"completed","output":[{"type":"message","id":"msg_1","role":"assistant","status":"completed","content":[{"type":"output_text","text":"ok"}]}],"usage":{"input_tokens":5,"output_tokens":2,"total_tokens":7,"input_tokens_details":{"cached_tokens":%d}}}}`, responseID, cachedTokens),
"",
"data: [DONE]",
"",
}, "\n")
return &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
Body: io.NopCloser(strings.NewReader(body)),
}
}
func TestHandleGrokAccountUpstreamErrorSpendingLimitUsesRecoverableProbeCool(t *testing.T) {
repo := &grokQuotaAccountRepo{}
svc := &OpenAIGatewayService{accountRepo: repo}
account := &Account{ID: 2570, Platform: PlatformGrok, Type: AccountTypeOAuth}
before := time.Now()
body := []byte(`{"code":"personal-team-blocked:spending-limit","error":"You have run out of credits"}`)
svc.handleGrokAccountUpstreamError(context.Background(), account, http.StatusForbidden, nil, body)
require.Equal(t, 1, repo.rateLimitedCalls)
require.WithinDuration(t, before.Add(grokSpendingLimitProbeCooldown), repo.lastRateLimitResetAt, 2*time.Second)
require.Zero(t, repo.tempUnschedCalls)
}
func TestHandleGrokAccountUpstreamErrorTempUnschedulesNonRateLimitStates(t *testing.T) {
tests := []struct {
name string
status int
headers http.Header
wantReason string
wantMinCooldown time.Duration
wantMaxCooldown time.Duration
}{
{
name: "unauthorized reauth",
status: http.StatusUnauthorized,
wantReason: "grok credentials unauthorized",
wantMinCooldown: 10*time.Minute - time.Second,
wantMaxCooldown: 10*time.Minute + time.Second,
},
{
name: "forbidden entitlement",
status: http.StatusForbidden,
wantReason: "grok access or entitlement denied",
wantMinCooldown: 30*time.Minute - time.Second,
wantMaxCooldown: 30*time.Minute + time.Second,
},
{
name: "payment required",
status: http.StatusPaymentRequired,
wantReason: "grok payment required",
wantMinCooldown: 30*time.Minute - time.Second,
wantMaxCooldown: 30*time.Minute + time.Second,
},
{
name: "upstream temporary error",
status: http.StatusInternalServerError,
wantReason: "grok upstream temporary error",
wantMinCooldown: 2*time.Minute - time.Second,
wantMaxCooldown: 2*time.Minute + time.Second,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
account := &Account{ID: 61, Platform: PlatformGrok, Type: AccountTypeOAuth}
repo := &grokQuotaAccountRepo{}
svc := &OpenAIGatewayService{accountRepo: repo}
before := time.Now()
svc.handleGrokAccountUpstreamError(context.Background(), account, tt.status, tt.headers, nil)
require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
require.Equal(t, 1, repo.tempUnschedCalls)
require.Zero(t, repo.rateLimitedCalls)
require.Equal(t, account.ID, repo.lastTempUnschedID)
require.Equal(t, tt.wantReason, repo.lastTempUnschedReason)
require.True(t, repo.lastTempUnschedUntil.After(before.Add(tt.wantMinCooldown)))
require.True(t, repo.lastTempUnschedUntil.Before(before.Add(tt.wantMaxCooldown)))
})
}
}
func TestHandleGrokAccountUpstreamErrorSpendingLimit403RateLimits(t *testing.T) {
account := &Account{ID: 614, Platform: PlatformGrok, Type: AccountTypeOAuth}
repo := &grokQuotaAccountRepo{}
svc := &OpenAIGatewayService{accountRepo: repo}
before := time.Now()
body := []byte(`{"code":"personal-team-blocked:spending-limit","error":"You have run out of credits"}`)
svc.handleGrokAccountUpstreamError(context.Background(), account, http.StatusForbidden, nil, body)
require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
require.Equal(t, 1, repo.rateLimitedCalls)
require.Equal(t, account.ID, repo.lastRateLimitedID)
require.WithinDuration(t, before.Add(grokSpendingLimitProbeCooldown), repo.lastRateLimitResetAt, 2*time.Second)
require.Zero(t, repo.tempUnschedCalls)
require.True(t, isGrokSpendingLimitError(body))
}
func TestHandleGrokAccountUpstreamError5xxRespectsPoolMode(t *testing.T) {
t.Run("pool mode keeps scheduling state", func(t *testing.T) {
account := &Account{
ID: 611,
Platform: PlatformGrok,
Type: AccountTypeAPIKey,
Credentials: map[string]any{
"pool_mode": true,
},
}
repo := &grokQuotaAccountRepo{}
svc := &OpenAIGatewayService{accountRepo: repo}
svc.handleGrokAccountUpstreamError(context.Background(), account, http.StatusBadGateway, nil, nil)
require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
require.Zero(t, repo.tempUnschedCalls)
require.Nil(t, account.TempUnschedulableUntil)
require.Empty(t, account.TempUnschedulableReason)
})
t.Run("non-pool mode keeps two minute cooldown", func(t *testing.T) {
account := &Account{ID: 612, Platform: PlatformGrok, Type: AccountTypeAPIKey}
repo := &grokQuotaAccountRepo{}
svc := &OpenAIGatewayService{accountRepo: repo}
before := time.Now()
svc.handleGrokAccountUpstreamError(context.Background(), account, http.StatusBadGateway, nil, nil)
require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
require.Equal(t, 1, repo.tempUnschedCalls)
require.Equal(t, account.ID, repo.lastTempUnschedID)
require.Equal(t, "grok upstream temporary error", repo.lastTempUnschedReason)
require.WithinDuration(t, before.Add(2*time.Minute), repo.lastTempUnschedUntil, time.Second)
})
}
func TestHandleGrokAccountUpstreamError429SetsRateLimitedFromRetryAfter(t *testing.T) {
account := &Account{ID: 61, Platform: PlatformGrok, Type: AccountTypeOAuth}
repo := &grokQuotaAccountRepo{}
svc := &OpenAIGatewayService{accountRepo: repo}
before := time.Now()
svc.handleGrokAccountUpstreamError(context.Background(), account, http.StatusTooManyRequests, http.Header{"Retry-After": []string{"45"}}, nil)
require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
require.Equal(t, 1, repo.rateLimitedCalls)
require.Equal(t, account.ID, repo.lastRateLimitedID)
require.WithinDuration(t, before.Add(45*time.Second), repo.lastRateLimitResetAt, time.Second)
require.Zero(t, repo.tempUnschedCalls)
}
func TestHandleGrokAccountUpstreamError429PoolModeKeepsSchedulingState(t *testing.T) {
account := &Account{
ID: 613,
Platform: PlatformGrok,
Type: AccountTypeAPIKey,
Credentials: map[string]any{
"pool_mode": true,
},
}
repo := &grokQuotaAccountRepo{}
svc := &OpenAIGatewayService{accountRepo: repo}
svc.handleGrokAccountUpstreamError(
context.Background(), account, http.StatusTooManyRequests,
http.Header{"Retry-After": []string{"45"}}, nil,
)
require.Equal(t, 1, repo.updateCalls, "pool mode should retain the quota snapshot for observability")
require.Zero(t, repo.rateLimitedCalls)
require.Zero(t, repo.tempUnschedCalls)
require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
require.Nil(t, account.RateLimitResetAt)
}
func TestHandleGrokAccountUpstreamError402RecoversAfterCooldownExpiry(t *testing.T) {
account := &Account{
ID: 610, Platform: PlatformGrok, Type: AccountTypeOAuth,
Status: StatusActive, Schedulable: true,
}
repo := &grokQuotaAccountRepo{}
svc := &OpenAIGatewayService{accountRepo: repo}
svc.handleGrokAccountUpstreamError(context.Background(), account, http.StatusPaymentRequired, nil, nil)
require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
require.Equal(t, 1, repo.tempUnschedCalls)
expired := time.Now().Add(-time.Second)
account.TempUnschedulableUntil = &expired
svc.openaiAccountRuntimeBlockUntil.Store(account.ID, expired)
require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
require.True(t, account.IsSchedulable())
}
func TestHandleGrokAccountUpstreamError429UsesLatestExhaustedWindowReset(t *testing.T) {
now := time.Now()
requestReset := now.Add(10 * time.Minute).Truncate(time.Second)
tokenReset := now.Add(20 * time.Minute).Truncate(time.Second)
headers := http.Header{
"X-Ratelimit-Limit-Requests": []string{"10"},
"X-Ratelimit-Remaining-Requests": []string{"0"},
"X-Ratelimit-Reset-Requests": []string{fmt.Sprintf("%d", requestReset.Unix())},
"X-Ratelimit-Limit-Tokens": []string{"1000"},
"X-Ratelimit-Remaining-Tokens": []string{"0"},
"X-Ratelimit-Reset-Tokens": []string{fmt.Sprintf("%d", tokenReset.Unix())},
}
account := &Account{ID: 62, Platform: PlatformGrok, Type: AccountTypeOAuth}
repo := &grokQuotaAccountRepo{}
svc := &OpenAIGatewayService{accountRepo: repo}
svc.handleGrokAccountUpstreamError(context.Background(), account, http.StatusTooManyRequests, headers, nil)
require.Equal(t, 1, repo.rateLimitedCalls)
require.WithinDuration(t, tokenReset, repo.lastRateLimitResetAt, time.Second)
require.Zero(t, repo.tempUnschedCalls)
}
func TestHandleGrokAccountUpstreamError429UsesFallbackReset(t *testing.T) {
account := &Account{ID: 63, Platform: PlatformGrok, Type: AccountTypeOAuth}
repo := &grokQuotaAccountRepo{}
svc := &OpenAIGatewayService{accountRepo: repo}
before := time.Now()
svc.handleGrokAccountUpstreamError(context.Background(), account, http.StatusTooManyRequests, nil, nil)
require.Equal(t, 1, repo.rateLimitedCalls)
require.WithinDuration(t, before.Add(grokRateLimitFallbackCooldown), repo.lastRateLimitResetAt, time.Second)
require.Zero(t, repo.tempUnschedCalls)
}
func TestGrokRateLimitResetAtForAccountEscalatesRepeated429s(t *testing.T) {
now := time.Now().UTC().Truncate(time.Second)
retryAfter := 45
snapshot := &xai.QuotaSnapshot{
StatusCode: http.StatusTooManyRequests,
RetryAfterSeconds: &retryAfter,
UpdatedAt: now.Format(time.RFC3339),
}
tests := []struct {
name string
previousCooldown time.Duration
wantCooldown time.Duration
}{
{name: "repeat after short boundary", previousCooldown: 45 * time.Second, wantCooldown: grokRateLimitRepeatCooldown},
{name: "sustained repeat", previousCooldown: grokRateLimitRepeatCooldown, wantCooldown: grokRateLimitSustainedCooldown},
{name: "capped repeat", previousCooldown: grokRateLimitSustainedCooldown, wantCooldown: grokRateLimitMaxAdaptiveCooldown},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
previousReset := now.Add(-time.Second)
previousLimited := previousReset.Add(-tt.previousCooldown)
account := &Account{
ID: 630,
Platform: PlatformGrok,
Type: AccountTypeOAuth,
RateLimitedAt: &previousLimited,
RateLimitResetAt: &previousReset,
}
resetAt, limited := grokRateLimitResetAtForAccount(account, snapshot, now)
require.True(t, limited)
require.WithinDuration(t, now.Add(tt.wantCooldown), resetAt, time.Second)
})
}
}
func TestGrokRateLimitResetAtForAccountPreservesAuthoritativeAndQuietRecovery(t *testing.T) {
now := time.Now().UTC().Truncate(time.Second)
retryAfter := 45
previousReset := now.Add(-grokRateLimitBackoffQuietPeriod - time.Second)
previousLimited := previousReset.Add(-grokRateLimitSustainedCooldown)
account := &Account{
ID: 631,
Platform: PlatformGrok,
Type: AccountTypeOAuth,
RateLimitedAt: &previousLimited,
RateLimitResetAt: &previousReset,
}
snapshot := &xai.QuotaSnapshot{
StatusCode: http.StatusTooManyRequests,
RetryAfterSeconds: &retryAfter,
UpdatedAt: now.Format(time.RFC3339),
}
resetAt, limited := grokRateLimitResetAtForAccount(account, snapshot, now)
require.True(t, limited)
require.WithinDuration(t, now.Add(45*time.Second), resetAt, time.Second)
authoritativeReset := now.Add(2 * time.Hour)
remaining := int64(0)
snapshot.Requests = &xai.QuotaWindow{Remaining: &remaining, ResetUnix: grokInt64PtrForTest(authoritativeReset.Unix())}
recentReset := now.Add(-time.Second)
recentLimited := recentReset.Add(-grokRateLimitSustainedCooldown)
account.RateLimitResetAt = &recentReset
account.RateLimitedAt = &recentLimited
resetAt, limited = grokRateLimitResetAtForAccount(account, snapshot, now)
require.True(t, limited)
require.WithinDuration(t, authoritativeReset, resetAt, time.Second)
}
func TestGrokRateLimitResetAtForAccountLeavesAPIKey429PolicyUnchanged(t *testing.T) {
now := time.Now().UTC().Truncate(time.Second)
retryAfter := 45
previousReset := now.Add(-time.Second)
previousLimited := previousReset.Add(-grokRateLimitSustainedCooldown)
account := &Account{
ID: 632,
Platform: PlatformGrok,
Type: AccountTypeAPIKey,
RateLimitedAt: &previousLimited,
RateLimitResetAt: &previousReset,
}
snapshot := &xai.QuotaSnapshot{
StatusCode: http.StatusTooManyRequests,
RetryAfterSeconds: &retryAfter,
UpdatedAt: now.Format(time.RFC3339),
}
resetAt, limited := grokRateLimitResetAtForAccount(account, snapshot, now)
require.True(t, limited)
require.WithinDuration(t, now.Add(45*time.Second), resetAt, time.Second)
}
func TestGrokRateLimitResetAtUsesFutureWindowAfterRetryAfterExpires(t *testing.T) {
now := time.Now().UTC().Truncate(time.Second)
observedAt := now.Add(-2 * time.Minute)
windowReset := now.Add(15 * time.Minute)
retryAfter := 30
snapshot := &xai.QuotaSnapshot{
StatusCode: http.StatusTooManyRequests,
UpdatedAt: observedAt.Format(time.RFC3339),
RetryAfterSeconds: &retryAfter,
Requests: &xai.QuotaWindow{
Limit: grokInt64PtrForTest(10),
Remaining: grokInt64PtrForTest(0),
ResetUnix: grokInt64PtrForTest(windowReset.Unix()),
},
}
resetAt, limited := grokRateLimitResetAt(snapshot, now)
require.True(t, limited)
require.WithinDuration(t, windowReset, resetAt, time.Second)
}
func TestHandleGrokAccountUpstreamError429DoesNotShortenExistingPause(t *testing.T) {
existingUntil := time.Now().Add(15 * time.Minute)
account := &Account{
ID: 64,
Platform: PlatformGrok,
Type: AccountTypeOAuth,
TempUnschedulableUntil: &existingUntil,
TempUnschedulableReason: "existing pause",
}
repo := &grokQuotaAccountRepo{}
svc := &OpenAIGatewayService{accountRepo: repo}
svc.handleGrokAccountUpstreamError(context.Background(), account, http.StatusTooManyRequests, http.Header{"Retry-After": []string{"45"}}, nil)
require.Equal(t, 1, repo.rateLimitedCalls)
require.WithinDuration(t, time.Now().Add(45*time.Second), repo.lastRateLimitResetAt, time.Second)
require.Zero(t, repo.tempUnschedCalls)
value, ok := svc.openaiAccountRuntimeBlockUntil.Load(account.ID)
require.True(t, ok)
runtimeUntil, ok := value.(time.Time)
require.True(t, ok)
require.WithinDuration(t, existingUntil, runtimeUntil, time.Second)
}
func TestUpdateGrokUsageSnapshotExhaustedSuccessBypassesThrottleAndSetsRateLimited(t *testing.T) {
account := &Account{ID: 65, Platform: PlatformGrok, Type: AccountTypeOAuth}
repo := &grokQuotaAccountRepo{}
svc := &OpenAIGatewayService{
accountRepo: repo,
codexSnapshotThrottle: newAccountWriteThrottle(time.Hour),
}
now := time.Now()
// Consume the normal snapshot write allowance first.
svc.updateGrokUsageSnapshot(context.Background(), account, &xai.QuotaSnapshot{
StatusCode: http.StatusOK,
Requests: &xai.QuotaWindow{
Limit: grokInt64PtrForTest(10),
Remaining: grokInt64PtrForTest(9),
},
UpdatedAt: now.UTC().Format(time.RFC3339),
})
resetAt := now.Add(30 * time.Minute).Truncate(time.Second)
svc.updateGrokUsageSnapshot(context.Background(), account, &xai.QuotaSnapshot{
StatusCode: http.StatusOK,
Requests: &xai.QuotaWindow{
Limit: grokInt64PtrForTest(10),
Remaining: grokInt64PtrForTest(0),
ResetUnix: grokInt64PtrForTest(resetAt.Unix()),
ResetAt: resetAt.UTC().Format(time.RFC3339),
},
UpdatedAt: now.UTC().Format(time.RFC3339),
})
require.Equal(t, 2, repo.updateCalls)
require.Equal(t, 1, repo.rateLimitedCalls)
require.Equal(t, account.ID, repo.lastRateLimitedID)
require.WithinDuration(t, resetAt, repo.lastRateLimitResetAt, time.Second)
require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
}
func TestUpdateGrokUsageSnapshotAvailableSuccessDoesNotSetRateLimited(t *testing.T) {
repo := &grokQuotaAccountRepo{}
svc := &OpenAIGatewayService{accountRepo: repo}
account := &Account{ID: 66, Platform: PlatformGrok, Type: AccountTypeOAuth}
svc.updateGrokUsageSnapshot(context.Background(), account, &xai.QuotaSnapshot{
StatusCode: http.StatusOK,
Requests: &xai.QuotaWindow{
Limit: grokInt64PtrForTest(10),
Remaining: grokInt64PtrForTest(1),
},
UpdatedAt: time.Now().UTC().Format(time.RFC3339),
})
require.Equal(t, 1, repo.updateCalls)
require.Zero(t, repo.rateLimitedCalls)
}
func TestUpdateGrokUsageFromResponseHeaderlessSuccessClearsObservedCooldown(t *testing.T) {
now := time.Now().UTC().Truncate(time.Second)
limitedAt := now.Add(-grokRateLimitRepeatCooldown)
observedResetAt := now.Add(-time.Second)
account := &Account{
ID: 660,
Platform: PlatformGrok,
Type: AccountTypeOAuth,
RateLimitedAt: &limitedAt,
RateLimitResetAt: &observedResetAt,
}
repo := &grokQuotaAccountRepo{recoveryClearResult: true}
svc := &OpenAIGatewayService{
accountRepo: repo,
codexSnapshotThrottle: newAccountWriteThrottle(time.Hour),
}
svc.updateGrokUsageFromResponse(context.Background(), account, nil, http.StatusOK)
require.Zero(t, repo.updateCalls, "headerless success must not overwrite an informative quota snapshot")
require.Equal(t, 1, repo.recoveryClearCalls)
require.Equal(t, limitedAt, repo.recoveryObservedAt)
require.Equal(t, observedResetAt, repo.recoveryObservedReset)
require.Same(t, &observedResetAt, account.RateLimitResetAt, "shared account snapshots must not be mutated in place")
}
func TestUpdateGrokUsageFromResponseRecoveryRespectsCancellationAndAPIKeyBoundary(t *testing.T) {
now := time.Now().UTC().Truncate(time.Second)
observedResetAt := now.Add(-time.Second)
observedLimitedAt := observedResetAt.Add(-grokRateLimitRepeatCooldown)
t.Run("parent cancellation does not mutate account state", func(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
cancel()
account := &Account{
ID: 661,
Platform: PlatformGrok,
Type: AccountTypeOAuth,
RateLimitedAt: &observedLimitedAt,
RateLimitResetAt: &observedResetAt,
}
repo := &grokQuotaAccountRepo{recoveryClearResult: true}
svc := &OpenAIGatewayService{accountRepo: repo}
svc.updateGrokUsageFromResponse(ctx, account, nil, http.StatusOK)
require.Zero(t, repo.recoveryClearCalls)
})
t.Run("API key success does not alter OAuth cooldown state", func(t *testing.T) {
account := &Account{
ID: 662,
Platform: PlatformGrok,
Type: AccountTypeAPIKey,
RateLimitedAt: &observedLimitedAt,
RateLimitResetAt: &observedResetAt,
}
repo := &grokQuotaAccountRepo{recoveryClearResult: true}
svc := &OpenAIGatewayService{accountRepo: repo}
svc.updateGrokUsageFromResponse(context.Background(), account, nil, http.StatusOK)
require.Zero(t, repo.recoveryClearCalls)
})
}
func TestUpdateGrokUsageSnapshotExhaustedSuccessWithoutResetUsesFallback(t *testing.T) {
repo := &grokQuotaAccountRepo{}
svc := &OpenAIGatewayService{accountRepo: repo}
account := &Account{ID: 67, Platform: PlatformGrok, Type: AccountTypeOAuth}
before := time.Now()
svc.updateGrokUsageSnapshot(context.Background(), account, &xai.QuotaSnapshot{
StatusCode: http.StatusOK,
Tokens: &xai.QuotaWindow{
Limit: grokInt64PtrForTest(2_000_000),
Remaining: grokInt64PtrForTest(0),
},
UpdatedAt: before.UTC().Format(time.RFC3339),
})
require.Equal(t, 1, repo.rateLimitedCalls)
require.WithinDuration(t, before.Add(grokRateLimitFallbackCooldown), repo.lastRateLimitResetAt, time.Second)
stored, ok := repo.updates[account.ID][grokQuotaSnapshotExtraKey].(*xai.QuotaSnapshot)
require.True(t, ok)
require.NotNil(t, stored.Tokens.ResetUnix)
paused, _ := shouldAutoPauseGrokQuotaWindow("tokens", stored.Tokens, before.Add(time.Second))
require.True(t, paused)
paused, _ = shouldAutoPauseGrokQuotaWindow("tokens", stored.Tokens, repo.lastRateLimitResetAt.Add(time.Second))
require.False(t, paused)
}
func TestOpenAIWSHTTPBridgeGrok429PersistsRateLimit(t *testing.T) {
repo := &grokQuotaAccountRepo{}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusTooManyRequests,
Header: http.Header{"Retry-After": []string{"45"}},
Body: io.NopCloser(strings.NewReader(`{"error":{"message":"rate limited"}}`)),
}}
svc := &OpenAIGatewayService{accountRepo: repo, httpUpstream: upstream}
account := &Account{ID: 68, Platform: PlatformGrok, Type: AccountTypeOAuth, Concurrency: 1}
before := time.Now()
result, err := svc.proxyOpenAIWSHTTPBridgeTurn(
context.Background(), nil, account, "token",
[]byte(`{"type":"response.create","model":"grok-4.3","input":"hi"}`),
64, "grok-4.3", "", "", "", "cache-id", 1,
func([]byte) error { return nil },
)
require.Error(t, err)
require.Nil(t, result)
require.Equal(t, 1, repo.rateLimitedCalls)
require.WithinDuration(t, before.Add(45*time.Second), repo.lastRateLimitResetAt, time.Second)
require.Zero(t, repo.tempUnschedCalls)
require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
}
func TestOpenAIWSHTTPBridgeSSEErrorSideEffectsRunOncePerPlatform(t *testing.T) {
gin.SetMode(gin.TestMode)
for _, platform := range []string{PlatformOpenAI, PlatformGrok} {
t.Run(platform, func(t *testing.T) {
repo := &grokQuotaAccountRepo{}
cfg := &config.Config{}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: make(http.Header),
Body: io.NopCloser(strings.NewReader(
"data: {\"type\":\"error\",\"error\":{\"type\":\"rate_limit_error\",\"code\":\"rate_limit_exceeded\",\"message\":\"limited\"}}\n\n",
)),
}}
svc := &OpenAIGatewayService{
cfg: cfg,
accountRepo: repo,
httpUpstream: upstream,
}
if platform == PlatformOpenAI {
svc.rateLimitService = NewRateLimitService(repo, nil, cfg, nil, nil)
}
account := &Account{ID: 70, Platform: platform, Type: AccountTypeOAuth, Concurrency: 1}
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
payload := []byte(`{"type":"response.create","model":"gpt-5","input":"hi"}`)
writes := 0
result, err := svc.proxyOpenAIWSHTTPBridgeTurn(
context.Background(), c, account, "sk-test", payload, len(payload),
"gpt-5", "", "", "", "", 1,
func([]byte) error {
writes++
return nil
},
)
require.Nil(t, result)
var failoverErr *UpstreamFailoverError
require.ErrorAs(t, err, &failoverErr)
require.Equal(t, http.StatusTooManyRequests, failoverErr.StatusCode)
require.Zero(t, writes)
require.Equal(t, 1, repo.rateLimitedCalls)
})
}
}
func TestOpenAIWSHTTPBridgeGrokExhaustedSuccessPersistsRateLimit(t *testing.T) {
repo := &grokQuotaAccountRepo{}
resetAt := time.Now().Add(20 * time.Minute).UTC().Truncate(time.Second)
resp := grokMessagesSSECompletedResponse("resp_ws_limited", 0)
resp.Header.Set("X-Ratelimit-Limit-Requests", "10")
resp.Header.Set("X-Ratelimit-Remaining-Requests", "0")
resp.Header.Set("X-Ratelimit-Reset-Requests", fmt.Sprintf("%d", resetAt.Unix()))
upstream := &httpUpstreamRecorder{resp: resp}
svc := &OpenAIGatewayService{accountRepo: repo, httpUpstream: upstream}
account := &Account{ID: 69, Platform: PlatformGrok, Type: AccountTypeOAuth, Concurrency: 1}
result, err := svc.proxyOpenAIWSHTTPBridgeTurn(
context.Background(), nil, account, "token",
[]byte(`{"type":"response.create","model":"grok-4.3","input":"hi"}`),
64, "grok-4.3", "", "", "", "cache-id", 1,
func([]byte) error { return nil },
)
require.NoError(t, err)
require.NotNil(t, result)
require.Equal(t, 1, repo.rateLimitedCalls)
require.WithinDuration(t, resetAt, repo.lastRateLimitResetAt, time.Second)
require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
}
func TestFailoverOpenAIUpstreamHTTPErrorUsesOnlyGrokRateLimitPolicy(t *testing.T) {
gin.SetMode(gin.TestMode)
repo := &grokQuotaAccountRepo{}
svc := &OpenAIGatewayService{accountRepo: repo}
account := &Account{ID: 70, Platform: PlatformGrok, Type: AccountTypeOAuth}
resp := &http.Response{
StatusCode: http.StatusTooManyRequests,
Header: http.Header{"Retry-After": []string{"45"}},
}
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
failoverErr := svc.failoverOpenAIUpstreamHTTPError(
context.Background(), c, account, resp,
[]byte(`{"error":{"message":"rate limited"}}`), "rate limited", "grok-4.3",
)
require.NotNil(t, failoverErr)
require.Equal(t, 1, repo.rateLimitedCalls)
require.Zero(t, repo.tempUnschedCalls)
}
func TestPatchGrokResponsesBody_StripsReasoningContentNull(t *testing.T) {
t.Parallel()
body := []byte(`{
"model": "grok-latest",
"input": [
{"type":"message","role":"user","content":[{"type":"input_text","text":"hi"}]},
{"type":"reasoning","summary":[{"type":"summary_text","text":"thinking..."}],"content":null,"encrypted_content":null},
{"type":"message","role":"assistant","content":[{"type":"output_text","text":"Hello!"}]}
]
}`)
patched, err := patchGrokResponsesBody(body, "grok-4.5")
require.NoError(t, err)
require.True(t, json.Valid(patched))
input := gjson.GetBytes(patched, "input")
require.True(t, input.IsArray())
items := input.Array()
require.Len(t, items, 3)
reasoning := items[1]
require.Equal(t, "reasoning", reasoning.Get("type").String())
require.True(t, reasoning.Get("summary").Exists(), "summary should be preserved")
require.False(t, reasoning.Get("content").Exists(), "content: null should be stripped")
}
func TestPatchGrokResponsesBody_KeepsReasoningContentNonNull(t *testing.T) {
t.Parallel()
body := []byte(`{
"model": "grok-latest",
"input": [
{"type":"reasoning","summary":[{"type":"summary_text","text":"ok"}],"content":"real content"}
]
}`)
patched, err := patchGrokResponsesBody(body, "grok-4.5")
require.NoError(t, err)
reasoning := gjson.GetBytes(patched, "input.0")
require.Equal(t, "real content", reasoning.Get("content").String(), "non-null content must not be stripped")
}
func TestPatchGrokResponsesBody_MultipleReasoningContentNull(t *testing.T) {
t.Parallel()
body := []byte(`{
"model": "grok-latest",
"input": [
{"type":"reasoning","summary":[{"type":"summary_text","text":"r1"}],"content":null},
{"type":"message","role":"user","content":"hi"},
{"type":"reasoning","summary":[{"type":"summary_text","text":"r2"}],"content":null}
]
}`)
patched, err := patchGrokResponsesBody(body, "grok-4.5")
require.NoError(t, err)
items := gjson.GetBytes(patched, "input").Array()
require.Len(t, items, 3)
require.False(t, items[0].Get("content").Exists())
require.False(t, items[2].Get("content").Exists())
}
func TestIsGrokImageGenerationModel(t *testing.T) {
t.Parallel()
tests := []struct {
model string
want bool
}{
{"grok-imagine", true},
{"grok-imagine-image-quality", true},
{"grok-imagine-edit", true},
{"grok-imagine-image-hd", true},
{"grok-4.5", false},
{"grok-composer", false},
}
for _, tt := range tests {
t.Run(tt.model, func(t *testing.T) {
require.Equal(t, tt.want, isGrokImageGenerationModel(tt.model))
})
}
}
func TestBuildGrokSchedulerExtraUpdates_FeedsThresholdEvaluator(t *testing.T) {
int64p := func(v int64) *int64 { return &v }
resetUnix := time.Now().Add(90 * time.Minute).Unix()
snapshot := &xai.QuotaSnapshot{
Requests: &xai.QuotaWindow{Limit: int64p(100), Remaining: int64p(30)}, // 70% used
Tokens: &xai.QuotaWindow{Limit: int64p(1000), Remaining: int64p(50), ResetUnix: &resetUnix}, // 95% used (most constrained)
}
updates := buildGrokSchedulerExtraUpdates(snapshot)
require.NotNil(t, updates)
require.InDelta(t, 95.0, updates["grok_sched_utilization"], 0.001, "picks the most-constrained window")
require.Contains(t, updates, "grok_sched_reset_at")
// The written extras must actually drive EvaluateAccountSchedulingThreshold
// (proves the previously-dead read side is now fed).
account := &Account{Platform: PlatformGrok, Extra: updates}
decision := EvaluateAccountSchedulingThreshold(account, map[string]int{PlatformGrok: 90}, time.Now())
require.True(t, decision.ShouldPause)
require.InDelta(t, 95.0, decision.UsedPercent, 0.001)
require.NotNil(t, decision.Until)
}
func TestBuildGrokSchedulerExtraUpdates_NilWhenNoQuotaWindows(t *testing.T) {
require.Nil(t, buildGrokSchedulerExtraUpdates(&xai.QuotaSnapshot{}))
require.Nil(t, buildGrokSchedulerExtraUpdates(nil))
}