diff --git a/backend/internal/service/openai_gateway_grok_cache.go b/backend/internal/service/openai_gateway_grok_cache.go index a7bb29f8e4..ca798d8141 100644 --- a/backend/internal/service/openai_gateway_grok_cache.go +++ b/backend/internal/service/openai_gateway_grok_cache.go @@ -102,24 +102,15 @@ func resolveGrokCacheIdentity(c *gin.Context, body []byte, explicitKey, upstream } func explicitGrokCacheSeed(c *gin.Context, body []byte, explicitKey string) string { - seed := "" - if c != nil { - // Claude Code session is the most stable multi-turn identity for - // /v1/messages → Grok bridges. Prefer it over generic session headers - // so prompt cache routing matches CPA behavior. - seed = extractClaudeCodeSessionID(c, body) - if seed == "" { - seed = strings.TrimSpace(c.GetHeader("session_id")) - } - if seed == "" { - seed = strings.TrimSpace(c.GetHeader("conversation_id")) - } - if seed == "" { - seed = strings.TrimSpace(c.GetHeader(grokConversationIDHeader)) - } + // Claude Code session is the most stable multi-turn identity for + // /v1/messages → Grok bridges. Prefer it over generic session headers so + // prompt cache routing matches CPA behavior. + seed := extractClaudeCodeSessionID(c, body) + if seed == "" { + seed = explicitOpenAIHeaderSessionID(c) } - if seed == "" && len(body) > 0 { - seed = extractClaudeCodeSessionIDFromPayload(body) + if seed == "" && c != nil { + seed = strings.TrimSpace(c.GetHeader(grokConversationIDHeader)) } if seed == "" && len(body) > 0 { seed = strings.TrimSpace(gjson.GetBytes(body, "prompt_cache_key").String()) diff --git a/backend/internal/service/openai_gateway_grok_cache_test.go b/backend/internal/service/openai_gateway_grok_cache_test.go index 2255f4a2b5..dc3beaa461 100644 --- a/backend/internal/service/openai_gateway_grok_cache_test.go +++ b/backend/internal/service/openai_gateway_grok_cache_test.go @@ -140,6 +140,124 @@ func TestResolveGrokCacheIdentityExplicitHeaderPriority(t *testing.T) { 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) + want := resolveGrokCacheIdentity(onlyCurrent, []byte(`{"model":"grok","input":"unrelated"}`), "", "grok-4.5") + require.Equal(t, want, got, header.name) + c.Request.Header.Del(header.name) + } +} + +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 { + require.Equal(t, header.value, 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) diff --git a/backend/internal/service/openai_gateway_grok_chat_bridge_test.go b/backend/internal/service/openai_gateway_grok_chat_bridge_test.go index 188e38f702..06f434341a 100644 --- a/backend/internal/service/openai_gateway_grok_chat_bridge_test.go +++ b/backend/internal/service/openai_gateway_grok_chat_bridge_test.go @@ -273,6 +273,72 @@ func TestForwardGrokChatViaResponsesNonStreamingCachesAndReturnsChat(t *testing. require.NotNil(t, repo.updates[account.ID][grokQuotaSnapshotExtraKey]) } +func TestForwardGrokChatViaResponsesCodeBuddyUsesStableConversationHeader(t *testing.T) { + gin.SetMode(gin.TestMode) + const conversationID = "codebuddy-session-42" + tests := []struct { + name string + requestID string + messageID string + body []byte + }{ + { + name: "first turn", + requestID: "request-one", + messageID: "message-one", + body: []byte(`{"model":"grok","messages":[{"role":"user","content":"first question"}],"stream":false,"tools":[]}`), + }, + { + name: "later turn with tools", + requestID: "request-two", + messageID: "message-two", + body: []byte(`{"model":"grok","messages":[{"role":"system","content":"be concise"},{"role":"user","content":"different later question"}],"stream":false,"tools":[{"type":"function","function":{"name":"lookup","description":"Lookup a value","parameters":{"type":"object","properties":{"key":{"type":"string"}}}}}],"tool_choice":"auto"}`), + }, + } + + var stableIdentity string + for index, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, grokChatRawEndpoint, bytes.NewReader(tt.body)) + c.Request.Header.Set(codeBuddyConversationHeader, conversationID) + c.Request.Header.Set("X-Conversation-Request-ID", tt.requestID) + c.Request.Header.Set("X-Conversation-Message-ID", tt.messageID) + c.Request.Header.Set("X-Request-ID", "generic-"+tt.requestID) + c.Set("api_key", &APIKey{ID: 7111}) + + identity := resolveGrokCacheIdentity(c, tt.body, "", "grok-4.5") + require.NotEmpty(t, identity) + if index == 0 { + stableIdentity = identity + } else { + require.Equal(t, stableIdentity, identity) + } + require.NotContains(t, identity, conversationID) + + account := grokChatBridgeTestAccount(int64(711 + index)) + repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{account.ID: account}, + }} + upstream := &httpUpstreamRecorder{resp: grokChatBridgeCompletedResponse("resp_codebuddy_"+strconv.Itoa(index), 4096)} + svc := &OpenAIGatewayService{ + httpUpstream: upstream, + grokTokenProvider: NewGrokTokenProvider(repo, nil), + accountRepo: repo, + } + + result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, tt.body, "", "") + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, xai.DefaultCLIBaseURL+"/responses", upstream.lastReq.URL.String()) + require.Equal(t, grokChatResponsesEndpoint, result.UpstreamEndpoint) + require.Equal(t, identity, gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String()) + require.Equal(t, identity, upstream.lastReq.Header.Get(grokConversationIDHeader)) + }) + } +} + func TestForwardGrokChatViaResponsesTraeToolHistoryKeepsCacheRoute(t *testing.T) { gin.SetMode(gin.TestMode) diff --git a/backend/internal/service/openai_gateway_scheduling.go b/backend/internal/service/openai_gateway_scheduling.go index 10d65e0aa8..ba350f5062 100644 --- a/backend/internal/service/openai_gateway_scheduling.go +++ b/backend/internal/service/openai_gateway_scheduling.go @@ -19,6 +19,37 @@ import ( "github.com/tidwall/gjson" ) +const ( + openCodeSessionAffinityHeader = "X-Session-Affinity" + openCodeSessionIDHeader = "X-Session-Id" + openCodeNativeSessionHeader = "X-OpenCode-Session" + codeBuddyConversationHeader = "X-Conversation-ID" +) + +// explicitOpenAIHeaderSessionID resolves stable conversation identifiers sent +// by OpenAI-compatible clients. Keep this list limited to session-scoped +// fields: request/message IDs rotate every turn and would defeat sticky routing +// and upstream prompt caching. +func explicitOpenAIHeaderSessionID(c *gin.Context) string { + if c == nil { + return "" + } + + for _, header := range []string{ + "session_id", + "conversation_id", + openCodeSessionAffinityHeader, + openCodeSessionIDHeader, + openCodeNativeSessionHeader, + codeBuddyConversationHeader, + } { + if sessionID := strings.TrimSpace(c.GetHeader(header)); sessionID != "" { + return sessionID + } + } + return "" +} + // ExtractSessionID extracts the raw session ID from headers or body without hashing. // Used by ForwardAsAnthropic to pass as prompt_cache_key for upstream cache. func (s *OpenAIGatewayService) ExtractSessionID(c *gin.Context, body []byte) string { @@ -30,10 +61,7 @@ func explicitOpenAISessionID(c *gin.Context, body []byte) string { return "" } - sessionID := strings.TrimSpace(c.GetHeader("session_id")) - if sessionID == "" { - sessionID = strings.TrimSpace(c.GetHeader("conversation_id")) - } + sessionID := explicitOpenAIHeaderSessionID(c) if sessionID == "" && len(body) > 0 { sessionID = strings.TrimSpace(gjson.GetBytes(body, "prompt_cache_key").String()) } @@ -49,10 +77,7 @@ func explicitOpenAIRequestSessionID(c *gin.Context, body []byte) string { return "" } - sessionID := strings.TrimSpace(c.GetHeader("session_id")) - if sessionID == "" { - sessionID = strings.TrimSpace(c.GetHeader("conversation_id")) - } + sessionID := explicitOpenAIHeaderSessionID(c) if sessionID == "" && isGrokRequestContext(c) { sessionID = strings.TrimSpace(c.GetHeader(grokConversationIDHeader)) } @@ -81,9 +106,11 @@ func (s *OpenAIGatewayService) GenerateExplicitSessionHash(c *gin.Context, body // Priority: // 1. Header: session_id // 2. Header: conversation_id -// 3. Header: x-grok-conv-id (Grok groups only) -// 4. Body: prompt_cache_key (opencode) -// 5. Body: content-based fallback (model + system + tools + first user message) +// 3. Header: x-session-affinity / x-session-id / x-opencode-session (OpenCode) +// 4. Header: x-conversation-id (CodeBuddy) +// 5. Header: x-grok-conv-id (Grok groups only) +// 6. Body: prompt_cache_key +// 7. Body: content-based fallback (model + system + tools + first user message) func (s *OpenAIGatewayService) GenerateSessionHash(c *gin.Context, body []byte) string { if c == nil { return "" diff --git a/backend/internal/service/openai_gateway_service_test.go b/backend/internal/service/openai_gateway_service_test.go index d0bebb2ef6..2c2725ea78 100644 --- a/backend/internal/service/openai_gateway_service_test.go +++ b/backend/internal/service/openai_gateway_service_test.go @@ -222,6 +222,61 @@ func TestOpenAIGatewayService_GenerateSessionHash_Priority(t *testing.T) { } } +func TestOpenAIGatewayService_ClientSessionHeaderPriority(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil) + c.Set("api_key", &APIKey{ID: 901, Group: &Group{Platform: PlatformGrok}}) + + headers := []struct { + name string + value string + }{ + {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) + } + + svc := &OpenAIGatewayService{} + body := []byte(`{"prompt_cache_key":"body-session"}`) + for _, header := range headers { + require.Equal(t, header.value, svc.ExtractSessionID(c, body), header.name) + require.Equal(t, fmt.Sprintf("%016x", xxhash.Sum64String(header.value)), svc.GenerateExplicitSessionHash(c, body), header.name) + if header.name != grokConversationIDHeader { + require.Equal(t, header.value, explicitOpenAISessionID(c, body), header.name) + } + c.Request.Header.Del(header.name) + } + require.Equal(t, "body-session", svc.ExtractSessionID(c, body)) +} + +func TestOpenAIGatewayService_ClientSessionHeadersIgnorePerRequestIDs(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil) + for name, value := range map[string]string{ + "X-Conversation-Request-ID": "request-rotates-every-turn", + "X-Conversation-Message-ID": "message-rotates-every-turn", + "X-Request-ID": "generic-request-id", + } { + c.Request.Header.Set(name, value) + } + + svc := &OpenAIGatewayService{} + require.Empty(t, explicitOpenAIHeaderSessionID(c)) + require.Empty(t, svc.ExtractSessionID(c, nil)) + require.Empty(t, svc.GenerateExplicitSessionHash(c, nil)) +} + func TestOpenAIGatewayService_GenerateSessionHash_UsesXXHash64(t *testing.T) { gin.SetMode(gin.TestMode) rec := httptest.NewRecorder()