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/agent/src/session-log-writer.test.ts b/packages/agent/src/session-log-writer.test.ts index eec8b3ed8a..182fab39b6 100644 --- a/packages/agent/src/session-log-writer.test.ts +++ b/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[]> => { + 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 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 } }; + }[]; + 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 be63d75e36..ef227e47d0 100644 --- a/packages/agent/src/session-log-writer.ts +++ b/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; + bufferedAt: number; +} + interface SessionState { context: SessionContext; chunkBuffer?: ChunkBuffer; lastAgentMessage?: string; currentTurnMessages: string[]; + toolUpdateCache: Map; } 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[] = []; 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): { + 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 (!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): boolean { return this.getSessionUpdateType(message) === "agent_message_chunk"; } diff --git a/packages/core/src/sessions/sessionEvents.test.ts b/packages/core/src/sessions/sessionEvents.test.ts index 332f630179..4706956a55 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,136 @@ describe("promptReferencesAbsoluteFolder", () => { expect(promptReferencesAbsoluteFolder("just text")).toBe(false); }); }); + +describe("collapseSupersededToolCallUpdates", () => { + const toolUpdateFields = ( + toolCallId: string, + fields: Record, + ): AcpMessage => + ({ + type: "acp_message", + ts: 1, + message: { + method: "session/update", + params: { + update: { + sessionUpdate: "tool_call_update", + toolCallId, + ...fields, + }, + }, + }, + }) as unknown as AcpMessage; + + const toolUpdate = (toolCallId: string, text: string): AcpMessage => + toolUpdateFields(toolCallId, { content: text }); + + 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("collapses to one update per toolCallId, at the last update's 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("collapses each distinct toolCallId independently", () => { + 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); + }); + + 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 aa6cd24148..7370397428 100644 --- a/packages/core/src/sessions/sessionEvents.ts +++ b/packages/core/src/sessions/sessionEvents.ts @@ -171,6 +171,94 @@ export function shellExecutesToContextBlocks( * Convert stored log entries to ACP messages. * Optionally prepends a user message with the task description. */ +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?: unknown } + | undefined; + if (update?.sessionUpdate !== "tool_call_update") return 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, + }); +} + +/** + * 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 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 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; +} + export function convertStoredEntriesToEvents( entries: StoredLogEntry[], taskDescription?: string, @@ -188,7 +276,7 @@ export function convertStoredEntriesToEvents( events.push(storedEntryToAcpMessage(entry)); } - return events; + return collapseSupersededToolCallUpdates(events); } /** diff --git a/packages/core/src/sessions/sessionService.ts b/packages/core/src/sessions/sessionService.ts index 5b5af8740e..c83b0c5e27 100644 --- a/packages/core/src/sessions/sessionService.ts +++ b/packages/core/src/sessions/sessionService.ts @@ -168,6 +168,12 @@ export interface SessionTrpc { }; logs: { readLocalLogs: TrpcQuery; + /** 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. */ readLocalLogsTail?: TrpcQuery; @@ -4845,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, @@ -4860,11 +4900,16 @@ export class SessionService { if (taskRunId) { try { - const localContent = await this.d.trpc.logs.readLocalLogs.query({ - taskRunId, - }); - if (localContent?.trim()) { - localResult = this.parseLogContent(localContent); + 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 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/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..b139d950a2 100644 --- a/packages/workspace-server/src/services/local-logs/service.ts +++ b/packages/workspace-server/src/services/local-logs/service.ts @@ -8,6 +8,108 @@ import type { ILogsService } from "./identifiers"; const DATA_DIR = ".posthog-code"; +const TOOL_CALL_UPDATE_MARKER = '"sessionUpdate":"tool_call_update"'; + +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 }; +} + +/** + * 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++) { + 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++) { + 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; + } + 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"); +} + interface WriteState { pending: string | undefined; lastWritten: string | undefined; @@ -52,6 +154,16 @@ export class LocalLogsService implements ILogsService { } } + 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..f1aa596cc8 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,109 @@ 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("collapses superseded tool_call_update lines to one 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 + 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"), + ).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)