Skip to content

Commit 5c0a7ab

Browse files
myk1ytedelauna
authored andcommitted
fix(xai): send body prompt_cache_key for Responses cache affinity (#2256)
* fix(xai): send body prompt_cache_key for Responses cache affinity What: add a task-scoped `prompt_cache_key` to the xAI Responses API request body in both `createMessage` and `completePrompt`, gated on the model's `supportsPromptCache` flag and clamped to 64 characters. Why: xAI's prompt cache is stored per server, so every turn of one conversation must be routed to the same server to reuse it. For the Responses API the documented affinity lever is the body field `prompt_cache_key`; the `x-grok-conv-id` header is the Chat Completions / gRPC equivalent and is not read on this path, so it was deliberately NOT added. Before this change `xai.ts` sent no affinity identifier at all, leaving server-side cache routing to load-balancer luck. Impact: xAI turns of one task now pin to one cache server, raising cache-read tokens and lowering cost on long conversations. The key is omitted (never sent as an empty string) when no identifier is available or the model does not support prompt caching. Usage accounting is unchanged: the shared normalizer already reads `input_tokens_details.cached_tokens`. * docs(api): attribute the 64-char prompt_cache_key cap honestly What: reword the XAI_MAX_PROMPT_CACHE_KEY_LENGTH comment. Why: review seat (PR #2256 mlarge P3) caught that xAI's docs do not state a 64-char bound for prompt_cache_key - the bound is OpenAI's documented one; we apply it as a conservative guard. Impact: comment-only; zero code change.
1 parent 62f6caa commit 5c0a7ab

2 files changed

Lines changed: 195 additions & 0 deletions

File tree

‎src/api/providers/__tests__/xai.spec.ts‎

Lines changed: 146 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,24 @@ import { asyncStreamFrom, collectStream } from "../../../test-utils/stream"
3333
import { makeCreateMessageMetadata } from "../../../test-utils/api"
3434
import { clearAllMocks } from "../../../test-utils/reset"
3535

36+
/**
37+
* Flip only `supportsPromptCache` on the handler's resolved model, so the rest of
38+
* the catalog entry (maxTokens, temperature, reasoning) stays as the real one.
39+
*
40+
* The xAI catalog is declared `as const`, so every entry's `supportsPromptCache`
41+
* is the literal `true` and no real uncached model exists to point at. The double
42+
* assertion is confined to this test double and is the only way to reach the
43+
* `supportsPromptCache === false` arm; the production code is untouched by it.
44+
*/
45+
function stubSupportsPromptCache(handler: XAIHandler, supportsPromptCache: boolean): void {
46+
const realModel = handler.getModel()
47+
vitest
48+
.spyOn(handler, "getModel")
49+
.mockReturnValue({ ...realModel, info: { ...realModel.info, supportsPromptCache } } as unknown as ReturnType<
50+
XAIHandler["getModel"]
51+
>)
52+
}
53+
3654
describe("XAIHandler", () => {
3755
let handler: XAIHandler
3856

@@ -565,6 +583,8 @@ describe("XAIHandler", () => {
565583
model: xaiDefaultModelId,
566584
input: [{ role: "user", content: [{ type: "input_text", text: "test prompt" }] }],
567585
store: false,
586+
// No task metadata exists on this path, so the handler session scopes the cache.
587+
prompt_cache_key: expect.any(String),
568588
})
569589
})
570590

@@ -912,4 +932,130 @@ describe("XAIHandler", () => {
912932
expect(removeEventListenerSpy).toHaveBeenCalledTimes(1)
913933
expect(removeEventListenerSpy).toHaveBeenCalledWith("abort", listener)
914934
})
935+
936+
describe("prompt cache affinity", () => {
937+
const lastRequestBody = () => mockResponsesCreate.mock.calls[mockResponsesCreate.mock.calls.length - 1][0]
938+
939+
it("should send the task id as the body prompt_cache_key", async () => {
940+
mockResponsesCreate.mockResolvedValueOnce(asyncStreamFrom([]))
941+
942+
const stream = handler.createMessage("test prompt", [], makeCreateMessageMetadata({ taskId: "task-abc" }))
943+
await stream.next()
944+
945+
expect(lastRequestBody().prompt_cache_key).toBe("task-abc")
946+
})
947+
948+
it("should keep the key stable across two turns of the same task", async () => {
949+
mockResponsesCreate.mockResolvedValue(asyncStreamFrom([]))
950+
951+
const first = handler.createMessage("first turn", [], makeCreateMessageMetadata({ taskId: "task-stable" }))
952+
await first.next()
953+
const second = handler.createMessage(
954+
"second turn",
955+
[],
956+
makeCreateMessageMetadata({ taskId: "task-stable" }),
957+
)
958+
await second.next()
959+
960+
// Two separate requests, one value: routing is pinned to one cache server.
961+
expect(mockResponsesCreate).toHaveBeenCalledTimes(2)
962+
expect(mockResponsesCreate.mock.calls[0][0].prompt_cache_key).toBe("task-stable")
963+
expect(mockResponsesCreate.mock.calls[1][0].prompt_cache_key).toBe(
964+
mockResponsesCreate.mock.calls[0][0].prompt_cache_key,
965+
)
966+
})
967+
968+
it("should give distinct keys to distinct tasks", async () => {
969+
mockResponsesCreate.mockResolvedValue(asyncStreamFrom([]))
970+
971+
for (const taskId of ["task-one", "task-two"]) {
972+
const stream = handler.createMessage("test prompt", [], makeCreateMessageMetadata({ taskId }))
973+
await stream.next()
974+
}
975+
976+
expect(mockResponsesCreate.mock.calls[0][0].prompt_cache_key).toBe("task-one")
977+
expect(mockResponsesCreate.mock.calls[1][0].prompt_cache_key).toBe("task-two")
978+
})
979+
980+
it("should fall back to the handler session id when no task id is supplied", async () => {
981+
mockResponsesCreate.mockResolvedValue(asyncStreamFrom([]))
982+
983+
const stream = handler.createMessage("test prompt", [])
984+
await stream.next()
985+
986+
const key = lastRequestBody().prompt_cache_key
987+
expect(typeof key).toBe("string")
988+
expect(key.length).toBeGreaterThan(0)
989+
990+
// The same handler keeps the same fallback key on the next turn.
991+
const second = handler.createMessage("test prompt", [])
992+
await second.next()
993+
expect(lastRequestBody().prompt_cache_key).toBe(key)
994+
})
995+
996+
it("should clamp an over-length task id to the 64-character limit", async () => {
997+
mockResponsesCreate.mockResolvedValueOnce(asyncStreamFrom([]))
998+
999+
const overLongTaskId = "t".repeat(80)
1000+
const stream = handler.createMessage(
1001+
"test prompt",
1002+
[],
1003+
makeCreateMessageMetadata({ taskId: overLongTaskId }),
1004+
)
1005+
await stream.next()
1006+
1007+
const key = lastRequestBody().prompt_cache_key
1008+
expect(key).toHaveLength(64)
1009+
expect(key).toBe(overLongTaskId.slice(0, 64))
1010+
})
1011+
1012+
it("should never send an empty-string prompt_cache_key", async () => {
1013+
mockResponsesCreate.mockResolvedValue(asyncStreamFrom([]))
1014+
1015+
const stream = handler.createMessage("test prompt", [], makeCreateMessageMetadata({ taskId: "" }))
1016+
await stream.next()
1017+
1018+
const key = lastRequestBody().prompt_cache_key
1019+
// An empty task id falls back to the session id instead of sending "".
1020+
expect(key).not.toBe("")
1021+
expect(typeof key).toBe("string")
1022+
})
1023+
1024+
it("should omit the key when the model does not support prompt caching", async () => {
1025+
const uncachedHandler = new XAIHandler({ apiModelId: xaiDefaultModelId })
1026+
stubSupportsPromptCache(uncachedHandler, false)
1027+
1028+
mockResponsesCreate.mockResolvedValueOnce(asyncStreamFrom([]))
1029+
const stream = uncachedHandler.createMessage(
1030+
"test prompt",
1031+
[],
1032+
makeCreateMessageMetadata({ taskId: "task-x" }),
1033+
)
1034+
await stream.next()
1035+
1036+
expect(lastRequestBody()).not.toHaveProperty("prompt_cache_key")
1037+
})
1038+
1039+
it("completePrompt should scope the prompt_cache_key to the handler session", async () => {
1040+
mockResponsesCreate.mockResolvedValue({ output_text: "response" })
1041+
1042+
await handler.completePrompt("first prompt")
1043+
await handler.completePrompt("second prompt")
1044+
1045+
const firstKey = mockResponsesCreate.mock.calls[0][0].prompt_cache_key
1046+
expect(typeof firstKey).toBe("string")
1047+
expect(firstKey.length).toBeGreaterThan(0)
1048+
expect(mockResponsesCreate.mock.calls[1][0].prompt_cache_key).toBe(firstKey)
1049+
})
1050+
1051+
it("completePrompt should omit the key when the model does not support prompt caching", async () => {
1052+
const uncachedHandler = new XAIHandler({ apiModelId: xaiDefaultModelId })
1053+
stubSupportsPromptCache(uncachedHandler, false)
1054+
1055+
mockResponsesCreate.mockResolvedValueOnce({ output_text: "response" })
1056+
await uncachedHandler.completePrompt("test prompt")
1057+
1058+
expect(mockResponsesCreate.mock.calls[0][0]).not.toHaveProperty("prompt_cache_key")
1059+
})
1060+
})
9151061
})

‎src/api/providers/xai.ts‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import { Anthropic } from "@anthropic-ai/sdk"
22
import OpenAI from "openai"
3+
import { v7 as uuidv7 } from "uuid"
34

45
import { type XAIModelId, xaiDefaultModelId, xaiModels, ApiProviderError } from "@roo-code/types"
56
import { TelemetryService } from "@roo-code/telemetry"
@@ -26,10 +27,16 @@ import { isMcpTool } from "../../utils/mcp-name"
2627

2728
const XAI_DEFAULT_TEMPERATURE = 0
2829

30+
// Conservative length guard for `prompt_cache_key`: OpenAI's documented cap is 64
31+
// chars; xAI does not document a bound for this field, so we apply the same cap.
32+
// Task ids are uuid v7 (36 chars), so this never truncates in practice.
33+
const XAI_MAX_PROMPT_CACHE_KEY_LENGTH = 64
34+
2935
export class XAIHandler extends BaseProvider implements SingleCompletionHandler {
3036
protected options: ApiHandlerOptions
3137
private client: OpenAI
3238
private readonly providerName = "xAI"
39+
private readonly sessionId: string = uuidv7()
3340

3441
constructor(options: ApiHandlerOptions) {
3542
super()
@@ -45,6 +52,38 @@ export class XAIHandler extends BaseProvider implements SingleCompletionHandler
4552
})
4653
}
4754

55+
/**
56+
* Cache routing affinity key for the Responses API.
57+
*
58+
* xAI's prompt cache lives on individual servers, so every turn of one
59+
* conversation has to reach the same server to reuse it. For the Responses API
60+
* the documented lever is the body field `prompt_cache_key`; the
61+
* `x-grok-conv-id` header is the Chat Completions / gRPC equivalent and is not
62+
* read on this path.
63+
*
64+
* The task id is the longest-lived identifier available here: `Task.taskId` is
65+
* assigned once and restored from history on resume, so it survives restarts
66+
* and provider switches, which a per-handler session id does not. The handler
67+
* session id is the fallback for callers that send no task id, and is the sole
68+
* key for `completePrompt`, which carries no task metadata at all.
69+
*
70+
* Returns undefined (never an empty string) when there is nothing to key on or
71+
* when the model does not support prompt caching, so the field is omitted
72+
* rather than sent blank.
73+
*/
74+
private getPromptCacheKey(model: ReturnType<XAIHandler["getModel"]>, taskId?: string): string | undefined {
75+
if (!model.info.supportsPromptCache) {
76+
return undefined
77+
}
78+
79+
const key = taskId || this.sessionId
80+
if (!key) {
81+
return undefined
82+
}
83+
84+
return key.length > XAI_MAX_PROMPT_CACHE_KEY_LENGTH ? key.slice(0, XAI_MAX_PROMPT_CACHE_KEY_LENGTH) : key
85+
}
86+
4887
override getModel() {
4988
const id =
5089
this.options.apiModelId && this.options.apiModelId in xaiModels
@@ -169,6 +208,12 @@ export class XAIHandler extends BaseProvider implements SingleCompletionHandler
169208
include: ["reasoning.encrypted_content"],
170209
}
171210

211+
// Route every turn of this task to the server holding its cache.
212+
const promptCacheKey = this.getPromptCacheKey(model, metadata?.taskId)
213+
if (promptCacheKey) {
214+
requestBody.prompt_cache_key = promptCacheKey
215+
}
216+
172217
// Model params are always resolved by getModel(); send them unconditionally.
173218
requestBody.max_output_tokens = model.maxTokens
174219
requestBody.temperature = model.temperature
@@ -255,11 +300,15 @@ export class XAIHandler extends BaseProvider implements SingleCompletionHandler
255300
requestOptions.timeout = toSdkTimeoutMs(timeoutMs)
256301
}
257302

303+
const promptCacheKey = this.getPromptCacheKey(model)
258304
const response = await this.client.responses.create(
259305
{
260306
model: model.id,
261307
input: [{ role: "user", content: [{ type: "input_text", text: prompt }] }],
262308
store: false,
309+
// A completion is its own cache scope: no task metadata exists on
310+
// this path, so the handler session id is the whole key.
311+
...(promptCacheKey ? { prompt_cache_key: promptCacheKey } : {}),
263312
},
264313
Object.keys(requestOptions).length > 0 ? requestOptions : undefined,
265314
)

0 commit comments

Comments
 (0)