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 }, 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/core/src/sessions/sessionQueueEditing.test.ts b/packages/core/src/sessions/sessionQueueEditing.test.ts new file mode 100644 index 0000000000..9c0f178c8a --- /dev/null +++ b/packages/core/src/sessions/sessionQueueEditing.test.ts @@ -0,0 +1,242 @@ +import { + type AgentSession, + type QueuedMessage, + sendableQueuePrefixLength, +} 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"); + }); +}); + +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([]); + }); +}); + +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 58b611dccd..16d6e1d647 100644 --- a/packages/core/src/sessions/sessionService.ts +++ b/packages/core/src/sessions/sessionService.ts @@ -29,6 +29,7 @@ import { type QueuedMessage, resolveBypassRevertMode, type StoredLogEntry, + sendableQueuePrefixLength, sessionSupportsNativeSteer, type TaskRunStatus, } from "@posthog/shared"; @@ -224,9 +225,22 @@ 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; + 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; max?: number }, + ): string | null; + dequeueMessages( + taskId: string, + options?: { stopAtEdited?: boolean; max?: number }, + ): QueuedMessage[]; prependQueuedMessages(taskId: string, messages: QueuedMessage[]): void; appendOptimisticItem( taskRunId: string, @@ -989,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; @@ -2009,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, @@ -2111,13 +2133,30 @@ 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(() => { + // 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, @@ -2127,7 +2166,7 @@ export class SessionService { }, 0); } - return hasQueuedMessages; + return hasSendableMessages; } private handlePermissionRequest( @@ -2399,15 +2438,20 @@ 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); + const combinedText = this.d.store.dequeueMessagesAsText(taskId, { + stopAtEdited: true, + max: 1, + }); if (!combinedText) { return { stopReason: "skipped" }; } @@ -2421,7 +2465,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, }); @@ -2662,6 +2706,104 @@ 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 + * 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. + * + * 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, + 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); + // 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, + rawPrompt: normalizedPrompt, + }); + } else { + this.d.store.updateQueuedMessage(taskId, messageId, { + content: extractPromptText(prompt), + }); + } + + // 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); + 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. */ @@ -2962,11 +3104,17 @@ export class SessionService { const authStatus = await this.getAuthCredentialsStatus(); if (authStatus.kind === "restoring") return; - const drained = this.d.store.dequeueMessages(taskId); + // 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 0e86a78c6a..f8300eb8af 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"; @@ -32,6 +33,55 @@ 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; +} + +/** + * 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) => { @@ -197,53 +247,101 @@ export const sessionStoreSetters = { }); }, - dequeueMessagesAsText: (taskId: string): 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 combined = session.messageQueue - .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 = []; + /** + * 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); }); - return combined; }, - dequeueMessages: (taskId: string): 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 queuedMessages = [...session.messageQueue]; - - sessionStore.setState((draft) => { - const trid = draft.taskIdIndex[taskId]; - if (!trid) return; - const draftSession = draft.sessions[trid]; - if (draftSession) { - draftSession.messageQueue = []; - } + /** + * 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; + }); + }, + + /** + * 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; }); + }, - return queuedMessages; + /** + * 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; max?: number }, + ): string | null => { + const drained = drainQueueHead(taskId, options); + if (drained.length === 0) return null; + return drained.map((msg) => msg.content).join("\n\n"); }, + /** + * 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; max?: number }, + ): QueuedMessage[] => drainQueueHead(taskId, options), + /** * Splice messages back at the head of the queue. Used to roll back a * dispatch attempt that drained the queue but failed before delivery. diff --git a/packages/shared/src/index.ts b/packages/shared/src/index.ts index 87b891dfe1..b03a71c46c 100644 --- a/packages/shared/src/index.ts +++ b/packages/shared/src/index.ts @@ -208,6 +208,7 @@ export { type QueuedMessage, resolveBypassRevertMode, type SessionStatus, + sendableQueuePrefixLength, sessionSupportsNativeSteer, } from "./sessions"; export type { 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/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/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/components/QueuedMessagesDock.test.tsx b/packages/ui/src/features/sessions/components/QueuedMessagesDock.test.tsx index 0813f010dd..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", () => ({ @@ -24,26 +40,59 @@ 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: storeSetters, useSessionIsCloud: () => false, useSessionSelector: ( _taskId: string, - select: (session: undefined) => T, - ) => select(undefined), + select: (session: { editingQueuedId?: string }) => T, + ) => select({ editingQueuedId: sessionState.editingQueuedId }), + useSessionStore: { + getState: () => ({ + taskIdIndex: new Proxy({}, { get: () => "run-1" }) as Record< + string, + string + >, + sessions: { "run-1": { messageQueue: queuedState.messages } }, + }), + }, })); 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( @@ -51,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, ), }; @@ -80,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", () => { @@ -131,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 ac463bf8cf..1642d6d855 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,12 +8,16 @@ 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 { useQueueCollapsed, @@ -20,15 +27,52 @@ 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, + type RefCallback, + useCallback, + useEffect, +} from "react"; interface QueuedMessagesDockProps { taskId: string; } +/** + * 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, + index, + taskId, + children, +}: { + id: string; + index: number; + taskId: string; + children: (dragHandleRef: RefCallback) => ReactNode; +}) { + const { ref, handleRef, isDragging } = useSortable({ + id, + index, + group: `queue:${taskId}`, + transition: { duration: 200, easing: "ease" }, + }); + + return ( +
+ {children(handleRef as RefCallback)} +
+ ); +} + /** * 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 +81,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 = 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( @@ -52,6 +98,35 @@ export function QueuedMessagesDock({ taskId }: QueuedMessagesDockProps) { const collapsed = useQueueCollapsed(taskId); const { setQueueCollapsed } = useSessionViewActions(); + // 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)) { + sessionService.clearEditingQueuedMessage(taskId); + } + }, [editingId, queued, taskId, sessionService]); + + 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; const isOpen = !collapsed; @@ -83,32 +158,58 @@ 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.", - ); - }); + + + {queued.map((message, index) => ( + + {(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, + ) } - : undefined - } - onReturnToEditor={() => returnToEditor(message)} - onRemove={() => - sessionStoreSetters.removeQueuedMessage(taskId, message.id) - } - /> - ))} - + /> + )} + + ))} + + 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} /> 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 5b0f600e1f..7096320c86 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,14 @@ 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 type { Ref } from "react"; import { MarkdownRenderer } from "../../../editor/components/MarkdownRenderer"; import type { QueuedMessage } from "../../sessionStore"; import { CollapsibleMessageContent } from "./CollapsibleMessageContent"; @@ -13,17 +16,23 @@ import { hasFileMentions, parseFileMentions } from "./parseFileMentions"; interface QueuedMessageViewProps { message: QueuedMessage; + dragHandleRef?: Ref; onSteer?: () => void; - onReturnToEditor?: () => void; + onEdit?: () => void; + onCancelEdit?: () => void; onRemove?: () => void; + isEditing?: boolean; supportsNativeSteer?: boolean; } export function QueuedMessageView({ message, + dragHandleRef, onSteer, - onReturnToEditor, + onEdit, + onCancelEdit, onRemove, + isEditing = false, supportsNativeSteer = false, }: QueuedMessageViewProps) { const steerTooltip = supportsNativeSteer @@ -31,9 +40,27 @@ export function QueuedMessageView({ : "Interrupt the current turn and resend with this message."; return ( - - - + + {/* 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. */} + + )} - - {onSteer && ( - - - - )} - {onReturnToEditor && ( - - - - - + + {isEditing ? ( + <> + + Editing in composer + + {onCancelEdit && ( + + + + + + )} + + ) : ( + <> + {onSteer && ( + + + + )} + {onEdit && ( + + + + + + )} + )} {onRemove && ( 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..3675fb1989 --- /dev/null +++ b/packages/ui/src/features/sessions/hooks/useEditQueuedMessage.test.tsx @@ -0,0 +1,185 @@ +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.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(() => + 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 new file mode 100644 index 0000000000..7eab3c95cd --- /dev/null +++ b/packages/ui/src/features/sessions/hooks/useEditQueuedMessage.ts @@ -0,0 +1,114 @@ +import { + type EditorContent, + isContentEmpty, + textToContent, + xmlToContent, +} from "@posthog/core/message-editor/content"; +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 { useCallback } from "react"; + +/** + * Empty editor content, used to clear the composer when a cancelled edit has no + * prior draft to restore. + */ +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. + * + * 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). + */ +export function useEditQueuedMessage( + taskId: string | undefined, +): (message: QueuedMessage) => void { + const { requestFocus, setPendingContent, setPreEditDraft, getDraft } = + useDraftStore((s) => s.actions); + const sessionService = useService(SESSION_SERVICE); + 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; + + // 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" + ? textToContent(priorDraft) + : priorDraft, + ); + + sessionService.setEditingQueuedMessage(taskId, message.id); + setPendingContent(taskId, pendingContent); + requestFocus(taskId); + }, + [ + 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 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, takePreEditDraft } = useDraftStore( + (s) => s.actions, + ); + const sessionService = useService(SESSION_SERVICE); + + return useCallback(() => { + if (!taskId) return; + sessionService.clearEditingQueuedMessage(taskId); + setPendingContent(taskId, takePreEditDraft(taskId) ?? EMPTY_CONTENT); + }, [taskId, sessionService, setPendingContent, takePreEditDraft]); +} 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.test.tsx b/packages/ui/src/features/sessions/hooks/useSessionCallbacks.test.tsx new file mode 100644 index 0000000000..394156e47d --- /dev/null +++ b/packages/ui/src/features/sessions/hooks/useSessionCallbacks.test.tsx @@ -0,0 +1,226 @@ +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(), + cancelPrompt: 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 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: () => sessionState, + dequeueMessages, + }, +})); + +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 while editing a queued message", () => { + beforeEach(() => { + vi.clearAllMocks(); + sessionState.editingQueuedId = "q-1"; + sessionState.messageQueue = []; + 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, + }); + }); + + 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 e0ce20a63d..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 { @@ -94,6 +95,42 @@ 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 = + sessionStoreSetters.getSessionByTaskId(taskId)?.editingQueuedId; + if (editingId) { + try { + const updated = await sessionService.updateQueuedMessage( + taskId, + editingId, + promptText ?? text, + ); + if (updated) { + 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); + // 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; + } + } + try { markAsViewed(taskId); markActivity(taskId); @@ -136,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 }); @@ -147,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 5d3f9022fb..7d89f2403e 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( @@ -4100,6 +4103,7 @@ describe("SessionService", () => { }); expect(mockSessionStoreSetters.dequeueMessages).toHaveBeenCalledWith( "task-123", + { stopAtEdited: true, max: 1 }, ); } finally { vi.useRealTimers(); @@ -4896,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 @@ -5117,6 +5328,274 @@ 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.updateQueuedMessage, + ).toHaveBeenCalledWith("task-123", "q-1", { content: "edited" }); + expect( + mockSessionStoreSetters.clearEditingQueuedMessage, + ).toHaveBeenCalledWith("task-123"); + expect( + mockSessionStoreSetters.dequeueMessagesAsText, + ).toHaveBeenCalledWith("task-123", { stopAtEdited: true, max: 1 }); + expect(mockTrpcAgent.prompt.mutate).toHaveBeenCalledWith( + expect.objectContaining({ sessionId: "run-123" }), + ); + } finally { + vi.useRealTimers(); + } + }); + + 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 { + 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, max: 1 }); + expect(mockTrpcAgent.prompt.mutate).toHaveBeenCalled(); + } finally { + 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", () => { + 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(); @@ -5894,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({