From 67118d5ab9129f6444dc8a65b982ebaec0b3e4ed Mon Sep 17 00:00:00 2001 From: lin-hongkuan Date: Thu, 25 Jun 2026 22:57:31 +0800 Subject: [PATCH] fix(trajectory): preserve codex completion usage --- .../codex/src/app-server/trajectory.test.ts | 62 +++++++++++++++++- extensions/codex/src/app-server/trajectory.ts | 63 +++++++++++++++---- 2 files changed, 112 insertions(+), 13 deletions(-) diff --git a/extensions/codex/src/app-server/trajectory.test.ts b/extensions/codex/src/app-server/trajectory.test.ts index 4eb99b9a6e0..d3797b1d254 100644 --- a/extensions/codex/src/app-server/trajectory.test.ts +++ b/extensions/codex/src/app-server/trajectory.test.ts @@ -5,6 +5,7 @@ import path from "node:path"; import { afterEach, describe, expect, it } from "vitest"; import { createCodexTrajectoryRecorder, + recordCodexTrajectoryCompletion, recordCodexTrajectoryContext, resolveCodexTrajectoryAppendFlags, resolveCodexTrajectoryPointerFlags, @@ -80,7 +81,9 @@ describe("Codex trajectory recorder", () => { expect(content).not.toContain("secret"); expect(content).not.toContain("sk-test-secret-token"); expect(content).not.toContain("sk-other-secret-token"); - expect(fs.statSync(filePath).mode & 0o777).toBe(0o600); + if (process.platform !== "win32") { + expect(fs.statSync(filePath).mode & 0o777).toBe(0o600); + } expect(fs.existsSync(path.join(tmpDir, "session.trajectory-path.json"))).toBe(true); }); @@ -253,4 +256,61 @@ describe("Codex trajectory recorder", () => { expect(parsed.data?.truncated).toBe(true); expect(parsed.data?.reason).toBe("trajectory-event-size-limit"); }); + + it("preserves usage when truncating oversized model completion events", async () => { + const tmpDir = makeTempDir(); + const sessionFile = path.join(tmpDir, "session.jsonl"); + const attempt = { + sessionFile, + sessionId: "session-1", + sessionKey: "agent:main:session-1", + runId: "run-1", + provider: "codex", + modelId: "gpt-5.4", + model: { api: "responses" }, + } as never; + const usage = { + input: 384_954, + output: 5_624, + cacheRead: 333_824, + reasoningTokens: 2_038, + total: 724_402, + }; + const recorder = createCodexTrajectoryRecorder({ + cwd: tmpDir, + attempt, + env: {}, + }); + + const trajectoryRecorder = expectTrajectoryRecorder(recorder); + recordCodexTrajectoryCompletion(trajectoryRecorder, { + attempt, + threadId: "thread-1", + turnId: "turn-1", + timedOut: false, + result: { + aborted: false, + attemptUsage: usage, + assistantTexts: ["done"], + messagesSnapshot: Array.from({ length: 20 }, (_value, index) => ({ + role: index % 2 === 0 ? "user" : "assistant", + content: `message-${index} ${"x".repeat(32_000)}`, + })), + } as never, + }); + await trajectoryRecorder.flush(); + + const parsed = JSON.parse( + fs.readFileSync(path.join(tmpDir, "session.trajectory.jsonl"), "utf8"), + ); + expect(parsed.type).toBe("model.completed"); + expect(parsed.data).toMatchObject({ + truncated: true, + reason: "trajectory-event-size-limit", + usage, + }); + expect(parsed.data.messagesSnapshot).toBeUndefined(); + expect(parsed.data.droppedFields).toContain("messagesSnapshot"); + expect(Buffer.byteLength(JSON.stringify(parsed), "utf8")).toBeLessThanOrEqual(256 * 1024); + }); }); diff --git a/extensions/codex/src/app-server/trajectory.ts b/extensions/codex/src/app-server/trajectory.ts index 4130e900db4..23a406166d1 100644 --- a/extensions/codex/src/app-server/trajectory.ts +++ b/extensions/codex/src/app-server/trajectory.ts @@ -40,6 +40,7 @@ const JWT_VALUE_RE = /\beyJ[A-Za-z0-9_-]{10,}\.[A-Za-z0-9_-]{10,}\.[A-Za-z0-9_-] const COOKIE_PAIR_RE = /\b([A-Za-z][A-Za-z0-9_.-]{1,64})=([A-Za-z0-9+/._~%=-]{16,})(?=;|\s|$)/gu; const TRAJECTORY_RUNTIME_FILE_MAX_BYTES = 50 * 1024 * 1024; const TRAJECTORY_RUNTIME_EVENT_MAX_BYTES = 256 * 1024; +const TRAJECTORY_RUNTIME_OVERSIZE_PRESERVED_DATA_KEYS = ["usage", "promptCache"] as const; type CodexTrajectoryOpenFlagConstants = Pick< typeof nodeFs.constants, @@ -82,19 +83,57 @@ function boundedTrajectoryLine(event: Record): string | undefin if (bytes <= TRAJECTORY_RUNTIME_EVENT_MAX_BYTES) { return `${line}\n`; } - const truncated = JSON.stringify({ - ...event, - data: { - truncated: true, - originalBytes: bytes, - limitBytes: TRAJECTORY_RUNTIME_EVENT_MAX_BYTES, - reason: "trajectory-event-size-limit", - }, - }); - if (Buffer.byteLength(truncated, "utf8") <= TRAJECTORY_RUNTIME_EVENT_MAX_BYTES) { - return `${truncated}\n`; + + const originalData = + event.data && typeof event.data === "object" && !Array.isArray(event.data) + ? (event.data as Record) + : {}; + const originalDataKeys = Object.keys(originalData); + const preservedDataKeys = new Set(); + const baseData = { + truncated: true, + originalBytes: bytes, + limitBytes: TRAJECTORY_RUNTIME_EVENT_MAX_BYTES, + reason: "trajectory-event-size-limit", + }; + const buildTruncatedLine = (includeDroppedFields: boolean): string | undefined => { + const data: Record = { ...baseData }; + for (const key of TRAJECTORY_RUNTIME_OVERSIZE_PRESERVED_DATA_KEYS) { + if (preservedDataKeys.has(key)) { + data[key] = originalData[key]; + } + } + if (includeDroppedFields) { + const droppedFields = originalDataKeys.filter((key) => !preservedDataKeys.has(key)); + if (droppedFields.length > 0) { + data.droppedFields = droppedFields; + } + } + const truncated = JSON.stringify({ ...event, data }); + if (Buffer.byteLength(truncated, "utf8") <= TRAJECTORY_RUNTIME_EVENT_MAX_BYTES) { + return `${truncated}\n`; + } + return undefined; + }; + + let best = buildTruncatedLine(true) ?? buildTruncatedLine(false); + if (!best) { + return undefined; } - return undefined; + + for (const key of TRAJECTORY_RUNTIME_OVERSIZE_PRESERVED_DATA_KEYS) { + if (!Object.hasOwn(originalData, key)) { + continue; + } + preservedDataKeys.add(key); + const next = buildTruncatedLine(true) ?? buildTruncatedLine(false); + if (next) { + best = next; + continue; + } + preservedDataKeys.delete(key); + } + return best; } function resolveTrajectoryPointerFilePath(sessionFile: string): string {