perf(sessions): collapse redundant tool_call_update snapshots

apps/code/src/main/di/container.ts

@@ -711,6 +711,8 @@ container.bind(LOGS_SERVICE).toDynamicValue((ctx) => {
     },
     readLocalLogs: (taskRunId: string) =>
       ws.localLogs.read.query({ taskRunId }),
+    readLocalLogsCollapsed: (taskRunId: string) =>
+      ws.localLogs.readCollapsed.query({ taskRunId }),
     readLocalLogsTail: (taskRunId: string, maxBytes: number) =>
       ws.localLogs.readTail.query({ taskRunId, maxBytes }),
     writeLocalLogs: (taskRunId: string, content: string) =>

packages/agent/src/session-log-writer.test.ts

@@ -514,3 +514,174 @@ describe("SessionLogWriter", () => {
     });
   });
 });
+
+describe("SessionLogWriter — local-cache tool_call_update coalescing", () => {
+  let tmp: string;
+  let writer: SessionLogWriter;
+  const RUN = "run-coalesce";
+
+  beforeEach(async () => {
+    const fs = await import("node:fs");
+    const os = await import("node:os");
+    const path = await import("node:path");
+    tmp = fs.mkdtempSync(path.join(os.tmpdir(), "slw-"));
+    writer = new SessionLogWriter({ localCachePath: tmp });
+    writer.register(RUN, { taskId: "t", runId: RUN });
+  });
+
+  afterEach(async () => {
+    const fs = await import("node:fs");
+    fs.rmSync(tmp, { recursive: true, force: true });
+  });
+
+  const readLog = async (): Promise<Record<string, unknown>[]> => {
+    const fs = await import("node:fs");
+    const path = await import("node:path");
+    const p = path.join(tmp, "sessions", RUN, "logs.ndjson");
+    if (!fs.existsSync(p)) return [];
+    return fs
+      .readFileSync(p, "utf-8")
+      .trim()
+      .split("\n")
+      .filter(Boolean)
+      .map((l) => JSON.parse(l));
+  };
+
+  const update = (extra: Record<string, unknown>) =>
+    makeSessionUpdate("tool_call_update", { toolCallId: "a", ...extra });
+
+  const sessionUpdateOf = (e: Record<string, unknown>) =>
+    // biome-ignore lint/suspicious/noExplicitAny: test introspection
+    (e.notification as any).params.update;
+
+  it("writes one merged update per call, flushed by a non-tool event", async () => {
+    writer.appendRawLine(RUN, update({ content: "a1" }));
+    writer.appendRawLine(RUN, update({ content: "a2" }));
+    writer.appendRawLine(RUN, update({ content: "a3" }));
+    // a non-tool event flushes the buffered union, then writes itself
+    writer.appendRawLine(RUN, makeSessionUpdate("agent_message"));
+
+    const log = await readLog();
+    expect(log).toHaveLength(2);
+    expect(sessionUpdateOf(log[0]).content).toBe("a3");
+    expect(sessionUpdateOf(log[1]).sessionUpdate).toBe("agent_message");
+  });
+
+  it("a terminal update merges into buffered snapshots, later fields winning", async () => {
+    writer.appendRawLine(RUN, update({ content: "a1" }));
+    writer.appendRawLine(RUN, update({ content: "a2" }));
+    writer.appendRawLine(
+      RUN,
+      update({ content: "final", status: "completed" }),
+    );
+
+    const log = await readLog();
+    expect(log).toHaveLength(1);
+    expect(sessionUpdateOf(log[0]).content).toBe("final");
+    expect(sessionUpdateOf(log[0]).status).toBe("completed");
+  });
+
+  it("fields carried only by earlier updates survive the terminal write", async () => {
+    // Mirrors the real emission shape: streamed rawInput snapshots carry the
+    // input, the terminal update carries only status/rawOutput.
+    writer.appendRawLine(RUN, update({ rawInput: { command: "ls" } }));
+    writer.appendRawLine(
+      RUN,
+      update({ rawInput: { command: "ls -la" }, title: "List files" }),
+    );
+    writer.appendRawLine(
+      RUN,
+      update({ status: "completed", rawOutput: "done" }),
+    );
+
+    const log = await readLog();
+    expect(log).toHaveLength(1);
+    expect(sessionUpdateOf(log[0])).toMatchObject({
+      toolCallId: "a",
+      rawInput: { command: "ls -la" },
+      title: "List files",
+      status: "completed",
+      rawOutput: "done",
+    });
+  });
+
+  it("keeps every uncoalesced update on the API path, unmutated by the merge", async () => {
+    const appendLog = vi.fn().mockResolvedValue(undefined);
+    const apiWriter = new SessionLogWriter({
+      localCachePath: tmp,
+      posthogAPI: { appendTaskRunLog: appendLog } as never,
+    });
+    const API_RUN = "run-api";
+    apiWriter.register(API_RUN, { taskId: "t", runId: API_RUN });
+
+    apiWriter.appendRawLine(
+      API_RUN,
+      makeSessionUpdate("tool_call_update", {
+        toolCallId: "a",
+        rawInput: { command: "ls" },
+      }),
+    );
+    apiWriter.appendRawLine(
+      API_RUN,
+      makeSessionUpdate("tool_call_update", {
+        toolCallId: "a",
+        status: "completed",
+      }),
+    );
+    await apiWriter.flush(API_RUN);
+
+    // The durable log receives both updates as emitted; the buffered merge
+    // must build its own object rather than write into the shared entries.
+    const entries = appendLog.mock.calls[0][2] as {
+      notification: { params: { update: Record<string, unknown> } };
+    }[];
+    expect(entries).toHaveLength(2);
+    expect(entries[0].notification.params.update).toEqual({
+      sessionUpdate: "tool_call_update",
+      toolCallId: "a",
+      rawInput: { command: "ls" },
+    });
+    expect(entries[1].notification.params.update).toEqual({
+      sessionUpdate: "tool_call_update",
+      toolCallId: "a",
+      status: "completed",
+    });
+  });
+
+  it("flushAll persists a still-buffered merged update", async () => {
+    writer.appendRawLine(RUN, update({ rawInput: { command: "ls" } }));
+    writer.appendRawLine(RUN, update({ content: "a2" }));
+    await writer.flushAll();
+
+    const log = await readLog();
+    expect(log).toHaveLength(1);
+    expect(sessionUpdateOf(log[0]).content).toBe("a2");
+    expect(sessionUpdateOf(log[0]).rawInput).toEqual({ command: "ls" });
+  });
+
+  it("hold-window flush writes the union so far and starts a new window", async () => {
+    vi.useFakeTimers();
+    try {
+      writer.appendRawLine(RUN, update({ rawInput: { command: "ls" } }));
+      vi.advanceTimersByTime(2500);
+      // Exceeds TOOL_UPDATE_MAX_HOLD_MS: the buffered union is written, this
+      // update starts a fresh window.
+      writer.appendRawLine(RUN, update({ content: "partial" }));
+      writer.appendRawLine(
+        RUN,
+        update({ status: "completed", rawOutput: "done" }),
+      );
+
+      const log = await readLog();
+      expect(log).toHaveLength(2);
+      expect(sessionUpdateOf(log[0]).rawInput).toEqual({ command: "ls" });
+      // The second line unions the post-window snapshot with the terminal
+      // update; a merge-on-read of both lines rebuilds the full call state.
+      expect(sessionUpdateOf(log[1]).content).toBe("partial");
+      expect(sessionUpdateOf(log[1]).status).toBe("completed");
+      expect(sessionUpdateOf(log[1]).rawOutput).toBe("done");
+    } finally {
+      vi.useRealTimers();
+    }
+  });
+});

packages/agent/src/session-log-writer.ts

@@ -21,14 +21,41 @@ interface ChunkBuffer {
   firstTimestamp: string;
 }
 
+/**
+ * In-progress `tool_call_update`s buffered per toolCallId, awaiting a
+ * coalesced write to the local cache (see appendRawLine). `mergedUpdate` is
+ * the shallow union of every buffered update (later fields win) — updates
+ * carry different fields at different times (streamed rawInput snapshots,
+ * input-derived title/content, terminal status/rawOutput), so writing only
+ * the newest one would permanently drop the rest from the local file.
+ * `latestEntry` supplies the envelope (timestamp etc.) for the merged write.
+ * The merged update is a copy; entries shared with the API path are never
+ * mutated.
+ */
+interface BufferedToolUpdate {
+  latestEntry: StoredNotification;
+  mergedUpdate: Record<string, unknown>;
+  bufferedAt: number;
+}
+
 interface SessionState {
   context: SessionContext;
   chunkBuffer?: ChunkBuffer;
   lastAgentMessage?: string;
   currentTurnMessages: string[];
+  toolUpdateCache: Map<string, BufferedToolUpdate>;
 }
 
 export class SessionLogWriter {
+  /**
+   * When consecutive in-progress tool updates for one call span more than this
+   * window, the buffered union is written and a new window starts, so the
+   * local cache keeps periodic snapshots during active streaming instead of only
+   * the final one. Not a durability bound: the API log receives every update
+   * uncoalesced and the local file is a load cache. A buffered union is
+   * otherwise written on a terminal update, any non-tool event, or flushAll.
+   */
+  private static readonly TOOL_UPDATE_MAX_HOLD_MS = 2000;
   private static readonly FLUSH_DEBOUNCE_MS = 500;
   private static readonly FLUSH_MAX_INTERVAL_MS = 5000;
   private static readonly MAX_FLUSH_RETRIES = 10;
@@ -61,6 +88,7 @@ export class SessionLogWriter {
     const flushPromises: Promise<void>[] = [];
     for (const [sessionId, session] of this.sessions) {
       this.emitCoalescedMessage(sessionId, session);
+      this.flushToolUpdateCache(sessionId, session);
       flushPromises.push(this.flush(sessionId));
     }
     await Promise.all(flushPromises);
@@ -71,7 +99,11 @@ export class SessionLogWriter {
       return;
     }
 
-    this.sessions.set(sessionId, { context, currentTurnMessages: [] });
+    this.sessions.set(sessionId, {
+      context,
+      currentTurnMessages: [],
+      toolUpdateCache: new Map(),
+    });
 
     this.lastFlushAttemptTime.set(sessionId, Date.now());
 
@@ -144,7 +176,60 @@ export class SessionLogWriter {
         notification: message,
       };
 
-      this.writeToLocalCache(sessionId, entry);
+      // Coalesce the local cache: buffer in-progress tool_call_update
+      // snapshots (they re-send the full growing output) and write one merged
+      // update per toolCallId. Written on a terminal update, any non-tool
+      // event, or — during a long run of updates — once the hold window is
+      // exceeded. The API path is untouched, so the durable log keeps every
+      // update.
+      const tcu = this.toolCallUpdateInfo(message);
+      if (tcu && !tcu.terminal) {
+        const cache = session.toolUpdateCache;
+        const existing = cache.get(tcu.toolCallId);
+        if (
+          existing &&
+          Date.now() - existing.bufferedAt >
+            SessionLogWriter.TOOL_UPDATE_MAX_HOLD_MS
+        ) {
+          // Window exceeded: persist the union buffered so far and start a
+          // fresh window from this update. The read path merges across lines,
+          // so splitting the union over periodic snapshots loses nothing.
+          this.writeToLocalCache(sessionId, this.buildMergedEntry(existing));
+          cache.set(tcu.toolCallId, {
+            latestEntry: entry,
+            mergedUpdate: { ...tcu.update },
+            bufferedAt: Date.now(),
+          });
+        } else if (existing) {
+          Object.assign(existing.mergedUpdate, tcu.update);
+          existing.latestEntry = entry;
+        } else {
+          cache.set(tcu.toolCallId, {
+            latestEntry: entry,
+            mergedUpdate: { ...tcu.update },
+            bufferedAt: Date.now(),
+          });
+        }
+      } else {
+        if (tcu?.terminal) {
+          // Merge the terminal update into any buffered union so fields only
+          // carried by earlier snapshots (rawInput, edit diffs) still reach
+          // the local cache; later fields win, so status/rawOutput come from
+          // the terminal update itself.
+          const buffered = session.toolUpdateCache.get(tcu.toolCallId);
+          session.toolUpdateCache.delete(tcu.toolCallId);
+          if (buffered) {
+            Object.assign(buffered.mergedUpdate, tcu.update);
+            buffered.latestEntry = entry;
+            this.writeToLocalCache(sessionId, this.buildMergedEntry(buffered));
+          } else {
+            this.writeToLocalCache(sessionId, entry);
+          }
+        } else {
+          this.flushToolUpdateCache(sessionId, session);
+          this.writeToLocalCache(sessionId, entry);
+        }
+      }
 
       if (this.posthogAPI) {
         const pending = this.pendingEntries.get(sessionId) ?? [];
@@ -258,6 +343,49 @@ export class SessionLogWriter {
     return this.getSessionUpdateType(message) === "agent_message";
   }
 
+  private toolCallUpdateInfo(message: Record<string, unknown>): {
+    toolCallId: string;
+    terminal: boolean;
+    update: Record<string, unknown>;
+  } | null {
+    if (this.getSessionUpdateType(message) !== "tool_call_update") return null;
+    const params = message.params as Record<string, unknown> | undefined;
+    const update = params?.update as Record<string, unknown> | undefined;
+    const toolCallId = update?.toolCallId;
+    if (!update || typeof toolCallId !== "string") return null;
+    const status = update.status;
+    return {
+      toolCallId,
+      terminal: status === "completed" || status === "failed",
+      update,
+    };
+  }
+
+  /**
+   * Rebuild the buffered update's entry around the merged union. Builds a
+   * fresh object: the buffered `latestEntry` is also queued on the API path
+   * and must not be mutated.
+   */
+  private buildMergedEntry(buffered: BufferedToolUpdate): StoredNotification {
+    const { notification } = buffered.latestEntry;
+    return {
+      ...buffered.latestEntry,
+      notification: {
+        ...notification,
+        params: { ...notification.params, update: buffered.mergedUpdate },
+      },
+    };
+  }
+
+  /** Write any buffered tool-update unions to the local cache, in order. */
+  private flushToolUpdateCache(sessionId: string, session: SessionState): void {
+    if (session.toolUpdateCache.size === 0) return;
+    for (const buffered of session.toolUpdateCache.values()) {
+      this.writeToLocalCache(sessionId, this.buildMergedEntry(buffered));
+    }
+    session.toolUpdateCache.clear();
+  }
+
   private isAgentMessageChunk(message: Record<string, unknown>): boolean {
     return this.getSessionUpdateType(message) === "agent_message_chunk";
   }