diff --git a/plugins/codex-security/scripts/workbench_scan_usage.py b/plugins/codex-security/scripts/workbench_scan_usage.py index f3eb2a0332..406208a3a1 100644 --- a/plugins/codex-security/scripts/workbench_scan_usage.py +++ b/plugins/codex-security/scripts/workbench_scan_usage.py @@ -456,7 +456,7 @@ def _read_rollout_usage( continue if event.get("type") != "event_msg" or payload.get("type") != "token_count": continue - if payload.get("info") is None: + if "info" in payload and payload["info"] is None: continue timestamp = _timestamp(event.get("timestamp")) snapshot = _token_snapshot(payload) diff --git a/plugins/codex-security/tests/test_workbench_scan_usage.py b/plugins/codex-security/tests/test_workbench_scan_usage.py index 06b7ea4250..a7a544ddee 100644 --- a/plugins/codex-security/tests/test_workbench_scan_usage.py +++ b/plugins/codex-security/tests/test_workbench_scan_usage.py @@ -323,6 +323,34 @@ def test_completion_counts_only_scan_owned_parent_and_descendants( assert json.loads(stored[0]) == {"usage": expected} +@pytest.mark.parametrize("measurement", ["unknown", "malformed", "missing"]) +def test_completion_distinguishes_unknown_usage_from_invalid_measurements( + tmp_path: Path, measurement: str +) -> None: + fixture = _start_scan(tmp_path) + counted = fixture.started_at + timedelta(microseconds=1) + payload: dict[str, Any] = {"type": "token_count"} + if measurement == "unknown": + payload["info"] = None + payload["rate_limits"] = {"primary": {"used_percent": 5}} + elif measurement == "malformed": + payload["info"] = {"total_token_usage": {"input_tokens": "invalid"}} + parent = _rollout( + tmp_path, + "scan-parent", + [ + _token_event(counted, 10, 2), + _event(counted, "event_msg", payload), + _token_event(counted, 20, 4), + ], + ) + _state_graph(fixture.environment, {"scan-parent": parent}, []) + usage = _complete_scan(fixture)["scan"]["usage"] + assert usage["totalTokens"] == 24 + assert usage["coverage"] == ("complete" if measurement == "unknown" else "partial") + assert ("token_record_invalid" in usage.get("warnings", [])) is (measurement != "unknown") + + @pytest.mark.parametrize("cache_write_field", ["cache_write_input_tokens", "cache_write_tokens"]) def test_scan_usage_helper_preserves_cached_token_totals( cache_write_field: str, diff --git a/sdk/typescript/src/cost.ts b/sdk/typescript/src/cost.ts index a27cdf998c..057c774c2f 100644 --- a/sdk/typescript/src/cost.ts +++ b/sdk/typescript/src/cost.ts @@ -1,3 +1,4 @@ +import { createHash } from "node:crypto"; import { open, readdir, stat } from "node:fs/promises"; import { join } from "node:path"; import { isRecord } from "./record.js"; @@ -38,6 +39,7 @@ interface SessionReasoning { } interface SessionUsage { + tracked: boolean; offset: number; pendingLine: Buffer[]; pendingLineBytes: number; @@ -51,7 +53,7 @@ interface SessionUsage { usage: ScanTokenUsage | null; calls: Map; activities: ScanActivity[]; - progress: ScanProgress[]; + progress?: ScanProgress[]; filesCompleted: number; filesTotal: number | null; prose: Set; @@ -85,6 +87,7 @@ const SESSION_READ_SIZE = 64 * 1_024; function createSessionUsage(): SessionUsage { return { + tracked: false, offset: 0, pendingLine: [], pendingLineBytes: 0, @@ -98,7 +101,6 @@ function createSessionUsage(): SessionUsage { usage: null, calls: new Map(), activities: [], - progress: [], filesCompleted: 0, filesTotal: null, prose: new Set(), @@ -202,7 +204,6 @@ export class ScanCostTracker { async #readSessions(): Promise { if (this.#threadId === null) return; - const unreadable: Array<{ session: SessionUsage; error: unknown }> = []; for await (const path of sessionFiles( join(this.#options.codexHome, "sessions"), )) { @@ -211,11 +212,10 @@ export class ScanCostTracker { session = createSessionUsage(); this.#sessions.set(path, session); } - try { - await readSessionUsage(path, session, this.#options.repository); - } catch (error) { - if (session.threadId === null) throw error; - unreadable.push({ session, error }); + if (session.threadId === null) { + // Index ownership before reading a transcript. Unrelated sessions need + // only their metadata, including parents discovered by a later poll. + await readSessionUsage(path, session, undefined, true); } } @@ -259,30 +259,34 @@ export class ScanCostTracker { } } } while (included.size !== previousSize); - for (const { session, error } of unreadable) { - if (included.has(session.threadId!)) throw error; - } - const usages = new Map(this.#receipts); for (const [path, tracked] of this.#sessions) { const threadId = tracked.threadId; if (threadId === null || !included.has(threadId)) continue; let session = tracked; - if ( - this.#options.onSessionEvent !== undefined && - session.events === undefined - ) { - // Replay only newly associated sessions, including their early events. - session = createSessionUsage(); - session.events = []; - await readSessionUsage(path, session, this.#options.repository); - this.#sessions.set(path, session); - } let worker: number | undefined; if (threadId !== this.#threadId) { worker = this.#workers.get(threadId) ?? this.#workers.size + 1; this.#workers.set(threadId, worker); } + if (!session.tracked) { + // Replay newly associated sessions from the start for every observer, + // not just raw session events. Their early usage and activity matter. + session = createSessionUsage(); + session.tracked = true; + if (this.#options.onSessionEvent !== undefined) session.events = []; + if (worker !== undefined && this.#options.onProgress !== undefined) { + session.progress = []; + } + } + await readSessionUsage( + path, + session, + worker !== undefined && this.#options.onActivity !== undefined + ? this.#options.repository + : undefined, + ); + this.#sessions.set(path, session); for (const event of session.events?.splice(0) ?? []) { this.#options.onSessionEvent?.({ threadId, @@ -331,7 +335,7 @@ export class ScanCostTracker { if (this.#options.onProgress === undefined || session.threadId === null) { return; } - for (const progress of session.progress.splice(0)) { + for (const progress of session.progress?.splice(0) ?? []) { const expectedFilesTotal = this.#expectedFilesTotal; if ( (expectedFilesTotal !== undefined && @@ -405,6 +409,7 @@ async function readSessionUsage( path: string, session: SessionUsage, repository?: string, + metadataOnly = false, ): Promise { if (session.unreadable) return; let file; @@ -427,13 +432,19 @@ async function readSessionUsage( if (bytesRead === 0) return; session.offset += bytesRead; try { - readSessionChunk(buffer.subarray(0, bytesRead), session, repository); + readSessionChunk( + buffer.subarray(0, bytesRead), + session, + repository, + metadataOnly, + ); } catch (error) { session.unreadable = true; session.pendingLine = []; session.pendingLineBytes = 0; throw error; } + if (metadataOnly && session.threadId !== null) return; } } finally { await file.close(); @@ -444,6 +455,7 @@ function readSessionChunk( contents: Buffer, session: SessionUsage, repository?: string, + metadataOnly = false, ): void { let lineStart = 0; while (lineStart < contents.length) { @@ -461,17 +473,24 @@ function readSessionChunk( } if (session.pendingLineBytes === 0) { - readSessionEvent(fragment.toString("utf8"), session, repository); + readSessionEvent( + fragment.toString("utf8"), + session, + repository, + metadataOnly, + ); } else { if (fragment.length > 0) session.pendingLine.push(Buffer.from(fragment)); readSessionEvent( Buffer.concat(session.pendingLine, lineBytes).toString("utf8"), session, repository, + metadataOnly, ); session.pendingLine = []; session.pendingLineBytes = 0; } + if (metadataOnly && session.threadId !== null) return; lineStart = newline + 1; } } @@ -480,6 +499,7 @@ function readSessionEvent( line: string, session: SessionUsage, repository?: string, + metadataOnly = false, ): void { if (line.length === 0) return; let event: unknown; @@ -507,6 +527,7 @@ function readSessionEvent( session.events?.push(event); return; } + if (metadataOnly) return; if (session.replaying) { if (event["type"] !== "event_msg") return; if (payload["type"] === "token_count" && isRecord(payload["info"])) { @@ -532,7 +553,7 @@ function readSessionEvent( } session.events?.push(event); if (event["type"] === "response_item") { - session.progress.push(...sessionProgressUpdates(payload)); + session.progress?.push(...sessionProgressUpdates(payload)); if (repository === undefined) return; if ( payload["type"] === "reasoning" && @@ -553,10 +574,7 @@ function readSessionEvent( }, repository, ); - if ( - activity === null || - session.prose.has(`${activity.kind}:${activity.description}`) - ) { + if (activity === null || session.prose.has(proseKey(activity))) { continue; } session.reasoning = { @@ -591,12 +609,12 @@ function readSessionEvent( session.reasoning = null; if ( activity.kind === "message" && - session.prose.has(`${activity.kind}:${activity.description}`) + session.prose.has(proseKey(activity)) ) { return; } if (activity.kind === "message") { - session.prose.add(`${activity.kind}:${activity.description}`); + session.prose.add(proseKey(activity)); } if (activity.status === "running") { session.calls.set(activity.id, activity); @@ -632,7 +650,9 @@ function readSessionEvent( payload["type"] === "agent_message" && typeof payload["message"] === "string" ) { - session.progress.push(...scanProgressUpdatesFromText(payload["message"])); + session.progress?.push( + ...scanProgressUpdatesFromText(payload["message"]), + ); } if (repository === undefined) return; if (payload["type"] !== "agent_message") { @@ -641,11 +661,8 @@ function readSessionEvent( } session.reasoning = null; const activity = scanActivityFromSessionEvent(event, repository); - if ( - activity !== null && - !session.prose.has(`${activity.kind}:${activity.description}`) - ) { - session.prose.add(`${activity.kind}:${activity.description}`); + if (activity !== null && !session.prose.has(proseKey(activity))) { + session.prose.add(proseKey(activity)); session.activities.push(activity); } return; @@ -742,10 +759,21 @@ function recordReasoningActivity( return; } reasoning.activity = activity; - session.prose.add(`${activity.kind}:${activity.description}`); + session.prose.add(proseKey(activity)); session.activities.push(activity); } +function proseKey(activity: ScanActivity): string { + // Deduplication needs an identity, not a retained copy of every transcript + // message and every expanding reasoning prefix. + // Hash UTF-16 code units to keep distinct lone surrogates distinct too. + return createHash("sha256") + .update(activity.kind) + .update(":") + .update(activity.description, "utf16le") + .digest("hex"); +} + function sessionProgressUpdates( payload: Readonly>, ): ScanProgress[] { diff --git a/sdk/typescript/tests-ts/cost.test.ts b/sdk/typescript/tests-ts/cost.test.ts index 94c7b244fc..d4627a7263 100644 --- a/sdk/typescript/tests-ts/cost.test.ts +++ b/sdk/typescript/tests-ts/cost.test.ts @@ -128,6 +128,185 @@ test.each([ }); describe("scan cost", () => { + test("reads only metadata from unrelated sessions and stops reopening them", async () => { + const home = await codexHome(); + const usage = { input_tokens: 100, output_tokens: 10 }; + await writeSession(home, "scan-thread", usage); + const unrelated = await writeSession(home, "unrelated-thread", usage); + await appendSessionItem(unrelated, { + type: "message", + role: "assistant", + content: [{ type: "output_text", text: "x".repeat(256 * 1_024) }], + }); + const originalOpen = fs.open; + let opens = 0; + let bytesRead = 0; + const restores: Array<() => void> = []; + const opening = spyOn(fs, "open").mockImplementation(async (...args) => { + const file = await originalOpen(...args); + if (String(args[0]) === unrelated) { + opens += 1; + const originalRead = file.read.bind(file); + const reading = spyOn(file, "read").mockImplementation((async ( + ...readArgs + ) => { + const result = await Reflect.apply(originalRead, file, readArgs); + bytesRead += result.bytesRead; + return result; + }) as typeof file.read); + restores.push(() => reading.mockRestore()); + } + return file; + }); + const tracker = new ScanCostTracker({ + codexHome: home, + model: "gpt-5.6-sol", + repository: home, + }); + tracker.start("scan-thread"); + try { + expect((await tracker.refresh()).cost?.inputTokens).toBe(100); + expect(bytesRead).toBeLessThan(256 * 1_024); + expect(opens).toBe(1); + await appendSessionItem(unrelated, { + type: "message", + role: "assistant", + content: [{ type: "output_text", text: "More unrelated output." }], + }); + await tracker.refresh(); + await tracker.stop(); + expect(opens).toBe(1); + } finally { + opening.mockRestore(); + for (const restore of restores) restore(); + } + }); + + test("retries an interrupted initial replay without losing usage or events", async () => { + const home = await codexHome(); + await writeSession(home, "scan-thread", { + input_tokens: 100, + output_tokens: 10, + }); + const events: ScanSessionEvent[] = []; + const errors: unknown[] = []; + const tracker = new ScanCostTracker({ + codexHome: home, + model: "gpt-5.6-sol", + onSessionEvent: (event) => events.push(event), + onError: (error) => errors.push(error), + }); + const originalOpen = fs.open; + let opens = 0; + const failure = Object.assign( + new Error("Synthetic transient read failure"), + { code: "EIO" }, + ); + const opening = spyOn(fs, "open").mockImplementation(async (...args) => { + if (++opens === 2) throw failure; + return originalOpen(...args); + }); + try { + tracker.start("scan-thread"); + const final = await tracker.stop(); + expect(errors).toEqual([failure]); + expect(final.cost?.inputTokens).toBe(100); + expect(events.map(({ event }) => event["type"])).toEqual([ + "session_meta", + "event_msg", + ]); + } finally { + opening.mockRestore(); + } + }); + + test("replays a worker after its metadata arrives across multiple reads", async () => { + const home = await codexHome(); + const usage = { input_tokens: 100, output_tokens: 10 }; + await writeSession(home, "scan-thread", usage); + const worker = join(home, "sessions", "partial-worker.jsonl"); + const metadata = JSON.stringify({ + type: "session_meta", + payload: { + padding: "x".repeat(128 * 1_024), + id: "worker-thread", + parent_thread_id: "scan-thread", + }, + }); + await writeFile(worker, metadata.slice(0, -2)); + const events: ScanSessionEvent[] = []; + const tracker = new ScanCostTracker({ + codexHome: home, + model: "gpt-5.6-sol", + onSessionEvent: (event) => events.push(event), + }); + tracker.start("scan-thread"); + try { + await tracker.refresh(); + expect(events.some((event) => event.threadId === "worker-thread")).toBe( + false, + ); + await appendFile(worker, `${metadata.slice(-2)}\n`); + await appendSessionItem(worker, { + type: "message", + role: "assistant", + content: [{ type: "output_text", text: "Early worker output." }], + }); + await tracker.refresh(); + await tracker.refresh(); + expect( + events + .filter((event) => event.threadId === "worker-thread") + .map(({ event }) => event["type"]), + ).toEqual(["session_meta", "response_item"]); + } finally { + await tracker.stop(); + } + }); + + test("deduplicates long worker messages without dropping distinct suffixes", async () => { + const home = await codexHome(); + const usage = { input_tokens: 100, output_tokens: 10 }; + await writeSession(home, "scan-thread", usage); + const worker = await writeSession(home, "worker-thread", usage, { + parent: "scan-thread", + }); + const first = `${"x".repeat(4_096)} first`; + const second = `${"x".repeat(4_096)} second`; + const distinct = [ + first, + second, + `${first}\ud800`, + `${first}\ud801`, + `${first}\ufffd`, + ]; + const activities: ScanActivity[] = []; + const tracker = new ScanCostTracker({ + codexHome: home, + model: "gpt-5.6-sol", + repository: home, + onActivity: (activity) => activities.push(activity), + }); + tracker.start("scan-thread"); + try { + for (const message of [...distinct, first]) { + await appendFile( + worker, + `${JSON.stringify({ + type: "event_msg", + payload: { type: "agent_message", message }, + })}\n`, + ); + await tracker.refresh(); + } + expect(activities.map((activity) => activity.description)).toEqual( + distinct, + ); + } finally { + await tracker.stop(); + } + }); + test("ignores token-count events without a usage snapshot in both readers", async () => { const home = await codexHome(); const path = await writeSession(