| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167 |
- //go:build unit
- package service
- import (
- "fmt"
- "net/http"
- "net/http/httptest"
- "sync"
- "testing"
- "github.com/Wei-Shaw/sub2api/internal/pkg/xai"
- "github.com/gin-gonic/gin"
- "github.com/stretchr/testify/require"
- "github.com/tidwall/gjson"
- )
- func newGrokCacheTestContext(apiKeyID int64) *gin.Context {
- recorder := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(recorder)
- c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
- if apiKeyID > 0 {
- c.Set("api_key", &APIKey{ID: apiKeyID, Group: &Group{Platform: PlatformGrok}})
- }
- return c
- }
- func TestGrokPreviousResponseSessionSeed(t *testing.T) {
- require.Equal(t, "grok-prev-resp:resp_abc123", grokPreviousResponseSessionSeed([]byte(`{"previous_response_id":"resp_abc123"}`)))
- require.Empty(t, grokPreviousResponseSessionSeed([]byte(`{"previous_response_id":"msg_abc123"}`)))
- require.Empty(t, grokPreviousResponseSessionSeed([]byte(`{"previous_response_id":""}`)))
- require.Empty(t, grokPreviousResponseSessionSeed([]byte(`{}`)))
- }
- func TestResolveGrokCacheIdentityUsesPreviousResponseIDWhenNoOtherSeed(t *testing.T) {
- gin.SetMode(gin.TestMode)
- c := newGrokCacheTestContext(301)
- // No prompt_cache_key / headers / reusable prefix — only previous_response_id.
- body := []byte(`{"model":"grok","input":[{"role":"user","content":"follow up"}],"previous_response_id":"resp_chain_001"}`)
- got := resolveGrokCacheIdentity(c, body, "", "grok-4.5")
- require.NotEmpty(t, got)
- // Same previous_response_id → same identity (model already in isolated seed).
- again := resolveGrokCacheIdentity(c, body, "", "grok-4.5")
- require.Equal(t, got, again)
- // Different model → different identity (model scope).
- otherModel := resolveGrokCacheIdentity(c, body, "", "grok-4.3")
- require.NotEqual(t, got, otherModel)
- // prompt_cache_key still wins over previous_response_id.
- withCache := []byte(`{"model":"grok","prompt_cache_key":"client-session","previous_response_id":"resp_chain_001","input":[{"role":"user","content":"x"}]}`)
- cacheID := resolveGrokCacheIdentity(c, withCache, "", "grok-4.5")
- require.NotEmpty(t, cacheID)
- require.NotEqual(t, got, cacheID)
- }
- func TestResolveGrokCacheIdentityStableAcrossAppendOnlyTurns(t *testing.T) {
- gin.SetMode(gin.TestMode)
- c := newGrokCacheTestContext(101)
- round1 := []byte(`{"model":"grok","instructions":"be concise","tools":[{"type":"function","name":"lookup","parameters":{"type":"object"}}],"input":[{"role":"user","content":"first question"}]}`)
- round2 := []byte(`{"model":"grok","instructions":"be concise","tools":[{"type":"function","name":"lookup","parameters":{"type":"object"}}],"input":[{"role":"user","content":"first question"},{"role":"assistant","content":"first answer"},{"role":"user","content":"second question"}]}`)
- first := resolveGrokCacheIdentity(c, round1, "", "grok-4.5")
- second := resolveGrokCacheIdentity(c, round2, "", "grok-4.5")
- require.NotEmpty(t, first)
- require.Len(t, first, 36)
- require.Equal(t, first, second)
- }
- func TestResolveGrokCacheIdentityStableAcrossIndependentPromptsWithSamePrefix(t *testing.T) {
- gin.SetMode(gin.TestMode)
- c := newGrokCacheTestContext(102)
- firstBody := []byte(`{"model":"grok","instructions":"be concise","tools":[{"type":"function","name":"lookup"}],"input":[{"role":"user","content":"Question A"}]}`)
- secondBody := []byte(`{"model":"grok","instructions":"be concise","tools":[{"type":"function","name":"lookup"}],"input":[{"role":"user","content":"Question B"}]}`)
- first := resolveGrokCacheIdentity(c, firstBody, "", "grok-4.5")
- second := resolveGrokCacheIdentity(c, secondBody, "", "grok-4.5")
- require.NotEmpty(t, first)
- require.Equal(t, first, second)
- }
- func TestResolveGrokCacheIdentityStablePrefixIsolation(t *testing.T) {
- gin.SetMode(gin.TestMode)
- baseBody := []byte(`{"model":"grok","instructions":"be concise","tools":[{"type":"function","name":"lookup"}],"input":[{"role":"system","content":"System A"},{"role":"user","content":"Question A"}]}`)
- differentInstructions := []byte(`{"model":"grok","instructions":"be detailed","tools":[{"type":"function","name":"lookup"}],"input":[{"role":"system","content":"System A"},{"role":"user","content":"Question B"}]}`)
- differentSystem := []byte(`{"model":"grok","instructions":"be concise","tools":[{"type":"function","name":"lookup"}],"input":[{"role":"system","content":"System B"},{"role":"user","content":"Question B"}]}`)
- differentTools := []byte(`{"model":"grok","instructions":"be concise","tools":[{"type":"function","name":"search"}],"input":[{"role":"system","content":"System A"},{"role":"user","content":"Question B"}]}`)
- base := resolveGrokCacheIdentity(newGrokCacheTestContext(103), baseBody, "", "grok-4.5")
- require.NotEqual(t, base, resolveGrokCacheIdentity(newGrokCacheTestContext(104), baseBody, "", "grok-4.5"))
- require.NotEqual(t, base, resolveGrokCacheIdentity(newGrokCacheTestContext(103), baseBody, "", "grok-4.3"))
- require.NotEqual(t, base, resolveGrokCacheIdentity(newGrokCacheTestContext(103), differentInstructions, "", "grok-4.5"))
- require.NotEqual(t, base, resolveGrokCacheIdentity(newGrokCacheTestContext(103), differentSystem, "", "grok-4.5"))
- require.NotEqual(t, base, resolveGrokCacheIdentity(newGrokCacheTestContext(103), differentTools, "", "grok-4.5"))
- }
- func TestResolveGrokCacheIdentityFallsBackWhenStablePrefixIsEmpty(t *testing.T) {
- gin.SetMode(gin.TestMode)
- c := newGrokCacheTestContext(105)
- firstBody := []byte(`{"model":"grok","tools":[],"input":"Question A"}`)
- secondBody := []byte(`{"model":"grok","tools":[],"input":"Question B"}`)
- first := resolveGrokCacheIdentity(c, firstBody, "", "grok-4.5")
- second := resolveGrokCacheIdentity(c, secondBody, "", "grok-4.5")
- require.NotEmpty(t, first)
- require.NotEmpty(t, second)
- require.NotEqual(t, first, second)
- }
- func TestResolveGrokCacheIdentitySkipsUnanchoredFallback(t *testing.T) {
- gin.SetMode(gin.TestMode)
- c := newGrokCacheTestContext(106)
- tests := [][]byte{
- []byte(`{"model":"grok"}`),
- []byte(`{"model":"grok","messages":[{"role":"assistant","content":"answer"}]}`),
- []byte(`{"model":"grok","messages":[{"role":"user","content":""}]}`),
- []byte(`{"model":"grok","input":" "}`),
- }
- for _, body := range tests {
- require.Empty(t, resolveGrokCacheIdentity(c, body, "", "grok-4.5"))
- }
- }
- func TestResolveGrokCacheIdentityIsolatesAPIKeyAndMappedModel(t *testing.T) {
- gin.SetMode(gin.TestMode)
- body := []byte(`{"model":"grok","input":"same prompt"}`)
- base := resolveGrokCacheIdentity(newGrokCacheTestContext(201), body, "", "grok-4.5")
- otherTenant := resolveGrokCacheIdentity(newGrokCacheTestContext(202), body, "", "grok-4.5")
- otherModel := resolveGrokCacheIdentity(newGrokCacheTestContext(201), body, "", "grok-4.3")
- require.NotEmpty(t, base)
- require.NotEqual(t, base, otherTenant)
- require.NotEqual(t, base, otherModel)
- }
- func TestResolveGrokCacheIdentityUsesAndIsolatesNativeConversationHeader(t *testing.T) {
- gin.SetMode(gin.TestMode)
- c := newGrokCacheTestContext(301)
- c.Request.Header.Set(grokConversationIDHeader, "raw-native-conversation")
- body1 := []byte(`{"model":"grok","input":"one"}`)
- body2 := []byte(`{"model":"grok","input":"different body that must not replace the explicit session"}`)
- first := resolveGrokCacheIdentity(c, body1, "body-cache-key", "grok-4.5")
- second := resolveGrokCacheIdentity(c, body2, "another-body-cache-key", "grok-4.5")
- require.Equal(t, "raw-native-conversation", (&OpenAIGatewayService{}).ExtractSessionID(c, body1))
- require.Equal(t, first, second)
- require.NotEqual(t, "raw-native-conversation", first)
- require.NotContains(t, first, "raw-native-conversation")
- }
- func TestResolveGrokCacheIdentityExplicitHeaderPriority(t *testing.T) {
- gin.SetMode(gin.TestMode)
- body := []byte(`{"model":"grok","prompt_cache_key":"body-key","input":"hi"}`)
- c := newGrokCacheTestContext(401)
- c.Request.Header.Set(grokConversationIDHeader, "grok-key")
- c.Request.Header.Set("conversation_id", "conversation-key")
- c.Request.Header.Set("session_id", "session-key")
- got := resolveGrokCacheIdentity(c, body, "explicit-argument", "grok-4.5")
- onlySession := newGrokCacheTestContext(401)
- onlySession.Request.Header.Set("session_id", "session-key")
- want := resolveGrokCacheIdentity(onlySession, []byte(`{"model":"grok","input":"unrelated"}`), "", "grok-4.5")
- require.Equal(t, want, got)
- }
- func TestResolveGrokCacheIdentityIDEHeaderPriority(t *testing.T) {
- gin.SetMode(gin.TestMode)
- body := []byte(`{"model":"grok","prompt_cache_key":"body-key","input":"hi"}`)
- headers := []struct {
- name string
- value string
- }{
- {name: openCodeSessionAffinityHeader, value: "opencode-affinity"},
- {name: openCodeSessionIDHeader, value: "opencode-session-id"},
- {name: openCodeNativeSessionHeader, value: "opencode-native-session"},
- {name: codeBuddyConversationHeader, value: "codebuddy-conversation"},
- {name: grokConversationIDHeader, value: "grok-conversation"},
- }
- c := newGrokCacheTestContext(402)
- for _, header := range headers {
- c.Request.Header.Set(header.name, header.value)
- }
- for _, header := range headers {
- got := resolveGrokCacheIdentity(c, body, "explicit-argument", "grok-4.5")
- onlyCurrent := newGrokCacheTestContext(402)
- onlyCurrent.Request.Header.Set(header.name, header.value)
- if header.name == grokConversationIDHeader {
- // prompt_cache_key beats X-Grok-Conv-Id: verify the identity is
- // derived from the body key, not the per-call conv label.
- withOnlyBody := newGrokCacheTestContext(402)
- want := resolveGrokCacheIdentity(withOnlyBody, body, "", "grok-4.5")
- require.Equal(t, want, got, header.name)
- withOnlyConv := newGrokCacheTestContext(402)
- withOnlyConv.Request.Header.Set(header.name, header.value)
- convOnly := resolveGrokCacheIdentity(withOnlyConv, []byte(`{"model":"grok","input":"unrelated"}`), "", "grok-4.5")
- require.NotEqual(t, convOnly, got, "body key must not be shadowed by X-Grok-Conv-Id")
- } else {
- want := resolveGrokCacheIdentity(onlyCurrent, body, "", "grok-4.5")
- require.Equal(t, want, got, header.name)
- }
- c.Request.Header.Del(header.name)
- }
- }
- // TestResolveGrokCacheIdentitySideCallSharesParentCacheKey locks in the
- // grok-build side-call fix: recap-style side-calls (turn-summary /
- // title-refresh) send a fresh X-Grok-Conv-Id label but the parent session id
- // as body prompt_cache_key. The derived identity must follow the body key so
- // side-calls share the main turn's server-side cache prefix instead of
- // replaying the full conversation at full price on every call.
- func TestResolveGrokCacheIdentitySideCallSharesParentCacheKey(t *testing.T) {
- gin.SetMode(gin.TestMode)
- parentSession := "6f1c2f46-0f5e-4f9d-9d4e-2f0f1c3d5b7a"
- mainTurn := newGrokCacheTestContext(910)
- mainTurn.Request.Header.Set(grokConversationIDHeader, parentSession)
- mainIdentity := resolveGrokCacheIdentity(mainTurn, []byte(`{"model":"grok-4.6","input":"main conversation"}`), "", "grok-4.6")
- require.NotEmpty(t, mainIdentity)
- sideCall := newGrokCacheTestContext(910)
- sideCall.Request.Header.Set(grokConversationIDHeader, "turn-summary-"+uuidNew())
- sideIdentity := resolveGrokCacheIdentity(sideCall, []byte(`{"model":"grok-4.6","prompt_cache_key":"`+parentSession+`","input":"summary replay"}`), "", "grok-4.6")
- require.Equal(t, mainIdentity, sideIdentity,
- "side-call with parent prompt_cache_key must share the main turn cache identity")
- titleRefresh := newGrokCacheTestContext(910)
- titleRefresh.Request.Header.Set(grokConversationIDHeader, "title-refresh-"+uuidNew())
- titleIdentity := resolveGrokCacheIdentity(titleRefresh, []byte(`{"model":"grok-4.6","prompt_cache_key":"`+parentSession+`","input":"title replay"}`), "", "grok-4.6")
- require.Equal(t, mainIdentity, titleIdentity)
- }
- func uuidNew() string {
- return fmt.Sprintf("%08x-%04x-4%03x-9%03x-%012x", 0x1234, 0x5678, 0x9abc, 0xdef0, 0x1234567890ab)
- }
- func TestExplicitGrokCacheSeedPriority(t *testing.T) {
- gin.SetMode(gin.TestMode)
- c := newGrokCacheTestContext(403)
- headers := []struct {
- name string
- value string
- }{
- {name: claudeCodeSessionHeader, value: "claude-session"},
- {name: "session_id", value: "generic-session"},
- {name: "conversation_id", value: "generic-conversation"},
- {name: openCodeSessionAffinityHeader, value: "opencode-affinity"},
- {name: openCodeSessionIDHeader, value: "opencode-session-id"},
- {name: openCodeNativeSessionHeader, value: "opencode-native-session"},
- {name: codeBuddyConversationHeader, value: "codebuddy-conversation"},
- {name: grokConversationIDHeader, value: "grok-conversation"},
- }
- for _, header := range headers {
- c.Request.Header.Set(header.name, header.value)
- }
- body := []byte(`{"model":"grok","prompt_cache_key":"body-key","input":"hi"}`)
- for _, header := range headers {
- var want string
- switch header.name {
- case grokConversationIDHeader:
- // Client-declared prompt_cache_key outranks X-Grok-Conv-Id: the
- // grok-build CLI pairs a fresh per-side-call conv label with the
- // stable parent session id in the body field, and the body field
- // is the official xAI cache-routing signal.
- want = "body-key"
- default:
- want = header.value
- }
- require.Equal(t, want, explicitGrokCacheSeed(c, body, "explicit-argument"), header.name)
- c.Request.Header.Del(header.name)
- }
- require.Equal(t, "body-key", explicitGrokCacheSeed(c, body, "explicit-argument"))
- require.Equal(t, "explicit-argument", explicitGrokCacheSeed(c, []byte(`{"model":"grok"}`), "explicit-argument"))
- }
- func TestResolveGrokCacheIdentityIDEHeadersAreStableIsolatedAndOpaque(t *testing.T) {
- gin.SetMode(gin.TestMode)
- tests := []struct {
- name string
- header string
- }{
- {name: "OpenCode affinity", header: openCodeSessionAffinityHeader},
- {name: "OpenCode session ID", header: openCodeSessionIDHeader},
- {name: "OpenCode native session", header: openCodeNativeSessionHeader},
- {name: "CodeBuddy conversation", header: codeBuddyConversationHeader},
- }
- for index, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- rawSession := "raw-ide-session-" + tt.name
- apiKeyID := int64(800 + index)
- c := newGrokCacheTestContext(apiKeyID)
- c.Request.Header.Set(tt.header, rawSession)
- firstBody := []byte(`{"model":"grok","prompt_cache_key":"turn-one-body-key","input":"first turn"}`)
- secondBody := []byte(`{"model":"grok","prompt_cache_key":"turn-two-body-key","input":"different second turn"}`)
- first := resolveGrokCacheIdentity(c, firstBody, "first-explicit-key", "grok-4.5")
- second := resolveGrokCacheIdentity(c, secondBody, "second-explicit-key", "grok-4.5")
- require.NotEmpty(t, first)
- require.Equal(t, first, second)
- require.NotEqual(t, rawSession, first)
- require.NotContains(t, first, rawSession)
- otherTenant := newGrokCacheTestContext(apiKeyID + 100)
- otherTenant.Request.Header.Set(tt.header, rawSession)
- require.NotEqual(t, first, resolveGrokCacheIdentity(otherTenant, firstBody, "", "grok-4.5"))
- require.NotEqual(t, first, resolveGrokCacheIdentity(c, firstBody, "", "grok-4.3"))
- })
- }
- }
- func TestOpenCodeResponsesHeaderAndBodyCacheSignalsConverge(t *testing.T) {
- gin.SetMode(gin.TestMode)
- const rawSession = "opencode-session-42"
- c := newGrokCacheTestContext(901)
- c.Request.Header.Set(openCodeSessionAffinityHeader, rawSession)
- c.Request.Header.Set(openCodeSessionIDHeader, rawSession)
- firstBody := []byte(`{"model":"grok","prompt_cache_key":"opencode-session-42","input":"first turn"}`)
- secondBody := []byte(`{"model":"grok","prompt_cache_key":"opencode-session-42","input":"different second turn"}`)
- first := resolveGrokCacheIdentity(c, firstBody, "", "grok-4.5")
- second := resolveGrokCacheIdentity(c, secondBody, "", "grok-4.5")
- bodyOnly := resolveGrokCacheIdentity(newGrokCacheTestContext(901), secondBody, "", "grok-4.5")
- require.NotEmpty(t, first)
- require.Equal(t, first, second)
- require.Equal(t, first, bodyOnly)
- patched, err := applyGrokResponsesCacheIdentity(secondBody, secondBody, second, false)
- require.NoError(t, err)
- require.Equal(t, second, gjson.GetBytes(patched, "prompt_cache_key").String())
- headers := make(http.Header)
- applyGrokCacheHeaders(headers, second)
- require.Equal(t, second, headers.Get(grokConversationIDHeader))
- require.NotContains(t, string(patched), rawSession)
- }
- func TestResolveGrokCacheIdentityPrefersClaudeCodeSession(t *testing.T) {
- gin.SetMode(gin.TestMode)
- c := newGrokCacheTestContext(701)
- c.Request.Header.Set(claudeCodeSessionHeader, "cc-session-abc")
- c.Request.Header.Set("session_id", "session-key")
- body1 := []byte(`{"model":"grok","input":"turn-1"}`)
- body2 := []byte(`{"model":"grok","input":"turn-2-different-user-text"}`)
- first := resolveGrokCacheIdentity(c, body1, "", "grok-4.5")
- second := resolveGrokCacheIdentity(c, body2, "unrelated-explicit", "grok-4.5")
- require.NotEmpty(t, first)
- require.Equal(t, first, second, "same Claude Code session must keep stable Grok cache identity across turns")
- // metadata.user_id JSON form used by Claude Code clients
- metaBody := []byte(`{"model":"grok","metadata":{"user_id":"{\"session_id\":\"meta-session-xyz\"}"},"input":"hi"}`)
- metaOnly := newGrokCacheTestContext(702)
- metaID := resolveGrokCacheIdentity(metaOnly, metaBody, "", "grok-4.5")
- require.NotEmpty(t, metaID)
- metaBody2 := []byte(`{"model":"grok","metadata":{"user_id":"{\"session_id\":\"meta-session-xyz\"}"},"input":"later turn"}`)
- require.Equal(t, metaID, resolveGrokCacheIdentity(metaOnly, metaBody2, "", "grok-4.5"))
- }
- func TestResolveGrokCacheIdentityFailsClosedWithoutAPIKeyContext(t *testing.T) {
- gin.SetMode(gin.TestMode)
- c := newGrokCacheTestContext(0)
- c.Request.Header.Set(grokConversationIDHeader, "native-session")
- require.Empty(t, resolveGrokCacheIdentity(c, []byte(`{"model":"grok","input":"hi"}`), "", "grok-4.5"))
- require.Empty(t, resolveGrokCacheIdentity(nil, []byte(`{"model":"grok","prompt_cache_key":"key"}`), "key", "grok-4.5"))
- }
- func TestGrokConversationHeaderIsScopedToGrokRequestScheduling(t *testing.T) {
- gin.SetMode(gin.TestMode)
- body := []byte(`{"model":"grok","prompt_cache_key":"body-session","input":"hi"}`)
- grokContext := newGrokCacheTestContext(601)
- grokContext.Request.Header.Set(grokConversationIDHeader, "native-grok-session")
- require.Equal(t, "native-grok-session", (&OpenAIGatewayService{}).ExtractSessionID(grokContext, body))
- openAIContext := newGrokCacheTestContext(601)
- openAIContext.Set("api_key", &APIKey{ID: 601, Group: &Group{Platform: PlatformOpenAI}})
- openAIContext.Request.Header.Set(grokConversationIDHeader, "must-be-ignored")
- require.Equal(t, "body-session", (&OpenAIGatewayService{}).ExtractSessionID(openAIContext, body))
- withoutGrokHeader := newGrokCacheTestContext(601)
- withoutGrokHeader.Set("api_key", &APIKey{ID: 601, Group: &Group{Platform: PlatformOpenAI}})
- require.Equal(t,
- (&OpenAIGatewayService{}).GenerateSessionHash(withoutGrokHeader, body),
- (&OpenAIGatewayService{}).GenerateSessionHash(openAIContext, body),
- )
- }
- func TestApplyGrokCacheIdentityWritesResponsesBodyAndHeader(t *testing.T) {
- sourceBody := []byte(`{"model":"grok-4.5","prompt_cache_key":"raw-client-key"}`)
- body, err := applyGrokResponsesCacheIdentity(sourceBody, sourceBody, "isolated-id", true)
- require.NoError(t, err)
- require.Equal(t, "isolated-id", gjson.GetBytes(body, "prompt_cache_key").String())
- require.Equal(t, "web_search", gjson.GetBytes(body, "tools.0.type").String())
- require.Equal(t, "x_search", gjson.GetBytes(body, "tools.1.type").String())
- require.Equal(t, grokFreeCacheDisabledToolChoice, gjson.GetBytes(body, "tool_choice").String())
- headers := make(http.Header)
- headers.Set(grokConversationIDHeader, "spoofed-client-value")
- applyGrokCacheHeaders(headers, "isolated-id")
- require.Equal(t, "isolated-id", headers.Get(grokConversationIDHeader))
- applyGrokCacheHeaders(headers, "")
- require.Empty(t, headers.Get(grokConversationIDHeader))
- chatBody, err := stripGrokChatPromptCacheKey(body)
- require.NoError(t, err)
- require.False(t, gjson.GetBytes(chatBody, "prompt_cache_key").Exists())
- unscopedSourceBody := []byte(`{"model":"grok","prompt_cache_key":"raw-client-key"}`)
- unscopedBody, err := applyGrokResponsesCacheIdentity(unscopedSourceBody, unscopedSourceBody, "", true)
- require.NoError(t, err)
- require.False(t, gjson.GetBytes(unscopedBody, "prompt_cache_key").Exists())
- require.False(t, gjson.GetBytes(unscopedBody, "tools").Exists())
- require.False(t, gjson.GetBytes(unscopedBody, "tool_choice").Exists())
- }
- func TestGrokFreeMessagesClientToolCacheDefaultsOnForKnownFree(t *testing.T) {
- account := healthyGrokOAuthGatewayTestAccount(901, "access-token")
- account.Credentials["subscription_tier"] = " FREE "
- tests := []struct {
- name string
- toolChoiceJSON string
- wantChoice bool
- }{
- {name: "missing tool choice"},
- {name: "automatic tool choice", toolChoiceJSON: `,"tool_choice":"auto"`, wantChoice: true},
- }
- for _, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- intentBody := []byte(`{"model":"grok","tools":[{"type":"function","name":"lookup","description":"look up a value","parameters":{"type":"object"}},{"type":"function","name":"save","parameters":{"type":"object"}}]` + tt.toolChoiceJSON + `}`)
- body, err := applyGrokResponsesCacheIdentity(intentBody, intentBody, "isolated-id", true)
- require.NoError(t, err)
- body, err = applyGrokFreeMessagesFunctionToolCacheRoute(body, intentBody, account, "isolated-id")
- require.NoError(t, err)
- require.Equal(t, "isolated-id", gjson.GetBytes(body, "prompt_cache_key").String())
- tools := gjson.GetBytes(body, "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, "function", tools[1].Get("type").String())
- require.Equal(t, "save", tools[1].Get("name").String())
- require.Equal(t, "web_search", tools[2].Get("type").String())
- require.Equal(t, "x_search", tools[3].Get("type").String())
- require.Equal(t, tt.wantChoice, gjson.GetBytes(body, "tool_choice").Exists())
- })
- }
- }
- func TestGrokFreeMessagesClientToolCacheDefaultsWithMissingAccountSetting(t *testing.T) {
- body := []byte(`{"model":"grok","tools":[{"type":"function","name":"view_image","parameters":{"type":"object"}}],"tool_choice":"auto"}`)
- tests := []struct {
- name string
- extra map[string]any
- }{
- {name: "nil extra"},
- {name: "empty extra", extra: map[string]any{}},
- {name: "unrelated extra", extra: map[string]any{"other_setting": true}},
- }
- for _, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- account := healthyGrokOAuthGatewayTestAccount(90101, "access-token")
- account.Credentials["subscription_tier"] = "free"
- account.Extra = tt.extra
- patched, err := applyGrokFreeMessagesFunctionToolCacheRoute(body, body, account, "isolated-id")
- require.NoError(t, err)
- tools := gjson.GetBytes(patched, "tools").Array()
- require.Len(t, tools, 3)
- require.Equal(t, "web_search", tools[1].Get("type").String())
- require.Equal(t, "x_search", tools[2].Get("type").String())
- })
- }
- }
- func TestApplyGrokCacheIdentityAppendsNativeToolsWhenSearchPresent(t *testing.T) {
- account := healthyGrokOAuthGatewayTestAccount(901, "access-token")
- account.Credentials["subscription_tier"] = " FREE "
- // Function tools INCLUDING web_search → convert + complement with x_search.
- intentBody := []byte(`{"model":"grok","tools":[{"type":"function","name":"lookup","description":"look up a value","parameters":{"type":"object"}},{"type":"function","name":"web_search","description":"search","parameters":{"type":"object"}}]}`)
- body, err := applyGrokResponsesCacheIdentity(intentBody, intentBody, "isolated-id", true)
- require.NoError(t, err)
- body, err = applyGrokFreeMessagesFunctionToolCacheRoute(body, intentBody, account, "isolated-id")
- require.NoError(t, err)
- tools := gjson.GetBytes(body, "tools").Array()
- require.Len(t, tools, 3, "lookup(function) + web_search(native) + x_search(native)")
- 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, "x_search", tools[2].Get("type").String())
- }
- func TestGrokFreeClientToolCacheAccountOptIn(t *testing.T) {
- account := healthyGrokOAuthGatewayTestAccount(9011, "access-token")
- account.Credentials["subscription_tier"] = "free"
- account.Extra = map[string]any{grokClientToolCacheOptInExtraKey: true}
- intentBody := []byte(`{"model":"grok","tools":[{"type":"function","name":"view_image","parameters":{"type":"object"}},{"type":"function","name":"read_file","parameters":{"type":"object"}}],"tool_choice":"auto"}`)
- body, err := applyGrokResponsesCacheIdentity(intentBody, intentBody, "isolated-id", true)
- require.NoError(t, err)
- body, err = applyGrokFreeMessagesFunctionToolCacheRoute(body, intentBody, account, "isolated-id")
- require.NoError(t, err)
- tools := gjson.GetBytes(body, "tools").Array()
- require.Len(t, tools, 4)
- require.Equal(t, "view_image", tools[0].Get("name").String())
- require.Equal(t, "read_file", tools[1].Get("name").String())
- require.Equal(t, "web_search", tools[2].Get("type").String())
- require.Equal(t, "x_search", tools[3].Get("type").String())
- }
- func TestGrokFreeMessagesClientToolCacheAccountOptOut(t *testing.T) {
- body := []byte(`{"model":"grok","tools":[{"type":"function","name":"view_image","parameters":{"type":"object"}}],"tool_choice":"auto"}`)
- tests := []struct {
- name string
- value any
- }{
- {name: "explicit false", value: false},
- {name: "string false is malformed", value: "false"},
- {name: "numeric value is malformed", value: 1},
- {name: "null value is malformed", value: nil},
- }
- for _, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- account := healthyGrokOAuthGatewayTestAccount(90111, "access-token")
- account.Credentials["subscription_tier"] = "free"
- account.Extra = map[string]any{grokClientToolCacheOptInExtraKey: tt.value}
- patched, err := applyGrokFreeMessagesFunctionToolCacheRoute(body, body, account, "isolated-id")
- require.NoError(t, err)
- require.JSONEq(t, string(body), string(patched))
- })
- }
- }
- func TestGrokFreeClientToolCacheRequestOptInOverridesAccountOptOut(t *testing.T) {
- account := healthyGrokOAuthGatewayTestAccount(9014, "access-token")
- account.Credentials["subscription_tier"] = "free"
- account.Extra = map[string]any{grokClientToolCacheOptInExtraKey: false}
- c := newGrokCacheTestContext(9014)
- c.Request.Header.Set(grokClientToolCacheOptInHeader, "prefer-cache")
- intentBody := []byte(`{"model":"grok","tools":[{"type":"function","name":"view_image","parameters":{"type":"object"}}],"tool_choice":"auto"}`)
- body, err := applyGrokResponsesCacheIdentity(intentBody, intentBody, "isolated-id", true)
- require.NoError(t, err)
- body, err = applyGrokFreeRequestToolCacheRoute(c, body, intentBody, account, "isolated-id")
- require.NoError(t, err)
- tools := gjson.GetBytes(body, "tools").Array()
- require.Len(t, tools, 3)
- require.Equal(t, "view_image", tools[0].Get("name").String())
- require.Equal(t, "web_search", tools[1].Get("type").String())
- require.Equal(t, "x_search", tools[2].Get("type").String())
- }
- func TestGrokFreeChatRequestClientToolCacheDefaultsOnWithoutClientFingerprint(t *testing.T) {
- account := healthyGrokOAuthGatewayTestAccount(90140, "access-token")
- account.Credentials["subscription_tier"] = "free"
- c := newGrokCacheTestContext(90140)
- c.Request.URL.Path = "/v1/chat/completions"
- body := []byte(`{"model":"grok","tools":[{"type":"function","name":"view_image","parameters":{"type":"object"}}],"tool_choice":"auto"}`)
- patched, err := applyGrokFreeRequestToolCacheRoute(c, body, body, account, "isolated-id")
- require.NoError(t, err)
- tools := gjson.GetBytes(patched, "tools").Array()
- require.Len(t, tools, 3)
- require.Equal(t, "view_image", tools[0].Get("name").String())
- require.Equal(t, "web_search", tools[1].Get("type").String())
- require.Equal(t, "x_search", tools[2].Get("type").String())
- }
- func TestGrokFreeClientToolCacheClaudeDesktopResponsesAutoOptIn(t *testing.T) {
- account := healthyGrokOAuthGatewayTestAccount(90141, "access-token")
- account.Credentials["subscription_tier"] = "free"
- body := []byte(`{"model":"grok","tools":[{"type":"function","name":"Read","parameters":{"type":"object"}},{"type":"function","name":"Edit","parameters":{"type":"object"}}],"tool_choice":"auto"}`)
- for _, xApp := range []string{"cli", "cli-bg"} {
- t.Run(xApp, func(t *testing.T) {
- c := newGrokCacheTestContext(90141)
- // The desktop marker text is intentionally not required; CC Switch and
- // Claude Desktop may change the descriptive User-Agent suffix.
- c.Request.Header.Set("User-Agent", "claude-cli/2.1.215 (external, future-desktop, agent-sdk/0.3.215)")
- c.Request.Header.Set("X-App", xApp)
- c.Request.Header.Set("anthropic-client-platform", "desktop_app")
- c.Request.Header.Set("X-Claude-Code-Session-Id", "desktop-session-1")
- patched, err := applyGrokFreeRequestToolCacheRoute(c, body, body, account, "isolated-id")
- require.NoError(t, err)
- tools := gjson.GetBytes(patched, "tools").Array()
- require.Len(t, tools, 4)
- require.Equal(t, "Read", tools[0].Get("name").String())
- require.Equal(t, "Edit", tools[1].Get("name").String())
- require.Equal(t, "web_search", tools[2].Get("type").String())
- require.Equal(t, "x_search", tools[3].Get("type").String())
- })
- }
- }
- func TestGrokFreeClientToolCacheClaudeDesktopFingerprintRequiresAllSignals(t *testing.T) {
- account := healthyGrokOAuthGatewayTestAccount(90142, "access-token")
- account.Credentials["subscription_tier"] = "free"
- account.Extra = map[string]any{grokClientToolCacheOptInExtraKey: false}
- body := []byte(`{"model":"grok","tools":[{"type":"function","name":"Read","parameters":{"type":"object"}}],"tool_choice":"auto"}`)
- tests := []struct {
- name string
- path string
- ua string
- xApp string
- platform string
- session string
- }{
- {
- name: "chat path",
- path: "/v1/chat/completions",
- ua: "claude-cli/2.1.215 (external, claude-desktop-3p, agent-sdk/0.3.215)",
- xApp: "cli",
- platform: "desktop_app",
- session: "desktop-session-1",
- },
- {
- name: "compact responses path",
- path: "/v1/responses/compact",
- ua: "claude-cli/2.1.215 (external, claude-desktop-3p, agent-sdk/0.3.215)",
- xApp: "cli",
- platform: "desktop_app",
- session: "desktop-session-1",
- },
- {
- name: "non claude cli user agent",
- path: "/v1/responses",
- ua: "Mozilla/5.0 (claude-desktop-3p)",
- xApp: "cli",
- platform: "desktop_app",
- session: "desktop-session-1",
- },
- {
- name: "missing x app",
- path: "/v1/responses",
- ua: "claude-cli/2.1.215 (external, claude-desktop-3p, agent-sdk/0.3.215)",
- platform: "desktop_app",
- session: "desktop-session-1",
- },
- {
- name: "wrong x app",
- path: "/v1/responses",
- ua: "claude-cli/2.1.215 (external, claude-desktop-3p, agent-sdk/0.3.215)",
- xApp: "desktop",
- platform: "desktop_app",
- session: "desktop-session-1",
- },
- {
- name: "missing client platform",
- path: "/v1/responses",
- ua: "claude-cli/2.1.215 (external, claude-desktop-3p, agent-sdk/0.3.215)",
- xApp: "cli",
- session: "desktop-session-1",
- },
- {
- name: "wrong client platform",
- path: "/v1/responses",
- ua: "claude-cli/2.1.215 (external, claude-desktop-3p, agent-sdk/0.3.215)",
- xApp: "cli",
- platform: "web",
- session: "desktop-session-1",
- },
- {
- name: "missing session header",
- path: "/v1/responses",
- ua: "claude-cli/2.1.215 (external, claude-desktop-3p, agent-sdk/0.3.215)",
- xApp: "cli",
- platform: "desktop_app",
- },
- }
- for _, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- c := newGrokCacheTestContext(90142)
- c.Request.URL.Path = tt.path
- if tt.ua != "" {
- c.Request.Header.Set("User-Agent", tt.ua)
- }
- if tt.xApp != "" {
- c.Request.Header.Set("X-App", tt.xApp)
- }
- if tt.platform != "" {
- c.Request.Header.Set("anthropic-client-platform", tt.platform)
- }
- if tt.session != "" {
- c.Request.Header.Set("X-Claude-Code-Session-Id", tt.session)
- }
- patched, err := applyGrokFreeRequestToolCacheRoute(c, body, body, account, "isolated-id")
- require.NoError(t, err)
- require.JSONEq(t, string(body), string(patched))
- })
- }
- }
- func TestGrokFreeClientToolCacheExplicitRequestOptOut(t *testing.T) {
- account := healthyGrokOAuthGatewayTestAccount(90144, "access-token")
- account.Credentials["subscription_tier"] = "free"
- body := []byte(`{"model":"grok","tools":[{"type":"function","name":"Read","parameters":{"type":"object"}}],"tool_choice":"auto"}`)
- for _, value := range []string{"0", "false", "no", "off"} {
- t.Run(value, func(t *testing.T) {
- c := newGrokCacheTestContext(90144)
- c.Request.URL.Path = "/v1/chat/completions"
- c.Request.Header.Set(grokClientToolCacheOptInHeader, value)
- patched, err := applyGrokFreeRequestToolCacheRoute(c, body, body, account, "isolated-id")
- require.NoError(t, err)
- require.JSONEq(t, string(body), string(patched))
- })
- }
- }
- func TestGrokFreeClientToolCacheClaudeDesktopAutoOptInDoesNotOverridePaidTier(t *testing.T) {
- account := healthyGrokOAuthGatewayTestAccount(90143, "access-token")
- account.Credentials["subscription_tier"] = "supergrok"
- c := newGrokCacheTestContext(90143)
- 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", "desktop-session-1")
- body := []byte(`{"model":"grok","tools":[{"type":"function","name":"Read","parameters":{"type":"object"}}],"tool_choice":"auto"}`)
- patched, err := applyGrokFreeRequestToolCacheRoute(c, body, body, account, "isolated-id")
- require.NoError(t, err)
- require.JSONEq(t, string(body), string(patched))
- }
- func TestGrokFreeRequestClientSearchFunctionUsesDefaultAccountPolicy(t *testing.T) {
- account := healthyGrokOAuthGatewayTestAccount(9015, "access-token")
- account.Credentials["subscription_tier"] = "free"
- c := newGrokCacheTestContext(9015)
- body := []byte(`{"model":"grok","tools":[{"type":"function","name":"view_image","parameters":{"type":"object"}},{"type":"function","name":"web_search","parameters":{"type":"object"}}],"tool_choice":"auto"}`)
- patched, err := applyGrokFreeRequestToolCacheRoute(c, body, body, account, "isolated-id")
- require.NoError(t, err)
- tools := gjson.GetBytes(patched, "tools").Array()
- require.Len(t, tools, 3)
- require.Equal(t, "view_image", tools[0].Get("name").String())
- require.Equal(t, "web_search", tools[1].Get("type").String())
- require.Equal(t, "x_search", tools[2].Get("type").String())
- }
- func TestGrokFreeRequestToolChoiceNoneUsesSafeCacheRoute(t *testing.T) {
- account := healthyGrokOAuthGatewayTestAccount(9016, "access-token")
- account.Credentials["subscription_tier"] = "free"
- c := newGrokCacheTestContext(9016)
- body := []byte(`{"model":"grok","tools":[{"type":"function","name":"view_image","parameters":{"type":"object"}}],"tool_choice":"none"}`)
- patched, err := applyGrokFreeRequestToolCacheRoute(c, body, body, account, "isolated-id")
- require.NoError(t, err)
- tools := gjson.GetBytes(patched, "tools").Array()
- require.Len(t, tools, 3)
- require.Equal(t, "view_image", tools[0].Get("name").String())
- require.Equal(t, "web_search", tools[1].Get("type").String())
- require.Equal(t, "x_search", tools[2].Get("type").String())
- require.Equal(t, "none", gjson.GetBytes(patched, "tool_choice").String())
- }
- func TestApplyGrokCacheIdentityRecognizesResponsesLiteAdditionalTools(t *testing.T) {
- intentBody := []byte(`{"model":"grok","input":[{"type":"additional_tools","tools":[{"type":"function","name":"lookup","parameters":{"type":"object"}}]},{"type":"message","role":"user","content":"hello"}]}`)
- patchedBody := []byte(`{"model":"grok-4.5","tools":[{"type":"function","name":"lookup","parameters":{"type":"object"}}],"input":[{"type":"message","role":"user","content":"hello"}]}`)
- patched, err := applyGrokResponsesCacheIdentity(patchedBody, intentBody, "isolated-id", true)
- require.NoError(t, err)
- tools := gjson.GetBytes(patched, "tools").Array()
- require.Len(t, tools, 1)
- require.Equal(t, "lookup", tools[0].Get("name").String())
- require.False(t, gjson.GetBytes(patched, "tool_choice").Exists())
- require.Equal(t, "isolated-id", gjson.GetBytes(patched, "prompt_cache_key").String())
- }
- func TestGrokFreeCacheRoutePreservesMixedSupportedToolsWithSearchIntent(t *testing.T) {
- account := healthyGrokOAuthGatewayTestAccount(9012, "access-token")
- account.Credentials["subscription_tier"] = "free"
- intentBody := []byte(`{"model":"grok","tools":[{"type":"function","name":"view_image","parameters":{"type":"object"}},{"type":"shell"},{"type":"web_search"}],"tool_choice":"auto"}`)
- body, err := applyGrokResponsesCacheIdentity(intentBody, intentBody, "isolated-id", true)
- require.NoError(t, err)
- body, err = applyGrokFreeMessagesFunctionToolCacheRoute(body, intentBody, account, "isolated-id")
- require.NoError(t, err)
- tools := gjson.GetBytes(body, "tools").Array()
- require.Len(t, tools, 4)
- require.Equal(t, "function", tools[0].Get("type").String())
- require.Equal(t, "shell", tools[1].Get("type").String())
- require.Equal(t, "web_search", tools[2].Get("type").String())
- require.Equal(t, "x_search", tools[3].Get("type").String())
- }
- func TestGrokClientToolCacheOptInDoesNotOverridePaidTier(t *testing.T) {
- account := healthyGrokOAuthGatewayTestAccount(9013, "access-token")
- account.Credentials["subscription_tier"] = "supergrok"
- account.Extra = map[string]any{grokClientToolCacheOptInExtraKey: true}
- body := []byte(`{"model":"grok","tools":[{"type":"function","name":"view_image","parameters":{"type":"object"}}],"tool_choice":"auto"}`)
- patched, err := applyGrokFreeMessagesFunctionToolCacheRoute(body, body, account, "isolated-id")
- require.NoError(t, err)
- require.JSONEq(t, string(body), string(patched))
- }
- func TestApplyGrokCacheIdentityRequiresPatchedFunctionTools(t *testing.T) {
- account := healthyGrokOAuthGatewayTestAccount(902, "access-token")
- account.Credentials["subscription_tier"] = "free"
- intentBody := []byte(`{"model":"grok","tools":[{"type":"function","name":"lookup"}],"tool_choice":"auto"}`)
- tests := []struct {
- name string
- patchedBody string
- }{
- {name: "missing tools", patchedBody: `{"model":"grok-4.5"}`},
- {name: "empty tools", patchedBody: `{"model":"grok-4.5","tools":[]}`},
- {name: "native tools only", patchedBody: `{"model":"grok-4.5","tools":[{"type":"web_search"}]}`},
- {name: "unexpected patched tool", patchedBody: `{"model":"grok-4.5","tools":[{"type":"function","name":"lookup"},{"type":"namespace","name":"server"}]}`},
- }
- for _, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- beforeTools := gjson.Get(tt.patchedBody, "tools")
- body, err := applyGrokResponsesCacheIdentity([]byte(tt.patchedBody), intentBody, "isolated-id", true)
- require.NoError(t, err)
- body, err = applyGrokFreeMessagesFunctionToolCacheRoute(body, intentBody, account, "isolated-id")
- require.NoError(t, err)
- require.Equal(t, "isolated-id", gjson.GetBytes(body, "prompt_cache_key").String())
- afterTools := gjson.GetBytes(body, "tools")
- require.Equal(t, beforeTools.Exists(), afterTools.Exists())
- require.Equal(t, beforeTools.Raw, afterTools.Raw)
- })
- }
- }
- func TestGrokFreeMessagesFunctionToolCacheRouteRequiresKnownFreeTier(t *testing.T) {
- // Include web_search as a function to also cover its conversion to a native tool.
- intentBody := []byte(`{"model":"grok","tools":[{"type":"function","name":"lookup"},{"type":"function","name":"web_search"}],"tool_choice":"auto"}`)
- tests := []struct {
- name string
- account *Account
- wantMix bool
- }{
- {
- name: "free credential tier",
- account: func() *Account {
- a := healthyGrokOAuthGatewayTestAccount(910, "access-token")
- a.Credentials["subscription_tier"] = "free"
- return a
- }(),
- wantMix: true,
- },
- {
- name: "free billing tier",
- account: func() *Account {
- a := healthyGrokOAuthGatewayTestAccount(911, "access-token")
- a.Extra = map[string]any{grokBillingExtraKey: map[string]any{"plan": "FREE"}}
- return a
- }(),
- wantMix: true,
- },
- {
- name: "free successful billing has blank plan",
- account: func() *Account {
- a := healthyGrokOAuthGatewayTestAccount(9111, "access-token")
- a.Extra = map[string]any{grokBillingExtraKey: map[string]any{
- "status_code": http.StatusOK,
- "source": "billing_probe",
- "monthly_updated_at": "2026-07-15T05:00:00Z",
- }}
- return a
- }(),
- wantMix: true,
- },
- {
- name: "free rolling token quota",
- account: func() *Account {
- a := healthyGrokOAuthGatewayTestAccount(9112, "access-token")
- a.Extra = map[string]any{grokQuotaSnapshotExtraKey: map[string]any{
- "headers_observed": true,
- "tokens": map[string]any{"limit": xai.GrokFreeRolling24hTokenLimit},
- }}
- return a
- }(),
- wantMix: true,
- },
- {
- name: "legacy free rolling token quota",
- account: func() *Account {
- a := healthyGrokOAuthGatewayTestAccount(9113, "access-token")
- a.Extra = map[string]any{grokQuotaSnapshotExtraKey: map[string]any{
- "headers_observed": true,
- "tokens": map[string]any{"limit": int64(2_000_000)},
- }}
- return a
- }(),
- wantMix: true,
- },
- {
- name: "supergrok remains unchanged",
- account: func() *Account {
- a := healthyGrokOAuthGatewayTestAccount(912, "access-token")
- a.Credentials["subscription_tier"] = "supergrok"
- return a
- }(),
- },
- {
- name: "paid billing overrides stale free quota",
- account: func() *Account {
- a := healthyGrokOAuthGatewayTestAccount(9121, "access-token")
- a.Extra = map[string]any{
- grokBillingExtraKey: map[string]any{"plan": "SuperGrok", "status_code": http.StatusOK},
- grokQuotaSnapshotExtraKey: map[string]any{
- "headers_observed": true,
- "tokens": map[string]any{"limit": int64(2_000_000)},
- },
- }
- return a
- }(),
- },
- {
- name: "paid billing overrides stale free credentials",
- account: func() *Account {
- a := healthyGrokOAuthGatewayTestAccount(9123, "access-token")
- a.Credentials["subscription_tier"] = "free"
- a.Extra = map[string]any{
- grokBillingExtraKey: map[string]any{"plan": "SuperGrok", "status_code": http.StatusOK},
- }
- return a
- }(),
- },
- {
- name: "partial billing without monthly evidence remains unknown",
- account: func() *Account {
- a := healthyGrokOAuthGatewayTestAccount(9122, "access-token")
- a.Extra = map[string]any{grokBillingExtraKey: map[string]any{
- "status_code": http.StatusOK,
- "source": "billing_probe",
- "partial": true,
- "failed_windows": []string{"monthly"},
- }}
- return a
- }(),
- },
- {
- name: "unknown tier remains unchanged",
- account: healthyGrokOAuthGatewayTestAccount(913, "access-token"),
- },
- {
- name: "api key remains unchanged",
- account: &Account{
- ID: 914,
- Platform: PlatformGrok,
- Type: AccountTypeAPIKey,
- },
- },
- }
- for _, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- body, err := applyGrokFreeMessagesFunctionToolCacheRoute(intentBody, intentBody, tt.account, "isolated-id")
- require.NoError(t, err)
- tools := gjson.GetBytes(body, "tools").Array()
- if tt.wantMix {
- require.Len(t, tools, 3)
- require.Equal(t, "web_search", tools[1].Get("type").String())
- require.Equal(t, "x_search", tools[2].Get("type").String())
- return
- }
- require.Len(t, tools, 2, "non-free accounts should not get native search injected")
- })
- }
- }
- func TestGrokFreeMessagesFunctionToolCacheRouteRequiresIdentity(t *testing.T) {
- account := healthyGrokOAuthGatewayTestAccount(915, "access-token")
- account.Credentials["subscription_tier"] = "free"
- body := []byte(`{"model":"grok","tools":[{"type":"function","name":"lookup"}],"tool_choice":"auto"}`)
- patched, err := applyGrokFreeMessagesFunctionToolCacheRoute(body, body, account, "")
- require.NoError(t, err)
- require.JSONEq(t, string(body), string(patched))
- require.Len(t, gjson.GetBytes(patched, "tools").Array(), 1)
- }
- func TestApplyGrokCacheIdentityPreservesIneligibleClientToolFields(t *testing.T) {
- tests := []struct {
- name string
- body string
- }{
- {
- name: "empty tools array",
- body: `{"model":"grok","tools":[]}`,
- },
- {
- name: "null tools",
- body: `{"model":"grok","tools":null}`,
- },
- {
- name: "tool choice only",
- body: `{"model":"grok","tool_choice":{"type":"function","name":"lookup"}}`,
- },
- {
- name: "null tool choice",
- body: `{"model":"grok","tool_choice":null}`,
- },
- {
- name: "native tool with auto choice",
- body: `{"model":"grok","tools":[{"type":"web_search"}],"tool_choice":"auto"}`,
- },
- {
- name: "function with required choice",
- body: `{"model":"grok","tools":[{"type":"function","name":"lookup"}],"tool_choice":"required"}`,
- },
- {
- name: "function with none choice",
- body: `{"model":"grok","tools":[{"type":"function","name":"lookup"}],"tool_choice":"none"}`,
- },
- {
- name: "function with specific choice",
- body: `{"model":"grok","tools":[{"type":"function","name":"lookup"}],"tool_choice":{"type":"function","name":"lookup"}}`,
- },
- {
- name: "function with object auto choice",
- body: `{"model":"grok","tools":[{"type":"function","name":"lookup"}],"tool_choice":{"type":"auto"}}`,
- },
- {
- name: "function mixed with unsupported tool",
- body: `{"model":"grok","tools":[{"type":"function","name":"lookup"},{"type":"namespace","name":"client_tools"}],"tool_choice":"auto"}`,
- },
- {
- name: "unsupported tool only",
- body: `{"model":"grok","tools":[{"type":"namespace","name":"client_tools"}]}`,
- },
- {
- name: "chat completions function shape",
- body: `{"model":"grok","tools":[{"type":"function","function":{"name":"lookup","parameters":{"type":"object"}}}],"tool_choice":"auto"}`,
- },
- {
- name: "incomplete responses function",
- body: `{"model":"grok","tools":[{"type":"function","parameters":{"type":"object"}}]}`,
- },
- }
- for _, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- beforeTools := gjson.Get(tt.body, "tools")
- beforeChoice := gjson.Get(tt.body, "tool_choice")
- body, err := applyGrokResponsesCacheIdentity([]byte(tt.body), []byte(tt.body), "isolated-id", true)
- require.NoError(t, err)
- require.Equal(t, "isolated-id", gjson.GetBytes(body, "prompt_cache_key").String())
- require.Equal(t, beforeTools.Exists(), gjson.GetBytes(body, "tools").Exists())
- require.Equal(t, beforeTools.Raw, gjson.GetBytes(body, "tools").Raw)
- require.Equal(t, beforeChoice.Exists(), gjson.GetBytes(body, "tool_choice").Exists())
- require.Equal(t, beforeChoice.Raw, gjson.GetBytes(body, "tool_choice").Raw)
- })
- }
- }
- func TestApplyGrokCacheIdentityUsesPreSanitizationToolIntent(t *testing.T) {
- tests := []struct {
- name string
- intentBody string
- }{
- {
- name: "unsupported tools removed by sanitizer",
- intentBody: `{"model":"grok","tools":[{"type":"namespace","name":"client_tools"}]}`,
- },
- {
- name: "tool choice removed with unsupported tool",
- intentBody: `{"model":"grok","tools":[{"type":"namespace","name":"client_tools"}],"tool_choice":{"type":"namespace","name":"client_tools"}}`,
- },
- }
- for _, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- // This is the shape apply receives after patchGrokResponsesBody has
- // removed unsupported tools and their associated tool_choice.
- patchedBody := []byte(`{"model":"grok-4.5","input":"hello"}`)
- body, err := applyGrokResponsesCacheIdentity(patchedBody, []byte(tt.intentBody), "isolated-id", true)
- require.NoError(t, err)
- require.Equal(t, "isolated-id", gjson.GetBytes(body, "prompt_cache_key").String())
- require.False(t, gjson.GetBytes(body, "tools").Exists())
- require.False(t, gjson.GetBytes(body, "tool_choice").Exists())
- })
- }
- }
- func TestApplyGrokCacheIdentityWithoutFreeTierRoutingOnlyWritesIdentity(t *testing.T) {
- sourceBody := []byte(`{"model":"grok-4.5","input":"hello"}`)
- body, err := applyGrokResponsesCacheIdentity(sourceBody, sourceBody, "isolated-id", false)
- require.NoError(t, err)
- require.Equal(t, "isolated-id", gjson.GetBytes(body, "prompt_cache_key").String())
- require.False(t, gjson.GetBytes(body, "tools").Exists())
- require.False(t, gjson.GetBytes(body, "tool_choice").Exists())
- }
- func TestGrokCompactRequestSkipsCacheIdentityAndNativeTools(t *testing.T) {
- gin.SetMode(gin.TestMode)
- c := newGrokCacheTestContext(701)
- c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses/compact", nil)
- body := []byte(`{"model":"grok","input":"compact this","prompt_cache_key":"raw-client-key"}`)
- identity := resolveGrokCacheIdentity(c, body, "", "grok-4.5")
- patched, err := applyGrokResponsesCacheIdentity(body, body, identity, true)
- require.NoError(t, err)
- require.Empty(t, identity)
- require.False(t, gjson.GetBytes(patched, "prompt_cache_key").Exists())
- require.False(t, gjson.GetBytes(patched, "tools").Exists())
- require.False(t, gjson.GetBytes(patched, "tool_choice").Exists())
- }
- func TestResolveGrokCacheIdentityConcurrentDeterminism(t *testing.T) {
- gin.SetMode(gin.TestMode)
- const workers = 50
- body := []byte(`{"model":"grok","messages":[{"role":"system","content":"stable"},{"role":"user","content":"hello"}]}`)
- identities := make(chan string, workers)
- var wg sync.WaitGroup
- for i := 0; i < workers; i++ {
- wg.Add(1)
- go func() {
- defer wg.Done()
- identities <- resolveGrokCacheIdentity(newGrokCacheTestContext(501), body, "", "grok-4.5")
- }()
- }
- wg.Wait()
- close(identities)
- var first string
- for identity := range identities {
- if first == "" {
- first = identity
- }
- require.Equal(t, first, identity)
- }
- require.NotEmpty(t, first)
- }
|