From 45f78c07d8f31356e951120f9800e5c2861876ef Mon Sep 17 00:00:00 2001
From: Max Sapelov <13932412+msapelov@users.noreply.github.com>
Date: Tue, 8 Sep 2026 21:28:18 +0200
Subject: [PATCH] fix(sdk): prevent transcript memory leaks in scan tracking

---
 sdk/typescript/src/cost.ts           | 104 +++++++++-----
 sdk/typescript/tests-ts/cost.test.ts | 205 +++++++++++++++++++++++++--
 2 files changed, 256 insertions(+), 53 deletions(-)

diff --git a/sdk/typescript/src/cost.ts b/sdk/typescript/src/cost.ts
index 81247c072..1b5df2369 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 } from "node:fs/promises";
 import { join } from "node:path";
 import {
@@ -37,6 +38,7 @@ interface SessionReasoning {
 }
 
 interface SessionUsage {
+  tracked: boolean;
   offset: number;
   pendingLine: Buffer[];
   pendingLineBytes: number;
@@ -50,7 +52,7 @@ interface SessionUsage {
   usage: ScanTokenUsage | null;
   calls: Map<string, ScanActivity>;
   activities: ScanActivity[];
-  progress: ScanProgress[];
+  progress?: ScanProgress[];
   filesCompleted: number;
   filesTotal: number | null;
   prose: Set<string>;
@@ -83,6 +85,7 @@ const SESSION_READ_SIZE = 64 * 1_024;
 
 function createSessionUsage(): SessionUsage {
   return {
+    tracked: false,
     offset: 0,
     pendingLine: [],
     pendingLineBytes: 0,
@@ -96,7 +99,6 @@ function createSessionUsage(): SessionUsage {
     usage: null,
     calls: new Map(),
     activities: [],
-    progress: [],
     filesCompleted: 0,
     filesTotal: null,
     prose: new Set(),
@@ -199,7 +201,6 @@ export class ScanCostTracker {
 
   async #readSessions(): Promise<void> {
     if (this.#threadId === null) return;
-    const unreadable: Array<{ session: SessionUsage; error: unknown }> = [];
     for await (const path of sessionFiles(
       join(this.#options.codexHome, "sessions"),
     )) {
@@ -208,11 +209,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);
       }
     }
 
@@ -258,30 +258,34 @@ export class ScanCostTracker {
         }
       }
     }
-    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 = [];
+        }
+        this.#sessions.set(path, session);
+      }
+      await readSessionUsage(
+        path,
+        session,
+        worker !== undefined && this.#options.onActivity !== undefined
+          ? this.#options.repository
+          : undefined,
+      );
       for (const event of session.events?.splice(0) ?? []) {
         this.#options.onSessionEvent?.({
           threadId,
@@ -325,7 +329,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 &&
@@ -390,6 +394,7 @@ async function readSessionUsage(
   path: string,
   session: SessionUsage,
   repository?: string,
+  metadataOnly = false,
 ): Promise<void> {
   if (session.unreadable) return;
   let file;
@@ -411,13 +416,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();
@@ -428,6 +439,7 @@ function readSessionChunk(
   contents: Buffer,
   session: SessionUsage,
   repository?: string,
+  metadataOnly = false,
 ): void {
   let lineStart = 0;
   while (lineStart < contents.length) {
@@ -445,17 +457,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;
   }
 }
@@ -464,6 +483,7 @@ function readSessionEvent(
   line: string,
   session: SessionUsage,
   repository?: string,
+  metadataOnly = false,
 ): void {
   if (line.length === 0) return;
   let event: unknown;
@@ -491,6 +511,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"])) {
@@ -516,7 +537,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" &&
@@ -537,10 +558,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 = {
@@ -575,12 +593,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);
@@ -616,7 +634,7 @@ function readSessionEvent(
       payload["type"] === "agent_message" &&
       typeof payload["message"] === "string"
     ) {
-      session.progress.push(
+      session.progress?.push(
         ...scanProgressUpdatesFromEvent({
           type: "item.completed",
           item: { type: "agent_message", text: payload["message"] },
@@ -630,11 +648,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;
@@ -731,10 +746,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<Record<string, unknown>>,
 ): ScanProgress[] {
diff --git a/sdk/typescript/tests-ts/cost.test.ts b/sdk/typescript/tests-ts/cost.test.ts
index c04c6fa07..f7ab8507d 100644
--- a/sdk/typescript/tests-ts/cost.test.ts
+++ b/sdk/typescript/tests-ts/cost.test.ts
@@ -1,4 +1,5 @@
 import { spawnSync } from "node:child_process";
+import * as filesystem from "node:fs/promises";
 import {
   appendFile,
   mkdir,
@@ -9,7 +10,7 @@ import {
 } from "node:fs/promises";
 import { tmpdir } from "node:os";
 import { join, parse, sep } from "node:path";
-import { afterEach, describe, expect, test } from "bun:test";
+import { afterEach, describe, expect, spyOn, test } from "bun:test";
 import {
   estimateScanCost,
   ScanCostTracker,
@@ -382,6 +383,152 @@ describe("scan cost", () => {
 });
 
 describe("live scan cost tracking", () => {
+  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 = filesystem.open;
+    let opens = 0;
+    let bytesRead = 0;
+    const restores: Array<() => void> = [];
+    const opening = spyOn(filesystem, "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("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,
+      "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("coalesces overlapping polling ticks and bounds final work", async () => {
     const home = await codexHome();
     await writeSession(home, "scan-thread", {
@@ -566,9 +713,14 @@ describe("live scan cost tracking", () => {
     expect(events).toHaveLength(6);
   });
 
-  test.each(["parent", "main"] as const)(
-    "replays early worker events when the %s session arrives later",
-    async (missing) => {
+  test.each([
+    ["parent", true],
+    ["parent", false],
+    ["main", true],
+    ["main", false],
+  ] as const)(
+    "replays early worker output when the %s session arrives later (raw events: %s)",
+    async (missing, rawEvents) => {
       const home = await codexHome();
       const scanDirectory = join(home, "scan");
       const usage = { input_tokens: 10, output_tokens: 1 };
@@ -598,19 +750,28 @@ describe("live scan cost tracking", () => {
         "2026-07-26T12:01:00Z",
       );
       const message = (text: string) => ({
+        id: text,
         type: "message",
         role: "assistant",
         content: [{ type: "output_text", text }],
       });
       await appendSessionItem(worker, message("Early worker output."));
+      await appendSessionItem(worker, progressMessage(2));
       const unrelated = await writeSession(home, "unrelated-thread", usage);
       await appendSessionItem(unrelated, message("Unrelated output."));
       const events: ScanSessionEvent[] = [];
+      const activities: ScanActivity[] = [];
+      const progress: ScanProgress[] = [];
       const tracker = new ScanCostTracker({
         codexHome: home,
         scanDirectory,
         model: "gpt-5.6-sol",
-        onSessionEvent: (event) => events.push(event),
+        repository: home,
+        onActivity: (activity) => activities.push(activity),
+        onProgress: (update) => progress.push(update),
+        ...(rawEvents
+          ? { onSessionEvent: (event: ScanSessionEvent) => events.push(event) }
+          : {}),
       });
       tracker.start("scan-thread");
       await tracker.refresh();
@@ -626,20 +787,36 @@ describe("live scan cost tracking", () => {
       await appendSessionItem(worker, message("Late worker output."));
       await tracker.refresh();
       await tracker.refresh();
-      await tracker.stop();
+      const snapshot = await tracker.stop();
+      expect(snapshot.usage).toMatchObject({
+        input_tokens: missing === "parent" ? 30 : 20,
+        output_tokens: missing === "parent" ? 3 : 2,
+      });
 
       const workerEvents = events.filter(
         (event) => event.threadId === "worker-thread",
       );
-      expect(workerEvents.map((event) => event.event)).toEqual([
-        expect.objectContaining({ type: "session_meta" }),
-        expect.objectContaining({ type: "event_msg" }),
-        { type: "response_item", payload: message("Early worker output.") },
-        { type: "response_item", payload: message("Late worker output.") },
+      if (rawEvents) {
+        expect(workerEvents.map((event) => event.event)).toEqual([
+          expect.objectContaining({ type: "session_meta" }),
+          expect.objectContaining({ type: "event_msg" }),
+          { type: "response_item", payload: message("Early worker output.") },
+          { type: "response_item", payload: progressMessage(2) },
+          { type: "response_item", payload: message("Late worker output.") },
+        ]);
+        expect(new Set(workerEvents.map((event) => event.worker))).toEqual(
+          new Set([1]),
+        );
+      } else {
+        expect(events).toEqual([]);
+      }
+      expect(activities.map((activity) => activity.description)).toEqual([
+        "Early worker output.",
+        "Late worker output.",
+      ]);
+      expect(progress).toEqual([
+        { phase: "discovery", filesCompleted: 2, filesTotal: 8 },
       ]);
-      expect(new Set(workerEvents.map((event) => event.worker))).toEqual(
-        new Set([1]),
-      );
       expect(
         events.some((event) => event.threadId === "unrelated-thread"),
       ).toBe(false);