From 1d8a138d0d2bd522ece6660806cda8ec80abd736 Mon Sep 17 00:00:00 2001 From: thevortex8 Date: Sat, 21 Mar 2026 16:49:42 -0400 Subject: [PATCH] fix: preserve queued message ids on replay --- packages/app/src/contexts/session-context.tsx | 35 +++++++++---- .../session-queued-message-replay.test.ts | 50 +++++++++++++++++++ .../contexts/session-queued-message-replay.ts | 30 +++++++++++ 3 files changed, 105 insertions(+), 10 deletions(-) create mode 100644 packages/app/src/contexts/session-queued-message-replay.test.ts create mode 100644 packages/app/src/contexts/session-queued-message-replay.ts diff --git a/packages/app/src/contexts/session-context.tsx b/packages/app/src/contexts/session-context.tsx index 86e711ce0f..1708c2b615 100644 --- a/packages/app/src/contexts/session-context.tsx +++ b/packages/app/src/contexts/session-context.tsx @@ -49,6 +49,7 @@ import { derivePendingPermissionKey, normalizeAgentSnapshot } from "@/utils/agen import { resolveProjectPlacement } from "@/utils/project-placement"; import { buildDraftStoreKey } from "@/stores/draft-keys"; import type { AttachmentMetadata } from "@/attachments/types"; +import { takeQueuedAgentMessageReplay } from "@/contexts/session-queued-message-replay"; // Re-export types from session-store and draft-store for backward compatibility export type { DraftInput } from "@/stores/draft-store"; @@ -272,7 +273,12 @@ function SessionProviderInternal({ children, serverId, client }: SessionProvider const previousAgentStatusRef = useRef>(new Map()); const sendAgentMessageRef = useRef< - ((agentId: string, message: string, images?: AttachmentMetadata[]) => Promise) | null + (( + agentId: string, + message: string, + images?: AttachmentMetadata[], + messageId?: string, + ) => Promise) | null >(null); const sessionStateTimeoutRef = useRef | null>(null); const attentionNotifiedRef = useRef>(new Map()); @@ -425,15 +431,19 @@ function SessionProviderInternal({ children, serverId, client }: SessionProvider const prevStatus = previousAgentStatusRef.current.get(agent.id); if (prevStatus === "running" && agent.status !== "running") { const session = useSessionStore.getState().sessions[serverId]; - const queue = session?.queuedMessages.get(agent.id); - if (queue && queue.length > 0) { - const [next, ...rest] = queue; + const replay = takeQueuedAgentMessageReplay(session?.queuedMessages.get(agent.id)); + if (replay) { if (sendAgentMessageRef.current) { - void sendAgentMessageRef.current(agent.id, next.text, next.images); + void sendAgentMessageRef.current( + agent.id, + replay.text, + replay.images, + replay.messageId, + ); } setQueuedMessages(serverId, (prev) => { const updated = new Map(prev); - updated.set(agent.id, rest); + updated.set(agent.id, replay.remainingQueue); return updated; }); } @@ -1469,11 +1479,16 @@ function SessionProviderInternal({ children, serverId, client }: SessionProvider ]); const sendAgentMessage = useCallback( - async (agentId: string, message: string, images?: AttachmentMetadata[]) => { - const messageId = generateMessageId(); + async ( + agentId: string, + message: string, + images?: AttachmentMetadata[], + messageId?: string, + ) => { + const resolvedMessageId = messageId ?? generateMessageId(); const userMessage: StreamItem = { kind: "user_message", - id: messageId, + id: resolvedMessageId, text: message, timestamp: new Date(), }; @@ -1507,7 +1522,7 @@ function SessionProviderInternal({ children, serverId, client }: SessionProvider } void client .sendAgentMessage(agentId, message, { - messageId, + messageId: resolvedMessageId, ...(imagesData && imagesData.length > 0 ? { images: imagesData } : {}), }) .catch((error) => { diff --git a/packages/app/src/contexts/session-queued-message-replay.test.ts b/packages/app/src/contexts/session-queued-message-replay.test.ts new file mode 100644 index 0000000000..67186947b4 --- /dev/null +++ b/packages/app/src/contexts/session-queued-message-replay.test.ts @@ -0,0 +1,50 @@ +import { describe, expect, it } from "vitest"; +import { takeQueuedAgentMessageReplay } from "./session-queued-message-replay"; + +describe("takeQueuedAgentMessageReplay", () => { + it("preserves the queued message id for idempotent replay", () => { + const replay = takeQueuedAgentMessageReplay([ + { + id: "msg-1", + text: "resume this", + images: [ + { + id: "img-1", + mimeType: "image/png", + storageType: "web-indexeddb", + storageKey: "attachment:img-1", + fileName: "test.png", + byteSize: 12, + createdAt: 1, + }, + ], + }, + { + id: "msg-2", + text: "second", + }, + ]); + + expect(replay).toEqual({ + messageId: "msg-1", + text: "resume this", + images: [ + { + id: "img-1", + mimeType: "image/png", + storageType: "web-indexeddb", + storageKey: "attachment:img-1", + fileName: "test.png", + byteSize: 12, + createdAt: 1, + }, + ], + remainingQueue: [{ id: "msg-2", text: "second" }], + }); + }); + + it("returns null when the queue is empty", () => { + expect(takeQueuedAgentMessageReplay([])).toBeNull(); + expect(takeQueuedAgentMessageReplay(undefined)).toBeNull(); + }); +}); diff --git a/packages/app/src/contexts/session-queued-message-replay.ts b/packages/app/src/contexts/session-queued-message-replay.ts new file mode 100644 index 0000000000..7b361ed90f --- /dev/null +++ b/packages/app/src/contexts/session-queued-message-replay.ts @@ -0,0 +1,30 @@ +import type { AttachmentMetadata } from "@/attachments/types"; + +export interface QueuedAgentMessageReplayItem { + id: string; + text: string; + images?: AttachmentMetadata[]; +} + +export interface QueuedAgentMessageReplay { + messageId: string; + text: string; + images?: AttachmentMetadata[]; + remainingQueue: QueuedAgentMessageReplayItem[]; +} + +export function takeQueuedAgentMessageReplay( + queue: readonly QueuedAgentMessageReplayItem[] | undefined, +): QueuedAgentMessageReplay | null { + if (!queue || queue.length === 0) { + return null; + } + + const [next, ...remainingQueue] = queue; + return { + messageId: next.id, + text: next.text, + images: next.images, + remainingQueue, + }; +}