From 2e5355e9471ac6b6d61e5e53fc2ee2256233e183 Mon Sep 17 00:00:00 2001 From: Georges-Antoine Assi Date: Tue, 7 Jul 2026 13:50:45 -0400 Subject: [PATCH 01/15] feat(sessions): reorder and edit queued messages in place Let users drag queued follow-ups to reorder them (the queue drains in array order, so the order shown is the order they send) and edit a queued message in the composer without pulling it out of the queue. - Add moveQueuedMessage and updateQueuedMessage store setters - Add SessionService.updateQueuedMessage mirroring enqueue normalization (cloud recomputes transport display + raw payload; local stores text) - Track the active edit target in sessionViewStore; the composer's next submit updates the message in place, falling back to a fresh send if it already drained or was discarded - Make each queued card a dnd-kit drag handle; highlight the card being edited with a cancel affordance Generated-By: PostHog Code Task-Id: 96d2978d-909b-4bbf-a45b-44ede597a1b8 --- .../src/sessions/sessionQueueEditing.test.ts | 93 ++++++++++ packages/core/src/sessions/sessionService.ts | 37 ++++ packages/core/src/sessions/sessionStore.ts | 48 ++++++ .../components/QueuedMessagesDock.test.tsx | 18 +- .../components/QueuedMessagesDock.tsx | 161 ++++++++++++++---- .../session-update/QueuedMessageView.tsx | 106 ++++++++---- .../sessions/hooks/useEditQueuedMessage.ts | 78 +++++++++ .../hooks/useReturnQueuedMessageToEditor.ts | 53 ------ .../sessions/hooks/useSessionCallbacks.ts | 27 +++ .../src/features/sessions/sessionViewStore.ts | 31 ++++ 10 files changed, 528 insertions(+), 124 deletions(-) create mode 100644 packages/core/src/sessions/sessionQueueEditing.test.ts create mode 100644 packages/ui/src/features/sessions/hooks/useEditQueuedMessage.ts delete mode 100644 packages/ui/src/features/sessions/hooks/useReturnQueuedMessageToEditor.ts diff --git a/packages/core/src/sessions/sessionQueueEditing.test.ts b/packages/core/src/sessions/sessionQueueEditing.test.ts new file mode 100644 index 0000000000..6afc3c18f1 --- /dev/null +++ b/packages/core/src/sessions/sessionQueueEditing.test.ts @@ -0,0 +1,93 @@ +import type { AgentSession, QueuedMessage } from "@posthog/shared"; +import { afterEach, describe, expect, it } from "vitest"; +import { sessionStore, sessionStoreSetters } from "./sessionStore"; + +const RUN = "run-queue"; +const TASK = "task-queue"; + +function seedQueue(messages: QueuedMessage[]) { + sessionStoreSetters.setSession({ + taskRunId: RUN, + taskId: TASK, + events: [], + messageQueue: messages, + pendingPermissions: new Map(), + status: "connected", + } as unknown as AgentSession); +} + +function queue(): QueuedMessage[] { + return sessionStore.getState().sessions[RUN].messageQueue; +} + +function msg(id: string, content: string): QueuedMessage { + return { id, content, queuedAt: 1 }; +} + +afterEach(() => sessionStoreSetters.removeSession(RUN)); + +describe("moveQueuedMessage", () => { + it("moves a message to a later position, preserving the others' order", () => { + seedQueue([msg("a", "A"), msg("b", "B"), msg("c", "C")]); + + sessionStoreSetters.moveQueuedMessage(TASK, 0, 2); + + expect(queue().map((m) => m.id)).toEqual(["b", "c", "a"]); + }); + + it("moves a message to an earlier position", () => { + seedQueue([msg("a", "A"), msg("b", "B"), msg("c", "C")]); + + sessionStoreSetters.moveQueuedMessage(TASK, 2, 0); + + expect(queue().map((m) => m.id)).toEqual(["c", "a", "b"]); + }); + + it.each([ + ["same index", 1, 1], + ["from out of range", 5, 0], + ["to out of range", 0, 9], + ["negative index", -1, 0], + ])("is a no-op for %s", (_label, from, to) => { + seedQueue([msg("a", "A"), msg("b", "B"), msg("c", "C")]); + + sessionStoreSetters.moveQueuedMessage(TASK, from, to); + + expect(queue().map((m) => m.id)).toEqual(["a", "b", "c"]); + }); +}); + +describe("updateQueuedMessage", () => { + it("replaces content and rawPrompt in place, keeping id and position", () => { + seedQueue([msg("a", "A"), msg("b", "B")]); + + sessionStoreSetters.updateQueuedMessage(TASK, "a", { + content: "edited A", + rawPrompt: "edited A raw", + }); + + expect(queue().map((m) => m.id)).toEqual(["a", "b"]); + expect(queue()[0].content).toBe("edited A"); + expect(queue()[0].rawPrompt).toBe("edited A raw"); + expect(queue()[1].content).toBe("B"); + }); + + it("clears a stale rawPrompt when the patch omits it (local edit)", () => { + seedQueue([{ id: "a", content: "A", rawPrompt: "old raw", queuedAt: 1 }]); + + sessionStoreSetters.updateQueuedMessage(TASK, "a", { content: "edited" }); + + expect(queue()[0].content).toBe("edited"); + expect(queue()[0].rawPrompt).toBeUndefined(); + }); + + it("is a no-op when the target id is not in the queue", () => { + seedQueue([msg("a", "A")]); + + sessionStoreSetters.updateQueuedMessage(TASK, "missing", { + content: "edited", + }); + + expect(queue()[0].content).toBe("A"); + }); +}); diff --git a/packages/core/src/sessions/sessionService.ts b/packages/core/src/sessions/sessionService.ts index 3c4f1b9559..7a9f832eb5 100644 --- a/packages/core/src/sessions/sessionService.ts +++ b/packages/core/src/sessions/sessionService.ts @@ -222,6 +222,11 @@ export interface ISessionStore { rawPrompt?: string | ContentBlock[], ): void; removeQueuedMessage(taskId: string, messageId: string): void; + updateQueuedMessage( + taskId: string, + messageId: string, + patch: { content: string; rawPrompt?: string | ContentBlock[] }, + ): void; clearMessageQueue(taskId: string): void; dequeueMessagesAsText(taskId: string): string | null; dequeueMessages(taskId: string): QueuedMessage[]; @@ -2546,6 +2551,38 @@ export class SessionService { } } + /** + * Update a queued message in place from an edited composer prompt, keeping it + * in the queue at its current position. Mirrors the enqueue normalization so + * the stored `content`/`rawPrompt` match what a freshly-queued prompt would + * hold (cloud recomputes the transport display + raw payload; local stores + * the serialized text). Returns false when the target is no longer queued so + * the caller can fall back to sending it as a new message. + */ + async updateQueuedMessage( + taskId: string, + messageId: string, + prompt: string | ContentBlock[], + ): Promise { + const session = this.d.store.getSessionByTaskId(taskId); + if (!session) return false; + if (!session.messageQueue.some((m) => m.id === messageId)) return false; + + if (session.isCloud) { + const normalizedPrompt = await this.resolveCloudPrompt(prompt); + const transport = this.d.h.getCloudPromptTransport(normalizedPrompt); + this.d.store.updateQueuedMessage(taskId, messageId, { + content: transport.promptText, + rawPrompt: normalizedPrompt, + }); + } else { + this.d.store.updateQueuedMessage(taskId, messageId, { + content: extractPromptText(prompt), + }); + } + return true; + } + /** * Cancel the current prompt. */ diff --git a/packages/core/src/sessions/sessionStore.ts b/packages/core/src/sessions/sessionStore.ts index 0e86a78c6a..80921cbdae 100644 --- a/packages/core/src/sessions/sessionStore.ts +++ b/packages/core/src/sessions/sessionStore.ts @@ -197,6 +197,54 @@ export const sessionStoreSetters = { }); }, + /** + * Move a queued message to a new position, preserving its identity. Used by + * the drag-to-reorder affordance: the queue drains in array order, so the + * order the user sees is the order the messages send. + */ + moveQueuedMessage: (taskId: string, fromIndex: number, toIndex: number) => { + sessionStore.setState((state) => { + const taskRunId = state.taskIdIndex[taskId]; + if (!taskRunId) return; + const session = state.sessions[taskRunId]; + if (!session) return; + const queue = session.messageQueue; + if ( + fromIndex < 0 || + toIndex < 0 || + fromIndex >= queue.length || + toIndex >= queue.length || + fromIndex === toIndex + ) { + return; + } + const [moved] = queue.splice(fromIndex, 1); + if (moved) queue.splice(toIndex, 0, moved); + }); + }, + + /** + * Replace a queued message's content in place, keeping its id, position, and + * queue timestamp. Used by edit-in-place: the message is edited in the + * composer and re-serialized here without leaving the queue. + */ + updateQueuedMessage: ( + taskId: string, + messageId: string, + patch: { content: string; rawPrompt?: string | ContentBlock[] }, + ) => { + sessionStore.setState((state) => { + const taskRunId = state.taskIdIndex[taskId]; + if (!taskRunId) return; + const session = state.sessions[taskRunId]; + if (!session) return; + const message = session.messageQueue.find((m) => m.id === messageId); + if (!message) return; + message.content = patch.content; + message.rawPrompt = patch.rawPrompt; + }); + }, + dequeueMessagesAsText: (taskId: string): string | null => { // Read the queue from the frozen committed state BEFORE entering the // immer draft — same rationale as `dequeueMessages`: anything captured diff --git a/packages/ui/src/features/sessions/components/QueuedMessagesDock.test.tsx b/packages/ui/src/features/sessions/components/QueuedMessagesDock.test.tsx index 0813f010dd..1307c4f529 100644 --- a/packages/ui/src/features/sessions/components/QueuedMessagesDock.test.tsx +++ b/packages/ui/src/features/sessions/components/QueuedMessagesDock.test.tsx @@ -24,20 +24,24 @@ vi.mock("@posthog/ui/features/sessions/hooks/useMessagingMode", () => ({ useSupportsNativeSteer: () => false, })); -vi.mock( - "@posthog/ui/features/sessions/hooks/useReturnQueuedMessageToEditor", - () => ({ - useReturnQueuedMessageToEditor: () => vi.fn(), - }), -); +vi.mock("@posthog/ui/features/sessions/hooks/useEditQueuedMessage", () => ({ + useEditQueuedMessage: () => vi.fn(), + useCancelQueuedMessageEdit: () => vi.fn(), +})); vi.mock("@posthog/ui/features/sessions/sessionStore", () => ({ - sessionStoreSetters: { removeQueuedMessage: vi.fn() }, + sessionStoreSetters: { + removeQueuedMessage: vi.fn(), + moveQueuedMessage: vi.fn(), + }, useSessionIsCloud: () => false, useSessionSelector: ( _taskId: string, select: (session: undefined) => T, ) => select(undefined), + useSessionStore: { + getState: () => ({ taskIdIndex: {}, sessions: {} }), + }, })); vi.mock("@posthog/ui/primitives/toast", () => ({ diff --git a/packages/ui/src/features/sessions/components/QueuedMessagesDock.tsx b/packages/ui/src/features/sessions/components/QueuedMessagesDock.tsx index ac463bf8cf..2d30364036 100644 --- a/packages/ui/src/features/sessions/components/QueuedMessagesDock.tsx +++ b/packages/ui/src/features/sessions/components/QueuedMessagesDock.tsx @@ -1,3 +1,6 @@ +import { PointerSensor } from "@dnd-kit/dom"; +import { type DragDropEvents, DragDropProvider } from "@dnd-kit/react"; +import { useSortable } from "@dnd-kit/react/sortable"; import { CaretDown, CaretRight, Stack } from "@phosphor-icons/react"; import { SESSION_SERVICE, @@ -5,14 +8,19 @@ import { } from "@posthog/core/sessions/sessionService"; import { useService } from "@posthog/di/react"; import { QueuedMessageView } from "@posthog/ui/features/sessions/components/session-update/QueuedMessageView"; +import { + useCancelQueuedMessageEdit, + useEditQueuedMessage, +} from "@posthog/ui/features/sessions/hooks/useEditQueuedMessage"; import { useSupportsNativeSteer } from "@posthog/ui/features/sessions/hooks/useMessagingMode"; -import { useReturnQueuedMessageToEditor } from "@posthog/ui/features/sessions/hooks/useReturnQueuedMessageToEditor"; import { sessionStoreSetters, useSessionIsCloud, useSessionSelector, + useSessionStore, } from "@posthog/ui/features/sessions/sessionStore"; import { + useEditingQueuedId, useQueueCollapsed, useSessionViewActions, } from "@posthog/ui/features/sessions/sessionViewStore"; @@ -20,15 +28,53 @@ import { useQueuedMessagesForTask } from "@posthog/ui/features/sessions/useSessi import { toast } from "@posthog/ui/primitives/toast"; import * as Collapsible from "@radix-ui/react-collapsible"; import { Box, Flex, Text } from "@radix-ui/themes"; +import { type ReactNode, useCallback, useEffect } from "react"; interface QueuedMessagesDockProps { taskId: string; } +/** + * A single queued card, wrapped so the whole card is a drag handle for + * reordering the queue. Drag activation waits for a small pointer move (see the + * provider's sensor) so the card's buttons still take clicks. + */ +function SortableQueuedMessage({ + id, + index, + taskId, + children, +}: { + id: string; + index: number; + taskId: string; + children: ReactNode; +}) { + const { ref, isDragging } = useSortable({ + id, + index, + group: `queue:${taskId}`, + transition: { duration: 200, easing: "ease" }, + }); + + return ( +
+ {children} +
+ ); +} + /** * Queued follow-ups pinned directly above the composer (outside the scrolling - * thread) with per-message actions: steer it into the running turn now, return - * it to the composer to re-read or edit, or discard it. + * thread) with per-message actions: steer it into the running turn now, edit it + * in the composer while it stays queued, or discard it. Cards can be dragged to + * reorder the queue — the order shown is the order they send. * * The list is bounded and scrolls internally so a long queue never pushes the * composer down or off-screen, and a header toggle lets the user collapse it. @@ -37,7 +83,9 @@ export function QueuedMessagesDock({ taskId }: QueuedMessagesDockProps) { const queued = useQueuedMessagesForTask(taskId); const sessionService = useService(SESSION_SERVICE); const supportsNativeSteer = useSupportsNativeSteer(taskId); - const returnToEditor = useReturnQueuedMessageToEditor(taskId); + const editMessage = useEditQueuedMessage(taskId); + const cancelEdit = useCancelQueuedMessageEdit(taskId); + const editingId = useEditingQueuedId(taskId); // Narrow reads (not the whole session) so the dock doesn't re-render on every // streamed token while a turn is running. const isCompacting = useSessionSelector( @@ -50,7 +98,35 @@ export function QueuedMessagesDock({ taskId }: QueuedMessagesDockProps) { // so hide it there too — the message stays queued and lands next turn. const canSteer = !isCompacting && !isCloud; const collapsed = useQueueCollapsed(taskId); - const { setQueueCollapsed } = useSessionViewActions(); + const { setQueueCollapsed, clearEditingQueuedId } = useSessionViewActions(); + + // If the message being edited leaves the queue (drained on turn end, or + // discarded), drop the stale edit target so the composer sends normally. + useEffect(() => { + if (editingId && !queued.some((m) => m.id === editingId)) { + clearEditingQueuedId(taskId); + } + }, [editingId, queued, taskId, clearEditingQueuedId]); + + const handleDragOver: DragDropEvents["dragover"] = useCallback( + (event) => { + const sourceId = event.operation.source?.id; + const targetId = event.operation.target?.id; + if (!sourceId || !targetId || sourceId === targetId) return; + + const state = useSessionStore.getState(); + const taskRunId = state.taskIdIndex[taskId]; + const queue = taskRunId + ? (state.sessions[taskRunId]?.messageQueue ?? []) + : []; + const fromIndex = queue.findIndex((m) => m.id === sourceId); + const toIndex = queue.findIndex((m) => m.id === targetId); + if (fromIndex === -1 || toIndex === -1 || fromIndex === toIndex) return; + + sessionStoreSetters.moveQueuedMessage(taskId, fromIndex, toIndex); + }, + [taskId], + ); if (queued.length === 0) return null; @@ -83,32 +159,55 @@ export function QueuedMessagesDock({ taskId }: QueuedMessagesDockProps) { - - {queued.map((message) => ( - { - void sessionService - .steerQueuedMessage(taskId, message.id) - .catch(() => { - toast.error( - "Couldn't steer this message. It's still queued.", - ); - }); - } - : undefined - } - onReturnToEditor={() => returnToEditor(message)} - onRemove={() => - sessionStoreSetters.removeQueuedMessage(taskId, message.id) - } - /> - ))} - + + + {queued.map((message, index) => ( + + { + void sessionService + .steerQueuedMessage(taskId, message.id) + .catch(() => { + toast.error( + "Couldn't steer this message. It's still queued.", + ); + }); + } + : undefined + } + onEdit={() => editMessage(message)} + onCancelEdit={cancelEdit} + onRemove={() => + sessionStoreSetters.removeQueuedMessage( + taskId, + message.id, + ) + } + /> + + ))} + + diff --git a/packages/ui/src/features/sessions/components/session-update/QueuedMessageView.tsx b/packages/ui/src/features/sessions/components/session-update/QueuedMessageView.tsx index 119245faf2..427598357f 100644 --- a/packages/ui/src/features/sessions/components/session-update/QueuedMessageView.tsx +++ b/packages/ui/src/features/sessions/components/session-update/QueuedMessageView.tsx @@ -1,11 +1,13 @@ import { ArrowBendDownLeft, + DotsSixVertical, PencilSimple, - Stack, Trash, + X, } from "@phosphor-icons/react"; import { Button } from "@posthog/quill"; -import { Box, Flex, IconButton, Tooltip } from "@radix-ui/themes"; +import { Box, Flex, IconButton, Text, Tooltip } from "@radix-ui/themes"; +import clsx from "clsx"; import { MarkdownRenderer } from "../../../editor/components/MarkdownRenderer"; import type { QueuedMessage } from "../../sessionStore"; import { CollapsibleMessageContent } from "./CollapsibleMessageContent"; @@ -14,16 +16,20 @@ import { hasFileMentions, parseFileMentions } from "./parseFileMentions"; interface QueuedMessageViewProps { message: QueuedMessage; onSteer?: () => void; - onReturnToEditor?: () => void; + onEdit?: () => void; + onCancelEdit?: () => void; onRemove?: () => void; + isEditing?: boolean; supportsNativeSteer?: boolean; } export function QueuedMessageView({ message, onSteer, - onReturnToEditor, + onEdit, + onCancelEdit, onRemove, + isEditing = false, supportsNativeSteer = false, }: QueuedMessageViewProps) { const steerTooltip = supportsNativeSteer @@ -31,13 +37,24 @@ export function QueuedMessageView({ : "Interrupt the current turn and resend with this message."; return ( - + - + {hasFileMentions(message.content) ? ( parseFileMentions(message.content) @@ -46,32 +63,55 @@ export function QueuedMessageView({ )} - {onSteer && ( - - - - )} - {onReturnToEditor && ( - - - - - + {isEditing ? ( + <> + + Editing in composer + + {onCancelEdit && ( + + + + + + )} + + ) : ( + <> + {onSteer && ( + + + + )} + {onEdit && ( + + + + + + )} + )} {onRemove && ( diff --git a/packages/ui/src/features/sessions/hooks/useEditQueuedMessage.ts b/packages/ui/src/features/sessions/hooks/useEditQueuedMessage.ts new file mode 100644 index 0000000000..4dd5f51df5 --- /dev/null +++ b/packages/ui/src/features/sessions/hooks/useEditQueuedMessage.ts @@ -0,0 +1,78 @@ +import { + type EditorContent, + xmlToContent, +} from "@posthog/core/message-editor/content"; +import { + combineQueuedCloudPrompts, + promptToQueuedEditorContent, +} from "@posthog/core/sessions/cloudPrompt"; +import { useDraftStore } from "@posthog/ui/features/message-editor/draftStore"; +import { + type QueuedMessage, + useSessionIsCloud, +} from "@posthog/ui/features/sessions/sessionStore"; +import { useSessionViewActions } from "@posthog/ui/features/sessions/sessionViewStore"; +import { useCallback } from "react"; + +/** + * Empty editor content, used to clear the composer when an edit is cancelled. + */ +const EMPTY_CONTENT: EditorContent = { segments: [] }; + +/** + * Load a queued message back into the composer for editing while keeping it in + * the queue at its current position. Marks the message as the active edit + * target so the composer's next submit updates it in place (see + * `useSessionCallbacks.handleSendPrompt`) rather than sending a new prompt. + * + * Content restore mirrors the cancel-to-composer path: cloud keeps its rich + * payload (mentions, attachments) via the queued-cloud conversion; local + * restores the serialized text (chips reparse from the `` tags). + */ +export function useEditQueuedMessage( + taskId: string | undefined, +): (message: QueuedMessage) => void { + const { requestFocus, setPendingContent } = useDraftStore((s) => s.actions); + const { setEditingQueuedId } = useSessionViewActions(); + const isCloud = useSessionIsCloud(taskId); + + return useCallback( + (message: QueuedMessage) => { + if (!taskId) return; + + let pendingContent: EditorContent | null; + if (isCloud) { + const combined = combineQueuedCloudPrompts([message]); + pendingContent = combined + ? promptToQueuedEditorContent(combined) + : null; + } else { + pendingContent = xmlToContent(message.content); + } + if (!pendingContent) return; + + setEditingQueuedId(taskId, message.id); + setPendingContent(taskId, pendingContent); + requestFocus(taskId); + }, + [taskId, isCloud, requestFocus, setPendingContent, setEditingQueuedId], + ); +} + +/** + * Abandon an in-progress queued-message edit: drop the edit target so the next + * submit sends normally again, and clear the composer so the abandoned draft + * doesn't linger. + */ +export function useCancelQueuedMessageEdit( + taskId: string | undefined, +): () => void { + const { setPendingContent } = useDraftStore((s) => s.actions); + const { clearEditingQueuedId } = useSessionViewActions(); + + return useCallback(() => { + if (!taskId) return; + clearEditingQueuedId(taskId); + setPendingContent(taskId, EMPTY_CONTENT); + }, [taskId, clearEditingQueuedId, setPendingContent]); +} diff --git a/packages/ui/src/features/sessions/hooks/useReturnQueuedMessageToEditor.ts b/packages/ui/src/features/sessions/hooks/useReturnQueuedMessageToEditor.ts deleted file mode 100644 index fccef985c1..0000000000 --- a/packages/ui/src/features/sessions/hooks/useReturnQueuedMessageToEditor.ts +++ /dev/null @@ -1,53 +0,0 @@ -import { - type EditorContent, - xmlToContent, -} from "@posthog/core/message-editor/content"; -import { - combineQueuedCloudPrompts, - promptToQueuedEditorContent, -} from "@posthog/core/sessions/cloudPrompt"; -import { useDraftStore } from "@posthog/ui/features/message-editor/draftStore"; -import { - type QueuedMessage, - sessionStoreSetters, - useSessionIsCloud, -} from "@posthog/ui/features/sessions/sessionStore"; -import { useCallback } from "react"; - -/** - * Pull a queued message out of the queue and back into the composer so it can be - * re-read or edited before re-sending. Mirrors the cancel-to-composer restore: - * cloud keeps its rich payload (mentions, attachments) via the queued-cloud - * conversion; local restores the plain text it was queued with. - */ -export function useReturnQueuedMessageToEditor( - taskId: string | undefined, -): (message: QueuedMessage) => void { - const { requestFocus, setPendingContent } = useDraftStore((s) => s.actions); - const isCloud = useSessionIsCloud(taskId); - - return useCallback( - (message: QueuedMessage) => { - if (!taskId) return; - - let pendingContent: EditorContent | null; - if (isCloud) { - const combined = combineQueuedCloudPrompts([message]); - pendingContent = combined - ? promptToQueuedEditorContent(combined) - : null; - } else { - // Local queued content is the serialized form (text + `` - // tags); parse it back into chip segments so attachments restore as - // chips, not raw XML. - pendingContent = xmlToContent(message.content); - } - if (!pendingContent) return; - - sessionStoreSetters.removeQueuedMessage(taskId, message.id); - setPendingContent(taskId, pendingContent); - requestFocus(taskId); - }, - [taskId, isCloud, requestFocus, setPendingContent], - ); -} diff --git a/packages/ui/src/features/sessions/hooks/useSessionCallbacks.ts b/packages/ui/src/features/sessions/hooks/useSessionCallbacks.ts index 53b0c06f93..cb6dbe2c38 100644 --- a/packages/ui/src/features/sessions/hooks/useSessionCallbacks.ts +++ b/packages/ui/src/features/sessions/hooks/useSessionCallbacks.ts @@ -20,6 +20,7 @@ import { type AgentSession, sessionStoreSetters, } from "@posthog/ui/features/sessions/sessionStore"; +import { sessionViewStore } from "@posthog/ui/features/sessions/sessionViewStore"; import { useTaskViewed } from "@posthog/ui/features/sidebar/useTaskViewed"; import { SHELL_CLIENT, @@ -90,6 +91,32 @@ export function useSessionCallbacks({ } } + // Editing a queued message in place: update it where it sits in the + // queue rather than sending a new prompt. If the target already drained + // or was discarded, fall through and send it as a fresh message. + const editingId = + sessionViewStore.getState().editingQueuedIdByTaskId[taskId]; + if (editingId) { + sessionViewStore.getState().actions.clearEditingQueuedId(taskId); + try { + const updated = await sessionService.updateQueuedMessage( + taskId, + editingId, + promptText ?? text, + ); + if (updated) { + markAsViewed(taskId); + return; + } + } catch (error) { + const message = + error instanceof Error ? error.message : "Failed to update message"; + toast.error(message); + log.error("Failed to update queued message", error); + return; + } + } + try { markAsViewed(taskId); markActivity(taskId); diff --git a/packages/ui/src/features/sessions/sessionViewStore.ts b/packages/ui/src/features/sessions/sessionViewStore.ts index ac87639745..6991a3b4ab 100644 --- a/packages/ui/src/features/sessions/sessionViewStore.ts +++ b/packages/ui/src/features/sessions/sessionViewStore.ts @@ -16,6 +16,12 @@ interface SessionViewState { * resets to expanded on app restart. */ queueCollapsedByTaskId: Record; + /** + * Ephemeral per-task id of the queued message currently being edited in the + * composer, keyed by taskId. When set, the composer's next submit updates + * that message in place instead of sending a new prompt. Not persisted. + */ + editingQueuedIdByTaskId: Record; } interface SessionViewActions { @@ -25,6 +31,8 @@ interface SessionViewActions { setGroupOverride: (id: string, expanded: boolean) => void; clearGroupOverrides: () => void; setQueueCollapsed: (taskId: string, collapsed: boolean) => void; + setEditingQueuedId: (taskId: string, messageId: string) => void; + clearEditingQueuedId: (taskId: string) => void; } type SessionViewStore = SessionViewState & { actions: SessionViewActions }; @@ -35,6 +43,7 @@ const useStore = create((set) => ({ showSearch: false, groupOverrides: {}, queueCollapsedByTaskId: {}, + editingQueuedIdByTaskId: {}, actions: { setShowRawLogs: (show) => set({ showRawLogs: show }), setSearchQuery: (query) => set({ searchQuery: query }), @@ -60,6 +69,20 @@ const useStore = create((set) => ({ [taskId]: collapsed, }, })), + setEditingQueuedId: (taskId, messageId) => + set((state) => ({ + editingQueuedIdByTaskId: { + ...state.editingQueuedIdByTaskId, + [taskId]: messageId, + }, + })), + clearEditingQueuedId: (taskId) => + set((state) => { + if (state.editingQueuedIdByTaskId[taskId] === undefined) return state; + const next = { ...state.editingQueuedIdByTaskId }; + delete next[taskId]; + return { editingQueuedIdByTaskId: next }; + }), }, })); @@ -69,4 +92,12 @@ export const useShowSearch = () => useStore((s) => s.showSearch); export const useGroupOverrides = () => useStore((s) => s.groupOverrides); export const useQueueCollapsed = (taskId: string) => useStore((s) => s.queueCollapsedByTaskId[taskId] ?? false); +export const useEditingQueuedId = (taskId: string) => + useStore((s) => s.editingQueuedIdByTaskId[taskId]); export const useSessionViewActions = () => useStore((s) => s.actions); + +/** + * Imperative handle to the view store for non-React call sites (the composer + * submit path reads/clears the editing target here). + */ +export const sessionViewStore = useStore; From 360147c31bbda85ca5cf31183b5ca8475688c225 Mon Sep 17 00:00:00 2001 From: Georges-Antoine Assi Date: Wed, 8 Jul 2026 16:33:37 -0400 Subject: [PATCH 02/15] add git to subpact --- apps/web/vite.config.ts | 1 + 1 file changed, 1 insertion(+) diff --git a/apps/web/vite.config.ts b/apps/web/vite.config.ts index 05f2f11d01..74abc3d03e 100644 --- a/apps/web/vite.config.ts +++ b/apps/web/vite.config.ts @@ -29,6 +29,7 @@ export default defineConfig({ subpath("api-client"), subpath("agent"), subpath("enricher"), + subpath("git"), ], }, server: { port: 5273 }, From 2ddb08008f488141a3a7d7965c10653c7a776e40 Mon Sep 17 00:00:00 2001 From: Georges-Antoine Assi Date: Mon, 13 Jul 2026 08:37:51 -0400 Subject: [PATCH 03/15] feat(sessions): hold queued sends past a message being edited When a queued message is open for an in-place edit, it now acts as a drain boundary: on turn end the messages queued before it still send, but it and everything after stay queued until the edit is saved or cancelled. Saving/cancelling while the agent is already idle re-drains the now-unblocked queue. The "which message is being edited" fact moves from the UI view store into domain state (AgentSession.editingQueuedId) since it now gates the send behavior, with a pure sendableQueuePrefixLength() helper. The auto-drain paths pass stopAtEdited; cancel/recall still drain everything. Generated-By: PostHog Code Task-Id: 5cf31c6c-1ba0-4137-88c1-3a369f3acf36 --- .../src/sessions/sessionQueueEditing.test.ts | 100 ++++++++++++++++- packages/core/src/sessions/sessionService.ts | 95 ++++++++++++++-- packages/core/src/sessions/sessionStore.ts | 79 +++++++++++-- packages/shared/src/index.ts | 1 + packages/shared/src/sessions.ts | 25 +++++ .../components/QueuedMessagesDock.tsx | 14 +-- .../sessions/hooks/useEditQueuedMessage.ts | 18 +-- .../sessions/hooks/useSessionCallbacks.ts | 7 +- .../sessions/sessionServiceHost.test.ts | 106 ++++++++++++++++++ .../src/features/sessions/sessionViewStore.ts | 31 ----- 10 files changed, 404 insertions(+), 72 deletions(-) diff --git a/packages/core/src/sessions/sessionQueueEditing.test.ts b/packages/core/src/sessions/sessionQueueEditing.test.ts index 6afc3c18f1..fb758003a4 100644 --- a/packages/core/src/sessions/sessionQueueEditing.test.ts +++ b/packages/core/src/sessions/sessionQueueEditing.test.ts @@ -1,4 +1,8 @@ -import type { AgentSession, QueuedMessage } from "@posthog/shared"; +import { + type AgentSession, + type QueuedMessage, + sendableQueuePrefixLength, +} from "@posthog/shared"; import { afterEach, describe, expect, it } from "vitest"; import { sessionStore, sessionStoreSetters } from "./sessionStore"; @@ -91,3 +95,97 @@ describe("updateQueuedMessage", () => { expect(queue()[0].content).toBe("A"); }); }); + +describe("sendableQueuePrefixLength", () => { + const q = (ids: string[]): Pick => ({ + messageQueue: ids.map((id) => msg(id, id)), + }); + + it("returns the full length when nothing is being edited", () => { + expect(sendableQueuePrefixLength(q(["a", "b", "c"]))).toBe(3); + }); + + it("stops at the edited message, so only earlier messages count", () => { + expect( + sendableQueuePrefixLength({ ...q(["a", "b", "c"]), editingQueuedId: "b" }), + ).toBe(1); + }); + + it("returns 0 when the head message is being edited", () => { + expect( + sendableQueuePrefixLength({ ...q(["a", "b", "c"]), editingQueuedId: "a" }), + ).toBe(0); + }); + + it("returns the full length when the edited id already left the queue", () => { + expect( + sendableQueuePrefixLength({ + ...q(["a", "b", "c"]), + editingQueuedId: "gone", + }), + ).toBe(3); + }); +}); + +describe("editing hold on the drain", () => { + it("set/clear stores and releases the edit hold", () => { + seedQueue([msg("a", "A"), msg("b", "B")]); + + sessionStoreSetters.setEditingQueuedMessage(TASK, "b"); + expect(sessionStore.getState().sessions[RUN].editingQueuedId).toBe("b"); + + sessionStoreSetters.clearEditingQueuedMessage(TASK); + expect( + sessionStore.getState().sessions[RUN].editingQueuedId, + ).toBeUndefined(); + }); + + it("stopAtEdited drains only the messages before the edited one", () => { + seedQueue([msg("a", "A"), msg("b", "B"), msg("c", "C")]); + sessionStoreSetters.setEditingQueuedMessage(TASK, "b"); + + const combined = sessionStoreSetters.dequeueMessagesAsText(TASK, { + stopAtEdited: true, + }); + + expect(combined).toBe("A"); + // The edited message and everything after it stay queued. + expect(queue().map((m) => m.id)).toEqual(["b", "c"]); + }); + + it("stopAtEdited sends nothing when the head message is being edited", () => { + seedQueue([msg("a", "A"), msg("b", "B")]); + sessionStoreSetters.setEditingQueuedMessage(TASK, "a"); + + expect( + sessionStoreSetters.dequeueMessagesAsText(TASK, { stopAtEdited: true }), + ).toBeNull(); + expect( + sessionStoreSetters.dequeueMessages(TASK, { stopAtEdited: true }), + ).toEqual([]); + expect(queue().map((m) => m.id)).toEqual(["a", "b"]); + }); + + it("dequeueMessages with stopAtEdited returns the sendable prefix as raw items", () => { + seedQueue([msg("a", "A"), msg("b", "B"), msg("c", "C")]); + sessionStoreSetters.setEditingQueuedMessage(TASK, "c"); + + const drained = sessionStoreSetters.dequeueMessages(TASK, { + stopAtEdited: true, + }); + + expect(drained.map((m) => m.id)).toEqual(["a", "b"]); + expect(queue().map((m) => m.id)).toEqual(["c"]); + }); + + it("drains the whole queue when stopAtEdited is not set, even mid-edit", () => { + seedQueue([msg("a", "A"), msg("b", "B"), msg("c", "C")]); + sessionStoreSetters.setEditingQueuedMessage(TASK, "b"); + + // Cancel / recall paths pull everything back regardless of the edit hold. + const combined = sessionStoreSetters.dequeueMessagesAsText(TASK); + + expect(combined).toBe("A\n\nB\n\nC"); + expect(queue()).toEqual([]); + }); +}); diff --git a/packages/core/src/sessions/sessionService.ts b/packages/core/src/sessions/sessionService.ts index 7a9f832eb5..4efe7a18fc 100644 --- a/packages/core/src/sessions/sessionService.ts +++ b/packages/core/src/sessions/sessionService.ts @@ -27,6 +27,7 @@ import { type PermissionRequest, type QueuedMessage, resolveBypassRevertMode, + sendableQueuePrefixLength, type StoredLogEntry, sessionSupportsNativeSteer, type TaskRunStatus, @@ -227,9 +228,17 @@ export interface ISessionStore { messageId: string, patch: { content: string; rawPrompt?: string | ContentBlock[] }, ): void; + setEditingQueuedMessage(taskId: string, messageId: string): void; + clearEditingQueuedMessage(taskId: string): void; clearMessageQueue(taskId: string): void; - dequeueMessagesAsText(taskId: string): string | null; - dequeueMessages(taskId: string): QueuedMessage[]; + dequeueMessagesAsText( + taskId: string, + options?: { stopAtEdited?: boolean }, + ): string | null; + dequeueMessages( + taskId: string, + options?: { stopAtEdited?: boolean }, + ): QueuedMessage[]; prependQueuedMessages(taskId: string, messages: QueuedMessage[]): void; appendOptimisticItem( taskRunId: string, @@ -2097,12 +2106,16 @@ export class SessionService { session: AgentSession, ): boolean { const freshSession = this.d.store.getSessions()[taskRunId]; - const hasQueuedMessages = - freshSession && - freshSession.messageQueue.length > 0 && - freshSession.status === "connected"; - - if (hasQueuedMessages) { + // A message being edited in place holds back itself and everything after + // it, so only the sendable prefix counts. When the whole queue is held, + // this returns false and the caller fires the "turn complete" notification + // (the agent is idle, waiting for the edit to be saved). + const hasSendableMessages = + !!freshSession && + freshSession.status === "connected" && + sendableQueuePrefixLength(freshSession) > 0; + + if (hasSendableMessages) { setTimeout(() => { this.sendQueuedMessages(session.taskId).catch((err) => { this.d.log.error("Failed to send queued messages", { @@ -2113,7 +2126,7 @@ export class SessionService { }, 0); } - return hasQueuedMessages; + return hasSendableMessages; } private handlePermissionRequest( @@ -2384,7 +2397,9 @@ export class SessionService { private async sendQueuedMessages( taskId: string, ): Promise<{ stopReason: string }> { - const combinedText = this.d.store.dequeueMessagesAsText(taskId); + const combinedText = this.d.store.dequeueMessagesAsText(taskId, { + stopAtEdited: true, + }); if (!combinedText) { return { stopReason: "skipped" }; } @@ -2551,6 +2566,26 @@ export class SessionService { } } + /** + * 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 + * it are held back when the turn ends (only the messages before it may send). + */ + setEditingQueuedMessage(taskId: string, messageId: string): void { + this.d.store.setEditingQueuedMessage(taskId, messageId); + } + + /** + * Release an in-place edit hold — the edit was cancelled, or the edited + * message left the queue. Drops the hold and, if the agent is now idle, sends + * the messages the hold was blocking (the turn may have ended while the user + * was still editing, so nothing else would trigger the drain). + */ + clearEditingQueuedMessage(taskId: string): void { + this.d.store.clearEditingQueuedMessage(taskId); + this.flushQueuedMessagesIfIdle(taskId); + } + /** * Update a queued message in place from an edited composer prompt, keeping it * in the queue at its current position. Mirrors the enqueue normalization so @@ -2558,6 +2593,10 @@ export class SessionService { * hold (cloud recomputes the transport display + raw payload; local stores * the serialized text). Returns false when the target is no longer queued so * the caller can fall back to sending it as a new message. + * + * Saving is also what releases the edit hold: the hold is cleared and, if the + * agent finished its turn while the user was editing, the now-unblocked queue + * is drained. */ async updateQueuedMessage( taskId: string, @@ -2580,9 +2619,41 @@ export class SessionService { content: extractPromptText(prompt), }); } + + if (session.editingQueuedId === messageId) { + this.d.store.clearEditingQueuedMessage(taskId); + } + this.flushQueuedMessagesIfIdle(taskId); return true; } + /** + * Nudge the queue after an in-place edit is saved or cancelled. The turn may + * have finished while the user was editing — in that case nothing else would + * trigger the normal turn-end drain, so the now-unblocked messages would sit + * stranded until the next turn. Only sends when the agent is actually idle; + * a mid-turn edit is left for the turn-end drain to pick up. + */ + private flushQueuedMessagesIfIdle(taskId: string): void { + const session = this.d.store.getSessionByTaskId(taskId); + if (!session || session.messageQueue.length === 0) return; + + if (session.isCloud) { + // The cloud flush re-checks run readiness itself, so scheduling while the + // run is still busy is a safe no-op. + this.scheduleCloudQueueFlush(taskId, "edit_released"); + return; + } + + if ( + session.status === "connected" && + !session.isPromptPending && + !session.isCompacting + ) { + this.drainQueuedMessages(session.taskRunId, session); + } + } + /** * Cancel the current prompt. */ @@ -2883,7 +2954,9 @@ export class SessionService { const authStatus = await this.getAuthCredentialsStatus(); if (authStatus.kind === "restoring") return; - const drained = this.d.store.dequeueMessages(taskId); + const drained = this.d.store.dequeueMessages(taskId, { + stopAtEdited: true, + }); const combined = this.d.h.combineQueuedCloudPrompts(drained); if (!combined) return; diff --git a/packages/core/src/sessions/sessionStore.ts b/packages/core/src/sessions/sessionStore.ts index 80921cbdae..b5f7f3355d 100644 --- a/packages/core/src/sessions/sessionStore.ts +++ b/packages/core/src/sessions/sessionStore.ts @@ -1,11 +1,12 @@ import type { ContentBlock } from "@agentclientprotocol/sdk"; -import type { - AcpMessage, - AgentSession, - OptimisticItem, - PermissionRequest, - QueuedMessage, - TaskRunStatus, +import { + type AcpMessage, + type AgentSession, + type OptimisticItem, + type PermissionRequest, + type QueuedMessage, + sendableQueuePrefixLength, + type TaskRunStatus, } from "@posthog/shared"; import { setAutoFreeze } from "immer"; import { immer } from "zustand/middleware/immer"; @@ -245,7 +246,41 @@ export const sessionStoreSetters = { }); }, - dequeueMessagesAsText: (taskId: string): string | null => { + /** + * Mark a queued message as being edited in the composer. While set it holds + * back the auto-drain: the message and everything after it stay queued (only + * the messages before it may send) until the edit is saved or cancelled. + */ + setEditingQueuedMessage: (taskId: string, messageId: string) => { + sessionStore.setState((state) => { + const taskRunId = state.taskIdIndex[taskId]; + if (!taskRunId) return; + const session = state.sessions[taskRunId]; + if (session) session.editingQueuedId = messageId; + }); + }, + + /** Release the in-place edit hold set by {@link setEditingQueuedMessage}. */ + clearEditingQueuedMessage: (taskId: string) => { + sessionStore.setState((state) => { + const taskRunId = state.taskIdIndex[taskId]; + if (!taskRunId) return; + const session = state.sessions[taskRunId]; + if (session) session.editingQueuedId = undefined; + }); + }, + + /** + * Drain the queue as one combined string. With `stopAtEdited`, only the + * messages before the one being edited are drained (the edited message and + * everything after stay queued); otherwise the whole queue drains — the + * default used by the cancel/recall paths that pull everything back into the + * composer. + */ + dequeueMessagesAsText: ( + taskId: string, + options?: { stopAtEdited?: boolean }, + ): string | null => { // Read the queue from the frozen committed state BEFORE entering the // immer draft — same rationale as `dequeueMessages`: anything captured // through a draft proxy can be revoked when setState exits. @@ -255,19 +290,34 @@ export const sessionStoreSetters = { const session = state.sessions[taskRunId]; if (!session || session.messageQueue.length === 0) return null; + const cutoff = options?.stopAtEdited + ? sendableQueuePrefixLength(session) + : session.messageQueue.length; + if (cutoff === 0) return null; + const combined = session.messageQueue + .slice(0, cutoff) .map((msg) => msg.content) .join("\n\n"); sessionStore.setState((draft) => { const trid = draft.taskIdIndex[taskId]; if (!trid) return; const draftSession = draft.sessions[trid]; - if (draftSession) draftSession.messageQueue = []; + if (draftSession) { + draftSession.messageQueue = draftSession.messageQueue.slice(cutoff); + } }); return combined; }, - dequeueMessages: (taskId: string): QueuedMessage[] => { + /** + * Drain the queue as raw messages. See {@link dequeueMessagesAsText} for the + * `stopAtEdited` semantics. + */ + dequeueMessages: ( + taskId: string, + options?: { stopAtEdited?: boolean }, + ): QueuedMessage[] => { // Read the queue from the frozen committed state BEFORE entering the // immer draft, otherwise the items returned are proxies that get // revoked when setState exits and any later access throws @@ -278,14 +328,19 @@ export const sessionStoreSetters = { const session = state.sessions[taskRunId]; if (!session || session.messageQueue.length === 0) return []; - const queuedMessages = [...session.messageQueue]; + const cutoff = options?.stopAtEdited + ? sendableQueuePrefixLength(session) + : session.messageQueue.length; + if (cutoff === 0) return []; + + const queuedMessages = session.messageQueue.slice(0, cutoff); sessionStore.setState((draft) => { const trid = draft.taskIdIndex[taskId]; if (!trid) return; const draftSession = draft.sessions[trid]; if (draftSession) { - draftSession.messageQueue = []; + draftSession.messageQueue = draftSession.messageQueue.slice(cutoff); } }); diff --git a/packages/shared/src/index.ts b/packages/shared/src/index.ts index 35fac6ff05..bf6a734d13 100644 --- a/packages/shared/src/index.ts +++ b/packages/shared/src/index.ts @@ -196,6 +196,7 @@ export { type PermissionRequest, type QueuedMessage, resolveBypassRevertMode, + sendableQueuePrefixLength, type SessionStatus, sessionSupportsNativeSteer, } from "./sessions"; diff --git a/packages/shared/src/sessions.ts b/packages/shared/src/sessions.ts index f938d1b9a4..165931ada6 100644 --- a/packages/shared/src/sessions.ts +++ b/packages/shared/src/sessions.ts @@ -79,6 +79,14 @@ export interface AgentSession { pendingPermissions: Map; pausedDurationMs: number; messageQueue: QueuedMessage[]; + /** + * Id of the queued message the user currently has open in the composer for an + * in-place edit, if any. While set it acts as a drain boundary: when the turn + * ends, everything queued *before* this message auto-sends, but this message + * and everything after it stay queued until the edit is saved or cancelled. + * See {@link sendableQueuePrefixLength}. + */ + editingQueuedId?: string; isCloud?: boolean; cloudStatus?: TaskRunStatus; cloudStage?: string | null; @@ -96,6 +104,23 @@ export interface AgentSession { agentIdleForRunId?: string; } +/** + * How many messages at the head of the queue are eligible to auto-send when the + * turn ends. A message being edited in place ({@link AgentSession.editingQueuedId}) + * is a hard boundary: the messages queued before it may send, but it and + * everything after it stay put until the edit is saved or cancelled. Returns the + * full queue length when nothing is being edited, or when the edited message has + * already left the queue (e.g. it was discarded). + */ +export function sendableQueuePrefixLength( + session: Pick, +): number { + const { messageQueue, editingQueuedId } = session; + if (!editingQueuedId) return messageQueue.length; + const editIndex = messageQueue.findIndex((m) => m.id === editingQueuedId); + return editIndex === -1 ? messageQueue.length : editIndex; +} + export function isSelectGroup( options: SessionConfigSelectOptions, ): options is SessionConfigSelectGroup[] { diff --git a/packages/ui/src/features/sessions/components/QueuedMessagesDock.tsx b/packages/ui/src/features/sessions/components/QueuedMessagesDock.tsx index 2d30364036..46cc7602ac 100644 --- a/packages/ui/src/features/sessions/components/QueuedMessagesDock.tsx +++ b/packages/ui/src/features/sessions/components/QueuedMessagesDock.tsx @@ -20,7 +20,6 @@ import { useSessionStore, } from "@posthog/ui/features/sessions/sessionStore"; import { - useEditingQueuedId, useQueueCollapsed, useSessionViewActions, } from "@posthog/ui/features/sessions/sessionViewStore"; @@ -85,7 +84,7 @@ export function QueuedMessagesDock({ taskId }: QueuedMessagesDockProps) { const supportsNativeSteer = useSupportsNativeSteer(taskId); const editMessage = useEditQueuedMessage(taskId); const cancelEdit = useCancelQueuedMessageEdit(taskId); - const editingId = useEditingQueuedId(taskId); + const editingId = useSessionSelector(taskId, (s) => s?.editingQueuedId); // Narrow reads (not the whole session) so the dock doesn't re-render on every // streamed token while a turn is running. const isCompacting = useSessionSelector( @@ -98,15 +97,16 @@ export function QueuedMessagesDock({ taskId }: QueuedMessagesDockProps) { // so hide it there too — the message stays queued and lands next turn. const canSteer = !isCompacting && !isCloud; const collapsed = useQueueCollapsed(taskId); - const { setQueueCollapsed, clearEditingQueuedId } = useSessionViewActions(); + const { setQueueCollapsed } = useSessionViewActions(); - // If the message being edited leaves the queue (drained on turn end, or - // discarded), drop the stale edit target so the composer sends normally. + // If the message being edited leaves the queue (e.g. discarded), drop the + // stale edit hold so the composer sends normally and any messages the hold + // was blocking can drain. useEffect(() => { if (editingId && !queued.some((m) => m.id === editingId)) { - clearEditingQueuedId(taskId); + sessionService.clearEditingQueuedMessage(taskId); } - }, [editingId, queued, taskId, clearEditingQueuedId]); + }, [editingId, queued, taskId, sessionService]); const handleDragOver: DragDropEvents["dragover"] = useCallback( (event) => { diff --git a/packages/ui/src/features/sessions/hooks/useEditQueuedMessage.ts b/packages/ui/src/features/sessions/hooks/useEditQueuedMessage.ts index 4dd5f51df5..37f8e37834 100644 --- a/packages/ui/src/features/sessions/hooks/useEditQueuedMessage.ts +++ b/packages/ui/src/features/sessions/hooks/useEditQueuedMessage.ts @@ -6,12 +6,16 @@ import { combineQueuedCloudPrompts, promptToQueuedEditorContent, } from "@posthog/core/sessions/cloudPrompt"; +import { + SESSION_SERVICE, + type SessionService, +} from "@posthog/core/sessions/sessionService"; +import { useService } from "@posthog/di/react"; import { useDraftStore } from "@posthog/ui/features/message-editor/draftStore"; import { type QueuedMessage, useSessionIsCloud, } from "@posthog/ui/features/sessions/sessionStore"; -import { useSessionViewActions } from "@posthog/ui/features/sessions/sessionViewStore"; import { useCallback } from "react"; /** @@ -33,7 +37,7 @@ export function useEditQueuedMessage( taskId: string | undefined, ): (message: QueuedMessage) => void { const { requestFocus, setPendingContent } = useDraftStore((s) => s.actions); - const { setEditingQueuedId } = useSessionViewActions(); + const sessionService = useService(SESSION_SERVICE); const isCloud = useSessionIsCloud(taskId); return useCallback( @@ -51,11 +55,11 @@ export function useEditQueuedMessage( } if (!pendingContent) return; - setEditingQueuedId(taskId, message.id); + sessionService.setEditingQueuedMessage(taskId, message.id); setPendingContent(taskId, pendingContent); requestFocus(taskId); }, - [taskId, isCloud, requestFocus, setPendingContent, setEditingQueuedId], + [taskId, isCloud, requestFocus, setPendingContent, sessionService], ); } @@ -68,11 +72,11 @@ export function useCancelQueuedMessageEdit( taskId: string | undefined, ): () => void { const { setPendingContent } = useDraftStore((s) => s.actions); - const { clearEditingQueuedId } = useSessionViewActions(); + const sessionService = useService(SESSION_SERVICE); return useCallback(() => { if (!taskId) return; - clearEditingQueuedId(taskId); + sessionService.clearEditingQueuedMessage(taskId); setPendingContent(taskId, EMPTY_CONTENT); - }, [taskId, clearEditingQueuedId, setPendingContent]); + }, [taskId, sessionService, setPendingContent]); } diff --git a/packages/ui/src/features/sessions/hooks/useSessionCallbacks.ts b/packages/ui/src/features/sessions/hooks/useSessionCallbacks.ts index cb6dbe2c38..8be5e18934 100644 --- a/packages/ui/src/features/sessions/hooks/useSessionCallbacks.ts +++ b/packages/ui/src/features/sessions/hooks/useSessionCallbacks.ts @@ -20,7 +20,6 @@ import { type AgentSession, sessionStoreSetters, } from "@posthog/ui/features/sessions/sessionStore"; -import { sessionViewStore } from "@posthog/ui/features/sessions/sessionViewStore"; import { useTaskViewed } from "@posthog/ui/features/sidebar/useTaskViewed"; import { SHELL_CLIENT, @@ -95,9 +94,8 @@ export function useSessionCallbacks({ // queue rather than sending a new prompt. If the target already drained // or was discarded, fall through and send it as a fresh message. const editingId = - sessionViewStore.getState().editingQueuedIdByTaskId[taskId]; + sessionStoreSetters.getSessionByTaskId(taskId)?.editingQueuedId; if (editingId) { - sessionViewStore.getState().actions.clearEditingQueuedId(taskId); try { const updated = await sessionService.updateQueuedMessage( taskId, @@ -108,11 +106,14 @@ export function useSessionCallbacks({ markAsViewed(taskId); return; } + // Target no longer queued — drop the stale hold and send as new. + sessionService.clearEditingQueuedMessage(taskId); } catch (error) { const message = error instanceof Error ? error.message : "Failed to update message"; toast.error(message); log.error("Failed to update queued message", error); + sessionService.clearEditingQueuedMessage(taskId); return; } } diff --git a/packages/ui/src/features/sessions/sessionServiceHost.test.ts b/packages/ui/src/features/sessions/sessionServiceHost.test.ts index be96ef8179..e1e6ef8867 100644 --- a/packages/ui/src/features/sessions/sessionServiceHost.test.ts +++ b/packages/ui/src/features/sessions/sessionServiceHost.test.ts @@ -67,6 +67,9 @@ const mockSessionStoreSetters = vi.hoisted(() => ({ appendEvents: vi.fn(), enqueueMessage: vi.fn(), removeQueuedMessage: vi.fn(), + updateQueuedMessage: vi.fn(), + setEditingQueuedMessage: vi.fn(), + clearEditingQueuedMessage: vi.fn(), clearMessageQueue: vi.fn(), dequeueMessagesAsText: vi.fn((): string | null => null), dequeueMessages: vi.fn( @@ -4049,6 +4052,7 @@ describe("SessionService", () => { }); expect(mockSessionStoreSetters.dequeueMessages).toHaveBeenCalledWith( "task-123", + { stopAtEdited: true }, ); } finally { vi.useRealTimers(); @@ -5038,6 +5042,108 @@ describe("SessionService", () => { }); }); + describe("in-place edit hold release", () => { + const seedEditedIdleSession = ( + overrides: Partial = {}, + ): AgentSession => { + const session = createMockSession({ + isCloud: false, + status: "connected", + isPromptPending: false, + messageQueue: [{ id: "q-1", content: "old", queuedAt: 1 }], + editingQueuedId: "q-1", + ...overrides, + }); + mockSessionStoreSetters.getSessionByTaskId.mockReturnValue(session); + mockSessionStoreSetters.getSessions.mockReturnValue({ + "run-123": session, + }); + // Mirror the real store so the flush's readiness check sees the hold gone. + mockSessionStoreSetters.clearEditingQueuedMessage.mockImplementation( + () => { + session.editingQueuedId = undefined; + }, + ); + return session; + }; + + it("saving an edit while the agent is idle clears the hold and drains the queue", async () => { + vi.useFakeTimers(); + try { + const service = getSessionService(); + seedEditedIdleSession(); + mockSessionStoreSetters.dequeueMessagesAsText.mockReturnValue("edited"); + mockTrpcAgent.prompt.mutate.mockResolvedValue({ + stopReason: "end_turn", + }); + + const updated = await service.updateQueuedMessage( + "task-123", + "q-1", + "edited", + ); + await vi.advanceTimersByTimeAsync(0); + + expect(updated).toBe(true); + expect( + mockSessionStoreSetters.clearEditingQueuedMessage, + ).toHaveBeenCalledWith("task-123"); + expect( + mockSessionStoreSetters.dequeueMessagesAsText, + ).toHaveBeenCalledWith("task-123", { stopAtEdited: true }); + expect(mockTrpcAgent.prompt.mutate).toHaveBeenCalledWith( + expect.objectContaining({ sessionId: "run-123" }), + ); + } finally { + vi.useRealTimers(); + } + }); + + it("saving an edit while the agent is still busy does not send immediately", async () => { + vi.useFakeTimers(); + try { + const service = getSessionService(); + seedEditedIdleSession({ isPromptPending: true }); + + await service.updateQueuedMessage("task-123", "q-1", "edited"); + await vi.advanceTimersByTimeAsync(0); + + expect( + mockSessionStoreSetters.clearEditingQueuedMessage, + ).toHaveBeenCalledWith("task-123"); + // Left for the turn-end drain — nothing sent mid-turn. + expect( + mockSessionStoreSetters.dequeueMessagesAsText, + ).not.toHaveBeenCalled(); + expect(mockTrpcAgent.prompt.mutate).not.toHaveBeenCalled(); + } finally { + vi.useRealTimers(); + } + }); + + it("cancelling an edit while idle drains the messages the hold was blocking", async () => { + vi.useFakeTimers(); + try { + const service = getSessionService(); + seedEditedIdleSession(); + mockSessionStoreSetters.dequeueMessagesAsText.mockReturnValue("q-1"); + mockTrpcAgent.prompt.mutate.mockResolvedValue({ + stopReason: "end_turn", + }); + + service.clearEditingQueuedMessage("task-123"); + await vi.advanceTimersByTimeAsync(0); + + expect( + mockSessionStoreSetters.dequeueMessagesAsText, + ).toHaveBeenCalledWith("task-123", { stopAtEdited: true }); + expect(mockTrpcAgent.prompt.mutate).toHaveBeenCalled(); + } finally { + vi.useRealTimers(); + } + }); + }); + describe("cancelPrompt", () => { it("returns false if no session exists", async () => { const service = getSessionService(); diff --git a/packages/ui/src/features/sessions/sessionViewStore.ts b/packages/ui/src/features/sessions/sessionViewStore.ts index 6991a3b4ab..ac87639745 100644 --- a/packages/ui/src/features/sessions/sessionViewStore.ts +++ b/packages/ui/src/features/sessions/sessionViewStore.ts @@ -16,12 +16,6 @@ interface SessionViewState { * resets to expanded on app restart. */ queueCollapsedByTaskId: Record; - /** - * Ephemeral per-task id of the queued message currently being edited in the - * composer, keyed by taskId. When set, the composer's next submit updates - * that message in place instead of sending a new prompt. Not persisted. - */ - editingQueuedIdByTaskId: Record; } interface SessionViewActions { @@ -31,8 +25,6 @@ interface SessionViewActions { setGroupOverride: (id: string, expanded: boolean) => void; clearGroupOverrides: () => void; setQueueCollapsed: (taskId: string, collapsed: boolean) => void; - setEditingQueuedId: (taskId: string, messageId: string) => void; - clearEditingQueuedId: (taskId: string) => void; } type SessionViewStore = SessionViewState & { actions: SessionViewActions }; @@ -43,7 +35,6 @@ const useStore = create((set) => ({ showSearch: false, groupOverrides: {}, queueCollapsedByTaskId: {}, - editingQueuedIdByTaskId: {}, actions: { setShowRawLogs: (show) => set({ showRawLogs: show }), setSearchQuery: (query) => set({ searchQuery: query }), @@ -69,20 +60,6 @@ const useStore = create((set) => ({ [taskId]: collapsed, }, })), - setEditingQueuedId: (taskId, messageId) => - set((state) => ({ - editingQueuedIdByTaskId: { - ...state.editingQueuedIdByTaskId, - [taskId]: messageId, - }, - })), - clearEditingQueuedId: (taskId) => - set((state) => { - if (state.editingQueuedIdByTaskId[taskId] === undefined) return state; - const next = { ...state.editingQueuedIdByTaskId }; - delete next[taskId]; - return { editingQueuedIdByTaskId: next }; - }), }, })); @@ -92,12 +69,4 @@ export const useShowSearch = () => useStore((s) => s.showSearch); export const useGroupOverrides = () => useStore((s) => s.groupOverrides); export const useQueueCollapsed = (taskId: string) => useStore((s) => s.queueCollapsedByTaskId[taskId] ?? false); -export const useEditingQueuedId = (taskId: string) => - useStore((s) => s.editingQueuedIdByTaskId[taskId]); export const useSessionViewActions = () => useStore((s) => s.actions); - -/** - * Imperative handle to the view store for non-React call sites (the composer - * submit path reads/clears the editing target here). - */ -export const sessionViewStore = useStore; From 30df9be06eaf0d76304f7dfbf8115ee8f27d0eeb Mon Sep 17 00:00:00 2001 From: Georges-Antoine Assi Date: Mon, 13 Jul 2026 09:57:11 -0400 Subject: [PATCH 04/15] feat(sessions): cancel queued-message edit on Escape MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit While editing a queued message in place, Escape now abandons the edit (clearing the hold and emptying the composer) instead of stopping the running turn. Cancel-edit takes priority over stop-run, and the escape hotkey is enabled whenever an edit is in progress — including while the agent is idle. Generated-By: PostHog Code Task-Id: 5cf31c6c-1ba0-4137-88c1-3a369f3acf36 --- .../components/PromptInput.test.tsx | 45 +++++++++++++++++++ .../message-editor/components/PromptInput.tsx | 21 ++++++++- .../sessions/components/SessionView.tsx | 10 +++++ 3 files changed, 74 insertions(+), 2 deletions(-) diff --git a/packages/ui/src/features/message-editor/components/PromptInput.test.tsx b/packages/ui/src/features/message-editor/components/PromptInput.test.tsx index bd0a282946..e56f646a67 100644 --- a/packages/ui/src/features/message-editor/components/PromptInput.test.tsx +++ b/packages/ui/src/features/message-editor/components/PromptInput.test.tsx @@ -137,3 +137,48 @@ describe("PromptInput submit/stop affordance", () => { expect(send).toBeDisabled(); }); }); + +describe("PromptInput escape handling", () => { + beforeEach(() => { + vi.clearAllMocks(); + editorState.isEmpty = false; + settingsState.slotMachineMode = false; + }); + + it("cancels the queued-message edit on Escape", async () => { + const user = userEvent.setup(); + const onCancelEdit = vi.fn(); + + renderInput({ isEditingQueued: true, onCancelEdit }); + + await user.keyboard("{Escape}"); + expect(onCancelEdit).toHaveBeenCalledOnce(); + }); + + it("prioritizes cancelling the edit over stopping the run on Escape", async () => { + const user = userEvent.setup(); + const onCancel = vi.fn(); + const onCancelEdit = vi.fn(); + + renderInput({ + isLoading: true, + onCancel, + isEditingQueued: true, + onCancelEdit, + }); + + await user.keyboard("{Escape}"); + expect(onCancelEdit).toHaveBeenCalledOnce(); + expect(onCancel).not.toHaveBeenCalled(); + }); + + it("still stops the run on Escape when not editing", async () => { + const user = userEvent.setup(); + const onCancel = vi.fn(); + + renderInput({ isLoading: true, onCancel, isEditingQueued: false }); + + await user.keyboard("{Escape}"); + expect(onCancel).toHaveBeenCalledOnce(); + }); +}); diff --git a/packages/ui/src/features/message-editor/components/PromptInput.tsx b/packages/ui/src/features/message-editor/components/PromptInput.tsx index 57044cf66c..1179fefc3c 100644 --- a/packages/ui/src/features/message-editor/components/PromptInput.tsx +++ b/packages/ui/src/features/message-editor/components/PromptInput.tsx @@ -82,6 +82,13 @@ export interface PromptInputProps { onBashCommand?: (command: string) => void; onBashModeChange?: (isBashMode: boolean) => void; onCancel?: () => void; + /** + * Whether the composer is currently editing a queued message in place. When + * true, Escape abandons the edit (via {@link onCancelEdit}) instead of + * stopping the running turn. + */ + isEditingQueued?: boolean; + onCancelEdit?: () => void; onToggleMessagingMode?: () => void; onAttachFiles?: (files: File[]) => void; onEmptyChange?: (isEmpty: boolean) => void; @@ -126,6 +133,8 @@ export const PromptInput = forwardRef( onBashCommand, onBashModeChange, onCancel, + isEditingQueued = false, + onCancelEdit, onToggleMessagingMode, onAttachFiles, onEmptyChange, @@ -245,6 +254,13 @@ export const PromptInput = forwardRef( (e) => { if (hasOpenOverlay()) return; if (!isActiveSession) return; + // Editing a queued message: Escape abandons the edit. It takes priority + // over stopping a running turn — while editing, Escape just cancels. + if (isEditingQueued && onCancelEdit) { + e.preventDefault(); + onCancelEdit(); + return; + } if (isLoading && onCancel) { e.preventDefault(); onCancel(); @@ -253,9 +269,10 @@ export const PromptInput = forwardRef( { enableOnFormTags: true, enableOnContentEditable: true, - enabled: isLoading && !!onCancel, + enabled: + (isEditingQueued && !!onCancelEdit) || (isLoading && !!onCancel), }, - [isActiveSession, isLoading, onCancel], + [isActiveSession, isLoading, onCancel, isEditingQueued, onCancelEdit], ); useHotkeys( diff --git a/packages/ui/src/features/sessions/components/SessionView.tsx b/packages/ui/src/features/sessions/components/SessionView.tsx index 5c84aaefa6..188b9f5978 100644 --- a/packages/ui/src/features/sessions/components/SessionView.tsx +++ b/packages/ui/src/features/sessions/components/SessionView.tsx @@ -32,12 +32,14 @@ import { SessionResourcesBar } from "@posthog/ui/features/sessions/components/Se import { SteerQueueToggle } from "@posthog/ui/features/sessions/components/SteerQueueToggle"; import { ThreadView } from "@posthog/ui/features/sessions/components/ThreadView"; import { CHAT_CONTENT_MAX_WIDTH } from "@posthog/ui/features/sessions/constants"; +import { useCancelQueuedMessageEdit } from "@posthog/ui/features/sessions/hooks/useEditQueuedMessage"; import { useSessionEventsResidency } from "@posthog/ui/features/sessions/hooks/useSessionEventsResidency"; import { useToggleMessagingMode } from "@posthog/ui/features/sessions/hooks/useToggleMessagingMode"; import { useAdapterForTask, useModeConfigOptionForTask, usePendingPermissionsForTask, + useSessionSelector, useThoughtLevelConfigOptionForTask, } from "@posthog/ui/features/sessions/sessionStore"; import { @@ -321,6 +323,12 @@ export function SessionView({ [isOnline, onBeforeSubmit], ); + const isEditingQueued = useSessionSelector( + taskId, + (s) => !!s?.editingQueuedId, + ); + const cancelQueuedEdit = useCancelQueuedMessageEdit(taskId); + const [isDraggingFile, setIsDraggingFile] = useState(false); const editorRef = useRef(null); const dragCounterRef = useRef(0); @@ -700,6 +708,8 @@ export function SessionView({ onSubmit={handleSubmit} onBashCommand={onBashCommand} onCancel={onCancelPrompt} + isEditingQueued={isEditingQueued} + onCancelEdit={cancelQueuedEdit} /> From 49c2d5e7275fbce6f5685dcf516abb069637287d Mon Sep 17 00:00:00 2001 From: Georges-Antoine Assi Date: Mon, 13 Jul 2026 10:22:33 -0400 Subject: [PATCH 05/15] fix(sessions): keep queued-message height stable and space its actions The editing (purple) state collapsed shorter than the default state because its actions are all ghost icon buttons (fit-content height), with no fixed-height Steer button to anchor the row. Pin the row to min-h-6 so both states match. Also widen the action gap (gap-1 -> gap-2) and add breathing room after the "Editing in composer" label. Generated-By: PostHog Code Task-Id: 5cf31c6c-1ba0-4137-88c1-3a369f3acf36 --- .../components/session-update/QueuedMessageView.tsx | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/packages/ui/src/features/sessions/components/session-update/QueuedMessageView.tsx b/packages/ui/src/features/sessions/components/session-update/QueuedMessageView.tsx index 5b476c5bee..9c88179e87 100644 --- a/packages/ui/src/features/sessions/components/session-update/QueuedMessageView.tsx +++ b/packages/ui/src/features/sessions/components/session-update/QueuedMessageView.tsx @@ -45,7 +45,10 @@ export function QueuedMessageView({ : "border-gray-5 bg-card", )} > - + {/* Pin the row height so it stays constant across states: the non-editing + Steer button (fixed height) anchors it, but the editing state's ghost + icon buttons are `fit-content` and would otherwise collapse shorter. */} + )} - + {isEditing ? ( <> - + Editing in composer {onCancelEdit && ( From 12a4ab39050f54c40cf0d705a61d6f28ced20453 Mon Sep 17 00:00:00 2001 From: Georges-Antoine Assi Date: Mon, 13 Jul 2026 10:32:40 -0400 Subject: [PATCH 06/15] feat(sessions): drain queued messages one turn at a time MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Previously the whole queue was merged into a single prompt when the turn ended, so multiple queued messages landed in one turn. Drain just the head message instead (max: 1) — when that turn completes, the existing turn-end drain fires again for the next one, so the queue sends sequentially. The in-place edit boundary still applies: the cutoff is min(max, sendable-prefix). Cancel/recall paths still pull the whole queue into the composer. Generated-By: PostHog Code Task-Id: 5cf31c6c-1ba0-4137-88c1-3a369f3acf36 --- .../src/sessions/sessionQueueEditing.test.ts | 55 ++++++++++++++++++- packages/core/src/sessions/sessionService.ts | 23 +++++--- packages/core/src/sessions/sessionStore.ts | 41 +++++++++----- .../sessions/sessionServiceHost.test.ts | 6 +- 4 files changed, 97 insertions(+), 28 deletions(-) diff --git a/packages/core/src/sessions/sessionQueueEditing.test.ts b/packages/core/src/sessions/sessionQueueEditing.test.ts index fb758003a4..9c0f178c8a 100644 --- a/packages/core/src/sessions/sessionQueueEditing.test.ts +++ b/packages/core/src/sessions/sessionQueueEditing.test.ts @@ -107,13 +107,19 @@ describe("sendableQueuePrefixLength", () => { it("stops at the edited message, so only earlier messages count", () => { expect( - sendableQueuePrefixLength({ ...q(["a", "b", "c"]), editingQueuedId: "b" }), + sendableQueuePrefixLength({ + ...q(["a", "b", "c"]), + editingQueuedId: "b", + }), ).toBe(1); }); it("returns 0 when the head message is being edited", () => { expect( - sendableQueuePrefixLength({ ...q(["a", "b", "c"]), editingQueuedId: "a" }), + sendableQueuePrefixLength({ + ...q(["a", "b", "c"]), + editingQueuedId: "a", + }), ).toBe(0); }); @@ -189,3 +195,48 @@ describe("editing hold on the drain", () => { expect(queue()).toEqual([]); }); }); + +describe("sequential drain (max: 1)", () => { + it("dequeueMessagesAsText drains only the head message, leaving the rest", () => { + seedQueue([msg("a", "A"), msg("b", "B"), msg("c", "C")]); + + const first = sessionStoreSetters.dequeueMessagesAsText(TASK, { max: 1 }); + expect(first).toBe("A"); + expect(queue().map((m) => m.id)).toEqual(["b", "c"]); + + // The turn-end drain fires again per turn; each call takes the next head. + const second = sessionStoreSetters.dequeueMessagesAsText(TASK, { max: 1 }); + expect(second).toBe("B"); + expect(queue().map((m) => m.id)).toEqual(["c"]); + }); + + it("dequeueMessages drains only the head message as a raw item", () => { + seedQueue([msg("a", "A"), msg("b", "B")]); + + const drained = sessionStoreSetters.dequeueMessages(TASK, { max: 1 }); + + expect(drained.map((m) => m.id)).toEqual(["a"]); + expect(queue().map((m) => m.id)).toEqual(["b"]); + }); + + it("takes min(max, edit boundary): head sends, edited message and rest stay", () => { + seedQueue([msg("a", "A"), msg("b", "B"), msg("c", "C")]); + sessionStoreSetters.setEditingQueuedMessage(TASK, "b"); + + const first = sessionStoreSetters.dequeueMessagesAsText(TASK, { + stopAtEdited: true, + max: 1, + }); + expect(first).toBe("A"); + expect(queue().map((m) => m.id)).toEqual(["b", "c"]); + + // Next drain sends nothing: the new head is the message being edited. + expect( + sessionStoreSetters.dequeueMessagesAsText(TASK, { + stopAtEdited: true, + max: 1, + }), + ).toBeNull(); + expect(queue().map((m) => m.id)).toEqual(["b", "c"]); + }); +}); diff --git a/packages/core/src/sessions/sessionService.ts b/packages/core/src/sessions/sessionService.ts index 29d36d74f4..03d63fbbcf 100644 --- a/packages/core/src/sessions/sessionService.ts +++ b/packages/core/src/sessions/sessionService.ts @@ -28,8 +28,8 @@ import { type PermissionRequest, type QueuedMessage, resolveBypassRevertMode, - sendableQueuePrefixLength, type StoredLogEntry, + sendableQueuePrefixLength, sessionSupportsNativeSteer, type TaskRunStatus, } from "@posthog/shared"; @@ -235,11 +235,11 @@ export interface ISessionStore { clearMessageQueue(taskId: string): void; dequeueMessagesAsText( taskId: string, - options?: { stopAtEdited?: boolean }, + options?: { stopAtEdited?: boolean; max?: number }, ): string | null; dequeueMessages( taskId: string, - options?: { stopAtEdited?: boolean }, + options?: { stopAtEdited?: boolean; max?: number }, ): QueuedMessage[]; prependQueuedMessages(taskId: string, messages: QueuedMessage[]): void; appendOptimisticItem( @@ -2417,16 +2417,19 @@ export class SessionService { } /** - * Send all queued messages as a single prompt. + * Send the next queued message as its own prompt. * Called internally when a turn completes and there are queued messages. - * Queue is cleared atomically before sending - if sending fails, messages are lost - * (this is acceptable since the user can re-type; avoiding complex retry logic). + * Only the head message is dequeued (`max: 1`) so a queue drains one turn at + * a time — when this turn completes, the drain fires again for the next one. + * The message is removed from the queue before sending; if sending fails it + * is lost (acceptable since the user can re-type; avoids complex retry logic). */ private async sendQueuedMessages( taskId: string, ): Promise<{ stopReason: string }> { const combinedText = this.d.store.dequeueMessagesAsText(taskId, { stopAtEdited: true, + max: 1, }); if (!combinedText) { return { stopReason: "skipped" }; @@ -2441,7 +2444,7 @@ export class SessionService { return { stopReason: "no_session" }; } - this.d.log.info("Sending queued messages as single prompt", { + this.d.log.info("Sending next queued message as prompt", { taskId, promptLength: combinedText.length, }); @@ -3070,13 +3073,17 @@ export class SessionService { const authStatus = await this.getAuthCredentialsStatus(); if (authStatus.kind === "restoring") return; + // Drain one message per turn (`max: 1`) so a queue sends sequentially: + // the next turn_complete flushes the next message. A later flush after + // this send finishes picks up the rest. const drained = this.d.store.dequeueMessages(taskId, { stopAtEdited: true, + max: 1, }); const combined = this.d.h.combineQueuedCloudPrompts(drained); if (!combined) return; - this.d.log.info("Sending queued cloud messages", { + this.d.log.info("Sending next queued cloud message", { taskId, drainedCount: drained.length, }); diff --git a/packages/core/src/sessions/sessionStore.ts b/packages/core/src/sessions/sessionStore.ts index b5f7f3355d..1a78b7844b 100644 --- a/packages/core/src/sessions/sessionStore.ts +++ b/packages/core/src/sessions/sessionStore.ts @@ -33,6 +33,23 @@ export const sessionStore = createStore()( })), ); +/** + * How many messages to drain off the head of a queue, honoring both options: + * `stopAtEdited` caps at the in-place edit boundary (nothing from the message + * being edited onward); `max` caps the count. The turn-end auto-drain passes + * `max: 1` so queued messages send one turn at a time instead of merged into + * one prompt; cancel/recall pass neither and take the whole queue. + */ +function drainCutoff( + session: AgentSession, + options?: { stopAtEdited?: boolean; max?: number }, +): number { + const sendable = options?.stopAtEdited + ? sendableQueuePrefixLength(session) + : session.messageQueue.length; + return options?.max != null ? Math.min(sendable, options.max) : sendable; +} + export const sessionStoreSetters = { setSession: (session: AgentSession) => { sessionStore.setState((state) => { @@ -271,15 +288,13 @@ export const sessionStoreSetters = { }, /** - * Drain the queue as one combined string. With `stopAtEdited`, only the - * messages before the one being edited are drained (the edited message and - * everything after stay queued); otherwise the whole queue drains — the - * default used by the cancel/recall paths that pull everything back into the - * composer. + * Drain messages off the head of the queue as one combined string. See + * {@link drainCutoff} for the `stopAtEdited`/`max` semantics; cancel/recall + * pass no options and pull the whole queue back into the composer. */ dequeueMessagesAsText: ( taskId: string, - options?: { stopAtEdited?: boolean }, + options?: { stopAtEdited?: boolean; max?: number }, ): string | null => { // Read the queue from the frozen committed state BEFORE entering the // immer draft — same rationale as `dequeueMessages`: anything captured @@ -290,9 +305,7 @@ export const sessionStoreSetters = { const session = state.sessions[taskRunId]; if (!session || session.messageQueue.length === 0) return null; - const cutoff = options?.stopAtEdited - ? sendableQueuePrefixLength(session) - : session.messageQueue.length; + const cutoff = drainCutoff(session, options); if (cutoff === 0) return null; const combined = session.messageQueue @@ -311,12 +324,12 @@ export const sessionStoreSetters = { }, /** - * Drain the queue as raw messages. See {@link dequeueMessagesAsText} for the - * `stopAtEdited` semantics. + * Drain messages off the head of the queue as raw messages. See + * {@link dequeueMessagesAsText} for the `stopAtEdited`/`max` semantics. */ dequeueMessages: ( taskId: string, - options?: { stopAtEdited?: boolean }, + options?: { stopAtEdited?: boolean; max?: number }, ): QueuedMessage[] => { // Read the queue from the frozen committed state BEFORE entering the // immer draft, otherwise the items returned are proxies that get @@ -328,9 +341,7 @@ export const sessionStoreSetters = { const session = state.sessions[taskRunId]; if (!session || session.messageQueue.length === 0) return []; - const cutoff = options?.stopAtEdited - ? sendableQueuePrefixLength(session) - : session.messageQueue.length; + const cutoff = drainCutoff(session, options); if (cutoff === 0) return []; const queuedMessages = session.messageQueue.slice(0, cutoff); diff --git a/packages/ui/src/features/sessions/sessionServiceHost.test.ts b/packages/ui/src/features/sessions/sessionServiceHost.test.ts index 07e59f2642..bde9c85934 100644 --- a/packages/ui/src/features/sessions/sessionServiceHost.test.ts +++ b/packages/ui/src/features/sessions/sessionServiceHost.test.ts @@ -4103,7 +4103,7 @@ describe("SessionService", () => { }); expect(mockSessionStoreSetters.dequeueMessages).toHaveBeenCalledWith( "task-123", - { stopAtEdited: true }, + { stopAtEdited: true, max: 1 }, ); } finally { vi.useRealTimers(); @@ -5169,7 +5169,7 @@ describe("SessionService", () => { ).toHaveBeenCalledWith("task-123"); expect( mockSessionStoreSetters.dequeueMessagesAsText, - ).toHaveBeenCalledWith("task-123", { stopAtEdited: true }); + ).toHaveBeenCalledWith("task-123", { stopAtEdited: true, max: 1 }); expect(mockTrpcAgent.prompt.mutate).toHaveBeenCalledWith( expect.objectContaining({ sessionId: "run-123" }), ); @@ -5215,7 +5215,7 @@ describe("SessionService", () => { expect( mockSessionStoreSetters.dequeueMessagesAsText, - ).toHaveBeenCalledWith("task-123", { stopAtEdited: true }); + ).toHaveBeenCalledWith("task-123", { stopAtEdited: true, max: 1 }); expect(mockTrpcAgent.prompt.mutate).toHaveBeenCalled(); } finally { vi.useRealTimers(); From fd413744467846b6357fde0587e2cd6e044fbf02 Mon Sep 17 00:00:00 2001 From: Georges-Antoine Assi Date: Mon, 13 Jul 2026 11:05:32 -0400 Subject: [PATCH 07/15] fix(shared): sort sendableQueuePrefixLength export (biome ci) The quality job's `biome ci` failed on organizeImports: the new export sorts after `type SessionStatus` (case-sensitive S < s), not before. Generated-By: PostHog Code Task-Id: 5cf31c6c-1ba0-4137-88c1-3a369f3acf36 --- packages/shared/src/index.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/shared/src/index.ts b/packages/shared/src/index.ts index a7d0fb91a8..b03a71c46c 100644 --- a/packages/shared/src/index.ts +++ b/packages/shared/src/index.ts @@ -207,8 +207,8 @@ export { type PermissionRequest, type QueuedMessage, resolveBypassRevertMode, - sendableQueuePrefixLength, type SessionStatus, + sendableQueuePrefixLength, sessionSupportsNativeSteer, } from "./sessions"; export type { From 15b9e96d6be62affa8a05b211c577ef8ea24c8c6 Mon Sep 17 00:00:00 2001 From: Georges-Antoine Assi Date: Mon, 13 Jul 2026 16:40:37 -0400 Subject: [PATCH 08/15] fix(sessions): restore prior composer draft when cancelling a queued edit Cancelling a queued-message edit unconditionally set the composer to empty content, wiping any draft the user had typed before clicking Edit. Snapshot the composer content when the edit begins and restore it on cancel, falling back to clearing only when there was no prior draft. The snapshot is refreshed at each edit and taken (removed) on cancel, so it can never restore stale content on a later edit. Generated-By: PostHog Code Task-Id: 5cf31c6c-1ba0-4137-88c1-3a369f3acf36 --- .../src/features/message-editor/draftStore.ts | 36 +++++ .../hooks/useEditQueuedMessage.test.tsx | 145 ++++++++++++++++++ .../sessions/hooks/useEditQueuedMessage.ts | 47 +++++- 3 files changed, 220 insertions(+), 8 deletions(-) create mode 100644 packages/ui/src/features/sessions/hooks/useEditQueuedMessage.test.tsx diff --git a/packages/ui/src/features/message-editor/draftStore.ts b/packages/ui/src/features/message-editor/draftStore.ts index 1a344e5ed0..bb155dca7e 100644 --- a/packages/ui/src/features/message-editor/draftStore.ts +++ b/packages/ui/src/features/message-editor/draftStore.ts @@ -23,6 +23,12 @@ interface DraftState { focusRequested: Record; pendingContent: Record; pendingInsert: Record; + /** + * Composer content captured when a queued-message edit begins, so cancelling + * the edit restores the user's prior draft instead of blanking the composer. + * Keyed by sessionId. Not persisted. + */ + preEditDraft: Record; _hasHydrated: boolean; } @@ -49,6 +55,16 @@ export interface DraftActions { /** Insert content at the cursor (append), unlike setPendingContent which replaces. */ insertPendingContent: (sessionId: SessionId, content: EditorContent) => void; clearPendingInsert: (sessionId: SessionId) => void; + /** + * Snapshot composer content before a queued-message edit overwrites it (see + * {@link DraftState.preEditDraft}). Passing null clears any existing snapshot. + */ + setPreEditDraft: ( + sessionId: SessionId, + content: EditorContent | null, + ) => void; + /** Return and remove the snapshot set by {@link setPreEditDraft} (null if none). */ + takePreEditDraft: (sessionId: SessionId) => EditorContent | null; } type DraftStore = DraftState & { actions: DraftActions }; @@ -62,6 +78,7 @@ export const useDraftStore = create()( focusRequested: {}, pendingContent: {}, pendingInsert: {}, + preEditDraft: {}, _hasHydrated: false, actions: { @@ -154,6 +171,25 @@ export const useDraftStore = create()( set((state) => { delete state.pendingInsert[sessionId]; }), + + setPreEditDraft: (sessionId, content) => + set((state) => { + if (content === null) { + delete state.preEditDraft[sessionId]; + } else { + state.preEditDraft[sessionId] = content; + } + }), + + takePreEditDraft: (sessionId) => { + const snapshot = get().preEditDraft[sessionId] ?? null; + if (snapshot) { + set((state) => { + delete state.preEditDraft[sessionId]; + }); + } + return snapshot; + }, }, })), { diff --git a/packages/ui/src/features/sessions/hooks/useEditQueuedMessage.test.tsx b/packages/ui/src/features/sessions/hooks/useEditQueuedMessage.test.tsx new file mode 100644 index 0000000000..5574e9bc72 --- /dev/null +++ b/packages/ui/src/features/sessions/hooks/useEditQueuedMessage.test.tsx @@ -0,0 +1,145 @@ +import type { EditorContent } from "@posthog/core/message-editor/content"; +import type { QueuedMessage } from "@posthog/ui/features/sessions/sessionStore"; +import { act, renderHook } from "@testing-library/react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +vi.mock("@posthog/ui/shell/rendererStorage", () => ({ + electronStorage: { + getItem: () => null, + setItem: () => {}, + removeItem: () => {}, + }, +})); + +const sessionService = vi.hoisted(() => ({ + setEditingQueuedMessage: vi.fn(), + clearEditingQueuedMessage: vi.fn(), +})); + +vi.mock("@posthog/core/sessions/sessionService", () => ({ + SESSION_SERVICE: Symbol.for("test.session-service"), +})); + +vi.mock("@posthog/di/react", () => ({ + useService: () => sessionService, +})); + +const cloudState = vi.hoisted(() => ({ isCloud: false })); +vi.mock("@posthog/ui/features/sessions/sessionStore", () => ({ + useSessionIsCloud: () => cloudState.isCloud, +})); + +import { useDraftStore } from "@posthog/ui/features/message-editor/draftStore"; +import { + useCancelQueuedMessageEdit, + useEditQueuedMessage, +} from "./useEditQueuedMessage"; + +const TASK = "task-1"; + +function pendingFor(sessionId: string): EditorContent | undefined { + return useDraftStore.getState().pendingContent[sessionId]; +} +function textContent(text: string): EditorContent { + return { segments: [{ type: "text", text }] }; +} + +const QUEUED: QueuedMessage = { + id: "q-1", + content: "queued body", + queuedAt: 1, +}; + +describe("useEditQueuedMessage / useCancelQueuedMessageEdit", () => { + beforeEach(() => { + vi.clearAllMocks(); + cloudState.isCloud = false; + useDraftStore.setState((state) => ({ + ...state, + drafts: {}, + pendingContent: {}, + preEditDraft: {}, + _hasHydrated: true, + })); + }); + + it("restores the pre-edit draft when the edit is cancelled", () => { + // The user already had a draft typed before clicking Edit. + act(() => { + useDraftStore.getState().actions.setDraft(TASK, textContent("my draft")); + }); + + const { result: edit } = renderHook(() => useEditQueuedMessage(TASK)); + const { result: cancel } = renderHook(() => + useCancelQueuedMessageEdit(TASK), + ); + + act(() => edit.current(QUEUED)); + + // Editing loads the queued message into the composer. + expect(sessionService.setEditingQueuedMessage).toHaveBeenCalledWith( + TASK, + "q-1", + ); + expect(pendingFor(TASK)?.segments).toEqual([ + { type: "text", text: "queued body" }, + ]); + + act(() => cancel.current()); + + // Cancel restores the draft the user had, not empty content. + expect(sessionService.clearEditingQueuedMessage).toHaveBeenCalledWith(TASK); + expect(pendingFor(TASK)).toEqual(textContent("my draft")); + }); + + it("clears the composer on cancel when there was no prior draft", () => { + const { result: edit } = renderHook(() => useEditQueuedMessage(TASK)); + const { result: cancel } = renderHook(() => + useCancelQueuedMessageEdit(TASK), + ); + + act(() => edit.current(QUEUED)); + act(() => cancel.current()); + + expect(pendingFor(TASK)).toEqual({ segments: [] }); + }); + + it("does not treat a whitespace-only draft as restorable", () => { + act(() => { + useDraftStore.getState().actions.setDraft(TASK, textContent(" ")); + }); + + const { result: edit } = renderHook(() => useEditQueuedMessage(TASK)); + const { result: cancel } = renderHook(() => + useCancelQueuedMessageEdit(TASK), + ); + + act(() => edit.current(QUEUED)); + act(() => cancel.current()); + + expect(pendingFor(TASK)).toEqual({ segments: [] }); + }); + + it("snapshots fresh each edit, so a stale draft is not restored later", () => { + const { result: edit } = renderHook(() => useEditQueuedMessage(TASK)); + const { result: cancel } = renderHook(() => + useCancelQueuedMessageEdit(TASK), + ); + + // First edit with a draft present, then cancel (restores + clears snapshot). + act(() => { + useDraftStore.getState().actions.setDraft(TASK, textContent("first")); + }); + act(() => edit.current(QUEUED)); + act(() => cancel.current()); + expect(pendingFor(TASK)).toEqual(textContent("first")); + + // Second edit with an empty composer must not resurrect the first draft. + act(() => { + useDraftStore.getState().actions.setDraft(TASK, null); + }); + act(() => edit.current(QUEUED)); + act(() => cancel.current()); + expect(pendingFor(TASK)).toEqual({ segments: [] }); + }); +}); diff --git a/packages/ui/src/features/sessions/hooks/useEditQueuedMessage.ts b/packages/ui/src/features/sessions/hooks/useEditQueuedMessage.ts index 37f8e37834..8a14ee757b 100644 --- a/packages/ui/src/features/sessions/hooks/useEditQueuedMessage.ts +++ b/packages/ui/src/features/sessions/hooks/useEditQueuedMessage.ts @@ -1,5 +1,6 @@ import { type EditorContent, + isContentEmpty, xmlToContent, } from "@posthog/core/message-editor/content"; import { @@ -19,7 +20,8 @@ import { import { useCallback } from "react"; /** - * Empty editor content, used to clear the composer when an edit is cancelled. + * Empty editor content, used to clear the composer when a cancelled edit has no + * prior draft to restore. */ const EMPTY_CONTENT: EditorContent = { segments: [] }; @@ -29,6 +31,10 @@ const EMPTY_CONTENT: EditorContent = { segments: [] }; * target so the composer's next submit updates it in place (see * `useSessionCallbacks.handleSendPrompt`) rather than sending a new prompt. * + * Snapshots whatever the user already had in the composer before overwriting it + * with the queued message, so cancelling the edit restores that draft rather + * than blanking it (see `useCancelQueuedMessageEdit`). + * * Content restore mirrors the cancel-to-composer path: cloud keeps its rich * payload (mentions, attachments) via the queued-cloud conversion; local * restores the serialized text (chips reparse from the `` tags). @@ -36,7 +42,8 @@ const EMPTY_CONTENT: EditorContent = { segments: [] }; export function useEditQueuedMessage( taskId: string | undefined, ): (message: QueuedMessage) => void { - const { requestFocus, setPendingContent } = useDraftStore((s) => s.actions); + const { requestFocus, setPendingContent, setPreEditDraft, getDraft } = + useDraftStore((s) => s.actions); const sessionService = useService(SESSION_SERVICE); const isCloud = useSessionIsCloud(taskId); @@ -55,28 +62,52 @@ export function useEditQueuedMessage( } if (!pendingContent) return; + // Capture the current composer draft before the queued message overwrites + // it, so a cancelled edit can put it back. Normalize the legacy string + // form to editor content; an empty draft snapshots as nothing to restore. + const priorDraft = getDraft(taskId); + setPreEditDraft( + taskId, + !priorDraft || isContentEmpty(priorDraft) + ? null + : typeof priorDraft === "string" + ? { segments: [{ type: "text", text: priorDraft }] } + : priorDraft, + ); + sessionService.setEditingQueuedMessage(taskId, message.id); setPendingContent(taskId, pendingContent); requestFocus(taskId); }, - [taskId, isCloud, requestFocus, setPendingContent, sessionService], + [ + taskId, + isCloud, + requestFocus, + setPendingContent, + setPreEditDraft, + getDraft, + sessionService, + ], ); } /** * Abandon an in-progress queued-message edit: drop the edit target so the next - * submit sends normally again, and clear the composer so the abandoned draft - * doesn't linger. + * submit sends normally again, and restore the draft that was in the composer + * before the edit began (blanking it only when there was none), so cancelling + * never discards unrelated work the user had typed. */ export function useCancelQueuedMessageEdit( taskId: string | undefined, ): () => void { - const { setPendingContent } = useDraftStore((s) => s.actions); + const { setPendingContent, takePreEditDraft } = useDraftStore( + (s) => s.actions, + ); const sessionService = useService(SESSION_SERVICE); return useCallback(() => { if (!taskId) return; sessionService.clearEditingQueuedMessage(taskId); - setPendingContent(taskId, EMPTY_CONTENT); - }, [taskId, sessionService, setPendingContent]); + setPendingContent(taskId, takePreEditDraft(taskId) ?? EMPTY_CONTENT); + }, [taskId, sessionService, setPendingContent, takePreEditDraft]); } From b99c446084568700bfd1fdfd29ed6431c12eef3c Mon Sep 17 00:00:00 2001 From: Georges-Antoine Assi Date: Tue, 14 Jul 2026 09:32:15 -0400 Subject: [PATCH 09/15] fix(sessions): re-check queue membership after cloud edit normalization updateQueuedMessage awaited resolveCloudPrompt before writing to the store. If the message drained during that await (a turn completed and sent it), the store update became a no-op but the method still returned true, so the caller treated the edit as saved and never fell back to sending it as a fresh message -- losing the edit. Re-read fresh session state after the await and return false if the target is no longer queued. Also read fresh state for the edit-hold clear, since the pre-await snapshot may be stale. Generated-By: PostHog Code Task-Id: 5cf31c6c-1ba0-4137-88c1-3a369f3acf36 --- packages/core/src/sessions/sessionService.ts | 12 +++- .../sessions/sessionServiceHost.test.ts | 69 +++++++++++++++++++ 2 files changed, 80 insertions(+), 1 deletion(-) diff --git a/packages/core/src/sessions/sessionService.ts b/packages/core/src/sessions/sessionService.ts index 03d63fbbcf..6a29c2544e 100644 --- a/packages/core/src/sessions/sessionService.ts +++ b/packages/core/src/sessions/sessionService.ts @@ -2728,6 +2728,13 @@ export class SessionService { if (session.isCloud) { const normalizedPrompt = await this.resolveCloudPrompt(prompt); + // Cloud normalization awaits, during which the message may have drained + // (a turn completed and sent it). Re-check against fresh state: without + // this, the store update below is a silent no-op yet we'd still report a + // successful save, so the caller wouldn't fall back to sending the edit + // as a fresh message and the edit would be lost. + const fresh = this.d.store.getSessionByTaskId(taskId); + if (!fresh?.messageQueue.some((m) => m.id === messageId)) return false; const transport = this.d.h.getCloudPromptTransport(normalizedPrompt); this.d.store.updateQueuedMessage(taskId, messageId, { content: transport.promptText, @@ -2739,7 +2746,10 @@ export class SessionService { }); } - if (session.editingQueuedId === messageId) { + // Read fresh: the cloud path awaited above, so the pre-await `session` + // snapshot may be stale for the edit-hold decision. + const latest = this.d.store.getSessionByTaskId(taskId); + if (latest?.editingQueuedId === messageId) { this.d.store.clearEditingQueuedMessage(taskId); } this.flushQueuedMessagesIfIdle(taskId); diff --git a/packages/ui/src/features/sessions/sessionServiceHost.test.ts b/packages/ui/src/features/sessions/sessionServiceHost.test.ts index bde9c85934..6cda0343e2 100644 --- a/packages/ui/src/features/sessions/sessionServiceHost.test.ts +++ b/packages/ui/src/features/sessions/sessionServiceHost.test.ts @@ -5223,6 +5223,75 @@ describe("SessionService", () => { }); }); + describe("updateQueuedMessage cloud normalization race", () => { + const cloudSession = ( + overrides: Partial = {}, + ): AgentSession => + createMockSession({ + isCloud: true, + cloudStatus: "in_progress", + status: "connected", + isPromptPending: true, + messageQueue: [ + { + id: "q-1", + content: "old", + rawPrompt: [{ type: "text", text: "old" }], + queuedAt: 1, + }, + ], + editingQueuedId: "q-1", + ...overrides, + }); + + it("returns false when the message drains while cloud normalization awaits", async () => { + const service = getSessionService(); + // Present for the initial membership check, gone for the post-await + // re-check (a turn completed and drained it during normalization). + mockSessionStoreSetters.getSessionByTaskId + .mockReturnValueOnce(cloudSession()) + .mockReturnValue( + cloudSession({ messageQueue: [], editingQueuedId: undefined }), + ); + + const updated = await service.updateQueuedMessage( + "task-123", + "q-1", + "edited", + ); + + // No-op store write must not be reported as a save, so the caller falls + // back to sending the edit as a fresh message instead of losing it. + expect(updated).toBe(false); + expect( + mockSessionStoreSetters.updateQueuedMessage, + ).not.toHaveBeenCalled(); + expect( + mockSessionStoreSetters.clearEditingQueuedMessage, + ).not.toHaveBeenCalled(); + }); + + it("updates in place when the message is still queued after normalization", async () => { + const service = getSessionService(); + mockSessionStoreSetters.getSessionByTaskId.mockReturnValue( + cloudSession(), + ); + + const updated = await service.updateQueuedMessage( + "task-123", + "q-1", + "edited", + ); + + expect(updated).toBe(true); + expect(mockSessionStoreSetters.updateQueuedMessage).toHaveBeenCalledWith( + "task-123", + "q-1", + expect.objectContaining({ content: expect.any(String) }), + ); + }); + }); + describe("cancelPrompt", () => { it("returns false if no session exists", async () => { const service = getSessionService(); From 9889d1bda91e7b9970bc52e5258536535b144b0c Mon Sep 17 00:00:00 2001 From: Georges-Antoine Assi Date: Tue, 14 Jul 2026 09:44:14 -0400 Subject: [PATCH 10/15] fix(sessions): keep the edit hold when a queued-message edit fails MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit On a failed updateQueuedMessage the catch cleared editingQueuedId, which releases the hold and can flush the queue when the run is idle — sending the original, unedited message, the opposite of what the user intended. Keep the hold on failure so the message stays queued and unsent, and restore the edited text to the composer (unless the user already started typing) so they can retry the save or cancel the edit explicitly. Generated-By: PostHog Code Task-Id: 5cf31c6c-1ba0-4137-88c1-3a369f3acf36 --- .../hooks/useSessionCallbacks.test.tsx | 140 ++++++++++++++++++ .../sessions/hooks/useSessionCallbacks.ts | 10 +- 2 files changed, 149 insertions(+), 1 deletion(-) create mode 100644 packages/ui/src/features/sessions/hooks/useSessionCallbacks.test.tsx diff --git a/packages/ui/src/features/sessions/hooks/useSessionCallbacks.test.tsx b/packages/ui/src/features/sessions/hooks/useSessionCallbacks.test.tsx new file mode 100644 index 0000000000..2537d12bbe --- /dev/null +++ b/packages/ui/src/features/sessions/hooks/useSessionCallbacks.test.tsx @@ -0,0 +1,140 @@ +import type { Task } from "@posthog/shared/domain-types"; +import { renderHook } from "@testing-library/react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +vi.mock("@posthog/ui/shell/rendererStorage", () => ({ + electronStorage: { + getItem: () => null, + setItem: () => {}, + removeItem: () => {}, + }, +})); + +const sessionService = vi.hoisted(() => ({ + updateQueuedMessage: vi.fn(), + clearEditingQueuedMessage: vi.fn(), + sendPrompt: vi.fn(), +})); + +vi.mock("@posthog/core/sessions/sessionService", () => ({ + SESSION_SERVICE: Symbol.for("test.session-service"), +})); + +vi.mock("@posthog/ui/features/terminal/shellClient", () => ({ + SHELL_CLIENT: Symbol.for("test.shell-client"), +})); + +vi.mock("@posthog/di/react", () => ({ + useService: (token: symbol) => + token === Symbol.for("test.session-service") ? sessionService : {}, +})); + +vi.mock("@posthog/host-router/react", () => ({ + useHostTRPCClient: () => ({ skills: { list: { query: async () => [] } } }), +})); + +const taskViewed = vi.hoisted(() => ({ + markActivity: vi.fn(), + markAsViewed: vi.fn(), +})); +vi.mock("@posthog/ui/features/sidebar/useTaskViewed", () => ({ + useTaskViewed: () => taskViewed, +})); + +vi.mock("@posthog/ui/features/sessions/hooks/useMessagingMode", () => ({ + useMessagingMode: () => "queue", +})); + +// No code command / skill rewrite; the raw text is used as the prompt. +vi.mock("@posthog/ui/features/message-editor/commands", () => ({ + tryExecuteCodeCommand: async () => false, + rewriteLocalSkillCommandPrompt: () => null, + resolveLocalSkillPrompt: async (text: string) => text, +})); + +const editingIdState = vi.hoisted(() => ({ + editingQueuedId: "q-1" as string | undefined, +})); +vi.mock("@posthog/ui/features/sessions/sessionStore", () => ({ + sessionStoreSetters: { + getSessionByTaskId: () => editingIdState, + }, +})); + +vi.mock("@posthog/ui/router/useAppView", () => ({ + getAppViewSnapshot: () => null, +})); + +const toastError = vi.hoisted(() => vi.fn()); +vi.mock("@posthog/ui/primitives/toast", () => ({ + toast: { error: toastError }, +})); + +import { useDraftStore } from "@posthog/ui/features/message-editor/draftStore"; +import { useSessionCallbacks } from "./useSessionCallbacks"; + +const TASK = "task-1"; +const task = { id: TASK, latest_run: null } as unknown as Task; + +function renderCallbacks() { + return renderHook(() => + useSessionCallbacks({ + taskId: TASK, + task, + session: undefined, + repoPath: "/repo", + }), + ); +} + +describe("useSessionCallbacks.handleSendPrompt — failed queued edit", () => { + beforeEach(() => { + vi.clearAllMocks(); + editingIdState.editingQueuedId = "q-1"; + useDraftStore.setState((state) => ({ + ...state, + drafts: {}, + pendingContent: {}, + _hasHydrated: true, + })); + }); + + it("keeps the edit hold and does not send when the edit fails", async () => { + sessionService.updateQueuedMessage.mockRejectedValue( + new Error("cloud edit failed"), + ); + + const { result } = renderCallbacks(); + await result.current.handleSendPrompt("my edit"); + + // The edit was attempted... + expect(sessionService.updateQueuedMessage).toHaveBeenCalledWith( + TASK, + "q-1", + "my edit", + ); + // ...and it failed, so the user is told. + expect(toastError).toHaveBeenCalled(); + // Critically: the hold is NOT released (which would drain and send the + // original, unedited message) and no fresh prompt is sent. + expect(sessionService.clearEditingQueuedMessage).not.toHaveBeenCalled(); + expect(sessionService.sendPrompt).not.toHaveBeenCalled(); + // The edited text is restored to the composer so the user can retry. + expect(useDraftStore.getState().pendingContent[TASK]).toBeDefined(); + }); + + it("releases the hold and sends fresh when the target is no longer queued", async () => { + // updateQueuedMessage resolves false: the message already drained. + sessionService.updateQueuedMessage.mockResolvedValue(false); + sessionService.sendPrompt.mockResolvedValue(undefined); + + const { result } = renderCallbacks(); + await result.current.handleSendPrompt("my edit"); + + // Stale hold dropped, and the edit is sent as a brand-new message. + expect(sessionService.clearEditingQueuedMessage).toHaveBeenCalledWith(TASK); + expect(sessionService.sendPrompt).toHaveBeenCalledWith(TASK, "my edit", { + steer: false, + }); + }); +}); diff --git a/packages/ui/src/features/sessions/hooks/useSessionCallbacks.ts b/packages/ui/src/features/sessions/hooks/useSessionCallbacks.ts index 3eb304d2cd..ee43a113ca 100644 --- a/packages/ui/src/features/sessions/hooks/useSessionCallbacks.ts +++ b/packages/ui/src/features/sessions/hooks/useSessionCallbacks.ts @@ -117,7 +117,15 @@ export function useSessionCallbacks({ error instanceof Error ? error.message : "Failed to update message"; toast.error(message); log.error("Failed to update queued message", error); - sessionService.clearEditingQueuedMessage(taskId); + // Keep the edit hold: releasing it would let the original, unedited + // message drain and send — the opposite of what the user intended by + // editing. The message stays held and the composer restores the + // edited text (unless the user already started typing) so they can + // retry the save or cancel the edit explicitly. + if (isContentEmpty(useDraftStore.getState().drafts[taskId] ?? null)) { + setPendingContent(taskId, xmlToContent(promptText ?? text)); + requestFocus(taskId); + } return; } } From 85b82f807736c0a2ef9aedcfafe9a172f979d1a8 Mon Sep 17 00:00:00 2001 From: Charles Vien Date: Tue, 14 Jul 2026 23:56:00 -0700 Subject: [PATCH 11/15] dedupe queued message drain in session store --- packages/core/src/sessions/sessionStore.ts | 88 +++++++++------------- 1 file changed, 36 insertions(+), 52 deletions(-) diff --git a/packages/core/src/sessions/sessionStore.ts b/packages/core/src/sessions/sessionStore.ts index 1a78b7844b..f8300eb8af 100644 --- a/packages/core/src/sessions/sessionStore.ts +++ b/packages/core/src/sessions/sessionStore.ts @@ -50,6 +50,38 @@ function drainCutoff( return options?.max != null ? Math.min(sendable, options.max) : sendable; } +/** + * Drain messages off the head of the queue, honoring {@link drainCutoff}. + * Reads the queue from the frozen committed state BEFORE entering the immer + * draft, otherwise the returned items are proxies that get revoked when + * setState exits and any later access throws "Cannot perform 'get' on a proxy + * that has been revoked". + */ +function drainQueueHead( + taskId: string, + options?: { stopAtEdited?: boolean; max?: number }, +): QueuedMessage[] { + const state = sessionStore.getState(); + const taskRunId = state.taskIdIndex[taskId]; + if (!taskRunId) return []; + const session = state.sessions[taskRunId]; + if (!session || session.messageQueue.length === 0) return []; + + const cutoff = drainCutoff(session, options); + if (cutoff === 0) return []; + + const drained = session.messageQueue.slice(0, cutoff); + sessionStore.setState((draft) => { + const trid = draft.taskIdIndex[taskId]; + if (!trid) return; + const draftSession = draft.sessions[trid]; + if (draftSession) { + draftSession.messageQueue = draftSession.messageQueue.slice(cutoff); + } + }); + return drained; +} + export const sessionStoreSetters = { setSession: (session: AgentSession) => { sessionStore.setState((state) => { @@ -296,31 +328,9 @@ export const sessionStoreSetters = { taskId: string, options?: { stopAtEdited?: boolean; max?: number }, ): string | null => { - // Read the queue from the frozen committed state BEFORE entering the - // immer draft — same rationale as `dequeueMessages`: anything captured - // through a draft proxy can be revoked when setState exits. - const state = sessionStore.getState(); - const taskRunId = state.taskIdIndex[taskId]; - if (!taskRunId) return null; - const session = state.sessions[taskRunId]; - if (!session || session.messageQueue.length === 0) return null; - - const cutoff = drainCutoff(session, options); - if (cutoff === 0) return null; - - const combined = session.messageQueue - .slice(0, cutoff) - .map((msg) => msg.content) - .join("\n\n"); - sessionStore.setState((draft) => { - const trid = draft.taskIdIndex[taskId]; - if (!trid) return; - const draftSession = draft.sessions[trid]; - if (draftSession) { - draftSession.messageQueue = draftSession.messageQueue.slice(cutoff); - } - }); - return combined; + const drained = drainQueueHead(taskId, options); + if (drained.length === 0) return null; + return drained.map((msg) => msg.content).join("\n\n"); }, /** @@ -330,33 +340,7 @@ export const sessionStoreSetters = { dequeueMessages: ( taskId: string, options?: { stopAtEdited?: boolean; max?: number }, - ): QueuedMessage[] => { - // Read the queue from the frozen committed state BEFORE entering the - // immer draft, otherwise the items returned are proxies that get - // revoked when setState exits and any later access throws - // "Cannot perform 'get' on a proxy that has been revoked". - const state = sessionStore.getState(); - const taskRunId = state.taskIdIndex[taskId]; - if (!taskRunId) return []; - const session = state.sessions[taskRunId]; - if (!session || session.messageQueue.length === 0) return []; - - const cutoff = drainCutoff(session, options); - if (cutoff === 0) return []; - - const queuedMessages = session.messageQueue.slice(0, cutoff); - - sessionStore.setState((draft) => { - const trid = draft.taskIdIndex[taskId]; - if (!trid) return; - const draftSession = draft.sessions[trid]; - if (draftSession) { - draftSession.messageQueue = draftSession.messageQueue.slice(cutoff); - } - }); - - return queuedMessages; - }, + ): QueuedMessage[] => drainQueueHead(taskId, options), /** * Splice messages back at the head of the queue. Used to roll back a From bae4e4d559cbfc84ddc31fbf2c4b8da9745b8105 Mon Sep 17 00:00:00 2001 From: Charles Vien Date: Tue, 14 Jul 2026 23:56:38 -0700 Subject: [PATCH 12/15] add textToContent editor content helper --- packages/core/src/message-editor/content.ts | 5 +++++ .../ui/src/features/sessions/hooks/useEditQueuedMessage.ts | 3 ++- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/packages/core/src/message-editor/content.ts b/packages/core/src/message-editor/content.ts index c997c11add..8bdf35b75b 100644 --- a/packages/core/src/message-editor/content.ts +++ b/packages/core/src/message-editor/content.ts @@ -195,6 +195,11 @@ export function xmlToPlainText(xml: string): string { return contentToPlainText(xmlToContent(xml)); } +/** Wrap plain text in editor content as a single text segment, no chip parsing. */ +export function textToContent(text: string): EditorContent { + return { segments: [{ type: "text", text }] }; +} + export function isContentEmpty( content: EditorContent | null | string, ): boolean { diff --git a/packages/ui/src/features/sessions/hooks/useEditQueuedMessage.ts b/packages/ui/src/features/sessions/hooks/useEditQueuedMessage.ts index 8a14ee757b..7eab3c95cd 100644 --- a/packages/ui/src/features/sessions/hooks/useEditQueuedMessage.ts +++ b/packages/ui/src/features/sessions/hooks/useEditQueuedMessage.ts @@ -1,6 +1,7 @@ import { type EditorContent, isContentEmpty, + textToContent, xmlToContent, } from "@posthog/core/message-editor/content"; import { @@ -71,7 +72,7 @@ export function useEditQueuedMessage( !priorDraft || isContentEmpty(priorDraft) ? null : typeof priorDraft === "string" - ? { segments: [{ type: "text", text: priorDraft }] } + ? textToContent(priorDraft) : priorDraft, ); From 767b08ac96768cbe34ddd5014ea6b20a6d78f4cd Mon Sep 17 00:00:00 2001 From: Charles Vien Date: Tue, 14 Jul 2026 23:56:44 -0700 Subject: [PATCH 13/15] fix drain and stop handling for queued edits --- packages/core/src/sessions/sessionService.ts | 29 +- .../hooks/useSessionCallbacks.test.tsx | 94 ++++- .../sessions/hooks/useSessionCallbacks.ts | 26 +- .../sessions/sessionServiceHost.test.ts | 335 ++++++++++++++++++ 4 files changed, 468 insertions(+), 16 deletions(-) diff --git a/packages/core/src/sessions/sessionService.ts b/packages/core/src/sessions/sessionService.ts index 6a29c2544e..16d6e1d647 100644 --- a/packages/core/src/sessions/sessionService.ts +++ b/packages/core/src/sessions/sessionService.ts @@ -1003,6 +1003,9 @@ export class SessionService { if (previous) { session.optimisticItems = previous.optimisticItems; session.messageQueue = previous.messageQueue; + // Keep the in-place edit hold with the queue it guards: dropping it here + // would let the edited message auto-send in its stale, pre-edit form. + session.editingQueuedId = previous.editingQueuedId; session.isPromptPending = previous.isPromptPending; session.promptStartedAt = previous.promptStartedAt; session.pausedDurationMs = previous.pausedDurationMs; @@ -2023,10 +2026,15 @@ export class SessionService { } const stopReason = (msg.result as { stopReason?: string }).stopReason; - const hasQueuedMessages = this.drainQueuedMessages(taskRunId, session); - - // Only notify when queue is empty - queued messages will start a new turn - if (stopReason && !hasQueuedMessages) { + // A cancelled turn is an explicit stop: auto-firing queued messages + // right after would restart the agent the user just halted. + const hasSendableMessages = + stopReason === "cancelled" + ? false + : this.drainQueuedMessages(taskRunId, session); + + // Only notify when nothing is sendable - queued messages start a new turn + if (stopReason && !hasSendableMessages) { this.d.notifyPromptComplete( session.taskTitle, stopReason, @@ -2136,6 +2144,19 @@ export class SessionService { if (hasSendableMessages) { setTimeout(() => { + // Re-check at fire time: the turn-end drain and the edit-release flush + // can each schedule a timer, and whichever fires second must not start + // a concurrent prompt (the first send flips isPromptPending before it + // awaits, so this check observes it). + const latest = this.d.store.getSessionByTaskId(session.taskId); + if ( + !latest || + latest.status !== "connected" || + latest.isPromptPending || + latest.isCompacting + ) { + return; + } this.sendQueuedMessages(session.taskId).catch((err) => { this.d.log.error("Failed to send queued messages", { taskId: session.taskId, diff --git a/packages/ui/src/features/sessions/hooks/useSessionCallbacks.test.tsx b/packages/ui/src/features/sessions/hooks/useSessionCallbacks.test.tsx index 2537d12bbe..394156e47d 100644 --- a/packages/ui/src/features/sessions/hooks/useSessionCallbacks.test.tsx +++ b/packages/ui/src/features/sessions/hooks/useSessionCallbacks.test.tsx @@ -14,6 +14,7 @@ const sessionService = vi.hoisted(() => ({ updateQueuedMessage: vi.fn(), clearEditingQueuedMessage: vi.fn(), sendPrompt: vi.fn(), + cancelPrompt: vi.fn(), })); vi.mock("@posthog/core/sessions/sessionService", () => ({ @@ -52,12 +53,17 @@ vi.mock("@posthog/ui/features/message-editor/commands", () => ({ resolveLocalSkillPrompt: async (text: string) => text, })); -const editingIdState = vi.hoisted(() => ({ +const sessionState = vi.hoisted(() => ({ editingQueuedId: "q-1" as string | undefined, + messageQueue: [] as Array<{ id: string; content: string; queuedAt: number }>, })); +const dequeueMessages = vi.hoisted(() => + vi.fn(() => [] as Array<{ id: string; content: string; queuedAt: number }>), +); vi.mock("@posthog/ui/features/sessions/sessionStore", () => ({ sessionStoreSetters: { - getSessionByTaskId: () => editingIdState, + getSessionByTaskId: () => sessionState, + dequeueMessages, }, })); @@ -87,10 +93,11 @@ function renderCallbacks() { ); } -describe("useSessionCallbacks.handleSendPrompt — failed queued edit", () => { +describe("useSessionCallbacks.handleSendPrompt while editing a queued message", () => { beforeEach(() => { vi.clearAllMocks(); - editingIdState.editingQueuedId = "q-1"; + sessionState.editingQueuedId = "q-1"; + sessionState.messageQueue = []; useDraftStore.setState((state) => ({ ...state, drafts: {}, @@ -137,4 +144,83 @@ describe("useSessionCallbacks.handleSendPrompt — failed queued edit", () => { steer: false, }); }); + + it("updates in place and never sends when the edit saves", async () => { + sessionService.updateQueuedMessage.mockResolvedValue(true); + + const { result } = renderCallbacks(); + await result.current.handleSendPrompt("my edit"); + + expect(sessionService.updateQueuedMessage).toHaveBeenCalledWith( + TASK, + "q-1", + "my edit", + ); + expect(taskViewed.markAsViewed).toHaveBeenCalledWith(TASK); + // Saving releases the hold inside the service, not the hook. + expect(sessionService.clearEditingQueuedMessage).not.toHaveBeenCalled(); + expect(sessionService.sendPrompt).not.toHaveBeenCalled(); + }); +}); + +describe("useSessionCallbacks.handleCancelPrompt", () => { + beforeEach(() => { + vi.clearAllMocks(); + sessionState.editingQueuedId = undefined; + sessionState.messageQueue = []; + dequeueMessages.mockReturnValue([]); + sessionService.cancelPrompt.mockResolvedValue(true); + useDraftStore.setState((state) => ({ + ...state, + drafts: {}, + pendingContent: {}, + _hasHydrated: true, + })); + }); + + it("recalls the queue into the composer when no edit is active", async () => { + dequeueMessages.mockReturnValue([ + { id: "q-1", content: "first", queuedAt: 1 }, + { id: "q-2", content: "second", queuedAt: 2 }, + ]); + + const { result } = renderCallbacks(); + await result.current.handleCancelPrompt(); + + expect(sessionService.cancelPrompt).toHaveBeenCalledWith(TASK); + expect(dequeueMessages).toHaveBeenCalledWith(TASK); + expect(useDraftStore.getState().pendingContent[TASK]).toEqual({ + segments: [{ type: "text", text: "first\n\nsecond" }], + }); + }); + + it("stops without touching the queue or composer while an edit is active", async () => { + sessionState.editingQueuedId = "q-1"; + sessionState.messageQueue = [{ id: "q-1", content: "old", queuedAt: 1 }]; + + const { result } = renderCallbacks(); + await result.current.handleCancelPrompt(); + + expect(sessionService.cancelPrompt).toHaveBeenCalledWith(TASK); + // The queue is left in place (the edit hold keeps it from auto-sending) + // and the composer keeps the in-progress edit. + expect(dequeueMessages).not.toHaveBeenCalled(); + expect(useDraftStore.getState().pendingContent[TASK]).toBeUndefined(); + }); + + it("falls back to the normal recall when the edit hold is stale", async () => { + sessionState.editingQueuedId = "q-gone"; + sessionState.messageQueue = [{ id: "q-1", content: "first", queuedAt: 1 }]; + dequeueMessages.mockReturnValue([ + { id: "q-1", content: "first", queuedAt: 1 }, + ]); + + const { result } = renderCallbacks(); + await result.current.handleCancelPrompt(); + + expect(dequeueMessages).toHaveBeenCalledWith(TASK); + expect(useDraftStore.getState().pendingContent[TASK]).toEqual({ + segments: [{ type: "text", text: "first" }], + }); + }); }); diff --git a/packages/ui/src/features/sessions/hooks/useSessionCallbacks.ts b/packages/ui/src/features/sessions/hooks/useSessionCallbacks.ts index ee43a113ca..188d470ce2 100644 --- a/packages/ui/src/features/sessions/hooks/useSessionCallbacks.ts +++ b/packages/ui/src/features/sessions/hooks/useSessionCallbacks.ts @@ -1,5 +1,6 @@ import { isContentEmpty, + textToContent, xmlToContent, } from "@posthog/core/message-editor/content"; import { @@ -172,6 +173,22 @@ export function useSessionCallbacks({ ); const handleCancelPrompt = useCallback(async () => { + // Stopping while a queued message is being edited: halt the turn but leave + // the queue and the composer alone, since recalling the queue into the + // 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 editingId = currentSession?.editingQueuedId; + if ( + editingId && + currentSession?.messageQueue.some((m) => m.id === editingId) + ) { + const result = await sessionService.cancelPrompt(taskId); + log.info("Prompt cancelled during queued edit", { success: result }); + requestFocus(taskId); + return; + } + const queuedMessages = sessionStoreSetters.dequeueMessages(taskId); const result = await sessionService.cancelPrompt(taskId); log.info("Prompt cancelled", { success: result }); @@ -183,14 +200,7 @@ export function useSessionCallbacks({ if (queuedPrompt) { const pendingContent = sessionRef.current?.isCloud ? promptToQueuedEditorContent(queuedPrompt) - : { - segments: [ - { - type: "text" as const, - text: typeof queuedPrompt === "string" ? queuedPrompt : "", - }, - ], - }; + : textToContent(typeof queuedPrompt === "string" ? queuedPrompt : ""); setPendingContent(taskId, pendingContent); } diff --git a/packages/ui/src/features/sessions/sessionServiceHost.test.ts b/packages/ui/src/features/sessions/sessionServiceHost.test.ts index 6cda0343e2..7d89f2403e 100644 --- a/packages/ui/src/features/sessions/sessionServiceHost.test.ts +++ b/packages/ui/src/features/sessions/sessionServiceHost.test.ts @@ -4900,6 +4900,213 @@ describe("SessionService", () => { }); }); + describe("turn-end queue drain gating", () => { + async function connectWithLiveSession() { + const service = getSessionService(); + + let session: AgentSession | undefined; + mockSessionStoreSetters.getSessionByTaskId.mockImplementation( + () => session, + ); + mockSessionStoreSetters.getSessions.mockImplementation(() => + session ? { "run-123": session } : {}, + ); + mockSessionStoreSetters.updateSession.mockImplementation( + (_taskRunId, updates) => { + if (session) + session = { ...session, ...(updates as Partial) }; + }, + ); + mockSessionStoreSetters.setSession.mockImplementation((next) => { + session = next as AgentSession; + }); + mockSessionStoreSetters.clearEditingQueuedMessage.mockImplementation( + () => { + if (session) session = { ...session, editingQueuedId: undefined }; + }, + ); + + mockBuildAuthenticatedClient.mockReturnValue({ + ...mockAuthenticatedClient, + createTaskRun: vi.fn().mockResolvedValue({ id: "run-123" }), + appendTaskRunLog: vi.fn(), + }); + mockTrpcAgent.start.mutate.mockResolvedValue({ + channel: "agent-event:run-123", + configOptions: [], + }); + + await service.connectToTask({ + task: createMockTask(), + repoPath: "/repo", + }); + + const onData = mockTrpcAgent.onSessionEvent.subscribe.mock.calls.at( + -1, + )?.[1]?.onData as ((payload: unknown) => void) | undefined; + expect(onData).toBeDefined(); + + return { + service, + onData: onData as (payload: unknown) => void, + setSession: (next: AgentSession) => { + session = next; + }, + }; + } + + const promptResponse = (id: number, stopReason: string) => ({ + type: "acp_message" as const, + ts: 1700000002, + message: { jsonrpc: "2.0" as const, id, result: { stopReason } }, + }); + + it("fires the turn-complete notification instead of draining when the head message is held by an edit", async () => { + const { onData, setSession } = await connectWithLiveSession(); + vi.useFakeTimers(); + try { + setSession( + createMockSession({ + currentPromptId: 42, + isPromptPending: true, + messageQueue: [{ id: "q-1", content: "held", queuedAt: 1 }], + editingQueuedId: "q-1", + }), + ); + + onData(promptResponse(42, "end_turn")); + await vi.advanceTimersByTimeAsync(20); + + expect( + mockNotificationService.notifyPromptComplete, + ).toHaveBeenCalledWith("Test Task", "end_turn", "task-123", undefined); + expect( + mockSessionStoreSetters.dequeueMessagesAsText, + ).not.toHaveBeenCalled(); + expect(mockTrpcAgent.prompt.mutate).not.toHaveBeenCalled(); + } finally { + vi.useRealTimers(); + } + }); + + it("does not drain the queue when the turn was cancelled", async () => { + const { onData, setSession } = await connectWithLiveSession(); + vi.useFakeTimers(); + try { + setSession( + createMockSession({ + currentPromptId: 42, + isPromptPending: true, + messageQueue: [{ id: "q-1", content: "queued", queuedAt: 1 }], + }), + ); + + onData(promptResponse(42, "cancelled")); + await vi.advanceTimersByTimeAsync(20); + + expect( + mockSessionStoreSetters.dequeueMessagesAsText, + ).not.toHaveBeenCalled(); + expect(mockTrpcAgent.prompt.mutate).not.toHaveBeenCalled(); + } finally { + vi.useRealTimers(); + } + }); + + it("drains queued messages one turn at a time", async () => { + const { onData, setSession } = await connectWithLiveSession(); + vi.useFakeTimers(); + try { + setSession( + createMockSession({ + currentPromptId: 42, + isPromptPending: true, + messageQueue: [ + { id: "q-1", content: "first", queuedAt: 1 }, + { id: "q-2", content: "second", queuedAt: 2 }, + ], + }), + ); + mockSessionStoreSetters.dequeueMessagesAsText + .mockReturnValueOnce("first") + .mockReturnValueOnce("second"); + mockTrpcAgent.prompt.mutate.mockResolvedValue({ + stopReason: "end_turn", + }); + + onData(promptResponse(42, "end_turn")); + await vi.advanceTimersByTimeAsync(20); + + expect(mockTrpcAgent.prompt.mutate).toHaveBeenCalledTimes(1); + expect( + mockSessionStoreSetters.dequeueMessagesAsText, + ).toHaveBeenLastCalledWith("task-123", { stopAtEdited: true, max: 1 }); + + // The sent message's turn runs and completes: its prompt echo claims a + // new id, then its response drains the next queued message. + onData({ + type: "acp_message", + ts: 1700000003, + message: { + jsonrpc: "2.0", + id: 43, + method: "session/prompt", + params: { prompt: [] }, + }, + }); + onData(promptResponse(43, "end_turn")); + await vi.advanceTimersByTimeAsync(20); + + expect(mockTrpcAgent.prompt.mutate).toHaveBeenCalledTimes(2); + expect( + mockSessionStoreSetters.dequeueMessagesAsText, + ).toHaveBeenCalledTimes(2); + } finally { + vi.useRealTimers(); + } + }); + + it("does not start a second prompt when an edit release races the turn-end drain", async () => { + const { service, onData, setSession } = await connectWithLiveSession(); + vi.useFakeTimers(); + try { + setSession( + createMockSession({ + currentPromptId: 42, + isPromptPending: true, + messageQueue: [ + { id: "q-1", content: "first", queuedAt: 1 }, + { id: "q-2", content: "second", queuedAt: 2 }, + ], + editingQueuedId: "q-2", + }), + ); + mockSessionStoreSetters.dequeueMessagesAsText + .mockReturnValueOnce("first") + .mockReturnValueOnce("second"); + // Keep the first send in flight so the raced timer must observe it. + mockTrpcAgent.prompt.mutate.mockImplementation( + () => new Promise(() => {}), + ); + + // The turn ends (the buffered event flush processes it and schedules a + // drain timer), then the user cancels the edit before that timer fires + // (scheduling a second drain via the idle flush). + onData(promptResponse(42, "end_turn")); + await vi.advanceTimersToNextTimerAsync(); + service.clearEditingQueuedMessage("task-123"); + await vi.advanceTimersByTimeAsync(20); + + expect(mockTrpcAgent.prompt.mutate).toHaveBeenCalledTimes(1); + expect( + mockSessionStoreSetters.dequeueMessagesAsText, + ).toHaveBeenCalledTimes(1); + } finally { + vi.useRealTimers(); + } + }); + }); + describe("steer echo routing", () => { async function connectAndCaptureOnData(): Promise< (payload: unknown) => void @@ -5164,6 +5371,9 @@ describe("SessionService", () => { await vi.advanceTimersByTimeAsync(0); expect(updated).toBe(true); + expect( + mockSessionStoreSetters.updateQueuedMessage, + ).toHaveBeenCalledWith("task-123", "q-1", { content: "edited" }); expect( mockSessionStoreSetters.clearEditingQueuedMessage, ).toHaveBeenCalledWith("task-123"); @@ -5178,6 +5388,27 @@ describe("SessionService", () => { } }); + it("returns false and keeps the hold when the target is no longer queued", async () => { + const service = getSessionService(); + seedEditedIdleSession({ + messageQueue: [{ id: "q-other", content: "x", queuedAt: 1 }], + }); + + const updated = await service.updateQueuedMessage( + "task-123", + "q-1", + "edited", + ); + + expect(updated).toBe(false); + expect( + mockSessionStoreSetters.updateQueuedMessage, + ).not.toHaveBeenCalled(); + expect( + mockSessionStoreSetters.clearEditingQueuedMessage, + ).not.toHaveBeenCalled(); + }); + it("saving an edit while the agent is still busy does not send immediately", async () => { vi.useFakeTimers(); try { @@ -5221,6 +5452,79 @@ describe("SessionService", () => { vi.useRealTimers(); } }); + + const seedEditedIdleCloudSession = () => { + const queuedMessage = { + id: "q-1", + content: "old", + rawPrompt: [{ type: "text" as const, text: "old" }], + queuedAt: 1, + }; + const session = createMockSession({ + isCloud: true, + cloudStatus: "in_progress", + status: "connected", + isPromptPending: false, + messageQueue: [queuedMessage], + editingQueuedId: "q-1", + }); + mockSessionStoreSetters.getSessionByTaskId.mockReturnValue(session); + mockSessionStoreSetters.getSessions.mockReturnValue({ + "run-123": session, + }); + mockSessionStoreSetters.clearEditingQueuedMessage.mockImplementation( + () => { + session.editingQueuedId = undefined; + }, + ); + mockSessionStoreSetters.dequeueMessages.mockReturnValue([queuedMessage]); + mockTrpcCloudTask.sendCommand.mutate.mockResolvedValue({ + success: true, + result: { stopReason: "end_turn" }, + }); + return session; + }; + + it("saving a cloud edit while the run is idle flushes the queue", async () => { + const service = getSessionService(); + seedEditedIdleCloudSession(); + + const updated = await service.updateQueuedMessage( + "task-123", + "q-1", + "edited", + ); + + expect(updated).toBe(true); + await vi.waitFor(() => { + expect(mockTrpcCloudTask.sendCommand.mutate).toHaveBeenCalledWith( + expect.objectContaining({ + taskId: "task-123", + method: "user_message", + }), + ); + }); + expect(mockSessionStoreSetters.dequeueMessages).toHaveBeenCalledWith( + "task-123", + { stopAtEdited: true, max: 1 }, + ); + }); + + it("cancelling a cloud edit while the run is idle flushes the queue", async () => { + const service = getSessionService(); + seedEditedIdleCloudSession(); + + service.clearEditingQueuedMessage("task-123"); + + await vi.waitFor(() => { + expect(mockTrpcCloudTask.sendCommand.mutate).toHaveBeenCalledWith( + expect.objectContaining({ + taskId: "task-123", + method: "user_message", + }), + ); + }); + }); }); describe("updateQueuedMessage cloud normalization race", () => { @@ -6069,6 +6373,37 @@ describe("SessionService", () => { expect(stored.events).toBe(previousEvents); }); + it("carries the queue and its edit hold across an in-place reconnect", async () => { + const service = getSessionService(); + const queued = [{ id: "q-1", content: "old", queuedAt: 1 }]; + const mockSession = createMockSession({ + status: "error", + logUrl: "https://logs.example.com/run-123", + messageQueue: queued, + editingQueuedId: "q-1", + }); + mockSessionStoreSetters.getSessionByTaskId.mockReturnValue(mockSession); + mockSessionStoreSetters.getSessions.mockReturnValue({ + "run-123": mockSession, + }); + mockTrpcAgent.reconnect.mutate.mockResolvedValue({ + sessionId: "run-123", + channel: "agent-event:run-123", + configOptions: [], + }); + mockTrpcWorkspace.verify.query.mockResolvedValue({ exists: true }); + mockTrpcLogs.readLocalLogs.query.mockResolvedValue(""); + mockTrpcLogs.fetchS3Logs.query.mockResolvedValue(""); + + await service.clearSessionError("task-123", "/repo"); + + // The rebuilt session must keep the hold with the queue it guards, or + // the edited message would auto-send in its stale, pre-edit form. + const stored = mockSessionStoreSetters.setSession.mock.calls.at(-1)?.[0]; + expect(stored.messageQueue).toBe(queued); + expect(stored.editingQueuedId).toBe("q-1"); + }); + it("creates fresh session when initialPrompt is set (prompt never delivered)", async () => { const service = getSessionService(); const mockSession = createMockSession({ From a4a74367a9464facd0f0bd8f3eb373cc4ba122ef Mon Sep 17 00:00:00 2001 From: Charles Vien Date: Tue, 14 Jul 2026 23:56:50 -0700 Subject: [PATCH 14/15] scope queued card drag to its grip handle --- .../components/QueuedMessagesDock.test.tsx | 143 ++++++++++++++++-- .../components/QueuedMessagesDock.tsx | 82 +++++----- .../session-update/QueuedMessageView.test.tsx | 110 ++++++++++++++ .../session-update/QueuedMessageView.tsx | 17 ++- 4 files changed, 295 insertions(+), 57 deletions(-) create mode 100644 packages/ui/src/features/sessions/components/session-update/QueuedMessageView.test.tsx diff --git a/packages/ui/src/features/sessions/components/QueuedMessagesDock.test.tsx b/packages/ui/src/features/sessions/components/QueuedMessagesDock.test.tsx index 1307c4f529..ff7e369fab 100644 --- a/packages/ui/src/features/sessions/components/QueuedMessagesDock.test.tsx +++ b/packages/ui/src/features/sessions/components/QueuedMessagesDock.test.tsx @@ -6,14 +6,30 @@ const queuedState = vi.hoisted(() => ({ messages: [] as Array<{ id: string; content: string; queuedAt: number }>, })); +const sessionState = vi.hoisted(() => ({ + editingQueuedId: undefined as string | undefined, +})); + +const sessionService = vi.hoisted(() => ({ + steerQueuedMessage: vi.fn().mockResolvedValue(undefined), + clearEditingQueuedMessage: vi.fn(), +})); + +const storeSetters = vi.hoisted(() => ({ + removeQueuedMessage: vi.fn(), + moveQueuedMessage: vi.fn(), +})); + +const dndCapture = vi.hoisted(() => ({ + onDragOver: undefined as ((event: unknown) => void) | undefined, +})); + vi.mock("@posthog/core/sessions/sessionService", () => ({ SESSION_SERVICE: Symbol.for("test.session-service"), })); vi.mock("@posthog/di/react", () => ({ - useService: () => ({ - steerQueuedMessage: vi.fn().mockResolvedValue(undefined), - }), + useService: () => sessionService, })); vi.mock("@posthog/ui/features/sessions/useSession", () => ({ @@ -30,17 +46,20 @@ vi.mock("@posthog/ui/features/sessions/hooks/useEditQueuedMessage", () => ({ })); vi.mock("@posthog/ui/features/sessions/sessionStore", () => ({ - sessionStoreSetters: { - removeQueuedMessage: vi.fn(), - moveQueuedMessage: vi.fn(), - }, + sessionStoreSetters: storeSetters, useSessionIsCloud: () => false, useSessionSelector: ( _taskId: string, - select: (session: undefined) => T, - ) => select(undefined), + select: (session: { editingQueuedId?: string }) => T, + ) => select({ editingQueuedId: sessionState.editingQueuedId }), useSessionStore: { - getState: () => ({ taskIdIndex: {}, sessions: {} }), + getState: () => ({ + taskIdIndex: new Proxy({}, { get: () => "run-1" }) as Record< + string, + string + >, + sessions: { "run-1": { messageQueue: queuedState.messages } }, + }), }, })); @@ -48,6 +67,32 @@ vi.mock("@posthog/ui/primitives/toast", () => ({ toast: { error: vi.fn() }, })); +// The dock's reorder handler is driven directly through the captured +// onDragOver prop; the sortable plumbing itself is @dnd-kit's to test. +vi.mock("@dnd-kit/react", async () => { + const React = await import("react"); + return { + DragDropProvider: ({ + onDragOver, + children, + }: { + onDragOver: (event: unknown) => void; + children: React.ReactNode; + }) => { + dndCapture.onDragOver = onDragOver; + return React.createElement(React.Fragment, null, children); + }, + }; +}); +vi.mock("@dnd-kit/react/sortable", () => ({ + useSortable: () => ({ + ref: () => {}, + handleRef: () => {}, + isDragging: false, + }), +})); +vi.mock("@dnd-kit/dom", () => ({ PointerSensor: class {} })); + // Stub the per-message card so the test exercises the dock's collapse/scroll // shell, not the markdown/steer internals it already owns. vi.mock( @@ -55,10 +100,19 @@ vi.mock( async () => { const React = await import("react"); return { - QueuedMessageView: ({ message }: { message: { content: string } }) => + QueuedMessageView: ({ + message, + isEditing, + }: { + message: { content: string }; + isEditing?: boolean; + }) => React.createElement( "div", - { "data-testid": "queued-card" }, + { + "data-testid": "queued-card", + "data-editing": String(!!isEditing), + }, message.content, ), }; @@ -84,7 +138,10 @@ function renderDock(taskId: string) { describe("QueuedMessagesDock", () => { beforeEach(() => { + vi.clearAllMocks(); queuedState.messages = []; + sessionState.editingQueuedId = undefined; + dndCapture.onDragOver = undefined; }); it("renders nothing when the queue is empty", () => { @@ -135,4 +192,66 @@ describe("QueuedMessagesDock", () => { fireEvent.click(trigger); expect(screen.getAllByTestId("queued-card")).toHaveLength(2); }); + + it("reorders the queue when a card is dragged over another", () => { + queuedState.messages = TWO_MESSAGES; + renderDock("task-drag"); + + dndCapture.onDragOver?.({ + operation: { source: { id: "q1" }, target: { id: "q2" } }, + }); + + expect(storeSetters.moveQueuedMessage).toHaveBeenCalledWith( + "task-drag", + 0, + 1, + ); + }); + + it.each([ + { + name: "source and target are the same card", + operation: { source: { id: "q1" }, target: { id: "q1" } }, + }, + { + name: "source is not in the queue", + operation: { source: { id: "missing" }, target: { id: "q2" } }, + }, + { + name: "target is not in the queue", + operation: { source: { id: "q1" }, target: { id: "missing" } }, + }, + { + name: "there is no drop target", + operation: { source: { id: "q1" }, target: undefined }, + }, + ])("does not reorder when $name", ({ operation }) => { + queuedState.messages = TWO_MESSAGES; + renderDock("task-drag-noop"); + + dndCapture.onDragOver?.({ operation }); + + expect(storeSetters.moveQueuedMessage).not.toHaveBeenCalled(); + }); + + it("marks only the edited message's card as editing", () => { + queuedState.messages = TWO_MESSAGES; + sessionState.editingQueuedId = "q1"; + renderDock("task-editing"); + + const [first, second] = screen.getAllByTestId("queued-card"); + expect(first).toHaveAttribute("data-editing", "true"); + expect(second).toHaveAttribute("data-editing", "false"); + expect(sessionService.clearEditingQueuedMessage).not.toHaveBeenCalled(); + }); + + it("clears a stale edit hold when the edited message leaves the queue", () => { + queuedState.messages = TWO_MESSAGES; + sessionState.editingQueuedId = "q-gone"; + renderDock("task-stale-hold"); + + expect(sessionService.clearEditingQueuedMessage).toHaveBeenCalledWith( + "task-stale-hold", + ); + }); }); diff --git a/packages/ui/src/features/sessions/components/QueuedMessagesDock.tsx b/packages/ui/src/features/sessions/components/QueuedMessagesDock.tsx index 46cc7602ac..1642d6d855 100644 --- a/packages/ui/src/features/sessions/components/QueuedMessagesDock.tsx +++ b/packages/ui/src/features/sessions/components/QueuedMessagesDock.tsx @@ -27,16 +27,21 @@ import { useQueuedMessagesForTask } from "@posthog/ui/features/sessions/useSessi import { toast } from "@posthog/ui/primitives/toast"; import * as Collapsible from "@radix-ui/react-collapsible"; import { Box, Flex, Text } from "@radix-ui/themes"; -import { type ReactNode, useCallback, useEffect } from "react"; +import { + type ReactNode, + type RefCallback, + useCallback, + useEffect, +} from "react"; interface QueuedMessagesDockProps { taskId: string; } /** - * A single queued card, wrapped so the whole card is a drag handle for - * reordering the queue. Drag activation waits for a small pointer move (see the - * provider's sensor) so the card's buttons still take clicks. + * A single queued card wrapped as a sortable item. Dragging is scoped to the + * card's grip button (the handle ref passed to `children`), so the card's own + * buttons never compete with a drag. */ function SortableQueuedMessage({ id, @@ -47,9 +52,9 @@ function SortableQueuedMessage({ id: string; index: number; taskId: string; - children: ReactNode; + children: (dragHandleRef: RefCallback) => ReactNode; }) { - const { ref, isDragging } = useSortable({ + const { ref, handleRef, isDragging } = useSortable({ id, index, group: `queue:${taskId}`, @@ -57,14 +62,8 @@ function SortableQueuedMessage({ }); return ( -
- {children} +
+ {children(handleRef as RefCallback)}
); } @@ -178,32 +177,35 @@ export function QueuedMessagesDock({ taskId }: QueuedMessagesDockProps) { index={index} taskId={taskId} > - { - void sessionService - .steerQueuedMessage(taskId, message.id) - .catch(() => { - toast.error( - "Couldn't steer this message. It's still queued.", - ); - }); - } - : undefined - } - onEdit={() => editMessage(message)} - onCancelEdit={cancelEdit} - onRemove={() => - sessionStoreSetters.removeQueuedMessage( - taskId, - message.id, - ) - } - /> + {(dragHandleRef) => ( + { + void sessionService + .steerQueuedMessage(taskId, message.id) + .catch(() => { + toast.error( + "Couldn't steer this message. It's still queued.", + ); + }); + } + : undefined + } + onEdit={() => editMessage(message)} + onCancelEdit={cancelEdit} + onRemove={() => + sessionStoreSetters.removeQueuedMessage( + taskId, + message.id, + ) + } + /> + )} ))} diff --git a/packages/ui/src/features/sessions/components/session-update/QueuedMessageView.test.tsx b/packages/ui/src/features/sessions/components/session-update/QueuedMessageView.test.tsx new file mode 100644 index 0000000000..34282c3fa0 --- /dev/null +++ b/packages/ui/src/features/sessions/components/session-update/QueuedMessageView.test.tsx @@ -0,0 +1,110 @@ +import { Theme } from "@radix-ui/themes"; +import { fireEvent, render, screen } from "@testing-library/react"; +import { describe, expect, it, vi } from "vitest"; +import type { QueuedMessage } from "../../sessionStore"; +import { QueuedMessageView } from "./QueuedMessageView"; + +const MESSAGE: QueuedMessage = { + id: "q-1", + content: "queued body", + queuedAt: 1, +}; + +function renderView( + props: Partial[0]> = {}, +) { + const handlers = { + onSteer: vi.fn(), + onEdit: vi.fn(), + onCancelEdit: vi.fn(), + onRemove: vi.fn(), + }; + render( + + + , + ); + return handlers; +} + +describe("QueuedMessageView", () => { + it.each([ + { + state: "queued", + isEditing: false, + visible: ["Steer this message", "Edit queued message"], + hidden: ["Cancel edit"], + }, + { + state: "editing", + isEditing: true, + visible: ["Cancel edit"], + hidden: ["Steer this message", "Edit queued message"], + }, + ])( + "shows the $state action set even when every handler is provided", + ({ isEditing, visible, hidden }) => { + renderView({ isEditing }); + + expect(screen.getByText("queued body")).toBeInTheDocument(); + for (const name of visible) { + expect(screen.getByRole("button", { name })).toBeInTheDocument(); + } + for (const name of hidden) { + expect(screen.queryByRole("button", { name })).not.toBeInTheDocument(); + } + expect( + screen.getByRole("button", { name: "Discard queued message" }), + ).toBeInTheDocument(); + }, + ); + + it("shows the editing hint while the message is open in the composer", () => { + renderView({ isEditing: true }); + expect(screen.getByText("Editing in composer")).toBeInTheDocument(); + }); + + it.each([ + { name: "Steer this message", handler: "onSteer" as const }, + { name: "Edit queued message", handler: "onEdit" as const }, + { name: "Discard queued message", handler: "onRemove" as const }, + ])( + "omits the $name button when $handler is not provided", + ({ name, handler }) => { + renderView({ [handler]: undefined }); + expect(screen.queryByRole("button", { name })).not.toBeInTheDocument(); + }, + ); + + it("omits the cancel button when onCancelEdit is not provided", () => { + renderView({ isEditing: true, onCancelEdit: undefined }); + expect( + screen.queryByRole("button", { name: "Cancel edit" }), + ).not.toBeInTheDocument(); + }); + + it("wires each visible action to its handler", () => { + const handlers = renderView(); + + fireEvent.click(screen.getByRole("button", { name: "Steer this message" })); + fireEvent.click( + screen.getByRole("button", { name: "Edit queued message" }), + ); + fireEvent.click( + screen.getByRole("button", { name: "Discard queued message" }), + ); + + expect(handlers.onSteer).toHaveBeenCalledTimes(1); + expect(handlers.onEdit).toHaveBeenCalledTimes(1); + expect(handlers.onRemove).toHaveBeenCalledTimes(1); + }); + + it("forwards the drag handle ref to the grip button", () => { + const dragHandleRef = vi.fn(); + renderView({ dragHandleRef }); + + const grip = screen.getByRole("button", { name: "Drag to reorder" }); + expect(grip).toBeInTheDocument(); + expect(dragHandleRef).toHaveBeenCalledWith(grip); + }); +}); diff --git a/packages/ui/src/features/sessions/components/session-update/QueuedMessageView.tsx b/packages/ui/src/features/sessions/components/session-update/QueuedMessageView.tsx index 9c88179e87..7096320c86 100644 --- a/packages/ui/src/features/sessions/components/session-update/QueuedMessageView.tsx +++ b/packages/ui/src/features/sessions/components/session-update/QueuedMessageView.tsx @@ -8,6 +8,7 @@ import { import { Button } from "@posthog/quill"; import { Box, Flex, IconButton, Text, Tooltip } from "@radix-ui/themes"; import clsx from "clsx"; +import type { Ref } from "react"; import { MarkdownRenderer } from "../../../editor/components/MarkdownRenderer"; import type { QueuedMessage } from "../../sessionStore"; import { CollapsibleMessageContent } from "./CollapsibleMessageContent"; @@ -15,6 +16,7 @@ import { hasFileMentions, parseFileMentions } from "./parseFileMentions"; interface QueuedMessageViewProps { message: QueuedMessage; + dragHandleRef?: Ref; onSteer?: () => void; onEdit?: () => void; onCancelEdit?: () => void; @@ -25,6 +27,7 @@ interface QueuedMessageViewProps { export function QueuedMessageView({ message, + dragHandleRef, onSteer, onEdit, onCancelEdit, @@ -49,11 +52,15 @@ export function QueuedMessageView({ Steer button (fixed height) anchors it, but the editing state's ghost icon buttons are `fit-content` and would otherwise collapse shorter. */} - + Date: Tue, 14 Jul 2026 23:56:59 -0700 Subject: [PATCH 15/15] add cloud coverage for queued edit loading --- .../hooks/useEditQueuedMessage.test.tsx | 40 +++++++++++++++++++ 1 file changed, 40 insertions(+) diff --git a/packages/ui/src/features/sessions/hooks/useEditQueuedMessage.test.tsx b/packages/ui/src/features/sessions/hooks/useEditQueuedMessage.test.tsx index 5574e9bc72..3675fb1989 100644 --- a/packages/ui/src/features/sessions/hooks/useEditQueuedMessage.test.tsx +++ b/packages/ui/src/features/sessions/hooks/useEditQueuedMessage.test.tsx @@ -120,6 +120,46 @@ describe("useEditQueuedMessage / useCancelQueuedMessageEdit", () => { expect(pendingFor(TASK)).toEqual({ segments: [] }); }); + it.each([ + { variant: "local", isCloud: false, message: QUEUED }, + { + variant: "cloud", + isCloud: true, + message: { + id: "q-1", + content: "queued body", + rawPrompt: [{ type: "text" as const, text: "queued body" }], + queuedAt: 1, + }, + }, + ])( + "loads a $variant queued message into the composer and marks it as the edit target", + ({ isCloud, message }) => { + cloudState.isCloud = isCloud; + + const { result: edit } = renderHook(() => useEditQueuedMessage(TASK)); + act(() => edit.current(message)); + + expect(sessionService.setEditingQueuedMessage).toHaveBeenCalledWith( + TASK, + "q-1", + ); + expect(pendingFor(TASK)?.segments).toEqual([ + { type: "text", text: "queued body" }, + ]); + }, + ); + + it("does not start an edit when a cloud message has no loadable content", () => { + cloudState.isCloud = true; + + const { result: edit } = renderHook(() => useEditQueuedMessage(TASK)); + act(() => edit.current({ id: "q-1", content: "", queuedAt: 1 })); + + expect(sessionService.setEditingQueuedMessage).not.toHaveBeenCalled(); + expect(pendingFor(TASK)).toBeUndefined(); + }); + it("snapshots fresh each edit, so a stale draft is not restored later", () => { const { result: edit } = renderHook(() => useEditQueuedMessage(TASK)); const { result: cancel } = renderHook(() =>