Add xAI Responses prompt_cache_key sticky routing - #23
Conversation
Mirror conversation affinity keys across Chat Completions headers and Responses body fields so Grok cache hits follow xAI's documented routing. Co-authored-by: Cursor <cursoragent@cursor.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughThe changes expand xAI prompt-cache key propagation and Responses prefix derivation, update related documentation and tests, and add live regression coverage for direct and gateway-routed Anthropic message streaming. ChangesxAI prompt caching
Anthropic streaming regressions
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
interceptors/promptcaching.go (1)
121-151: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winRedundant JSON parsing of the same body.
xaiAPIType(via itsDetectAPIType(rawBody)fallback),xaiBodyPromptCacheKey, andsetXAIPromptCacheKeyeach independentlyjson.Unmarshalthe samerawBody— up to 3 full parses per request on this hot path. For large conversation histories this adds avoidable CPU/allocation overhead.Consider parsing once (e.g., into a
map[string]interface{}or a shared struct) and threading the result throughxaiBodyPromptCacheKey/setXAIPromptCacheKey, falling back toapiTypedetection from the same parse where possible.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@interceptors/promptcaching.go` around lines 121 - 151, Avoid repeatedly unmarshalling rawBody in the prompt-caching flow by parsing it once and reusing the parsed representation across xaiAPIType, xaiBodyPromptCacheKey, and setXAIPromptCacheKey. Thread the shared parsed data through these helpers, derive API type and cache keys from it where possible, and preserve existing behavior when parsing fails or no cache key is available.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@interceptors/promptcaching.go`:
- Around line 205-219: The setXAIPromptCacheKey function must avoid unmarshaling
and re-marshaling the entire request body, which can alter numeric values.
Replace the map[string]interface{} round trip with targeted JSON editing, such
as map[string]json.RawMessage, while preserving existing non-empty
prompt_cache_key values and injecting only the prompt_cache_key field.
---
Nitpick comments:
In `@interceptors/promptcaching.go`:
- Around line 121-151: Avoid repeatedly unmarshalling rawBody in the
prompt-caching flow by parsing it once and reusing the parsed representation
across xaiAPIType, xaiBodyPromptCacheKey, and setXAIPromptCacheKey. Thread the
shared parsed data through these helpers, derive API type and cache keys from it
where possible, and preserve existing behavior when parsing fails or no cache
key is available.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 56aaeac7-d3d8-473d-bb55-b193dbfa8897
📒 Files selected for processing (4)
DESIGN.mdREADME.mdinterceptors/promptcaching.gointerceptors/promptcaching_test.go
📜 Review details
🔇 Additional comments (5)
interceptors/promptcaching.go (2)
174-193: 🗄️ Data Integrity & IntegrationCorrectness hinges on
DetectAPITypeFromPath/DetectAPITypeaccuracy.If either misclassifies a Responses request as non-Responses (e.g., unusual base path, proxied route),
setXAIPromptCacheKeyis skipped entirely and only the header gets mirrored — which per xAI docs does not provide sticky routing for the Responses API. Always set x-grok-conv-id (or prompt_cache_key for Responses API) — Routes requests to the same server, maximizing cache hits. Worth confirming these two functions robustly handle all deployed base-path/routing variants for Responses traffic.
683-734: LGTM!interceptors/promptcaching_test.go (1)
2091-2269: LGTM! Solid end-to-end coverage of the new precedence/mirroring branches and the Responsesinstructionsprefix hashing.README.md (1)
342-342: LGTM!DESIGN.md (1)
1024-1027: LGTM!
Preserve request JSON numbers via RawMessage injection, parse the body once in the xAI caching path, and fold Anthropic AI Gateway stream regressions with corrected event counting and lower max_tokens. Co-authored-by: Cursor <cursoragent@cursor.com>
There was a problem hiding this comment.
🧹 Nitpick comments (2)
aigateway_regression_live_test.go (2)
136-147: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate "read response / check status / assert stream" boilerplate across both live tests.
The read-body/status-check/truncate/assert sequence is nearly identical in
TestLiveAnthropicMessagesStreamCompletes(Lines 136-147) andTestLiveAgentuityAIGatewayAnthropicMessagesStreamCompletes(Lines 211-221). Extracting a small helper (e.g.readAndAssertLiveStream(t, resp)) would remove the duplication and keep both tests focused on request construction.Also applies to: 211-221
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@aigateway_regression_live_test.go` around lines 136 - 147, The live tests duplicate response-body reading, status validation, and stream assertion. Extract this sequence into a shared helper such as readAndAssertLiveStream, then call it from TestLiveAnthropicMessagesStreamCompletes and TestLiveAgentuityAIGatewayAnthropicMessagesStreamCompletes while preserving the existing error messages and truncation behavior.
128-128: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePrefer structured JSON marshaling over raw string interpolation.
This test builds the request body via string concatenation with
modelspliced directly into a JSON literal, while the sibling test (line 175) builds the same shape viamap[string]any+json.Marshal. Aligning both on structured marshaling avoids relying onmodelnever containing a character that needs escaping and keeps the two tests consistent.♻️ Suggested consistency fix
- body := `{"model":"` + model + `","stream":true,"max_tokens":512,"thinking":{"type":"disabled"},"messages":[{"role":"user","content":[{"type":"text","text":"Reply with GENESIS_DRIVER_SMOKE_OK and nothing else."}]}]}` - req := httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader([]byte(body))) + bodyMap := map[string]any{ + "model": model, + "stream": true, + "max_tokens": 512, + "thinking": map[string]any{"type": "disabled"}, + "messages": []map[string]any{ + { + "role": "user", + "content": []map[string]any{ + {"type": "text", "text": "Reply with GENESIS_DRIVER_SMOKE_OK and nothing else."}, + }, + }, + }, + } + rawBody, err := json.Marshal(bodyMap) + if err != nil { + t.Fatalf("marshal request: %v", err) + } + req := httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(rawBody))🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@aigateway_regression_live_test.go` at line 128, Update the request-body construction in the affected test to use a structured map matching the sibling test’s payload shape, then serialize it with json.Marshal. Preserve the existing model, streaming, token limit, thinking, and message values while removing raw interpolation into the JSON literal.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@aigateway_regression_live_test.go`:
- Around line 136-147: The live tests duplicate response-body reading, status
validation, and stream assertion. Extract this sequence into a shared helper
such as readAndAssertLiveStream, then call it from
TestLiveAnthropicMessagesStreamCompletes and
TestLiveAgentuityAIGatewayAnthropicMessagesStreamCompletes while preserving the
existing error messages and truncation behavior.
- Line 128: Update the request-body construction in the affected test to use a
structured map matching the sibling test’s payload shape, then serialize it with
json.Marshal. Preserve the existing model, streaming, token limit, thinking, and
message values while removing raw interpolation into the JSON literal.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 98ad9b54-ebfc-4f43-b305-3302c8d04de1
📒 Files selected for processing (3)
aigateway_regression_live_test.gointerceptors/promptcaching.gointerceptors/promptcaching_test.go
🚧 Files skipped from review as they are similar to previous changes (2)
- interceptors/promptcaching_test.go
- interceptors/promptcaching.go
📜 Review details
🔇 Additional comments (6)
aigateway_regression_live_test.go (6)
8-20: LGTM!
100-147: LGTM!
149-222: LGTM!
223-238: LGTM!
240-298: LGTM!
300-315: LGTM!
Share request body marshaling and response read/assert logic across the local router and deployed AI Gateway Anthropic stream regressions. Co-authored-by: Cursor <cursoragent@cursor.com>
1 participant
Summary
prompt_cache_keyin addition to Chat Completionsx-grok-conv-idprompt_cache_keyviamap[string]json.RawMessageto avoid numeric JSON corruption, parsing the body once per requestmax_tokensTest plan
go test ./interceptors/ -run 'PromptCachingInterceptor_XAI|DeriveCacheKeyFromPrefix|SetXAIPromptCacheKey'go test -tags=integration . -run 'TestLiveAIGatewayRegressionModels|TestLiveAnthropic'(skips without live env)github.com/agentuity/llmproxyin ion and verify Grok Responses sticky routingx-grok-conv-idon WSSummary by CodeRabbit
New Features
Bug Fixes
Tests