fix(xai): send body prompt_cache_key for Responses cache affinity

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

@@ -33,6 +33,24 @@ import { asyncStreamFrom, collectStream } from "../../../test-utils/stream"
 import { makeCreateMessageMetadata } from "../../../test-utils/api"
 import { clearAllMocks } from "../../../test-utils/reset"
 
+/**
+ * Flip only `supportsPromptCache` on the handler's resolved model, so the rest of
+ * the catalog entry (maxTokens, temperature, reasoning) stays as the real one.
+ *
+ * The xAI catalog is declared `as const`, so every entry's `supportsPromptCache`
+ * is the literal `true` and no real uncached model exists to point at. The double
+ * assertion is confined to this test double and is the only way to reach the
+ * `supportsPromptCache === false` arm; the production code is untouched by it.
+ */
+function stubSupportsPromptCache(handler: XAIHandler, supportsPromptCache: boolean): void {
+	const realModel = handler.getModel()
+	vitest
+		.spyOn(handler, "getModel")
+		.mockReturnValue({ ...realModel, info: { ...realModel.info, supportsPromptCache } } as unknown as ReturnType<
+			XAIHandler["getModel"]
+		>)
+}
+
 describe("XAIHandler", () => {
 	let handler: XAIHandler
 
@@ -565,6 +583,8 @@ describe("XAIHandler", () => {
 			model: xaiDefaultModelId,
 			input: [{ role: "user", content: [{ type: "input_text", text: "test prompt" }] }],
 			store: false,
+			// No task metadata exists on this path, so the handler session scopes the cache.
+			prompt_cache_key: expect.any(String),
 		})
 	})
 
@@ -912,4 +932,130 @@ describe("XAIHandler", () => {
 		expect(removeEventListenerSpy).toHaveBeenCalledTimes(1)
 		expect(removeEventListenerSpy).toHaveBeenCalledWith("abort", listener)
 	})
+
+	describe("prompt cache affinity", () => {
+		const lastRequestBody = () => mockResponsesCreate.mock.calls[mockResponsesCreate.mock.calls.length - 1][0]
+
+		it("should send the task id as the body prompt_cache_key", async () => {
+			mockResponsesCreate.mockResolvedValueOnce(asyncStreamFrom([]))
+
+			const stream = handler.createMessage("test prompt", [], makeCreateMessageMetadata({ taskId: "task-abc" }))
+			await stream.next()
+
+			expect(lastRequestBody().prompt_cache_key).toBe("task-abc")
+		})
+
+		it("should keep the key stable across two turns of the same task", async () => {
+			mockResponsesCreate.mockResolvedValue(asyncStreamFrom([]))
+
+			const first = handler.createMessage("first turn", [], makeCreateMessageMetadata({ taskId: "task-stable" }))
+			await first.next()
+			const second = handler.createMessage(
+				"second turn",
+				[],
+				makeCreateMessageMetadata({ taskId: "task-stable" }),
+			)
+			await second.next()
+
+			// Two separate requests, one value: routing is pinned to one cache server.
+			expect(mockResponsesCreate).toHaveBeenCalledTimes(2)
+			expect(mockResponsesCreate.mock.calls[0][0].prompt_cache_key).toBe("task-stable")
+			expect(mockResponsesCreate.mock.calls[1][0].prompt_cache_key).toBe(
+				mockResponsesCreate.mock.calls[0][0].prompt_cache_key,
+			)
+		})
+
+		it("should give distinct keys to distinct tasks", async () => {
+			mockResponsesCreate.mockResolvedValue(asyncStreamFrom([]))
+
+			for (const taskId of ["task-one", "task-two"]) {
+				const stream = handler.createMessage("test prompt", [], makeCreateMessageMetadata({ taskId }))
+				await stream.next()
+			}
+
+			expect(mockResponsesCreate.mock.calls[0][0].prompt_cache_key).toBe("task-one")
+			expect(mockResponsesCreate.mock.calls[1][0].prompt_cache_key).toBe("task-two")
+		})
+
+		it("should fall back to the handler session id when no task id is supplied", async () => {
+			mockResponsesCreate.mockResolvedValue(asyncStreamFrom([]))
+
+			const stream = handler.createMessage("test prompt", [])
+			await stream.next()
+
+			const key = lastRequestBody().prompt_cache_key
+			expect(typeof key).toBe("string")
+			expect(key.length).toBeGreaterThan(0)
+
+			// The same handler keeps the same fallback key on the next turn.
+			const second = handler.createMessage("test prompt", [])
+			await second.next()
+			expect(lastRequestBody().prompt_cache_key).toBe(key)
+		})
+
+		it("should clamp an over-length task id to the 64-character limit", async () => {
+			mockResponsesCreate.mockResolvedValueOnce(asyncStreamFrom([]))
+
+			const overLongTaskId = "t".repeat(80)
+			const stream = handler.createMessage(
+				"test prompt",
+				[],
+				makeCreateMessageMetadata({ taskId: overLongTaskId }),
+			)
+			await stream.next()
+
+			const key = lastRequestBody().prompt_cache_key
+			expect(key).toHaveLength(64)
+			expect(key).toBe(overLongTaskId.slice(0, 64))
+		})
+
+		it("should never send an empty-string prompt_cache_key", async () => {
+			mockResponsesCreate.mockResolvedValue(asyncStreamFrom([]))
+
+			const stream = handler.createMessage("test prompt", [], makeCreateMessageMetadata({ taskId: "" }))
+			await stream.next()
+
+			const key = lastRequestBody().prompt_cache_key
+			// An empty task id falls back to the session id instead of sending "".
+			expect(key).not.toBe("")
+			expect(typeof key).toBe("string")
+		})
+
+		it("should omit the key when the model does not support prompt caching", async () => {
+			const uncachedHandler = new XAIHandler({ apiModelId: xaiDefaultModelId })
+			stubSupportsPromptCache(uncachedHandler, false)
+
+			mockResponsesCreate.mockResolvedValueOnce(asyncStreamFrom([]))
+			const stream = uncachedHandler.createMessage(
+				"test prompt",
+				[],
+				makeCreateMessageMetadata({ taskId: "task-x" }),
+			)
+			await stream.next()
+
+			expect(lastRequestBody()).not.toHaveProperty("prompt_cache_key")
+		})
+
+		it("completePrompt should scope the prompt_cache_key to the handler session", async () => {
+			mockResponsesCreate.mockResolvedValue({ output_text: "response" })
+
+			await handler.completePrompt("first prompt")
+			await handler.completePrompt("second prompt")
+
+			const firstKey = mockResponsesCreate.mock.calls[0][0].prompt_cache_key
+			expect(typeof firstKey).toBe("string")
+			expect(firstKey.length).toBeGreaterThan(0)
+			expect(mockResponsesCreate.mock.calls[1][0].prompt_cache_key).toBe(firstKey)
+		})
+
+		it("completePrompt should omit the key when the model does not support prompt caching", async () => {
+			const uncachedHandler = new XAIHandler({ apiModelId: xaiDefaultModelId })
+			stubSupportsPromptCache(uncachedHandler, false)
+
+			mockResponsesCreate.mockResolvedValueOnce({ output_text: "response" })
+			await uncachedHandler.completePrompt("test prompt")
+
+			expect(mockResponsesCreate.mock.calls[0][0]).not.toHaveProperty("prompt_cache_key")
+		})
+	})
 })

src/api/providers/xai.ts

@@ -1,5 +1,6 @@
 import { Anthropic } from "@anthropic-ai/sdk"
 import OpenAI from "openai"
+import { v7 as uuidv7 } from "uuid"
 
 import { type XAIModelId, xaiDefaultModelId, xaiModels, ApiProviderError } from "@roo-code/types"
 import { TelemetryService } from "@roo-code/telemetry"
@@ -26,10 +27,16 @@
 
 const XAI_DEFAULT_TEMPERATURE = 0
 
+// Conservative length guard for `prompt_cache_key`: OpenAI's documented cap is 64
+// chars; xAI does not document a bound for this field, so we apply the same cap.
+// Task ids are uuid v7 (36 chars), so this never truncates in practice.
+const XAI_MAX_PROMPT_CACHE_KEY_LENGTH = 64
+
 export class XAIHandler extends BaseProvider implements SingleCompletionHandler {
 	protected options: ApiHandlerOptions
 	private client: OpenAI
 	private readonly providerName = "xAI"
+	private readonly sessionId: string = uuidv7()
 
 	constructor(options: ApiHandlerOptions) {
 		super()
@@ -45,6 +52,38 @@
 		})
 	}
 
+	/**
+	 * Cache routing affinity key for the Responses API.
+	 *
+	 * xAI's prompt cache lives on individual servers, so every turn of one
+	 * conversation has to reach the same server to reuse it. For the Responses API
+	 * the documented 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.
+	 *
+	 * The task id is the longest-lived identifier available here: `Task.taskId` is
+	 * assigned once and restored from history on resume, so it survives restarts
+	 * and provider switches, which a per-handler session id does not. The handler
+	 * session id is the fallback for callers that send no task id, and is the sole
+	 * key for `completePrompt`, which carries no task metadata at all.
+	 *
+	 * Returns undefined (never an empty string) when there is nothing to key on or
+	 * when the model does not support prompt caching, so the field is omitted
+	 * rather than sent blank.
+	 */
+	private getPromptCacheKey(model: ReturnType<XAIHandler["getModel"]>, taskId?: string): string | undefined {
+		if (!model.info.supportsPromptCache) {
+			return undefined
+		}
+
+		const key = taskId || this.sessionId
+		if (!key) {
+			return undefined
+		}
+
+		return key.length > XAI_MAX_PROMPT_CACHE_KEY_LENGTH ? key.slice(0, XAI_MAX_PROMPT_CACHE_KEY_LENGTH) : key
+	}
+
 	override getModel() {
 		const id =
 			this.options.apiModelId && this.options.apiModelId in xaiModels
@@ -169,6 +208,12 @@
 			include: ["reasoning.encrypted_content"],
 		}
 
+		// Route every turn of this task to the server holding its cache.
+		const promptCacheKey = this.getPromptCacheKey(model, metadata?.taskId)
+		if (promptCacheKey) {
+			requestBody.prompt_cache_key = promptCacheKey
+		}
+
 		// Model params are always resolved by getModel(); send them unconditionally.
 		requestBody.max_output_tokens = model.maxTokens
 		requestBody.temperature = model.temperature
@@ -255,11 +300,15 @@
 				requestOptions.timeout = toSdkTimeoutMs(timeoutMs)
 			}
 
+			const promptCacheKey = this.getPromptCacheKey(model)
 			const response = await this.client.responses.create(
 				{
 					model: model.id,
 					input: [{ role: "user", content: [{ type: "input_text", text: prompt }] }],
 					store: false,
+					// A completion is its own cache scope: no task metadata exists on
+					// this path, so the handler session id is the whole key.
+					...(promptCacheKey ? { prompt_cache_key: promptCacheKey } : {}),
 				},
 				Object.keys(requestOptions).length > 0 ? requestOptions : undefined,
 			)