From 8354f1f0d53aa84bc7150f8188ec2cc194b6d71b Mon Sep 17 00:00:00 2001 From: Matt Pua Date: Mon, 29 Jun 2026 17:41:49 -0400 Subject: [PATCH 1/2] feat(sessions): refill composer when cancelling before the agent starts Cancelling a turn before the agent produces any output now refills the just-sent message back into the composer so a quick "oops" doesn't lose it. The message stays in conversation history (it's already in the agent's context and the persisted log). Refill is skipped once the agent has started, and when you're not focused on that chat. Attachments and mentions are restored as chips via xmlToContent. Generated-By: PostHog Code Task-Id: ef6829e7-daed-47c3-abdd-59340c9d493c --- .../src/sessions/hasAgentStartedTurn.test.ts | 89 +++++++++++++++++++ packages/core/src/sessions/sessionService.ts | 58 ++++++++++++ .../sessions/hooks/useSessionCallbacks.ts | 35 ++++++-- 3 files changed, 175 insertions(+), 7 deletions(-) create mode 100644 packages/core/src/sessions/hasAgentStartedTurn.test.ts diff --git a/packages/core/src/sessions/hasAgentStartedTurn.test.ts b/packages/core/src/sessions/hasAgentStartedTurn.test.ts new file mode 100644 index 0000000000..8b26ce4810 --- /dev/null +++ b/packages/core/src/sessions/hasAgentStartedTurn.test.ts @@ -0,0 +1,89 @@ +import type { AcpMessage } from "@posthog/shared"; +import { describe, expect, it } from "vitest"; +import { createUserMessageEvent } from "./sessionEvents"; +import { hasAgentStartedTurn } from "./sessionService"; + +function agentUpdate( + sessionUpdate: string, + ts: number, + content?: { type: string }, +): AcpMessage { + return { + type: "acp_message", + ts, + message: { + jsonrpc: "2.0", + method: "session/update", + params: { update: { sessionUpdate, ...(content ? { content } : {}) } }, + }, + } as AcpMessage; +} + +describe("hasAgentStartedTurn", () => { + it("is false while still optimistic (echo not landed)", () => { + // A prior completed turn is present, but the just-sent prompt is still an + // optimistic item — its echo has not landed, so the agent has not started. + const events = [ + createUserMessageEvent("first", 1), + agentUpdate("agent_message_chunk", 2, { type: "text" }), + ]; + expect(hasAgentStartedTurn({ events, hasOptimisticItems: true })).toBe( + false, + ); + }); + + it("is false with no events", () => { + expect(hasAgentStartedTurn({ events: [], hasOptimisticItems: false })).toBe( + false, + ); + }); + + it("is false when only the prompt echo has landed", () => { + expect( + hasAgentStartedTurn({ + events: [createUserMessageEvent("hi", 1)], + hasOptimisticItems: false, + }), + ).toBe(false); + }); + + it.each([ + [ + "agent text chunk", + agentUpdate("agent_message_chunk", 2, { type: "text" }), + ], + ["agent thought chunk", agentUpdate("agent_thought_chunk", 2)], + ["tool call", agentUpdate("tool_call", 2)], + ["tool call update", agentUpdate("tool_call_update", 2)], + ])("is true after %s", (_label, output) => { + expect( + hasAgentStartedTurn({ + events: [createUserMessageEvent("hi", 1), output], + hasOptimisticItems: false, + }), + ).toBe(true); + }); + + it("ignores output from a previous turn (looks only after the latest prompt)", () => { + const events = [ + createUserMessageEvent("first", 1), + agentUpdate("agent_message_chunk", 2, { type: "text" }), + createUserMessageEvent("second", 3), + ]; + expect(hasAgentStartedTurn({ events, hasOptimisticItems: false })).toBe( + false, + ); + }); + + it("is true when the latest turn has output after an earlier turn", () => { + const events = [ + createUserMessageEvent("first", 1), + agentUpdate("agent_message_chunk", 2, { type: "text" }), + createUserMessageEvent("second", 3), + agentUpdate("tool_call", 4), + ]; + expect(hasAgentStartedTurn({ events, hasOptimisticItems: false })).toBe( + true, + ); + }); +}); diff --git a/packages/core/src/sessions/sessionService.ts b/packages/core/src/sessions/sessionService.ts index 30e08e3e86..147262d271 100644 --- a/packages/core/src/sessions/sessionService.ts +++ b/packages/core/src/sessions/sessionService.ts @@ -472,6 +472,42 @@ function classifyTurnEventKind( return "other"; } +/** + * Whether the agent has begun working on the turn currently in flight. + * + * While the just-sent prompt is still optimistic (`hasOptimisticItems`, i.e. the + * `session/prompt` echo has not landed) no output is possible — refill is safe. + * Skipping this guard would let a prior completed turn's output look like the + * current turn's. Once the echo has landed, returns true if any agent + * text/output event follows the latest prompt. Reads only persisted state, so + * it is the same decision for local and cloud. + */ +export function hasAgentStartedTurn({ + events, + hasOptimisticItems, +}: { + events: AcpMessage[]; + hasOptimisticItems: boolean; +}): boolean { + if (hasOptimisticItems) return false; + + let lastPromptIndex = -1; + for (let i = events.length - 1; i >= 0; i--) { + const msg = events[i].message; + if (isJsonRpcRequest(msg) && msg.method === "session/prompt") { + lastPromptIndex = i; + break; + } + } + if (lastPromptIndex === -1) return false; + + for (let i = lastPromptIndex + 1; i < events.length; i++) { + const kind = classifyTurnEventKind(events[i].message); + if (kind === "text" || kind === "output") return true; + } + return false; +} + export class SessionService { private connectingTasks = new Map>(); private reconcilingTasks = new Set(); @@ -2206,6 +2242,28 @@ export class SessionService { } } + /** + * Whether the agent has produced any output for the latest turn. + * + * Used to decide if a just-cancelled message can be safely refilled into the + * composer: if the agent never started, the message is effectively lost from + * the user's point of view, so we put it back. + * + * While the prompt is still optimistic (the `session/prompt` echo has not + * landed) no output is possible, so this returns false — refill is safe. + * Once the echo has landed it looks for any agent text/output event after the + * latest prompt. Reads only persisted store events, so it works for both + * local and cloud (avoids the local-only `liveTurnContent`). + */ + hasAgentStartedCurrentTurn(taskId: string): boolean { + const session = this.d.store.getSessionByTaskId(taskId); + if (!session) return false; + return hasAgentStartedTurn({ + events: session.events, + hasOptimisticItems: session.optimisticItems.length > 0, + }); + } + /** * Cancel the current prompt. */ diff --git a/packages/ui/src/features/sessions/hooks/useSessionCallbacks.ts b/packages/ui/src/features/sessions/hooks/useSessionCallbacks.ts index 53b0c06f93..efe7699496 100644 --- a/packages/ui/src/features/sessions/hooks/useSessionCallbacks.ts +++ b/packages/ui/src/features/sessions/hooks/useSessionCallbacks.ts @@ -1,3 +1,4 @@ +import { xmlToContent } from "@posthog/core/message-editor/content"; import { combineQueuedCloudPrompts, promptToQueuedEditorContent, @@ -32,6 +33,11 @@ import { useCallback, useRef } from "react"; const log = logger.scope("session-callbacks"); +function isViewingTask(taskId: string): boolean { + const view = getAppViewSnapshot(); + return view?.type === "task-detail" && view?.taskId === taskId; +} + interface UseSessionCallbacksOptions { taskId: string; task: Task; @@ -54,6 +60,10 @@ export function useSessionCallbacks({ const sessionRef = useRef(session); sessionRef.current = session; + // Serialized text of the in-flight message, captured on send so a cancel + // before the agent produces output can refill it into the composer. + const lastSentTextRef = useRef(null); + const messagingMode = useMessagingMode(taskId); const handleSendPrompt = useCallback( @@ -93,14 +103,15 @@ export function useSessionCallbacks({ try { markAsViewed(taskId); markActivity(taskId); - await sessionService.sendPrompt(taskId, promptText ?? text, { - steer: messagingMode === "steer", - }); - const view = getAppViewSnapshot(); - const isViewingTask = - view?.type === "task-detail" && view?.taskId === taskId; - if (isViewingTask) { + const steer = messagingMode === "steer"; + // A steer folds into the running turn rather than starting its own, so + // it isn't a candidate for refill-on-cancel. + lastSentTextRef.current = steer ? null : text; + + await sessionService.sendPrompt(taskId, promptText ?? text, { steer }); + + if (isViewingTask(taskId)) { markAsViewed(taskId); } } catch (error) { @@ -123,6 +134,13 @@ export function useSessionCallbacks({ ); const handleCancelPrompt = useCallback(async () => { + // Snapshot the refill inputs at cancel time: whether the agent had started, + // and the just-sent text. Null the stash so a repeated cancel can't refill + // the same message twice. + const agentStarted = sessionService.hasAgentStartedCurrentTurn(taskId); + const lastSentText = lastSentTextRef.current; + lastSentTextRef.current = null; + const queuedMessages = sessionStoreSetters.dequeueMessages(taskId); const result = await sessionService.cancelPrompt(taskId); log.info("Prompt cancelled", { success: result }); @@ -144,6 +162,9 @@ export function useSessionCallbacks({ }; setPendingContent(taskId, pendingContent); + } else if (!agentStarted && lastSentText && isViewingTask(taskId)) { + // xmlToContent restores attachments and mentions as chips, not raw XML. + setPendingContent(taskId, xmlToContent(lastSentText)); } requestFocus(taskId); }, [taskId, setPendingContent, requestFocus, sessionService]); From ec6170145fc557b944ffe85c369d5db672ea9de1 Mon Sep 17 00:00:00 2001 From: Matt Pua Date: Tue, 30 Jun 2026 10:16:20 -0400 Subject: [PATCH 2/2] test(sessions): parametrise hasAgentStartedTurn false and isolation cases Collapse the three standalone false-returning tests into one it.each and the two turn-isolation cases into another, matching the codebase's preference for parametrised tests. Addresses greptile review feedback. Generated-By: PostHog Code Task-Id: fb9502ee-29c5-41cb-8c21-649fea69ca3d --- .../src/sessions/hasAgentStartedTurn.test.ts | 78 +++++++++---------- 1 file changed, 38 insertions(+), 40 deletions(-) diff --git a/packages/core/src/sessions/hasAgentStartedTurn.test.ts b/packages/core/src/sessions/hasAgentStartedTurn.test.ts index 8b26ce4810..57a449679f 100644 --- a/packages/core/src/sessions/hasAgentStartedTurn.test.ts +++ b/packages/core/src/sessions/hasAgentStartedTurn.test.ts @@ -20,31 +20,25 @@ function agentUpdate( } describe("hasAgentStartedTurn", () => { - it("is false while still optimistic (echo not landed)", () => { + it.each<[string, AcpMessage[], boolean]>([ // A prior completed turn is present, but the just-sent prompt is still an // optimistic item — its echo has not landed, so the agent has not started. - const events = [ - createUserMessageEvent("first", 1), - agentUpdate("agent_message_chunk", 2, { type: "text" }), - ]; - expect(hasAgentStartedTurn({ events, hasOptimisticItems: true })).toBe( - false, - ); - }); - - it("is false with no events", () => { - expect(hasAgentStartedTurn({ events: [], hasOptimisticItems: false })).toBe( + [ + "still optimistic (echo not landed)", + [ + createUserMessageEvent("first", 1), + agentUpdate("agent_message_chunk", 2, { type: "text" }), + ], + true, + ], + ["there are no events", [], false], + [ + "only the prompt echo has landed", + [createUserMessageEvent("hi", 1)], false, - ); - }); - - it("is false when only the prompt echo has landed", () => { - expect( - hasAgentStartedTurn({ - events: [createUserMessageEvent("hi", 1)], - hasOptimisticItems: false, - }), - ).toBe(false); + ], + ])("is false when %s", (_label, events, hasOptimisticItems) => { + expect(hasAgentStartedTurn({ events, hasOptimisticItems })).toBe(false); }); it.each([ @@ -64,26 +58,30 @@ describe("hasAgentStartedTurn", () => { ).toBe(true); }); - it("ignores output from a previous turn (looks only after the latest prompt)", () => { - const events = [ - createUserMessageEvent("first", 1), - agentUpdate("agent_message_chunk", 2, { type: "text" }), - createUserMessageEvent("second", 3), - ]; - expect(hasAgentStartedTurn({ events, hasOptimisticItems: false })).toBe( + it.each<[string, AcpMessage[], boolean]>([ + // Only looks after the latest prompt: the prior turn's output is ignored. + [ + "prior turn has output but the latest prompt does not", + [ + createUserMessageEvent("first", 1), + agentUpdate("agent_message_chunk", 2, { type: "text" }), + createUserMessageEvent("second", 3), + ], false, - ); - }); - - it("is true when the latest turn has output after an earlier turn", () => { - const events = [ - createUserMessageEvent("first", 1), - agentUpdate("agent_message_chunk", 2, { type: "text" }), - createUserMessageEvent("second", 3), - agentUpdate("tool_call", 4), - ]; - expect(hasAgentStartedTurn({ events, hasOptimisticItems: false })).toBe( + ], + [ + "the latest turn has output after an earlier turn", + [ + createUserMessageEvent("first", 1), + agentUpdate("agent_message_chunk", 2, { type: "text" }), + createUserMessageEvent("second", 3), + agentUpdate("tool_call", 4), + ], true, + ], + ])("turn isolation: %s", (_label, events, expected) => { + expect(hasAgentStartedTurn({ events, hasOptimisticItems: false })).toBe( + expected, ); }); });