From 2cb505bd80407ebc71f7ada34aa4b5081b02703b Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Wed, 19 Aug 2026 14:03:43 +0300 Subject: [PATCH 1/8] perf(desktop): reduce task runtime CPU and process leaks --- .../pi-runtime/piSessionController.test.ts | 104 +++++++++ .../src/pi-runtime/piSessionController.ts | 115 ++++++++- .../src/sessions/sessionEventBatching.test.ts | 75 +++++- .../sessions/sessionEventDiagnostics.test.ts | 4 +- .../core/src/sessions/sessionService.ts | 108 +++++++-- .../core/src/sidebar/buildSidebarData.ts | 9 +- .../src/task-detail/cloudToolChanges.test.ts | 82 +++++++ .../core/src/task-detail/cloudToolChanges.ts | 149 +++++++++--- .../canvas/freeform/FreeformCanvasView.tsx | 19 +- .../freeform/useCanvasGenerationToasts.ts | 24 +- .../canvas/hooks/useChannelTaskData.ts | 19 +- .../hooks/useCommandCenterData.ts | 15 +- .../editor/components/useSmoothedText.test.ts | 15 ++ .../editor/components/useSmoothedText.ts | 8 + .../features/git-interaction/cloudPrUrl.ts | 9 +- .../components/CloudGitInteractionHeader.tsx | 17 +- .../features/git-interaction/useCloudPrUrl.ts | 14 +- .../sessions/components/ConversationView.tsx | 4 +- .../components/EmbeddedSessionView.tsx | 2 +- .../components/GeneratingIndicator.test.ts | 23 +- .../components/GeneratingIndicator.tsx | 8 +- .../components/buildConversationItems.ts | 85 ++++++- .../components/chat-thread/ChatThread.tsx | 12 +- .../chat-thread/ChatThreadGrouping.test.ts | 120 ++++++++++ .../components/chat-thread/chatRowGrouping.ts | 57 +++++ .../incrementalConversationItems.test.ts | 87 +++++-- .../incrementalConversationItems.ts | 53 +---- .../session-update/StatusNotificationView.tsx | 6 +- .../session-update/SubagentToolView.tsx | 1 + .../sessions/hooks/useChatTitleGenerator.ts | 29 ++- .../hooks/useSessionCallbacks.test.tsx | 6 +- .../sessions/hooks/useSessionCallbacks.ts | 21 +- .../components/TaskHeaderActions.tsx | 7 +- .../task-detail/components/TaskLogsPanel.tsx | 2 +- .../task-detail/hooks/useCloudEventSummary.ts | 23 +- .../task-detail/hooks/useCloudRunState.ts | 42 +++- .../packages/ui/src/styles/globals.css | 8 +- .../process-tracking/process-tracking.test.ts | 17 +- .../process-tracking/process-tracking.ts | 23 +- .../process-tracking/process-utils.test.ts | 124 ++++++++++ .../process-tracking/process-utils.ts | 219 +++++++++++++++--- .../src/services/watcher/service.test.ts | 31 ++- .../src/services/watcher/service.ts | 20 +- 43 files changed, 1516 insertions(+), 300 deletions(-) create mode 100644 products/desktop/packages/ui/src/features/sessions/components/chat-thread/ChatThreadGrouping.test.ts create mode 100644 products/desktop/packages/ui/src/features/sessions/components/chat-thread/chatRowGrouping.ts create mode 100644 products/desktop/packages/workspace-server/src/services/process-tracking/process-utils.test.ts diff --git a/products/desktop/packages/core/src/pi-runtime/piSessionController.test.ts b/products/desktop/packages/core/src/pi-runtime/piSessionController.test.ts index 6215d46cdaf8..db95840458b7 100644 --- a/products/desktop/packages/core/src/pi-runtime/piSessionController.test.ts +++ b/products/desktop/packages/core/src/pi-runtime/piSessionController.test.ts @@ -1678,6 +1678,46 @@ describe("PiSessionController", () => { ]); }); + it("does not append a chunk twice when it arrives during initial load", async () => { + vi.useFakeTimers(); + let resolveConversation: (events: AgentConversationEvent[]) => void = + () => {}; + const conversation = new Promise((resolve) => { + resolveConversation = resolve; + }); + const chunk: AgentConversationEvent = { + type: "assistant_message_chunk", + timestamp: 1, + content: { type: "text", text: "hello" }, + }; + let onEvent: (event: AgentConversationEvent) => void = () => {}; + const session = createSession(); + vi.mocked(session.getConversation).mockReturnValue(conversation); + vi.mocked(session.client.getState).mockResolvedValue({ + ...(await session.client.getState()), + isStreaming: true, + }); + vi.mocked(session.onConversationEvent).mockImplementation((handler) => { + onEvent = handler; + return () => {}; + }); + const controller = createController(session); + + const connection = controller.connect("task-1"); + await vi.waitFor(() => + expect(session.onConversationEvent).toHaveBeenCalledOnce(), + ); + onEvent(chunk); + resolveConversation([]); + await connection; + await vi.advanceTimersByTimeAsync(50); + + expect(controller.store.getState().sessions["task-1"].events).toEqual([ + chunk, + ]); + vi.useRealTimers(); + }); + it("loads session state and appends normalized runtime events", async () => { const initialEvent: AgentConversationEvent = { type: "assistant_message_chunk", @@ -1706,4 +1746,68 @@ describe("PiSessionController", () => { status: { isCompacting: true }, }); }); + + it("batches streamed chunks into one store update", async () => { + vi.useFakeTimers(); + const first: AgentConversationEvent = { + type: "assistant_message_chunk", + timestamp: 1, + content: { type: "text", text: "hello" }, + }; + const second: AgentConversationEvent = { + type: "assistant_message_chunk", + timestamp: 2, + content: { type: "text", text: " world" }, + }; + let onEvent: (event: AgentConversationEvent) => void = () => {}; + const session = createSession(); + vi.mocked(session.onConversationEvent).mockImplementation((handler) => { + onEvent = handler; + return () => {}; + }); + const controller = createController(session); + await controller.connect("task-1"); + const listener = vi.fn(); + controller.store.subscribe(listener); + + onEvent(first); + onEvent(second); + + expect(listener).not.toHaveBeenCalled(); + await vi.advanceTimersByTimeAsync(50); + expect(listener).toHaveBeenCalledOnce(); + expect(controller.store.getState().sessions["task-1"].events).toEqual([ + first, + second, + ]); + vi.useRealTimers(); + }); + + it("flushes streamed chunks before a turn completes", async () => { + const chunk: AgentConversationEvent = { + type: "assistant_message_chunk", + timestamp: 1, + content: { type: "text", text: "done" }, + }; + const completed: AgentConversationEvent = { + type: "turn_completed", + timestamp: 2, + }; + let onEvent: (event: AgentConversationEvent) => void = () => {}; + const session = createSession(); + vi.mocked(session.onConversationEvent).mockImplementation((handler) => { + onEvent = handler; + return () => {}; + }); + const controller = createController(session); + await controller.connect("task-1"); + + onEvent(chunk); + onEvent(completed); + + expect(controller.store.getState().sessions["task-1"].events).toEqual([ + chunk, + completed, + ]); + }); }); diff --git a/products/desktop/packages/core/src/pi-runtime/piSessionController.ts b/products/desktop/packages/core/src/pi-runtime/piSessionController.ts index 4ed75a71da24..29771179c3e3 100644 --- a/products/desktop/packages/core/src/pi-runtime/piSessionController.ts +++ b/products/desktop/packages/core/src/pi-runtime/piSessionController.ts @@ -121,6 +121,8 @@ type PiTurnState = | { phase: "active"; startedAt?: number; stopReason?: string } | { phase: "completed" }; +const STREAM_UPDATE_INTERVAL_MS = 50; + type PiOperation = | "prompt" | "compact" @@ -168,6 +170,11 @@ export class PiSessionController { { trusted: boolean; promise: Promise } >(); private readonly liveEvents = new Map(); + private readonly pendingEvents = new Map(); + private readonly pendingEventTimers = new Map< + string, + ReturnType + >(); private readonly connections = new Map>(); private readonly readiness = new Map>(); private readonly sessionVersions = new Map(); @@ -793,11 +800,15 @@ export class PiSessionController { const conversationEvents = events.filter( (event) => event.type !== "queue_update", ); - const liveEvents = this.liveEvents.get(taskId) ?? []; + const liveEvents = [ + ...(this.liveEvents.get(taskId) ?? []), + ...(this.pendingEvents.get(taskId) ?? []), + ]; const newLiveEvents = this.reconcileLiveEvents( conversationEvents, liveEvents, ); + this.discardPendingEvents(taskId); this.liveEvents.set(taskId, newLiveEvents); const historyUserMessageIds = new Set( conversationEvents.flatMap((event) => @@ -908,6 +919,18 @@ export class PiSessionController { event: AgentConversationEvent, context?: PiConversationEventContext, ): void { + const isLive = context?.isLive ?? true; + if ( + isLive && + (event.type === "assistant_message_chunk" || + event.type === "assistant_thought_chunk" || + event.type === "tool_call_updated") + ) { + this.queueEvent(taskId, event); + return; + } + + this.flushPendingEvents(taskId); if (event.type === "queue_update") { const queue = { steering: event.steering, @@ -925,7 +948,6 @@ export class PiSessionController { return; } - const isLive = context?.isLive ?? true; this.applyTurnEvent(taskId, event, isLive); if (event.type === "runtime_error") { @@ -952,16 +974,9 @@ export class PiSessionController { ); } } - const isDirectBashEvent = - (event.type === "tool_call_started" || - event.type === "tool_call_updated") && - event.toolCall.origin === "user_shell"; const hasTurnActivity = - !isDirectBashEvent && - (event.type === "assistant_message_chunk" || - event.type === "assistant_thought_chunk" || - event.type === "tool_call_started" || - event.type === "tool_call_updated"); + event.type === "tool_call_started" && + event.toolCall.origin !== "user_shell"; if (status && hasTurnActivity) { status = { ...status, isStreaming: true }; } @@ -1106,6 +1121,82 @@ export class PiSessionController { this.updateSession(taskId, { cloudStatus }); } + private queueEvent(taskId: string, event: AgentConversationEvent): void { + const pending = this.pendingEvents.get(taskId) ?? []; + pending.push(event); + this.pendingEvents.set(taskId, pending); + + if (this.pendingEventTimers.has(taskId)) return; + this.pendingEventTimers.set( + taskId, + setTimeout(() => { + this.pendingEventTimers.delete(taskId); + this.flushPendingEvents(taskId); + }, STREAM_UPDATE_INTERVAL_MS), + ); + } + + private flushPendingEvents(taskId: string): void { + const timer = this.pendingEventTimers.get(taskId); + if (timer) { + clearTimeout(timer); + this.pendingEventTimers.delete(taskId); + } + + const pending = this.pendingEvents.get(taskId); + if (!pending?.length) return; + this.pendingEvents.delete(taskId); + + const session = this.getSession(taskId); + const seenSourceIds = new Set( + session.events.flatMap((event) => + event.sourceId ? [event.sourceId] : [], + ), + ); + const events = pending.filter((event) => { + if (!event.sourceId || !seenSourceIds.has(event.sourceId)) { + if (event.sourceId) seenSourceIds.add(event.sourceId); + return true; + } + return false; + }); + if (events.length === 0) return; + + for (const event of events) { + this.applyTurnEvent(taskId, event, true); + } + + this.liveEvents.set(taskId, [ + ...(this.liveEvents.get(taskId) ?? []), + ...events, + ]); + const hasTurnActivity = events.some( + (event) => + event.type !== "tool_call_updated" || + event.toolCall.origin !== "user_shell", + ); + const latestSession = this.getSession(taskId); + this.updateSession(taskId, { + connectionState: "connected", + events: [...latestSession.events, ...events], + status: + latestSession.status && hasTurnActivity + ? { ...latestSession.status, isStreaming: true } + : latestSession.status, + error: + latestSession.error?.scope === "operation" + ? latestSession.error + : undefined, + }); + } + + private discardPendingEvents(taskId: string): void { + const timer = this.pendingEventTimers.get(taskId); + if (timer) clearTimeout(timer); + this.pendingEventTimers.delete(taskId); + this.pendingEvents.delete(taskId); + } + private async refreshStats(taskId: string): Promise { const sessionVersion = this.getSessionVersion(taskId); try { @@ -1536,6 +1627,7 @@ export class PiSessionController { private disposeTask(taskId: string): void { this.cancelAuthRestoration.get(taskId)?.(); + this.flushPendingEvents(taskId); this.resetTransport(taskId); this.taskRunIds.delete(taskId); this.liveEvents.delete(taskId); @@ -1547,6 +1639,7 @@ export class PiSessionController { } private resetTransport(taskId: string): void { + this.discardPendingEvents(taskId); this.advanceSessionVersion(taskId); this.disposeConversationSubscription(taskId); this.sessions.delete(taskId); diff --git a/products/desktop/packages/core/src/sessions/sessionEventBatching.test.ts b/products/desktop/packages/core/src/sessions/sessionEventBatching.test.ts index c0799d89741e..2008bc46f3aa 100644 --- a/products/desktop/packages/core/src/sessions/sessionEventBatching.test.ts +++ b/products/desktop/packages/core/src/sessions/sessionEventBatching.test.ts @@ -4,7 +4,7 @@ import { SessionService, type SessionServiceDeps } from "./sessionService"; const TASK_ID = "task-1"; const RUN_ID = "run-1"; -const FLUSH_MS = 16; +const FLUSH_MS = 50; /** A plain streamed agent-message chunk — the common per-token event that just * gets appended to the transcript. */ @@ -32,6 +32,51 @@ function chunkText(event: AcpMessage): string { return params.update.content.text; } +function toolCall(id: string): AcpMessage { + return { + ts: 2, + message: { + jsonrpc: "2.0", + method: "session/update", + params: { + sessionId: RUN_ID, + update: { + sessionUpdate: "tool_call", + toolCallId: id, + title: "Read file", + status: "in_progress", + }, + }, + }, + } as unknown as AcpMessage; +} + +function configOptionUpdate(): AcpMessage { + return { + ts: 3, + message: { + jsonrpc: "2.0", + method: "session/update", + params: { + sessionId: RUN_ID, + update: { + sessionUpdate: "config_option_update", + configOptions: [ + { + id: "mode", + name: "Mode", + category: "mode", + type: "select", + currentValue: "default", + options: [], + }, + ], + }, + }, + }, + } as unknown as AcpMessage; +} + function createHarness() { const sessions: Record = { [RUN_ID]: { @@ -138,6 +183,7 @@ function createHarness() { service, appendEvents, notifyPromptComplete, + setPersistedConfigOptions: deps.setPersistedConfigOptions, updateSession: store.updateSession, emit: (event: AcpMessage) => onEvent?.(event), events: () => sessions[RUN_ID].events, @@ -189,6 +235,7 @@ describe("streamed event batching", () => { // A single flush tick drains the whole burst, in arrival order. vi.advanceTimersByTime(FLUSH_MS); expect(h.events().map(chunkText)).toEqual(["a", "b", "c"]); + expect(h.appendEvents).toHaveBeenCalledOnce(); }); it("flushes buffered events synchronously on teardown", () => { @@ -207,6 +254,32 @@ describe("streamed event batching", () => { expect(h.events()).toHaveLength(2); }); + it("batches interleaved text and tool updates in order", () => { + const h = createHarness(); + const streamed = chunk("a"); + const active = toolCall("tool-1"); + + h.emit(streamed); + h.emit(active); + + expect(h.events()).toEqual([]); + vi.advanceTimersByTime(FLUSH_MS); + expect(h.events()).toEqual([streamed, active]); + expect(h.appendEvents).toHaveBeenCalledOnce(); + }); + + it("applies and persists config option updates immediately", () => { + const h = createHarness(); + const streamed = chunk("a"); + const configUpdate = configOptionUpdate(); + + h.emit(streamed); + h.emit(configUpdate); + + expect(h.events()).toEqual([streamed, configUpdate]); + expect(h.setPersistedConfigOptions).toHaveBeenCalledOnce(); + }); + it("keeps the turn duration when the prompt mutation clears state before the response flushes", () => { const h = createHarness(); diff --git a/products/desktop/packages/core/src/sessions/sessionEventDiagnostics.test.ts b/products/desktop/packages/core/src/sessions/sessionEventDiagnostics.test.ts index 8842c9a0bb17..aadb1fa1354b 100644 --- a/products/desktop/packages/core/src/sessions/sessionEventDiagnostics.test.ts +++ b/products/desktop/packages/core/src/sessions/sessionEventDiagnostics.test.ts @@ -98,9 +98,9 @@ describe("session event diagnostics", () => { events.subscriptions[0].handlers.onData(chunk("first")); events.subscriptions[0].handlers.onData(chunk("second")); - vi.advanceTimersByTime(20); + vi.advanceTimersByTime(50); - expect(appendEvents).toHaveBeenCalledTimes(2); + expect(appendEvents).toHaveBeenCalledTimes(3); expect(log.error).toHaveBeenCalledWith( "Session event handling failed", expect.objectContaining({ diff --git a/products/desktop/packages/core/src/sessions/sessionService.ts b/products/desktop/packages/core/src/sessions/sessionService.ts index c6d9ff3f8cf8..b7e6d477bfa7 100644 --- a/products/desktop/packages/core/src/sessions/sessionService.ts +++ b/products/desktop/packages/core/src/sessions/sessionService.ts @@ -175,10 +175,10 @@ function describeAcpMethod(acpMsg: AcpMessage): string { * Streamed events are buffered and flushed on this cadence so a burst of tokens * coalesces into one processing pass (and roughly one render) instead of one * per event. Electron IPC delivers each event as its own task, so a microtask - * flush wouldn't batch across them — a short timer does. One frame is - * imperceptible for streamed text. + * flush wouldn't batch across them — a short timer does. A 50ms presentation + * delay is imperceptible for streamed text and cuts update churn substantially. */ -const SESSION_EVENT_FLUSH_MS = 16; +const SESSION_EVENT_FLUSH_MS = 50; /** * Steering an adapter that can't fold a message into a running turn leaves only * one way in: interrupt it. Cancelling the instant the user hits send cuts off @@ -1670,6 +1670,14 @@ function classifyTurnEventKind( return "other"; } +function isStreamUpdateEvent(acpMsg: AcpMessage): boolean { + const msg = acpMsg.message; + if (!("method" in msg) || msg.method !== "session/update") return false; + const update = (msg as { params?: { update?: { sessionUpdate?: string } } }) + .params?.update; + return update?.sessionUpdate !== "config_option_update"; +} + export class SessionService { private connectingTasks = new Map>(); private reconcilingTasks = new Set(); @@ -2753,9 +2761,7 @@ export class SessionService { const batches = this.pendingSessionEvents; this.pendingSessionEvents = new Map(); for (const [taskRunId, events] of batches) { - for (const acpMsg of events) { - this.applySessionEvent(taskRunId, acpMsg); - } + this.applySessionEventBatch(taskRunId, events); } } @@ -2765,9 +2771,7 @@ export class SessionService { const events = this.pendingSessionEvents.get(taskRunId); if (!events) return; this.pendingSessionEvents.delete(taskRunId); - for (const acpMsg of events) { - this.applySessionEvent(taskRunId, acpMsg); - } + this.applySessionEventBatch(taskRunId, events); } /** @@ -2779,18 +2783,76 @@ export class SessionService { try { this.handleSessionEvent(taskRunId, acpMsg); } catch (error) { - const stats = this.statsFor(taskRunId); - stats.failed += 1; - if (stats.failed === 1 || stats.failed % 100 === 0) { - this.d.log.error("Session event handling failed", { - taskRunId, - method: describeAcpMethod(acpMsg), - failed: stats.failed, - received: stats.received, - error, - }); + this.recordSessionEventFailure(taskRunId, acpMsg, error); + } + } + + private recordSessionEventFailure( + taskRunId: string, + acpMsg: AcpMessage, + error: unknown, + ): void { + const stats = this.statsFor(taskRunId); + stats.failed += 1; + if (stats.failed === 1 || stats.failed % 100 === 0) { + this.d.log.error("Session event handling failed", { + taskRunId, + method: describeAcpMethod(acpMsg), + failed: stats.failed, + received: stats.received, + error, + }); + } + } + + private applySessionEventBatch( + taskRunId: string, + events: AcpMessage[], + ): void { + let passive: AcpMessage[] = []; + const flushPassive = () => { + if (passive.length === 0) return; + const session = this.d.store.getSessions()[taskRunId]; + if (session) { + try { + if (session.initialPrompt?.length) { + this.d.store.updateSession(taskRunId, { + initialPrompt: undefined, + }); + } + this.d.store.appendEvents(taskRunId, passive); + } catch (error) { + this.recordSessionEventFailure(taskRunId, passive[0], error); + for (const event of passive) { + this.applySessionEvent(taskRunId, event); + } + passive = []; + return; + } + try { + this.updatePromptStateFromEvents(taskRunId, passive, { + isLive: true, + }); + } catch (error) { + this.recordSessionEventFailure( + taskRunId, + passive.at(-1) ?? passive[0], + error, + ); + } + } + passive = []; + }; + + for (const event of events) { + if (isStreamUpdateEvent(event)) { + passive.push(event); + } else { + flushPassive(); + this.applySessionEvent(taskRunId, event); } } + flushPassive(); } private statsFor(taskRunId: string): { @@ -2954,7 +3016,13 @@ export class SessionService { { taskRunId }, { onData: (payload: unknown) => { - this.enqueueSessionEvent(taskRunId, payload as AcpMessage); + const event = payload as AcpMessage; + if (isStreamUpdateEvent(event)) { + this.enqueueSessionEvent(taskRunId, event); + } else { + this.flushSessionEventsForTask(taskRunId); + this.handleSessionEvent(taskRunId, event); + } }, onError: (err) => { this.d.log.error("Session subscription error", { diff --git a/products/desktop/packages/core/src/sidebar/buildSidebarData.ts b/products/desktop/packages/core/src/sidebar/buildSidebarData.ts index c9f76c9589e9..430a1487da8f 100644 --- a/products/desktop/packages/core/src/sidebar/buildSidebarData.ts +++ b/products/desktop/packages/core/src/sidebar/buildSidebarData.ts @@ -1,4 +1,8 @@ -import { readPrUrls, type WorkspaceMode } from "@posthog/shared"; +import { + readPrUrls, + type SessionStatus, + type WorkspaceMode, +} from "@posthog/shared"; import type { Task, TaskRunStatus } from "@posthog/shared/domain-types"; import { taskActivityAt } from "../tasks/taskActivity"; import { getRepositoryInfo } from "./groupTasks"; @@ -111,6 +115,7 @@ export function filterVisibleTasks( } export interface TaskSession { + status?: SessionStatus; isPromptPending?: boolean; pendingPermissions?: { size: number }; cloudStatus?: TaskRunStatus; @@ -133,7 +138,7 @@ export function computeSidebarSessionSignature( typeof session.cloudOutput?.pr_url === "string" ? session.cloudOutput.pr_url : ""; - signature += `${session.taskId}:${session.isPromptPending ? 1 : 0}:${ + signature += `${session.taskId}:${session.status ?? ""}:${session.isPromptPending ? 1 : 0}:${ session.pendingPermissions?.size ?? 0 }:${session.cloudStatus ?? ""}:${prUrl};`; } diff --git a/products/desktop/packages/core/src/task-detail/cloudToolChanges.test.ts b/products/desktop/packages/core/src/task-detail/cloudToolChanges.test.ts index a098cdb19127..b5f46e548382 100644 --- a/products/desktop/packages/core/src/task-detail/cloudToolChanges.test.ts +++ b/products/desktop/packages/core/src/task-detail/cloudToolChanges.test.ts @@ -1,12 +1,94 @@ +import type { AcpMessage } from "@posthog/shared"; import { describe, expect, it } from "vitest"; import { cachedDiffStats, + createCloudEventSummaryTracker, extractCloudFileContent, extractCloudToolChangedFiles, type ParsedToolCall, } from "./cloudToolChanges"; +function toolEvent( + toolCallId: string, + update: Record, +): AcpMessage { + return { + ts: 1, + message: { + jsonrpc: "2.0", + method: "session/update", + params: { + sessionId: "run-1", + update: { sessionUpdate: "tool_call_update", toolCallId, ...update }, + }, + }, + } as AcpMessage; +} + +function textEvent(text: string): AcpMessage { + return { + ts: 1, + message: { + jsonrpc: "2.0", + method: "session/update", + params: { + sessionId: "run-1", + update: { + sessionUpdate: "agent_message_chunk", + content: { type: "text", text }, + }, + }, + }, + } as AcpMessage; +} + +describe("createCloudEventSummaryTracker", () => { + it("merges appended tool updates and resets when the transcript is replaced", () => { + const tracker = createCloudEventSummaryTracker(); + const started = toolEvent("tool-1", { title: "Edit file" }); + const completed = toolEvent("tool-1", { status: "completed" }); + + tracker.update([started]); + const appended = tracker.update([started, completed]); + expect(appended.toolCalls.get("tool-1")).toMatchObject({ + title: "Edit file", + status: "completed", + }); + + const replacement = tracker.update([ + toolEvent("tool-2", { title: "Write file" }), + ]); + expect([...replacement.toolCalls.keys()]).toEqual(["tool-2"]); + }); + + it("reuses the projected summary when appended events do not change tools", () => { + const tracker = createCloudEventSummaryTracker(); + const started = toolEvent("tool-1", { title: "Edit file" }); + const first = tracker.update([started]); + + expect(tracker.update([started, textEvent("hello")])).toBe(first); + }); + + it("retains the changed-files revision for irrelevant streamed content", () => { + const tracker = createCloudEventSummaryTracker(); + const started = toolEvent("tool-1", { + title: "Edit file", + locations: [{ path: "src/file.ts", line: null }], + }); + const first = tracker.update([started]); + const input = toolEvent("tool-1", { + content: [ + { type: "content", content: { type: "text", text: "partial" } }, + ], + }); + + const second = tracker.update([started, input]); + + expect(second.changedFilesRevision).toBe(first.changedFilesRevision); + }); +}); + function diffObj( newText: string, oldText: string, diff --git a/products/desktop/packages/core/src/task-detail/cloudToolChanges.ts b/products/desktop/packages/core/src/task-detail/cloudToolChanges.ts index 479d0035b48a..eb74512fe1f6 100644 --- a/products/desktop/packages/core/src/task-detail/cloudToolChanges.ts +++ b/products/desktop/packages/core/src/task-detail/cloudToolChanges.ts @@ -8,6 +8,7 @@ import { isJsonRpcNotification, } from "@posthog/shared"; import type { ChangedFile } from "@posthog/shared/domain-types"; +import { createAppendOnlyTracker } from "../sessions/appendOnlyTracker"; function getContentText( content: ToolCallContent[] | undefined, @@ -178,6 +179,73 @@ export function cachedDiffStats( export interface CloudEventSummary { toolCalls: Map; + revision: number; + changedFilesRevision: number; +} + +function changedFilesKey(toolCall: ParsedToolCall): string { + const diff = getDiffContent(toolCall.content); + return JSON.stringify({ + kind: inferKind(toolCall.kind, toolCall.title), + failed: toolCall.status === "failed", + locations: toolCall.locations?.map((location) => location.path), + diff: diff + ? { + path: diff.path, + oldText: diff.oldText, + newText: diff.newText, + } + : undefined, + }); +} + +function applyCloudEvent( + toolCalls: Map, + event: AcpMessage, +): { toolChanged: boolean; changedFilesChanged: boolean } { + const message = event.message; + if (!isJsonRpcNotification(message) || message.method !== "session/update") { + return { toolChanged: false, changedFilesChanged: false }; + } + const params = message.params as + | { update?: Record } + | undefined; + const update = params?.update; + if (!update || typeof update !== "object") { + return { toolChanged: false, changedFilesChanged: false }; + } + + const sessionUpdate = update.sessionUpdate; + if (sessionUpdate !== "tool_call" && sessionUpdate !== "tool_call_update") { + return { toolChanged: false, changedFilesChanged: false }; + } + + const toolCallId = + typeof update.toolCallId === "string" ? update.toolCallId : undefined; + if (!toolCallId) return { toolChanged: false, changedFilesChanged: false }; + + const patch: Partial = { + toolCallId, + kind: typeof update.kind === "string" ? update.kind : null, + title: typeof update.title === "string" ? update.title : undefined, + status: typeof update.status === "string" ? update.status : null, + locations: Array.isArray(update.locations) + ? (update.locations as ToolCallLocation[]) + : undefined, + content: Array.isArray(update.content) + ? (update.content as ToolCallContent[]) + : undefined, + rawOutput: update.rawOutput, + }; + + const existing = toolCalls.get(toolCallId); + const merged = mergeToolCall(existing, patch); + toolCalls.set(toolCallId, merged); + return { + toolChanged: true, + changedFilesChanged: + !existing || changedFilesKey(existing) !== changedFilesKey(merged), + }; } /** @@ -189,48 +257,53 @@ export function buildCloudEventSummary( const toolCalls = new Map(); for (const event of events) { - const message = event.message; - if (!isJsonRpcNotification(message)) continue; - - if (message.method === "session/update") { - const params = message.params as - | { update?: Record } - | undefined; - const update = params?.update; - if (!update || typeof update !== "object") continue; - - const sessionUpdate = update.sessionUpdate; - if ( - sessionUpdate !== "tool_call" && - sessionUpdate !== "tool_call_update" - ) { - continue; - } + applyCloudEvent(toolCalls, event); + } - const toolCallId = - typeof update.toolCallId === "string" ? update.toolCallId : undefined; - if (!toolCallId) continue; - - const patch: Partial = { - toolCallId, - kind: typeof update.kind === "string" ? update.kind : null, - title: typeof update.title === "string" ? update.title : undefined, - status: typeof update.status === "string" ? update.status : null, - locations: Array.isArray(update.locations) - ? (update.locations as ToolCallLocation[]) - : undefined, - content: Array.isArray(update.content) - ? (update.content as ToolCallContent[]) - : undefined, - rawOutput: update.rawOutput, - }; + return { toolCalls, revision: 0, changedFilesRevision: 0 }; +} - const merged = mergeToolCall(toolCalls.get(toolCallId), patch); - toolCalls.set(toolCallId, merged); - } +export function createCloudEventSummaryTracker(): { + update(events: AcpMessage[]): CloudEventSummary; +} { + interface TrackerState { + toolCalls: Map; + revision: number; + changedFilesRevision: number; } - return { toolCalls }; + let projectedState: TrackerState | undefined; + let projectedRevision = -1; + let projectedResult: CloudEventSummary = { + toolCalls: new Map(), + revision: 0, + changedFilesRevision: 0, + }; + + return createAppendOnlyTracker({ + init: () => ({ + toolCalls: new Map(), + revision: 0, + changedFilesRevision: 0, + }), + processEvent: (state, event) => { + const change = applyCloudEvent(state.toolCalls, event); + if (change.toolChanged) state.revision++; + if (change.changedFilesChanged) state.changedFilesRevision++; + }, + getResult: (state) => { + if (state !== projectedState || state.revision !== projectedRevision) { + projectedState = state; + projectedRevision = state.revision; + projectedResult = { + toolCalls: state.toolCalls, + revision: state.revision, + changedFilesRevision: state.changedFilesRevision, + }; + } + return projectedResult; + }, + }); } export function extractCloudFileDiff( diff --git a/products/desktop/packages/ui/src/features/canvas/freeform/FreeformCanvasView.tsx b/products/desktop/packages/ui/src/features/canvas/freeform/FreeformCanvasView.tsx index 199c9c0166a5..2bbf6e29fa78 100644 --- a/products/desktop/packages/ui/src/features/canvas/freeform/FreeformCanvasView.tsx +++ b/products/desktop/packages/ui/src/features/canvas/freeform/FreeformCanvasView.tsx @@ -74,7 +74,7 @@ import { readCommentContext, } from "@posthog/ui/features/sessions/components/commentViewTypes"; import { useCommentsQuery } from "@posthog/ui/features/sessions/components/useComments"; -import { useSessionForTask } from "@posthog/ui/features/sessions/useSession"; +import { useSessionSelector } from "@posthog/ui/features/sessions/useSession"; import { taskDetailQuery } from "@posthog/ui/features/tasks/queries"; import { ResizableSidebar } from "@posthog/ui/primitives/ResizableSidebar"; import { toast } from "@posthog/ui/primitives/toast"; @@ -91,6 +91,7 @@ import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query"; import { Link } from "@tanstack/react-router"; import { AnimatePresence, motion } from "framer-motion"; import { useCallback, useEffect, useMemo, useRef, useState } from "react"; +import { shallow } from "zustand/shallow"; import { BuiltCanvas } from "./BuiltCanvas"; import { CanvasAgentRequestDialog } from "./CanvasAgentRequestDialog"; import { CanvasBuildStatus } from "./CanvasBuildStatus"; @@ -223,7 +224,21 @@ export function FreeformCanvasView({ enabled: !!effectiveTaskId, refetchInterval: effectiveTaskId ? 5000 : false, }); - const genSession = useSessionForTask(effectiveTaskId ?? undefined); + const genSession = useSessionSelector( + effectiveTaskId ?? undefined, + (session) => + session + ? { + status: session.status, + cloudStatus: session.cloudStatus, + isPromptPending: session.isPromptPending, + taskRunId: session.taskRunId, + agentIdleForRunId: session.agentIdleForRunId, + isCloud: session.isCloud, + } + : undefined, + shallow, + ); // Whether the run's session is still alive. Drives record + build polling so // a freshly published version and its queued build get picked up. A local ACP // session stays "connected" after its generation prompt finishes, so this diff --git a/products/desktop/packages/ui/src/features/canvas/freeform/useCanvasGenerationToasts.ts b/products/desktop/packages/ui/src/features/canvas/freeform/useCanvasGenerationToasts.ts index 266e32f0a920..0825bc1442d9 100644 --- a/products/desktop/packages/ui/src/features/canvas/freeform/useCanvasGenerationToasts.ts +++ b/products/desktop/packages/ui/src/features/canvas/freeform/useCanvasGenerationToasts.ts @@ -127,15 +127,23 @@ export function useCanvasGenerationToasts(): void { })), }); - // The live ACP sessions — for local runs this, not the run record, is what - // tells us generation has actually finished. - const sessions = useSessionStore((s) => s.sessions); - const taskIdIndex = useSessionStore((s) => s.taskIdIndex); + // Subscribe only to fields that affect generation state. Transcript appends + // must not rerender this persistent root-mounted watcher. + const sessionSignature = useSessionStore((state) => + taskIds + .map((id) => { + const runId = state.taskIdIndex[id]; + const session = runId ? state.sessions[runId] : undefined; + return `${id}:${session?.status ?? ""}:${session?.cloudStatus ?? ""}:${session?.isPromptPending ? 1 : 0}`; + }) + .join("|"), + ); + const sessionState = useSessionStore.getState(); // Compute the "still generating?" signal per tracked task each render. const states = taskIds.map((id, i) => { - const runId = taskIdIndex[id]; - const session = runId ? sessions[runId] : undefined; + const runId = sessionState.taskIdIndex[id]; + const session = runId ? sessionState.sessions[runId] : undefined; const latestRun = details[i]?.data?.latest_run; const generating = isCanvasGenerating({ genTaskId: id, @@ -147,12 +155,12 @@ export function useCanvasGenerationToasts(): void { }); // A stable signature so the transition effect only runs on real changes. - const sig = states + const sig = `${sessionSignature}|${states .map( (s) => `${s.id}:${s.generating ? 1 : 0}:${s.latestRun?.status ?? ""}:${s.session?.status ?? ""}:${s.session?.cloudStatus ?? ""}:${s.session?.isPromptPending ? 1 : 0}`, ) - .join("|"); + .join("|")}`; const statesRef = useRef(states); statesRef.current = states; diff --git a/products/desktop/packages/ui/src/features/canvas/hooks/useChannelTaskData.ts b/products/desktop/packages/ui/src/features/canvas/hooks/useChannelTaskData.ts index e1e685f534f1..16ec94f53133 100644 --- a/products/desktop/packages/ui/src/features/canvas/hooks/useChannelTaskData.ts +++ b/products/desktop/packages/ui/src/features/canvas/hooks/useChannelTaskData.ts @@ -5,12 +5,13 @@ import { } from "@posthog/core/sidebar/buildSidebarData"; import type { TaskData } from "@posthog/core/sidebar/sidebarData.types"; import type { Task } from "@posthog/shared/domain-types"; -import { useSessionForTask } from "@posthog/ui/features/sessions/useSession"; +import { useSessionSelector } from "@posthog/ui/features/sessions/useSession"; import { usePinnedTasks } from "@posthog/ui/features/sidebar/usePinnedTasks"; import { useTaskViewed } from "@posthog/ui/features/sidebar/useTaskViewed"; import { useSuspendedTaskIds } from "@posthog/ui/features/suspension/useSuspendedTaskIds"; import { useWorkspace } from "@posthog/ui/features/workspace/useWorkspace"; import { useMemo } from "react"; +import { shallow } from "zustand/shallow"; const EMPTY_SET: ReadonlySet = new Set(); const EMPTY_MAP: ReadonlyMap = new Map(); @@ -21,7 +22,19 @@ const EMPTY_MAP: ReadonlyMap = new Map(); export function useChannelTaskData( task: Task | undefined, ): TaskData | undefined { - const session = useSessionForTask(task?.id); + const session = useSessionSelector( + task?.id, + (current): TaskSession | undefined => + current + ? { + isPromptPending: current.isPromptPending, + pendingPermissions: current.pendingPermissions, + cloudStatus: current.cloudStatus, + cloudOutput: current.cloudOutput, + } + : undefined, + shallow, + ); const workspace = useWorkspace(task?.id); const { pinnedTaskIds } = usePinnedTasks(); const suspendedTaskIds = useSuspendedTaskIds(); @@ -31,7 +44,7 @@ export function useChannelTaskData( if (!task) return undefined; const sidebarTask = narrowFullTask(task); return deriveTaskData(sidebarTask, { - session: session as TaskSession | undefined, + session, workspace: workspace ?? undefined, timestamp: timestamps[task.id], pinnedIds: pinnedTaskIds, diff --git a/products/desktop/packages/ui/src/features/command-center/hooks/useCommandCenterData.ts b/products/desktop/packages/ui/src/features/command-center/hooks/useCommandCenterData.ts index 372053b391ae..cbbddaa055f5 100644 --- a/products/desktop/packages/ui/src/features/command-center/hooks/useCommandCenterData.ts +++ b/products/desktop/packages/ui/src/features/command-center/hooks/useCommandCenterData.ts @@ -9,8 +9,7 @@ import { } from "@posthog/core/command-center/status"; import type { Task } from "@posthog/shared/domain-types"; import { useMemo } from "react"; -import type { AgentSession } from "../../sessions/sessionStore"; -import { useSessions } from "../../sessions/useSession"; +import { useSidebarSessionMap } from "../../sidebar/useSidebarSessionMap"; import { useTasks } from "../../tasks/useTasks"; import { useWorkspaces } from "../../workspace/useWorkspace"; import { useCommandCenterStore } from "../commandCenterStore"; @@ -24,7 +23,7 @@ export function useCommandCenterData(): { } { const storeCells = useCommandCenterStore((s) => s.cells); const { data: tasks = [] } = useTasks(); - const sessions = useSessions(); + const sessionByTaskId = useSidebarSessionMap(); const { data: workspaces } = useWorkspaces(); const taskById = useMemo(() => { @@ -35,16 +34,6 @@ export function useCommandCenterData(): { return map; }, [tasks]); - const sessionByTaskId = useMemo(() => { - const map = new Map(); - for (const session of Object.values(sessions)) { - if (session.taskId) { - map.set(session.taskId, session); - } - } - return map; - }, [sessions]); - const cells = useMemo( () => buildCommandCenterCells(storeCells, { diff --git a/products/desktop/packages/ui/src/features/editor/components/useSmoothedText.test.ts b/products/desktop/packages/ui/src/features/editor/components/useSmoothedText.test.ts index 3d2571d31dc6..720075457106 100644 --- a/products/desktop/packages/ui/src/features/editor/components/useSmoothedText.test.ts +++ b/products/desktop/packages/ui/src/features/editor/components/useSmoothedText.test.ts @@ -115,6 +115,21 @@ describe("useSmoothedText", () => { expect(result.current.length).toBe(11); }); + it("caps streaming renders at 20 frames per second", () => { + const { result, rerender } = renderHook( + ({ t }) => useSmoothedText(t, 100), + { initialProps: { t: "" } }, + ); + rerender({ t: "x".repeat(50) }); + + flushFrame(0); + expect(result.current.length).toBe(1); + flushFrame(16); + expect(result.current.length).toBe(1); + flushFrame(34); + expect(result.current.length).toBe(6); + }); + it("cancels the pending frame on unmount", () => { const { rerender, unmount } = renderHook( ({ t }) => useSmoothedText(t, 100), diff --git a/products/desktop/packages/ui/src/features/editor/components/useSmoothedText.ts b/products/desktop/packages/ui/src/features/editor/components/useSmoothedText.ts index a2083ee4d16b..af01db72ba18 100644 --- a/products/desktop/packages/ui/src/features/editor/components/useSmoothedText.ts +++ b/products/desktop/packages/ui/src/features/editor/components/useSmoothedText.ts @@ -4,6 +4,7 @@ import { useEffect, useRef, useState } from "react"; // one, so the cadence reads as even typing instead of speeding up to clear a // backlog. Matches the feel of #2685. See https://upstash.com/blog/smooth-streaming. const DEFAULT_CHARS_PER_SECOND = 120; +const FRAME_INTERVAL_MS = 50; // Past this backlog we stop easing and snap, so a large buffered chunk (e.g. a // reconnect replaying a long message) never crawls. const MAX_LAG_CHARS = 600; @@ -75,6 +76,13 @@ export function useSmoothedText( lastTsRef.current = null; const tick = (ts: number) => { const last = lastTsRef.current ?? ts; + if ( + lastTsRef.current !== null && + ts - lastTsRef.current < FRAME_INTERVAL_MS + ) { + rafRef.current = requestAnimationFrame(tick); + return; + } lastTsRef.current = ts; shownLenRef.current = nextRevealLength( shownLenRef.current, diff --git a/products/desktop/packages/ui/src/features/git-interaction/cloudPrUrl.ts b/products/desktop/packages/ui/src/features/git-interaction/cloudPrUrl.ts index a6a942e6d109..6e3fb0d41811 100644 --- a/products/desktop/packages/ui/src/features/git-interaction/cloudPrUrl.ts +++ b/products/desktop/packages/ui/src/features/git-interaction/cloudPrUrl.ts @@ -1,10 +1,11 @@ import { mergePrUrls, readPrSummaries, readPrUrls } from "@posthog/shared"; import type { Task } from "@posthog/shared/domain-types"; -import type { AgentSession } from "@posthog/ui/features/sessions/sessionStore"; + +type CloudPrSession = { cloudOutput?: Record | null }; export function resolveCloudPrUrls( task: Task | undefined, - session: AgentSession | undefined, + session: CloudPrSession | undefined, ): string[] { return mergePrUrls( readPrUrls(task?.latest_run?.output), @@ -14,7 +15,7 @@ export function resolveCloudPrUrls( export function resolveCloudPrSummaries( task: Task | undefined, - session: AgentSession | undefined, + session: CloudPrSession | undefined, ): Record { return { ...readPrSummaries(session?.cloudOutput), @@ -24,7 +25,7 @@ export function resolveCloudPrSummaries( export function resolveCloudPrUrl( task: Task | undefined, - session: AgentSession | undefined, + session: CloudPrSession | undefined, ): string | null { return resolveCloudPrUrls(task, session)[0] ?? null; } diff --git a/products/desktop/packages/ui/src/features/git-interaction/components/CloudGitInteractionHeader.tsx b/products/desktop/packages/ui/src/features/git-interaction/components/CloudGitInteractionHeader.tsx index 564db72b7957..0ad44e0fa6de 100644 --- a/products/desktop/packages/ui/src/features/git-interaction/components/CloudGitInteractionHeader.tsx +++ b/products/desktop/packages/ui/src/features/git-interaction/components/CloudGitInteractionHeader.tsx @@ -14,7 +14,10 @@ import { useFeatureFlag } from "../../feature-flags/useFeatureFlag"; import { DirtyTreeDialog } from "../../sessions/components/DirtyTreeDialog"; import { HandoffConfirmDialog } from "../../sessions/components/HandoffConfirmDialog"; import { useHandoffDialogStore } from "../../sessions/handoffDialogStore"; -import { useSessionForTask } from "../../sessions/useSession"; +import { + useSessionHandoffInProgress, + useSessionSelector, +} from "../../sessions/useSession"; import { GIT_CACHE_KEY_PROVIDER, type GitCacheKeyProvider, @@ -35,7 +38,11 @@ export function CloudGitInteractionHeader({ taskId, task, }: CloudGitInteractionHeaderProps) { - const session = useSessionForTask(taskId); + const inProgress = useSessionHandoffInProgress(taskId); + const cloudBranch = useSessionSelector( + taskId, + (session) => session?.cloudBranch ?? null, + ); const queryClient = useQueryClient(); const cacheKeyProvider = useService( GIT_CACHE_KEY_PROVIDER, @@ -112,17 +119,13 @@ export function CloudGitInteractionHeader({ if (!cloudHandoffEnabled || !localWorkspaces) return null; if (task.origin_product === "image_builder") return null; - const inProgress = session?.handoffInProgress ?? false; - return ( <>
- localHandoff.openConfirm(taskId, session?.cloudBranch ?? null) - } + onClick={() => localHandoff.openConfirm(taskId, cloudBranch)} > {inProgress ? ( diff --git a/products/desktop/packages/ui/src/features/git-interaction/useCloudPrUrl.ts b/products/desktop/packages/ui/src/features/git-interaction/useCloudPrUrl.ts index eb6de92938ab..1b4080430f48 100644 --- a/products/desktop/packages/ui/src/features/git-interaction/useCloudPrUrl.ts +++ b/products/desktop/packages/ui/src/features/git-interaction/useCloudPrUrl.ts @@ -1,4 +1,4 @@ -import { useSessionForTask } from "../sessions/useSession"; +import { useSessionSelector } from "../sessions/useSession"; import { useTasks } from "../tasks/useTasks"; import { resolveCloudPrSummaries, @@ -16,13 +16,21 @@ export function useCloudPrUrl(taskId: string): string | null { export function useCloudPrUrls(taskId: string): string[] { const { data: tasks = [] } = useTasks(); const task = tasks.find((t) => t.id === taskId); - const session = useSessionForTask(taskId); + const cloudOutput = useSessionSelector( + taskId, + (session) => session?.cloudOutput, + ); + const session = cloudOutput ? { cloudOutput } : undefined; return resolveCloudPrUrls(task, session); } export function useCloudPrSummaries(taskId: string): Record { const { data: tasks = [] } = useTasks(); const task = tasks.find((t) => t.id === taskId); - const session = useSessionForTask(taskId); + const cloudOutput = useSessionSelector( + taskId, + (session) => session?.cloudOutput, + ); + const session = cloudOutput ? { cloudOutput } : undefined; return resolveCloudPrSummaries(task, session); } diff --git a/products/desktop/packages/ui/src/features/sessions/components/ConversationView.tsx b/products/desktop/packages/ui/src/features/sessions/components/ConversationView.tsx index d2f74705ddb9..6b1ded4cf8f1 100644 --- a/products/desktop/packages/ui/src/features/sessions/components/ConversationView.tsx +++ b/products/desktop/packages/ui/src/features/sessions/components/ConversationView.tsx @@ -142,8 +142,8 @@ export function ConversationView({ // Streaming appends one event per token. The parse is incremental — each // event is handled once and completed turns are reused by reference — so per // token the work tracks the active turn, not the whole thread. We feed - // `events` directly (no frame-throttle) so a sent message's optimistic->real - // swap is never delayed past the frame the store commits it. + // `events` directly; the controller batches streaming-only updates while + // terminal and status events still flush immediately. const { items: conversationItems, lastTurnInfo, diff --git a/products/desktop/packages/ui/src/features/sessions/components/EmbeddedSessionView.tsx b/products/desktop/packages/ui/src/features/sessions/components/EmbeddedSessionView.tsx index 5652addb14ed..ff8da870ff45 100644 --- a/products/desktop/packages/ui/src/features/sessions/components/EmbeddedSessionView.tsx +++ b/products/desktop/packages/ui/src/features/sessions/components/EmbeddedSessionView.tsx @@ -47,7 +47,7 @@ export function EmbeddedSessionView({ handleRetry, handleNewSession, handleBashCommand, - } = useSessionCallbacks({ taskId, task, session, repoPath }); + } = useSessionCallbacks({ taskId, task, repoPath }); useEffect(() => { requestFocus(taskId); diff --git a/products/desktop/packages/ui/src/features/sessions/components/GeneratingIndicator.test.ts b/products/desktop/packages/ui/src/features/sessions/components/GeneratingIndicator.test.ts index 82602871be97..4e6f517a40ee 100644 --- a/products/desktop/packages/ui/src/features/sessions/components/GeneratingIndicator.test.ts +++ b/products/desktop/packages/ui/src/features/sessions/components/GeneratingIndicator.test.ts @@ -1,7 +1,14 @@ -import { formatDuration } from "@posthog/ui/features/sessions/components/GeneratingIndicator"; -import { describe, expect, it } from "vitest"; +import { + formatDuration, + GeneratingIndicator, +} from "@posthog/ui/features/sessions/components/GeneratingIndicator"; +import { render, screen } from "@testing-library/react"; +import { createElement } from "react"; +import { afterEach, describe, expect, it, vi } from "vitest"; + +describe("GeneratingIndicator", () => { + afterEach(() => vi.useRealTimers()); -describe("formatDuration", () => { it("formats sub-minute durations with configurable precision", () => { expect(formatDuration(12_340)).toBe("12.34s"); expect(formatDuration(12_340, 1)).toBe("12.3s"); @@ -11,4 +18,14 @@ describe("formatDuration", () => { expect(formatDuration(62_340)).toBe("1m 02s"); expect(formatDuration(62_340, 1)).toBe("1m 02s"); }); + it("shows elapsed time immediately when remounted", () => { + vi.useFakeTimers(); + vi.setSystemTime(new Date("2026-07-28T12:00:30Z")); + + render( + createElement(GeneratingIndicator, { startedAt: Date.now() - 30_000 }), + ); + + expect(screen.getByText("30s")).toBeInTheDocument(); + }); }); diff --git a/products/desktop/packages/ui/src/features/sessions/components/GeneratingIndicator.tsx b/products/desktop/packages/ui/src/features/sessions/components/GeneratingIndicator.tsx index ba7b32287411..b57907fbfd24 100644 --- a/products/desktop/packages/ui/src/features/sessions/components/GeneratingIndicator.tsx +++ b/products/desktop/packages/ui/src/features/sessions/components/GeneratingIndicator.tsx @@ -72,14 +72,16 @@ export function GeneratingIndicator({ useEffect(() => { const startTime = startedAt ?? Date.now(); - const interval = setInterval(() => { + const tick = () => { const now = Date.now(); setElapsed(Math.max(0, now - startTime - pausedRef.current)); // Measured from the last event rather than the turn's start, so a turn // that streamed for a minute and then went silent still reads as quiet. const since = lastActivityRef.current; setQuietFor(since === null ? 0 : Math.max(0, now - since)); - }, 100); + }; + tick(); + const interval = setInterval(tick, 1000); return () => clearInterval(interval); }, [startedAt]); @@ -119,7 +121,7 @@ export function GeneratingIndicator({ (Esc to stop {dot} - {formatDuration(elapsed, 1)} + {formatDuration(elapsed, 0)} {quietFor >= QUIET_AFTER_MS && ( <> diff --git a/products/desktop/packages/ui/src/features/sessions/components/buildConversationItems.ts b/products/desktop/packages/ui/src/features/sessions/components/buildConversationItems.ts index 940a23de43a9..5f48d15c704e 100644 --- a/products/desktop/packages/ui/src/features/sessions/components/buildConversationItems.ts +++ b/products/desktop/packages/ui/src/features/sessions/components/buildConversationItems.ts @@ -627,6 +627,7 @@ function completePromptTurn( const wasCancelled = turn.stopReason === "cancelled"; turn.context.turnCancelled = wasCancelled; + replaceTurnContextRows(b, turn.context); if (turn.gitAction.isGitAction && turn.gitAction.actionType) { b.items.push({ @@ -650,6 +651,25 @@ function completePromptTurn( } } +function replaceTurnContextRows(b: ItemBuilder, context: TurnContext): void { + const visited = new Set(); + const replaceRows = (items: ConversationItem[]): void => { + if (visited.has(items)) return; + visited.add(items); + for (let index = 0; index < items.length; index++) { + const item = items[index]; + if (item.type !== "session_update") continue; + if (item.turnContext === context) { + items[index] = { ...item }; + } + for (const children of item.turnContext.childItems.values()) { + replaceRows(children); + } + } + }; + replaceRows(b.items); +} + function handleNotification( b: ItemBuilder, msg: { method: string; params?: unknown }, @@ -1041,6 +1061,7 @@ function pushChildItem(b: ItemBuilder, parentId: string, update: RenderItem) { update, turnContext: turn.context, }); + reissueToolCallRow(b, parentId); } function appendTextChunkToChildren( @@ -1079,6 +1100,7 @@ function appendTextChunkToChildren( }, }, }; + reissueToolCallRow(b, parentId); } else { turn.itemCount++; children.push({ @@ -1087,9 +1109,46 @@ function appendTextChunkToChildren( update: { ...update, content: { ...update.content } }, turnContext: turn.context, }); + reissueToolCallRow(b, parentId); } } +function reissueToolCallRow( + b: ItemBuilder, + toolCallId: string, + nextUpdate?: ToolCall, +): void { + const turn = b.currentTurn; + if (!turn) return; + const current = turn.toolCalls.get(toolCallId); + const update = nextUpdate ?? (current ? { ...current } : undefined); + if (!update) return; + turn.toolCalls.set(toolCallId, update); + + const visited = new Set(); + const replace = (items: ConversationItem[]): void => { + if (visited.has(items)) return; + visited.add(items); + for (let index = 0; index < items.length; index++) { + const item = items[index]; + if (item.type !== "session_update") continue; + if ( + item.update.sessionUpdate === "tool_call" && + item.update.toolCallId === toolCallId + ) { + items[index] = { + ...item, + update: { ...update, sessionUpdate: "tool_call" }, + }; + } + for (const children of item.turnContext.childItems.values()) { + replace(children); + } + } + }; + replace(b.items); +} + function processSessionUpdate( b: ItemBuilder, update: ConversationSessionUpdate, @@ -1119,8 +1178,9 @@ function processSessionUpdate( const existing = turn.toolCalls.get(update.toolCallId); if (existing) { const wasTerminal = isTerminalToolStatus(existing.status); - Object.assign(existing, update); - if (!wasTerminal && isTerminalToolStatus(existing.status)) { + const merged = { ...existing, ...update }; + reissueToolCallRow(b, update.toolCallId, merged); + if (!wasTerminal && isTerminalToolStatus(merged.status)) { b.completedToolCallCount++; } } else { @@ -1146,8 +1206,25 @@ function processSessionUpdate( if (existing) { const wasTerminal = isTerminalToolStatus(existing.status); const { sessionUpdate: _, ...rest } = update; - Object.assign(existing, rest); - if (!wasTerminal && isTerminalToolStatus(existing.status)) { + const merged: ToolCall = { + ...existing, + ...rest, + toolCallId: existing.toolCallId, + title: rest.title ?? existing.title, + content: + rest.content === null + ? undefined + : (rest.content ?? existing.content), + kind: rest.kind === null ? undefined : (rest.kind ?? existing.kind), + locations: + rest.locations === null + ? undefined + : (rest.locations ?? existing.locations), + status: + rest.status === null ? undefined : (rest.status ?? existing.status), + }; + reissueToolCallRow(b, update.toolCallId, merged); + if (!wasTerminal && isTerminalToolStatus(merged.status)) { b.completedToolCallCount++; } } diff --git a/products/desktop/packages/ui/src/features/sessions/components/chat-thread/ChatThread.tsx b/products/desktop/packages/ui/src/features/sessions/components/chat-thread/ChatThread.tsx index 7f310b5449c5..085dffd0b9bd 100644 --- a/products/desktop/packages/ui/src/features/sessions/components/chat-thread/ChatThread.tsx +++ b/products/desktop/packages/ui/src/features/sessions/components/chat-thread/ChatThread.tsx @@ -59,6 +59,7 @@ import { ChatStreamingMarkdown, } from "@posthog/ui/features/sessions/components/chat-thread/ChatMarkdown"; import { ChatThreadFooter } from "@posthog/ui/features/sessions/components/chat-thread/ChatThreadFooter"; +import { createIncrementalChatRowGrouper } from "@posthog/ui/features/sessions/components/chat-thread/chatRowGrouping"; import { ChatThreadChromeProvider } from "@posthog/ui/features/sessions/components/chat-thread/chatThreadChrome"; import type { PromptRecallHandler } from "@posthog/ui/features/sessions/components/chat-thread/composerPromptRecall"; import { MessageJumpPicker } from "@posthog/ui/features/sessions/components/chat-thread/MessageJumpPicker"; @@ -1371,9 +1372,16 @@ function ChatThreadRenderer({ [conversationItems, optimisticItems, isCloud], ); + const rowGrouper = useMemo( + () => + createIncrementalChatRowGrouper((nextItems) => + groupIntoTurns(groupToolCalls ? groupToolRuns(nextItems) : nextItems), + ), + [groupToolCalls], + ); const rows = useMemo( - () => groupIntoTurns(groupToolCalls ? groupToolRuns(items) : items), - [items, groupToolCalls], + () => rowGrouper.update(items), + [items, rowGrouper], ); // Virtualization ratchet: past the threshold the thread switches to the windowed body and diff --git a/products/desktop/packages/ui/src/features/sessions/components/chat-thread/ChatThreadGrouping.test.ts b/products/desktop/packages/ui/src/features/sessions/components/chat-thread/ChatThreadGrouping.test.ts new file mode 100644 index 000000000000..c62b9c5355f5 --- /dev/null +++ b/products/desktop/packages/ui/src/features/sessions/components/chat-thread/ChatThreadGrouping.test.ts @@ -0,0 +1,120 @@ +import type { ConversationItem } from "@posthog/ui/features/sessions/components/buildConversationItems"; +import { createIncrementalChatRowGrouper } from "@posthog/ui/features/sessions/components/chat-thread/chatRowGrouping"; +import type { TurnRow } from "@posthog/ui/features/sessions/components/chat-thread/threadVirtualization"; +import { isUserInitiatedConversationItem } from "@posthog/ui/features/sessions/components/isUserInitiatedConversationItem"; +import { describe, expect, it } from "vitest"; + +function userMessage(id: string): ConversationItem { + return { type: "user_message", id, content: id, timestamp: 1 }; +} + +function agentMessage(id: string): ConversationItem { + return { + type: "session_update", + id, + turnContext: { + toolCalls: new Map(), + childItems: new Map(), + turnCancelled: false, + turnComplete: false, + }, + update: { + sessionUpdate: "agent_message_chunk", + content: { type: "text", text: id }, + }, + }; +} + +function groupRows(items: ConversationItem[]): TurnRow[] { + const rows: TurnRow[] = []; + let agentItems: ConversationItem[] = []; + const flush = () => { + if (agentItems.length > 0) { + rows.push({ + type: "agent_turn", + id: agentItems[0].id, + items: agentItems, + }); + agentItems = []; + } + }; + for (const item of items) { + if (isUserInitiatedConversationItem(item)) { + flush(); + rows.push(item); + } else { + agentItems.push(item); + } + } + flush(); + return rows; +} + +describe("createIncrementalChatRowGrouper", () => { + it("reuses completed turns while rebuilding the active turn", () => { + const grouper = createIncrementalChatRowGrouper(groupRows); + const firstItems = [userMessage("u1"), agentMessage("a1")]; + const first = grouper.update(firstItems); + const secondItems = [...firstItems, userMessage("u2"), agentMessage("a2")]; + const second = grouper.update(secondItems); + const third = grouper.update([...secondItems, agentMessage("a3")]); + + expect(second[0]).toBe(first[0]); + expect(second[1]).toBe(first[1]); + expect(third[0]).toBe(second[0]); + expect(third[1]).toBe(second[1]); + expect(third.at(-1)).toMatchObject({ + type: "agent_turn", + items: [{ id: "a2" }, { id: "a3" }], + }); + }); + + it("fully rebuilds after a non-append replacement", () => { + const grouper = createIncrementalChatRowGrouper(groupRows); + grouper.update([userMessage("u1"), agentMessage("a1")]); + + expect( + grouper.update([userMessage("x1"), agentMessage("x2")]), + ).toMatchObject([ + { id: "x1" }, + { type: "agent_turn", items: [{ id: "x2" }] }, + ]); + }); + + it("rebuilds when a row inside the retained prefix is replaced in place", () => { + // The conversation builder swaps row objects at arbitrary indices (a status + // completing, a shell result arriving) — including inside turns the grouper + // already cached. The cached rows must not survive such a replacement. + const grouper = createIncrementalChatRowGrouper(groupRows); + const u1 = userMessage("u1"); + const status = agentMessage("s1"); + const u2 = userMessage("u2"); + grouper.update([u1, status, u2]); + + const replaced = { ...status }; + const rows = grouper.update([u1, replaced, u2, agentMessage("a2")]); + + const turn = rows[1]; + if (turn.type !== "agent_turn") throw new Error("expected an agent turn"); + expect(turn.items[0]).toBe(replaced); + }); + + it("replaces an optimistic boundary whose confirmed item has a new id", () => { + const grouper = createIncrementalChatRowGrouper(groupRows); + const prefix = [userMessage("u1"), agentMessage("a1")]; + grouper.update([...prefix, userMessage("optimistic-u2")]); + + expect( + grouper.update([ + ...prefix, + userMessage("confirmed-u2"), + agentMessage("a2"), + ]), + ).toMatchObject([ + { id: "u1" }, + { type: "agent_turn", items: [{ id: "a1" }] }, + { id: "confirmed-u2" }, + { type: "agent_turn", items: [{ id: "a2" }] }, + ]); + }); +}); diff --git a/products/desktop/packages/ui/src/features/sessions/components/chat-thread/chatRowGrouping.ts b/products/desktop/packages/ui/src/features/sessions/components/chat-thread/chatRowGrouping.ts new file mode 100644 index 000000000000..3be915f40d66 --- /dev/null +++ b/products/desktop/packages/ui/src/features/sessions/components/chat-thread/chatRowGrouping.ts @@ -0,0 +1,57 @@ +import type { ConversationItem } from "@posthog/ui/features/sessions/components/buildConversationItems"; +import type { TurnRow } from "@posthog/ui/features/sessions/components/chat-thread/threadVirtualization"; +import { isUserInitiatedConversationItem } from "@posthog/ui/features/sessions/components/isUserInitiatedConversationItem"; + +type GroupRows = (items: ConversationItem[]) => TurnRow[]; + +export function createIncrementalChatRowGrouper(groupRows: GroupRows) { + let cachedItems: ConversationItem[] = []; + let cachedRows: TurnRow[] = []; + + return { + update(items: ConversationItem[]): TurnRow[] { + if (items === cachedItems) return cachedRows; + + let rebuildStart = 0; + for (let index = items.length - 1; index >= 0; index--) { + if (isUserInitiatedConversationItem(items[index])) { + rebuildStart = index; + break; + } + } + + for (let index = 0; index < rebuildStart; index++) { + if (cachedItems[index] !== items[index]) { + rebuildStart = 0; + break; + } + } + + let boundaryId = items[rebuildStart]?.id; + let cachedBoundaryIndex = boundaryId + ? cachedRows.findIndex((row) => row.id === boundaryId) + : -1; + if ( + rebuildStart > 0 && + rebuildStart < cachedItems.length && + cachedBoundaryIndex < 0 + ) { + rebuildStart = 0; + boundaryId = items[0]?.id; + cachedBoundaryIndex = boundaryId + ? cachedRows.findIndex((row) => row.id === boundaryId) + : -1; + } + const prefixRowCount = + rebuildStart === 0 + ? 0 + : cachedBoundaryIndex >= 0 + ? cachedBoundaryIndex + : cachedRows.length; + const suffixRows = groupRows(items.slice(rebuildStart)); + cachedItems = items; + cachedRows = [...cachedRows.slice(0, prefixRowCount), ...suffixRows]; + return cachedRows; + }, + }; +} diff --git a/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.test.ts b/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.test.ts index a93814b7ac8a..a8dd087adfd4 100644 --- a/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.test.ts +++ b/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.test.ts @@ -502,31 +502,32 @@ describe("createIncrementalConversationBuilder", () => { } }); - it("gives the active turn a fresh toolCalls Map identity each call so an in-place tool update re-renders", () => { - // SessionUpdateView is memoized on `turnContext.toolCalls`, and tool_call_update - // mutates the tool entry in place. If the Map reference is reused across calls - // the memo bails and the completed status (and streamed output) stay hidden - // until the turn ends. Guard that the Map identity changes every event. + it("replaces the tool row's update object when a tool_call_update lands so its memoized view re-renders", () => { + // Memoized rows compare the row's `update` object; a tool_call_update must + // surface as a fresh object carrying the merged state, or the completed + // status (and streamed output) stay hidden until the turn ends. const inc = createIncrementalConversationBuilder(); const base = [userPromptMsg(1, 1, "go"), toolCallMsg(2, "t1")]; const r1 = inc.update(base, true); const next = [...base, toolUpdateMsg(3, "t1", { status: "completed" })]; const r2 = inc.update(next, true); - const ctx1 = r1.items.find((i) => i.type === "session_update"); - const ctx2 = r2.items.find((i) => i.type === "session_update"); - if (ctx1?.type !== "session_update" || ctx2?.type !== "session_update") { + const row1 = r1.items.find((i) => i.type === "session_update"); + const row2 = r2.items.find((i) => i.type === "session_update"); + if (row1?.type !== "session_update" || row2?.type !== "session_update") { throw new Error("expected tool-call session_update rows"); } - expect(ctx2.turnContext.toolCalls).not.toBe(ctx1.turnContext.toolCalls); - expect(ctx2.turnContext.childItems).not.toBe(ctx1.turnContext.childItems); + expect(row2.update).not.toBe(row1.update); + expect((row2.update as { status?: string }).status).toBe("completed"); + // The shared toolCalls Map holds the merged entry the view resolves. + expect(row2.turnContext.toolCalls.get("t1")?.status).toBe("completed"); }); it("surfaces a running agent's child tool calls live, before the turn completes", () => { // A subagent appends child tool calls (parentToolCallId) while it runs. The - // parent row is memoized on `turnContext.childItems`; without a fresh Map ref - // the new children stay invisible until turn end. Guard that a child appended - // mid-turn changes the childItems Map identity and is present. + // parent row renders them, so the parent's update object must be re-issued + // when a child arrives — otherwise its memoized view bails and the new + // children stay invisible until turn end. const inc = createIncrementalConversationBuilder(); const base = [ userPromptMsg(1, 1, "go"), @@ -543,11 +544,67 @@ describe("createIncrementalConversationBuilder", () => { if (row1?.type !== "session_update" || row2?.type !== "session_update") { throw new Error("expected agent session_update rows"); } - // New child arrived mid-turn: fresh Map identity so the memoized parent re-renders. - expect(row2.turnContext.childItems).not.toBe(row1.turnContext.childItems); + // New child arrived mid-turn: fresh parent update so the memoized row re-renders. + expect(row2).not.toBe(row1); + expect(row2.update).not.toBe(row1.update); expect(row2.turnContext.childItems.get("agent1")?.length).toBe(1); }); + it("re-issues a thought row's identity when its turn completes in the same batch a new turn starts", () => { + // The settled thought then sits before the next turn's boundary, where + // row caches retain by identity — only a fresh object surfaces the flip. + const inc = createIncrementalConversationBuilder(); + const base = [userPromptMsg(1, 1, "go"), thoughtChunk(2, "hmm")]; + const r1 = inc.update(base, true); + const next = [ + ...base, + promptResponseMsg(3, 1), + userPromptMsg(4, 2, "next"), + ]; + const r2 = inc.update(next, true); + + const isThought = (i: ConversationItem) => + i.type === "session_update" && + i.update.sessionUpdate === "agent_thought_chunk"; + const thought1 = r1.items.find(isThought); + const thought2 = r2.items.find(isThought); + if ( + thought1?.type !== "session_update" || + thought2?.type !== "session_update" + ) { + throw new Error("expected thought rows"); + } + expect(thought1.thoughtComplete).toBe(false); + expect(thought2.thoughtComplete).toBe(true); + expect(thought2).not.toBe(thought1); + }); + + it("settles an older turn's thought after a newer turn starts", () => { + const inc = createIncrementalConversationBuilder(); + const beforeCompletion = [ + userPromptMsg(1, 1, "first"), + thoughtChunk(2, "hmm"), + userPromptMsg(3, 2, "second"), + agentChunk(4, "working"), + ]; + inc.update(beforeCompletion, true); + + const result = inc.update( + [...beforeCompletion, promptResponseMsg(5, 1, "cancelled")], + true, + ); + const thought = result.items.find( + (item) => + item.type === "session_update" && + item.update.sessionUpdate === "agent_thought_chunk", + ); + + if (thought?.type !== "session_update") { + throw new Error("expected thought row"); + } + expect(thought.thoughtComplete).toBe(true); + }); + it("groups canonical PostHog child metadata under its subagent", () => { const inc = createIncrementalConversationBuilder(); const messages = [ diff --git a/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.ts b/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.ts index ec0f951287fc..f189f27219c4 100644 --- a/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.ts +++ b/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.ts @@ -3,7 +3,6 @@ import { type BuildConversationOptions, type BuildResult, buildConversationItems, - type ConversationItem, createItemBuilder, finalizeBuilder, type ItemBuilder, @@ -11,7 +10,6 @@ import { orderEventsByTimestamp, processEvent, readLastTurnInfo, - type TurnContext, } from "./buildConversationItems"; /** @@ -176,8 +174,13 @@ export function createIncrementalConversationBuilder() { markThoughtCompletion(builder.items); + // Rows keep their identity across calls — the builder replaces a row + // object whenever its content changes (tool merges, child streams, + // progress cards), so memoized views re-render exactly the changed rows. + // Turn flags and `thoughtComplete` are surfaced as value props by the + // renderers, not via row identity. return { - items: assembleItems(builder, activeStart), + items: builder.items.slice(), lastTurnInfo: readLastTurnInfoForOutput(builder), isCompacting: builder.isCompacting, isClearing: builder.isClearing, @@ -189,50 +192,6 @@ export function createIncrementalConversationBuilder() { return { update, reset }; } -function assembleItems( - b: ItemBuilder, - activeStart: number, -): ConversationItem[] { - // Completed turns: reuse the builder's own objects. They aren't rebuilt - // across calls, so their identity is stable and memoized rows skip work. - const out = b.items.slice(0, activeStart); - if (activeStart >= b.items.length) return out; - - const turn = b.currentTurn; - // The active turn streams: clone its rows onto a fresh shared context each - // call so their memoized views re-render and read the latest tool/child - // state — matching the all-new-objects behavior a full rebuild gives the - // live turn. Non-update rows (the user message, git actions) never change, - // so pass them through by reference. - const activeContext: TurnContext | null = turn - ? { - toolCalls: new Map(turn.context.toolCalls), - childItems: new Map(turn.context.childItems), - turnCancelled: turn.context.turnCancelled, - turnComplete: turn.context.turnComplete, - } - : null; - - for (let i = activeStart; i < b.items.length; i++) { - const item = b.items[i]; - // Only rows of the active turn get the fresh context. A prompt can open - // its turn *before* a trailing progress card from the previous turn (see - // `handlePromptRequest`), so the active range may hold older-turn rows — - // those keep their own (frozen) context, matching a full rebuild. - if ( - item.type === "session_update" && - activeContext && - turn && - item.turnContext === turn.context - ) { - out.push({ ...item, turnContext: activeContext }); - } else { - out.push(item); - } - } - return out; -} - function readLastTurnInfoForOutput(b: ItemBuilder) { const info = readLastTurnInfo(b); if (!info) return null; diff --git a/products/desktop/packages/ui/src/features/sessions/components/session-update/StatusNotificationView.tsx b/products/desktop/packages/ui/src/features/sessions/components/session-update/StatusNotificationView.tsx index 10fcef5fd9fc..9d3afce8ae30 100644 --- a/products/desktop/packages/ui/src/features/sessions/components/session-update/StatusNotificationView.tsx +++ b/products/desktop/packages/ui/src/features/sessions/components/session-update/StatusNotificationView.tsx @@ -233,7 +233,7 @@ function RetryingStatusView({ setRemainingMs(Math.max(0, delayMs - (Date.now() - start))); }; tick(); - const interval = setInterval(tick, 100); + const interval = setInterval(tick, 1000); return () => clearInterval(interval); }, [delayMs, startedAt]); @@ -282,7 +282,7 @@ function CompactingStatusView({ const start = startedAt ?? Date.now(); const tick = () => setElapsed(Date.now() - start); tick(); - const interval = setInterval(tick, 100); + const interval = setInterval(tick, 1000); return () => clearInterval(interval); }, [startedAt]); @@ -292,7 +292,7 @@ function CompactingStatusView({ {label} - {formatDuration(elapsed, 1)} + {formatDuration(elapsed, 0)} {/* Decorative: the spinner and the text above carry the accessible status. */} diff --git a/products/desktop/packages/ui/src/features/sessions/components/session-update/SubagentToolView.tsx b/products/desktop/packages/ui/src/features/sessions/components/session-update/SubagentToolView.tsx index ca1308fc78c3..7113d0131802 100644 --- a/products/desktop/packages/ui/src/features/sessions/components/session-update/SubagentToolView.tsx +++ b/products/desktop/packages/ui/src/features/sessions/components/session-update/SubagentToolView.tsx @@ -54,6 +54,7 @@ export function SubagentToolView({ childItems={turnContext.childItems} turnCancelled={turnContext.turnCancelled} turnComplete={turnContext.turnComplete} + thoughtComplete={child.thoughtComplete} /> ) : null, ) diff --git a/products/desktop/packages/ui/src/features/sessions/hooks/useChatTitleGenerator.ts b/products/desktop/packages/ui/src/features/sessions/hooks/useChatTitleGenerator.ts index 827a0637d585..122480be3dc2 100644 --- a/products/desktop/packages/ui/src/features/sessions/hooks/useChatTitleGenerator.ts +++ b/products/desktop/packages/ui/src/features/sessions/hooks/useChatTitleGenerator.ts @@ -1,4 +1,5 @@ import type { Schemas } from "@posthog/api-client"; +import { createAppendOnlyTracker } from "@posthog/core/sessions/appendOnlyTracker"; import { canApplyTitleFromPrompts, decideTitleGeneration, @@ -28,10 +29,20 @@ import { taskKeys } from "@posthog/ui/features/tasks/taskKeys"; import { logger } from "@posthog/ui/shell/logger"; import { titleAttachmentStoreApi } from "@posthog/ui/shell/titleAttachmentStore"; import { type QueryClient, useQueryClient } from "@tanstack/react-query"; -import { useEffect } from "react"; +import { useEffect, useMemo, useRef } from "react"; const log = logger.scope("chat-title-generator"); +function createPromptCountTracker() { + return createAppendOnlyTracker<{ count: number }, number>({ + init: () => ({ count: 0 }), + processEvent: (state, event) => { + state.count += extractUserPromptsFromEvents([event]).length; + }, + getResult: (state) => state.count, + }); +} + function getCachedTask( queryClient: QueryClient, taskId: string, @@ -55,13 +66,19 @@ export function useChatTitleGenerator(task: Task): void { (state) => state.status === "authenticated" && !!state.cloudRegion, ); - const promptCount = useSessionStore((state) => { + const events = useSessionStore((state) => { const taskRunId = state.taskIdIndex[taskId]; - if (!taskRunId) return 0; - const session = state.sessions[taskRunId]; - if (!session?.events) return 0; - return extractUserPromptsFromEvents(session.events).length; + return taskRunId ? state.sessions[taskRunId]?.events : undefined; }); + const promptCountTrackerRef = useRef | null>(null); + promptCountTrackerRef.current ??= createPromptCountTracker(); + const promptCountTracker = promptCountTrackerRef.current; + const promptCount = useMemo( + () => (events ? promptCountTracker.update(events) : 0), + [events, promptCountTracker], + ); useEffect(() => { if (!isAuthenticated || (task.created_by && !currentUser)) return; diff --git a/products/desktop/packages/ui/src/features/sessions/hooks/useSessionCallbacks.test.tsx b/products/desktop/packages/ui/src/features/sessions/hooks/useSessionCallbacks.test.tsx index 1bb918f22469..07ecb5035ca3 100644 --- a/products/desktop/packages/ui/src/features/sessions/hooks/useSessionCallbacks.test.tsx +++ b/products/desktop/packages/ui/src/features/sessions/hooks/useSessionCallbacks.test.tsx @@ -1,4 +1,3 @@ -import type { AgentSession } from "@posthog/shared"; import type { Task } from "@posthog/shared/domain-types"; import { renderHook } from "@testing-library/react"; import { beforeEach, describe, expect, it, vi } from "vitest"; @@ -92,12 +91,11 @@ import { useSessionCallbacks } from "./useSessionCallbacks"; const TASK = "task-1"; const task = { id: TASK, latest_run: null } as unknown as Task; -function renderCallbacks(session?: AgentSession) { +function renderCallbacks() { return renderHook(() => useSessionCallbacks({ taskId: TASK, task, - session, repoPath: "/repo", }), ); @@ -197,7 +195,7 @@ describe("useSessionCallbacks.handleSendPrompt", () => { messagingMode.value = "steer"; sessionService.sendPrompt.mockResolvedValue({ stopReason: "steered" }); - const { result } = renderCallbacks(sessionState as unknown as AgentSession); + const { result } = renderCallbacks(); await result.current.handleSendPrompt("change direction"); expect(sessionService.sendPrompt).toHaveBeenCalledWith( diff --git a/products/desktop/packages/ui/src/features/sessions/hooks/useSessionCallbacks.ts b/products/desktop/packages/ui/src/features/sessions/hooks/useSessionCallbacks.ts index f4b29a3e90ba..48807b469532 100644 --- a/products/desktop/packages/ui/src/features/sessions/hooks/useSessionCallbacks.ts +++ b/products/desktop/packages/ui/src/features/sessions/hooks/useSessionCallbacks.ts @@ -21,10 +21,7 @@ import { } from "@posthog/ui/features/message-editor/commands"; import { useDraftStore } from "@posthog/ui/features/message-editor/draftStore"; import { useMessagingMode } from "@posthog/ui/features/sessions/hooks/useMessagingMode"; -import { - type AgentSession, - sessionStoreSetters, -} from "@posthog/ui/features/sessions/sessionStore"; +import { sessionStoreSetters } from "@posthog/ui/features/sessions/sessionStore"; import { useTaskViewed } from "@posthog/ui/features/sidebar/useTaskViewed"; import { SHELL_CLIENT, @@ -33,21 +30,19 @@ import { import { toast } from "@posthog/ui/primitives/toast"; import { getAppViewSnapshot } from "@posthog/ui/router/useAppView"; import { logger } from "@posthog/ui/shell/logger"; -import { useCallback, useRef } from "react"; +import { useCallback } from "react"; const log = logger.scope("session-callbacks"); interface UseSessionCallbacksOptions { taskId: string; task: Task; - session: AgentSession | undefined; repoPath: string | null; } export function useSessionCallbacks({ taskId, task, - session, repoPath, }: UseSessionCallbacksOptions) { const sessionService = useService(SESSION_SERVICE); @@ -56,14 +51,11 @@ export function useSessionCallbacks({ const { markActivity, markAsViewed } = useTaskViewed(); const { requestFocus, setPendingContent } = useDraftStore((s) => s.actions); - const sessionRef = useRef(session); - sessionRef.current = session; - const messagingMode = useMessagingMode(taskId); const handleSendPrompt = useCallback( async (text: string): Promise => { - const currentSession = sessionRef.current; + const currentSession = sessionStoreSetters.getSessionByTaskId(taskId); const currentEvents = currentSession?.events ?? []; const handled = await tryExecuteCodeCommand(text, { taskId, @@ -173,6 +165,7 @@ export function useSessionCallbacks({ // composer would clobber the in-progress edit. The edit hold keeps the // queue from auto-sending until the edit is saved or cancelled. const currentSession = sessionStoreSetters.getSessionByTaskId(taskId); + const isCloud = currentSession?.isCloud ?? false; const editingId = currentSession?.editingQueuedId; if ( editingId && @@ -188,12 +181,12 @@ export function useSessionCallbacks({ const result = await sessionService.cancelPrompt(taskId); log.info("Prompt cancelled", { success: result }); - const queuedPrompt = sessionRef.current?.isCloud + const queuedPrompt = isCloud ? combineQueuedCloudPrompts(queuedMessages) : queuedMessages.map((message) => message.content).join("\n\n"); if (queuedPrompt) { - const pendingContent = sessionRef.current?.isCloud + const pendingContent = isCloud ? promptToQueuedEditorContent(queuedPrompt) : textToContent(typeof queuedPrompt === "string" ? queuedPrompt : ""); @@ -204,7 +197,7 @@ export function useSessionCallbacks({ const handleRetry = useCallback(async () => { try { - if (sessionRef.current?.isCloud) { + if (sessionStoreSetters.getSessionByTaskId(taskId)?.isCloud) { await sessionService.retryCloudTaskWatch(taskId); return; } diff --git a/products/desktop/packages/ui/src/features/task-detail/components/TaskHeaderActions.tsx b/products/desktop/packages/ui/src/features/task-detail/components/TaskHeaderActions.tsx index c9d2dea3b70c..e06c6c8e2a29 100644 --- a/products/desktop/packages/ui/src/features/task-detail/components/TaskHeaderActions.tsx +++ b/products/desktop/packages/ui/src/features/task-detail/components/TaskHeaderActions.tsx @@ -18,7 +18,7 @@ import { HandoffConfirmDialog } from "@posthog/ui/features/sessions/components/H import { StopCloudRunButton } from "@posthog/ui/features/sessions/components/StopCloudRunButton"; import { useHandoffDialogStore } from "@posthog/ui/features/sessions/handoffDialogStore"; import { useSessionCallbacks } from "@posthog/ui/features/sessions/hooks/useSessionCallbacks"; -import { useSessionForTask } from "@posthog/ui/features/sessions/useSession"; +import { useSessionHandoffInProgress } from "@posthog/ui/features/sessions/useSession"; import { useIsCloudTask, useWorkspace, @@ -31,7 +31,7 @@ import { useState } from "react"; const CLOUD_HANDOFF_FLAG = "phc-cloud-handoff"; function LocalHandoffButton({ taskId, task }: { taskId: string; task: Task }) { - const session = useSessionForTask(taskId); + const inProgress = useSessionHandoffInProgress(taskId); const workspace = useWorkspace(taskId); const repoPath = workspace?.folderPath ?? null; const authStatus = useAuthStateValue((s) => s.status); @@ -40,7 +40,6 @@ function LocalHandoffButton({ taskId, task }: { taskId: string; task: Task }) { const { initiateHandoffToCloud } = useSessionCallbacks({ taskId, task, - session: session ?? undefined, repoPath, }); @@ -68,8 +67,6 @@ function LocalHandoffButton({ taskId, task }: { taskId: string; task: Task }) { } }; - const inProgress = session?.handoffInProgress ?? false; - return ( <>
diff --git a/products/desktop/packages/ui/src/features/task-detail/components/TaskLogsPanel.tsx b/products/desktop/packages/ui/src/features/task-detail/components/TaskLogsPanel.tsx index 1cc72833562d..6c827eeb79dd 100644 --- a/products/desktop/packages/ui/src/features/task-detail/components/TaskLogsPanel.tsx +++ b/products/desktop/packages/ui/src/features/task-detail/components/TaskLogsPanel.tsx @@ -85,7 +85,7 @@ export function TaskLogsPanel({ taskId, task, hideInput }: TaskLogsPanelProps) { handleRetry, handleNewSession, handleBashCommand, - } = useSessionCallbacks({ taskId, task, session, repoPath }); + } = useSessionCallbacks({ taskId, task, repoPath }); const { handleBeforeSubmit, dialogProps } = useBranchMismatchDialog({ taskId, diff --git a/products/desktop/packages/ui/src/features/task-detail/hooks/useCloudEventSummary.ts b/products/desktop/packages/ui/src/features/task-detail/hooks/useCloudEventSummary.ts index 4a1d1effa779..8cfcdc8cdcc9 100644 --- a/products/desktop/packages/ui/src/features/task-detail/hooks/useCloudEventSummary.ts +++ b/products/desktop/packages/ui/src/features/task-detail/hooks/useCloudEventSummary.ts @@ -1,22 +1,31 @@ import { - buildCloudEventSummary, type CloudEventSummary, + createCloudEventSummaryTracker, } from "@posthog/core/task-detail/cloudToolChanges"; -import { useMemo } from "react"; -import { useSessionForTask } from "../../sessions/useSession"; +import { useMemo, useRef } from "react"; +import { useSessionSelector } from "../../sessions/useSession"; const EMPTY_SUMMARY: CloudEventSummary = { toolCalls: new Map(), + revision: 0, + changedFilesRevision: 0, }; export function useCloudEventSummary( taskId: string, enabled = true, ): CloudEventSummary { - const session = useSessionForTask(enabled ? taskId : undefined); - const events = session?.events; + const events = useSessionSelector( + enabled ? taskId : undefined, + (session) => session?.events, + ); + const trackerRef = useRef | null>(null); + trackerRef.current ??= createCloudEventSummaryTracker(); + const tracker = trackerRef.current; return useMemo( - () => (events ? buildCloudEventSummary(events) : EMPTY_SUMMARY), - [events], + () => (events ? tracker.update(events) : EMPTY_SUMMARY), + [events, tracker], ); } diff --git a/products/desktop/packages/ui/src/features/task-detail/hooks/useCloudRunState.ts b/products/desktop/packages/ui/src/features/task-detail/hooks/useCloudRunState.ts index cff2f5cec040..6cbf209cf37c 100644 --- a/products/desktop/packages/ui/src/features/task-detail/hooks/useCloudRunState.ts +++ b/products/desktop/packages/ui/src/features/task-detail/hooks/useCloudRunState.ts @@ -1,9 +1,10 @@ import { deriveCloudRunState } from "@posthog/core/task-detail/cloudRunState"; import { extractCloudToolChangedFiles } from "@posthog/core/task-detail/cloudToolChanges"; import type { Task } from "@posthog/shared/domain-types"; -import { useMemo } from "react"; +import { useMemo, useRef } from "react"; +import { shallow } from "zustand/shallow"; import { resolveCloudPrUrl } from "../../git-interaction/cloudPrUrl"; -import { useSessionForTask } from "../../sessions/useSession"; +import { useSessionSelector } from "../../sessions/useSession"; import { pickFreshestTask } from "../../tasks/taskFreshness"; import { useTasks } from "../../tasks/useTasks"; import { useCloudEventSummary } from "./useCloudEventSummary"; @@ -19,17 +20,44 @@ export function useCloudRunState(taskId: string, task: Task) { [task, taskId, tasks], ); - const session = useSessionForTask(taskId); + const session = useSessionSelector( + taskId, + (current) => + current + ? { + taskRunId: current.taskRunId, + cloudBranch: current.cloudBranch, + cloudStatus: current.cloudStatus, + cloudOutput: current.cloudOutput, + } + : undefined, + shallow, + ); const prUrl = resolveCloudPrUrl(freshTask, session); const { effectiveBranch, repo, cloudStatus, isRunActive } = deriveCloudRunState(freshTask, session, prUrl); const summary = useCloudEventSummary(taskId); - const fallbackFiles = useMemo( - () => extractCloudToolChangedFiles(summary.toolCalls), - [summary], - ); + const fallbackFilesRef = useRef< + | { + taskId: string; + revision: number; + files: ReturnType; + } + | undefined + >(undefined); + if ( + fallbackFilesRef.current?.taskId !== taskId || + fallbackFilesRef.current.revision !== summary.changedFilesRevision + ) { + fallbackFilesRef.current = { + taskId, + revision: summary.changedFilesRevision, + files: extractCloudToolChangedFiles(summary.toolCalls), + }; + } + const fallbackFiles = fallbackFilesRef.current.files; return { freshTask, diff --git a/products/desktop/packages/ui/src/styles/globals.css b/products/desktop/packages/ui/src/styles/globals.css index 0b7a55adecff..8ba300839571 100644 --- a/products/desktop/packages/ui/src/styles/globals.css +++ b/products/desktop/packages/ui/src/styles/globals.css @@ -454,16 +454,15 @@ body:has(.rt-DialogOverlay[data-state="open"]) [data-quill-portal] { 0%, 100% { opacity: 1; - color: var(--accent-9); } 50% { opacity: 0.5; - color: var(--accent-11); } } .ph-pulse { - animation: ph-pulse 1s ease-in-out infinite; + color: var(--accent-9); + animation: ph-pulse 1s step-end infinite; } /* Shimmer for an item whose run is still in flight. A highlight band sweeps @@ -508,8 +507,7 @@ body:has(.rt-DialogOverlay[data-state="open"]) [data-quill-portal] { /* Freeze CSS keyframe animations while the app window is unfocused or hidden. Perpetual indicators (spinners, pulses, floats) otherwise keep the compositor - and GPU busy in the background for animations nobody is looking at — and some, - like ph-pulse's color shift, force main-thread repaints every frame. The class + and GPU busy in the background for animations nobody is looking at. The class is toggled from the renderer focus store (see GlobalEventHandlers). Transitions are unaffected, and animations resume where they left off on refocus. */ body.ph-window-blurred *, diff --git a/products/desktop/packages/workspace-server/src/services/process-tracking/process-tracking.test.ts b/products/desktop/packages/workspace-server/src/services/process-tracking/process-tracking.test.ts index c7b8eb8e6db9..f6c8864a3ae5 100644 --- a/products/desktop/packages/workspace-server/src/services/process-tracking/process-tracking.test.ts +++ b/products/desktop/packages/workspace-server/src/services/process-tracking/process-tracking.test.ts @@ -3,6 +3,7 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; const mockPlatform = vi.hoisted(() => vi.fn(() => "darwin")); const mockIsProcessAlive = vi.hoisted(() => vi.fn((_pid: number) => true)); const mockKillProcessTree = vi.hoisted(() => vi.fn()); +const mockKillProcessTrees = vi.hoisted(() => vi.fn()); const mockExecAsync = vi.hoisted(() => vi.fn()); vi.mock("node:child_process", () => ({ @@ -23,6 +24,7 @@ vi.mock("node:os", () => ({ vi.mock("./process-utils", () => ({ isProcessAlive: mockIsProcessAlive, killProcessTree: mockKillProcessTree, + killProcessTrees: mockKillProcessTrees, })); import { ProcessTrackingService } from "./process-tracking"; @@ -314,9 +316,7 @@ describe("ProcessTrackingService", () => { service.killByCategory("shell"); - expect(mockKillProcessTree).toHaveBeenCalledWith(1); - expect(mockKillProcessTree).toHaveBeenCalledWith(2); - expect(mockKillProcessTree).not.toHaveBeenCalledWith(3); + expect(mockKillProcessTrees).toHaveBeenCalledWith([1, 2]); expect(service.getByCategory("shell")).toHaveLength(0); expect(service.getByCategory("agent")).toHaveLength(1); }); @@ -327,6 +327,7 @@ describe("ProcessTrackingService", () => { service.killByCategory("shell"); expect(mockKillProcessTree).not.toHaveBeenCalled(); + expect(mockKillProcessTrees).toHaveBeenCalledWith([]); expect(service.getAll()).toHaveLength(1); }); }); @@ -363,9 +364,7 @@ describe("ProcessTrackingService", () => { service.killByTaskId("task-1"); - expect(mockKillProcessTree).toHaveBeenCalledWith(1); - expect(mockKillProcessTree).toHaveBeenCalledWith(2); - expect(mockKillProcessTree).not.toHaveBeenCalledWith(3); + expect(mockKillProcessTrees).toHaveBeenCalledWith([1, 2]); expect(service.getByTaskId("task-1")).toEqual([]); expect(service.getByTaskId("task-2")).toHaveLength(1); }); @@ -376,6 +375,7 @@ describe("ProcessTrackingService", () => { service.killByTaskId("task-999"); expect(mockKillProcessTree).not.toHaveBeenCalled(); + expect(mockKillProcessTrees).toHaveBeenCalledWith([]); expect(service.getAll()).toHaveLength(1); }); }); @@ -426,9 +426,7 @@ describe("ProcessTrackingService", () => { service.killAll(); - expect(mockKillProcessTree).toHaveBeenCalledWith(1); - expect(mockKillProcessTree).toHaveBeenCalledWith(2); - expect(mockKillProcessTree).toHaveBeenCalledWith(3); + expect(mockKillProcessTrees).toHaveBeenCalledWith([1, 2, 3]); expect(service.getAll()).toHaveLength(0); }); @@ -436,6 +434,7 @@ describe("ProcessTrackingService", () => { service.killAll(); expect(mockKillProcessTree).not.toHaveBeenCalled(); + expect(mockKillProcessTrees).toHaveBeenCalledWith([]); }); }); }); diff --git a/products/desktop/packages/workspace-server/src/services/process-tracking/process-tracking.ts b/products/desktop/packages/workspace-server/src/services/process-tracking/process-tracking.ts index 809ee079aa80..c2ce6543c90d 100644 --- a/products/desktop/packages/workspace-server/src/services/process-tracking/process-tracking.ts +++ b/products/desktop/packages/workspace-server/src/services/process-tracking/process-tracking.ts @@ -2,7 +2,11 @@ import { exec } from "node:child_process"; import { platform } from "node:os"; import { promisify } from "node:util"; import { injectable, preDestroy } from "inversify"; -import { isProcessAlive, killProcessTree } from "./process-utils"; +import { + isProcessAlive, + killProcessTree, + killProcessTrees, +} from "./process-utils"; import type { DiscoveredProcess, ProcessCategory, @@ -195,25 +199,24 @@ export class ProcessTrackingService { killByCategory(category: ProcessCategory): void { const procs = this.getByCategory(category); - for (const proc of procs) { - this.kill(proc.pid); - } + this.killProcesses(procs); } killByTaskId(taskId: string): void { const procs = this.getByTaskId(taskId); - for (const proc of procs) { - this.kill(proc.pid); - } + this.killProcesses(procs); + } + + private killProcesses(procs: readonly TrackedProcess[]): void { + killProcessTrees(procs.map((proc) => proc.pid)); + for (const proc of procs) this.unregister(proc.pid, "killed"); } @preDestroy() killAll(): void { this._isShuttingDown = true; - for (const proc of this.processes.values()) { - killProcessTree(proc.pid); - } + killProcessTrees(Array.from(this.processes.keys())); this.processes.clear(); this.taskProcesses.clear(); } diff --git a/products/desktop/packages/workspace-server/src/services/process-tracking/process-utils.test.ts b/products/desktop/packages/workspace-server/src/services/process-tracking/process-utils.test.ts new file mode 100644 index 000000000000..16ff9e607004 --- /dev/null +++ b/products/desktop/packages/workspace-server/src/services/process-tracking/process-utils.test.ts @@ -0,0 +1,124 @@ +import { describe, expect, it, vi } from "vitest"; +import { + findMatchingProcessTargets, + findProcessTree, + killUnixProcessTrees, +} from "./process-utils"; + +const startedAt = "Sat Jul 25 00:00:00 2026"; + +describe("findProcessTree", () => { + it("returns descendants deepest-first across nested process groups", () => { + const result = findProcessTree(10, [ + { pid: 1, ppid: 0, pgid: 1, startedAt }, + { pid: 10, ppid: 1, pgid: 10, startedAt }, + { pid: 11, ppid: 10, pgid: 10, startedAt }, + { pid: 12, ppid: 11, pgid: 12, startedAt }, + { pid: 13, ppid: 12, pgid: 12, startedAt }, + { pid: 20, ppid: 1, pgid: 20, startedAt }, + ]); + + expect(result).toEqual([ + { pid: 13, ppid: 12, pgid: 12, startedAt }, + { pid: 12, ppid: 11, pgid: 12, startedAt }, + { pid: 11, ppid: 10, pgid: 10, startedAt }, + { pid: 10, ppid: 1, pgid: 10, startedAt }, + ]); + }); + + it("returns no unrelated processes when the root has exited", () => { + expect( + findProcessTree(10, [{ pid: 20, ppid: 1, pgid: 20, startedAt }]), + ).toEqual([]); + }); +}); + +describe("findMatchingProcessTargets", () => { + it("excludes reused process and group ids", () => { + const original = [ + { pid: 10, ppid: 1, pgid: 10, startedAt }, + { pid: 11, ppid: 10, pgid: 11, startedAt }, + ]; + const current = [ + { pid: 10, ppid: 1, pgid: 10, startedAt: "later" }, + { pid: 11, ppid: 1, pgid: 11, startedAt: "later" }, + ]; + + expect(findMatchingProcessTargets(original, current, undefined)).toEqual( + [], + ); + }); + + it("targets surviving descendants after they are reparented", () => { + const original = [ + { pid: 11, ppid: 10, pgid: 11, startedAt }, + { pid: 10, ppid: 1, pgid: 10, startedAt }, + ]; + const current = [{ pid: 11, ppid: 1, pgid: 11, startedAt }]; + + expect(findMatchingProcessTargets(original, current, undefined)).toEqual([ + -11, 11, + ]); + }); +}); + +describe("killUnixProcessTrees", () => { + it("revalidates identities before the delayed kill", () => { + const original = [ + { pid: 1, ppid: 0, pgid: 1, startedAt }, + { pid: 10, ppid: 1, pgid: 10, startedAt }, + { pid: 11, ppid: 10, pgid: 11, startedAt }, + ]; + const current = [ + { pid: 10, ppid: 1, pgid: 10, startedAt: "reused" }, + { pid: 11, ppid: 1, pgid: 11, startedAt }, + ]; + const signal = vi.fn(); + let delayed: (() => void) | undefined; + + killUnixProcessTrees([10], original, 1, { + currentProcesses: () => current, + signal, + schedule: (callback) => { + delayed = callback; + }, + }); + delayed?.(); + + expect(signal.mock.calls).toEqual([ + [[-10, -11, 10, 11], "SIGTERM"], + [[-11, 11], "SIGKILL"], + ]); + }); + + it("falls back to the raw group when the root identity is unavailable", () => { + const signal = vi.fn(); + const schedule = vi.fn(); + + killUnixProcessTrees([10], [{ pid: 20, ppid: 1, pgid: 20, startedAt }], 1, { + currentProcesses: () => [], + signal, + schedule, + }); + + expect(signal).toHaveBeenCalledWith([-10, 10], "SIGTERM"); + expect(schedule).not.toHaveBeenCalled(); + }); + + it("falls back for missing roots in a mixed batch", () => { + const signal = vi.fn(); + + killUnixProcessTrees( + [10, 20], + [{ pid: 10, ppid: 1, pgid: 10, startedAt }], + 1, + { + currentProcesses: () => [], + signal, + schedule: vi.fn(), + }, + ); + + expect(signal).toHaveBeenCalledWith([-20, 20, -10, 10], "SIGTERM"); + }); +}); diff --git a/products/desktop/packages/workspace-server/src/services/process-tracking/process-utils.ts b/products/desktop/packages/workspace-server/src/services/process-tracking/process-utils.ts index 5cd9d4e68652..4318b964c081 100644 --- a/products/desktop/packages/workspace-server/src/services/process-tracking/process-utils.ts +++ b/products/desktop/packages/workspace-server/src/services/process-tracking/process-utils.ts @@ -1,44 +1,197 @@ -import { execSync } from "node:child_process"; +import { execFileSync, execSync } from "node:child_process"; import { platform } from "node:os"; const SIGKILL_GRACE_MS = 5_000; +export interface ProcessEntry { + pid: number; + ppid: number; + pgid: number; + startedAt: string; +} + +export function findProcessTree( + rootPid: number, + processTable: readonly ProcessEntry[], +): ProcessEntry[] { + return findProcessTreeFromIndex( + rootPid, + processTable, + indexProcesses(processTable), + ); +} + +function indexProcesses( + processTable: readonly ProcessEntry[], +): Map { + const children = new Map(); + for (const entry of processTable) { + const siblings = children.get(entry.ppid) ?? []; + siblings.push(entry); + children.set(entry.ppid, siblings); + } + return children; +} + +function findProcessTreeFromIndex( + rootPid: number, + processTable: readonly ProcessEntry[], + children: ReadonlyMap, +): ProcessEntry[] { + const tree: ProcessEntry[] = []; + const visit = (pid: number): void => { + for (const child of children.get(pid) ?? []) { + visit(child.pid); + tree.push(child); + } + }; + visit(rootPid); + + const root = processTable.find((entry) => entry.pid === rootPid); + if (root) tree.push(root); + return tree; +} + +function snapshotUnixProcesses(): ProcessEntry[] { + try { + const output = execFileSync("ps", ["-axo", "pid=,ppid=,pgid=,lstart="], { + encoding: "utf8", + }); + return output + .trim() + .split("\n") + .map((line): ProcessEntry | null => { + const match = line.trim().match(/^(\d+)\s+(\d+)\s+(\d+)\s+(.+)$/); + if (!match) return null; + return { + pid: Number(match[1]), + ppid: Number(match[2]), + pgid: Number(match[3]), + startedAt: match[4], + }; + }) + .filter((entry): entry is ProcessEntry => entry !== null); + } catch { + return []; + } +} + +export function findMatchingProcessTargets( + originalTree: readonly ProcessEntry[], + currentProcesses: readonly ProcessEntry[], + excludedPgid: number | undefined, +): number[] { + const originalByPid = new Map( + originalTree.map((entry) => [entry.pid, entry]), + ); + const matching = currentProcesses.filter((entry) => { + const original = originalByPid.get(entry.pid); + return ( + original?.pgid === entry.pgid && original.startedAt === entry.startedAt + ); + }); + const groups = new Set( + matching + .map((entry) => entry.pgid) + .filter((pgid) => pgid > 0 && pgid !== excludedPgid), + ); + return [ + ...Array.from(groups, (pgid) => -pgid), + ...matching.map((entry) => entry.pid), + ]; +} + +export interface UnixProcessKillerDeps { + currentProcesses: () => ProcessEntry[]; + signal: (targets: readonly number[], signal: NodeJS.Signals) => void; + schedule: (callback: () => void, delayMs: number) => void; +} + +export function killUnixProcessTrees( + rootPids: readonly number[], + initialProcesses: readonly ProcessEntry[], + ownPgid: number | undefined, + deps: UnixProcessKillerDeps, +): void { + const children = indexProcesses(initialProcesses); + const missingRootPids: number[] = []; + const trees = rootPids.flatMap((pid) => { + const tree = findProcessTreeFromIndex(pid, initialProcesses, children); + if (tree.length === 0) missingRootPids.push(pid); + return tree; + }); + const originalTree = Array.from( + new Map(trees.flat().map((entry) => [entry.pid, entry])).values(), + ); + if (originalTree.length === 0 || ownPgid === undefined) { + deps.signal( + rootPids.flatMap((pid) => [-pid, pid]), + "SIGTERM", + ); + return; + } + + const targets = [ + ...missingRootPids.flatMap((pid) => [-pid, pid]), + ...findMatchingProcessTargets(originalTree, initialProcesses, ownPgid), + ]; + deps.signal(targets, "SIGTERM"); + deps.schedule(() => { + deps.signal( + findMatchingProcessTargets( + originalTree, + deps.currentProcesses(), + ownPgid, + ), + "SIGKILL", + ); + }, SIGKILL_GRACE_MS); +} + +function signalTargets( + targets: readonly number[], + signal: NodeJS.Signals, +): void { + for (const target of targets) { + try { + process.kill(target, signal); + } catch {} + } +} + /** - * Kill a process and all its children by killing the process group. - * On Unix, we use process.kill(-pid) to kill the entire process group. + * Kill a process and all its descendants, including children that created + * their own process groups. * On Windows, we use taskkill with /T flag to kill the process tree. */ -export function killProcessTree(pid: number): void { - try { - if (platform() === "win32") { - // Windows: use taskkill with /T to kill process tree - execSync(`taskkill /PID ${pid} /T /F`, { stdio: "ignore" }); - } else { - // SIGTERM the process group first, fall back to individual process - let sent = false; - for (const target of [-pid, pid]) { - try { - process.kill(target, "SIGTERM"); - sent = true; - break; - } catch {} - } - - if (!sent) return; - - // Force kill after a grace period — unref so the timer doesn't delay app exit. - // We skip the liveness check since isProcessAlive only tests the group leader; - // orphaned children in the same group would be missed. The catch blocks - // handle ESRCH if everything already exited. - setTimeout(() => { - for (const target of [-pid, pid]) { - try { - process.kill(target, "SIGKILL"); - } catch {} - } - }, SIGKILL_GRACE_MS).unref(); +export function killProcessTrees(pids: readonly number[]): void { + if (pids.length === 0) return; + if (platform() === "win32") { + // Windows: use taskkill with /T to kill process tree + for (const pid of pids) { + try { + execSync(`taskkill /PID ${pid} /T /F`, { stdio: "ignore" }); + } catch {} } - } catch {} + } else { + try { + const processes = snapshotUnixProcesses(); + const ownPgid = processes.find( + (entry) => entry.pid === process.pid, + )?.pgid; + killUnixProcessTrees(pids, processes, ownPgid, { + currentProcesses: snapshotUnixProcesses, + signal: signalTargets, + schedule: (callback, delayMs) => { + setTimeout(callback, delayMs).unref(); + }, + }); + } catch {} + } +} + +export function killProcessTree(pid: number): void { + killProcessTrees([pid]); } /** diff --git a/products/desktop/packages/workspace-server/src/services/watcher/service.test.ts b/products/desktop/packages/workspace-server/src/services/watcher/service.test.ts index ce3280e85832..4239522dacc1 100644 --- a/products/desktop/packages/workspace-server/src/services/watcher/service.test.ts +++ b/products/desktop/packages/workspace-server/src/services/watcher/service.test.ts @@ -1,6 +1,12 @@ import { afterEach, describe, expect, it, vi } from "vitest"; import type { FileWatcherEvent, WatcherEvent } from "./schemas"; -import { DEBOUNCE_MS, MAX_WAIT_MS, WatcherService } from "./service"; +import { + accumulateFsEvents, + DEBOUNCE_MS, + drainPending, + MAX_WAIT_MS, + WatcherService, +} from "./service"; afterEach(() => { vi.useRealTimers(); @@ -209,3 +215,26 @@ describe("WatcherService.watchRepo debounce", () => { await done; }); }); + +describe("WatcherService bulk changes", () => { + it("stops retaining paths after crossing the bulk threshold", () => { + const pending: Parameters[0] = { + dirs: new Set(), + files: new Set(), + deletes: new Set(), + bulkChanged: false, + }; + const events = Array.from({ length: 1000 }, (_, index) => ({ + type: "update" as const, + path: `/repo/dist/${index}.js`, + })); + + accumulateFsEvents(pending, events); + + expect(pending.files.size).toBe(0); + expect(drainPending("/repo", pending)).toEqual([ + { kind: "working-tree-changed", repoPath: "/repo" }, + ]); + expect(pending.bulkChanged).toBe(false); + }); +}); diff --git a/products/desktop/packages/workspace-server/src/services/watcher/service.ts b/products/desktop/packages/workspace-server/src/services/watcher/service.ts index 44d364f99da7..a382de7a3e56 100644 --- a/products/desktop/packages/workspace-server/src/services/watcher/service.ts +++ b/products/desktop/packages/workspace-server/src/services/watcher/service.ts @@ -29,7 +29,7 @@ export const DEBOUNCE_MS = 500; // wise never trip it until it paused, freezing the diff panel/stats mid-run. // The max-wait forces a flush at least this often during sustained activity so // the UI keeps advancing while the agent works. -export const MAX_WAIT_MS = 1000; +export const MAX_WAIT_MS = 3000; const BULK_THRESHOLD = 100; const dirname = (p: string): string => { @@ -51,12 +51,14 @@ interface Pending { dirs: Set; files: Set; deletes: Set; + bulkChanged: boolean; } const createPending = (): Pending => ({ dirs: new Set(), files: new Set(), deletes: new Set(), + bulkChanged: false, }); export const accumulateFsEvents = ( @@ -64,9 +66,16 @@ export const accumulateFsEvents = ( events: WatcherEvent[], ): void => { for (const event of events) { + if (pending.bulkChanged) continue; pending.dirs.add(dirname(event.path)); if (event.type === "delete") pending.deletes.add(event.path); else pending.files.add(event.path); + if (pending.files.size + pending.deletes.size > BULK_THRESHOLD) { + pending.dirs.clear(); + pending.files.clear(); + pending.deletes.clear(); + pending.bulkChanged = true; + } } }; @@ -76,12 +85,14 @@ export const drainPending = ( ): FileWatcherEvent[] => { const totalChanges = pending.files.size + pending.deletes.size; const out: FileWatcherEvent[] = []; - if (totalChanges === 0 && pending.dirs.size === 0) return out; + if (!pending.bulkChanged && totalChanges === 0 && pending.dirs.size === 0) { + return out; + } - if (totalChanges > 0) { + if (pending.bulkChanged || totalChanges > 0) { out.push({ kind: "working-tree-changed", repoPath }); } - if (totalChanges <= BULK_THRESHOLD) { + if (!pending.bulkChanged) { for (const dirPath of pending.dirs) out.push({ kind: "directory-changed", repoPath, dirPath }); for (const filePath of pending.files) @@ -92,6 +103,7 @@ export const drainPending = ( pending.dirs.clear(); pending.files.clear(); pending.deletes.clear(); + pending.bulkChanged = false; return out; }; From 5379a6474d58f0fac3ec10191b0fa8bc7ef0df2d Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Wed, 19 Aug 2026 14:36:37 +0300 Subject: [PATCH 2/8] fix(desktop): address runtime review findings --- .../src/sessions/sessionEventBatching.test.ts | 80 ++++++++++++++++++- .../core/src/sessions/sessionService.ts | 77 +++++++++++++----- .../src/task-detail/cloudToolChanges.test.ts | 47 +++++++++++ .../core/src/task-detail/cloudToolChanges.ts | 26 +++--- .../freeform/useCanvasGenerationToasts.ts | 2 +- .../components/buildConversationItems.ts | 51 ++++++------ .../task-detail/hooks/useCloudEventSummary.ts | 9 +-- .../task-detail/hooks/useCloudRunState.ts | 23 +----- .../process-tracking/process-utils.test.ts | 31 +++++++ .../process-tracking/process-utils.ts | 20 +++-- 10 files changed, 270 insertions(+), 96 deletions(-) diff --git a/products/desktop/packages/core/src/sessions/sessionEventBatching.test.ts b/products/desktop/packages/core/src/sessions/sessionEventBatching.test.ts index 2008bc46f3aa..2305f36c0322 100644 --- a/products/desktop/packages/core/src/sessions/sessionEventBatching.test.ts +++ b/products/desktop/packages/core/src/sessions/sessionEventBatching.test.ts @@ -77,6 +77,42 @@ function configOptionUpdate(): AcpMessage { } as unknown as AcpMessage; } +function usageUpdate(used: number, size: number): AcpMessage { + return { + ts: 4, + message: { + jsonrpc: "2.0", + method: "session/update", + params: { + sessionId: RUN_ID, + update: { sessionUpdate: "usage_update", used, size }, + }, + }, + } as unknown as AcpMessage; +} + +function completedSpeechToolCall(): AcpMessage { + return { + ts: 5, + message: { + jsonrpc: "2.0", + method: "session/update", + params: { + sessionId: RUN_ID, + update: { + sessionUpdate: "tool_call", + toolCallId: "speak-1", + status: "completed", + _meta: { + claudeCode: { toolName: "mcp__posthog-code-tools__speak" }, + }, + rawInput: { text: "Finished the task", kind: "done" }, + }, + }, + }, + } as unknown as AcpMessage; +} + function createHarness() { const sessions: Record = { [RUN_ID]: { @@ -127,6 +163,8 @@ function createHarness() { }; const notifyPromptComplete = vi.fn(); + const enqueueSpeech = vi.fn(); + const setPersistedConfigOptions = vi.fn(); const deps = { store, log: noopLog, @@ -146,10 +184,10 @@ function createHarness() { ); } }, - enqueueSpeech: vi.fn(), + enqueueSpeech, taskViewedApi: { markActivity: vi.fn() }, getPersistedConfigOptions: () => undefined, - setPersistedConfigOptions: vi.fn(), + setPersistedConfigOptions, trpc: { agent: { onSessionEvent: { @@ -182,11 +220,13 @@ function createHarness() { return { service, appendEvents, + enqueueSpeech, notifyPromptComplete, - setPersistedConfigOptions: deps.setPersistedConfigOptions, + setPersistedConfigOptions, updateSession: store.updateSession, emit: (event: AcpMessage) => onEvent?.(event), events: () => sessions[RUN_ID].events, + session: () => sessions[RUN_ID], }; } @@ -268,6 +308,31 @@ describe("streamed event batching", () => { expect(h.appendEvents).toHaveBeenCalledOnce(); }); + it("applies context usage from a batched update", () => { + const h = createHarness(); + + h.emit(usageUpdate(25, 100)); + vi.advanceTimersByTime(FLUSH_MS); + + expect(h.session()).toMatchObject({ contextUsed: 25, contextSize: 100 }); + }); + + it("enqueues completed speech from a batched tool update", () => { + const h = createHarness(); + + h.emit(completedSpeechToolCall()); + vi.advanceTimersByTime(FLUSH_MS); + + expect(h.enqueueSpeech).toHaveBeenCalledWith({ + text: "Finished the task", + taskTitle: "Local Task", + taskId: TASK_ID, + kind: "done", + source: "agent", + addressByName: true, + }); + }); + it("applies and persists config option updates immediately", () => { const h = createHarness(); const streamed = chunk("a"); @@ -280,6 +345,15 @@ describe("streamed event batching", () => { expect(h.setPersistedConfigOptions).toHaveBeenCalledOnce(); }); + it("contains errors from immediate session updates", () => { + const h = createHarness(); + h.setPersistedConfigOptions.mockImplementation(() => { + throw new Error("persist failed"); + }); + + expect(() => h.emit(configOptionUpdate())).not.toThrow(); + }); + it("keeps the turn duration when the prompt mutation clears state before the response flushes", () => { const h = createHarness(); diff --git a/products/desktop/packages/core/src/sessions/sessionService.ts b/products/desktop/packages/core/src/sessions/sessionService.ts index b7e6d477bfa7..47235668bb92 100644 --- a/products/desktop/packages/core/src/sessions/sessionService.ts +++ b/products/desktop/packages/core/src/sessions/sessionService.ts @@ -1678,6 +1678,11 @@ function isStreamUpdateEvent(acpMsg: AcpMessage): boolean { return update?.sessionUpdate !== "config_option_update"; } +interface SessionEventHandlingOptions { + appendedInBatch?: boolean; + promptStateUpdatedInBatch?: boolean; +} + export class SessionService { private connectingTasks = new Map>(); private reconcilingTasks = new Set(); @@ -2732,16 +2737,7 @@ export class SessionService { private lastAgentTextAt = new Map(); private enqueueSessionEvent(taskRunId: string, acpMsg: AcpMessage): void { - this.hostEndedResubscribes.delete(taskRunId); - this.lastSessionEventAt.set(taskRunId, Date.now()); - this.silenceLogged.delete(taskRunId); - this.ensureSilenceCheck(); - const stats = this.statsFor(taskRunId); - stats.received += 1; - stats.lastMethod = describeAcpMethod(acpMsg); - if (isAgentTextStreamEvent(acpMsg)) { - this.lastAgentTextAt.set(taskRunId, Date.now()); - } + this.recordSessionEventArrival(taskRunId, acpMsg); const buffered = this.pendingSessionEvents.get(taskRunId); if (buffered) { buffered.push(acpMsg); @@ -2756,6 +2752,22 @@ export class SessionService { } } + private recordSessionEventArrival( + taskRunId: string, + acpMsg: AcpMessage, + ): void { + this.hostEndedResubscribes.delete(taskRunId); + this.lastSessionEventAt.set(taskRunId, Date.now()); + this.silenceLogged.delete(taskRunId); + this.ensureSilenceCheck(); + const stats = this.statsFor(taskRunId); + stats.received += 1; + stats.lastMethod = describeAcpMethod(acpMsg); + if (isAgentTextStreamEvent(acpMsg)) { + this.lastAgentTextAt.set(taskRunId, Date.now()); + } + } + private flushSessionEvents(): void { if (this.pendingSessionEvents.size === 0) return; const batches = this.pendingSessionEvents; @@ -2780,8 +2792,16 @@ export class SessionService { * frozen with nothing in the log. Log it, count it, keep going. */ private applySessionEvent(taskRunId: string, acpMsg: AcpMessage): void { + this.applySessionEventWithOptions(taskRunId, acpMsg); + } + + private applySessionEventWithOptions( + taskRunId: string, + acpMsg: AcpMessage, + options: SessionEventHandlingOptions = {}, + ): void { try { - this.handleSessionEvent(taskRunId, acpMsg); + this.handleSessionEvent(taskRunId, acpMsg, options); } catch (error) { this.recordSessionEventFailure(taskRunId, acpMsg, error); } @@ -2840,6 +2860,12 @@ export class SessionService { error, ); } + for (const event of passive) { + this.applySessionEventWithOptions(taskRunId, event, { + appendedInBatch: true, + promptStateUpdatedInBatch: true, + }); + } } passive = []; }; @@ -3020,8 +3046,9 @@ export class SessionService { if (isStreamUpdateEvent(event)) { this.enqueueSessionEvent(taskRunId, event); } else { + this.recordSessionEventArrival(taskRunId, event); this.flushSessionEventsForTask(taskRunId); - this.handleSessionEvent(taskRunId, event); + this.applySessionEvent(taskRunId, event); } }, onError: (err) => { @@ -3467,7 +3494,11 @@ export class SessionService { return turnStart > 0 && spokeAt >= turnStart; } - private handleSessionEvent(taskRunId: string, acpMsg: AcpMessage): void { + private handleSessionEvent( + taskRunId: string, + acpMsg: AcpMessage, + options: SessionEventHandlingOptions = {}, + ): void { const session = this.d.store.getSessions()[taskRunId]; if (!session) return; @@ -3477,21 +3508,29 @@ export class SessionService { // Once the agent starts responding, clear initialPrompt so that // retry reconnects to this session instead of creating a new one. - if (!isUserPromptEcho && session.initialPrompt?.length) { + if ( + !options.appendedInBatch && + !isUserPromptEcho && + session.initialPrompt?.length + ) { this.d.store.updateSession(taskRunId, { initialPrompt: undefined, }); } - if (isUserPromptEcho && !this.isSteerMessage(acpMsg.message)) { - this.d.store.replaceOptimisticWithEvent(taskRunId, acpMsg); - } else { - this.d.store.appendEvents(taskRunId, [acpMsg]); + if (!options.appendedInBatch) { + if (isUserPromptEcho && !this.isSteerMessage(acpMsg.message)) { + this.d.store.replaceOptimisticWithEvent(taskRunId, acpMsg); + } else { + this.d.store.appendEvents(taskRunId, [acpMsg]); + } } const turnStartedAtTs = this.liveTurnContent.get(taskRunId)?.startedAtTs ?? session.promptStartedAt; - this.updatePromptStateFromEvents(taskRunId, [acpMsg], { isLive: true }); + if (!options.promptStateUpdatedInBatch) { + this.updatePromptStateFromEvents(taskRunId, [acpMsg], { isLive: true }); + } const msg = acpMsg.message; diff --git a/products/desktop/packages/core/src/task-detail/cloudToolChanges.test.ts b/products/desktop/packages/core/src/task-detail/cloudToolChanges.test.ts index b5f46e548382..df88df857d2d 100644 --- a/products/desktop/packages/core/src/task-detail/cloudToolChanges.test.ts +++ b/products/desktop/packages/core/src/task-detail/cloudToolChanges.test.ts @@ -87,6 +87,53 @@ describe("createCloudEventSummaryTracker", () => { expect(second.changedFilesRevision).toBe(first.changedFilesRevision); }); + + it("increments the changed-files revision and refreshes extracted files", () => { + const tracker = createCloudEventSummaryTracker(); + const started = toolEvent("tool-1", { + kind: "edit", + locations: [{ path: "src/file.ts", line: null }], + content: [ + { + type: "diff", + path: "src/file.ts", + oldText: "old", + newText: "first", + }, + ], + }); + const first = tracker.update([started]); + const changed = toolEvent("tool-1", { + content: [ + { + type: "diff", + path: "src/file.ts", + oldText: "old", + newText: "second\nline", + }, + ], + }); + + const second = tracker.update([started, changed]); + + expect({ + revision: second.changedFilesRevision, + files: second.changedFiles, + }).toMatchObject({ + revision: first.changedFilesRevision + 1, + files: [{ path: "src/file.ts", linesAdded: 2, linesRemoved: 1 }], + }); + }); + + it("does not mutate an earlier projected summary", () => { + const tracker = createCloudEventSummaryTracker(); + const started = toolEvent("tool-1", { status: "in_progress" }); + const first = tracker.update([started]); + + tracker.update([started, toolEvent("tool-1", { status: "completed" })]); + + expect(first.toolCalls.get("tool-1")?.status).toBe("in_progress"); + }); }); function diffObj( diff --git a/products/desktop/packages/core/src/task-detail/cloudToolChanges.ts b/products/desktop/packages/core/src/task-detail/cloudToolChanges.ts index eb74512fe1f6..bfefc2aec159 100644 --- a/products/desktop/packages/core/src/task-detail/cloudToolChanges.ts +++ b/products/desktop/packages/core/src/task-detail/cloudToolChanges.ts @@ -179,6 +179,7 @@ export function cachedDiffStats( export interface CloudEventSummary { toolCalls: Map; + changedFiles: ChangedFile[]; revision: number; changedFilesRevision: number; } @@ -248,21 +249,6 @@ function applyCloudEvent( }; } -/** - * Single-pass extraction of tool calls from events. - */ -export function buildCloudEventSummary( - events: AcpMessage[], -): CloudEventSummary { - const toolCalls = new Map(); - - for (const event of events) { - applyCloudEvent(toolCalls, event); - } - - return { toolCalls, revision: 0, changedFilesRevision: 0 }; -} - export function createCloudEventSummaryTracker(): { update(events: AcpMessage[]): CloudEventSummary; } { @@ -274,8 +260,11 @@ export function createCloudEventSummaryTracker(): { let projectedState: TrackerState | undefined; let projectedRevision = -1; + let projectedChangedFilesRevision = -1; + let projectedChangedFiles: ChangedFile[] = []; let projectedResult: CloudEventSummary = { toolCalls: new Map(), + changedFiles: [], revision: 0, changedFilesRevision: 0, }; @@ -292,11 +281,16 @@ export function createCloudEventSummaryTracker(): { if (change.changedFilesChanged) state.changedFilesRevision++; }, getResult: (state) => { + if (state.changedFilesRevision !== projectedChangedFilesRevision) { + projectedChangedFilesRevision = state.changedFilesRevision; + projectedChangedFiles = extractCloudToolChangedFiles(state.toolCalls); + } if (state !== projectedState || state.revision !== projectedRevision) { projectedState = state; projectedRevision = state.revision; projectedResult = { - toolCalls: state.toolCalls, + toolCalls: new Map(state.toolCalls), + changedFiles: projectedChangedFiles, revision: state.revision, changedFilesRevision: state.changedFilesRevision, }; diff --git a/products/desktop/packages/ui/src/features/canvas/freeform/useCanvasGenerationToasts.ts b/products/desktop/packages/ui/src/features/canvas/freeform/useCanvasGenerationToasts.ts index 0825bc1442d9..9f5acdfa1698 100644 --- a/products/desktop/packages/ui/src/features/canvas/freeform/useCanvasGenerationToasts.ts +++ b/products/desktop/packages/ui/src/features/canvas/freeform/useCanvasGenerationToasts.ts @@ -134,7 +134,7 @@ export function useCanvasGenerationToasts(): void { .map((id) => { const runId = state.taskIdIndex[id]; const session = runId ? state.sessions[runId] : undefined; - return `${id}:${session?.status ?? ""}:${session?.cloudStatus ?? ""}:${session?.isPromptPending ? 1 : 0}`; + return `${id}:${session?.taskRunId ?? ""}:${session?.status ?? ""}:${session?.isCloud ? 1 : 0}:${session?.cloudStatus ?? ""}:${session?.isPromptPending ? 1 : 0}:${session?.agentIdleForRunId ?? ""}`; }) .join("|"), ); diff --git a/products/desktop/packages/ui/src/features/sessions/components/buildConversationItems.ts b/products/desktop/packages/ui/src/features/sessions/components/buildConversationItems.ts index 5f48d15c704e..9d388e93fda2 100644 --- a/products/desktop/packages/ui/src/features/sessions/components/buildConversationItems.ts +++ b/products/desktop/packages/ui/src/features/sessions/components/buildConversationItems.ts @@ -124,6 +124,7 @@ interface TurnState { export interface ItemBuilder { items: ConversationItem[]; + toolCallRows: Map; currentTurn: TurnState | null; /** Index in `items` where the current turn's first item sits. Lets an * incremental consumer treat everything before it (completed turns) as @@ -159,6 +160,7 @@ export interface ItemBuilder { export function createItemBuilder(): ItemBuilder { return { items: [], + toolCallRows: new Map(), currentTurn: null, currentTurnStartIndex: 0, pendingPrompts: new Map(), @@ -234,13 +236,17 @@ function pushItem(b: ItemBuilder, update: RenderItem, ts?: number) { const turn = b.currentTurn; if (!turn) return; turn.itemCount++; - b.items.push({ + const item: ConversationItem = { type: "session_update", id: `${turn.id}-item-${turn.nextItemId++}`, update, turnContext: turn.context, timestamp: ts, - }); + }; + const index = b.items.push(item) - 1; + if (update.sessionUpdate === "tool_call") { + b.toolCallRows.set(update.toolCallId, { items: b.items, index }); + } } export interface BuildConversationOptions { @@ -1055,12 +1061,16 @@ function pushChildItem(b: ItemBuilder, parentId: string, update: RenderItem) { turn.context.childItems.set(parentId, children); } turn.itemCount++; - children.push({ + const item: ConversationItem = { type: "session_update", id: `${turn.id}-child-${turn.nextItemId++}`, update, turnContext: turn.context, - }); + }; + const index = children.push(item) - 1; + if (update.sessionUpdate === "tool_call") { + b.toolCallRows.set(update.toolCallId, { items: children, index }); + } reissueToolCallRow(b, parentId); } @@ -1125,28 +1135,19 @@ function reissueToolCallRow( if (!update) return; turn.toolCalls.set(toolCallId, update); - const visited = new Set(); - const replace = (items: ConversationItem[]): void => { - if (visited.has(items)) return; - visited.add(items); - for (let index = 0; index < items.length; index++) { - const item = items[index]; - if (item.type !== "session_update") continue; - if ( - item.update.sessionUpdate === "tool_call" && - item.update.toolCallId === toolCallId - ) { - items[index] = { - ...item, - update: { ...update, sessionUpdate: "tool_call" }, - }; - } - for (const children of item.turnContext.childItems.values()) { - replace(children); - } - } + const row = b.toolCallRows.get(toolCallId); + if (!row) return; + const item = row.items[row.index]; + if ( + item?.type !== "session_update" || + item.update.sessionUpdate !== "tool_call" + ) { + return; + } + row.items[row.index] = { + ...item, + update: { ...update, sessionUpdate: "tool_call" }, }; - replace(b.items); } function processSessionUpdate( diff --git a/products/desktop/packages/ui/src/features/task-detail/hooks/useCloudEventSummary.ts b/products/desktop/packages/ui/src/features/task-detail/hooks/useCloudEventSummary.ts index 8cfcdc8cdcc9..15afce6192de 100644 --- a/products/desktop/packages/ui/src/features/task-detail/hooks/useCloudEventSummary.ts +++ b/products/desktop/packages/ui/src/features/task-detail/hooks/useCloudEventSummary.ts @@ -2,11 +2,12 @@ import { type CloudEventSummary, createCloudEventSummaryTracker, } from "@posthog/core/task-detail/cloudToolChanges"; -import { useMemo, useRef } from "react"; +import { useMemo, useState } from "react"; import { useSessionSelector } from "../../sessions/useSession"; const EMPTY_SUMMARY: CloudEventSummary = { toolCalls: new Map(), + changedFiles: [], revision: 0, changedFilesRevision: 0, }; @@ -19,11 +20,7 @@ export function useCloudEventSummary( enabled ? taskId : undefined, (session) => session?.events, ); - const trackerRef = useRef | null>(null); - trackerRef.current ??= createCloudEventSummaryTracker(); - const tracker = trackerRef.current; + const [tracker] = useState(createCloudEventSummaryTracker); return useMemo( () => (events ? tracker.update(events) : EMPTY_SUMMARY), [events, tracker], diff --git a/products/desktop/packages/ui/src/features/task-detail/hooks/useCloudRunState.ts b/products/desktop/packages/ui/src/features/task-detail/hooks/useCloudRunState.ts index 6cbf209cf37c..e25197bead2d 100644 --- a/products/desktop/packages/ui/src/features/task-detail/hooks/useCloudRunState.ts +++ b/products/desktop/packages/ui/src/features/task-detail/hooks/useCloudRunState.ts @@ -1,7 +1,6 @@ import { deriveCloudRunState } from "@posthog/core/task-detail/cloudRunState"; -import { extractCloudToolChangedFiles } from "@posthog/core/task-detail/cloudToolChanges"; import type { Task } from "@posthog/shared/domain-types"; -import { useMemo, useRef } from "react"; +import { useMemo } from "react"; import { shallow } from "zustand/shallow"; import { resolveCloudPrUrl } from "../../git-interaction/cloudPrUrl"; import { useSessionSelector } from "../../sessions/useSession"; @@ -39,25 +38,7 @@ export function useCloudRunState(taskId: string, task: Task) { deriveCloudRunState(freshTask, session, prUrl); const summary = useCloudEventSummary(taskId); - const fallbackFilesRef = useRef< - | { - taskId: string; - revision: number; - files: ReturnType; - } - | undefined - >(undefined); - if ( - fallbackFilesRef.current?.taskId !== taskId || - fallbackFilesRef.current.revision !== summary.changedFilesRevision - ) { - fallbackFilesRef.current = { - taskId, - revision: summary.changedFilesRevision, - files: extractCloudToolChangedFiles(summary.toolCalls), - }; - } - const fallbackFiles = fallbackFilesRef.current.files; + const fallbackFiles = summary.changedFiles; return { freshTask, diff --git a/products/desktop/packages/workspace-server/src/services/process-tracking/process-utils.test.ts b/products/desktop/packages/workspace-server/src/services/process-tracking/process-utils.test.ts index 16ff9e607004..f4a4cb3262a8 100644 --- a/products/desktop/packages/workspace-server/src/services/process-tracking/process-utils.test.ts +++ b/products/desktop/packages/workspace-server/src/services/process-tracking/process-utils.test.ts @@ -60,6 +60,17 @@ describe("findMatchingProcessTargets", () => { -11, 11, ]); }); + + it("does not target the desktop process group", () => { + const original = [ + { pid: 10, ppid: 1, pgid: 1, startedAt }, + { pid: 11, ppid: 10, pgid: 11, startedAt }, + ]; + + expect(findMatchingProcessTargets(original, original, 1)).toEqual([ + -11, 10, 11, + ]); + }); }); describe("killUnixProcessTrees", () => { @@ -105,6 +116,26 @@ describe("killUnixProcessTrees", () => { expect(schedule).not.toHaveBeenCalled(); }); + it("escalates for identifiable members of an orphaned group", () => { + const orphan = { pid: 11, ppid: 1, pgid: 10, startedAt }; + const signal = vi.fn(); + let delayed: (() => void) | undefined; + + killUnixProcessTrees([10], [orphan], 1, { + currentProcesses: () => [orphan], + signal, + schedule: (callback) => { + delayed = callback; + }, + }); + delayed?.(); + + expect(signal.mock.calls).toEqual([ + [[-10, 10, 11], "SIGTERM"], + [[-10, 11], "SIGKILL"], + ]); + }); + it("falls back for missing roots in a mixed batch", () => { const signal = vi.fn(); diff --git a/products/desktop/packages/workspace-server/src/services/process-tracking/process-utils.ts b/products/desktop/packages/workspace-server/src/services/process-tracking/process-utils.ts index 4318b964c081..4e6c3ab7d647 100644 --- a/products/desktop/packages/workspace-server/src/services/process-tracking/process-utils.ts +++ b/products/desktop/packages/workspace-server/src/services/process-tracking/process-utils.ts @@ -120,8 +120,16 @@ export function killUnixProcessTrees( if (tree.length === 0) missingRootPids.push(pid); return tree; }); + const orphanedGroupMembers = missingRootPids.flatMap((pid) => + initialProcesses.filter((entry) => entry.pgid === pid), + ); const originalTree = Array.from( - new Map(trees.flat().map((entry) => [entry.pid, entry])).values(), + new Map( + [...trees.flat(), ...orphanedGroupMembers].map((entry) => [ + entry.pid, + entry, + ]), + ).values(), ); if (originalTree.length === 0 || ownPgid === undefined) { deps.signal( @@ -131,10 +139,12 @@ export function killUnixProcessTrees( return; } - const targets = [ - ...missingRootPids.flatMap((pid) => [-pid, pid]), - ...findMatchingProcessTargets(originalTree, initialProcesses, ownPgid), - ]; + const targets = Array.from( + new Set([ + ...missingRootPids.flatMap((pid) => [-pid, pid]), + ...findMatchingProcessTargets(originalTree, initialProcesses, ownPgid), + ]), + ); deps.signal(targets, "SIGTERM"); deps.schedule(() => { deps.signal( From b5b84b8815dcd9617ca1c0f22fb8e26ec8c3b67b Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Wed, 19 Aug 2026 14:52:50 +0300 Subject: [PATCH 3/8] fix(desktop): preserve incremental transcript updates --- .../src/task-detail/cloudToolChanges.test.ts | 19 ++++ .../core/src/task-detail/cloudToolChanges.ts | 5 +- .../sessions/components/ConversationView.tsx | 5 +- .../components/buildConversationItems.ts | 102 ++++++++++++++---- .../components/chat-thread/ChatThread.tsx | 16 ++- .../chat-thread/ChatThreadGrouping.test.ts | 39 +++++-- .../components/chat-thread/chatRowGrouping.ts | 44 +++++--- .../incrementalConversationItems.test.ts | 82 +++++++++++++- .../incrementalConversationItems.ts | 13 +-- .../incrementalThreadGrouping.test.ts | 4 +- .../new-thread/incrementalThreadGrouping.ts | 25 +---- 11 files changed, 271 insertions(+), 83 deletions(-) diff --git a/products/desktop/packages/core/src/task-detail/cloudToolChanges.test.ts b/products/desktop/packages/core/src/task-detail/cloudToolChanges.test.ts index df88df857d2d..c7e69a34874b 100644 --- a/products/desktop/packages/core/src/task-detail/cloudToolChanges.test.ts +++ b/products/desktop/packages/core/src/task-detail/cloudToolChanges.test.ts @@ -134,6 +134,25 @@ describe("createCloudEventSummaryTracker", () => { expect(first.toolCalls.get("tool-1")?.status).toBe("in_progress"); }); + + it("refreshes changed files when a replacement has the same revision", () => { + const tracker = createCloudEventSummaryTracker(); + tracker.update([ + toolEvent("tool-1", { + kind: "write", + locations: [{ path: "src/old.ts", line: null }], + }), + ]); + + const replacement = tracker.update([ + toolEvent("tool-2", { + kind: "write", + locations: [{ path: "src/new.ts", line: null }], + }), + ]); + + expect(replacement.changedFiles).toMatchObject([{ path: "src/new.ts" }]); + }); }); function diffObj( diff --git a/products/desktop/packages/core/src/task-detail/cloudToolChanges.ts b/products/desktop/packages/core/src/task-detail/cloudToolChanges.ts index bfefc2aec159..a8e8914ea017 100644 --- a/products/desktop/packages/core/src/task-detail/cloudToolChanges.ts +++ b/products/desktop/packages/core/src/task-detail/cloudToolChanges.ts @@ -281,7 +281,10 @@ export function createCloudEventSummaryTracker(): { if (change.changedFilesChanged) state.changedFilesRevision++; }, getResult: (state) => { - if (state.changedFilesRevision !== projectedChangedFilesRevision) { + if ( + state !== projectedState || + state.changedFilesRevision !== projectedChangedFilesRevision + ) { projectedChangedFilesRevision = state.changedFilesRevision; projectedChangedFiles = extractCloudToolChangedFiles(state.toolCalls); } diff --git a/products/desktop/packages/ui/src/features/sessions/components/ConversationView.tsx b/products/desktop/packages/ui/src/features/sessions/components/ConversationView.tsx index 6b1ded4cf8f1..2b41f7aa6564 100644 --- a/products/desktop/packages/ui/src/features/sessions/components/ConversationView.tsx +++ b/products/desktop/packages/ui/src/features/sessions/components/ConversationView.tsx @@ -146,6 +146,7 @@ export function ConversationView({ // terminal and status events still flush immediately. const { items: conversationItems, + stablePrefixItemCount, lastTurnInfo, isCompacting, isClearing, @@ -200,8 +201,8 @@ export function ConversationView({ threadGrouperRef.current ??= createIncrementalThreadGrouper(); const threadGrouper = threadGrouperRef.current; const grouping = useMemo( - () => threadGrouper.update(items, groupOverrides), - [items, groupOverrides, threadGrouper], + () => threadGrouper.update(items, groupOverrides, stablePrefixItemCount), + [items, groupOverrides, stablePrefixItemCount, threadGrouper], ); const threadRows = grouping.rows; const rowKeepMounted = grouping.keepMounted; diff --git a/products/desktop/packages/ui/src/features/sessions/components/buildConversationItems.ts b/products/desktop/packages/ui/src/features/sessions/components/buildConversationItems.ts index 9d388e93fda2..ab5358364e94 100644 --- a/products/desktop/packages/ui/src/features/sessions/components/buildConversationItems.ts +++ b/products/desktop/packages/ui/src/features/sessions/components/buildConversationItems.ts @@ -77,6 +77,7 @@ export interface LastTurnInfo { export interface BuildResult { items: ConversationItem[]; + stablePrefixItemCount: number; lastTurnInfo: LastTurnInfo | null; isCompacting: boolean; /** A `/clear` is in flight (its status row shows the dedicated spinner), so @@ -124,7 +125,15 @@ interface TurnState { export interface ItemBuilder { items: ConversationItem[]; - toolCallRows: Map; + toolCallRows: Map< + string, + { + items: ConversationItem[]; + index: number; + ancestors: { items: ConversationItem[]; index: number }[]; + rootIndex: number; + } + >; currentTurn: TurnState | null; /** Index in `items` where the current turn's first item sits. Lets an * incremental consumer treat everything before it (completed turns) as @@ -139,11 +148,7 @@ export interface ItemBuilder { * event for the same id mutates the same card, regardless of which turn is * currently active. */ progressCards: Map; - /** Lowest item index touched by a progress event since it was last reset. - * An incremental consumer resets this before feeding a batch of events and - * reads it after to detect a card being mutated inside an already frozen - * (completed) turn, which would otherwise go unseen. */ - lowestTouchedProgressIndex: number; + lowestTouchedItemIndex: number; /** Count of tool calls that have reached a terminal status (completed / * failed / cancelled). Increments once per tool call when it first settles. * Drives the generating indicator's status word so it advances on real work @@ -168,7 +173,7 @@ export function createItemBuilder(): ItemBuilder { isCompacting: false, isClearing: false, progressCards: new Map(), - lowestTouchedProgressIndex: Number.POSITIVE_INFINITY, + lowestTouchedItemIndex: Number.POSITIVE_INFINITY, completedToolCallCount: 0, lastActivityAt: null, runStartedRunIds: new Set(), @@ -245,7 +250,12 @@ function pushItem(b: ItemBuilder, update: RenderItem, ts?: number) { }; const index = b.items.push(item) - 1; if (update.sessionUpdate === "tool_call") { - b.toolCallRows.set(update.toolCallId, { items: b.items, index }); + b.toolCallRows.set(update.toolCallId, { + items: b.items, + index, + ancestors: [], + rootIndex: index, + }); } } @@ -296,6 +306,7 @@ export function buildConversationItems( return { items: b.items, + stablePrefixItemCount: 0, lastTurnInfo, isCompacting: b.isCompacting, isClearing: b.isClearing, @@ -353,6 +364,7 @@ export function buildAgentConversationItems( return { items: b.items, + stablePrefixItemCount: 0, lastTurnInfo: readLastTurnInfo(b), isCompacting: b.isCompacting, isClearing: b.isClearing, @@ -556,8 +568,8 @@ function handlePromptRequest( } // The shifted cards may live inside a turn the incremental builder already // froze; flag the mutation so it falls back to a full rebuild. - if (insertIndex < b.lowestTouchedProgressIndex) { - b.lowestTouchedProgressIndex = insertIndex; + if (insertIndex < b.lowestTouchedItemIndex) { + b.lowestTouchedItemIndex = insertIndex; } } @@ -659,21 +671,42 @@ function completePromptTurn( function replaceTurnContextRows(b: ItemBuilder, context: TurnContext): void { const visited = new Set(); - const replaceRows = (items: ConversationItem[]): void => { - if (visited.has(items)) return; + const replaceRows = ( + items: ConversationItem[], + rootIndex: number, + ): boolean => { + if (visited.has(items)) return false; visited.add(items); + let replaced = false; for (let index = 0; index < items.length; index++) { const item = items[index]; if (item.type !== "session_update") continue; if (item.turnContext === context) { items[index] = { ...item }; + replaced = true; } for (const children of item.turnContext.childItems.values()) { - replaceRows(children); + replaced = replaceRows(children, rootIndex) || replaced; } } + if (replaced && rootIndex < b.lowestTouchedItemIndex) { + b.lowestTouchedItemIndex = rootIndex; + } + return replaced; }; - replaceRows(b.items); + for (let index = 0; index < b.items.length; index++) { + const item = b.items[index]; + if (item.type !== "session_update") continue; + if (item.turnContext === context) { + b.items[index] = { ...item }; + if (index < b.lowestTouchedItemIndex) { + b.lowestTouchedItemIndex = index; + } + } + for (const children of item.turnContext.childItems.values()) { + replaceRows(children, index); + } + } } function handleNotification( @@ -686,7 +719,12 @@ function handleNotification( const params = msg.params as UserShellExecuteParams; const existing = b.shellExecutes.get(params.id); if (existing) { - existing.item.result = params.result; + const item = { ...existing.item, result: params.result }; + b.items[existing.index] = item; + b.shellExecutes.set(params.id, { item, index: existing.index }); + if (existing.index < b.lowestTouchedItemIndex) { + b.lowestTouchedItemIndex = existing.index; + } } else { const item: UserShellExecute = { type: "user_shell_execute", @@ -750,8 +788,8 @@ function handleNotification( b.runStartedRunIds.add(runId); const card = b.progressCards.get(`setup:${runId}`); if (card) { - if (card.itemIndex < b.lowestTouchedProgressIndex) { - b.lowestTouchedProgressIndex = card.itemIndex; + if (card.itemIndex < b.lowestTouchedItemIndex) { + b.lowestTouchedItemIndex = card.itemIndex; } syncProgressCard(card, b); } @@ -952,8 +990,8 @@ function handleProgress( const status = normalizeStepStatus(params.status); const card = ensureProgressCardForGroup(b, params.group, ts); if (!card) return; - if (card.itemIndex < b.lowestTouchedProgressIndex) { - b.lowestTouchedProgressIndex = card.itemIndex; + if (card.itemIndex < b.lowestTouchedItemIndex) { + b.lowestTouchedItemIndex = card.itemIndex; } card.steps.set(params.step, { key: params.step, @@ -996,6 +1034,9 @@ function markRuntimeStatusComplete(b: ItemBuilder, status: string) { // stuck with its spinner and a still-ticking timer. A new reference forces // the completion to render (and the row to unmount). b.items[i] = { ...item, update: { ...item.update, isComplete: true } }; + if (i < b.lowestTouchedItemIndex) { + b.lowestTouchedItemIndex = i; + } return; } } @@ -1069,7 +1110,18 @@ function pushChildItem(b: ItemBuilder, parentId: string, update: RenderItem) { }; const index = children.push(item) - 1; if (update.sessionUpdate === "tool_call") { - b.toolCallRows.set(update.toolCallId, { items: children, index }); + const parentRow = b.toolCallRows.get(parentId); + b.toolCallRows.set(update.toolCallId, { + items: children, + index, + ancestors: parentRow + ? [ + ...parentRow.ancestors, + { items: parentRow.items, index: parentRow.index }, + ] + : [], + rootIndex: parentRow?.rootIndex ?? b.currentTurnStartIndex, + }); } reissueToolCallRow(b, parentId); } @@ -1148,6 +1200,16 @@ function reissueToolCallRow( ...item, update: { ...update, sessionUpdate: "tool_call" }, }; + for (let index = row.ancestors.length - 1; index >= 0; index--) { + const ancestor = row.ancestors[index]; + const ancestorItem = ancestor.items[ancestor.index]; + if (ancestorItem) { + ancestor.items[ancestor.index] = { ...ancestorItem }; + } + } + if (row.rootIndex < b.lowestTouchedItemIndex) { + b.lowestTouchedItemIndex = row.rootIndex; + } } function processSessionUpdate( diff --git a/products/desktop/packages/ui/src/features/sessions/components/chat-thread/ChatThread.tsx b/products/desktop/packages/ui/src/features/sessions/components/chat-thread/ChatThread.tsx index 085dffd0b9bd..fe7f9672e134 100644 --- a/products/desktop/packages/ui/src/features/sessions/components/chat-thread/ChatThread.tsx +++ b/products/desktop/packages/ui/src/features/sessions/components/chat-thread/ChatThread.tsx @@ -1313,15 +1313,18 @@ export function ChatThread({ events, ...props }: ChatThreadProps) { conversationItems={items} footerEvents={[]} footerState={footerState} + stablePrefixItemCount={0} /> ); } export function AcpChatThread({ events, ...props }: AcpChatThreadProps) { const showDebugLogs = useSettingsStore((state) => state.debugLogsCloudRuns); - const { items } = useConversationItems(events, props.isPromptPending, { - showDebugLogs, - }); + const { items, stablePrefixItemCount } = useConversationItems( + events, + props.isPromptPending, + { showDebugLogs }, + ); return ( ); } @@ -1336,11 +1340,13 @@ export function AcpChatThread({ events, ...props }: AcpChatThreadProps) { interface ChatThreadRendererProps extends SharedChatThreadProps { conversationItems: ConversationItem[]; footerEvents: AcpMessage[]; + stablePrefixItemCount: number; } function ChatThreadRenderer({ conversationItems, footerEvents, + stablePrefixItemCount, groupToolCalls = true, isPromptPending, promptStartedAt, @@ -1380,8 +1386,8 @@ function ChatThreadRenderer({ [groupToolCalls], ); const rows = useMemo( - () => rowGrouper.update(items), - [items, rowGrouper], + () => rowGrouper.update(items, stablePrefixItemCount), + [items, rowGrouper, stablePrefixItemCount], ); // Virtualization ratchet: past the threshold the thread switches to the windowed body and diff --git a/products/desktop/packages/ui/src/features/sessions/components/chat-thread/ChatThreadGrouping.test.ts b/products/desktop/packages/ui/src/features/sessions/components/chat-thread/ChatThreadGrouping.test.ts index c62b9c5355f5..d38536f307f9 100644 --- a/products/desktop/packages/ui/src/features/sessions/components/chat-thread/ChatThreadGrouping.test.ts +++ b/products/desktop/packages/ui/src/features/sessions/components/chat-thread/ChatThreadGrouping.test.ts @@ -56,8 +56,8 @@ describe("createIncrementalChatRowGrouper", () => { const firstItems = [userMessage("u1"), agentMessage("a1")]; const first = grouper.update(firstItems); const secondItems = [...firstItems, userMessage("u2"), agentMessage("a2")]; - const second = grouper.update(secondItems); - const third = grouper.update([...secondItems, agentMessage("a3")]); + const second = grouper.update(secondItems, 2); + const third = grouper.update([...secondItems, agentMessage("a3")], 2); expect(second[0]).toBe(first[0]); expect(second[1]).toBe(first[1]); @@ -102,14 +102,13 @@ describe("createIncrementalChatRowGrouper", () => { it("replaces an optimistic boundary whose confirmed item has a new id", () => { const grouper = createIncrementalChatRowGrouper(groupRows); const prefix = [userMessage("u1"), agentMessage("a1")]; - grouper.update([...prefix, userMessage("optimistic-u2")]); + grouper.update([...prefix, userMessage("optimistic-u2")], 2); expect( - grouper.update([ - ...prefix, - userMessage("confirmed-u2"), - agentMessage("a2"), - ]), + grouper.update( + [...prefix, userMessage("confirmed-u2"), agentMessage("a2")], + 2, + ), ).toMatchObject([ { id: "u1" }, { type: "agent_turn", items: [{ id: "a1" }] }, @@ -117,4 +116,28 @@ describe("createIncrementalChatRowGrouper", () => { { type: "agent_turn", items: [{ id: "a2" }] }, ]); }); + + it("does not inspect the completed prefix on a streamed append", () => { + const grouper = createIncrementalChatRowGrouper(groupRows); + const items = [ + userMessage("u1"), + agentMessage("a1"), + userMessage("u2"), + agentMessage("a2"), + ]; + grouper.update(items, 2); + const inspected = new Set(); + const next = new Proxy([...items, agentMessage("a3")], { + get(target, property, receiver) { + if (typeof property === "string" && /^\d+$/.test(property)) { + inspected.add(Number(property)); + } + return Reflect.get(target, property, receiver); + }, + }); + + grouper.update(next, 2); + + expect(inspected.has(0)).toBe(false); + }); }); diff --git a/products/desktop/packages/ui/src/features/sessions/components/chat-thread/chatRowGrouping.ts b/products/desktop/packages/ui/src/features/sessions/components/chat-thread/chatRowGrouping.ts index 3be915f40d66..3d879ab7dcbf 100644 --- a/products/desktop/packages/ui/src/features/sessions/components/chat-thread/chatRowGrouping.ts +++ b/products/desktop/packages/ui/src/features/sessions/components/chat-thread/chatRowGrouping.ts @@ -7,40 +7,44 @@ type GroupRows = (items: ConversationItem[]) => TurnRow[]; export function createIncrementalChatRowGrouper(groupRows: GroupRows) { let cachedItems: ConversationItem[] = []; let cachedRows: TurnRow[] = []; + const cachedRowIndexes = new Map(); + + const rebuildAll = (items: ConversationItem[]): TurnRow[] => { + cachedItems = items; + cachedRows = groupRows(items); + cachedRowIndexes.clear(); + cachedRows.forEach((row, index) => { + cachedRowIndexes.set(row.id, index); + }); + return cachedRows; + }; return { - update(items: ConversationItem[]): TurnRow[] { + update(items: ConversationItem[], stablePrefixItemCount = 0): TurnRow[] { if (items === cachedItems) return cachedRows; let rebuildStart = 0; - for (let index = items.length - 1; index >= 0; index--) { + for ( + let index = Math.min(stablePrefixItemCount, items.length - 1); + index >= 0; + index-- + ) { if (isUserInitiatedConversationItem(items[index])) { rebuildStart = index; break; } } - for (let index = 0; index < rebuildStart; index++) { - if (cachedItems[index] !== items[index]) { - rebuildStart = 0; - break; - } - } - - let boundaryId = items[rebuildStart]?.id; - let cachedBoundaryIndex = boundaryId - ? cachedRows.findIndex((row) => row.id === boundaryId) + const boundaryId = items[rebuildStart]?.id; + const cachedBoundaryIndex = boundaryId + ? (cachedRowIndexes.get(boundaryId) ?? -1) : -1; if ( rebuildStart > 0 && rebuildStart < cachedItems.length && cachedBoundaryIndex < 0 ) { - rebuildStart = 0; - boundaryId = items[0]?.id; - cachedBoundaryIndex = boundaryId - ? cachedRows.findIndex((row) => row.id === boundaryId) - : -1; + return rebuildAll(items); } const prefixRowCount = rebuildStart === 0 @@ -49,6 +53,12 @@ export function createIncrementalChatRowGrouper(groupRows: GroupRows) { ? cachedBoundaryIndex : cachedRows.length; const suffixRows = groupRows(items.slice(rebuildStart)); + for (let index = prefixRowCount; index < cachedRows.length; index++) { + cachedRowIndexes.delete(cachedRows[index].id); + } + suffixRows.forEach((row, index) => { + cachedRowIndexes.set(row.id, prefixRowCount + index); + }); cachedItems = items; cachedRows = [...cachedRows.slice(0, prefixRowCount), ...suffixRows]; return cachedRows; diff --git a/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.test.ts b/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.test.ts index a8dd087adfd4..52a35040a96d 100644 --- a/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.test.ts +++ b/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.test.ts @@ -85,14 +85,19 @@ function progressMsg( }; } -function shellExecuteMsg(ts: number, id: string, command: string): AcpMessage { +function shellExecuteMsg( + ts: number, + id: string, + command: string, + result?: { stdout: string; stderr: string; exitCode: number }, +): AcpMessage { return { type: "acp_message", ts, message: { jsonrpc: "2.0", method: "_array/user_shell_execute", - params: { id, command, cwd: "/repo" }, + params: { id, command, cwd: "/repo", result }, }, }; } @@ -550,6 +555,79 @@ describe("createIncrementalConversationBuilder", () => { expect(row2.turnContext.childItems.get("agent1")?.length).toBe(1); }); + it("reissues every ancestor when a nested child tool call arrives", () => { + const inc = createIncrementalConversationBuilder(); + const base = [ + userPromptMsg(1, 1, "go"), + toolCallMsg(2, "agent1"), + childToolCallMsg(3, "child1", "agent1"), + ]; + const first = inc.update(base, true); + const rootBefore = first.items.find( + (item) => + item.type === "session_update" && + item.update.sessionUpdate === "tool_call" && + item.update.toolCallId === "agent1", + ); + if (rootBefore?.type !== "session_update") { + throw new Error("expected root tool row"); + } + const childBefore = rootBefore.turnContext.childItems.get("agent1")?.[0]; + + const second = inc.update( + [...base, childToolCallMsg(4, "child2", "child1")], + true, + ); + const rootAfter = second.items.find( + (item) => + item.type === "session_update" && + item.update.sessionUpdate === "tool_call" && + item.update.toolCallId === "agent1", + ); + if (rootAfter?.type !== "session_update") { + throw new Error("expected root tool row"); + } + const childAfter = rootAfter.turnContext.childItems.get("agent1")?.[0]; + if (childAfter?.type !== "session_update") { + throw new Error("expected child tool row"); + } + + expect(rootAfter).not.toBe(rootBefore); + expect(childAfter).not.toBe(childBefore); + expect(childAfter.turnContext.childItems.get("child1")?.length).toBe(1); + }); + + it("replaces an older shell row when its result crosses a turn boundary", () => { + const inc = createIncrementalConversationBuilder(); + const base = [ + shellExecuteMsg(1, "shell-1", "pnpm test"), + userPromptMsg(2, 1, "next"), + agentChunk(3, "working"), + ]; + const first = inc.update(base, true); + const shellBefore = first.items[0]; + + const second = inc.update( + [ + ...base, + shellExecuteMsg(4, "shell-1", "pnpm test", { + stdout: "passed", + stderr: "", + exitCode: 0, + }), + ], + true, + ); + const shellAfter = second.items[0]; + + expect(shellAfter).not.toBe(shellBefore); + expect(shellAfter).toMatchObject({ + type: "user_shell_execute", + result: { stdout: "passed", exitCode: 0 }, + }); + expect(second.stablePrefixItemCount).toBe(0); + }); + it("re-issues a thought row's identity when its turn completes in the same batch a new turn starts", () => { // The settled thought then sits before the next turn's boundary, where // row caches retain by identity — only a fresh object surfaces the flip. diff --git a/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.ts b/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.ts index f189f27219c4..1a9e8800cd28 100644 --- a/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.ts +++ b/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.ts @@ -94,6 +94,7 @@ export function createIncrementalConversationBuilder() { finalizeBuilder(builder, isPromptPending); const result: BuildResult = { items: builder.items, + stablePrefixItemCount: 0, lastTurnInfo: readLastTurnInfo(builder), isCompacting: builder.isCompacting, isClearing: builder.isClearing, @@ -134,7 +135,8 @@ export function createIncrementalConversationBuilder() { } const builder = b as ItemBuilder; - builder.lowestTouchedProgressIndex = Number.POSITIVE_INFINITY; + const hadProcessedEvents = processedCount > 0; + builder.lowestTouchedItemIndex = Number.POSITIVE_INFINITY; // A rebuild re-reads the whole array, so order it first. An append continues // a sequence already in ts-order and takes the new events as they came. const ordered = @@ -156,11 +158,9 @@ export function createIncrementalConversationBuilder() { ? builder.currentTurnStartIndex : builder.items.length; - // A progress card living in the frozen region was mutated by this batch — - // an event reached back across a turn boundary. The append-only view can't - // show that, so rebuild fully this frame (the persistent builder stays - // valid for the next one). - if (builder.lowestTouchedProgressIndex < activeStart) { + // An event reached back across a turn boundary. Rebuild the visible snapshot + // so cached rows observe it; the persistent builder remains valid. + if (hadProcessedEvents && builder.lowestTouchedItemIndex < activeStart) { return buildConversationItems(events, isPromptPending, options); } @@ -181,6 +181,7 @@ export function createIncrementalConversationBuilder() { // renderers, not via row identity. return { items: builder.items.slice(), + stablePrefixItemCount: activeStart, lastTurnInfo: readLastTurnInfoForOutput(builder), isCompacting: builder.isCompacting, isClearing: builder.isClearing, diff --git a/products/desktop/packages/ui/src/features/sessions/components/new-thread/incrementalThreadGrouping.test.ts b/products/desktop/packages/ui/src/features/sessions/components/new-thread/incrementalThreadGrouping.test.ts index 395d77a1a1e1..03b8d2cc1c78 100644 --- a/products/desktop/packages/ui/src/features/sessions/components/new-thread/incrementalThreadGrouping.test.ts +++ b/products/desktop/packages/ui/src/features/sessions/components/new-thread/incrementalThreadGrouping.test.ts @@ -97,10 +97,10 @@ describe("createIncrementalThreadGrouper", () => { const grouper = createIncrementalThreadGrouper(); const overrides = {}; const items = [userMessage("u1"), toolItem("t1", completeContext)]; - const first = grouper.update(items, overrides); + const first = grouper.update(items, overrides, items.length); const next = [...items, agentMessage("m1")]; - const second = grouper.update(next, overrides); + const second = grouper.update(next, overrides, items.length); expectGroupingEquivalent(second, buildThreadGroups(next, overrides)); expect(second.rows[0]).toBe(first.rows[0]); diff --git a/products/desktop/packages/ui/src/features/sessions/components/new-thread/incrementalThreadGrouping.ts b/products/desktop/packages/ui/src/features/sessions/components/new-thread/incrementalThreadGrouping.ts index c8b6189752f2..05cbde60502c 100644 --- a/products/desktop/packages/ui/src/features/sessions/components/new-thread/incrementalThreadGrouping.ts +++ b/products/desktop/packages/ui/src/features/sessions/components/new-thread/incrementalThreadGrouping.ts @@ -24,13 +24,14 @@ export function createIncrementalThreadGrouper() { const rebuildAll = ( items: ConversationItem[], overrides: Record, + stablePrefixItemCount: number, ): ThreadGrouping => { const grouping = buildThreadGroups(items, overrides); cache = { items, overrides, grouping, - stablePrefixItemCount: findStablePrefixItemCount(items), + stablePrefixItemCount, }; return grouping; }; @@ -38,16 +39,16 @@ export function createIncrementalThreadGrouper() { const update = ( items: ConversationItem[], overrides: Record, + stablePrefixItemCount = 0, ): ThreadGrouping => { if (!cache || cache.overrides !== overrides) { - return rebuildAll(items, overrides); + return rebuildAll(items, overrides, stablePrefixItemCount); } if (cache.items === items) { return cache.grouping; } - const stablePrefixItemCount = findStablePrefixItemCount(items); const rebuildStart = groupBoundaryAtOrBefore( items, Math.min(cache.stablePrefixItemCount, stablePrefixItemCount), @@ -61,7 +62,7 @@ export function createIncrementalThreadGrouper() { rebuildStart > 0 && cache.items[rebuildStart - 1] !== items[rebuildStart - 1] ) { - return rebuildAll(items, overrides); + return rebuildAll(items, overrides, stablePrefixItemCount); } const prefixRowCount = getPrefixRowCount( @@ -102,22 +103,6 @@ export function createIncrementalThreadGrouper() { return { update }; } -/** - * Index of the first item belonging to the still-streaming tail: walk back over - * the trailing run of active (not turn-complete) session updates. - */ -function findStablePrefixItemCount(items: ConversationItem[]): number { - let count = items.length; - while (count > 0) { - const item = items[count - 1]; - if (item.type !== "session_update" || item.turnContext.turnComplete) { - break; - } - count--; - } - return count; -} - /** * A foldable group is a run of groupable items broken only by item type, never * by turn completion, so a single group can straddle the completed/active From 845ffda2a3b3de097a5d10cd7737351e73ea9728 Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Wed, 19 Aug 2026 15:06:21 +0300 Subject: [PATCH 4/8] fix(desktop): invalidate stale transcript groups --- .../sessions/components/ConversationView.tsx | 27 +++++-- .../components/chat-thread/ChatThread.tsx | 26 ++++-- .../chat-thread/ChatThreadGrouping.test.ts | 26 +++++- .../incrementalConversationItems.test.ts | 27 +++++++ .../incrementalConversationItems.ts | 5 +- .../components/mergeConversationItems.test.ts | 38 ++++++++- .../components/mergeConversationItems.ts | 80 +++++++++++++++++++ 7 files changed, 212 insertions(+), 17 deletions(-) diff --git a/products/desktop/packages/ui/src/features/sessions/components/ConversationView.tsx b/products/desktop/packages/ui/src/features/sessions/components/ConversationView.tsx index 2b41f7aa6564..f661dc2e346d 100644 --- a/products/desktop/packages/ui/src/features/sessions/components/ConversationView.tsx +++ b/products/desktop/packages/ui/src/features/sessions/components/ConversationView.tsx @@ -26,7 +26,7 @@ import { type ConversationTurn, groupRowsIntoTurns, } from "@posthog/ui/features/sessions/components/groupConversationTurns"; -import { mergeConversationItems } from "@posthog/ui/features/sessions/components/mergeConversationItems"; +import { createIncrementalConversationMerger } from "@posthog/ui/features/sessions/components/mergeConversationItems"; import type { ThreadGrouping, ThreadRow, @@ -182,15 +182,29 @@ export function ConversationView({ const isCloud = session?.isCloud ?? false; - const items = useMemo( + const mergerRef = useRef | null>(null); + mergerRef.current ??= createIncrementalConversationMerger(); + const merger = mergerRef.current; + const mergedConversation = useMemo( () => - mergeConversationItems({ + merger.update({ conversationItems, optimisticItems, isCloud, + stablePrefixItemCount, }), - [conversationItems, optimisticItems, isCloud], + [ + conversationItems, + optimisticItems, + isCloud, + merger, + stablePrefixItemCount, + ], ); + const items = mergedConversation.items; + const mergedStablePrefixItemCount = mergedConversation.stablePrefixItemCount; // Fold each completed turn's tool-call work into a collapsible chip, and emit // the keepMounted indices (standalone MCP-app rows, whose iframes must survive @@ -201,8 +215,9 @@ export function ConversationView({ threadGrouperRef.current ??= createIncrementalThreadGrouper(); const threadGrouper = threadGrouperRef.current; const grouping = useMemo( - () => threadGrouper.update(items, groupOverrides, stablePrefixItemCount), - [items, groupOverrides, stablePrefixItemCount, threadGrouper], + () => + threadGrouper.update(items, groupOverrides, mergedStablePrefixItemCount), + [items, groupOverrides, mergedStablePrefixItemCount, threadGrouper], ); const threadRows = grouping.rows; const rowKeepMounted = grouping.keepMounted; diff --git a/products/desktop/packages/ui/src/features/sessions/components/chat-thread/ChatThread.tsx b/products/desktop/packages/ui/src/features/sessions/components/chat-thread/ChatThread.tsx index fe7f9672e134..2982f92e4a03 100644 --- a/products/desktop/packages/ui/src/features/sessions/components/chat-thread/ChatThread.tsx +++ b/products/desktop/packages/ui/src/features/sessions/components/chat-thread/ChatThread.tsx @@ -94,7 +94,7 @@ import { import { GitActionMessage } from "@posthog/ui/features/sessions/components/GitActionMessage"; import { GitActionResult } from "@posthog/ui/features/sessions/components/GitActionResult"; import { isUserInitiatedConversationItem } from "@posthog/ui/features/sessions/components/isUserInitiatedConversationItem"; -import { mergeConversationItems } from "@posthog/ui/features/sessions/components/mergeConversationItems"; +import { createIncrementalConversationMerger } from "@posthog/ui/features/sessions/components/mergeConversationItems"; import { extractCanvasInstructions } from "@posthog/ui/features/sessions/components/session-update/canvasInstructions"; import { extractChannelContext } from "@posthog/ui/features/sessions/components/session-update/channelContext"; import { extractCustomInstructions } from "@posthog/ui/features/sessions/components/session-update/customInstructions"; @@ -1372,11 +1372,25 @@ function ChatThreadRenderer({ const optimisticItems = useOptimisticItemsForTask(taskId); const isCloud = useSessionIsCloud(taskId); - const items = useMemo( + const merger = useMemo(() => createIncrementalConversationMerger(), []); + const mergedConversation = useMemo( () => - mergeConversationItems({ conversationItems, optimisticItems, isCloud }), - [conversationItems, optimisticItems, isCloud], + merger.update({ + conversationItems, + optimisticItems, + isCloud, + stablePrefixItemCount, + }), + [ + conversationItems, + optimisticItems, + isCloud, + merger, + stablePrefixItemCount, + ], ); + const items = mergedConversation.items; + const mergedStablePrefixItemCount = mergedConversation.stablePrefixItemCount; const rowGrouper = useMemo( () => @@ -1386,8 +1400,8 @@ function ChatThreadRenderer({ [groupToolCalls], ); const rows = useMemo( - () => rowGrouper.update(items, stablePrefixItemCount), - [items, rowGrouper, stablePrefixItemCount], + () => rowGrouper.update(items, mergedStablePrefixItemCount), + [items, mergedStablePrefixItemCount, rowGrouper], ); // Virtualization ratchet: past the threshold the thread switches to the windowed body and diff --git a/products/desktop/packages/ui/src/features/sessions/components/chat-thread/ChatThreadGrouping.test.ts b/products/desktop/packages/ui/src/features/sessions/components/chat-thread/ChatThreadGrouping.test.ts index d38536f307f9..ac08db42f297 100644 --- a/products/desktop/packages/ui/src/features/sessions/components/chat-thread/ChatThreadGrouping.test.ts +++ b/products/desktop/packages/ui/src/features/sessions/components/chat-thread/ChatThreadGrouping.test.ts @@ -8,7 +8,7 @@ function userMessage(id: string): ConversationItem { return { type: "user_message", id, content: id, timestamp: 1 }; } -function agentMessage(id: string): ConversationItem { +function agentMessage(id: string, text = id): ConversationItem { return { type: "session_update", id, @@ -20,7 +20,7 @@ function agentMessage(id: string): ConversationItem { }, update: { sessionUpdate: "agent_message_chunk", - content: { type: "text", text: id }, + content: { type: "text", text }, }, }; } @@ -81,6 +81,28 @@ describe("createIncrementalChatRowGrouper", () => { ]); }); + it("rebuilds same-id rows when the stable prefix is reset", () => { + const grouper = createIncrementalChatRowGrouper(groupRows); + grouper.update([ + userMessage("u1"), + agentMessage("a1", "old response"), + userMessage("u2"), + agentMessage("a2"), + ]); + + const replacement = agentMessage("a1", "replacement response"); + const rows = grouper.update( + [userMessage("u1"), replacement, userMessage("u2"), agentMessage("a2")], + 0, + ); + const firstAgentTurn = rows[1]; + + if (firstAgentTurn.type !== "agent_turn") { + throw new Error("expected an agent turn"); + } + expect(firstAgentTurn.items[0]).toBe(replacement); + }); + it("rebuilds when a row inside the retained prefix is replaced in place", () => { // The conversation builder swaps row objects at arbitrary indices (a status // completing, a shell result arriving) — including inside turns the grouper diff --git a/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.test.ts b/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.test.ts index 52a35040a96d..d631c59d241f 100644 --- a/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.test.ts +++ b/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.test.ts @@ -301,6 +301,33 @@ const EQUIVALENCE_CASES = Object.entries(SCENARIOS).flatMap(([name, events]) => ); describe("createIncrementalConversationBuilder", () => { + it("resets the stable prefix when a transcript is replaced with the same prompt ids", () => { + const inc = createIncrementalConversationBuilder(); + inc.update( + [ + userPromptMsg(1, 1, "first"), + agentChunk(2, "old response"), + promptResponseMsg(3, 1), + userPromptMsg(4, 2, "second"), + agentChunk(5, "streaming"), + ], + true, + ); + + const replacement = inc.update( + [ + userPromptMsg(1, 1, "first"), + agentChunk(2, "replacement response"), + promptResponseMsg(3, 1), + userPromptMsg(4, 2, "second"), + agentChunk(5, "streaming"), + ], + true, + ); + + expect(replacement.stablePrefixItemCount).toBe(0); + }); + it.each(EQUIVALENCE_CASES)( "matches buildConversationItems at every prefix — $name (pending=$pending)", ({ events, pending }) => { diff --git a/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.ts b/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.ts index 1a9e8800cd28..c4970eaf2463 100644 --- a/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.ts +++ b/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.ts @@ -127,7 +127,8 @@ export function createIncrementalConversationBuilder() { events[processedCount - 1] === boundaryEventRef) && extendsTimestampOrder(events, processedCount, lastProcessedTs); - if (!canAppend) { + const didRebuild = !canAppend; + if (didRebuild) { b = createItemBuilder(); processedCount = 0; lastProcessedTs = Number.NEGATIVE_INFINITY; @@ -181,7 +182,7 @@ export function createIncrementalConversationBuilder() { // renderers, not via row identity. return { items: builder.items.slice(), - stablePrefixItemCount: activeStart, + stablePrefixItemCount: didRebuild ? 0 : activeStart, lastTurnInfo: readLastTurnInfoForOutput(builder), isCompacting: builder.isCompacting, isClearing: builder.isClearing, diff --git a/products/desktop/packages/ui/src/features/sessions/components/mergeConversationItems.test.ts b/products/desktop/packages/ui/src/features/sessions/components/mergeConversationItems.test.ts index fe98b2d706b3..b75019809501 100644 --- a/products/desktop/packages/ui/src/features/sessions/components/mergeConversationItems.test.ts +++ b/products/desktop/packages/ui/src/features/sessions/components/mergeConversationItems.test.ts @@ -1,6 +1,9 @@ import { describe, expect, it } from "vitest"; import type { ConversationItem } from "./buildConversationItems"; -import { mergeConversationItems } from "./mergeConversationItems"; +import { + createIncrementalConversationMerger, + mergeConversationItems, +} from "./mergeConversationItems"; function progressGroup(id: string): ConversationItem { return { @@ -237,4 +240,37 @@ describe("mergeConversationItems", () => { }); expect(result.map((i) => i.id)).toEqual(["old", "setup", "opt"]); }); + + it("invalidates the merged prefix when a streaming echo upgrades the pinned prompt", () => { + const merger = createIncrementalConversationMerger(); + const optimisticItems = [userMessage("opt", "hello")]; + const plainEcho = userMessage("plain", "hello", undefined, 100); + const completedItem = userMessage("other", "different", undefined, 200); + merger.update({ + conversationItems: [plainEcho, completedItem], + optimisticItems, + isCloud: true, + stablePrefixItemCount: 2, + }); + + const shadowEcho = userMessage( + "shadow", + 'hello\n\nbackground', + undefined, + 300, + ); + const upgraded = merger.update({ + conversationItems: [plainEcho, completedItem, shadowEcho], + optimisticItems, + isCloud: true, + stablePrefixItemCount: 2, + }); + + expect(upgraded.stablePrefixItemCount).toBe(0); + expect(upgraded.items[0]).toMatchObject({ + id: "opt", + content: shadowEcho.content, + timestamp: 300, + }); + }); }); diff --git a/products/desktop/packages/ui/src/features/sessions/components/mergeConversationItems.ts b/products/desktop/packages/ui/src/features/sessions/components/mergeConversationItems.ts index 820aa83e572d..ede3ece8b371 100644 --- a/products/desktop/packages/ui/src/features/sessions/components/mergeConversationItems.ts +++ b/products/desktop/packages/ui/src/features/sessions/components/mergeConversationItems.ts @@ -15,6 +15,16 @@ interface MergeConversationItemsArgs { isCloud: boolean; } +interface IncrementalMergeConversationItemsArgs + extends MergeConversationItemsArgs { + stablePrefixItemCount: number; +} + +interface IncrementalMergeResult { + items: ConversationItem[]; + stablePrefixItemCount: number; +} + type UserMessageItem = Extract; // The pinned optimistic bubble is seeded from the bare task description, but the @@ -148,3 +158,73 @@ export function mergeConversationItems({ ...dedupedConversation.slice(tailInsertionIndex), ]; } + +function sameUserMessage( + left: ConversationItem, + right: ConversationItem, +): boolean { + return ( + left.type === "user_message" && + right.type === "user_message" && + left.id === right.id && + left.content === right.content && + left.timestamp === right.timestamp && + left.attachments === right.attachments && + left.pinToTop === right.pinToTop + ); +} + +export function createIncrementalConversationMerger() { + let previousItems: ConversationItem[] = []; + + return { + update({ + conversationItems, + optimisticItems, + isCloud, + stablePrefixItemCount, + }: IncrementalMergeConversationItemsArgs): IncrementalMergeResult { + let items = mergeConversationItems({ + conversationItems, + optimisticItems, + isCloud, + }); + if (stablePrefixItemCount === 0 || previousItems.length === 0) { + previousItems = items; + return { items, stablePrefixItemCount: 0 }; + } + + const stableItems = mergeConversationItems({ + conversationItems: conversationItems.slice(0, stablePrefixItemCount), + optimisticItems, + isCloud, + }); + let mergedStablePrefixItemCount = 0; + while ( + mergedStablePrefixItemCount < items.length && + mergedStablePrefixItemCount < stableItems.length && + mergedStablePrefixItemCount < previousItems.length + ) { + const item = items[mergedStablePrefixItemCount]; + const stableItem = stableItems[mergedStablePrefixItemCount]; + const previousItem = previousItems[mergedStablePrefixItemCount]; + if (item !== stableItem && !sameUserMessage(item, stableItem)) { + break; + } + if (item !== previousItem && !sameUserMessage(item, previousItem)) { + break; + } + if (item !== previousItem) { + if (items === conversationItems) { + items = items.slice(); + } + items[mergedStablePrefixItemCount] = previousItem; + } + mergedStablePrefixItemCount++; + } + + previousItems = items; + return { items, stablePrefixItemCount: mergedStablePrefixItemCount }; + }, + }; +} From 6f184beb6c4283abc10a183c2b9f564d89fa42aa Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Wed, 19 Aug 2026 15:19:48 +0300 Subject: [PATCH 5/8] fix(desktop): isolate transcript snapshots --- .../incrementalConversationItems.test.ts | 17 ++- .../incrementalConversationItems.ts | 113 +++++++++++++++++- 2 files changed, 123 insertions(+), 7 deletions(-) diff --git a/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.test.ts b/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.test.ts index d631c59d241f..16008667b6d4 100644 --- a/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.test.ts +++ b/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.test.ts @@ -551,7 +551,7 @@ describe("createIncrementalConversationBuilder", () => { } expect(row2.update).not.toBe(row1.update); expect((row2.update as { status?: string }).status).toBe("completed"); - // The shared toolCalls Map holds the merged entry the view resolves. + expect(row1.turnContext.toolCalls.get("t1")?.status).toBe("pending"); expect(row2.turnContext.toolCalls.get("t1")?.status).toBe("completed"); }); @@ -579,7 +579,22 @@ describe("createIncrementalConversationBuilder", () => { // New child arrived mid-turn: fresh parent update so the memoized row re-renders. expect(row2).not.toBe(row1); expect(row2.update).not.toBe(row1.update); + expect(row1.turnContext.childItems.get("agent1")).toBeUndefined(); expect(row2.turnContext.childItems.get("agent1")?.length).toBe(1); + + const withSecondChild = [...next, childToolCallMsg(4, "child2", "agent1")]; + const r3 = inc.update(withSecondChild, true); + const row3 = r3.items.find( + (item) => + item.type === "session_update" && + item.update.sessionUpdate === "tool_call" && + item.update.toolCallId === "agent1", + ); + if (row3?.type !== "session_update") { + throw new Error("expected agent session_update row"); + } + expect(row2.turnContext.childItems.get("agent1")?.length).toBe(1); + expect(row3.turnContext.childItems.get("agent1")?.length).toBe(2); }); it("reissues every ancestor when a nested child tool call arrives", () => { diff --git a/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.ts b/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.ts index c4970eaf2463..a6a72986a8a2 100644 --- a/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.ts +++ b/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.ts @@ -3,6 +3,7 @@ import { type BuildConversationOptions, type BuildResult, buildConversationItems, + type ConversationItem, createItemBuilder, finalizeBuilder, type ItemBuilder, @@ -10,6 +11,7 @@ import { orderEventsByTimestamp, processEvent, readLastTurnInfo, + type TurnContext, } from "./buildConversationItems"; /** @@ -53,6 +55,8 @@ export function createIncrementalConversationBuilder() { let firstEventRef: AcpMessage | null = null; let boundaryEventRef: AcpMessage | null = null; let showDebugLogs: boolean | undefined; + let publishedItems = new WeakMap(); + let publishedContexts = new WeakMap(); /** Timestamp of the last event fed to `b`, so a late arrival is detectable. */ let lastProcessedTs = Number.NEGATIVE_INFINITY; @@ -62,6 +66,8 @@ export function createIncrementalConversationBuilder() { firstEventRef = null; boundaryEventRef = null; lastProcessedTs = Number.NEGATIVE_INFINITY; + publishedItems = new WeakMap(); + publishedContexts = new WeakMap(); } function update( @@ -175,13 +181,15 @@ export function createIncrementalConversationBuilder() { markThoughtCompletion(builder.items); - // Rows keep their identity across calls — the builder replaces a row - // object whenever its content changes (tool merges, child streams, - // progress cards), so memoized views re-render exactly the changed rows. - // Turn flags and `thoughtComplete` are surfaced as value props by the - // renderers, not via row identity. + // Published rows retain identity until their content or turn context + // changes. Snapshot contexts keep later builder mutations out of results + // already owned by a committed or abandoned render. return { - items: builder.items.slice(), + items: publishConversationItems( + builder.items, + publishedItems, + publishedContexts, + ), stablePrefixItemCount: didRebuild ? 0 : activeStart, lastTurnInfo: readLastTurnInfoForOutput(builder), isCompacting: builder.isCompacting, @@ -194,6 +202,99 @@ export function createIncrementalConversationBuilder() { return { update, reset }; } +interface PublishedTurnContext { + context: TurnContext; + childSources: Map; +} + +function publishConversationItems( + items: ConversationItem[], + publishedItems: WeakMap, + publishedContexts: WeakMap, +): ConversationItem[] { + const currentContexts = new WeakMap(); + + const publishContext = (context: TurnContext): TurnContext => { + const current = currentContexts.get(context); + if (current) return current; + + const existing = publishedContexts.get(context); + if (existing && isPublishedContextCurrent(existing, context)) { + currentContexts.set(context, existing.context); + return existing.context; + } + + const childSources = new Map(); + for (const [parentId, children] of context.childItems) { + childSources.set(parentId, children.slice()); + } + + const published: TurnContext = { + toolCalls: new Map(context.toolCalls), + childItems: new Map(), + turnCancelled: context.turnCancelled, + turnComplete: context.turnComplete, + }; + currentContexts.set(context, published); + publishedContexts.set(context, { context: published, childSources }); + for (const [parentId, children] of context.childItems) { + published.childItems.set(parentId, publishItems(children)); + } + return published; + }; + + const publishItems = (sourceItems: ConversationItem[]): ConversationItem[] => + sourceItems.map((item) => { + if (item.type !== "session_update") return item; + const existing = publishedItems.get(item); + const publishedContext = publishContext(item.turnContext); + if ( + existing?.type === "session_update" && + existing.thoughtComplete === item.thoughtComplete && + existing.turnContext === publishedContext + ) { + return existing; + } + const published = { + ...item, + turnContext: publishedContext, + }; + publishedItems.set(item, published); + return published; + }); + + return publishItems(items); +} + +function isPublishedContextCurrent( + published: PublishedTurnContext, + source: TurnContext, +): boolean { + const context = published.context; + if ( + context.turnCancelled !== source.turnCancelled || + context.turnComplete !== source.turnComplete || + context.toolCalls.size !== source.toolCalls.size || + published.childSources.size !== source.childItems.size + ) { + return false; + } + for (const [toolCallId, toolCall] of source.toolCalls) { + if (context.toolCalls.get(toolCallId) !== toolCall) return false; + } + for (const [parentId, children] of source.childItems) { + const publishedChildren = published.childSources.get(parentId); + if ( + !publishedChildren || + publishedChildren.length !== children.length || + publishedChildren.some((item, index) => item !== children[index]) + ) { + return false; + } + } + return true; +} + function readLastTurnInfoForOutput(b: ItemBuilder) { const info = readLastTurnInfo(b); if (!info) return null; From 4fe06209e60572ea2274cb05d03135d6fb9b5cb0 Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Wed, 19 Aug 2026 15:32:18 +0300 Subject: [PATCH 6/8] fix(desktop): avoid transcript history rescans --- .../components/buildConversationItems.ts | 58 +++++++++++++++---- .../incrementalConversationItems.ts | 50 ++++------------ 2 files changed, 58 insertions(+), 50 deletions(-) diff --git a/products/desktop/packages/ui/src/features/sessions/components/buildConversationItems.ts b/products/desktop/packages/ui/src/features/sessions/components/buildConversationItems.ts index ab5358364e94..146b2a2ce6a2 100644 --- a/products/desktop/packages/ui/src/features/sessions/components/buildConversationItems.ts +++ b/products/desktop/packages/ui/src/features/sessions/components/buildConversationItems.ts @@ -41,6 +41,16 @@ export interface TurnContext { turnComplete: boolean; } +const turnContextRevisions = new WeakMap(); + +export function readTurnContextRevision(context: TurnContext): number { + return turnContextRevisions.get(context) ?? 0; +} + +export function markTurnContextChanged(context: TurnContext): void { + turnContextRevisions.set(context, readTurnContextRevision(context) + 1); +} + export type ConversationItem = | { type: "user_message"; @@ -194,29 +204,35 @@ function isTerminalToolStatus(status: string | null | undefined): boolean { return status != null && TERMINAL_TOOL_STATUSES.has(status); } -function isThoughtItem( - item: ConversationItem, -): item is ConversationItem & { type: "session_update" } { +type ThoughtItem = Extract & { + update: Extract; +}; + +function isThoughtItem(item: ConversationItem): item is ThoughtItem { return ( item.type === "session_update" && item.update.sessionUpdate === "agent_thought_chunk" ); } -export function markThoughtCompletion(items: ConversationItem[]) { - markThoughtCompletionInItems(items, new Set()); +export function markThoughtCompletion( + items: ConversationItem[], + startIndex = 0, +) { + markThoughtCompletionInItems(items, new Set(), startIndex); } function markThoughtCompletionInItems( items: ConversationItem[], visited: Set, + startIndex = 0, ) { if (visited.has(items)) return; visited.add(items); const seenContexts = new Set(); const itemContexts = new Set(); - for (let i = items.length - 1; i >= 0; i--) { + for (let i = items.length - 1; i >= startIndex; i--) { const item = items[i]; if (isThoughtItem(item)) { @@ -492,14 +508,20 @@ export function finalizeBuilder( for (const turn of b.pendingPrompts.values()) { turn.isComplete = true; turn.durationMs = 0; - turn.context.turnComplete = true; + if (!turn.context.turnComplete) { + turn.context.turnComplete = true; + markTurnContextChanged(turn.context); + } } } // Mark implicit turn complete if it's still the current turn after all events if (b.currentTurn?.promptId === -1) { b.currentTurn.isComplete = true; - b.currentTurn.context.turnComplete = true; + if (!b.currentTurn.context.turnComplete) { + b.currentTurn.context.turnComplete = true; + markTurnContextChanged(b.currentTurn.context); + } } markThoughtCompletion(b.items); @@ -523,7 +545,11 @@ function handlePromptRequest( // If the current turn is the implicit one, mark it complete before starting a real turn if (b.currentTurn && b.currentTurn.promptId === -1) { b.currentTurn.isComplete = true; - b.currentTurn.context.turnComplete = true; + if (!b.currentTurn.context.turnComplete) { + b.currentTurn.context.turnComplete = true; + markTurnContextChanged(b.currentTurn.context); + replaceTurnContextRows(b, b.currentTurn.context); + } } const userPrompt = extractUserPrompt(msg.params); @@ -645,6 +671,7 @@ function completePromptTurn( const wasCancelled = turn.stopReason === "cancelled"; turn.context.turnCancelled = wasCancelled; + markTurnContextChanged(turn.context); replaceTurnContextRows(b, turn.context); if (turn.gitAction.isGitAction && turn.gitAction.actionType) { @@ -682,7 +709,9 @@ function replaceTurnContextRows(b: ItemBuilder, context: TurnContext): void { const item = items[index]; if (item.type !== "session_update") continue; if (item.turnContext === context) { - items[index] = { ...item }; + items[index] = isThoughtItem(item) + ? { ...item, thoughtComplete: true } + : { ...item }; replaced = true; } for (const children of item.turnContext.childItems.values()) { @@ -698,7 +727,9 @@ function replaceTurnContextRows(b: ItemBuilder, context: TurnContext): void { const item = b.items[index]; if (item.type !== "session_update") continue; if (item.turnContext === context) { - b.items[index] = { ...item }; + b.items[index] = isThoughtItem(item) + ? { ...item, thoughtComplete: true } + : { ...item }; if (index < b.lowestTouchedItemIndex) { b.lowestTouchedItemIndex = index; } @@ -1123,6 +1154,7 @@ function pushChildItem(b: ItemBuilder, parentId: string, update: RenderItem) { rootIndex: parentRow?.rootIndex ?? b.currentTurnStartIndex, }); } + markTurnContextChanged(turn.context); reissueToolCallRow(b, parentId); } @@ -1162,6 +1194,7 @@ function appendTextChunkToChildren( }, }, }; + markTurnContextChanged(turn.context); reissueToolCallRow(b, parentId); } else { turn.itemCount++; @@ -1171,6 +1204,7 @@ function appendTextChunkToChildren( update: { ...update, content: { ...update.content } }, turnContext: turn.context, }); + markTurnContextChanged(turn.context); reissueToolCallRow(b, parentId); } } @@ -1186,6 +1220,7 @@ function reissueToolCallRow( const update = nextUpdate ?? (current ? { ...current } : undefined); if (!update) return; turn.toolCalls.set(toolCallId, update); + markTurnContextChanged(turn.context); const row = b.toolCallRows.get(toolCallId); if (!row) return; @@ -1249,6 +1284,7 @@ function processSessionUpdate( } else { const toolCall = { ...update }; turn.toolCalls.set(update.toolCallId, toolCall); + markTurnContextChanged(turn.context); if (isTerminalToolStatus(toolCall.status)) { b.completedToolCallCount++; } diff --git a/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.ts b/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.ts index a6a72986a8a2..1d4de8f19b26 100644 --- a/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.ts +++ b/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.ts @@ -8,9 +8,11 @@ import { finalizeBuilder, type ItemBuilder, markThoughtCompletion, + markTurnContextChanged, orderEventsByTimestamp, processEvent, readLastTurnInfo, + readTurnContextRevision, type TurnContext, } from "./buildConversationItems"; @@ -176,10 +178,13 @@ export function createIncrementalConversationBuilder() { // it's safe to persist (a later real completion still flows through // `completePromptTurn`, which gates on `isComplete`, left untouched here). if (turn && turn.promptId === -1) { - turn.context.turnComplete = true; + if (!turn.context.turnComplete) { + turn.context.turnComplete = true; + markTurnContextChanged(turn.context); + } } - markThoughtCompletion(builder.items); + markThoughtCompletion(builder.items, activeStart); // Published rows retain identity until their content or turn context // changes. Snapshot contexts keep later builder mutations out of results @@ -204,7 +209,7 @@ export function createIncrementalConversationBuilder() { interface PublishedTurnContext { context: TurnContext; - childSources: Map; + revision: number; } function publishConversationItems( @@ -219,16 +224,12 @@ function publishConversationItems( if (current) return current; const existing = publishedContexts.get(context); - if (existing && isPublishedContextCurrent(existing, context)) { + const revision = readTurnContextRevision(context); + if (existing?.revision === revision) { currentContexts.set(context, existing.context); return existing.context; } - const childSources = new Map(); - for (const [parentId, children] of context.childItems) { - childSources.set(parentId, children.slice()); - } - const published: TurnContext = { toolCalls: new Map(context.toolCalls), childItems: new Map(), @@ -236,7 +237,7 @@ function publishConversationItems( turnComplete: context.turnComplete, }; currentContexts.set(context, published); - publishedContexts.set(context, { context: published, childSources }); + publishedContexts.set(context, { context: published, revision }); for (const [parentId, children] of context.childItems) { published.childItems.set(parentId, publishItems(children)); } @@ -266,35 +267,6 @@ function publishConversationItems( return publishItems(items); } -function isPublishedContextCurrent( - published: PublishedTurnContext, - source: TurnContext, -): boolean { - const context = published.context; - if ( - context.turnCancelled !== source.turnCancelled || - context.turnComplete !== source.turnComplete || - context.toolCalls.size !== source.toolCalls.size || - published.childSources.size !== source.childItems.size - ) { - return false; - } - for (const [toolCallId, toolCall] of source.toolCalls) { - if (context.toolCalls.get(toolCallId) !== toolCall) return false; - } - for (const [parentId, children] of source.childItems) { - const publishedChildren = published.childSources.get(parentId); - if ( - !publishedChildren || - publishedChildren.length !== children.length || - publishedChildren.some((item, index) => item !== children[index]) - ) { - return false; - } - } - return true; -} - function readLastTurnInfoForOutput(b: ItemBuilder) { const info = readLastTurnInfo(b); if (!info) return null; From a6d68801eedd83858898c68bf6d932bd4b3c66f0 Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Wed, 19 Aug 2026 15:43:53 +0300 Subject: [PATCH 7/8] fix(desktop): preserve thought completion on rebuilds --- .../sessions/components/ConversationView.tsx | 32 +++++-------------- .../components/buildConversationItems.ts | 8 ++--- .../incrementalConversationItems.test.ts | 21 ++++++++++++ .../incrementalConversationItems.ts | 17 ++++++---- .../incrementalThreadGrouping.test.ts | 4 +-- .../new-thread/incrementalThreadGrouping.ts | 25 ++++++++++++--- 6 files changed, 65 insertions(+), 42 deletions(-) diff --git a/products/desktop/packages/ui/src/features/sessions/components/ConversationView.tsx b/products/desktop/packages/ui/src/features/sessions/components/ConversationView.tsx index f661dc2e346d..d2f74705ddb9 100644 --- a/products/desktop/packages/ui/src/features/sessions/components/ConversationView.tsx +++ b/products/desktop/packages/ui/src/features/sessions/components/ConversationView.tsx @@ -26,7 +26,7 @@ import { type ConversationTurn, groupRowsIntoTurns, } from "@posthog/ui/features/sessions/components/groupConversationTurns"; -import { createIncrementalConversationMerger } from "@posthog/ui/features/sessions/components/mergeConversationItems"; +import { mergeConversationItems } from "@posthog/ui/features/sessions/components/mergeConversationItems"; import type { ThreadGrouping, ThreadRow, @@ -142,11 +142,10 @@ export function ConversationView({ // Streaming appends one event per token. The parse is incremental — each // event is handled once and completed turns are reused by reference — so per // token the work tracks the active turn, not the whole thread. We feed - // `events` directly; the controller batches streaming-only updates while - // terminal and status events still flush immediately. + // `events` directly (no frame-throttle) so a sent message's optimistic->real + // swap is never delayed past the frame the store commits it. const { items: conversationItems, - stablePrefixItemCount, lastTurnInfo, isCompacting, isClearing, @@ -182,29 +181,15 @@ export function ConversationView({ const isCloud = session?.isCloud ?? false; - const mergerRef = useRef | null>(null); - mergerRef.current ??= createIncrementalConversationMerger(); - const merger = mergerRef.current; - const mergedConversation = useMemo( + const items = useMemo( () => - merger.update({ + mergeConversationItems({ conversationItems, optimisticItems, isCloud, - stablePrefixItemCount, }), - [ - conversationItems, - optimisticItems, - isCloud, - merger, - stablePrefixItemCount, - ], + [conversationItems, optimisticItems, isCloud], ); - const items = mergedConversation.items; - const mergedStablePrefixItemCount = mergedConversation.stablePrefixItemCount; // Fold each completed turn's tool-call work into a collapsible chip, and emit // the keepMounted indices (standalone MCP-app rows, whose iframes must survive @@ -215,9 +200,8 @@ export function ConversationView({ threadGrouperRef.current ??= createIncrementalThreadGrouper(); const threadGrouper = threadGrouperRef.current; const grouping = useMemo( - () => - threadGrouper.update(items, groupOverrides, mergedStablePrefixItemCount), - [items, groupOverrides, mergedStablePrefixItemCount, threadGrouper], + () => threadGrouper.update(items, groupOverrides), + [items, groupOverrides, threadGrouper], ); const threadRows = grouping.rows; const rowKeepMounted = grouping.keepMounted; diff --git a/products/desktop/packages/ui/src/features/sessions/components/buildConversationItems.ts b/products/desktop/packages/ui/src/features/sessions/components/buildConversationItems.ts index 146b2a2ce6a2..24051f10a7eb 100644 --- a/products/desktop/packages/ui/src/features/sessions/components/buildConversationItems.ts +++ b/products/desktop/packages/ui/src/features/sessions/components/buildConversationItems.ts @@ -592,10 +592,10 @@ function handlePromptRequest( for (const card of b.progressCards.values()) { if (card.itemIndex >= insertIndex) card.itemIndex++; } - // The shifted cards may live inside a turn the incremental builder already - // froze; flag the mutation so it falls back to a full rebuild. - if (insertIndex < b.lowestTouchedItemIndex) { - b.lowestTouchedItemIndex = insertIndex; + // The shifted cards can settle a thought anywhere in the preceding turn. + // Flag that whole turn so the incremental builder reclassifies it once. + if (b.currentTurnStartIndex < b.lowestTouchedItemIndex) { + b.lowestTouchedItemIndex = b.currentTurnStartIndex; } } diff --git a/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.test.ts b/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.test.ts index 16008667b6d4..0a18eef9365a 100644 --- a/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.test.ts +++ b/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.test.ts @@ -301,6 +301,27 @@ const EQUIVALENCE_CASES = Object.entries(SCENARIOS).flatMap(([name, events]) => ); describe("createIncrementalConversationBuilder", () => { + it("classifies historical thoughts during hydration and transcript replacement", () => { + const transcript = () => [ + userPromptMsg(1, 1, "first"), + thoughtChunk(2, "thinking"), + progressMsg(3, "setup", "completed", "Ready"), + userPromptMsg(4, 2, "second"), + agentChunk(5, "working"), + ]; + const inc = createIncrementalConversationBuilder(); + + const hydrated = transcript(); + expect(normalize(inc.update(hydrated, true))).toEqual( + normalize(buildConversationItems(hydrated, true)), + ); + + const replacement = transcript(); + expect(normalize(inc.update(replacement, true))).toEqual( + normalize(buildConversationItems(replacement, true)), + ); + }); + it("resets the stable prefix when a transcript is replaced with the same prompt ids", () => { const inc = createIncrementalConversationBuilder(); inc.update( diff --git a/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.ts b/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.ts index 1d4de8f19b26..bc2e573a0052 100644 --- a/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.ts +++ b/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.ts @@ -167,12 +167,6 @@ export function createIncrementalConversationBuilder() { ? builder.currentTurnStartIndex : builder.items.length; - // An event reached back across a turn boundary. Rebuild the visible snapshot - // so cached rows observe it; the persistent builder remains valid. - if (hadProcessedEvents && builder.lowestTouchedItemIndex < activeStart) { - return buildConversationItems(events, isPromptPending, options); - } - // `buildConversationItems` always marks a trailing implicit turn complete. // Replicate that on the live turn's context so thought-completion matches; // it's safe to persist (a later real completion still flows through @@ -184,7 +178,16 @@ export function createIncrementalConversationBuilder() { } } - markThoughtCompletion(builder.items, activeStart); + const thoughtScanStart = didRebuild + ? 0 + : Math.min(activeStart, builder.lowestTouchedItemIndex); + markThoughtCompletion(builder.items, thoughtScanStart); + + // An event reached back across a turn boundary. Rebuild the visible snapshot + // so cached rows observe it; the persistent builder remains valid. + if (hadProcessedEvents && builder.lowestTouchedItemIndex < activeStart) { + return buildConversationItems(events, isPromptPending, options); + } // Published rows retain identity until their content or turn context // changes. Snapshot contexts keep later builder mutations out of results diff --git a/products/desktop/packages/ui/src/features/sessions/components/new-thread/incrementalThreadGrouping.test.ts b/products/desktop/packages/ui/src/features/sessions/components/new-thread/incrementalThreadGrouping.test.ts index 03b8d2cc1c78..395d77a1a1e1 100644 --- a/products/desktop/packages/ui/src/features/sessions/components/new-thread/incrementalThreadGrouping.test.ts +++ b/products/desktop/packages/ui/src/features/sessions/components/new-thread/incrementalThreadGrouping.test.ts @@ -97,10 +97,10 @@ describe("createIncrementalThreadGrouper", () => { const grouper = createIncrementalThreadGrouper(); const overrides = {}; const items = [userMessage("u1"), toolItem("t1", completeContext)]; - const first = grouper.update(items, overrides, items.length); + const first = grouper.update(items, overrides); const next = [...items, agentMessage("m1")]; - const second = grouper.update(next, overrides, items.length); + const second = grouper.update(next, overrides); expectGroupingEquivalent(second, buildThreadGroups(next, overrides)); expect(second.rows[0]).toBe(first.rows[0]); diff --git a/products/desktop/packages/ui/src/features/sessions/components/new-thread/incrementalThreadGrouping.ts b/products/desktop/packages/ui/src/features/sessions/components/new-thread/incrementalThreadGrouping.ts index 05cbde60502c..c8b6189752f2 100644 --- a/products/desktop/packages/ui/src/features/sessions/components/new-thread/incrementalThreadGrouping.ts +++ b/products/desktop/packages/ui/src/features/sessions/components/new-thread/incrementalThreadGrouping.ts @@ -24,14 +24,13 @@ export function createIncrementalThreadGrouper() { const rebuildAll = ( items: ConversationItem[], overrides: Record, - stablePrefixItemCount: number, ): ThreadGrouping => { const grouping = buildThreadGroups(items, overrides); cache = { items, overrides, grouping, - stablePrefixItemCount, + stablePrefixItemCount: findStablePrefixItemCount(items), }; return grouping; }; @@ -39,16 +38,16 @@ export function createIncrementalThreadGrouper() { const update = ( items: ConversationItem[], overrides: Record, - stablePrefixItemCount = 0, ): ThreadGrouping => { if (!cache || cache.overrides !== overrides) { - return rebuildAll(items, overrides, stablePrefixItemCount); + return rebuildAll(items, overrides); } if (cache.items === items) { return cache.grouping; } + const stablePrefixItemCount = findStablePrefixItemCount(items); const rebuildStart = groupBoundaryAtOrBefore( items, Math.min(cache.stablePrefixItemCount, stablePrefixItemCount), @@ -62,7 +61,7 @@ export function createIncrementalThreadGrouper() { rebuildStart > 0 && cache.items[rebuildStart - 1] !== items[rebuildStart - 1] ) { - return rebuildAll(items, overrides, stablePrefixItemCount); + return rebuildAll(items, overrides); } const prefixRowCount = getPrefixRowCount( @@ -103,6 +102,22 @@ export function createIncrementalThreadGrouper() { return { update }; } +/** + * Index of the first item belonging to the still-streaming tail: walk back over + * the trailing run of active (not turn-complete) session updates. + */ +function findStablePrefixItemCount(items: ConversationItem[]): number { + let count = items.length; + while (count > 0) { + const item = items[count - 1]; + if (item.type !== "session_update" || item.turnContext.turnComplete) { + break; + } + count--; + } + return count; +} + /** * A foldable group is a run of groupable items broken only by item type, never * by turn completion, so a single group can straddle the completed/active From 14ad73a693f88047fde07e53944e56e8228c112c Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Wed, 19 Aug 2026 15:47:57 +0300 Subject: [PATCH 8/8] test(desktop): cover displaced transcript progress --- .../incrementalConversationItems.test.ts | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.test.ts b/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.test.ts index 0a18eef9365a..1c226f7ebc5c 100644 --- a/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.test.ts +++ b/products/desktop/packages/ui/src/features/sessions/components/incrementalConversationItems.test.ts @@ -322,6 +322,26 @@ describe("createIncrementalConversationBuilder", () => { ); }); + it("settles a prior thought when the next prompt precedes trailing progress", () => { + const inc = createIncrementalConversationBuilder(); + const initial = [userPromptMsg(1, 1, "first"), thoughtChunk(2, "thinking")]; + inc.update(initial, true); + + const displaced = [ + ...initial, + progressMsg(3, "setup", "completed", "Ready"), + userPromptMsg(4, 2, "second"), + ]; + expect(normalize(inc.update(displaced, true))).toEqual( + normalize(buildConversationItems(displaced, true)), + ); + + const continued = [...displaced, agentChunk(5, "working")]; + expect(normalize(inc.update(continued, true))).toEqual( + normalize(buildConversationItems(continued, true)), + ); + }); + it("resets the stable prefix when a transcript is replaced with the same prompt ids", () => { const inc = createIncrementalConversationBuilder(); inc.update(