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 3bdc1662582b..6d754723410f 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,87 @@ 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 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 toolUpdate( toolCallId: string, fields: Record, @@ -110,6 +191,8 @@ function createHarness() { }; const notifyPromptComplete = vi.fn(); + const enqueueSpeech = vi.fn(); + const setPersistedConfigOptions = vi.fn(); const deps = { store, log: noopLog, @@ -129,10 +212,10 @@ function createHarness() { ); } }, - enqueueSpeech: vi.fn(), + enqueueSpeech, taskViewedApi: { markActivity: vi.fn() }, getPersistedConfigOptions: () => undefined, - setPersistedConfigOptions: vi.fn(), + setPersistedConfigOptions, trpc: { agent: { onSessionEvent: { @@ -165,10 +248,13 @@ function createHarness() { return { service, appendEvents, + enqueueSpeech, notifyPromptComplete, + setPersistedConfigOptions, updateSession: store.updateSession, emit: (event: AcpMessage) => onEvent?.(event), events: () => sessions[RUN_ID].events, + session: () => sessions[RUN_ID], }; } @@ -217,6 +303,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("stores one merged snapshot for buffered updates to the same tool call", () => { @@ -260,6 +347,66 @@ 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 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"); + const configUpdate = configOptionUpdate(); + + h.emit(streamed); + h.emit(configUpdate); + + expect(h.events()).toEqual([streamed, configUpdate]); + 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/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 db6a23cd5201..8e31b9f24136 100644 --- a/products/desktop/packages/core/src/sessions/sessionService.ts +++ b/products/desktop/packages/core/src/sessions/sessionService.ts @@ -178,10 +178,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 @@ -1701,6 +1701,19 @@ 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"; +} + +interface SessionEventHandlingOptions { + appendedInBatch?: boolean; + promptStateUpdatedInBatch?: boolean; +} + export class SessionService { private connectingTasks = new Map>(); private reconcilingTasks = new Set(); @@ -2755,16 +2768,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); @@ -2779,14 +2783,31 @@ 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; this.pendingSessionEvents = new Map(); for (const [taskRunId, events] of batches) { - for (const acpMsg of collapseSupersededToolCallUpdates(events)) { - this.applySessionEvent(taskRunId, acpMsg); - } + this.applySessionEventBatch( + taskRunId, + collapseSupersededToolCallUpdates(events), + ); } } @@ -2796,9 +2817,10 @@ export class SessionService { const events = this.pendingSessionEvents.get(taskRunId); if (!events) return; this.pendingSessionEvents.delete(taskRunId); - for (const acpMsg of collapseSupersededToolCallUpdates(events)) { - this.applySessionEvent(taskRunId, acpMsg); - } + this.applySessionEventBatch( + taskRunId, + collapseSupersededToolCallUpdates(events), + ); } /** @@ -2807,21 +2829,93 @@ 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) { - 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, + ); + } + for (const event of passive) { + this.applySessionEventWithOptions(taskRunId, event, { + appendedInBatch: true, + promptStateUpdatedInBatch: true, + }); + } + } + passive = []; + }; + + for (const event of events) { + if (isStreamUpdateEvent(event)) { + passive.push(event); + } else { + flushPassive(); + this.applySessionEvent(taskRunId, event); } } + flushPassive(); } private statsFor(taskRunId: string): { @@ -2985,7 +3079,14 @@ 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.recordSessionEventArrival(taskRunId, event); + this.flushSessionEventsForTask(taskRunId); + this.applySessionEvent(taskRunId, event); + } }, onError: (err) => { this.d.log.error("Session subscription error", { @@ -3430,7 +3531,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; @@ -3440,21 +3545,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/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..c7e69a34874b 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,160 @@ +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); + }); + + 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"); + }); + + 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( 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..a8e8914ea017 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,59 +179,128 @@ export function cachedDiffStats( export interface CloudEventSummary { toolCalls: Map; + changedFiles: ChangedFile[]; + revision: number; + changedFilesRevision: number; } -/** - * Single-pass extraction of tool calls from events. - */ -export function buildCloudEventSummary( - events: AcpMessage[], -): CloudEventSummary { - 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; - } +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, + }); +} - 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, - }; +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 merged = mergeToolCall(toolCalls.get(toolCallId), patch); - toolCalls.set(toolCallId, merged); - } + 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), + }; +} + +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 projectedChangedFilesRevision = -1; + let projectedChangedFiles: ChangedFile[] = []; + let projectedResult: CloudEventSummary = { + toolCalls: new Map(), + changedFiles: [], + 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.changedFilesRevision !== projectedChangedFilesRevision + ) { + projectedChangedFilesRevision = state.changedFilesRevision; + projectedChangedFiles = extractCloudToolChangedFiles(state.toolCalls); + } + if (state !== projectedState || state.revision !== projectedRevision) { + projectedState = state; + projectedRevision = state.revision; + projectedResult = { + toolCalls: new Map(state.toolCalls), + changedFiles: projectedChangedFiles, + 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..9f5acdfa1698 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?.taskRunId ?? ""}:${session?.status ?? ""}:${session?.isCloud ? 1 : 0}:${session?.cloudStatus ?? ""}:${session?.isPromptPending ? 1 : 0}:${session?.agentIdleForRunId ?? ""}`; + }) + .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/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..24051f10a7eb 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"; @@ -77,6 +87,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,6 +135,15 @@ interface TurnState { export interface ItemBuilder { items: ConversationItem[]; + 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 @@ -138,11 +158,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 @@ -159,6 +175,7 @@ export interface ItemBuilder { export function createItemBuilder(): ItemBuilder { return { items: [], + toolCallRows: new Map(), currentTurn: null, currentTurnStartIndex: 0, pendingPrompts: new Map(), @@ -166,7 +183,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(), @@ -187,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)) { @@ -234,13 +257,22 @@ 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, + ancestors: [], + rootIndex: index, + }); + } } export interface BuildConversationOptions { @@ -290,6 +322,7 @@ export function buildConversationItems( return { items: b.items, + stablePrefixItemCount: 0, lastTurnInfo, isCompacting: b.isCompacting, isClearing: b.isClearing, @@ -347,6 +380,7 @@ export function buildAgentConversationItems( return { items: b.items, + stablePrefixItemCount: 0, lastTurnInfo: readLastTurnInfo(b), isCompacting: b.isCompacting, isClearing: b.isClearing, @@ -474,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); @@ -505,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); @@ -548,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.lowestTouchedProgressIndex) { - b.lowestTouchedProgressIndex = 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; } } @@ -627,6 +671,8 @@ 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) { b.items.push({ @@ -650,6 +696,50 @@ function completePromptTurn( } } +function replaceTurnContextRows(b: ItemBuilder, context: TurnContext): void { + const visited = new Set(); + 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] = isThoughtItem(item) + ? { ...item, thoughtComplete: true } + : { ...item }; + replaced = true; + } + for (const children of item.turnContext.childItems.values()) { + replaced = replaceRows(children, rootIndex) || replaced; + } + } + if (replaced && rootIndex < b.lowestTouchedItemIndex) { + b.lowestTouchedItemIndex = rootIndex; + } + return replaced; + }; + 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] = isThoughtItem(item) + ? { ...item, thoughtComplete: true } + : { ...item }; + if (index < b.lowestTouchedItemIndex) { + b.lowestTouchedItemIndex = index; + } + } + for (const children of item.turnContext.childItems.values()) { + replaceRows(children, index); + } + } +} + function handleNotification( b: ItemBuilder, msg: { method: string; params?: unknown }, @@ -660,7 +750,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", @@ -724,8 +819,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); } @@ -926,8 +1021,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, @@ -970,6 +1065,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; } } @@ -1035,12 +1133,29 @@ 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") { + 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, + }); + } + markTurnContextChanged(turn.context); + reissueToolCallRow(b, parentId); } function appendTextChunkToChildren( @@ -1079,6 +1194,8 @@ function appendTextChunkToChildren( }, }, }; + markTurnContextChanged(turn.context); + reissueToolCallRow(b, parentId); } else { turn.itemCount++; children.push({ @@ -1087,6 +1204,46 @@ function appendTextChunkToChildren( update: { ...update, content: { ...update.content } }, turnContext: turn.context, }); + markTurnContextChanged(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); + markTurnContextChanged(turn.context); + + 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" }, + }; + 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; } } @@ -1119,13 +1276,15 @@ 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 { const toolCall = { ...update }; turn.toolCalls.set(update.toolCallId, toolCall); + markTurnContextChanged(turn.context); if (isTerminalToolStatus(toolCall.status)) { b.completedToolCallCount++; } @@ -1146,8 +1305,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..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 @@ -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"; @@ -93,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"; @@ -1312,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 ( ); } @@ -1335,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, @@ -1365,15 +1372,36 @@ 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( + () => + createIncrementalChatRowGrouper((nextItems) => + groupIntoTurns(groupToolCalls ? groupToolRuns(nextItems) : nextItems), + ), + [groupToolCalls], + ); const rows = useMemo( - () => groupIntoTurns(groupToolCalls ? groupToolRuns(items) : items), - [items, groupToolCalls], + () => 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 new file mode 100644 index 000000000000..ac08db42f297 --- /dev/null +++ b/products/desktop/packages/ui/src/features/sessions/components/chat-thread/ChatThreadGrouping.test.ts @@ -0,0 +1,165 @@ +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, text = id): 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 }, + }, + }; +} + +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, 2); + const third = grouper.update([...secondItems, agentMessage("a3")], 2); + + 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 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 + // 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")], 2); + + expect( + grouper.update( + [...prefix, userMessage("confirmed-u2"), agentMessage("a2")], + 2, + ), + ).toMatchObject([ + { id: "u1" }, + { type: "agent_turn", items: [{ id: "a1" }] }, + { id: "confirmed-u2" }, + { 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 new file mode 100644 index 000000000000..3d879ab7dcbf --- /dev/null +++ b/products/desktop/packages/ui/src/features/sessions/components/chat-thread/chatRowGrouping.ts @@ -0,0 +1,67 @@ +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[] = []; + 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[], stablePrefixItemCount = 0): TurnRow[] { + if (items === cachedItems) return cachedRows; + + let rebuildStart = 0; + for ( + let index = Math.min(stablePrefixItemCount, items.length - 1); + index >= 0; + index-- + ) { + if (isUserInitiatedConversationItem(items[index])) { + rebuildStart = index; + break; + } + } + + const boundaryId = items[rebuildStart]?.id; + const cachedBoundaryIndex = boundaryId + ? (cachedRowIndexes.get(boundaryId) ?? -1) + : -1; + if ( + rebuildStart > 0 && + rebuildStart < cachedItems.length && + cachedBoundaryIndex < 0 + ) { + return rebuildAll(items); + } + const prefixRowCount = + rebuildStart === 0 + ? 0 + : cachedBoundaryIndex >= 0 + ? 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 a93814b7ac8a..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 @@ -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 }, }, }; } @@ -296,6 +301,74 @@ 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("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( + [ + 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 }) => { @@ -502,31 +575,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"); + expect(row1.turnContext.toolCalls.get("t1")?.status).toBe("pending"); + 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,9 +617,153 @@ 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(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", () => { + 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. + 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", () => { 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..bc2e573a0052 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"; @@ -55,6 +57,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; @@ -64,6 +68,8 @@ export function createIncrementalConversationBuilder() { firstEventRef = null; boundaryEventRef = null; lastProcessedTs = Number.NEGATIVE_INFINITY; + publishedItems = new WeakMap(); + publishedContexts = new WeakMap(); } function update( @@ -96,6 +102,7 @@ export function createIncrementalConversationBuilder() { finalizeBuilder(builder, isPromptPending); const result: BuildResult = { items: builder.items, + stablePrefixItemCount: 0, lastTurnInfo: readLastTurnInfo(builder), isCompacting: builder.isCompacting, isClearing: builder.isClearing, @@ -128,7 +135,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; @@ -136,7 +144,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 = @@ -158,26 +167,38 @@ 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) { - 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 // `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); + 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 + // already owned by a committed or abandoned render. return { - items: assembleItems(builder, activeStart), + items: publishConversationItems( + builder.items, + publishedItems, + publishedContexts, + ), + stablePrefixItemCount: didRebuild ? 0 : activeStart, lastTurnInfo: readLastTurnInfoForOutput(builder), isCompacting: builder.isCompacting, isClearing: builder.isClearing, @@ -189,48 +210,64 @@ export function createIncrementalConversationBuilder() { return { update, reset }; } -function assembleItems( - b: ItemBuilder, - activeStart: number, +interface PublishedTurnContext { + context: TurnContext; + revision: number; +} + +function publishConversationItems( + items: ConversationItem[], + publishedItems: WeakMap, + publishedContexts: WeakMap, ): 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); + const currentContexts = new WeakMap(); + + const publishContext = (context: TurnContext): TurnContext => { + const current = currentContexts.get(context); + if (current) return current; + + const existing = publishedContexts.get(context); + const revision = readTurnContextRevision(context); + if (existing?.revision === revision) { + currentContexts.set(context, existing.context); + return existing.context; } - } - return out; + + 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, revision }); + 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 readLastTurnInfoForOutput(b: ItemBuilder) { 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 }; + }, + }; +} 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..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 @@ -1,22 +1,28 @@ import { - buildCloudEventSummary, type CloudEventSummary, + createCloudEventSummaryTracker, } from "@posthog/core/task-detail/cloudToolChanges"; -import { useMemo } from "react"; -import { useSessionForTask } from "../../sessions/useSession"; +import { useMemo, useState } from "react"; +import { useSessionSelector } from "../../sessions/useSession"; const EMPTY_SUMMARY: CloudEventSummary = { toolCalls: new Map(), + changedFiles: [], + 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 [tracker] = useState(createCloudEventSummaryTracker); 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..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,9 +1,9 @@ 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 { 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 +19,26 @@ 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 fallbackFiles = summary.changedFiles; 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..f4a4cb3262a8 --- /dev/null +++ b/products/desktop/packages/workspace-server/src/services/process-tracking/process-utils.test.ts @@ -0,0 +1,155 @@ +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, + ]); + }); + + 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", () => { + 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("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(); + + 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..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 @@ -1,44 +1,207 @@ -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 orphanedGroupMembers = missingRootPids.flatMap((pid) => + initialProcesses.filter((entry) => entry.pgid === pid), + ); + const originalTree = Array.from( + new Map( + [...trees.flat(), ...orphanedGroupMembers].map((entry) => [ + entry.pid, + entry, + ]), + ).values(), + ); + if (originalTree.length === 0 || ownPgid === undefined) { + deps.signal( + rootPids.flatMap((pid) => [-pid, pid]), + "SIGTERM", + ); + return; + } + + const targets = Array.from( + new Set([ + ...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; };