diff --git a/packages/core/src/sessions/hasAgentStartedTurn.test.ts b/packages/core/src/sessions/hasAgentStartedTurn.test.ts new file mode 100644 index 0000000000..57a449679f --- /dev/null +++ b/packages/core/src/sessions/hasAgentStartedTurn.test.ts @@ -0,0 +1,87 @@ +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.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. + [ + "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, + ], + ])("is false when %s", (_label, events, hasOptimisticItems) => { + expect(hasAgentStartedTurn({ events, hasOptimisticItems })).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.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, + ], + [ + "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, + ); + }); +}); diff --git a/packages/core/src/sessions/sessionService.ts b/packages/core/src/sessions/sessionService.ts index 8c29049924..dcc3bbd007 100644 --- a/packages/core/src/sessions/sessionService.ts +++ b/packages/core/src/sessions/sessionService.ts @@ -601,6 +601,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(); @@ -2866,6 +2902,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, + }); + } + /** * Begin an in-place edit of a queued message: mark it as the edit target so * that, until the edit is saved or cancelled, it and everything queued after diff --git a/packages/ui/src/features/sessions/hooks/useSessionCallbacks.ts b/packages/ui/src/features/sessions/hooks/useSessionCallbacks.ts index 188d470ce2..6573bf7ca7 100644 --- a/packages/ui/src/features/sessions/hooks/useSessionCallbacks.ts +++ b/packages/ui/src/features/sessions/hooks/useSessionCallbacks.ts @@ -37,6 +37,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; @@ -59,6 +64,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( @@ -134,14 +143,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) { @@ -183,12 +193,20 @@ export function useSessionCallbacks({ editingId && currentSession?.messageQueue.some((m) => m.id === editingId) ) { + lastSentTextRef.current = null; const result = await sessionService.cancelPrompt(taskId); log.info("Prompt cancelled during queued edit", { success: result }); requestFocus(taskId); return; } + // 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 }); @@ -203,6 +221,9 @@ export function useSessionCallbacks({ : textToContent(typeof queuedPrompt === "string" ? queuedPrompt : ""); 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]);