From 485392fca848e27a84bcb26db1d654f3243ce3e2 Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Thu, 2 Jul 2026 05:43:07 +0300 Subject: [PATCH 1/6] perf(sessions): collapse superseded tool_call_update snapshots on load MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Agents re-send the full accumulated tool output on every tool_call_update, so a long-running tool (e.g. a property-based test run: 9.4k updates for one call, each growing to ~230KB) leaves hundreds of MB of near-identical snapshots in a loaded transcript — only the last is ever rendered. Keep just the latest update per toolCallId when converting stored entries to events. Measured across 7 large ai-gateway tasks: retained JS heap after GC dropped from 1220MB to 177MB (~7x), renderer RSS from 2.5GB to 0.9GB. Live streaming is unaffected (collapse runs only on the disk-load path). --- .../core/src/sessions/sessionEvents.test.ts | 71 +++++++++++++++++++ packages/core/src/sessions/sessionEvents.ts | 42 ++++++++++- 2 files changed, 112 insertions(+), 1 deletion(-) diff --git a/packages/core/src/sessions/sessionEvents.test.ts b/packages/core/src/sessions/sessionEvents.test.ts index 332f630179..a1fca16467 100644 --- a/packages/core/src/sessions/sessionEvents.test.ts +++ b/packages/core/src/sessions/sessionEvents.test.ts @@ -4,6 +4,7 @@ import { describe, expect, it } from "vitest"; import { makeAttachmentUri } from "./promptContent"; import { + collapseSupersededToolCallUpdates, convertStoredEntriesToEvents, extractUserPromptsFromEvents, hasSessionPromptEvent, @@ -288,3 +289,73 @@ describe("promptReferencesAbsoluteFolder", () => { expect(promptReferencesAbsoluteFolder("just text")).toBe(false); }); }); + +describe("collapseSupersededToolCallUpdates", () => { + const toolUpdate = (toolCallId: string, text: string): AcpMessage => + ({ + type: "acp_message", + ts: 1, + message: { + method: "session/update", + params: { + update: { + sessionUpdate: "tool_call_update", + toolCallId, + content: text, + }, + }, + }, + }) as unknown as AcpMessage; + + const other = (text: string): AcpMessage => + ({ + type: "acp_message", + ts: 1, + message: { + method: "session/update", + params: { + update: { + sessionUpdate: "agent_message_chunk", + content: { type: "text", text }, + }, + }, + }, + }) as unknown as AcpMessage; + + // biome-ignore lint/suspicious/noExplicitAny: test introspection + const sessionUpdate = (e: AcpMessage) => (e.message as any).params.update; + + it("keeps only the latest update per toolCallId, in its original position", () => { + const events = [ + toolUpdate("a", "a1"), + toolUpdate("a", "a2"), + other("hi"), + toolUpdate("a", "a3"), + ]; + const collapsed = collapseSupersededToolCallUpdates(events); + expect(collapsed).toHaveLength(2); + expect(sessionUpdate(collapsed[0]).sessionUpdate).toBe( + "agent_message_chunk", + ); + expect(sessionUpdate(collapsed[1]).content).toBe("a3"); + }); + + it("keeps the latest of each distinct toolCallId", () => { + const events = [ + toolUpdate("a", "a1"), + toolUpdate("b", "b1"), + toolUpdate("a", "a2"), + toolUpdate("b", "b2"), + ]; + const collapsed = collapseSupersededToolCallUpdates(events); + expect(collapsed.map((e) => sessionUpdate(e).content)).toEqual([ + "a2", + "b2", + ]); + }); + + it("leaves transcripts without tool updates untouched", () => { + const events = [other("one"), other("two")]; + expect(collapseSupersededToolCallUpdates(events)).toBe(events); + }); +}); diff --git a/packages/core/src/sessions/sessionEvents.ts b/packages/core/src/sessions/sessionEvents.ts index aa6cd24148..ca90336bbd 100644 --- a/packages/core/src/sessions/sessionEvents.ts +++ b/packages/core/src/sessions/sessionEvents.ts @@ -171,6 +171,46 @@ export function shellExecutesToContextBlocks( * Convert stored log entries to ACP messages. * Optionally prepends a user message with the task description. */ +function toolCallUpdateId(event: AcpMessage): string | undefined { + const msg = event.message; + if (!isJsonRpcNotification(msg) || msg.method !== "session/update") { + return undefined; + } + const update = (msg.params as SessionNotification | undefined)?.update as + | { sessionUpdate?: string; toolCallId?: string } + | undefined; + if (update?.sessionUpdate !== "tool_call_update") return undefined; + return typeof update.toolCallId === "string" ? update.toolCallId : undefined; +} + +/** + * Drop superseded `tool_call_update` snapshots, keeping only the latest per + * `toolCallId`. Agents re-send the full accumulated tool output on every + * update, so a long-running tool leaves thousands of near-identical growing + * snapshots in a loaded transcript — only the last is ever rendered. Keeping + * just that one (in its original position, carrying the full final output) + * caps resident memory: a 9k-update tool run drops from ~312MB to ~3MB of + * tool content with no visible change. + */ +export function collapseSupersededToolCallUpdates( + events: AcpMessage[], +): AcpMessage[] { + const lastIndexById = new Map(); + for (let i = 0; i < events.length; i++) { + const id = toolCallUpdateId(events[i]); + if (id !== undefined) lastIndexById.set(id, i); + } + if (lastIndexById.size === 0) return events; + + const collapsed: AcpMessage[] = []; + for (let i = 0; i < events.length; i++) { + const id = toolCallUpdateId(events[i]); + if (id !== undefined && lastIndexById.get(id) !== i) continue; + collapsed.push(events[i]); + } + return collapsed; +} + export function convertStoredEntriesToEvents( entries: StoredLogEntry[], taskDescription?: string, @@ -188,7 +228,7 @@ export function convertStoredEntriesToEvents( events.push(storedEntryToAcpMessage(entry)); } - return events; + return collapseSupersededToolCallUpdates(events); } /** From 341473676f672bc90d2f4c376df79a9f1bd9930c Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Thu, 2 Jul 2026 05:54:27 +0300 Subject: [PATCH 2/6] perf(sessions): collapse tool_call_update snapshots server-side before transfer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The renderer-side collapse capped retained memory but the full redundant log still crossed IPC and got parsed in the renderer, spiking heap to ~2GB and janking the load. Add readLocalLogsCollapsed on the workspace-server: it drops superseded tool_call_update lines (line-based, no JSON.parse) and returns only the latest per toolCallId, keeping the original line count for resume/gap tracking. fetchSessionLogs uses it when available. Measured on the 327MB property-based-test task: cold-open peak heap ~2GB -> ~190MB (no spike); renderer receives ~3MB instead of 327MB. Load wall-clock is still gated by Node reading+scanning the 327MB on disk (~2.9s) — the source fix (agent emitting tool_call_update deltas) is what shrinks the log itself. --- apps/code/src/main/di/container.ts | 2 + packages/core/src/sessions/sessionService.ts | 47 +++++++++++---- .../src/services/local-logs/identifiers.ts | 7 +++ .../src/services/local-logs/schemas.ts | 10 ++++ .../src/services/local-logs/service.ts | 50 ++++++++++++++++ .../services/local-logs/serviceTail.test.ts | 57 +++++++++++++++++++ packages/workspace-server/src/trpc.ts | 9 +++ 7 files changed, 172 insertions(+), 10 deletions(-) diff --git a/apps/code/src/main/di/container.ts b/apps/code/src/main/di/container.ts index d2381f0276..ab47b06f73 100644 --- a/apps/code/src/main/di/container.ts +++ b/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) => diff --git a/packages/core/src/sessions/sessionService.ts b/packages/core/src/sessions/sessionService.ts index 5b5af8740e..0c2a73bb43 100644 --- a/packages/core/src/sessions/sessionService.ts +++ b/packages/core/src/sessions/sessionService.ts @@ -168,6 +168,9 @@ export interface SessionTrpc { }; logs: { readLocalLogs: TrpcQuery; + /** Optional: collapses superseded tool_call_update snapshots server-side so + * a tool-heavy log doesn't ship its full redundant history over IPC. */ + readLocalLogsCollapsed?: TrpcQuery; /** Optional: only the Electron host exposes the tail read. Core feature- * detects and falls back to a full read when it's absent. */ readLocalLogsTail?: TrpcQuery; @@ -4860,16 +4863,40 @@ export class SessionService { if (taskRunId) { try { - const localContent = await this.d.trpc.logs.readLocalLogs.query({ - taskRunId, - }); - if (localContent?.trim()) { - localResult = this.parseLogContent(localContent); - if ( - !options.minEntryCount || - localResult.totalLineCount >= options.minEntryCount - ) { - return localResult; + const collapsedQuery = this.d.trpc.logs.readLocalLogsCollapsed; + if (collapsedQuery) { + // Collapse tool_call_update snapshots server-side so a tool-heavy + // log doesn't ship its full redundant history over IPC. The parsed + // line count is the collapsed count, so use the server's original + // totalLineCount for resume/gap tracking. + const res = (await collapsedQuery.query({ taskRunId })) as { + content: string; + totalLineCount: number; + } | null; + if (res?.content?.trim()) { + localResult = { + ...this.parseLogContent(res.content), + totalLineCount: res.totalLineCount, + }; + if ( + !options.minEntryCount || + localResult.totalLineCount >= options.minEntryCount + ) { + return localResult; + } + } + } else { + const localContent = await this.d.trpc.logs.readLocalLogs.query({ + taskRunId, + }); + if (localContent?.trim()) { + localResult = this.parseLogContent(localContent); + if ( + !options.minEntryCount || + localResult.totalLineCount >= options.minEntryCount + ) { + return localResult; + } } } } catch { diff --git a/packages/workspace-server/src/services/local-logs/identifiers.ts b/packages/workspace-server/src/services/local-logs/identifiers.ts index 212ef8155b..517b1bb4c6 100644 --- a/packages/workspace-server/src/services/local-logs/identifiers.ts +++ b/packages/workspace-server/src/services/local-logs/identifiers.ts @@ -3,6 +3,13 @@ export const LOGS_SERVICE = Symbol.for("posthog.workspace.logsService"); export interface ILogsService { fetchS3Logs(logUrl: string): Promise; readLocalLogs(taskRunId: string): Promise; + /** + * Like `readLocalLogs`, but collapses superseded `tool_call_update` snapshots + * before returning. `totalLineCount` is the original pre-collapse line count. + */ + readLocalLogsCollapsed( + taskRunId: string, + ): Promise<{ content: string; totalLineCount: number } | null>; /** * Read only the last `maxBytes` of the log for a fast initial paint. Returns * `truncated: true` when older history was skipped (the partial first line is diff --git a/packages/workspace-server/src/services/local-logs/schemas.ts b/packages/workspace-server/src/services/local-logs/schemas.ts index 708671d973..dc6f36eace 100644 --- a/packages/workspace-server/src/services/local-logs/schemas.ts +++ b/packages/workspace-server/src/services/local-logs/schemas.ts @@ -6,6 +6,16 @@ export const fetchS3LogsOutput = z.string().nullable(); export const readLocalLogsInput = z.object({ taskRunId: z.string().min(1) }); export const readLocalLogsOutput = z.string().nullable(); +export const readLocalLogsCollapsedInput = z.object({ + taskRunId: z.string().min(1), +}); +export const readLocalLogsCollapsedOutput = z + .object({ + content: z.string(), + totalLineCount: z.number().int().nonnegative(), + }) + .nullable(); + export const readLocalLogsTailInput = z.object({ taskRunId: z.string().min(1), maxBytes: z.number().int().positive(), diff --git a/packages/workspace-server/src/services/local-logs/service.ts b/packages/workspace-server/src/services/local-logs/service.ts index 058e142785..500f65bdf0 100644 --- a/packages/workspace-server/src/services/local-logs/service.ts +++ b/packages/workspace-server/src/services/local-logs/service.ts @@ -8,6 +8,40 @@ import type { ILogsService } from "./identifiers"; const DATA_DIR = ".posthog-code"; +const TOOL_CALL_UPDATE_MARKER = '"sessionUpdate":"tool_call_update"'; +const TOOL_CALL_ID_RE = /"toolCallId":"([^"]+)"/; + +/** + * Drop superseded `tool_call_update` lines, keeping only the last per + * `toolCallId`. Agents re-send the full accumulated tool output on every + * update, so a long-running tool leaves hundreds of MB of near-identical + * growing snapshots in the log; only the last is ever rendered. Collapsing + * here (before the log crosses to the renderer) keeps the transfer + parse + * proportional to the real content, not the redundant history. Whole lines are + * dropped, so the result is still valid NDJSON. Line-based (no JSON.parse) to + * stay cheap on a 300MB log. + */ +function collapseToolCallUpdateLines(ndjson: string): string { + const lines = ndjson.split("\n"); + const lastIndexById = new Map(); + for (let i = 0; i < lines.length; i++) { + if (!lines[i].includes(TOOL_CALL_UPDATE_MARKER)) continue; + const id = lines[i].match(TOOL_CALL_ID_RE)?.[1]; + if (id) lastIndexById.set(id, i); + } + if (lastIndexById.size === 0) return ndjson; + + const kept: string[] = []; + for (let i = 0; i < lines.length; i++) { + if (lines[i].includes(TOOL_CALL_UPDATE_MARKER)) { + const id = lines[i].match(TOOL_CALL_ID_RE)?.[1]; + if (id && lastIndexById.get(id) !== i) continue; + } + kept.push(lines[i]); + } + return kept.join("\n"); +} + interface WriteState { pending: string | undefined; lastWritten: string | undefined; @@ -52,6 +86,22 @@ export class LocalLogsService implements ILogsService { } } + /** + * Like `readLocalLogs`, but collapses superseded `tool_call_update` snapshots + * before returning so a tool-heavy log doesn't ship its full redundant + * history across IPC. `totalLineCount` is the original (pre-collapse) line + * count, which callers use for resume/gap tracking. + */ + async readLocalLogsCollapsed( + taskRunId: string, + ): Promise<{ content: string; totalLineCount: number } | null> { + const raw = await this.readLocalLogs(taskRunId); + if (raw === null) return null; + const trimmed = raw.trim(); + const totalLineCount = trimmed ? trimmed.split("\n").length : 0; + return { content: collapseToolCallUpdateLines(raw), totalLineCount }; + } + async readLocalLogsTail( taskRunId: string, maxBytes: number, diff --git a/packages/workspace-server/src/services/local-logs/serviceTail.test.ts b/packages/workspace-server/src/services/local-logs/serviceTail.test.ts index 446d96ed92..c2be22e70a 100644 --- a/packages/workspace-server/src/services/local-logs/serviceTail.test.ts +++ b/packages/workspace-server/src/services/local-logs/serviceTail.test.ts @@ -77,3 +77,60 @@ describe("LocalLogsService.readLocalLogsTail", () => { ).toBeNull(); }); }); + +describe("LocalLogsService.readLocalLogsCollapsed", () => { + let tmpHome: string; + + beforeEach(() => { + tmpHome = fs.mkdtempSync(path.join(os.tmpdir(), "phlogs-")); + vi.spyOn(os, "homedir").mockReturnValue(tmpHome); + fs.mkdirSync(path.join(tmpHome, ".posthog-code", "sessions", RUN), { + recursive: true, + }); + }); + + afterEach(() => { + vi.restoreAllMocks(); + fs.rmSync(tmpHome, { recursive: true, force: true }); + }); + + const logPath = () => + path.join(tmpHome, ".posthog-code", "sessions", RUN, "logs.ndjson"); + + const toolUpdate = (toolCallId: string, out: string) => + JSON.stringify({ + notification: { + method: "session/update", + params: { + update: { sessionUpdate: "tool_call_update", toolCallId, out }, + }, + }, + }); + + it("drops superseded tool_call_update lines, keeps latest per toolCallId, preserves original line count", async () => { + const lines = [ + `{"i":0}`, + toolUpdate("a", "a1"), + toolUpdate("a", "a2"), + `{"i":1}`, + toolUpdate("a", "a3"), + ]; + fs.writeFileSync(logPath(), `${lines.join("\n")}\n`); + + const res = await new LocalLogsService().readLocalLogsCollapsed(RUN); + + expect(res?.totalLineCount).toBe(5); + const kept = res?.content.trim().split("\n") ?? []; + // both non-tool lines + only the latest "a" update remain + expect(kept).toHaveLength(3); + expect(kept[2]).toContain(`"out":"a3"`); + expect(res?.content).not.toContain(`"out":"a1"`); + expect(res?.content).not.toContain(`"out":"a2"`); + }); + + it("returns null when the log doesn't exist", async () => { + expect( + await new LocalLogsService().readLocalLogsCollapsed("missing"), + ).toBeNull(); + }); +}); diff --git a/packages/workspace-server/src/trpc.ts b/packages/workspace-server/src/trpc.ts index 2769c52014..a1241a576c 100644 --- a/packages/workspace-server/src/trpc.ts +++ b/packages/workspace-server/src/trpc.ts @@ -130,6 +130,8 @@ import { countLocalLogEntriesInput, countLocalLogEntriesOutput, deleteLocalLogCacheInput, + readLocalLogsCollapsedInput, + readLocalLogsCollapsedOutput, readLocalLogsInput, readLocalLogsOutput, readLocalLogsTailInput, @@ -856,6 +858,13 @@ export function createAppRouter({ localLogsService().readLocalLogs(input.taskRunId), ), + readCollapsed: t.procedure + .input(readLocalLogsCollapsedInput) + .output(readLocalLogsCollapsedOutput) + .query(({ input }) => + localLogsService().readLocalLogsCollapsed(input.taskRunId), + ), + readTail: t.procedure .input(readLocalLogsTailInput) .output(readLocalLogsTailOutput) From acbe8f60d11709c2395aa03c5f2a47a1aa562a6e Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Thu, 2 Jul 2026 06:06:04 +0300 Subject: [PATCH 3/6] perf(agent): coalesce in-progress tool_call_update snapshots in the local log MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The session log writer appended every tool_call_update, and agents re-send the full accumulated tool output on each one, so a long-running tool wrote hundreds of MB of near-identical growing snapshots to the local cache. Buffer the latest in-progress update per toolCallId and write only that — flushed by a terminal (completed/failed) update, any non-tool event, a 2s max-hold window (bounds crash loss), or flushAll on shutdown. The API/S3 path is unchanged; the reader already collapses on load. New local logs are born proportional to real content instead of O(updates x size). Unit-tested; needs a live agent run to confirm born-small logs end to end. --- packages/agent/src/session-log-writer.test.ts | 77 ++++++++++++++++++ packages/agent/src/session-log-writer.ts | 79 ++++++++++++++++++- 2 files changed, 154 insertions(+), 2 deletions(-) diff --git a/packages/agent/src/session-log-writer.test.ts b/packages/agent/src/session-log-writer.test.ts index eec8b3ed8a..831835cb69 100644 --- a/packages/agent/src/session-log-writer.test.ts +++ b/packages/agent/src/session-log-writer.test.ts @@ -514,3 +514,80 @@ 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[]> => { + 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) => + makeSessionUpdate("tool_call_update", { toolCallId: "a", ...extra }); + + const sessionUpdateOf = (e: Record) => + // biome-ignore lint/suspicious/noExplicitAny: test introspection + (e.notification as any).params.update; + + it("writes only the latest in-progress update, 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 latest, 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 supersedes buffered in-progress snapshots", 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("flushAll persists a still-buffered update", async () => { + writer.appendRawLine(RUN, update({ content: "a1" })); + writer.appendRawLine(RUN, update({ content: "a2" })); + await writer.flushAll(); + + const log = await readLog(); + expect(log).toHaveLength(1); + expect(sessionUpdateOf(log[0]).content).toBe("a2"); + }); +}); diff --git a/packages/agent/src/session-log-writer.ts b/packages/agent/src/session-log-writer.ts index be63d75e36..39bf8fc324 100644 --- a/packages/agent/src/session-log-writer.ts +++ b/packages/agent/src/session-log-writer.ts @@ -26,9 +26,25 @@ interface SessionState { chunkBuffer?: ChunkBuffer; lastAgentMessage?: string; currentTurnMessages: string[]; + /** + * Latest in-progress `tool_call_update` per toolCallId, not yet written to + * the local cache. Agents re-send the full accumulated tool output on every + * update; only the last before a terminal/other event needs persisting, so + * we coalesce here to keep the on-disk log proportional to real content. + */ + toolUpdateCache: Map< + string, + { entry: StoredNotification; bufferedAt: number } + >; } export class SessionLogWriter { + /** + * Max wall-clock a coalesced in-progress tool update is held before being + * written to the local cache anyway. Bounds crash loss to this window while + * still collapsing rapid re-sends of a growing tool output. + */ + 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 +77,7 @@ export class SessionLogWriter { const flushPromises: Promise[] = []; 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 +88,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 +165,37 @@ export class SessionLogWriter { notification: message, }; - this.writeToLocalCache(sessionId, entry); + // Coalesce the local cache: hold in-progress tool_call_update snapshots + // (they re-send the full growing output) and write only the latest per + // toolCallId. A terminal update, any non-tool event, or the hold window + // elapsing flushes them. The API path is untouched. + 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 + ) { + this.writeToLocalCache(sessionId, existing.entry); + cache.set(tcu.toolCallId, { entry, bufferedAt: Date.now() }); + } else { + cache.set(tcu.toolCallId, { + entry, + bufferedAt: existing?.bufferedAt ?? Date.now(), + }); + } + } else { + if (tcu?.terminal) { + // The terminal update carries the full final output and supersedes + // any buffered in-progress snapshot for this call. + session.toolUpdateCache.delete(tcu.toolCallId); + } else { + this.flushToolUpdateCache(sessionId, session); + } + this.writeToLocalCache(sessionId, entry); + } if (this.posthogAPI) { const pending = this.pendingEntries.get(sessionId) ?? []; @@ -258,6 +309,30 @@ export class SessionLogWriter { return this.getSessionUpdateType(message) === "agent_message"; } + private toolCallUpdateInfo( + message: Record, + ): { toolCallId: string; terminal: boolean } | null { + if (this.getSessionUpdateType(message) !== "tool_call_update") return null; + const params = message.params as Record | undefined; + const update = params?.update as Record | undefined; + const toolCallId = update?.toolCallId; + if (typeof toolCallId !== "string") return null; + const status = update?.status; + return { + toolCallId, + terminal: status === "completed" || status === "failed", + }; + } + + /** Write any buffered in-progress tool updates to the local cache, in order. */ + private flushToolUpdateCache(sessionId: string, session: SessionState): void { + if (session.toolUpdateCache.size === 0) return; + for (const { entry } of session.toolUpdateCache.values()) { + this.writeToLocalCache(sessionId, entry); + } + session.toolUpdateCache.clear(); + } + private isAgentMessageChunk(message: Record): boolean { return this.getSessionUpdateType(message) === "agent_message_chunk"; } From 6a78d436f808dcad31d640082ac7e41126b49e3f Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Thu, 2 Jul 2026 06:19:47 +0300 Subject: [PATCH 4/6] docs(sessions): trim repeated tool_call_update-collapse rationale --- packages/agent/src/session-log-writer.ts | 8 ++------ packages/core/src/sessions/sessionService.ts | 6 ++---- .../src/services/local-logs/service.ts | 19 +++++-------------- 3 files changed, 9 insertions(+), 24 deletions(-) diff --git a/packages/agent/src/session-log-writer.ts b/packages/agent/src/session-log-writer.ts index 39bf8fc324..e5b9572bfa 100644 --- a/packages/agent/src/session-log-writer.ts +++ b/packages/agent/src/session-log-writer.ts @@ -26,12 +26,8 @@ interface SessionState { chunkBuffer?: ChunkBuffer; lastAgentMessage?: string; currentTurnMessages: string[]; - /** - * Latest in-progress `tool_call_update` per toolCallId, not yet written to - * the local cache. Agents re-send the full accumulated tool output on every - * update; only the last before a terminal/other event needs persisting, so - * we coalesce here to keep the on-disk log proportional to real content. - */ + /** Latest in-progress `tool_call_update` per toolCallId, awaiting a + * coalesced write to the local cache (see appendRawLine). */ toolUpdateCache: Map< string, { entry: StoredNotification; bufferedAt: number } diff --git a/packages/core/src/sessions/sessionService.ts b/packages/core/src/sessions/sessionService.ts index 0c2a73bb43..18b205be32 100644 --- a/packages/core/src/sessions/sessionService.ts +++ b/packages/core/src/sessions/sessionService.ts @@ -4865,10 +4865,8 @@ export class SessionService { try { const collapsedQuery = this.d.trpc.logs.readLocalLogsCollapsed; if (collapsedQuery) { - // Collapse tool_call_update snapshots server-side so a tool-heavy - // log doesn't ship its full redundant history over IPC. The parsed - // line count is the collapsed count, so use the server's original - // totalLineCount for resume/gap tracking. + // The collapsed content has fewer lines, so use the server's + // original totalLineCount for resume/gap tracking, not the parsed one. const res = (await collapsedQuery.query({ taskRunId })) as { content: string; totalLineCount: number; diff --git a/packages/workspace-server/src/services/local-logs/service.ts b/packages/workspace-server/src/services/local-logs/service.ts index 500f65bdf0..4956ea0d52 100644 --- a/packages/workspace-server/src/services/local-logs/service.ts +++ b/packages/workspace-server/src/services/local-logs/service.ts @@ -12,14 +12,11 @@ const TOOL_CALL_UPDATE_MARKER = '"sessionUpdate":"tool_call_update"'; const TOOL_CALL_ID_RE = /"toolCallId":"([^"]+)"/; /** - * Drop superseded `tool_call_update` lines, keeping only the last per - * `toolCallId`. Agents re-send the full accumulated tool output on every - * update, so a long-running tool leaves hundreds of MB of near-identical - * growing snapshots in the log; only the last is ever rendered. Collapsing - * here (before the log crosses to the renderer) keeps the transfer + parse - * proportional to the real content, not the redundant history. Whole lines are - * dropped, so the result is still valid NDJSON. Line-based (no JSON.parse) to - * stay cheap on a 300MB log. + * Drop superseded `tool_call_update` lines (keep the last per `toolCallId`) + * before the log crosses to the renderer. Agents re-send the full accumulated + * tool output on every update, so the transfer + parse would otherwise carry + * hundreds of MB of redundant snapshots. Whole lines are dropped so the result + * stays valid NDJSON; line-based (no JSON.parse) to stay cheap on a 300MB log. */ function collapseToolCallUpdateLines(ndjson: string): string { const lines = ndjson.split("\n"); @@ -86,12 +83,6 @@ export class LocalLogsService implements ILogsService { } } - /** - * Like `readLocalLogs`, but collapses superseded `tool_call_update` snapshots - * before returning so a tool-heavy log doesn't ship its full redundant - * history across IPC. `totalLineCount` is the original (pre-collapse) line - * count, which callers use for resume/gap tracking. - */ async readLocalLogsCollapsed( taskRunId: string, ): Promise<{ content: string; totalLineCount: number } | null> { From c05af114faa7fec2d2421caf1c554988e7336277 Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Thu, 2 Jul 2026 06:32:46 +0300 Subject: [PATCH 5/6] docs(agent): correct tool-update hold-window comment (not a crash bound) --- packages/agent/src/session-log-writer.ts | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/packages/agent/src/session-log-writer.ts b/packages/agent/src/session-log-writer.ts index e5b9572bfa..af278bd8c3 100644 --- a/packages/agent/src/session-log-writer.ts +++ b/packages/agent/src/session-log-writer.ts @@ -36,9 +36,12 @@ interface SessionState { export class SessionLogWriter { /** - * Max wall-clock a coalesced in-progress tool update is held before being - * written to the local cache anyway. Bounds crash loss to this window while - * still collapsing rapid re-sends of a growing tool output. + * When consecutive in-progress tool updates for one call span more than this + * window, the buffered snapshot 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 snapshot 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; @@ -163,8 +166,9 @@ export class SessionLogWriter { // Coalesce the local cache: hold in-progress tool_call_update snapshots // (they re-send the full growing output) and write only the latest per - // toolCallId. A terminal update, any non-tool event, or the hold window - // elapsing flushes them. The API path is untouched. + // toolCallId. Flushed 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; From a3bf5ab7c1ce8ea23775c6a8aee65fd31018abff Mon Sep 17 00:00:00 2001 From: Adam Bowker Date: Thu, 2 Jul 2026 14:08:31 -0400 Subject: [PATCH 6/6] fix(sessions): wire collapsed log read through the host router; merge, don't drop, superseded tool updates MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two fixes on top of the tool_call_update collapse: Host-router wiring: the renderer's SessionService reaches logs through the host router, which lacked the readLocalLogsCollapsed procedure. A tRPC proxy client is truthy for any path, so the feature-detect passed and every local read threw NOT_FOUND at call time, silently falling back to S3 — and loading empty transcripts for tasks with no cloud logs. Adds the procedure, and fetchSessionLogs now falls back to the plain read when the collapsed call fails instead of misreporting the local log as unreadable. Merge semantics: the conversation reducer Object.assigns every tool_call_update into the tool call, and updates carry disjoint fields (streamed rawInput snapshots, input-derived title/diff content, then a terminal update with only status/rawOutput). Keeping only the last update left ~half of all tool calls with no rawInput anywhere and stripped Edit/Write diffs on reload. All three layers (core read, workspace-server transfer, agent write buffer) now merge updates per toolCallId — shallow, later fields win — which reproduces exactly the state a full replay builds. The agent's API path still receives every update unmutated; the transfer layer now parses tool-update lines (other lines pass through untouched), which also removes the regex mis-bucketing risk. Validated replay-equivalent against the five biggest real local logs (up to 88MB): merged output deep-equals a full replay per tool call, with ~95% of the size win retained (88MB -> 17MB). Generated-By: PostHog Code Task-Id: fe96a3c1-92c7-41ef-9bd1-f2b890ebe566 --- packages/agent/src/session-log-writer.test.ts | 104 ++++++++++++++++- packages/agent/src/session-log-writer.ts | 109 +++++++++++++----- .../core/src/sessions/sessionEvents.test.ts | 71 +++++++++++- packages/core/src/sessions/sessionEvents.ts | 76 +++++++++--- packages/core/src/sessions/sessionService.ts | 88 ++++++++------ .../host-router/src/routers/logs.router.ts | 11 ++ .../src/services/local-logs/service.ts | 97 +++++++++++++--- .../services/local-logs/serviceTail.test.ts | 53 ++++++++- 8 files changed, 509 insertions(+), 100 deletions(-) diff --git a/packages/agent/src/session-log-writer.test.ts b/packages/agent/src/session-log-writer.test.ts index 831835cb69..182fab39b6 100644 --- a/packages/agent/src/session-log-writer.test.ts +++ b/packages/agent/src/session-log-writer.test.ts @@ -554,11 +554,11 @@ describe("SessionLogWriter — local-cache tool_call_update coalescing", () => { // biome-ignore lint/suspicious/noExplicitAny: test introspection (e.notification as any).params.update; - it("writes only the latest in-progress update, flushed by a non-tool event", async () => { + 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 latest, then writes itself + // a non-tool event flushes the buffered union, then writes itself writer.appendRawLine(RUN, makeSessionUpdate("agent_message")); const log = await readLog(); @@ -567,7 +567,7 @@ describe("SessionLogWriter — local-cache tool_call_update coalescing", () => { expect(sessionUpdateOf(log[1]).sessionUpdate).toBe("agent_message"); }); - it("a terminal update supersedes buffered in-progress snapshots", async () => { + 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( @@ -581,13 +581,107 @@ describe("SessionLogWriter — local-cache tool_call_update coalescing", () => { expect(sessionUpdateOf(log[0]).status).toBe("completed"); }); - it("flushAll persists a still-buffered update", async () => { - writer.appendRawLine(RUN, update({ content: "a1" })); + 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 } }; + }[]; + 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(); + } }); }); diff --git a/packages/agent/src/session-log-writer.ts b/packages/agent/src/session-log-writer.ts index af278bd8c3..ef227e47d0 100644 --- a/packages/agent/src/session-log-writer.ts +++ b/packages/agent/src/session-log-writer.ts @@ -21,26 +21,38 @@ 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; + bufferedAt: number; +} + interface SessionState { context: SessionContext; chunkBuffer?: ChunkBuffer; lastAgentMessage?: string; currentTurnMessages: string[]; - /** Latest in-progress `tool_call_update` per toolCallId, awaiting a - * coalesced write to the local cache (see appendRawLine). */ - toolUpdateCache: Map< - string, - { entry: StoredNotification; bufferedAt: number } - >; + toolUpdateCache: Map; } export class SessionLogWriter { /** * When consecutive in-progress tool updates for one call span more than this - * window, the buffered snapshot is written and a new window starts, so the + * 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 snapshot is + * 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; @@ -164,11 +176,12 @@ export class SessionLogWriter { notification: message, }; - // Coalesce the local cache: hold in-progress tool_call_update snapshots - // (they re-send the full growing output) and write only the latest per - // toolCallId. Flushed 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. + // 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; @@ -178,23 +191,44 @@ export class SessionLogWriter { Date.now() - existing.bufferedAt > SessionLogWriter.TOOL_UPDATE_MAX_HOLD_MS ) { - this.writeToLocalCache(sessionId, existing.entry); - cache.set(tcu.toolCallId, { entry, bufferedAt: Date.now() }); + // 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, { - entry, - bufferedAt: existing?.bufferedAt ?? Date.now(), + latestEntry: entry, + mergedUpdate: { ...tcu.update }, + bufferedAt: Date.now(), }); } } else { if (tcu?.terminal) { - // The terminal update carries the full final output and supersedes - // any buffered in-progress snapshot for this call. + // 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); } - this.writeToLocalCache(sessionId, entry); } if (this.posthogAPI) { @@ -309,26 +343,45 @@ export class SessionLogWriter { return this.getSessionUpdateType(message) === "agent_message"; } - private toolCallUpdateInfo( - message: Record, - ): { toolCallId: string; terminal: boolean } | null { + private toolCallUpdateInfo(message: Record): { + toolCallId: string; + terminal: boolean; + update: Record; + } | null { if (this.getSessionUpdateType(message) !== "tool_call_update") return null; const params = message.params as Record | undefined; const update = params?.update as Record | undefined; const toolCallId = update?.toolCallId; - if (typeof toolCallId !== "string") return null; - const status = update?.status; + 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 in-progress tool updates to the local cache, in order. */ + /** 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 { entry } of session.toolUpdateCache.values()) { - this.writeToLocalCache(sessionId, entry); + for (const buffered of session.toolUpdateCache.values()) { + this.writeToLocalCache(sessionId, this.buildMergedEntry(buffered)); } session.toolUpdateCache.clear(); } diff --git a/packages/core/src/sessions/sessionEvents.test.ts b/packages/core/src/sessions/sessionEvents.test.ts index a1fca16467..4706956a55 100644 --- a/packages/core/src/sessions/sessionEvents.test.ts +++ b/packages/core/src/sessions/sessionEvents.test.ts @@ -291,7 +291,10 @@ describe("promptReferencesAbsoluteFolder", () => { }); describe("collapseSupersededToolCallUpdates", () => { - const toolUpdate = (toolCallId: string, text: string): AcpMessage => + const toolUpdateFields = ( + toolCallId: string, + fields: Record, + ): AcpMessage => ({ type: "acp_message", ts: 1, @@ -301,12 +304,15 @@ describe("collapseSupersededToolCallUpdates", () => { update: { sessionUpdate: "tool_call_update", toolCallId, - content: text, + ...fields, }, }, }, }) as unknown as AcpMessage; + const toolUpdate = (toolCallId: string, text: string): AcpMessage => + toolUpdateFields(toolCallId, { content: text }); + const other = (text: string): AcpMessage => ({ type: "acp_message", @@ -325,7 +331,7 @@ describe("collapseSupersededToolCallUpdates", () => { // biome-ignore lint/suspicious/noExplicitAny: test introspection const sessionUpdate = (e: AcpMessage) => (e.message as any).params.update; - it("keeps only the latest update per toolCallId, in its original position", () => { + it("collapses to one update per toolCallId, at the last update's position", () => { const events = [ toolUpdate("a", "a1"), toolUpdate("a", "a2"), @@ -340,7 +346,7 @@ describe("collapseSupersededToolCallUpdates", () => { expect(sessionUpdate(collapsed[1]).content).toBe("a3"); }); - it("keeps the latest of each distinct toolCallId", () => { + it("collapses each distinct toolCallId independently", () => { const events = [ toolUpdate("a", "a1"), toolUpdate("b", "b1"), @@ -358,4 +364,61 @@ describe("collapseSupersededToolCallUpdates", () => { const events = [other("one"), other("two")]; expect(collapseSupersededToolCallUpdates(events)).toBe(events); }); + + it("merges fields across updates so nothing a replay would keep is lost", () => { + // Mirrors the real emission shape: streamed rawInput snapshots, then an + // input-complete update with title/content, then a terminal update that + // carries only status/rawOutput. + const events = [ + toolUpdateFields("a", { rawInput: { command: "ls" } }), + toolUpdateFields("a", { + rawInput: { command: "ls -la" }, + title: "List files", + content: "input-derived", + }), + toolUpdateFields("a", { status: "completed", rawOutput: "done" }), + ]; + const collapsed = collapseSupersededToolCallUpdates(events); + expect(collapsed).toHaveLength(1); + expect(sessionUpdate(collapsed[0])).toEqual({ + sessionUpdate: "tool_call_update", + toolCallId: "a", + rawInput: { command: "ls -la" }, + title: "List files", + content: "input-derived", + status: "completed", + rawOutput: "done", + }); + }); + + it("later fields win when re-sent (matching the reducer's Object.assign)", () => { + const events = [ + toolUpdateFields("a", { content: "stale", status: "in_progress" }), + toolUpdateFields("a", { content: "fresh", status: "completed" }), + ]; + const collapsed = collapseSupersededToolCallUpdates(events); + expect(collapsed).toHaveLength(1); + expect(sessionUpdate(collapsed[0]).content).toBe("fresh"); + expect(sessionUpdate(collapsed[0]).status).toBe("completed"); + }); + + it("keeps a single-update call by reference, no synthetic clone", () => { + const only = toolUpdate("a", "a1"); + const events = [other("hi"), only]; + const collapsed = collapseSupersededToolCallUpdates(events); + expect(collapsed).toHaveLength(2); + expect(collapsed[1]).toBe(only); + }); + + it("does not mutate the original (frozen) events when merging", () => { + const first = toolUpdateFields("a", { rawInput: { command: "ls" } }); + const last = toolUpdateFields("a", { status: "completed" }); + Object.freeze(first); + Object.freeze(last); + const collapsed = collapseSupersededToolCallUpdates([first, last]); + expect(sessionUpdate(first)).not.toHaveProperty("status"); + expect(sessionUpdate(last)).not.toHaveProperty("rawInput"); + expect(sessionUpdate(collapsed[0])).toHaveProperty("rawInput"); + expect(sessionUpdate(collapsed[0]).status).toBe("completed"); + }); }); diff --git a/packages/core/src/sessions/sessionEvents.ts b/packages/core/src/sessions/sessionEvents.ts index ca90336bbd..7370397428 100644 --- a/packages/core/src/sessions/sessionEvents.ts +++ b/packages/core/src/sessions/sessionEvents.ts @@ -171,41 +171,89 @@ export function shellExecutesToContextBlocks( * Convert stored log entries to ACP messages. * Optionally prepends a user message with the task description. */ -function toolCallUpdateId(event: AcpMessage): string | undefined { +function toolCallUpdateOf( + event: AcpMessage, +): (Record & { toolCallId: string }) | undefined { const msg = event.message; if (!isJsonRpcNotification(msg) || msg.method !== "session/update") { return undefined; } const update = (msg.params as SessionNotification | undefined)?.update as - | { sessionUpdate?: string; toolCallId?: string } + | { sessionUpdate?: string; toolCallId?: unknown } | undefined; if (update?.sessionUpdate !== "tool_call_update") return undefined; - return typeof update.toolCallId === "string" ? update.toolCallId : undefined; + if (typeof update.toolCallId !== "string") return undefined; + return update as Record & { toolCallId: string }; +} + +/** Rebuild a (frozen) tool_call_update event around a replacement update. */ +function withToolCallUpdate( + event: AcpMessage, + update: Record, +): AcpMessage { + const msg = event.message as { params?: SessionNotification }; + return Object.freeze({ + ...event, + message: { + ...msg, + params: { ...msg.params, update }, + } as JsonRpcMessage, + }); } /** - * Drop superseded `tool_call_update` snapshots, keeping only the latest per - * `toolCallId`. Agents re-send the full accumulated tool output on every - * update, so a long-running tool leaves thousands of near-identical growing - * snapshots in a loaded transcript — only the last is ever rendered. Keeping - * just that one (in its original position, carrying the full final output) - * caps resident memory: a 9k-update tool run drops from ~312MB to ~3MB of - * tool content with no visible change. + * Collapse superseded `tool_call_update` snapshots into one merged update per + * `toolCallId`, kept at the last update's position. Agents re-send the full + * accumulated tool output on every update, so a long-running tool leaves + * thousands of near-identical growing snapshots in a loaded transcript; one + * 9k-update tool run carried ~312MB of tool content of which the merged + * result is ~3MB. + * + * Updates are merged (shallow, later fields win) rather than dropped because + * they carry different fields at different times — streamed `rawInput` + * snapshots, input-derived title/content, edit diffs, then terminal + * status/rawOutput — and the conversation reducer `Object.assign`s each one + * into the tool call. Merging here reproduces exactly the state a full replay + * would build; keeping only the last update would lose any field it doesn't + * re-send (e.g. `rawInput` for every streamed call). */ export function collapseSupersededToolCallUpdates( events: AcpMessage[], ): AcpMessage[] { + const firstIndexById = new Map(); const lastIndexById = new Map(); + const mergedById = new Map>(); for (let i = 0; i < events.length; i++) { - const id = toolCallUpdateId(events[i]); - if (id !== undefined) lastIndexById.set(id, i); + const update = toolCallUpdateOf(events[i]); + if (!update) continue; + lastIndexById.set(update.toolCallId, i); + const merged = mergedById.get(update.toolCallId); + if (merged) { + Object.assign(merged, update); + } else { + firstIndexById.set(update.toolCallId, i); + mergedById.set(update.toolCallId, { ...update }); + } } if (lastIndexById.size === 0) return events; const collapsed: AcpMessage[] = []; for (let i = 0; i < events.length; i++) { - const id = toolCallUpdateId(events[i]); - if (id !== undefined && lastIndexById.get(id) !== i) continue; + const update = toolCallUpdateOf(events[i]); + if (update) { + const id = update.toolCallId; + if (lastIndexById.get(id) !== i) continue; + // A call with a single update needs no synthetic merge. + if (firstIndexById.get(id) === i) { + collapsed.push(events[i]); + } else { + const merged = mergedById.get(id); + collapsed.push( + merged ? withToolCallUpdate(events[i], merged) : events[i], + ); + } + continue; + } collapsed.push(events[i]); } return collapsed; diff --git a/packages/core/src/sessions/sessionService.ts b/packages/core/src/sessions/sessionService.ts index 18b205be32..c83b0c5e27 100644 --- a/packages/core/src/sessions/sessionService.ts +++ b/packages/core/src/sessions/sessionService.ts @@ -168,8 +168,11 @@ export interface SessionTrpc { }; logs: { readLocalLogs: TrpcQuery; - /** Optional: collapses superseded tool_call_update snapshots server-side so - * a tool-heavy log doesn't ship its full redundant history over IPC. */ + /** Optional: merges superseded tool_call_update snapshots server-side so + * a tool-heavy log doesn't ship its full redundant history over IPC. + * Presence can't be trusted on proxy-based hosts (a tRPC client fabricates + * a query for any path), so callers fall back to `readLocalLogs` when the + * call itself fails. */ readLocalLogsCollapsed?: TrpcQuery; /** Optional: only the Electron host exposes the tail read. Core feature- * detects and falls back to a full read when it's absent. */ @@ -4848,6 +4851,40 @@ export class SessionService { } } + /** + * Read the local log, preferring the collapsed read (superseded + * tool_call_update snapshots merged server-side, so a tool-heavy log + * doesn't cross the transport at full size). `originalLineCount` is the + * pre-collapse line count when the collapsed read served the content. + * + * A tRPC proxy client fabricates a query object for any path, so a host + * whose router lacks the procedure only fails at call time — fall back to + * the plain read then, instead of misreporting the local log as unreadable. + */ + private async readLocalLogsPreferCollapsed( + taskRunId: string, + ): Promise<{ content: string | null; originalLineCount?: number }> { + const collapsedQuery = this.d.trpc.logs.readLocalLogsCollapsed; + if (collapsedQuery) { + try { + const res = (await collapsedQuery.query({ taskRunId })) as { + content: string; + totalLineCount: number; + } | null; + return { + content: res?.content ?? null, + originalLineCount: res?.totalLineCount, + }; + } catch { + this.d.log.warn("Collapsed local log read failed, using plain read", { + taskRunId, + }); + } + } + const content = await this.d.trpc.logs.readLocalLogs.query({ taskRunId }); + return { content }; + } + private async fetchSessionLogs( logUrl: string | undefined, taskRunId?: string, @@ -4863,38 +4900,21 @@ export class SessionService { if (taskRunId) { try { - const collapsedQuery = this.d.trpc.logs.readLocalLogsCollapsed; - if (collapsedQuery) { - // The collapsed content has fewer lines, so use the server's - // original totalLineCount for resume/gap tracking, not the parsed one. - const res = (await collapsedQuery.query({ taskRunId })) as { - content: string; - totalLineCount: number; - } | null; - if (res?.content?.trim()) { - localResult = { - ...this.parseLogContent(res.content), - totalLineCount: res.totalLineCount, - }; - if ( - !options.minEntryCount || - localResult.totalLineCount >= options.minEntryCount - ) { - return localResult; - } - } - } else { - const localContent = await this.d.trpc.logs.readLocalLogs.query({ - taskRunId, - }); - if (localContent?.trim()) { - localResult = this.parseLogContent(localContent); - if ( - !options.minEntryCount || - localResult.totalLineCount >= options.minEntryCount - ) { - return localResult; - } + const { content, originalLineCount } = + await this.readLocalLogsPreferCollapsed(taskRunId); + if (content?.trim()) { + const parsed = this.parseLogContent(content); + // Collapsed content has fewer lines than the file, so keep the + // server's original line count for resume/gap tracking. + localResult = + originalLineCount === undefined + ? parsed + : { ...parsed, totalLineCount: originalLineCount }; + if ( + !options.minEntryCount || + localResult.totalLineCount >= options.minEntryCount + ) { + return localResult; } } } catch { diff --git a/packages/host-router/src/routers/logs.router.ts b/packages/host-router/src/routers/logs.router.ts index 2c4e177be4..94a42bd378 100644 --- a/packages/host-router/src/routers/logs.router.ts +++ b/packages/host-router/src/routers/logs.router.ts @@ -4,6 +4,8 @@ import { LOGS_SERVICE } from "@posthog/workspace-server/services/local-logs/iden import { fetchS3LogsInput, fetchS3LogsOutput, + readLocalLogsCollapsedInput, + readLocalLogsCollapsedOutput, readLocalLogsInput, readLocalLogsOutput, readLocalLogsTailInput, @@ -28,6 +30,15 @@ export const logsRouter = router({ .readLocalLogs(input.taskRunId), ), + readLocalLogsCollapsed: publicProcedure + .input(readLocalLogsCollapsedInput) + .output(readLocalLogsCollapsedOutput) + .query(({ ctx, input }) => + ctx.container + .get(LOGS_SERVICE) + .readLocalLogsCollapsed(input.taskRunId), + ), + readLocalLogsTail: publicProcedure .input(readLocalLogsTailInput) .output(readLocalLogsTailOutput) diff --git a/packages/workspace-server/src/services/local-logs/service.ts b/packages/workspace-server/src/services/local-logs/service.ts index 4956ea0d52..b139d950a2 100644 --- a/packages/workspace-server/src/services/local-logs/service.ts +++ b/packages/workspace-server/src/services/local-logs/service.ts @@ -9,32 +9,103 @@ import type { ILogsService } from "./identifiers"; const DATA_DIR = ".posthog-code"; const TOOL_CALL_UPDATE_MARKER = '"sessionUpdate":"tool_call_update"'; -const TOOL_CALL_ID_RE = /"toolCallId":"([^"]+)"/; + +interface StoredEntryShape { + notification?: { + params?: Record; + } & Record; +} + +function toolCallUpdateOfLine( + line: string, +): { entry: StoredEntryShape; update: Record } | null { + // Substring pre-filter: an unescaped marker can only occur as JSON + // structure (a quote inside a string value is always escaped), so + // non-matching lines are skipped without parsing. + if (!line.includes(TOOL_CALL_UPDATE_MARKER)) return null; + let entry: StoredEntryShape; + try { + entry = JSON.parse(line) as StoredEntryShape; + } catch { + return null; + } + const update = entry?.notification?.params?.update as + | Record + | undefined; + if (update?.sessionUpdate !== "tool_call_update") return null; + if (typeof update.toolCallId !== "string") return null; + return { entry, update }; +} /** - * Drop superseded `tool_call_update` lines (keep the last per `toolCallId`) - * before the log crosses to the renderer. Agents re-send the full accumulated - * tool output on every update, so the transfer + parse would otherwise carry - * hundreds of MB of redundant snapshots. Whole lines are dropped so the result - * stays valid NDJSON; line-based (no JSON.parse) to stay cheap on a 300MB log. + * Collapse superseded `tool_call_update` lines into one merged line per + * `toolCallId` (at the last update's position) before the log crosses to the + * renderer. Agents re-send the full accumulated tool output on every update, + * so the transfer + parse would otherwise carry hundreds of MB of redundant + * snapshots. + * + * Updates are merged (shallow, later fields win) rather than dropped: they + * carry different fields at different times (streamed `rawInput` snapshots, + * input-derived title/content, edit diffs, terminal status/rawOutput) and the + * renderer reducer `Object.assign`s each one, so a merged update reproduces + * exactly what replaying every line would build. Only `tool_call_update` + * lines are parsed; other lines pass through untouched, so the result stays + * valid NDJSON. Parsing those lines here trades one pass of workspace-server + * CPU for not shipping and parsing the same bytes in the renderer. */ function collapseToolCallUpdateLines(ndjson: string): string { const lines = ndjson.split("\n"); + const idByIndex = new Array(lines.length); + const firstIndexById = new Map(); const lastIndexById = new Map(); + const lastEntryById = new Map(); + const mergedById = new Map>(); + for (let i = 0; i < lines.length; i++) { - if (!lines[i].includes(TOOL_CALL_UPDATE_MARKER)) continue; - const id = lines[i].match(TOOL_CALL_ID_RE)?.[1]; - if (id) lastIndexById.set(id, i); + const parsed = toolCallUpdateOfLine(lines[i]); + if (!parsed) continue; + const id = parsed.update.toolCallId as string; + idByIndex[i] = id; + lastIndexById.set(id, i); + lastEntryById.set(id, parsed.entry); + const merged = mergedById.get(id); + if (merged) { + Object.assign(merged, parsed.update); + } else { + firstIndexById.set(id, i); + mergedById.set(id, { ...parsed.update }); + } } if (lastIndexById.size === 0) return ndjson; const kept: string[] = []; for (let i = 0; i < lines.length; i++) { - if (lines[i].includes(TOOL_CALL_UPDATE_MARKER)) { - const id = lines[i].match(TOOL_CALL_ID_RE)?.[1]; - if (id && lastIndexById.get(id) !== i) continue; + const id = idByIndex[i]; + if (id === undefined) { + kept.push(lines[i]); + continue; + } + if (lastIndexById.get(id) !== i) continue; + if (firstIndexById.get(id) === i) { + // Single update for this call: the original line already is the merge. + kept.push(lines[i]); + continue; } - kept.push(lines[i]); + const entry = lastEntryById.get(id); + const merged = mergedById.get(id); + if (!entry || !merged) { + kept.push(lines[i]); + continue; + } + kept.push( + JSON.stringify({ + ...entry, + notification: { + ...entry.notification, + params: { ...entry.notification?.params, update: merged }, + }, + }), + ); } return kept.join("\n"); } diff --git a/packages/workspace-server/src/services/local-logs/serviceTail.test.ts b/packages/workspace-server/src/services/local-logs/serviceTail.test.ts index c2be22e70a..f1aa596cc8 100644 --- a/packages/workspace-server/src/services/local-logs/serviceTail.test.ts +++ b/packages/workspace-server/src/services/local-logs/serviceTail.test.ts @@ -107,7 +107,7 @@ describe("LocalLogsService.readLocalLogsCollapsed", () => { }, }); - it("drops superseded tool_call_update lines, keeps latest per toolCallId, preserves original line count", async () => { + it("collapses superseded tool_call_update lines to one per toolCallId, preserves original line count", async () => { const lines = [ `{"i":0}`, toolUpdate("a", "a1"), @@ -121,13 +121,62 @@ describe("LocalLogsService.readLocalLogsCollapsed", () => { expect(res?.totalLineCount).toBe(5); const kept = res?.content.trim().split("\n") ?? []; - // both non-tool lines + only the latest "a" update remain + // both non-tool lines + one merged "a" update remain expect(kept).toHaveLength(3); expect(kept[2]).toContain(`"out":"a3"`); expect(res?.content).not.toContain(`"out":"a1"`); expect(res?.content).not.toContain(`"out":"a2"`); }); + it("merges fields across updates instead of dropping them", async () => { + const withFields = (fields: Record) => + JSON.stringify({ + ts: 1, + notification: { + method: "session/update", + params: { + update: { + sessionUpdate: "tool_call_update", + toolCallId: "a", + ...fields, + }, + }, + }, + }); + const lines = [ + withFields({ rawInput: { command: "ls -la" }, title: "List files" }), + withFields({ status: "completed", rawOutput: "done" }), + ]; + fs.writeFileSync(logPath(), `${lines.join("\n")}\n`); + + const res = await new LocalLogsService().readLocalLogsCollapsed(RUN); + + const kept = res?.content.trim().split("\n") ?? []; + expect(kept).toHaveLength(1); + const update = JSON.parse(kept[0]).notification.params.update; + expect(update).toEqual({ + sessionUpdate: "tool_call_update", + toolCallId: "a", + rawInput: { command: "ls -la" }, + title: "List files", + status: "completed", + rawOutput: "done", + }); + }); + + it("passes through lines that fail to parse", async () => { + const broken = `{"truncated": "sessionUpdate":"tool_call_update" not-json`; + const lines = [toolUpdate("a", "a1"), broken, toolUpdate("a", "a2")]; + fs.writeFileSync(logPath(), `${lines.join("\n")}\n`); + + const res = await new LocalLogsService().readLocalLogsCollapsed(RUN); + + const kept = res?.content.trim().split("\n") ?? []; + expect(kept).toHaveLength(2); + expect(kept[0]).toBe(broken); + expect(kept[1]).toContain(`"out":"a2"`); + }); + it("returns null when the log doesn't exist", async () => { expect( await new LocalLogsService().readLocalLogsCollapsed("missing"),