shumengya/sproutclaw
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: align OpenAI cache affinity and use uuidv7 session ids

Browse Source
This commit is contained in:
Mario Zechner
2026-04-14 23:20:13 +02:00
parent d62d22173a
commit 018b40c30c
11 changed files with 245 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
packages/ai/CHANGELOG.md
View File

@@ -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

1
packages/ai/src/providers/openai-codex-responses.ts
View File

@@ -950,6 +950,7 @@ function buildSSEHeaders(
if (sessionId) {
headers.set("session_id", sessionId);
headers.set("x-client-request-id", sessionId);
}
return headers;

10
packages/ai/src/providers/openai-responses.ts
View File

@@ -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<string, string>,
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);

36
packages/ai/test/openai-codex-cache-affinity-e2e.test.ts Normal file
View File

@@ -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",
);
});
});

9
packages/ai/test/openai-codex-stream.test.ts
View File

@@ -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<string, unknown>) : 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,

32
packages/ai/test/openai-responses-cache-affinity-e2e.test.ts Normal file
View File

@@ -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",
);
});
});

73
packages/ai/test/openai-responses-copilot-provider.test.ts
View File

@@ -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<string, string | readonly string[]> | 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<typeof streamOpenAIResponses>[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 });
});
});

1
packages/coding-agent/CHANGELOG.md
View File

@@ -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))

1
packages/coding-agent/package.json
View File

@@ -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": {

13
packages/coding-agent/src/core/session-manager.ts
View File

@@ -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`);

78
packages/coding-agent/test/session-manager/custom-session-id.test.ts
View File

@@ -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);
});
});