diff --git a/packages/agent/src/agent.ts b/packages/agent/src/agent.ts index 147edc8396d..a6352f0bcd4 100644 --- a/packages/agent/src/agent.ts +++ b/packages/agent/src/agent.ts @@ -114,6 +114,7 @@ export interface AgentOptions { steeringMode?: QueueMode; followUpMode?: QueueMode; sessionId?: string; + promptCacheKey?: string; thinkingBudgets?: ThinkingBudgets; transport?: Transport; maxRetryDelayMs?: number; @@ -198,6 +199,8 @@ export class Agent { private activeRun?: ActiveRun; /** Session identifier forwarded to providers for cache-aware backends. */ public sessionId?: string; + /** Optional override for the provider prompt cache key, replacing `sessionId` as the cache routing key. */ + public promptCacheKey?: string; /** Optional per-level thinking token budgets forwarded to the stream function. */ public thinkingBudgets?: ThinkingBudgets; /** Preferred transport forwarded to the stream function. */ @@ -224,6 +227,7 @@ export class Agent { this.steeringQueue = new PendingMessageQueue(runtimeOptions.steeringMode ?? "one-at-a-time"); this.followUpQueue = new PendingMessageQueue(runtimeOptions.followUpMode ?? "one-at-a-time"); this.sessionId = runtimeOptions.sessionId; + this.promptCacheKey = runtimeOptions.promptCacheKey; this.thinkingBudgets = runtimeOptions.thinkingBudgets; this.transport = runtimeOptions.transport ?? "auto"; this.maxRetryDelayMs = runtimeOptions.maxRetryDelayMs; @@ -437,6 +441,7 @@ export class Agent { model: this._state.model, reasoning: this._state.thinkingLevel === "off" ? undefined : this._state.thinkingLevel, sessionId: this.sessionId, + promptCacheKey: this.promptCacheKey, onPayload: this.onPayload, onResponse: this.onResponse, transport: this.transport, diff --git a/packages/agent/src/harness/agent-harness.ts b/packages/agent/src/harness/agent-harness.ts index cbc2d0a4adb..1378428331b 100644 --- a/packages/agent/src/harness/agent-harness.ts +++ b/packages/agent/src/harness/agent-harness.ts @@ -103,6 +103,7 @@ function applyStreamOptionsPatch( if (Object.hasOwn(patch, "maxRetries")) result.maxRetries = patch.maxRetries; if (Object.hasOwn(patch, "maxRetryDelayMs")) result.maxRetryDelayMs = patch.maxRetryDelayMs; if (Object.hasOwn(patch, "cacheRetention")) result.cacheRetention = patch.cacheRetention; + if (Object.hasOwn(patch, "promptCacheKey")) result.promptCacheKey = patch.promptCacheKey; if (Object.hasOwn(patch, "headers")) { if (patch.headers === undefined) { @@ -421,6 +422,7 @@ export class AgentHarness< reasoning: streamOptions?.reasoning, signal: streamOptions?.signal, sessionId: turnState.sessionId, + promptCacheKey: requestOptions.promptCacheKey, timeoutMs: requestOptions.timeoutMs, transport: requestOptions.transport, }); diff --git a/packages/agent/src/harness/types.ts b/packages/agent/src/harness/types.ts index 7a217ae349a..1c0ffe76e6b 100644 --- a/packages/agent/src/harness/types.ts +++ b/packages/agent/src/harness/types.ts @@ -132,6 +132,8 @@ export interface AgentHarnessStreamOptions { metadata?: SimpleStreamOptions["metadata"]; /** Provider cache retention hint. */ cacheRetention?: SimpleStreamOptions["cacheRetention"]; + /** Optional override for the provider prompt cache key, replacing the session id as the cache routing key. */ + promptCacheKey?: string; } /** Per-request stream option patch returned by provider hooks. */ diff --git a/packages/agent/src/proxy.ts b/packages/agent/src/proxy.ts index 5f0925c9148..a41b3a1377c 100644 --- a/packages/agent/src/proxy.ts +++ b/packages/agent/src/proxy.ts @@ -63,6 +63,7 @@ type ProxySerializableStreamOptions = Pick< | "reasoning" | "cacheRetention" | "sessionId" + | "promptCacheKey" | "headers" | "metadata" | "transport" @@ -105,6 +106,7 @@ function buildProxyRequestOptions(options: ProxyStreamOptions): ProxySerializabl reasoning: options.reasoning, cacheRetention: options.cacheRetention, sessionId: options.sessionId, + promptCacheKey: options.promptCacheKey, headers: options.headers, metadata: options.metadata, transport: options.transport, diff --git a/packages/ai/src/api/azure-openai-responses.ts b/packages/ai/src/api/azure-openai-responses.ts index 1e409650a8f..769f17aa4ad 100644 --- a/packages/ai/src/api/azure-openai-responses.ts +++ b/packages/ai/src/api/azure-openai-responses.ts @@ -281,7 +281,7 @@ function buildParams( model: deploymentName, input: messages, stream: true, - prompt_cache_key: clampOpenAIPromptCacheKey(options?.sessionId), + prompt_cache_key: clampOpenAIPromptCacheKey(options?.promptCacheKey ?? options?.sessionId), store: false, }; diff --git a/packages/ai/src/api/openai-codex-responses.ts b/packages/ai/src/api/openai-codex-responses.ts index cf89ee03430..dd192f6869c 100644 --- a/packages/ai/src/api/openai-codex-responses.ts +++ b/packages/ai/src/api/openai-codex-responses.ts @@ -269,7 +269,11 @@ export const stream: StreamFunction<"openai-codex-responses", OpenAICodexRespons ); const cacheSessionId = options?.cacheRetention === "none" ? undefined : options?.sessionId; const codexSessionId = clampOpenAIPromptCacheKey(cacheSessionId); - let body = buildRequestBody(model, context, options, codexSessionId, grammarToolInputProperties); + const promptCacheKey = + options?.cacheRetention === "none" + ? undefined + : clampOpenAIPromptCacheKey(options?.promptCacheKey ?? options?.sessionId); + let body = buildRequestBody(model, context, options, promptCacheKey, grammarToolInputProperties); const nextBody = await options?.onPayload?.(body, model); if (nextBody !== undefined) { body = nextBody as RequestBody; @@ -517,7 +521,7 @@ function buildRequestBody( model: Model<"openai-codex-responses">, context: Context, options: OpenAICodexResponsesOptions | undefined, - cacheSessionId: string | undefined, + promptCacheKey: string | undefined, grammarToolInputProperties: ReadonlyMap = createGrammarToolInputProperties( context.tools, model.compat?.supportsOpenAIGrammarTools ?? false, @@ -545,7 +549,7 @@ function buildRequestBody( input: messages, text: { verbosity: options?.textVerbosity || "low" }, include: ["reasoning.encrypted_content"], - prompt_cache_key: cacheSessionId, + prompt_cache_key: promptCacheKey, tool_choice: options?.toolChoice ?? "auto", parallel_tool_calls: true, }; diff --git a/packages/ai/src/api/openai-completions.ts b/packages/ai/src/api/openai-completions.ts index 29bb064f9bb..d9f57dd62b6 100644 --- a/packages/ai/src/api/openai-completions.ts +++ b/packages/ai/src/api/openai-completions.ts @@ -689,7 +689,7 @@ function buildParams( prompt_cache_key: (model.baseUrl.includes("api.openai.com") && cacheRetention !== "none") || (cacheRetention === "long" && compat.supportsLongCacheRetention) - ? clampOpenAIPromptCacheKey(options?.sessionId) + ? clampOpenAIPromptCacheKey(options?.promptCacheKey ?? options?.sessionId) : undefined, prompt_cache_retention: cacheRetention === "long" && compat.supportsLongCacheRetention ? "24h" : undefined, }; diff --git a/packages/ai/src/api/openai-responses.ts b/packages/ai/src/api/openai-responses.ts index 59f92a25cd8..ec3610fd396 100644 --- a/packages/ai/src/api/openai-responses.ts +++ b/packages/ai/src/api/openai-responses.ts @@ -275,7 +275,10 @@ function buildParams( model: model.id, input: messages, stream: true, - prompt_cache_key: cacheRetention === "none" ? undefined : clampOpenAIPromptCacheKey(options?.sessionId), + prompt_cache_key: + cacheRetention === "none" + ? undefined + : clampOpenAIPromptCacheKey(options?.promptCacheKey ?? options?.sessionId), prompt_cache_retention: getPromptCacheRetention(compat, cacheRetention), prompt_cache_options: disableImplicitPromptCache ? { mode: "explicit" } : undefined, store: false, diff --git a/packages/ai/src/api/simple-options.ts b/packages/ai/src/api/simple-options.ts index 55bfe8df8e1..0162f9d7284 100644 --- a/packages/ai/src/api/simple-options.ts +++ b/packages/ai/src/api/simple-options.ts @@ -32,6 +32,7 @@ export function buildBaseOptions( transport: options?.transport, cacheRetention: options?.cacheRetention, sessionId: options?.sessionId, + promptCacheKey: options?.promptCacheKey, headers: options?.headers, onPayload: options?.onPayload, onResponse: options?.onResponse, diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index 1d8fcfd3af2..3c99d1ea12f 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -133,6 +133,15 @@ export interface StreamOptions { * session-aware features. Ignored by providers that don't support it. */ sessionId?: string; + /** + * Optional override for the provider prompt cache key, which defaults to + * `sessionId`. Set one shared key on sessions with an identical prompt prefix + * so their requests route to the same cache and hit each other's cached + * prefix. OpenAI recommends roughly 15 requests per minute per key, so + * callers shard the key when concurrency exceeds that. Session-affinity + * headers remain tied to `sessionId`. + */ + promptCacheKey?: string; /** * Optional callback for inspecting or replacing provider payloads before sending. * Return undefined to keep the payload unchanged. diff --git a/packages/ai/test/azure-openai-base-url.test.ts b/packages/ai/test/azure-openai-base-url.test.ts index 12e3fb25e6a..2ebbc1a505e 100644 --- a/packages/ai/test/azure-openai-base-url.test.ts +++ b/packages/ai/test/azure-openai-base-url.test.ts @@ -157,6 +157,18 @@ describe("azure-openai-responses base URL normalization", () => { expect(azureMock.lastParams?.prompt_cache_key).toBe("x".repeat(64)); }); + it("prefers promptCacheKey over sessionId for prompt_cache_key", async () => { + const model = getModel("azure-openai-responses", "gpt-4o-mini"); + await streamAzureOpenAIResponses(model, context, { + apiKey: "test-api-key", + azureBaseUrl: "https://my-resource.openai.azure.com", + sessionId: "session-azure", + promptCacheKey: "shared-key", + }).result(); + + expect(azureMock.lastParams?.prompt_cache_key).toBe("shared-key"); + }); + it("disables server-side response storage", async () => { const model = getModel("azure-openai-responses", "gpt-4o-mini"); await streamAzureOpenAIResponses(model, context, { diff --git a/packages/ai/test/cache-retention.test.ts b/packages/ai/test/cache-retention.test.ts index 5721dd3797f..794b06f6023 100644 --- a/packages/ai/test/cache-retention.test.ts +++ b/packages/ai/test/cache-retention.test.ts @@ -423,6 +423,31 @@ describe("Cache Retention (PI_CACHE_RETENTION)", () => { expect(capturedPayload.prompt_cache_key).toBe("session-2"); expect(capturedPayload.prompt_cache_retention).toBe("24h"); }); + + it("should prefer promptCacheKey over sessionId for prompt_cache_key", async () => { + const model = getModel("openai", "gpt-4o-mini"); + let capturedPayload: any = null; + + try { + const s = streamOpenAIResponses(model, context, { + apiKey: "fake-key", + sessionId: "session-3", + promptCacheKey: "shared-key", + onPayload: stopAfterPayload((payload) => { + capturedPayload = payload; + }), + }); + + for await (const event of s) { + if (event.type === "error") break; + } + } catch { + // Expected to fail + } + + expect(capturedPayload).not.toBeNull(); + expect(capturedPayload.prompt_cache_key).toBe("shared-key"); + }); }); describe("OpenAI Completions Provider", () => { diff --git a/packages/ai/test/openai-codex-stream.test.ts b/packages/ai/test/openai-codex-stream.test.ts index 582ec9a1c5a..e4268daa42b 100644 --- a/packages/ai/test/openai-codex-stream.test.ts +++ b/packages/ai/test/openai-codex-stream.test.ts @@ -739,6 +739,56 @@ describe("openai-codex streaming", () => { expect(capturedHeaders?.get("x-client-request-id")).toBe("x".repeat(64)); }); + it("prefers promptCacheKey over sessionId for prompt_cache_key", async () => { + const token = mockToken(); + let capturedPayload: { prompt_cache_key?: string } | undefined; + const encoder = new TextEncoder(); + vi.stubGlobal( + "fetch", + vi.fn( + async () => + new Response( + new ReadableStream({ + start(controller) { + controller.enqueue(encoder.encode(buildSSEPayload({ status: "completed" }))); + controller.close(); + }, + }), + { status: 200, headers: { "content-type": "text/event-stream" } }, + ), + ), + ); + + const model: Model<"openai-codex-responses"> = { + id: "gpt-5.1-codex", + name: "GPT-5.1 Codex", + api: "openai-codex-responses", + provider: "openai-codex", + baseUrl: "https://chatgpt.com/backend-api", + reasoning: true, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 400000, + maxTokens: 128000, + }; + const context: Context = { + systemPrompt: "You are a helpful assistant.", + messages: [{ role: "user", content: "Say hello", timestamp: Date.now() }], + }; + + await streamOpenAICodexResponses(model, context, { + apiKey: token, + transport: "sse", + sessionId: "session-codex", + promptCacheKey: "shared-key", + onPayload: (payload) => { + capturedPayload = payload as { prompt_cache_key?: string }; + }, + }).result(); + + expect(capturedPayload?.prompt_cache_key).toBe("shared-key"); + }); + it("preserves gpt-5.5 xhigh reasoning effort from simple options", async () => { const tempDir = mkdtempSync(join(tmpdir(), "pi-codex-stream-")); process.env.PI_CODING_AGENT_DIR = tempDir; diff --git a/packages/ai/test/openai-completions-prompt-cache.test.ts b/packages/ai/test/openai-completions-prompt-cache.test.ts index b697ff0f457..5c67514dedc 100644 --- a/packages/ai/test/openai-completions-prompt-cache.test.ts +++ b/packages/ai/test/openai-completions-prompt-cache.test.ts @@ -93,6 +93,7 @@ describe("openai-completions prompt caching", () => { options?: { cacheRetention?: "none" | "short" | "long"; sessionId?: string; + promptCacheKey?: string; headers?: Record; }, model: Model<"openai-completions"> = createModel(), @@ -133,6 +134,23 @@ describe("openai-completions prompt caching", () => { expect(payload?.prompt_cache_key).toBe("x".repeat(64)); }); + it("prefers promptCacheKey over sessionId and clamps it", async () => { + const promptCacheKey = "k".repeat(67); + const { payload } = await captureRequest({ sessionId: "session-123", promptCacheKey }); + + expect(payload?.prompt_cache_key).toBe("k".repeat(64)); + }); + + it("omits prompt_cache_key when cacheRetention is none even with promptCacheKey set", async () => { + const { payload } = await captureRequest({ + cacheRetention: "none", + sessionId: "session-123", + promptCacheKey: "shared-key", + }); + + expect(payload?.prompt_cache_key).toBeUndefined(); + }); + it("omits prompt cache fields when cacheRetention is none", async () => { const { payload } = await captureRequest({ cacheRetention: "none", sessionId: "session-789" });