diff --git a/src/api/providers/__tests__/openai-native-usage.spec.ts b/src/api/providers/__tests__/openai-native-usage.spec.ts index 8922c2cf9c..f299abf1d4 100644 --- a/src/api/providers/__tests__/openai-native-usage.spec.ts +++ b/src/api/providers/__tests__/openai-native-usage.spec.ts @@ -1,6 +1,22 @@ import { describe, it, expect, beforeEach } from "vitest" import { OpenAiNativeHandler } from "../openai-native" import { OpenAiServiceTier, openAiNativeModels } from "@roo-code/types" +import { asyncStreamFrom, collectStream } from "../../../test-utils/stream" + +// The streaming entry point reaches the SDK client, so stub the client surface the +// handler actually calls. Without this the real SDK would attempt a network request. +const mockResponsesCreate = vitest.fn() + +vitest.mock("openai", () => ({ + __esModule: true, + default: vitest.fn().mockImplementation(function () { + return { + responses: { + create: mockResponsesCreate, + }, + } + }), +})) describe("OpenAiNativeHandler - normalizeUsage", () => { let handler: OpenAiNativeHandler @@ -440,6 +456,97 @@ describe("OpenAiNativeHandler - normalizeUsage", () => { }) }) + describe("OpenAiNativeHandler - prompt cache key", () => { + let handler: OpenAiNativeHandler + + beforeEach(() => { + handler = new OpenAiNativeHandler({ + openAiNativeApiKey: "test-key", + }) + }) + + const buildRequestBody = (modelId: string, taskId?: string) => { + handler["options"].apiModelId = modelId + const model = handler.getModel() + const metadata = taskId === undefined ? undefined : { taskId } + return handler["buildRequestBody"](model, [], "", model.verbosity, undefined, metadata) + } + + it("should send the task id as the body prompt_cache_key", () => { + const body = buildRequestBody("gpt-5.1", "task-abc") + expect(body.prompt_cache_key).toBe("task-abc") + }) + + it("should keep the key stable across turns of the same task", () => { + const firstTurn = buildRequestBody("gpt-5.1", "task-stable") + const secondTurn = buildRequestBody("gpt-5.1", "task-stable") + expect(firstTurn.prompt_cache_key).toBe("task-stable") + expect(secondTurn.prompt_cache_key).toBe(firstTurn.prompt_cache_key) + }) + + it("should give distinct keys to distinct tasks", () => { + const first = buildRequestBody("gpt-5.1", "task-one") + const second = buildRequestBody("gpt-5.1", "task-two") + expect(first.prompt_cache_key).toBe("task-one") + expect(second.prompt_cache_key).toBe("task-two") + expect(first.prompt_cache_key).not.toBe(second.prompt_cache_key) + }) + + it("should omit the key when no task id is supplied", () => { + // No task-scoped identity means no accumulating prefix, so the field is omitted + // rather than filled with a per-call random value that would override default routing. + const body = buildRequestBody("gpt-5.1") + expect(body).not.toHaveProperty("prompt_cache_key") + }) + + it("should omit the key when the model does not support prompt caching", () => { + const modelId = "codex-mini-latest" + expect(openAiNativeModels[modelId as keyof typeof openAiNativeModels].supportsPromptCache).toBe(false) + + const body = buildRequestBody(modelId, "task-abc") + expect(body.prompt_cache_key).toBeUndefined() + }) + + it("should omit the key for a non-caching model through the streaming entry point", async () => { + // The builder-level test above cannot see the streaming path: createMessage forwards + // metadata into buildRequestBody, so a key added anywhere along that path would still + // pass it. Drive the public entry point with a task id to pin the model gate there too. + const modelId = "codex-mini-latest" + expect(openAiNativeModels[modelId as keyof typeof openAiNativeModels].supportsPromptCache).toBe(false) + + mockResponsesCreate.mockClear() + mockResponsesCreate.mockResolvedValue(asyncStreamFrom([])) + + const streamingHandler = new OpenAiNativeHandler({ + openAiNativeApiKey: "test-key", + apiModelId: modelId, + }) + + await collectStream( + streamingHandler.createMessage("You are a helpful assistant.", [{ role: "user", content: "Hello!" }], { + taskId: "task-abc", + }), + ) + + expect(mockResponsesCreate).toHaveBeenCalledTimes(1) + const [body] = mockResponsesCreate.mock.calls[0] + expect(body).not.toHaveProperty("prompt_cache_key") + }) + + it("should clamp an over-length task id to OpenAI's 64-character limit", () => { + const longTaskId = "t".repeat(80) + const body = buildRequestBody("gpt-5.1", longTaskId) + expect(body.prompt_cache_key).toHaveLength(64) + expect(body.prompt_cache_key).toBe(longTaskId.slice(0, 64)) + }) + + it("should keep prompt_cache_retention alongside the key for 24h models", () => { + const body = buildRequestBody("gpt-5.1", "task-abc") + expect(body.prompt_cache_key).toBe("task-abc") + expect(body.prompt_cache_retention).toBe("24h") + }) + }) + describe("cost calculation", () => { it.each([ { tier: OpenAiServiceTier.Default, expectedCost: 4.525 }, diff --git a/src/api/providers/__tests__/openai-native.spec.ts b/src/api/providers/__tests__/openai-native.spec.ts index 8a086b48db..56aacc7ba5 100644 --- a/src/api/providers/__tests__/openai-native.spec.ts +++ b/src/api/providers/__tests__/openai-native.spec.ts @@ -534,6 +534,25 @@ describe("OpenAiNativeHandler", () => { expect(result).toBe("") }) + + it("should omit the prompt_cache_key because no task-scoped identity exists", async () => { + mockResponsesCreate.mockResolvedValue({ output: [] }) + + await handler.completePrompt("Test prompt") + + const [request] = mockResponsesCreate.mock.calls[0] + expect(request).not.toHaveProperty("prompt_cache_key") + }) + + it("should omit the prompt_cache_key when the model does not support prompt caching", async () => { + mockResponsesCreate.mockResolvedValue({ output: [] }) + handler = new OpenAiNativeHandler({ ...mockOptions, apiModelId: "codex-mini-latest" }) + + await handler.completePrompt("Test prompt") + + const [request] = mockResponsesCreate.mock.calls[0] + expect(request).not.toHaveProperty("prompt_cache_key") + }) }) describe("getModel", () => { @@ -1446,6 +1465,101 @@ describe("OpenAiNativeHandler", () => { expect(callBody.previous_response_id).toBeUndefined() }) + it("sends a task-scoped prompt_cache_key in the body and no session_id header", async () => { + const mockFetch = vitest.fn().mockResolvedValue({ + ok: true, + body: new ReadableStream({ + start(controller) { + controller.enqueue(new TextEncoder().encode("data: [DONE]\n\n")) + controller.close() + }, + }), + }) + global.fetch = mockFetch as typeof fetch + // Force the hand-built SSE path so the wire body and headers are inspectable. + mockResponsesCreate.mockRejectedValue(new Error("SDK not available")) + + const cacheHandler = new OpenAiNativeHandler({ ...mockOptions, apiModelId: "gpt-5.1" }) + + for (const taskId of ["task-cache-one", "task-cache-one", "task-cache-two"]) { + for await (const _chunk of cacheHandler.createMessage(systemPrompt, messages, { taskId })) { + // consume + } + } + + const bodies = mockFetch.mock.calls.map(([, init]) => JSON.parse(init.body)) + expect(bodies).toHaveLength(3) + expect(bodies[0].prompt_cache_key).toBe("task-cache-one") + // Same task, second turn: the key must not drift, or the growing prefix loses affinity. + expect(bodies[1].prompt_cache_key).toBe(bodies[0].prompt_cache_key) + // A different task must not share the bucket. + expect(bodies[2].prompt_cache_key).toBe("task-cache-two") + expect(bodies[2].prompt_cache_key).not.toBe(bodies[0].prompt_cache_key) + + // The retention control is untouched by this change. + expect(bodies[0].prompt_cache_retention).toBe("24h") + + // The affinity header is gone: OpenAI ignores it for caching and hardened + // gateways reject underscore header names outright. + for (const [, init] of mockFetch.mock.calls) { + expect(init.headers).not.toHaveProperty("session_id") + } + }) + + it("sends the task-scoped prompt_cache_key through the SDK path", async () => { + mockResponsesCreate.mockResolvedValue(asyncStreamFrom([])) + ;(OpenAI as unknown as ReturnType).mockClear() + + const cacheHandler = new OpenAiNativeHandler({ ...mockOptions, apiModelId: "gpt-5.1" }) + + // The second turn of the same task carries a longer conversation. Reusing one + // messages array would let a key derived from the conversation pass, so the two + // calls must differ in input while the key holds. + const followUp: Anthropic.Messages.MessageParam[] = [ + ...messages, + { role: "assistant", content: "Hello! How can I help?" }, + { role: "user", content: "Tell me about prompt caching." }, + ] + const turns: Array<{ taskId: string; messages: Anthropic.Messages.MessageParam[] }> = [ + { taskId: "sdk-task-one", messages }, + { taskId: "sdk-task-one", messages: followUp }, + { taskId: "sdk-task-two", messages }, + ] + + for (const turn of turns) { + await collectStream(cacheHandler.createMessage(systemPrompt, turn.messages, { taskId: turn.taskId })) + } + + const bodies = mockResponsesCreate.mock.calls.map(([body]) => body) + expect(bodies[0].prompt_cache_key).toBe("sdk-task-one") + expect(bodies[1].input).not.toEqual(bodies[0].input) + expect(bodies[1].prompt_cache_key).toBe(bodies[0].prompt_cache_key) + expect(bodies[2].prompt_cache_key).toBe("sdk-task-two") + expect(bodies[0].prompt_cache_retention).toBe("24h") + expect(mockResponsesCreate.mock.calls[0][1].headers).not.toHaveProperty("session_id") + + // Per-request headers are only half the surface: the SDK constructor also takes + // defaultHeaders, and a restored affinity header there would ship on every request + // while the assertion above still passed. Cache affinity is a body field, so neither + // the header name nor the body field name may appear in the constructor's defaults. + const constructorArgs = (OpenAI as unknown as ReturnType).mock.calls[0][0] + expect(constructorArgs.defaultHeaders).not.toHaveProperty("session_id") + expect(constructorArgs.defaultHeaders).not.toHaveProperty("prompt_cache_key") + }) + + it("omits the prompt_cache_key through createMessage when no metadata is supplied", async () => { + mockResponsesCreate.mockResolvedValue(asyncStreamFrom([])) + + const cacheHandler = new OpenAiNativeHandler({ ...mockOptions, apiModelId: "gpt-5.1" }) + + // No metadata at all: the streaming entry point must not synthesize a fallback + // task id, or the field would replace OpenAI's default routing for no affinity. + await collectStream(cacheHandler.createMessage(systemPrompt, messages)) + + const [body] = mockResponsesCreate.mock.calls[0] + expect(body).not.toHaveProperty("prompt_cache_key") + }) + it("should provide helpful error messages for different error codes", async () => { const testCases = [ { status: 400, expectedMessage: "Invalid request to Responses API" }, diff --git a/src/api/providers/openai-native.ts b/src/api/providers/openai-native.ts index fbbecf5760..e61080c4a2 100644 --- a/src/api/providers/openai-native.ts +++ b/src/api/providers/openai-native.ts @@ -1,5 +1,4 @@ import * as os from "os" -import { v7 as uuidv7 } from "uuid" import { Anthropic } from "@anthropic-ai/sdk" import OpenAI from "openai" @@ -34,12 +33,16 @@ import { sanitizeOpenAiCallId } from "../../utils/tool-id" export type OpenAiNativeModel = ReturnType +/** + * OpenAI caps `prompt_cache_key` at 64 characters; longer values are rejected with a 400. + * Task ids are uuidv7 today (36 chars), so this is a guard rather than a routine truncation. + */ +const PROMPT_CACHE_KEY_MAX_LENGTH = 64 + export class OpenAiNativeHandler extends BaseProvider implements SingleCompletionHandler { protected options: ApiHandlerOptions private client: OpenAI private readonly providerName = "OpenAI Native" - // Session ID for request tracking (persists for the lifetime of the handler) - private readonly sessionId: string /** * Some Responses streams emit tool-call argument deltas without stable call id/name. * Track the last observed tool identity from output_item events so we can still @@ -88,22 +91,22 @@ export class OpenAiNativeHandler extends BaseProvider implements SingleCompletio constructor(options: ApiHandlerOptions) { super() this.options = options - // Generate a session ID for request tracking - this.sessionId = uuidv7() // Default to including reasoning.summary: "auto" for models that support Responses API // reasoning summaries unless explicitly disabled. if (this.options.enableResponsesReasoningSummary === undefined) { this.options.enableResponsesReasoningSummary = true } const apiKey = this.options.openAiNativeApiKey ?? NOT_PROVIDED - // Include originator, session_id, and User-Agent headers for API tracking and debugging + // Include originator and User-Agent headers for API tracking and debugging. + // Cache affinity rides on the body `prompt_cache_key`, not on a header: OpenAI documents + // the `session_id` header as ignored on the public Responses API, and underscore header + // names are rejected outright by hardened gateways (Envoy/nginx `underscores_in_headers`). const userAgent = `zoo-code/${Package.version} (${os.platform()} ${os.release()}; ${os.arch()}) node/${process.version.slice(1)}` this.client = new OpenAI({ baseURL: this.options.openAiNativeBaseUrl || undefined, apiKey, defaultHeaders: { originator: "zoo-code", - session_id: this.sessionId, "User-Agent": userAgent, }, timeout: this.timeoutMs, @@ -330,6 +333,11 @@ export class OpenAiNativeHandler extends BaseProvider implements SingleCompletio include?: string[] /** Prompt cache retention policy: "in_memory" (default) or "24h" for extended caching */ prompt_cache_retention?: "in_memory" | "24h" + /** + * Cache routing affinity key. OpenAI documents this body field as the routing lever; + * the `session_id` request header is not part of the prompt-caching API. + */ + prompt_cache_key?: string tools?: Array<{ type: "function" name: string @@ -347,6 +355,10 @@ export class OpenAiNativeHandler extends BaseProvider implements SingleCompletio // Decide whether to enable extended prompt cache retention for this request const promptCacheRetention = this.getPromptCacheRetention(model) + // Cache routing affinity: a stable, task-scoped key lets OpenAI route every turn of the + // same task to the machine holding its cached prefix. + const promptCacheKey = this.getPromptCacheKey(model, metadata) + const body: ResponsesRequestBody = { model: model.id, input: formattedInput, @@ -379,6 +391,8 @@ export class OpenAiNativeHandler extends BaseProvider implements SingleCompletio // Enable extended prompt cache retention for models that support it. // This uses the OpenAI Responses API `prompt_cache_retention` parameter. ...(promptCacheRetention ? { prompt_cache_retention: promptCacheRetention } : {}), + // Cache routing affinity key (body field, not the inert `session_id` header). + ...(promptCacheKey ? { prompt_cache_key: promptCacheKey } : {}), tools: (metadata?.tools ?? []) .filter((tool) => tool.type === "function") .map((tool) => { @@ -418,12 +432,10 @@ export class OpenAiNativeHandler extends BaseProvider implements SingleCompletio // Create AbortController for cancellation this.abortController = new AbortController() - // Build per-request headers using taskId when available, falling back to sessionId - const taskId = metadata?.taskId + // Build per-request headers for API tracking and debugging const userAgent = `zoo-code/${Package.version} (${os.platform()} ${os.release()}; ${os.arch()}) node/${process.version.slice(1)}` const requestHeaders: Record = { originator: "zoo-code", - session_id: taskId || this.sessionId, "User-Agent": userAgent, } @@ -557,7 +569,7 @@ export class OpenAiNativeHandler extends BaseProvider implements SingleCompletio private async *makeResponsesApiRequest( requestBody: any, model: OpenAiNativeModel, - metadata?: ApiHandlerCreateMessageMetadata, + _metadata?: ApiHandlerCreateMessageMetadata, _systemPrompt?: string, _messages?: Anthropic.Messages.MessageParam[], ): ApiStream { @@ -568,8 +580,7 @@ export class OpenAiNativeHandler extends BaseProvider implements SingleCompletio // Create AbortController for cancellation this.abortController = new AbortController() - // Build per-request headers using taskId when available, falling back to sessionId - const taskId = metadata?.taskId + // Build per-request headers for API tracking and debugging const userAgent = `zoo-code/${Package.version} (${os.platform()} ${os.release()}; ${os.arch()}) node/${process.version.slice(1)}` try { @@ -579,7 +590,6 @@ export class OpenAiNativeHandler extends BaseProvider implements SingleCompletio "Content-Type": "application/json", Authorization: `Bearer ${apiKey}`, originator: "zoo-code", - session_id: taskId || this.sessionId, "User-Agent": userAgent, }, body: JSON.stringify(requestBody), @@ -1420,6 +1430,34 @@ export class OpenAiNativeHandler extends BaseProvider implements SingleCompletio return undefined } + /** + * Returns the prompt cache routing key for this request, if one applies. + * + * OpenAI keys cache affinity off the `prompt_cache_key` body field; the `session_id` + * request header is not part of the prompt-caching API. The key must stay identical + * across every turn of one task (so its growing prefix keeps landing on the same + * machine) and differ between tasks (so unrelated prefixes do not compete for one + * routing bucket). `metadata.taskId` is the longest-lived identifier the handler + * receives: it is the persisted task id, stable across turns, restarts and + * rehydration. + * + * The key is sent ONLY when a task-scoped identity exists. With no task id there is + * no conversation whose prefix could accumulate, so a per-call random key would buy + * no affinity while replacing whatever routing OpenAI applies by default. The + * fail-safe is therefore to omit the field and leave default routing untouched. + */ + private getPromptCacheKey( + model: OpenAiNativeModel, + metadata?: ApiHandlerCreateMessageMetadata, + ): string | undefined { + if (!model.info.supportsPromptCache) return undefined + + const taskId = metadata?.taskId + if (!taskId) return undefined + + return taskId.length > PROMPT_CACHE_KEY_MAX_LENGTH ? taskId.slice(0, PROMPT_CACHE_KEY_MAX_LENGTH) : taskId + } + /** * Returns a shallow-cloned ModelInfo with pricing overridden for the given tier, if available. * If no tier or no overrides exist, the original ModelInfo is returned. @@ -1531,6 +1569,12 @@ export class OpenAiNativeHandler extends BaseProvider implements SingleCompletio requestBody[SERVICE_TIER_KEY] = serviceTier } + // No prompt_cache_key here on purpose. `completePrompt` carries no task id + // (`CompletePromptOptions` has only abortSignal/timeoutMs) and its callers build a + // throwaway handler per call, so there is no conversation whose prefix could + // accumulate. A per-call random key would buy no affinity while overriding + // OpenAI's default routing, so the field is omitted. + // Add reasoning if supported if (reasoningEffort) { requestBody.reasoning = {