diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index acc47805..3a8a2c9a 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Fixed + +- Fixed direct OpenAI Responses requests to send aligned `prompt_cache_key`, `session_id`, and `x-client-request-id` values when `sessionId` is provided, improving prompt cache affinity for append-only sessions ([#3018](https://github.com/badlogic/pi-mono/pull/3018) by [@steipete](https://github.com/steipete)) + ## [0.67.1] - 2026-04-13 ## [0.67.0] - 2026-04-13 diff --git a/packages/ai/src/providers/openai-codex-responses.ts b/packages/ai/src/providers/openai-codex-responses.ts index fd8ba1f0..4cca39b7 100644 --- a/packages/ai/src/providers/openai-codex-responses.ts +++ b/packages/ai/src/providers/openai-codex-responses.ts @@ -950,6 +950,7 @@ function buildSSEHeaders( if (sessionId) { headers.set("session_id", sessionId); + headers.set("x-client-request-id", sessionId); } return headers; diff --git a/packages/ai/src/providers/openai-responses.ts b/packages/ai/src/providers/openai-responses.ts index 26abf50e..35bfeffe 100644 --- a/packages/ai/src/providers/openai-responses.ts +++ b/packages/ai/src/providers/openai-responses.ts @@ -88,7 +88,9 @@ export const streamOpenAIResponses: StreamFunction<"openai-responses", OpenAIRes try { // Create OpenAI client const apiKey = options?.apiKey || getEnvApiKey(model.provider) || ""; - const client = createClient(model, context, apiKey, options?.headers); + const cacheRetention = resolveCacheRetention(options?.cacheRetention); + const cacheSessionId = cacheRetention === "none" ? undefined : options?.sessionId; + const client = createClient(model, context, apiKey, options?.headers, cacheSessionId); let params = buildParams(model, context, options); const nextParams = await options?.onPayload?.(params, model); if (nextParams !== undefined) { @@ -155,6 +157,7 @@ function createClient( context: Context, apiKey?: string, optionsHeaders?: Record, + sessionId?: string, ) { if (!apiKey) { if (!process.env.OPENAI_API_KEY) { @@ -175,6 +178,11 @@ function createClient( Object.assign(headers, copilotHeaders); } + if (sessionId && model.provider === "openai" && model.baseUrl.includes("api.openai.com")) { + headers.session_id = sessionId; + headers["x-client-request-id"] = sessionId; + } + // Merge options headers last so they can override defaults if (optionsHeaders) { Object.assign(headers, optionsHeaders); diff --git a/packages/ai/test/openai-codex-cache-affinity-e2e.test.ts b/packages/ai/test/openai-codex-cache-affinity-e2e.test.ts new file mode 100644 index 00000000..cf0ddf39 --- /dev/null +++ b/packages/ai/test/openai-codex-cache-affinity-e2e.test.ts @@ -0,0 +1,36 @@ +import { describe, expect, it } from "vitest"; +import { getModel } from "../src/models.js"; +import { complete } from "../src/stream.js"; +import type { Context } from "../src/types.js"; +import { resolveApiKey } from "./oauth.js"; + +const codexToken = await resolveApiKey("openai-codex"); + +describe("openai-codex cache affinity e2e", () => { + it.skipIf(!codexToken)("handles SSE requests with aligned cache-affinity identifiers", async () => { + const model = getModel("openai-codex", "gpt-5.3-codex"); + const sessionId = "0195d6e4-4cf9-7f44-a2d8-f8f7f49ee9d3"; + const context: Context = { + systemPrompt: "You are a helpful assistant. Reply exactly as requested.", + messages: [ + { + role: "user", + content: "Reply with exactly: cache affinity e2e success", + timestamp: Date.now(), + }, + ], + }; + + const response = await complete(model, context, { + apiKey: codexToken, + sessionId, + transport: "sse", + }); + + expect(response.stopReason, response.errorMessage).not.toBe("error"); + expect(response.errorMessage).toBeUndefined(); + expect(response.content.map((block) => (block.type === "text" ? block.text : "")).join("")).toContain( + "cache affinity e2e success", + ); + }); +}); diff --git a/packages/ai/test/openai-codex-stream.test.ts b/packages/ai/test/openai-codex-stream.test.ts index d3db8c29..58496988 100644 --- a/packages/ai/test/openai-codex-stream.test.ts +++ b/packages/ai/test/openai-codex-stream.test.ts @@ -303,7 +303,7 @@ describe("openai-codex streaming", () => { expect(result.stopReason).toBe("length"); }); - it("sets conversation_id/session_id headers and prompt_cache_key when sessionId is provided", async () => { + it("sets session_id/x-client-request-id headers and prompt_cache_key when sessionId is provided", async () => { const tempDir = mkdtempSync(join(tmpdir(), "pi-codex-stream-")); process.env.PI_CODING_AGENT_DIR = tempDir; @@ -364,13 +364,12 @@ describe("openai-codex streaming", () => { if (url === "https://chatgpt.com/backend-api/codex/responses") { const headers = init?.headers instanceof Headers ? init.headers : undefined; // Verify sessionId is set in headers - expect(headers?.get("conversation_id")).toBe(sessionId); expect(headers?.get("session_id")).toBe(sessionId); + expect(headers?.get("x-client-request-id")).toBe(sessionId); // Verify sessionId is set in request body as prompt_cache_key const body = typeof init?.body === "string" ? (JSON.parse(init.body) as Record) : null; expect(body?.prompt_cache_key).toBe(sessionId); - expect(body?.prompt_cache_retention).toBe("in-memory"); return new Response(stream, { status: 200, @@ -500,7 +499,7 @@ describe("openai-codex streaming", () => { await streamResult.result(); }); - it("does not set conversation_id/session_id headers when sessionId is not provided", async () => { + it("does not set session_id/x-client-request-id headers when sessionId is not provided", async () => { const tempDir = mkdtempSync(join(tmpdir(), "pi-codex-stream-")); process.env.PI_CODING_AGENT_DIR = tempDir; @@ -560,8 +559,8 @@ describe("openai-codex streaming", () => { if (url === "https://chatgpt.com/backend-api/codex/responses") { const headers = init?.headers instanceof Headers ? init.headers : undefined; // Verify headers are not set when sessionId is not provided - expect(headers?.has("conversation_id")).toBe(false); expect(headers?.has("session_id")).toBe(false); + expect(headers?.has("x-client-request-id")).toBe(false); return new Response(stream, { status: 200, diff --git a/packages/ai/test/openai-responses-cache-affinity-e2e.test.ts b/packages/ai/test/openai-responses-cache-affinity-e2e.test.ts new file mode 100644 index 00000000..a2fd0d0e --- /dev/null +++ b/packages/ai/test/openai-responses-cache-affinity-e2e.test.ts @@ -0,0 +1,32 @@ +import { describe, expect, it } from "vitest"; +import { getModel } from "../src/models.js"; +import { complete } from "../src/stream.js"; +import type { Context } from "../src/types.js"; + +describe.skipIf(!process.env.OPENAI_API_KEY)("openai responses cache affinity e2e", () => { + it("handles direct OpenAI Responses requests with aligned cache-affinity identifiers", { retry: 2 }, async () => { + const model = getModel("openai", "gpt-5.4"); + const sessionId = "0195d6e4-4cf9-7f44-a2d8-f8f7f49ee9d3"; + const context: Context = { + systemPrompt: "You are a helpful assistant. Reply exactly as requested.", + messages: [ + { + role: "user", + content: "Reply with exactly: openai cache affinity e2e success", + timestamp: Date.now(), + }, + ], + }; + + const response = await complete(model, context, { + apiKey: process.env.OPENAI_API_KEY!, + sessionId, + }); + + expect(response.stopReason, response.errorMessage).not.toBe("error"); + expect(response.errorMessage).toBeUndefined(); + expect(response.content.map((block) => (block.type === "text" ? block.text : "")).join("")).toContain( + "openai cache affinity e2e success", + ); + }); +}); diff --git a/packages/ai/test/openai-responses-copilot-provider.test.ts b/packages/ai/test/openai-responses-copilot-provider.test.ts index 04b2856c..df1ccb8a 100644 --- a/packages/ai/test/openai-responses-copilot-provider.test.ts +++ b/packages/ai/test/openai-responses-copilot-provider.test.ts @@ -2,7 +2,54 @@ import { afterEach, describe, expect, it, vi } from "vitest"; import { getModel } from "../src/models.js"; import { streamOpenAIResponses } from "../src/providers/openai-responses.js"; -describe("openai-responses github-copilot defaults", () => { +type CapturedHeaders = Headers | string[][] | Record | undefined; + +function getHeader(headers: CapturedHeaders, name: string): string | null { + if (!headers) return null; + if (headers instanceof Headers) return headers.get(name); + + const lowerName = name.toLowerCase(); + if (Array.isArray(headers)) { + const match = headers.find(([key]) => key?.toLowerCase() === lowerName); + return match?.[1] ?? null; + } + + for (const [key, value] of Object.entries(headers)) { + if (key.toLowerCase() === lowerName) return typeof value === "string" ? value : value.join(", "); + } + return null; +} + +async function captureOpenAIResponseHeaders( + options: Parameters[2], +): Promise<{ sessionId: string | null; clientRequestId: string | null }> { + const captured = { sessionId: null as string | null, clientRequestId: null as string | null }; + vi.spyOn(globalThis, "fetch").mockImplementation(async (_input, init) => { + captured.sessionId = getHeader(init?.headers, "session_id"); + captured.clientRequestId = getHeader(init?.headers, "x-client-request-id"); + return new Response("data: [DONE]\n\n", { + status: 200, + headers: { "content-type": "text/event-stream" }, + }); + }); + + const stream = streamOpenAIResponses( + getModel("openai", "gpt-5.4"), + { + systemPrompt: "sys", + messages: [{ role: "user", content: "hi", timestamp: Date.now() }], + }, + { apiKey: "test-key", ...options }, + ); + + for await (const event of stream) { + if (event.type === "done" || event.type === "error") break; + } + + return captured; +} + +describe("openai-responses provider defaults", () => { afterEach(() => { vi.restoreAllMocks(); }); @@ -41,4 +88,28 @@ describe("openai-responses github-copilot defaults", () => { reasoning: expect.anything(), }); }); + + it("sets cache-affinity headers for official OpenAI Responses requests with a sessionId", async () => { + const captured = await captureOpenAIResponseHeaders({ sessionId: "session-123" }); + + expect(captured).toEqual({ sessionId: "session-123", clientRequestId: "session-123" }); + }); + + it("lets explicit headers override the default OpenAI cache-affinity headers", async () => { + const captured = await captureOpenAIResponseHeaders({ + sessionId: "session-123", + headers: { + session_id: "override-session", + "x-client-request-id": "override-request", + }, + }); + + expect(captured).toEqual({ sessionId: "override-session", clientRequestId: "override-request" }); + }); + + it("omits OpenAI cache-affinity headers when cacheRetention is none", async () => { + const captured = await captureOpenAIResponseHeaders({ cacheRetention: "none", sessionId: "session-123" }); + + expect(captured).toEqual({ sessionId: null, clientRequestId: null }); + }); }); diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 5b5775af..651a64e9 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -55,6 +55,7 @@ How to disable it: - Fixed Gemma 4 thinking level mapping to route between `MINIMAL` and `HIGH`, and map Pi reasoning levels to the model's supported thinking levels ([#2903](https://github.com/badlogic/pi-mono/pull/2903) by [@aadishv](https://github.com/aadishv)) - Fixed Gemini 2.5 Flash Lite minimal thinking budget to use the model's supported 512-token minimum instead of the regular Flash 128-token minimum, avoiding invalid thinking budget errors ([#2861](https://github.com/badlogic/pi-mono/pull/2861) by [@JasonOA888](https://github.com/JasonOA888)) - Fixed OpenAI Codex Responses requests to forward configured `serviceTier` values, restoring service-tier selection for Codex sessions ([#2996](https://github.com/badlogic/pi-mono/pull/2996) by [@markusylisiurunen](https://github.com/markusylisiurunen)) +- Fixed newly generated session IDs to use UUIDv7, improving time locality for session-based request routing ([#3018](https://github.com/badlogic/pi-mono/pull/3018) by [@steipete](https://github.com/steipete)) - Fixed `Container.render()` stack overflow on long sessions by replacing `Array.push(...spread)` with a loop-based push, preventing `RangeError: Maximum call stack size exceeded` when child output exceeds the V8 call stack argument limit ([#2651](https://github.com/badlogic/pi-mono/issues/2651)) - Fixed editor sticky-column tracking around paste markers so vertical cursor navigation restores the column from before the cursor entered a paste marker instead of jumping inside or past pasted content ([#3092](https://github.com/badlogic/pi-mono/pull/3092) by [@Perlence](https://github.com/Perlence)) - Fixed queued messages typed during `/tree` branch summarization to flush automatically after navigation completes, so they no longer remain stuck in the steering queue ([#3091](https://github.com/badlogic/pi-mono/pull/3091) by [@Perlence](https://github.com/Perlence)) diff --git a/packages/coding-agent/package.json b/packages/coding-agent/package.json index 4b701a4a..ff9db2e0 100644 --- a/packages/coding-agent/package.json +++ b/packages/coding-agent/package.json @@ -58,6 +58,7 @@ "proper-lockfile": "^4.1.2", "strip-ansi": "^7.1.0", "undici": "^7.19.1", + "uuid": "^11.1.0", "yaml": "^2.8.2" }, "overrides": { diff --git a/packages/coding-agent/src/core/session-manager.ts b/packages/coding-agent/src/core/session-manager.ts index 926f7809..f366274b 100644 --- a/packages/coding-agent/src/core/session-manager.ts +++ b/packages/coding-agent/src/core/session-manager.ts @@ -15,6 +15,7 @@ import { } from "fs"; import { readdir, readFile, stat } from "fs/promises"; import { join, resolve } from "path"; +import { v7 as uuidv7 } from "uuid"; import { getAgentDir as getDefaultAgentDir, getSessionsDir } from "../config.js"; import { type BashExecutionMessage, @@ -197,6 +198,10 @@ export type ReadonlySessionManager = Pick< | "getSessionName" >; +function createSessionId(): string { + return uuidv7(); +} + /** Generate a unique short ID (8 hex chars, collision-checked) */ function generateId(byId: { has(id: string): boolean }): string { for (let i = 0; i < 100; i++) { @@ -707,7 +712,7 @@ export class SessionManager { } const header = this.fileEntries.find((e) => e.type === "session") as SessionHeader | undefined; - this.sessionId = header?.id ?? randomUUID(); + this.sessionId = header?.id ?? createSessionId(); if (migrateToCurrentVersion(this.fileEntries)) { this._rewriteFile(); @@ -723,7 +728,7 @@ export class SessionManager { } newSession(options?: NewSessionOptions): string | undefined { - this.sessionId = options?.id ?? randomUUID(); + this.sessionId = options?.id ?? createSessionId(); const timestamp = new Date().toISOString(); const header: SessionHeader = { type: "session", @@ -1172,7 +1177,7 @@ export class SessionManager { // Filter out LabelEntry from path - we'll recreate them from the resolved map const pathWithoutLabels = path.filter((e) => e.type !== "label"); - const newSessionId = randomUUID(); + const newSessionId = createSessionId(); const timestamp = new Date().toISOString(); const fileTimestamp = timestamp.replace(/[:.]/g, "-"); const newSessionFile = join(this.getSessionDir(), `${fileTimestamp}_${newSessionId}.jsonl`); @@ -1325,7 +1330,7 @@ export class SessionManager { } // Create new session file with new ID but forked content - const newSessionId = randomUUID(); + const newSessionId = createSessionId(); const timestamp = new Date().toISOString(); const fileTimestamp = timestamp.replace(/[:.]/g, "-"); const newSessionFile = join(dir, `${fileTimestamp}_${newSessionId}.jsonl`); diff --git a/packages/coding-agent/test/session-manager/custom-session-id.test.ts b/packages/coding-agent/test/session-manager/custom-session-id.test.ts index 8cded2d2..72eed465 100644 --- a/packages/coding-agent/test/session-manager/custom-session-id.test.ts +++ b/packages/coding-agent/test/session-manager/custom-session-id.test.ts @@ -1,6 +1,11 @@ +import { mkdtempSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; import { describe, expect, it } from "vitest"; import { SessionManager } from "../../src/core/session-manager.js"; +const UUID_V7_RE = /^[0-9a-f]{8}-[0-9a-f]{4}-7[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$/; + describe("SessionManager.newSession with custom id", () => { it("uses the provided id instead of generating one", () => { const session = SessionManager.inMemory(); @@ -8,20 +13,22 @@ describe("SessionManager.newSession with custom id", () => { expect(session.getSessionId()).toBe("my-custom-id"); }); - it("generates a random id when no id is provided", () => { + it("generates a UUIDv7 id when no id is provided", () => { const session = SessionManager.inMemory(); session.newSession(); const id = session.getSessionId(); expect(id).toBeDefined(); expect(id).not.toBe(""); + expect(id).toMatch(UUID_V7_RE); }); - it("generates a random id when options is provided without id", () => { + it("generates a UUIDv7 id when options is provided without id", () => { const session = SessionManager.inMemory(); session.newSession({ parentSession: "parent.jsonl" }); const id = session.getSessionId(); expect(id).toBeDefined(); expect(id).not.toBe(""); + expect(id).toMatch(UUID_V7_RE); }); it("includes the custom id in the session header", () => { @@ -32,4 +39,71 @@ describe("SessionManager.newSession with custom id", () => { expect(header).not.toBeNull(); expect(header!.id).toBe("header-test-id"); }); + + it("generates a UUIDv7 id when constructed without an explicit id", () => { + const session = SessionManager.inMemory(); + expect(session.getSessionId()).toMatch(UUID_V7_RE); + expect(session.getHeader()!.id).toBe(session.getSessionId()); + }); + + it("generates a UUIDv7 id when creating a branched session", () => { + const session = SessionManager.inMemory(); + const firstId = session.appendMessage({ + role: "user", + content: [{ type: "text", text: "hello" }], + timestamp: Date.now(), + }); + + session.createBranchedSession(firstId); + + expect(session.getSessionId()).toMatch(UUID_V7_RE); + expect(session.getHeader()!.id).toBe(session.getSessionId()); + }); + + it("generates a UUIDv7 id when forking from another session file", () => { + const tempDir = mkdtempSync(join(tmpdir(), "pi-session-manager-")); + const sourcePath = join(tempDir, "source.jsonl"); + writeFileSync( + sourcePath, + `${[ + JSON.stringify({ + type: "session", + version: 3, + id: "legacy-session-id", + timestamp: new Date().toISOString(), + cwd: tempDir, + }), + JSON.stringify({ + type: "message", + id: "entry-1", + parentId: null, + timestamp: new Date().toISOString(), + message: { + role: "assistant", + content: [{ type: "text", text: "hello" }], + api: "openai-responses", + provider: "openai", + model: "gpt-5.4", + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "stop", + timestamp: Date.now(), + }, + }), + ].join("\n")} +`, + ); + + const forked = SessionManager.forkFrom(sourcePath, tempDir, tempDir); + const header = forked.getHeader(); + expect(header).not.toBeNull(); + expect(header!.id).toMatch(UUID_V7_RE); + expect(header!.parentSession).toBe(sourcePath); + }); });