From 63d746c5b4f5a816a01f18f1d543c81b34460b2c Mon Sep 17 00:00:00 2001 From: fdkevin <25277753+fdkevin0@users.noreply.github.com> Date: Sat, 7 Nov 2026 17:39:57 +0800 Subject: [PATCH] fix(app,client): preserve clientMessageId across queue drains and timeout retries --- packages/app/src/composer/actions.test.ts | 68 ++++++++++++++++- packages/app/src/composer/actions.ts | 15 +++- packages/app/src/composer/index.tsx | 74 +++++++++++++------ packages/app/src/runtime/host-runtime.ts | 3 +- packages/app/src/stores/session-store.test.ts | 39 ++++++++++ packages/app/src/stores/session-store.ts | 47 +++++++++++- packages/client/src/daemon-client.test.ts | 7 +- packages/client/src/daemon-client.ts | 10 ++- 8 files changed, 230 insertions(+), 33 deletions(-) diff --git a/packages/app/src/composer/actions.test.ts b/packages/app/src/composer/actions.test.ts index 697f848a09..1c8d2f1d96 100644 --- a/packages/app/src/composer/actions.test.ts +++ b/packages/app/src/composer/actions.test.ts @@ -806,12 +806,16 @@ describe("sendQueuedComposerMessageNow", () => { expect(submitted).toEqual([]); }); - it("removes the queued entry and submits its text + attachments", async () => { + it("removes the queued entry and submits its messageId + text + attachments", async () => { const review = reviewWorkspaceAttachment("Queued for send."); const queue = createFakeQueue( new Map([["agent", [{ id: "msg-1", text: "send me", attachments: [review] }]]]), ); - const submitted: Array<{ text: string; attachments: ComposerAttachment[] }> = []; + const submitted: Array<{ + messageId: string; + text: string; + attachments: ComposerAttachment[]; + }> = []; const result = await sendQueuedComposerMessageNow({ agentId: "agent", messageId: "msg-1", @@ -822,7 +826,65 @@ describe("sendQueuedComposerMessageNow", () => { }); expect(result).toEqual({ status: "submitted" }); expect(queue.state.get("agent")).toEqual([]); - expect(submitted).toEqual([{ text: "send me", attachments: [review] }]); + expect(submitted).toEqual([{ messageId: "msg-1", text: "send me", attachments: [review] }]); + }); + + it("preserves the queued messageId as clientMessageId across failed drain retries", async () => { + const queue = createFakeQueue( + new Map([["agent", [{ id: "msg-queued-1", text: "retry me", attachments: [] }]]]), + ); + const stream = createFakeStream(); + const sentIds: string[] = []; + let shouldFail = true; + const client: ComposerSendClient = { + sendAgentMessage: async (_agentId, _text, options) => { + sentIds.push(options.messageId); + if (shouldFail) { + shouldFail = false; + throw new Error("Timeout waiting for message (60000ms)"); + } + }, + uploadFile: async () => ({ requestId: "req", file: null, error: "unused" }), + }; + + const submitMessage = ({ + messageId, + text, + attachments, + }: { + messageId: string; + text: string; + attachments: ComposerAttachment[]; + }) => + dispatchComposerAgentMessage({ + client, + agentId: "agent", + clientMessageId: messageId, + text, + attachments, + encodeImages: passthroughEncodeImages, + submission: stream, + }); + + const first = await sendQueuedComposerMessageNow({ + agentId: "agent", + messageId: "msg-queued-1", + queue, + submitMessage, + }); + expect(first).toEqual({ + status: "failed", + errorMessage: "Timeout waiting for message (60000ms)", + }); + + const second = await sendQueuedComposerMessageNow({ + agentId: "agent", + messageId: "msg-queued-1", + queue, + submitMessage, + }); + expect(second).toEqual({ status: "submitted" }); + expect(sentIds).toEqual(["msg-queued-1", "msg-queued-1"]); }); it("restores the queued entry to the front and surfaces the error message on failure", async () => { diff --git a/packages/app/src/composer/actions.ts b/packages/app/src/composer/actions.ts index f810ef19c6..2599fdde96 100644 --- a/packages/app/src/composer/actions.ts +++ b/packages/app/src/composer/actions.ts @@ -188,6 +188,7 @@ export function cancelComposerAgent(input: CancelComposerAgentInput): Promise Promise; + submitMessage: (input: { + messageId: string; + text: string; + attachments: ComposerAttachment[]; + }) => Promise; failedToSendMessage?: string; } @@ -318,7 +323,11 @@ export async function sendQueuedComposerMessageNow( return next; }); try { - await input.submitMessage({ text: item.text, attachments: item.attachments }); + await input.submitMessage({ + messageId: item.id, + text: item.text, + attachments: item.attachments, + }); return { status: "submitted" }; } catch (error) { input.queue.write((prev) => { diff --git a/packages/app/src/composer/index.tsx b/packages/app/src/composer/index.tsx index 60b05be23e..b0bb35a4f0 100644 --- a/packages/app/src/composer/index.tsx +++ b/packages/app/src/composer/index.tsx @@ -82,6 +82,7 @@ import { type MessageInputRef, } from "./input/input"; import type { ImageAttachment, MessagePayload, TextReplacement } from "./types"; +import { generateMessageId } from "@/types/stream"; import { ICON_SIZE, type Theme } from "@/styles/theme"; import type { DraftCommandTarget } from "@/agent-commands"; import { encodeImages } from "@/utils/encode-images"; @@ -1487,12 +1488,18 @@ function ComposerContentImpl({ const { pickImages } = useImageAttachmentPicker(); const { pickFiles } = useFilePicker(); const agentIdRef = useRef(agentId); + const failedDirectSendRetryRef = useRef<{ + agentId: string; + text: string; + clientMessageId: string; + } | null>(null); const sendAgentMessageRef = useRef< | (( agentId: string, text: string, attachments: ComposerAttachment[], activeTurnBehavior: "interrupt" | "steer", + clientMessageId?: string, ) => Promise) | null >(null); @@ -1546,7 +1553,7 @@ function ComposerContentImpl({ }, [focusInput, onFocusInput]); const submitMessage = useCallback( - async (text: string, submitAttachments: ComposerAttachment[]) => { + async (text: string, submitAttachments: ComposerAttachment[], messageId?: string) => { onMessageSent?.(); if (onSubmitMessageRef.current) { await onSubmitMessageRef.current({ text, attachments: submitAttachments, cwd }); @@ -1560,6 +1567,7 @@ function ComposerContentImpl({ text, submitAttachments, appSettings.sendBehavior === "steer" ? "steer" : "interrupt", + messageId, ); }, [appSettings.sendBehavior, cwd, onMessageSent, t], @@ -1575,29 +1583,51 @@ function ComposerContentImpl({ text: string, sendAttachments: ComposerAttachment[], activeTurnBehavior: "interrupt" | "steer", + explicitClientMessageId?: string, ) => { if (!client) { throw new Error(t("workspace.terminal.hostDisconnected")); } - await dispatchComposerAgentMessage({ - client, - agentId: targetAgentId, - text, - attachments: sendAttachments, - attachmentSubmitFormat: resolveComposerAttachmentSubmitFormat({ - supportsForgeAttachments: supportsForgeSearch, - }), - encodeImages, - submission: createMessageSubmissionWriter(serverId), - activeTurnBehavior, - activeTurnId: - activeTurnBehavior === "steer" - ? (selectAgentTurnPresentation( - useSessionStore.getState().sessions[serverId], - targetAgentId, - ).turnId ?? undefined) - : undefined, - }); + const pendingRetry = failedDirectSendRetryRef.current; + const resolvedClientMessageId = + explicitClientMessageId ?? + (pendingRetry?.agentId === targetAgentId && pendingRetry.text === text + ? pendingRetry.clientMessageId + : generateMessageId()); + try { + await dispatchComposerAgentMessage({ + client, + agentId: targetAgentId, + clientMessageId: resolvedClientMessageId, + text, + attachments: sendAttachments, + attachmentSubmitFormat: resolveComposerAttachmentSubmitFormat({ + supportsForgeAttachments: supportsForgeSearch, + }), + encodeImages, + submission: createMessageSubmissionWriter(serverId), + activeTurnBehavior, + activeTurnId: + activeTurnBehavior === "steer" + ? (selectAgentTurnPresentation( + useSessionStore.getState().sessions[serverId], + targetAgentId, + ).turnId ?? undefined) + : undefined, + }); + if (!explicitClientMessageId) { + failedDirectSendRetryRef.current = null; + } + } catch (error) { + if (!explicitClientMessageId) { + failedDirectSendRetryRef.current = { + agentId: targetAgentId, + text, + clientMessageId: resolvedClientMessageId, + }; + } + throw error; + } onAttentionPromptSend?.(); }; }, [appSettings.sendBehavior, client, onAttentionPromptSend, serverId, supportsForgeSearch, t]); @@ -1963,8 +1993,8 @@ function ComposerContentImpl({ agentId, messageId: id, queue: queueWriter, - submitMessage: ({ text, attachments: queuedAttachments }) => - submitMessage(text, queuedAttachments), + submitMessage: ({ messageId: queuedMessageId, text, attachments: queuedAttachments }) => + submitMessage(text, queuedAttachments, queuedMessageId), failedToSendMessage: t("composer.errors.failedToSend"), }); if (result.status === "failed") { diff --git a/packages/app/src/runtime/host-runtime.ts b/packages/app/src/runtime/host-runtime.ts index b6d10dc8c7..ec88e43768 100644 --- a/packages/app/src/runtime/host-runtime.ts +++ b/packages/app/src/runtime/host-runtime.ts @@ -2604,12 +2604,13 @@ export class HostRuntimeStore { useSessionStore.getState().sessions[serverId]?.queuedMessages.get(queuedAgentId) ?? [], write: (update) => useSessionStore.getState().setQueuedMessages(serverId, update), }, - submitMessage: async ({ text, attachments }) => { + submitMessage: async ({ messageId, text, attachments }) => { const supportsForgeAttachments = useSessionStore.getState().sessions[serverId]?.serverInfo?.features?.forgeSearch === true; await dispatchComposerAgentMessage({ client, agentId, + clientMessageId: messageId, text, attachments, attachmentSubmitFormat: resolveComposerAttachmentSubmitFormat({ diff --git a/packages/app/src/stores/session-store.test.ts b/packages/app/src/stores/session-store.test.ts index 1b9052c4c1..2aed69f986 100644 --- a/packages/app/src/stores/session-store.test.ts +++ b/packages/app/src/stores/session-store.test.ts @@ -849,3 +849,42 @@ describe("removeWorkspace", () => { expect(after.workspaces).toBe(before.workspaces); }); }); + +describe("acknowledgedClientMessageIds pruning", () => { + it("prunes acknowledged queued messages in setAgentStreamState and applyAgentTimelineResponseState", () => { + const store = useSessionStore.getState(); + initializeTestSession(); + store.setQueuedMessages( + "test-server", + new Map([ + [ + "agent-1", + [ + { id: "cmsg-1", text: "first", attachments: [] }, + { id: "cmsg-2", text: "second", attachments: [] }, + ], + ], + ]), + ); + + store.setAgentStreamState("test-server", "agent-1", { + acknowledgedClientMessageIds: ["cmsg-1"], + }); + expect( + useSessionStore.getState().sessions["test-server"]?.queuedMessages.get("agent-1"), + ).toEqual([{ id: "cmsg-2", text: "second", attachments: [] }]); + + store.applyAgentTimelineResponseState("test-server", "agent-1", { + items: [], + head: [], + range: null, + older: "unchanged", + newer: false, + synchronized: true, + acknowledgedClientMessageIds: ["cmsg-2"], + }); + expect( + useSessionStore.getState().sessions["test-server"]?.queuedMessages.get("agent-1"), + ).toBeUndefined(); + }); +}); diff --git a/packages/app/src/stores/session-store.ts b/packages/app/src/stores/session-store.ts index 7c52269e07..370b1811ea 100644 --- a/packages/app/src/stores/session-store.ts +++ b/packages/app/src/stores/session-store.ts @@ -732,6 +732,29 @@ function isSessionServerInfoUnchanged(input: { ); } +function pruneAcknowledgedQueuedMessages( + queuedMessages: Map< + string, + Array<{ id: string; text: string; attachments: ComposerAttachment[] }> + >, + agentId: string, + acknowledgedIds: readonly string[], +): Map> { + if (acknowledgedIds.length === 0) return queuedMessages; + const current = queuedMessages.get(agentId); + if (!current?.length) return queuedMessages; + const acknowledgedSet = new Set(acknowledgedIds); + const filtered = current.filter((item) => !acknowledgedSet.has(item.id)); + if (filtered.length === current.length) return queuedMessages; + const next = new Map(queuedMessages); + if (filtered.length > 0) { + next.set(agentId, filtered); + } else { + next.delete(agentId); + } + return next; +} + export const useSessionStore = create()( subscribeWithSelector((set, get) => { const commitActivityUpdates: AgentLastActivityCommitter = (updates) => { @@ -1041,16 +1064,29 @@ export const useSessionStore = create()( } } + const acknowledgedIds = state.acknowledgedClientMessageIds ?? []; const currentSubmissions = session.messageSubmissions.get(agentId) ?? []; const observedSubmissions = observeMessageSubmissionCanonical( currentSubmissions, - state.acknowledgedClientMessageIds ?? [], + acknowledgedIds, ); const changedSubmissions = observedSubmissions !== currentSubmissions; + const queuedMessages = pruneAcknowledgedQueuedMessages( + session.queuedMessages, + agentId, + acknowledgedIds, + ); + const changedQueuedMessages = queuedMessages !== session.queuedMessages; const agentTasks = updateAgentTasks(session.agentTasks, agentId, state.taskSnapshot); const changedTasks = agentTasks !== session.agentTasks; - if (!changedTail && !changedHead && !changedSubmissions && !changedTasks) { + if ( + !changedTail && + !changedHead && + !changedSubmissions && + !changedQueuedMessages && + !changedTasks + ) { return prev; } @@ -1074,6 +1110,7 @@ export const useSessionStore = create()( agentStreamHead: nextHead, agentTasks, messageSubmissions, + queuedMessages, }, }, }; @@ -1447,6 +1484,11 @@ export const useSessionStore = create()( messageSubmissions.delete(agentId); } } + const queuedMessages = pruneAcknowledgedQueuedMessages( + session.queuedMessages, + agentId, + state.acknowledgedClientMessageIds, + ); if (state.synchronized) { nextAuthoritative.set(agentId, true); nextSyncGeneration.set(agentId, session.historySyncGeneration); @@ -1471,6 +1513,7 @@ export const useSessionStore = create()( agentAuthoritativeHistoryApplied: nextAuthoritative, agentHistorySyncGeneration: nextSyncGeneration, messageSubmissions, + queuedMessages, }, }, }; diff --git a/packages/client/src/daemon-client.test.ts b/packages/client/src/daemon-client.test.ts index 87719a0a56..6706e6bac7 100644 --- a/packages/client/src/daemon-client.test.ts +++ b/packages/client/src/daemon-client.test.ts @@ -1589,7 +1589,12 @@ test("defaults session RPC waiters to sixty seconds", async () => { expect(settled).toBe(false); await vi.advanceTimersByTimeAsync(1); - await expect(responsePromise).rejects.toThrow("Timeout waiting for message (60000ms)"); + await expect(responsePromise).rejects.toMatchObject({ + message: "Timeout waiting for message (60000ms)", + code: "DAEMON_REQUEST_TIMEOUT", + outcome: "unknown", + serverId: "srv_test_1", + }); }); test("honors explicit fetchAgent timeout below the session RPC default", async () => { diff --git a/packages/client/src/daemon-client.ts b/packages/client/src/daemon-client.ts index 017c302f7d..544a361409 100644 --- a/packages/client/src/daemon-client.ts +++ b/packages/client/src/daemon-client.ts @@ -7745,7 +7745,15 @@ export class DaemonClient { if (waiter) { this.waiters.delete(waiter); } - wrappedReject(timeoutError); + wrappedReject( + waiter?.submitted + ? new DaemonConnectionError(timeoutError.message, timeoutError.code, { + serverId: this.destinationServerId, + outcome: "unknown", + cause: timeoutError, + }) + : timeoutError, + ); }, timeout) : null;