From d17aa098a6b1aa3c432ab4b550ca0de37b83c2da Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Mon, 27 Jul 2026 12:44:32 +0200 Subject: [PATCH] fix(app): stop completed turns from appearing stuck Turn activity was inferred from a user-message row flag, so a stale duplicate row could keep the working footer active after the turn ended. Track submission lifecycle separately and route every send path through the shared submission flow. --- docs/timeline-sync.md | 41 +- .../app/e2e/agent-message-submission.spec.ts | 312 +++++++++++ .../app/e2e/helpers/agent-message-gate.ts | 74 +++ packages/app/src/agent-stream/view.tsx | 13 +- packages/app/src/components/message.tsx | 18 +- .../rewind/use-rewind-agent-mutation.ts | 8 - packages/app/src/composer/actions.test.ts | 183 +++++- packages/app/src/composer/actions.ts | 66 +-- .../src/composer/draft/create-flow.test.ts | 6 +- .../app/src/composer/draft/create-flow.ts | 23 +- .../app/src/composer/draft/workspace-tab.tsx | 6 +- packages/app/src/composer/index.tsx | 22 +- .../app/src/composer/submission/model.test.ts | 203 +++++++ packages/app/src/composer/submission/model.ts | 134 +++++ .../app/src/composer/submission/writer.ts | 18 + packages/app/src/composer/submit.ts | 2 +- packages/app/src/contexts/session-context.tsx | 5 + packages/app/src/panels/agent-panel.tsx | 7 + packages/app/src/runtime/host-runtime.test.ts | 60 ++ packages/app/src/runtime/host-runtime.ts | 26 +- packages/app/src/stores/session-store.ts | 184 ++++++- .../timeline/session-stream-reducers.test.ts | 301 +++++++--- .../src/timeline/session-stream-reducers.ts | 229 ++------ packages/app/src/timeline/turn-time.test.ts | 26 +- packages/app/src/timeline/turn-time.ts | 8 +- packages/app/src/types/stream.test.ts | 332 +++++++---- packages/app/src/types/stream.ts | 519 +++++++++++++----- .../src/utils/agent-directory-sync.test.ts | 70 +++ .../providers/mock-load-test-agent.test.ts | 19 + .../agent/providers/mock-load-test-agent.ts | 12 + 30 files changed, 2246 insertions(+), 681 deletions(-) create mode 100644 packages/app/e2e/agent-message-submission.spec.ts create mode 100644 packages/app/e2e/helpers/agent-message-gate.ts create mode 100644 packages/app/src/composer/submission/model.test.ts create mode 100644 packages/app/src/composer/submission/model.ts create mode 100644 packages/app/src/composer/submission/writer.ts diff --git a/docs/timeline-sync.md b/docs/timeline-sync.md index b881744f50..6893f19ddb 100644 --- a/docs/timeline-sync.md +++ b/docs/timeline-sync.md @@ -87,14 +87,43 @@ its completion advances `seqEnd`, followed by a merged assistant message. The ap remaining page through the existing stream reducer. It must not append full projected text to a live prefix. -Optimistic user prompts occupy stable timeline slots. Catch-up never extracts, delays, or reinserts -them. A canonical user row replaces its matching slot in place; an unmatched prompt stays exactly -where the user submitted it. Other canonical rows are applied after the already-present timeline -instead of relocating visible user messages around newly fetched history. +Every path that sends a message to an agent — composer send, dictation accept-and-send, queued +send-now, and the automatic queue drain in `HostRuntime` — goes through +`dispatchComposerAgentMessage` with a submission writer. There is no second transport for the same +product action: calling `client.sendAgentMessage` directly skips the submitted row and the pending +footer, and permanently drops attachments because the daemon does not echo them back. + +A submitted prompt is one `UserMessageItem` row. That row is the authoritative local presentation: +its stable identity, text, timestamp, images, and attachments do not change when the provider +acknowledges it. Submission lifecycle is a separate record keyed by agent, not another row shape or +a property inferred from message identity. It begins in `waiting-for-rpc`. Submissions made while a +stored agent is not running require both the RPC and an authoritative non-running-to-running +transition; RPC-first moves to `waiting-for-running`, while running-first moves to +`running-observed`. The second observation removes the record. A force-send that begins while the +agent is already running requires only RPC acceptance, so prior running state cannot count as new +acceptance proof. + +An accepted send response does not project agent lifecycle. If it arrives while the replica is +still non-running, the lifecycle record remains `waiting-for-running` until the real transition, so +the footer has no blank frame. A transport error rolls back only `waiting-for-rpc`; +`running-observed` proves the run started and preserves the submitted row. Canonical submitted user rows carry the provider's `messageId` and Paseo's optional -`clientMessageId`. Clients reconcile optimistic prompts by `clientMessageId`. Content matching is -limited to the dated compatibility path for daemon timelines created before that field existed. +`clientMessageId`. The user-message producer reconciles them by `clientMessageId`, adds provider +identity to the existing row, and keeps the local presentation in its original timeline slot. +Content matching is limited to the dated compatibility path for daemon timelines created before +that field existed, and it is direction-independent: either the submitted row or its legacy +canonical twin can arrive second. Generic reducers and consumers do not reimplement message +identity matching. + +Ordinary bootstrap, same-epoch reset, and catch-up replacement preserve unmatched locally submitted +rows because a provider may never echo them. A known epoch change or rewind replaces history and +drops acknowledged local rows omitted by the new canonical epoch; only the explicitly pending +submission crosses that destructive boundary. + +Canonical replacement owns both timeline lanes. It reconciles matching submitted head rows in the +head lane and omits their canonical copies from the tail. If a live assistant head is the canonical +assistant prefix, it also stays in the head lane. No row may be returned in both lanes. ## Relevant code diff --git a/packages/app/e2e/agent-message-submission.spec.ts b/packages/app/e2e/agent-message-submission.spec.ts new file mode 100644 index 0000000000..5ee83ab5fa --- /dev/null +++ b/packages/app/e2e/agent-message-submission.spec.ts @@ -0,0 +1,312 @@ +import type { Locator, Page } from "@playwright/test"; +import { expect, test as baseTest } from "./fixtures"; +import { expectAgentIdle } from "./helpers/agent-stream"; +import { gateNextAgentMessage } from "./helpers/agent-message-gate"; +import { + attachImageFromMenu, + expectComposerDraft, + expectComposerEditable, + expectAttachmentPill, + expectComposerVisible, +} from "./helpers/composer"; +import { openAgentRoute, seedMockAgentWorkspace } from "./helpers/mock-agent"; +import { seedWorkspace } from "./helpers/seed-client"; +import { waitForWorkspaceTabsVisible } from "./helpers/workspace-tabs"; +import { getServerId } from "./helpers/server-id"; +import { buildHostWorkspaceRoute } from "@/utils/host-routes"; +import { delayBrowserAgentCreatedStatus } from "./helpers/new-workspace"; + +const IMAGE = { + name: "message-submission.png", + mimeType: "image/png", + buffer: Buffer.from( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==", + "base64", + ), +}; + +interface MessageGeometry { + x: number; + y: number; + width: number; + height: number; +} + +interface SubmissionScenario { + gate: Awaited>; +} + +interface DraftCreateScenario { + workspaceId: string; + agentCreatedDelay: Awaited>; +} + +interface RejectionScenario { + errorMessage: string; +} + +const test = baseTest.extend<{ + submissionScenario: SubmissionScenario; + draftCreateScenario: DraftCreateScenario; + rejectionScenario: RejectionScenario; +}>({ + submissionScenario: async ({ page }, provide, testInfo) => { + const gate = await gateNextAgentMessage(page); + const agent = await seedMockAgentWorkspace({ + repoPrefix: `message-submission-${testInfo.workerIndex}-`, + title: "Message submission regression", + model: "ten-second-stream", + }); + await openAgentRoute(page, { workspaceId: agent.workspaceId, agentId: agent.agentId }); + await expectComposerVisible(page); + await expectAgentIdle(page); + await provide({ gate }); + await agent.cleanup(); + }, + draftCreateScenario: async ({ page }, provide, testInfo) => { + const agentCreatedDelay = await delayBrowserAgentCreatedStatus(page); + const workspace = await seedWorkspace({ + repoPrefix: `message-create-handoff-${testInfo.workerIndex}-`, + }); + await provide({ workspaceId: workspace.workspaceId, agentCreatedDelay }); + agentCreatedDelay.release(); + await workspace.cleanup(); + }, + rejectionScenario: async ({ page }, provide, testInfo) => { + const errorMessage = "Requested mock prompt rejection"; + const agent = await seedMockAgentWorkspace({ + repoPrefix: `message-rejection-${testInfo.workerIndex}-`, + title: "Message rejection regression", + model: "ten-second-stream", + featureValues: { mockPromptRejections: 1 }, + }); + await openAgentRoute(page, { workspaceId: agent.workspaceId, agentId: agent.agentId }); + await expectComposerVisible(page); + await expectAgentIdle(page); + await provide({ errorMessage }); + await agent.cleanup(); + }, +}); + +async function submitMessageWithImage(page: Page, prompt: string): Promise { + await attachImageFromMenu(page, IMAGE); + await expectAttachmentPill(page, "composer-image-attachment-pill"); + const composer = page.getByRole("textbox", { name: "Message agent..." }).first(); + await composer.fill(prompt); + await composer.press("Enter"); + const nextFrame = await composer.evaluate( + (composerElement, submittedPrompt) => + new Promise<{ + rowPresent: boolean; + workingPresent: boolean; + composerValue: string | null; + attachmentPresent: boolean; + }>((resolve) => { + requestAnimationFrame(() => { + const rows = Array.from(document.querySelectorAll('[data-testid="user-message"]')); + const composerInput = composerElement as HTMLInputElement | HTMLTextAreaElement; + resolve({ + rowPresent: rows.some((row) => row.textContent?.includes(submittedPrompt)), + workingPresent: Boolean( + document.querySelector('[data-testid="turn-working-indicator"]'), + ), + composerValue: composerInput.value, + attachmentPresent: Boolean( + document.querySelector('[data-testid="composer-image-attachment-pill"]'), + ), + }); + }); + }), + prompt, + ); + expect(nextFrame).toEqual({ + rowPresent: true, + workingPresent: true, + composerValue: "", + attachmentPresent: false, + }); + return page.getByTestId("user-message").filter({ hasText: prompt }).last(); +} + +async function expectPendingSubmission(page: Page, userMessage: Locator): Promise { + await expect(userMessage).toBeVisible(); + await expect(page.getByTestId("turn-working-indicator")).toBeVisible(); + await expect(page.getByRole("textbox", { name: "Message agent..." }).first()).toHaveValue(""); + await expect(page.getByTestId("composer-image-attachment-pill")).toHaveCount(0); + await expect(userMessage.getByTestId("user-message-timestamp")).toBeAttached(); + await expect(userMessage.getByTestId("user-message-trailing-row")).toHaveCSS("opacity", "0"); + await expect(userMessage).toHaveAttribute("aria-busy", "true"); + await expect(userMessage.getByRole("button", { name: "Open image attachment" })).toBeVisible(); +} + +async function readMessageGeometry(userMessage: Locator): Promise { + const box = await userMessage.boundingBox(); + if (!box) throw new Error("Submitted user message has no browser geometry"); + return { x: box.x, y: box.y, width: box.width, height: box.height }; +} + +async function beginWorkingFooterContinuityCheck(page: Page): Promise<() => Promise> { + await expect(page.getByTestId("turn-working-indicator")).toBeVisible(); + await page.evaluate(() => { + const state = { active: true, sawMissing: false }; + const windowState = window as unknown as Record; + windowState.__messageSubmissionFooterContinuity = state; + const checkFrame = () => { + if (!state.active) return; + if (!document.querySelector('[data-testid="turn-working-indicator"]')) { + state.sawMissing = true; + } + requestAnimationFrame(checkFrame); + }; + requestAnimationFrame(checkFrame); + }); + + return async () => { + const sawMissing = await page.evaluate(() => { + const windowState = window as unknown as Record; + const state = windowState.__messageSubmissionFooterContinuity as + | { active: boolean; sawMissing: boolean } + | undefined; + if (!state) throw new Error("Working-footer continuity check was not started"); + state.active = false; + delete windowState.__messageSubmissionFooterContinuity; + return state.sawMissing; + }); + expect(sawMissing).toBe(false); + }; +} + +async function expectAcceptedSubmission( + page: Page, + userMessage: Locator, + submittedGeometry: MessageGeometry, +): Promise { + await expect(page.getByTestId("turn-working-indicator")).toBeVisible(); + await expect(userMessage).toHaveAttribute("aria-busy", "false", { timeout: 30_000 }); + await expect(page.getByTestId("turn-working-indicator")).toBeVisible(); + expect(await readMessageGeometry(userMessage)).toEqual(submittedGeometry); +} + +async function submitMessageThatWillBeRejected(page: Page, prompt: string): Promise { + await attachImageFromMenu(page, IMAGE); + await expectAttachmentPill(page, "composer-image-attachment-pill"); + const composer = page.getByRole("textbox", { name: "Message agent..." }).first(); + await composer.fill(prompt); + await composer.press("Enter"); +} + +async function expectRejectedSubmissionRestored( + page: Page, + input: { prompt: string; errorMessage: string }, +): Promise { + await expect(page.getByText(input.errorMessage)).toBeVisible({ timeout: 30_000 }); + await expectComposerDraft(page, input.prompt); + await expectComposerEditable(page); + await expectAttachmentPill(page, "composer-image-attachment-pill"); + await expect(page.getByRole("button", { name: "Send message" })).toBeEnabled(); + await expect(page.getByTestId("user-message").filter({ hasText: input.prompt })).toHaveCount(0); + await expect(page.getByTestId("turn-working-indicator")).toHaveCount(0); +} + +async function retryRestoredSubmission(page: Page, prompt: string): Promise { + await page.getByRole("textbox", { name: "Message agent..." }).first().press("Enter"); + const userMessage = page.getByTestId("user-message").filter({ hasText: prompt }); + await expect(userMessage).toHaveCount(1); + await expect(userMessage).toHaveAttribute("aria-busy", "false", { timeout: 30_000 }); + await expect(userMessage.getByRole("button", { name: "Open image attachment" })).toBeVisible(); + await expect(page.getByTestId("composer-image-attachment-pill")).toHaveCount(0); +} + +async function openWorkspaceDraft(page: Page, workspaceId: string): Promise { + await page.goto(buildHostWorkspaceRoute(getServerId(), workspaceId)); + await waitForWorkspaceTabsVisible(page); + await page.getByTestId("workspace-new-agent-tab-inline").click(); + await expectComposerVisible(page); +} + +async function expectCreatedAgentHandoff( + page: Page, + prompt: string, + userMessage: Locator, + submittedGeometry: MessageGeometry, +): Promise { + await expect(page.getByTestId("turn-working-indicator")).toBeVisible(); + await expect(page.getByTestId(/^workspace-tab-agent_/).first()).toBeVisible({ timeout: 30_000 }); + await expect(userMessage).toHaveAttribute("aria-busy", "false", { timeout: 30_000 }); + await expect(page.getByTestId("turn-working-indicator")).toBeVisible(); + await expect(page.getByTestId("user-message").filter({ hasText: prompt })).toHaveCount(1); + await expect(userMessage.getByRole("button", { name: "Open image attachment" })).toBeVisible(); + expect(await readMessageGeometry(userMessage)).toEqual(submittedGeometry); +} + +interface DraftCreatePendingSubmission { + prompt: string; + userMessage: Locator; + submittedGeometry: MessageGeometry; + finishFooterContinuityCheck: () => Promise; +} + +async function beginDraftCreateSubmission( + page: Page, + scenario: DraftCreateScenario, +): Promise { + await openWorkspaceDraft(page, scenario.workspaceId); + const prompt = "Keep this row through create handoff."; + const userMessage = await submitMessageWithImage(page, prompt); + await scenario.agentCreatedDelay.waitForCreateRequest(); + await scenario.agentCreatedDelay.waitForDelayedCreatedStatus(); + await expectPendingSubmission(page, userMessage); + const submittedGeometry = await readMessageGeometry(userMessage); + const finishFooterContinuityCheck = await beginWorkingFooterContinuityCheck(page); + return { prompt, userMessage, submittedGeometry, finishFooterContinuityCheck }; +} + +async function completeDraftCreateSubmission( + page: Page, + scenario: DraftCreateScenario, + pending: DraftCreatePendingSubmission, +): Promise { + scenario.agentCreatedDelay.release(); + await expectCreatedAgentHandoff( + page, + pending.prompt, + pending.userMessage, + pending.submittedGeometry, + ); + await pending.finishFooterContinuityCheck(); +} + +test.describe("Agent message submission", () => { + test("keeps the submitted row stable when the host accepts", async ({ + page, + submissionScenario, + }) => { + const userMessage = await submitMessageWithImage(page, "Hold this submission."); + await expectPendingSubmission(page, userMessage); + await submissionScenario.gate.waitForRequest(); + const submittedGeometry = await readMessageGeometry(userMessage); + const finishFooterContinuityCheck = await beginWorkingFooterContinuityCheck(page); + submissionScenario.gate.accept(); + await expectAcceptedSubmission(page, userMessage, submittedGeometry); + await finishFooterContinuityCheck(); + }); + + test("keeps the submitted row stable through draft create handoff", async ({ + page, + draftCreateScenario, + }) => { + test.setTimeout(120_000); + const pending = await beginDraftCreateSubmission(page, draftCreateScenario); + await completeDraftCreateSubmission(page, draftCreateScenario, pending); + }); + + test("restores a rejected submission and accepts its retry", async ({ + page, + rejectionScenario, + }) => { + const prompt = "Restore this rejected submission."; + await submitMessageThatWillBeRejected(page, prompt); + await expectRejectedSubmissionRestored(page, { prompt, ...rejectionScenario }); + await retryRestoredSubmission(page, prompt); + }); +}); diff --git a/packages/app/e2e/helpers/agent-message-gate.ts b/packages/app/e2e/helpers/agent-message-gate.ts new file mode 100644 index 0000000000..ec9d4a346e --- /dev/null +++ b/packages/app/e2e/helpers/agent-message-gate.ts @@ -0,0 +1,74 @@ +import type { Page, WebSocketRoute } from "@playwright/test"; +import { daemonWsRoutePattern } from "./daemon-port"; + +type WebSocketMessage = string | Buffer; + +interface SendAgentMessageRequest { + type: "send_agent_message_request"; + requestId: string; + agentId: string; +} + +function readSendRequest(message: WebSocketMessage): SendAgentMessageRequest | null { + if (typeof message !== "string") return null; + try { + const envelope = JSON.parse(message) as { + type?: unknown; + message?: Record; + }; + const request = envelope.type === "session" ? envelope.message : null; + if ( + request?.type !== "send_agent_message_request" || + typeof request.requestId !== "string" || + typeof request.agentId !== "string" + ) { + return null; + } + return { + type: "send_agent_message_request", + requestId: request.requestId, + agentId: request.agentId, + }; + } catch { + return null; + } +} + +export async function gateNextAgentMessage(page: Page) { + let serverSocket: WebSocketRoute | null = null; + let heldMessage: WebSocketMessage | null = null; + let hasHeldRequest = false; + let resolveRequest: ((request: SendAgentMessageRequest) => void) | null = null; + const requestSeen = new Promise((resolve) => { + resolveRequest = resolve; + }); + + await page.routeWebSocket(daemonWsRoutePattern(), (ws) => { + const server = ws.connectToServer(); + serverSocket = server; + + ws.onMessage((message) => { + const request = readSendRequest(message); + if (request && !hasHeldRequest) { + heldMessage = message; + hasHeldRequest = true; + resolveRequest?.(request); + return; + } + server.send(message); + }); + + server.onMessage((message) => ws.send(message)); + }); + + return { + waitForRequest: () => requestSeen, + accept() { + if (!serverSocket || !heldMessage) { + throw new Error("No held send-agent-message request to accept"); + } + serverSocket.send(heldMessage); + heldMessage = null; + }, + }; +} diff --git a/packages/app/src/agent-stream/view.tsx b/packages/app/src/agent-stream/view.tsx index 5af665d40a..b12b5b6c7e 100644 --- a/packages/app/src/agent-stream/view.tsx +++ b/packages/app/src/agent-stream/view.tsx @@ -41,6 +41,7 @@ import { } from "@/components/message"; import { PlanCard } from "@/components/plan-card"; import type { StreamItem } from "@/types/stream"; +import type { PendingMessageSubmission } from "@/composer/submission/model"; import type { PendingPermission } from "@/types/shared"; import type { AgentCapabilityFlags, @@ -239,6 +240,7 @@ export interface AgentStreamViewProps { streamItems: StreamItem[]; streamHead?: StreamItem[]; pendingPermissions: Map; + pendingMessageSubmission?: PendingMessageSubmission | null; routeBottomAnchorRequest?: BottomAnchorRouteRequest | null; isAuthoritativeHistoryReady?: boolean; toast?: ToastApi | null; @@ -327,6 +329,7 @@ const AgentStreamViewComponent = forwardRef ); }, - [context.capabilities, agentId, client, resolvedServerId], + [context.capabilities, agentId, client, pendingMessageSubmission, resolvedServerId], ); const renderAssistantMessageItem = useCallback( @@ -878,7 +882,7 @@ const AgentStreamViewComponent = forwardRef renderPendingPermissionsNode({ @@ -1157,6 +1161,9 @@ function agentStreamViewPropsEqual( if (left.streamItems !== right.streamItems) reasons.push("streamItems"); if (left.streamHead !== right.streamHead) reasons.push("streamHead"); if (left.pendingPermissions !== right.pendingPermissions) reasons.push("pendingPermissions"); + if (left.pendingMessageSubmission !== right.pendingMessageSubmission) { + reasons.push("pendingMessageSubmission"); + } if ( !bottomAnchorRouteRequestsEqual(left.routeBottomAnchorRequest, right.routeBottomAnchorRequest) ) { diff --git a/packages/app/src/components/message.tsx b/packages/app/src/components/message.tsx index 8abe679281..e1870bd729 100644 --- a/packages/app/src/components/message.tsx +++ b/packages/app/src/components/message.tsx @@ -132,6 +132,7 @@ interface UserMessageProps { client?: DaemonClient | null; isFirstInGroup?: boolean; isLastInGroup?: boolean; + isPending?: boolean; disableOuterSpacing?: boolean; } @@ -430,6 +431,7 @@ export const UserMessage = memo(function UserMessage({ client, isFirstInGroup = true, isLastInGroup = true, + isPending = false, disableOuterSpacing, }: UserMessageProps) { const isCompact = useIsCompactFormFactor(); @@ -441,7 +443,7 @@ export const UserMessage = memo(function UserMessage({ const hasText = message.trim().length > 0; const hasImages = images.length > 0; const hasAttachments = attachments.length > 0; - const showTrailingRow = hasText && (isCompact || isNative || isHovered); + const showTrailingRow = !isPending && hasText && (isCompact || isNative || isHovered); const formattedTimestamp = useMemo( () => formatMessageTimestamp(new Date(timestamp)), [timestamp], @@ -494,7 +496,7 @@ export const UserMessage = memo(function UserMessage({ ); return ( - + {hasText ? ( - - {formattedTimestamp} - {capabilities ? ( + + + {formattedTimestamp} + + {capabilities && messageId ? ( void } = {}, ): ComposerSendClient & { calls: FakeSendCall[] } { const calls: FakeSendCall[] = []; return { calls, sendAgentMessage: async (agentId, text, opts) => { - calls.push({ agentId, text, options: opts }); + const call = { agentId, text, options: opts }; + calls.push(call); if (options.rejection) { + options.beforeRejection?.(call); throw options.rejection; } }, @@ -183,27 +193,80 @@ function createFakeSendClient( }; } -interface FakeStream extends AgentStreamWriter { +interface FakeStream extends MessageSubmissionWriter { head: Map; tail: Map; + acceptance: MessageSubmissionAcceptance; } -function createFakeStream(initialHead: Map = new Map()): FakeStream { +function createFakeStream( + initialHead: Map = new Map(), + options: { acceptance?: MessageSubmissionAcceptance } = {}, +): FakeStream { const fake: FakeStream = { head: new Map(initialHead), tail: new Map(), - getTail: (agentId) => fake.tail.get(agentId), - getHead: (agentId) => fake.head.get(agentId), - setHead: (updater) => { - fake.head = updater(fake.head); + acceptance: options.acceptance ?? "rpc-and-running", + begin: (agentId, message) => { + const next = beginMessageSubmission(readSubmission(fake, agentId), { + message, + submittedAt: message.timestamp, + acceptance: fake.acceptance, + }); + writeSubmission(fake, agentId, next); }, - setTail: (updater) => { - fake.tail = updater(fake.tail); + accept: (agentId, clientMessageId) => { + writeSubmission( + fake, + agentId, + acceptMessageSubmission(readSubmission(fake, agentId), clientMessageId), + ); + }, + reject: (agentId, clientMessageId) => { + const result = rejectMessageSubmission(readSubmission(fake, agentId), clientMessageId); + writeSubmission(fake, agentId, result.state); + return result.outcome; }, }; return fake; } +const submissionsByFakeStream = new WeakMap< + FakeStream, + Map +>(); + +function readSubmission(fake: FakeStream, agentId: string): MessageSubmissionState { + return { + tail: fake.tail.get(agentId) ?? [], + head: fake.head.get(agentId) ?? [], + submission: submissionsByFakeStream.get(fake)?.get(agentId) ?? null, + }; +} + +function writeSubmission(fake: FakeStream, agentId: string, state: MessageSubmissionState): void { + fake.tail = new Map(fake.tail).set(agentId, state.tail); + fake.head = new Map(fake.head).set(agentId, state.head); + const submissions = submissionsByFakeStream.get(fake) ?? new Map(); + submissions.set(agentId, state.submission); + submissionsByFakeStream.set(fake, submissions); +} + +function applyFakeAuthoritativeRunning( + fake: FakeStream, + agentId: string, + clientMessageId: string, +): void { + const state = readSubmission(fake, agentId); + if (state.submission?.clientMessageId !== clientMessageId) { + return; + } + writeSubmission(fake, agentId, { + ...state, + submission: observeMessageSubmissionRunning(state.submission), + }); +} + function createFakeQueue( initial: Map = new Map(), ): QueueWriter & { state: Map } { @@ -337,7 +400,7 @@ describe("pickAndPersistImages", () => { }); describe("dispatchComposerAgentMessage", () => { - it("removes the optimistic prompt when the host rejects it", async () => { + it("removes the submitted prompt when the host rejects it", async () => { const rejection = new Error("Host rejected prompt"); const client = createFakeSendClient({ rejection }); const stream = createFakeStream(); @@ -349,14 +412,87 @@ describe("dispatchComposerAgentMessage", () => { text: "rejected prompt", attachments: [], encodeImages: passthroughEncodeImages, - stream, + submission: stream, }), ).rejects.toBe(rejection); - expect(stream.head.get("agent")).toBeUndefined(); + expect(stream.head.get("agent")).toEqual([]); expect(stream.tail.get("agent") ?? []).toEqual([]); }); + it("keeps the submitted prompt when authoritative running wins a late transport error", async () => { + const stream = createFakeStream(); + const transportError = new Error("Connection lost after authoritative running"); + const client = createFakeSendClient({ + rejection: transportError, + beforeRejection: ({ agentId, options }) => { + applyFakeAuthoritativeRunning(stream, agentId, options.messageId); + }, + }); + + await expect( + dispatchComposerAgentMessage({ + client, + agentId: "agent", + text: "accepted before transport loss", + attachments: [], + encodeImages: passthroughEncodeImages, + submission: stream, + }), + ).resolves.toBeUndefined(); + + expect(stream.tail.get("agent")).toHaveLength(1); + expect(stream.tail.get("agent")?.[0]).toMatchObject({ + kind: "user_message", + text: "accepted before transport loss", + }); + }); + + it("rolls back an already-running force send when its RPC fails", async () => { + const stream = createFakeStream(new Map(), { acceptance: "rpc-only" }); + const transportError = new Error("Force send failed while the prior turn was running"); + const client = createFakeSendClient({ + rejection: transportError, + beforeRejection: ({ agentId, options }) => { + applyFakeAuthoritativeRunning(stream, agentId, options.messageId); + }, + }); + + await expect( + dispatchComposerAgentMessage({ + client, + agentId: "agent", + text: "force send", + attachments: [], + encodeImages: passthroughEncodeImages, + submission: stream, + }), + ).rejects.toBe(transportError); + + expect(stream.tail.get("agent") ?? []).toEqual([]); + }); + + it("does not swallow a transport error when submission state is missing", async () => { + const transportError = new Error("Connection lost with unknown submission state"); + const client = createFakeSendClient({ rejection: transportError }); + const submission: MessageSubmissionWriter = { + begin: () => {}, + accept: () => {}, + reject: () => "unknown", + }; + + await expect( + dispatchComposerAgentMessage({ + client, + agentId: "agent", + text: "unknown state", + attachments: [], + encodeImages: passthroughEncodeImages, + submission, + }), + ).rejects.toBe(transportError); + }); + it("sends text + image data + structured attachments and appends user_message to the tail when head is empty", async () => { const client = createFakeSendClient(); const stream = createFakeStream(); @@ -371,7 +507,7 @@ describe("dispatchComposerAgentMessage", () => { { kind: "github_pr", item: prItem }, ], encodeImages: passthroughEncodeImages, - stream, + submission: stream, }); expect(client.calls).toHaveLength(1); @@ -393,7 +529,7 @@ describe("dispatchComposerAgentMessage", () => { }, ]); - expect(stream.head.get("agent")).toBeUndefined(); + expect(stream.head.get("agent")).toEqual([]); const tail = stream.tail.get("agent"); expect(tail).toHaveLength(1); const userMessage = tail?.[0] as Extract; @@ -402,7 +538,8 @@ describe("dispatchComposerAgentMessage", () => { expect(userMessage.images).toEqual([image]); expect(userMessage.attachments).toEqual(call.options.attachments); expect(userMessage.id).toBe(call.options.messageId); - expect(userMessage.optimistic).toBe(true); + expect(userMessage.clientMessageId).toBe(call.options.messageId); + expect(userMessage.messageId).toBeUndefined(); }); it("can send legacy GitHub attachment payloads for old daemons", async () => { @@ -416,7 +553,7 @@ describe("dispatchComposerAgentMessage", () => { attachments: [{ kind: "forge_change_request", item: prItem }], attachmentSubmitFormat: "legacy-github", encodeImages: passthroughEncodeImages, - stream, + submission: stream, }); expect(client.calls[0].options.attachments).toEqual([ @@ -449,11 +586,11 @@ describe("dispatchComposerAgentMessage", () => { text: "next message", attachments: [], encodeImages: passthroughEncodeImages, - stream, + submission: stream, }); expect(stream.head.get("agent")).toHaveLength(2); - expect(stream.tail.get("agent")).toBeUndefined(); + expect(stream.tail.get("agent")).toEqual([]); }); it("submits empty wire arrays when no attachments are provided", async () => { @@ -466,7 +603,7 @@ describe("dispatchComposerAgentMessage", () => { text: "plain message", attachments: [], encodeImages: passthroughEncodeImages, - stream, + submission: stream, }); expect(client.calls[0]?.options).toMatchObject({ @@ -486,7 +623,7 @@ describe("dispatchComposerAgentMessage", () => { text: "review this", attachments: [review], encodeImages: passthroughEncodeImages, - stream, + submission: stream, }); expect(client.calls[0]?.options.attachments).toEqual([review.attachment]); @@ -504,7 +641,7 @@ describe("dispatchComposerAgentMessage", () => { text: "inspect element", attachments: [browserElement], encodeImages: passthroughEncodeImages, - stream, + submission: stream, }); expect(client.calls[0]?.options.attachments).toEqual([ diff --git a/packages/app/src/composer/actions.ts b/packages/app/src/composer/actions.ts index 44fc317478..d622269aa2 100644 --- a/packages/app/src/composer/actions.ts +++ b/packages/app/src/composer/actions.ts @@ -12,13 +12,8 @@ import { splitComposerAttachmentsForSubmit, type ComposerAttachmentSubmitFormat, } from "@/composer/attachments/submit"; -import { - appendOptimisticUserMessageToStream, - buildOptimisticUserMessage, - generateMessageId, - type StreamItem, - type UserMessageItem, -} from "@/types/stream"; +import { createUserMessage, generateMessageId, type UserMessageItem } from "@/types/stream"; +import type { MessageSubmissionRejectionOutcome } from "@/composer/submission/model"; import type { PickedImageAttachmentInput } from "@/hooks/image-attachment-picker"; import { i18n } from "@/i18n/i18next"; @@ -70,11 +65,10 @@ export interface ComposerCancelClient { cancelAgent: (agentId: string) => Promise | void; } -export interface AgentStreamWriter { - getTail: (agentId: string) => StreamItem[] | undefined; - getHead: (agentId: string) => StreamItem[] | undefined; - setHead: (updater: (prev: Map) => Map) => void; - setTail: (updater: (prev: Map) => Map) => void; +export interface MessageSubmissionWriter { + begin: (agentId: string, message: UserMessageItem) => void; + accept: (agentId: string, clientMessageId: string) => void; + reject: (agentId: string, clientMessageId: string) => MessageSubmissionRejectionOutcome; } export interface QueueWriter { @@ -169,7 +163,7 @@ export interface DispatchComposerAgentMessageInput { encodeImages: ( images: AttachmentMetadata[], ) => Promise | undefined>; - stream: AgentStreamWriter; + submission: MessageSubmissionWriter; } export async function dispatchComposerAgentMessage( @@ -178,60 +172,30 @@ export async function dispatchComposerAgentMessage( const wirePayload = splitComposerAttachmentsForSubmit(input.attachments, { format: input.attachmentSubmitFormat, }); - const messageId = generateMessageId(); - const userMessage = buildOptimisticUserMessage({ - id: messageId, + const clientMessageId = generateMessageId(); + const userMessage = createUserMessage({ + clientMessageId, text: input.text, timestamp: new Date(), images: wirePayload.images, attachments: wirePayload.attachments, }); - const rollbackOptimisticMessage = appendUserMessageToStream( - input.agentId, - userMessage, - input.stream, - ); + input.submission.begin(input.agentId, userMessage); try { const imagesData = await input.encodeImages(wirePayload.images); await input.client.sendAgentMessage(input.agentId, input.text, { - messageId, + messageId: clientMessageId, images: imagesData ?? [], attachments: wirePayload.attachments, }); + input.submission.accept(input.agentId, clientMessageId); } catch (error) { - rollbackOptimisticMessage(); + const outcome = input.submission.reject(input.agentId, clientMessageId); + if (outcome === "accepted") return; throw error; } } -function appendUserMessageToStream( - agentId: string, - userMessage: UserMessageItem, - stream: AgentStreamWriter, -): () => void { - const result = appendOptimisticUserMessageToStream({ - tail: stream.getTail(agentId) ?? [], - head: stream.getHead(agentId) ?? [], - message: userMessage, - placement: "active-head", - }); - const write = result.changedHead ? stream.setHead : stream.setTail; - const items = result.changedHead ? result.head : result.tail; - write((prev) => new Map(prev).set(agentId, items)); - - return () => { - write((prev) => { - const current = prev.get(agentId); - if (!current) return prev; - const nextItems = current.filter( - (item) => item.id !== userMessage.id || item.kind !== "user_message" || !item.optimistic, - ); - if (nextItems.length === current.length) return prev; - return new Map(prev).set(agentId, nextItems); - }); - }; -} - export interface QueueComposerMessageInput { agentId: string; text: string; diff --git a/packages/app/src/composer/draft/create-flow.test.ts b/packages/app/src/composer/draft/create-flow.test.ts index 41ea2a4d94..471d1fc8b9 100644 --- a/packages/app/src/composer/draft/create-flow.test.ts +++ b/packages/app/src/composer/draft/create-flow.test.ts @@ -13,7 +13,7 @@ describe("useDraftAgentCreateFlow", () => { useCreateFlowStore.setState({ pendingByDraftId: {} }); }); - it("renders a prepared new-workspace create attempt as optimistic chat before continuing it", async () => { + it("renders a prepared new-workspace submission before continuing it", async () => { const image: UserMessageImageAttachment = { id: "image-1", mimeType: "image/png", @@ -60,13 +60,13 @@ describe("useDraftAgentCreateFlow", () => { expect(result.current.isSubmitting).toBe(true); expect(result.current.draftAgent).toEqual({ currentAttempt: attempt }); - expect(result.current.optimisticStreamItems).toEqual([ + expect(result.current.submittedStreamItems).toEqual([ { kind: "user_message", id: "msg-prepared", + clientMessageId: "msg-prepared", text: "build this", timestamp: attempt.timestamp, - optimistic: true, images: [image], attachments: [attachment], }, diff --git a/packages/app/src/composer/draft/create-flow.ts b/packages/app/src/composer/draft/create-flow.ts index d4fddc412e..eaaddefe37 100644 --- a/packages/app/src/composer/draft/create-flow.ts +++ b/packages/app/src/composer/draft/create-flow.ts @@ -8,12 +8,13 @@ import { import { useCreateFlowStore } from "@/stores/create-flow-store"; import { useSessionStore } from "@/stores/session-store"; import { - buildOptimisticUserMessage, + createUserMessage, generateMessageId, type StreamItem, type UserMessageImageAttachment, } from "@/types/stream"; import type { AgentAttachment } from "@getpaseo/protocol/messages"; +import type { PendingMessageSubmission } from "@/composer/submission/model"; const EMPTY_STREAM_ITEMS: StreamItem[] = []; @@ -133,7 +134,7 @@ export function useDraftAgentCreateFlow({ const formErrorMessage = machine.tag === "draft" ? machine.errorMessage : ""; const isSubmitting = machine.tag === "creating"; - const optimisticStreamItems = useMemo(() => { + const submittedStreamItems = useMemo(() => { if (machine.tag !== "creating") { return EMPTY_STREAM_ITEMS; } @@ -147,8 +148,8 @@ export function useDraftAgentCreateFlow({ } return [ - buildOptimisticUserMessage({ - id: machine.attempt.clientMessageId, + createUserMessage({ + clientMessageId: machine.attempt.clientMessageId, text: machine.attempt.text, timestamp: machine.attempt.timestamp, images: machine.attempt.images, @@ -156,6 +157,13 @@ export function useDraftAgentCreateFlow({ }), ]; }, [machine]); + const pendingMessageSubmission = useMemo(() => { + if (machine.tag !== "creating") return null; + return { + clientMessageId: machine.attempt.clientMessageId, + submittedAt: machine.attempt.timestamp, + }; + }, [machine]); const draftAgent = useMemo(() => { if (machine.tag !== "creating") { @@ -195,8 +203,8 @@ export function useDraftAgentCreateFlow({ handoffCreatedAgentUserMessage( pendingServerId, createResult.agentId, - buildOptimisticUserMessage({ - id: attempt.clientMessageId, + createUserMessage({ + clientMessageId: attempt.clientMessageId, text: attempt.text, timestamp: attempt.timestamp, images: attempt.images, @@ -326,7 +334,8 @@ export function useDraftAgentCreateFlow({ machine, formErrorMessage, isSubmitting, - optimisticStreamItems, + submittedStreamItems, + pendingMessageSubmission, draftAgent, handleCreateFromInput, continueCreateFromAttempt, diff --git a/packages/app/src/composer/draft/workspace-tab.tsx b/packages/app/src/composer/draft/workspace-tab.tsx index 490b0a4951..155bcbdcfb 100644 --- a/packages/app/src/composer/draft/workspace-tab.tsx +++ b/packages/app/src/composer/draft/workspace-tab.tsx @@ -479,7 +479,8 @@ export function WorkspaceDraftAgentTab({ const { formErrorMessage, isSubmitting, - optimisticStreamItems, + submittedStreamItems, + pendingMessageSubmission, draftAgent, handleCreateFromInput, continueCreateFromAttempt, @@ -642,7 +643,8 @@ export function WorkspaceDraftAgentTab({ agentId={tabId} serverId={serverId} context={draftAgent} - streamItems={optimisticStreamItems} + streamItems={submittedStreamItems} + pendingMessageSubmission={pendingMessageSubmission} pendingPermissions={EMPTY_PENDING_PERMISSIONS} onOpenWorkspaceFile={onOpenWorkspaceFile} /> diff --git a/packages/app/src/composer/index.tsx b/packages/app/src/composer/index.tsx index 35232cd023..12be754227 100644 --- a/packages/app/src/composer/index.tsx +++ b/packages/app/src/composer/index.tsx @@ -64,7 +64,6 @@ import { sendQueuedComposerMessageNow, toggleGithubAttachmentFromPicker, uploadFileAttachments, - type AgentStreamWriter, type QueueWriter, type QueuedComposerMessage, } from "@/composer/actions"; @@ -91,6 +90,7 @@ import { useKeyboardActionHandler } from "@/hooks/use-keyboard-action-handler"; import type { KeyboardActionDefinition } from "@/keyboard/keyboard-action-dispatcher"; import type { MessageInputKeyboardActionKind } from "@/keyboard/actions"; import { submitAgentInput } from "@/composer/submit"; +import { createMessageSubmissionWriter } from "@/composer/submission/writer"; import { ComposerKeyboardScopeProvider } from "@/composer/keyboard-scope"; import { useAppSettings } from "@/hooks/use-settings"; import { isWeb, isNative } from "@/constants/platform"; @@ -1079,8 +1079,6 @@ export function Composer({ const queuedMessages = queuedMessagesRaw ?? EMPTY_ARRAY; const setQueuedMessages = useSessionStore((state) => state.setQueuedMessages); - const setAgentStreamTail = useSessionStore((state) => state.setAgentStreamTail); - const setAgentStreamHead = useSessionStore((state) => state.setAgentStreamHead); const isCompactFormFactor = useIsCompactFormFactor(); const isCompactLayout = resolveCompactLayout(isCompactLayoutOverride, isCompactFormFactor); @@ -1283,12 +1281,6 @@ export function Composer({ if (!client) { throw new Error(t("workspace.terminal.hostDisconnected")); } - const stream: AgentStreamWriter = { - getTail: (id) => useSessionStore.getState().sessions[serverId]?.agentStreamTail?.get(id), - getHead: (id) => useSessionStore.getState().sessions[serverId]?.agentStreamHead?.get(id), - setHead: (updater) => setAgentStreamHead(serverId, updater), - setTail: (updater) => setAgentStreamTail(serverId, updater), - }; await dispatchComposerAgentMessage({ client, agentId: targetAgentId, @@ -1298,19 +1290,11 @@ export function Composer({ supportsForgeAttachments: supportsForgeSearch, }), encodeImages, - stream, + submission: createMessageSubmissionWriter(serverId), }); onAttentionPromptSend?.(); }; - }, [ - client, - onAttentionPromptSend, - serverId, - setAgentStreamTail, - setAgentStreamHead, - supportsForgeSearch, - t, - ]); + }, [client, onAttentionPromptSend, serverId, supportsForgeSearch, t]); useEffect(() => { onSubmitMessageRef.current = onSubmitMessage; diff --git a/packages/app/src/composer/submission/model.test.ts b/packages/app/src/composer/submission/model.test.ts new file mode 100644 index 0000000000..1a18dfe724 --- /dev/null +++ b/packages/app/src/composer/submission/model.test.ts @@ -0,0 +1,203 @@ +import { describe, expect, it } from "vitest"; +import { createUserMessage, type StreamItem } from "@/types/stream"; +import { + acceptMessageSubmission, + beginMessageSubmission, + getPendingMessageSubmission, + observeMessageSubmissionRunning, + rejectMessageSubmission, + type MessageSubmissionState, +} from "./model"; + +function initialState(head: StreamItem[] = []): MessageSubmissionState { + return { tail: [], head, submission: null }; +} + +describe("message submission model", () => { + it("waits for authoritative running after RPC acceptance", () => { + const message = createUserMessage({ + clientMessageId: "client-1", + text: "hello", + timestamp: new Date("2026-07-26T10:00:00.000Z"), + }); + const begun = beginMessageSubmission(initialState(), { + message, + submittedAt: message.timestamp, + acceptance: "rpc-and-running", + }); + + const accepted = acceptMessageSubmission(begun, "client-1"); + const runningObserved = observeMessageSubmissionRunning(accepted.submission); + + expect(accepted.submission).toEqual({ + clientMessageId: "client-1", + submittedAt: message.timestamp, + phase: "waiting-for-running", + }); + expect(getPendingMessageSubmission(accepted.submission)).toEqual(accepted.submission); + expect(runningObserved).toBeNull(); + }); + + it("keeps a row when authoritative running precedes a late RPC error", () => { + const message = createUserMessage({ + clientMessageId: "client-1", + text: "hello", + timestamp: new Date("2026-07-26T10:00:00.000Z"), + }); + const begun = beginMessageSubmission(initialState(), { + message, + submittedAt: message.timestamp, + acceptance: "rpc-and-running", + }); + + const runningObserved = observeMessageSubmissionRunning(begun.submission); + const rejected = rejectMessageSubmission({ ...begun, submission: runningObserved }, "client-1"); + + expect(runningObserved).toEqual({ + clientMessageId: "client-1", + submittedAt: message.timestamp, + phase: "running-observed", + }); + expect(getPendingMessageSubmission(runningObserved)).toBeNull(); + expect(rejected).toEqual({ + outcome: "accepted", + state: { ...begun, submission: null }, + }); + }); + + it("does not accept an already-running force send from agent status alone", () => { + const message = createUserMessage({ + clientMessageId: "client-1", + text: "hello", + timestamp: new Date("2026-07-26T10:00:00.000Z"), + }); + const begun = beginMessageSubmission(initialState(), { + message, + submittedAt: message.timestamp, + acceptance: "rpc-only", + }); + + expect(observeMessageSubmissionRunning(begun.submission)).toBe(begun.submission); + }); + + it("begins and accepts without changing the submitted row", () => { + const submittedAt = new Date("2026-07-26T10:00:00.000Z"); + const message = createUserMessage({ + clientMessageId: "client-1", + text: "hello", + timestamp: submittedAt, + }); + + const begun = beginMessageSubmission(initialState(), { + message, + submittedAt, + acceptance: "rpc-only", + }); + const accepted = acceptMessageSubmission(begun, "client-1"); + + expect(begun).toEqual({ + tail: [message], + head: [], + submission: { + clientMessageId: "client-1", + submittedAt, + phase: "waiting-for-rpc", + acceptance: "rpc-only", + }, + }); + expect(accepted.tail).toBe(begun.tail); + expect(accepted.head).toBe(begun.head); + expect(accepted.submission).toBeNull(); + }); + + it("rejects by removing the submitted row and closing pending", () => { + const submittedAt = new Date("2026-07-26T10:00:00.000Z"); + const message = createUserMessage({ + clientMessageId: "client-1", + text: "hello", + timestamp: submittedAt, + }); + const existingHead = createUserMessage({ + id: "provider-prior", + messageId: "provider-prior", + text: "prior", + timestamp: new Date(0), + }); + const begun = beginMessageSubmission(initialState([existingHead]), { + message, + submittedAt, + acceptance: "rpc-and-running", + }); + + expect(rejectMessageSubmission(begun, "client-1")).toEqual({ + outcome: "rejected", + state: { + tail: [], + head: [existingHead], + submission: null, + }, + }); + }); + + it("treats running-observed as authoritatively accepted", () => { + const message = createUserMessage({ + clientMessageId: "client-1", + text: "hello", + timestamp: new Date("2026-07-26T10:00:00.000Z"), + }); + const state = initialState([message]); + + const runningObserved = { + ...state, + submission: { + clientMessageId: "client-1", + submittedAt: message.timestamp, + phase: "running-observed" as const, + }, + }; + + expect(rejectMessageSubmission(runningObserved, "client-1")).toEqual({ + outcome: "accepted", + state: { ...runningObserved, submission: null }, + }); + }); + + it("reports unknown when neither pending nor its submitted row exists", () => { + const state = initialState(); + + expect(rejectMessageSubmission(state, "client-1")).toEqual({ + outcome: "unknown", + state, + }); + }); +}); + +describe("submitting alongside repeated history", () => { + it("appends a resubmitted prompt instead of adopting an identical older row", () => { + // A row from a daemon that never recorded clientMessageId, e.g. history created + // before the field existed. Sending the same text again must be a new message. + const olderRow = createUserMessage({ + id: "provider-1", + messageId: "provider-1", + text: "continue", + timestamp: new Date("2026-07-26T10:00:00.000Z"), + }); + const resubmitted = createUserMessage({ + clientMessageId: "client-2", + text: "continue", + timestamp: new Date("2026-07-26T10:05:00.000Z"), + attachments: [{ type: "text", mimeType: "text/plain", text: "attachment" }], + }); + + const begun = beginMessageSubmission( + { tail: [olderRow], head: [], submission: null }, + { + message: resubmitted, + submittedAt: resubmitted.timestamp, + acceptance: "rpc-and-running", + }, + ); + + expect(begun.tail).toEqual([olderRow, resubmitted]); + }); +}); diff --git a/packages/app/src/composer/submission/model.ts b/packages/app/src/composer/submission/model.ts new file mode 100644 index 0000000000..b9147fe0e1 --- /dev/null +++ b/packages/app/src/composer/submission/model.ts @@ -0,0 +1,134 @@ +import type { StreamItem, UserMessageItem } from "@/types/stream"; + +export interface PendingMessageSubmission { + clientMessageId: string; + submittedAt: Date; +} + +export type MessageSubmissionAcceptance = "rpc-only" | "rpc-and-running"; + +export type MessageSubmissionRecord = PendingMessageSubmission & + ( + | { phase: "waiting-for-rpc"; acceptance: MessageSubmissionAcceptance } + | { phase: "waiting-for-running" } + | { phase: "running-observed" } + ); + +export function getPendingMessageSubmission( + submission: MessageSubmissionRecord | null | undefined, +): PendingMessageSubmission | null { + return submission && submission.phase !== "running-observed" ? submission : null; +} + +export interface MessageSubmissionState { + tail: StreamItem[]; + head: StreamItem[]; + submission: MessageSubmissionRecord | null; +} + +export type MessageSubmissionRejectionOutcome = "rejected" | "accepted" | "unknown"; + +export interface MessageSubmissionRejectionResult { + state: MessageSubmissionState; + outcome: MessageSubmissionRejectionOutcome; +} + +export function beginMessageSubmission( + state: MessageSubmissionState, + input: { + message: UserMessageItem; + submittedAt: Date; + acceptance: MessageSubmissionAcceptance; + }, +): MessageSubmissionState { + if (!input.message.clientMessageId) { + throw new Error("Submitted user message requires client identity"); + } + const pending = { + clientMessageId: input.message.clientMessageId, + submittedAt: input.submittedAt, + phase: "waiting-for-rpc" as const, + acceptance: input.acceptance, + }; + // A submission carries a freshly generated client identity, so it cannot already be in + // the timeline. Append it. Reconciling here would let the legacy content fallback adopt + // an older row that happens to repeat the same text. + if (state.head.length > 0) { + return { + tail: state.tail, + head: [...state.head, input.message], + submission: pending, + }; + } + return { + tail: [...state.tail, input.message], + head: state.head, + submission: pending, + }; +} + +export function acceptMessageSubmission( + state: MessageSubmissionState, + clientMessageId: string, +): MessageSubmissionState { + if (state.submission?.clientMessageId !== clientMessageId) { + return state; + } + if (state.submission.phase === "running-observed") { + return { ...state, submission: null }; + } + if (state.submission.phase === "waiting-for-running") return state; + return state.submission.acceptance === "rpc-only" + ? { ...state, submission: null } + : { + ...state, + submission: { + clientMessageId: state.submission.clientMessageId, + submittedAt: state.submission.submittedAt, + phase: "waiting-for-running", + }, + }; +} + +export function observeMessageSubmissionRunning( + submission: MessageSubmissionRecord | null, +): MessageSubmissionRecord | null { + if (!submission || submission.phase === "running-observed") return submission; + if (submission.phase === "waiting-for-running") return null; + if (submission.acceptance === "rpc-only") return submission; + return { + clientMessageId: submission.clientMessageId, + submittedAt: submission.submittedAt, + phase: "running-observed", + }; +} + +function removeSubmittedMessage(items: StreamItem[], clientMessageId: string): StreamItem[] { + const next = items.filter( + (item) => item.kind !== "user_message" || item.clientMessageId !== clientMessageId, + ); + return next.length === items.length ? items : next; +} + +export function rejectMessageSubmission( + state: MessageSubmissionState, + clientMessageId: string, +): MessageSubmissionRejectionResult { + if (state.submission?.clientMessageId !== clientMessageId) { + return { outcome: "unknown", state }; + } + if ( + state.submission.phase === "running-observed" || + state.submission.phase === "waiting-for-running" + ) { + return { outcome: "accepted", state: { ...state, submission: null } }; + } + return { + outcome: "rejected", + state: { + tail: removeSubmittedMessage(state.tail, clientMessageId), + head: removeSubmittedMessage(state.head, clientMessageId), + submission: null, + }, + }; +} diff --git a/packages/app/src/composer/submission/writer.ts b/packages/app/src/composer/submission/writer.ts new file mode 100644 index 0000000000..755eed1564 --- /dev/null +++ b/packages/app/src/composer/submission/writer.ts @@ -0,0 +1,18 @@ +import type { MessageSubmissionWriter } from "@/composer/actions"; +import { useSessionStore } from "@/stores/session-store"; + +/** + * Binds the submission lifecycle to a host session. Every path that sends a message to an + * agent — composer send, queued send-now, automatic queue drain — goes through this so a + * submitted row and its pending state are always created together. + */ +export function createMessageSubmissionWriter(serverId: string): MessageSubmissionWriter { + return { + begin: (agentId, message) => + useSessionStore.getState().beginAgentMessageSubmission(serverId, agentId, message), + accept: (agentId, clientMessageId) => + useSessionStore.getState().acceptAgentMessageSubmission(serverId, agentId, clientMessageId), + reject: (agentId, clientMessageId) => + useSessionStore.getState().rejectAgentMessageSubmission(serverId, agentId, clientMessageId), + }; +} diff --git a/packages/app/src/composer/submit.ts b/packages/app/src/composer/submit.ts index fb892d9850..c0c39171c5 100644 --- a/packages/app/src/composer/submit.ts +++ b/packages/app/src/composer/submit.ts @@ -51,7 +51,7 @@ export async function submitAgentInput( return "queued"; } - // Clear immediately so optimistic stream updates and composer state stay in sync. + // Clear immediately so the submitted timeline row and composer state stay in sync. if (shouldClearOnSubmit) { input.setUserInput(""); input.setAttachments([]); diff --git a/packages/app/src/contexts/session-context.tsx b/packages/app/src/contexts/session-context.tsx index 14d36f183a..964ecf150f 100644 --- a/packages/app/src/contexts/session-context.tsx +++ b/packages/app/src/contexts/session-context.tsx @@ -53,6 +53,7 @@ import { } from "@/utils/agent-initialization"; import { encodeImages } from "@/utils/encode-images"; import { derivePendingPermissionKey } from "@/utils/agent-snapshots"; +import { getPendingMessageSubmission } from "@/composer/submission/model"; import type { AttachmentMetadata } from "@/attachments/types"; import { patchWorkspaceScripts } from "@/contexts/session-workspace-scripts"; import { useToast } from "@/contexts/toast-context"; @@ -658,6 +659,9 @@ function SessionProviderInternal({ children, serverId, client }: SessionProvider const currentCursor = session?.agentTimelineCursor.get(agentId); const currentTail = session?.agentStreamTail.get(agentId) ?? []; const currentHead = session?.agentStreamHead.get(agentId) ?? []; + const pendingClientMessageId = + getPendingMessageSubmission(session?.messageSubmissions.get(agentId))?.clientMessageId ?? + null; setAgentTimelineHasOlder(serverId, (prev) => { if (prev.get(agentId) === payload.hasOlder) { @@ -677,6 +681,7 @@ function SessionProviderInternal({ children, serverId, client }: SessionProvider isInitializing, hasActiveInitDeferred, initRequestDirection: activeInitDeferred?.requestDirection ?? "tail", + pendingClientMessageId, }); if (result.error) { diff --git a/packages/app/src/panels/agent-panel.tsx b/packages/app/src/panels/agent-panel.tsx index 240f300dbb..565278ee9a 100644 --- a/packages/app/src/panels/agent-panel.tsx +++ b/packages/app/src/panels/agent-panel.tsx @@ -25,6 +25,7 @@ import { FileDropZone } from "@/components/file-drop/file-drop-zone"; import { useRetainedPanelActive } from "@/components/retained-panel"; import { SidebarCallout } from "@/components/sidebar-callout"; import { Composer } from "@/composer"; +import { getPendingMessageSubmission } from "@/composer/submission/model"; import { RewindComposerRestoreProvider } from "@/components/rewind/composer-restore"; import { getProviderIcon } from "@/components/provider-icons"; import { @@ -1288,6 +1289,11 @@ const AgentStreamSection = memo(function AgentStreamSection({ const streamItemsRaw = useSessionStore((state) => agentId ? state.sessions[serverId]?.agentStreamTail?.get(agentId) : undefined, ); + const pendingMessageSubmission = useSessionStore((state) => + agentId + ? getPendingMessageSubmission(state.sessions[serverId]?.messageSubmissions.get(agentId)) + : null, + ); const streamItems = streamItemsRaw ?? EMPTY_STREAM_ITEMS; const pendingPermissionList = useStoreWithEqualityFn( useSessionStore, @@ -1327,6 +1333,7 @@ const AgentStreamSection = memo(function AgentStreamSection({ routeBottomAnchorRequest={routeBottomAnchorRequest} isAuthoritativeHistoryReady={hasAppliedAuthoritativeHistory} toast={toast} + pendingMessageSubmission={pendingMessageSubmission} onOpenWorkspaceFile={onOpenWorkspaceFile} /> ); diff --git a/packages/app/src/runtime/host-runtime.test.ts b/packages/app/src/runtime/host-runtime.test.ts index 0f4946a4b9..64207753cc 100644 --- a/packages/app/src/runtime/host-runtime.test.ts +++ b/packages/app/src/runtime/host-runtime.test.ts @@ -2259,6 +2259,66 @@ describe("HostRuntimeStore", () => { useSessionStore.getState().clearSession(host.serverId); }); + it("submits an automatically drained message through the submission producer", async () => { + const host = makeHost({ serverId: "srv_drain_submission" }); + const fakeClient = new FakeDaemonClient(); + const send = new Deferred(); + fakeClient.sendAgentMessageResponses.push(send.promise); + const store = new HostRuntimeStore({ + deps: { + createClient: () => fakeClient as unknown as DaemonClient, + connectToDaemon: async () => ({ + client: fakeClient as unknown as DaemonClient, + serverId: host.serverId, + hostname: null, + }), + getClientId: async () => "cid_drain_submission", + }, + }); + const sessionStore = useSessionStore.getState(); + sessionStore.initializeSession(host.serverId, fakeClient as unknown as DaemonClient, 1); + sessionStore.setQueuedMessages( + host.serverId, + new Map([ + [ + "agent", + [ + { + id: "queued-with-attachment", + text: "read this file", + attachments: [ + { + kind: "workspace_file" as const, + path: "src/main.ts", + selection: { kind: "whole_file" as const }, + }, + ], + }, + ], + ], + ]), + ); + + store.drainQueuedAgentMessage(host.serverId, "agent"); + await fakeClient.waitForSentMessages(1); + + // The row and the pending submission must exist while the RPC is still in flight — + // the user sees their message and the working footer immediately, exactly as when + // they press send. + const session = useSessionStore.getState().sessions[host.serverId]; + const tail = session?.agentStreamTail.get("agent") ?? []; + expect(tail).toHaveLength(1); + expect(tail[0]).toMatchObject({ + kind: "user_message", + text: "read this file", + attachments: [{ type: "text", title: "main.ts", text: "Workspace file: src/main.ts" }], + }); + expect(session?.messageSubmissions.get("agent")).toBeDefined(); + + send.resolve(); + useSessionStore.getState().clearSession(host.serverId); + }); + it("restores an automatically drained message when sending fails", async () => { const host = makeHost({ serverId: "srv_failed_queue_drain" }); const fakeClient = new FakeDaemonClient(); diff --git a/packages/app/src/runtime/host-runtime.ts b/packages/app/src/runtime/host-runtime.ts index 1446328177..d77682a2f8 100644 --- a/packages/app/src/runtime/host-runtime.ts +++ b/packages/app/src/runtime/host-runtime.ts @@ -49,11 +49,9 @@ import { } from "@/data/push-router"; import { mountBrowserAutomationDaemonClientHandler } from "@/browser-automation/handler"; import { schedulesQueryBaseKey } from "@/schedules/aggregated-schedules"; -import { sendQueuedComposerMessageNow } from "@/composer/actions"; -import { - resolveComposerAttachmentSubmitFormat, - splitComposerAttachmentsForSubmit, -} from "@/composer/attachments/submit"; +import { dispatchComposerAgentMessage, sendQueuedComposerMessageNow } from "@/composer/actions"; +import { createMessageSubmissionWriter } from "@/composer/submission/writer"; +import { resolveComposerAttachmentSubmitFormat } from "@/composer/attachments/submit"; import { encodeImages } from "@/utils/encode-images"; import { DirectorySync, type RefreshAgentDirectoryResult } from "@/runtime/directory-sync"; import { ReplicaCache } from "@/runtime/replica-cache"; @@ -2071,14 +2069,16 @@ export class HostRuntimeStore { submitMessage: async ({ text, attachments }) => { const supportsForgeAttachments = useSessionStore.getState().sessions[serverId]?.serverInfo?.features?.forgeSearch === true; - const wirePayload = splitComposerAttachmentsForSubmit(attachments, { - format: resolveComposerAttachmentSubmitFormat({ supportsForgeAttachments }), - }); - const images = await encodeImages(wirePayload.images); - await client.sendAgentMessage(agentId, text, { - messageId: next.id, - ...(images && images.length > 0 ? { images } : {}), - attachments: wirePayload.attachments, + await dispatchComposerAgentMessage({ + client, + agentId, + text, + attachments, + attachmentSubmitFormat: resolveComposerAttachmentSubmitFormat({ + supportsForgeAttachments, + }), + encodeImages, + submission: createMessageSubmissionWriter(serverId), }); }, }) diff --git a/packages/app/src/stores/session-store.ts b/packages/app/src/stores/session-store.ts index 1d61867f26..89beef3b89 100644 --- a/packages/app/src/stores/session-store.ts +++ b/packages/app/src/stores/session-store.ts @@ -9,6 +9,14 @@ import { type StreamItem, type UserMessageItem, } from "@/types/stream"; +import { + acceptMessageSubmission, + beginMessageSubmission, + observeMessageSubmissionRunning, + rejectMessageSubmission, + type MessageSubmissionRecord, + type MessageSubmissionRejectionOutcome, +} from "@/composer/submission/model"; import type { PendingPermission } from "@/types/shared"; import type { ComposerAttachment } from "@/attachments/types"; import type { AgentLifecycleStatus } from "@getpaseo/protocol/agent-lifecycle"; @@ -368,6 +376,7 @@ export interface SessionState { // Stream state (head/tail model) agentStreamTail: Map; agentStreamHead: Map; + messageSubmissions: Map; agentTimelineCursor: Map; agentTimelineHasOlder: Map; agentTimelineOlderFetchInFlight: Map; @@ -461,6 +470,21 @@ interface SessionStoreActions { agentId: string, state: { tail?: StreamItem[]; head?: StreamItem[] }, ) => void; + beginAgentMessageSubmission: ( + serverId: string, + agentId: string, + message: UserMessageItem, + ) => void; + acceptAgentMessageSubmission: ( + serverId: string, + agentId: string, + clientMessageId: string, + ) => void; + rejectAgentMessageSubmission: ( + serverId: string, + agentId: string, + clientMessageId: string, + ) => MessageSubmissionRejectionOutcome; handoffCreatedAgentUserMessage: ( serverId: string, agentId: string, @@ -567,6 +591,39 @@ type SessionStore = SessionStoreState & SessionStoreActions; const agentLastActivityCoalescer = createAgentLastActivityCoalescer(); +function applyAgentStateToMessageSubmissions(input: { + previousAgents: Map; + nextAgents: Map; + submissions: Map; +}): Map { + let nextSubmissions = input.submissions; + const write = (agentId: string, submission: MessageSubmissionRecord | null): void => { + if (nextSubmissions === input.submissions) nextSubmissions = new Map(input.submissions); + if (submission) { + nextSubmissions.set(agentId, submission); + } else { + nextSubmissions.delete(agentId); + } + }; + + for (const [agentId, submission] of input.submissions) { + const nextAgent = input.nextAgents.get(agentId); + if (!nextAgent) { + write(agentId, null); + continue; + } + const previousAgent = input.previousAgents.get(agentId); + const observedRunningTransition = + previousAgent !== undefined && + previousAgent.status !== "running" && + nextAgent.status === "running"; + if (!observedRunningTransition) continue; + const nextSubmission = observeMessageSubmissionRunning(submission); + if (nextSubmission !== submission) write(agentId, nextSubmission); + } + return nextSubmissions; +} + // Helper to create initial session state function createInitialSessionState( serverId: string, @@ -588,6 +645,7 @@ function createInitialSessionState( currentAssistantMessage: "", agentStreamTail: new Map(), agentStreamHead: new Map(), + messageSubmissions: new Map(), agentTimelineCursor: new Map(), agentTimelineHasOlder: new Map(), agentTimelineOlderFetchInFlight: new Map(), @@ -1065,6 +1123,116 @@ export const useSessionStore = create()( }); }, + beginAgentMessageSubmission: (serverId, agentId, message) => { + set((prev) => { + const session = prev.sessions[serverId]; + if (!session) return prev; + const submittedAt = message.timestamp; + const currentAgent = session.agents.get(agentId); + const next = beginMessageSubmission( + { + tail: session.agentStreamTail.get(agentId) ?? [], + head: session.agentStreamHead.get(agentId) ?? [], + submission: session.messageSubmissions.get(agentId) ?? null, + }, + { + message, + submittedAt, + acceptance: + currentAgent !== undefined && currentAgent.status !== "running" + ? "rpc-and-running" + : "rpc-only", + }, + ); + const messageSubmissions = new Map(session.messageSubmissions); + if (!next.submission) { + throw new Error("Beginning a message submission must create lifecycle state"); + } + messageSubmissions.set(agentId, next.submission); + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { + ...session, + agentStreamTail: new Map(session.agentStreamTail).set(agentId, next.tail), + agentStreamHead: new Map(session.agentStreamHead).set(agentId, next.head), + messageSubmissions, + }, + }, + }; + }); + }, + + acceptAgentMessageSubmission: (serverId, agentId, clientMessageId) => { + set((prev) => { + const session = prev.sessions[serverId]; + if (!session) return prev; + const currentSubmission = session.messageSubmissions.get(agentId) ?? null; + const next = acceptMessageSubmission( + { + tail: session.agentStreamTail.get(agentId) ?? [], + head: session.agentStreamHead.get(agentId) ?? [], + submission: currentSubmission, + }, + clientMessageId, + ); + if (next.submission === currentSubmission) return prev; + const messageSubmissions = new Map(session.messageSubmissions); + if (next.submission) { + messageSubmissions.set(agentId, next.submission); + } else { + messageSubmissions.delete(agentId); + } + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { ...session, messageSubmissions }, + }, + }; + }); + }, + + rejectAgentMessageSubmission: (serverId, agentId, clientMessageId) => { + let outcome: MessageSubmissionRejectionOutcome = "unknown"; + set((prev) => { + const session = prev.sessions[serverId]; + if (!session) return prev; + const currentTail = session.agentStreamTail.get(agentId) ?? []; + const currentHead = session.agentStreamHead.get(agentId) ?? []; + const currentSubmission = session.messageSubmissions.get(agentId) ?? null; + const result = rejectMessageSubmission( + { tail: currentTail, head: currentHead, submission: currentSubmission }, + clientMessageId, + ); + outcome = result.outcome; + const next = result.state; + if (outcome === "unknown") return prev; + const messageSubmissions = new Map(session.messageSubmissions); + messageSubmissions.delete(agentId); + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { + ...session, + agentStreamTail: + next.tail === currentTail + ? session.agentStreamTail + : new Map(session.agentStreamTail).set(agentId, next.tail), + agentStreamHead: + next.head === currentHead + ? session.agentStreamHead + : new Map(session.agentStreamHead).set(agentId, next.head), + messageSubmissions, + }, + }, + }; + }); + return outcome; + }, + handoffCreatedAgentUserMessage: (serverId, agentId, message) => { let didHandoff = false; set((prev) => { @@ -1298,7 +1466,12 @@ export const useSessionStore = create()( return prev; } const nextAgents = typeof agents === "function" ? agents(session.agents) : agents; - if (session.agents === nextAgents) { + const messageSubmissions = applyAgentStateToMessageSubmissions({ + previousAgents: session.agents, + nextAgents, + submissions: session.messageSubmissions, + }); + if (session.agents === nextAgents && session.messageSubmissions === messageSubmissions) { return prev; } return { @@ -1308,10 +1481,11 @@ export const useSessionStore = create()( [serverId]: { ...session, agents: nextAgents, - workspaceAgentActivity: buildWorkspaceAgentActivityIndex( - nextAgents, - session.workspaceAgentActivity, - ), + messageSubmissions, + workspaceAgentActivity: + nextAgents === session.agents + ? session.workspaceAgentActivity + : buildWorkspaceAgentActivityIndex(nextAgents, session.workspaceAgentActivity), }, }, }; diff --git a/packages/app/src/timeline/session-stream-reducers.test.ts b/packages/app/src/timeline/session-stream-reducers.test.ts index 491da67636..d3604caeca 100644 --- a/packages/app/src/timeline/session-stream-reducers.test.ts +++ b/packages/app/src/timeline/session-stream-reducers.test.ts @@ -1,7 +1,7 @@ import { describe, expect, it } from "vitest"; import type { AgentStreamEventPayload } from "@getpaseo/protocol/messages"; import { - buildOptimisticUserMessage, + createUserMessage, hydrateStreamState, type AgentToolCallItem, type StreamItem, @@ -119,12 +119,12 @@ function makeAssistantItem( }; } -function makeOptimisticUserMessage( +function makeSubmittedUserMessage( text: string, - id = `optimistic-${text.length}`, + id = `submitted-${text.length}`, ): Extract { - return buildOptimisticUserMessage({ - id, + return createUserMessage({ + clientMessageId: id, text, timestamp: new Date(1000), }); @@ -172,6 +172,7 @@ const baseTimelineInput: ProcessTimelineResponseInput = { isInitializing: false, hasActiveInitDeferred: false, initRequestDirection: "tail", + pendingClientMessageId: null, }; const baseStreamInput: ProcessAgentStreamEventInput = { @@ -293,6 +294,42 @@ describe("processTimelineResponse", () => { expect(result.sideEffects.some((e) => e.type === "flush_pending_updates")).toBe(true); }); + it("keeps a live assistant and submitted head prompt in one lane during replacement", () => { + const submitted = makeSubmittedUserMessage("New prompt", "client-new-prompt"); + const liveAssistant = { + ...makeAssistantItem("Live answer", "answer-1"), + messageId: "answer-1", + }; + + const result = processTimelineResponse({ + ...baseTimelineInput, + currentTail: [], + currentHead: [liveAssistant, submitted], + currentCursor: { epoch: "epoch-1", startSeq: 1, endSeq: 1 }, + pendingClientMessageId: "client-new-prompt", + payload: { + ...baseTimelineInput.payload, + reset: true, + epoch: "epoch-1", + startCursor: { seq: 1 }, + endCursor: { seq: 1 }, + entries: [ + { + ...makeTimelineEntry(1, "Live", "assistant_message"), + item: { + type: "assistant_message", + text: "Live", + messageId: "answer-1", + }, + }, + ], + }, + }); + + expect(result.tail).toEqual([]); + expect(result.head).toEqual([{ ...liveAssistant, text: "Live" }, submitted]); + }); + it("uses the timeline entry timestamp as canonical", () => { const result = processTimelineResponse({ ...baseTimelineInput, @@ -321,12 +358,12 @@ describe("processTimelineResponse", () => { expect(assistant?.timestamp.toISOString()).toBe("2025-01-01T12:00:04.000Z"); }); - it("reconciles an optimistic user message during tail replacement", () => { + it("reconciles a submitted user message during tail replacement", () => { const image = { - id: "optimistic-image", + id: "submitted-image", mimeType: "image/png", storageType: "web-indexeddb" as const, - storageKey: "optimistic-image", + storageKey: "submitted-image", createdAt: 1000, }; const attachment = { @@ -335,8 +372,8 @@ describe("processTimelineResponse", () => { text: "attached context", title: "context.txt", }; - const optimistic = buildOptimisticUserMessage({ - id: "optimistic-create-user", + const submitted = createUserMessage({ + clientMessageId: "submitted-create-user", text: "Analyze this", timestamp: new Date(1000), images: [image], @@ -345,7 +382,7 @@ describe("processTimelineResponse", () => { const result = processTimelineResponse({ ...baseTimelineInput, - currentTail: [optimistic], + currentTail: [submitted], payload: { ...baseTimelineInput.payload, reset: true, @@ -358,6 +395,7 @@ describe("processTimelineResponse", () => { type: "user_message", text: "server-rendered attachment text", messageId: "canonical-create-user", + clientMessageId: "submitted-create-user", }, }, ], @@ -367,13 +405,14 @@ describe("processTimelineResponse", () => { const userMessages = result.tail.filter((item) => item.kind === "user_message"); expect(userMessages).toHaveLength(1); expect(userMessages[0]).toMatchObject({ - id: "canonical-create-user", + id: "submitted-create-user", + clientMessageId: "submitted-create-user", + messageId: "canonical-create-user", text: "Analyze this", timestamp: new Date(1000), images: [image], attachments: [attachment], }); - expect(userMessages[0]?.optimistic).toBeUndefined(); const repeated = processTimelineResponse({ ...baseTimelineInput, @@ -390,6 +429,7 @@ describe("processTimelineResponse", () => { type: "user_message", text: "server-rendered attachment text", messageId: "canonical-create-user", + clientMessageId: "submitted-create-user", }, }, ], @@ -399,24 +439,73 @@ describe("processTimelineResponse", () => { expect(repeated.tail.filter((item) => item.kind === "user_message")).toEqual(userMessages); }); - it("keeps an unmatched optimistic user message during tail replacement", () => { - const optimistic = makeOptimisticUserMessage("still sending", "optimistic-unmatched"); + it("keeps an unmatched submitted user message during tail replacement", () => { + const submitted = makeSubmittedUserMessage("still sending", "submitted-unmatched"); const result = processTimelineResponse({ ...baseTimelineInput, - currentTail: [optimistic], + currentTail: [submitted], + currentCursor: { epoch: "epoch-1", startSeq: 1, endSeq: 1 }, + pendingClientMessageId: "submitted-unmatched", payload: { ...baseTimelineInput.payload, reset: true, + epoch: "epoch-2", entries: [], }, }); - expect(result.tail).toEqual([optimistic]); + expect(result.tail).toEqual([submitted]); + }); + + it("keeps an acknowledged local row omitted by a same-epoch replacement", () => { + const acknowledged = createUserMessage({ + clientMessageId: "client-local-only", + text: "provider may not echo this", + timestamp: new Date(1000), + }); + + const result = processTimelineResponse({ + ...baseTimelineInput, + currentTail: [acknowledged], + currentCursor: { epoch: "epoch-1", startSeq: 1, endSeq: 1 }, + pendingClientMessageId: null, + payload: { + ...baseTimelineInput.payload, + reset: true, + epoch: "epoch-1", + entries: [], + }, + }); + + expect(result.tail).toEqual([acknowledged]); + }); + + it("drops an acknowledged local row omitted by a known epoch change", () => { + const acknowledged = createUserMessage({ + clientMessageId: "client-prior-epoch", + text: "prior prompt", + timestamp: new Date(1000), + }); + + const result = processTimelineResponse({ + ...baseTimelineInput, + currentTail: [acknowledged], + currentCursor: { epoch: "epoch-1", startSeq: 1, endSeq: 1 }, + pendingClientMessageId: null, + payload: { + ...baseTimelineInput.payload, + reset: true, + epoch: "epoch-2", + entries: [], + }, + }); + + expect(result.tail).toEqual([]); }); it("does not move an unmatched submission during timeline replacement", () => { - const unmatched = makeOptimisticUserMessage("first submission", "client-first"); + const unmatched = makeSubmittedUserMessage("first submission", "client-first"); const acknowledged: StreamItem[] = [ { kind: "user_message", @@ -437,6 +526,7 @@ describe("processTimelineResponse", () => { const result = processTimelineResponse({ ...baseTimelineInput, currentTail: [unmatched, ...acknowledged], + pendingClientMessageId: "client-first", payload: { ...baseTimelineInput.payload, reset: true, @@ -620,11 +710,11 @@ describe("processTimelineResponse", () => { startSeq: 1, endSeq: 1, }; - const optimistic = makeOptimisticUserMessage("sent while catching up", "optimistic-after"); + const submitted = makeSubmittedUserMessage("sent while catching up", "submitted-after"); const result = processTimelineResponse({ ...baseTimelineInput, - currentTail: [optimistic], + currentTail: [submitted], currentCursor: existingCursor, payload: { ...baseTimelineInput.payload, @@ -644,16 +734,19 @@ describe("processTimelineResponse", () => { const userMessages = result.tail.filter((item) => item.kind === "user_message"); expect(userMessages).toHaveLength(1); - expect(userMessages[0]?.id).toBe("canonical-after"); - expect(userMessages[0]?.optimistic).toBeUndefined(); + expect(userMessages[0]).toMatchObject({ + id: "submitted-after", + clientMessageId: "submitted-after", + messageId: "canonical-after", + }); }); - it("reconciles an optimistic user message by client message id", () => { - const optimistic = makeOptimisticUserMessage("local presentation", "client-message"); + it("reconciles a submitted user message by client message id", () => { + const submitted = makeSubmittedUserMessage("local presentation", "client-message"); const result = processTimelineResponse({ ...baseTimelineInput, - currentTail: [optimistic], + currentTail: [submitted], currentCursor: { epoch: "epoch-1", startSeq: 1, endSeq: 1 }, payload: { ...baseTimelineInput.payload, @@ -675,20 +768,20 @@ describe("processTimelineResponse", () => { const userMessages = result.tail.filter((item) => item.kind === "user_message"); expect(userMessages).toEqual([ expect.objectContaining({ - id: "provider-message", + id: "client-message", clientMessageId: "client-message", + messageId: "provider-message", text: "local presentation", }), ]); - expect(userMessages[0]?.optimistic).toBeUndefined(); }); - it("reconciles multiple optimistic user messages in canonical order", () => { + it("reconciles multiple submitted user messages in canonical order", () => { const result = processTimelineResponse({ ...baseTimelineInput, currentTail: [ - makeOptimisticUserMessage("first prompt", "optimistic-first"), - makeOptimisticUserMessage("second prompt", "optimistic-second"), + makeSubmittedUserMessage("first prompt", "submitted-first"), + makeSubmittedUserMessage("second prompt", "submitted-second"), ], currentCursor: { epoch: "epoch-1", startSeq: 1, endSeq: 1 }, payload: { @@ -699,14 +792,14 @@ describe("processTimelineResponse", () => { entries: [ { ...makeTimelineEntry(2, "first prompt", "user_message"), - item: { type: "user_message", text: "first prompt", messageId: "optimistic-first" }, + item: { type: "user_message", text: "first prompt", messageId: "submitted-first" }, }, { ...makeTimelineEntry(3, "second prompt", "user_message"), item: { type: "user_message", text: "second prompt", - messageId: "optimistic-second", + messageId: "submitted-second", }, }, ], @@ -716,15 +809,15 @@ describe("processTimelineResponse", () => { expect( result.tail .filter((item) => item.kind === "user_message") - .map((item) => ({ id: item.id, text: item.text, optimistic: item.optimistic })), + .map((item) => ({ id: item.id, text: item.text, messageId: item.messageId })), ).toEqual([ - { id: "optimistic-first", text: "first prompt", optimistic: undefined }, - { id: "optimistic-second", text: "second prompt", optimistic: undefined }, + { id: "submitted-first", text: "first prompt", messageId: "submitted-first" }, + { id: "submitted-second", text: "second prompt", messageId: "submitted-second" }, ]); }); - it("keeps a tail optimistic prompt before a reconciled live assistant head", () => { - const prompt = makeOptimisticUserMessage("new prompt", "optimistic-new-prompt"); + it("keeps a tail submitted prompt before a reconciled live assistant head", () => { + const prompt = makeSubmittedUserMessage("new prompt", "submitted-new-prompt"); const result = processTimelineResponse({ ...baseTimelineInput, @@ -763,8 +856,8 @@ describe("processTimelineResponse", () => { ).toEqual(["new prompt", "Hello"]); }); - it("keeps a tail optimistic prompt before a live head flushed by catch-up", () => { - const prompt = makeOptimisticUserMessage("new prompt", "optimistic-new-prompt"); + it("keeps a tail submitted prompt before a live head flushed by catch-up", () => { + const prompt = makeSubmittedUserMessage("new prompt", "submitted-new-prompt"); const result = processTimelineResponse({ ...baseTimelineInput, @@ -1000,7 +1093,7 @@ describe("processTimelineResponse", () => { }); it("does not move a submitted prompt when catch-up history arrives", () => { - const prompt = makeOptimisticUserMessage("New prompt", "new-prompt"); + const prompt = makeSubmittedUserMessage("New prompt", "new-prompt"); const result = processTimelineResponse({ ...baseTimelineInput, @@ -1024,7 +1117,7 @@ describe("processTimelineResponse", () => { }); it("does not move an unmatched head prompt when catch-up history arrives", () => { - const prompt = makeOptimisticUserMessage("New prompt", "new-prompt"); + const prompt = makeSubmittedUserMessage("New prompt", "new-prompt"); const result = processTimelineResponse({ ...baseTimelineInput, @@ -1053,7 +1146,7 @@ describe("processTimelineResponse", () => { }); it("acknowledges a head prompt in place while catch-up history arrives", () => { - const prompt = makeOptimisticUserMessage("New prompt", "new-prompt"); + const prompt = makeSubmittedUserMessage("New prompt", "new-prompt"); const result = processTimelineResponse({ ...baseTimelineInput, @@ -1090,12 +1183,12 @@ describe("processTimelineResponse", () => { expect( [...result.tail, ...result.head] .filter((item) => item.kind === "user_message") - .map((item) => item.optimistic), - ).toEqual([undefined]); + .map((item) => item.clientMessageId), + ).toEqual(["new-prompt"]); }); it("does not move a prompt around unrelated catch-up history", () => { - const prompt = makeOptimisticUserMessage("New prompt", "new-prompt"); + const prompt = makeSubmittedUserMessage("New prompt", "new-prompt"); const result = processTimelineResponse({ ...baseTimelineInput, @@ -1141,7 +1234,7 @@ describe("processTimelineResponse", () => { }); it("does not move a prompt or its live answer around catch-up history", () => { - const prompt = makeOptimisticUserMessage("New prompt", "new-prompt"); + const prompt = makeSubmittedUserMessage("New prompt", "new-prompt"); const live = processAgentStreamEvents({ events: [ makeStreamReducerEvent( @@ -1186,7 +1279,7 @@ describe("processTimelineResponse", () => { }); it("does not move a prompt or its live answer around catch-up tool history", () => { - const prompt = makeOptimisticUserMessage("New prompt", "new-prompt"); + const prompt = makeSubmittedUserMessage("New prompt", "new-prompt"); const live = processAgentStreamEvents({ events: [ makeStreamReducerEvent( @@ -1228,7 +1321,7 @@ describe("processTimelineResponse", () => { }); it("never moves submitted messages behind a later assistant response", () => { - const unmatched = makeOptimisticUserMessage("first submission", "client-first"); + const unmatched = makeSubmittedUserMessage("first submission", "client-first"); const acknowledged: StreamItem[] = [ { kind: "user_message", @@ -1287,7 +1380,7 @@ describe("processTimelineResponse", () => { }); it("acknowledges a local prompt in place when a remote user row also arrives", () => { - const prompt = makeOptimisticUserMessage("Local prompt", "local-prompt"); + const prompt = makeSubmittedUserMessage("Local prompt", "local-prompt"); const result = processTimelineResponse({ ...baseTimelineInput, @@ -1320,17 +1413,12 @@ describe("processTimelineResponse", () => { }); expect( - result.tail - .filter((item) => item.kind === "user_message") - .map((item) => ({ text: item.text, optimistic: item.optimistic })), - ).toEqual([ - { text: "Local prompt", optimistic: undefined }, - { text: "Remote prompt", optimistic: undefined }, - ]); + result.tail.filter((item) => item.kind === "user_message").map((item) => item.text), + ).toEqual(["Local prompt", "Remote prompt"]); }); - it("keeps an unmatched optimistic prompt when catch-up contains only a remote user row", () => { - const prompt = makeOptimisticUserMessage("Local prompt", "local-prompt"); + it("keeps an unmatched submitted prompt when catch-up contains only a remote user row", () => { + const prompt = makeSubmittedUserMessage("Local prompt", "local-prompt"); const result = processTimelineResponse({ ...baseTimelineInput, @@ -1355,17 +1443,12 @@ describe("processTimelineResponse", () => { }); expect( - result.tail - .filter((item) => item.kind === "user_message") - .map((item) => ({ text: item.text, optimistic: item.optimistic })), - ).toEqual([ - { text: "Local prompt", optimistic: true }, - { text: "Remote prompt", optimistic: undefined }, - ]); + result.tail.filter((item) => item.kind === "user_message").map((item) => item.text), + ).toEqual(["Local prompt", "Remote prompt"]); }); it("does not match equal prompt text when canonical client message ids differ", () => { - const prompt = makeOptimisticUserMessage("continue", "local-prompt"); + const prompt = makeSubmittedUserMessage("continue", "local-prompt"); const result = processTimelineResponse({ ...baseTimelineInput, @@ -1393,10 +1476,10 @@ describe("processTimelineResponse", () => { expect( result.tail .filter((item) => item.kind === "user_message") - .map((item) => ({ id: item.id, optimistic: item.optimistic })), + .map((item) => ({ id: item.id, messageId: item.messageId })), ).toEqual([ - { id: "local-prompt", optimistic: true }, - { id: "remote-prompt", optimistic: undefined }, + { id: "local-prompt", messageId: undefined }, + { id: "remote-prompt", messageId: "remote-prompt" }, ]); }); @@ -1639,8 +1722,8 @@ describe("processTimelineResponse", () => { }); }); - it("does not reconcile an active optimistic user message from a before-page response", () => { - const optimistic = makeOptimisticUserMessage("active prompt", "optimistic-active"); + it("does not reconcile an active submitted user message from a before-page response", () => { + const submitted = makeSubmittedUserMessage("active prompt", "submitted-active"); const existingCursor: TimelineCursor = { epoch: "epoch-1", startSeq: 3, @@ -1649,7 +1732,7 @@ describe("processTimelineResponse", () => { const result = processTimelineResponse({ ...baseTimelineInput, - currentTail: [optimistic], + currentTail: [submitted], currentCursor: existingCursor, payload: { ...baseTimelineInput.payload, @@ -1672,8 +1755,8 @@ describe("processTimelineResponse", () => { const userMessages = result.tail.filter((item) => item.kind === "user_message"); expect(userMessages).toHaveLength(2); - expect(userMessages.map((item) => item.id)).toEqual(["canonical-before", "optimistic-active"]); - expect(userMessages[1]?.optimistic).toBe(true); + expect(userMessages.map((item) => item.id)).toEqual(["canonical-before", "submitted-active"]); + expect(userMessages[1]?.clientMessageId).toBe("submitted-active"); }); it("leaves the cursor alone when a before page makes no progress", () => { @@ -1813,6 +1896,68 @@ describe("processTimelineResponse", () => { ]); }); + it("removes a reconciled submitted prompt before coalescing a tool call at the pagination seam", () => { + const clientMessageId = "client-boundary-prompt"; + const callId = "toolu_submitted_boundary"; + const currentTail = [ + makeSubmittedUserMessage("Inspect the file", clientMessageId), + ...hydrateStreamState( + [ + { + event: { + type: "timeline", + provider: "claude", + item: makeToolCallTimelineEntry(3, callId, "completed", { + type: "read", + filePath: "/tmp/example.ts", + }).item, + } as AgentStreamEventPayload, + timestamp: new Date(3000), + }, + ], + { source: "canonical" }, + ), + ]; + + const result = processTimelineResponse({ + ...baseTimelineInput, + currentTail, + currentCursor: { epoch: "epoch-1", startSeq: 3, endSeq: 5 }, + payload: { + ...baseTimelineInput.payload, + direction: "before", + epoch: "epoch-1", + startCursor: { seq: 1 }, + endCursor: { seq: 2 }, + entries: [ + { + ...makeTimelineEntry(1, "Inspect the file", "user_message"), + item: { + type: "user_message", + text: "Inspect the file", + messageId: "provider-boundary-prompt", + clientMessageId, + }, + }, + makeToolCallTimelineEntry(2, callId, "running", { + type: "unknown", + input: { file_path: "/tmp/example.ts" }, + output: null, + }), + ], + }, + }); + + expect(result.tail.filter((item) => item.kind === "user_message")).toHaveLength(1); + expect(getAgentToolCalls(result.tail)).toEqual([ + expect.objectContaining({ + payload: expect.objectContaining({ + data: expect.objectContaining({ callId, status: "completed" }), + }), + }), + ]); + }); + it("does not coalesce tool call lifecycle rows away from the prepend boundary", () => { const callId = "toolu_not_boundary"; const currentTail = hydrateStreamState( @@ -2225,7 +2370,7 @@ describe("processAgentStreamEvent", () => { }); }); - it("derives optimistic idle status on turn_completed for running agent", () => { + it("derives submitted idle status on turn_completed for running agent", () => { const turnCompletedEvent: AgentStreamEventPayload = { type: "turn_completed", provider: "claude", @@ -2249,7 +2394,7 @@ describe("processAgentStreamEvent", () => { expect(result.agent!.lastActivityAt.getTime()).toBe(2000); }); - it("derives optimistic error status on turn_failed for running agent", () => { + it("derives submitted error status on turn_failed for running agent", () => { const turnFailedEvent: AgentStreamEventPayload = { type: "turn_failed", provider: "claude", @@ -2271,7 +2416,7 @@ describe("processAgentStreamEvent", () => { expect(result.agent!.status).toBe("error"); }); - it("does not derive optimistic idle status on turn_canceled for running agent", () => { + it("does not derive submitted idle status on turn_canceled for running agent", () => { const turnCanceledEvent: AgentStreamEventPayload = { type: "turn_canceled", provider: "codex", @@ -2655,7 +2800,7 @@ describe("processAgentStreamEvents", () => { ]); }); - it("returns the final optimistic lifecycle patch across a batch", () => { + it("returns the final submitted lifecycle patch across a batch", () => { const result = processAgentStreamEvents({ events: [ makeStreamReducerEvent(makeTimelineEvent("Done"), 1), diff --git a/packages/app/src/timeline/session-stream-reducers.ts b/packages/app/src/timeline/session-stream-reducers.ts index 7458e855d3..533896e95c 100644 --- a/packages/app/src/timeline/session-stream-reducers.ts +++ b/packages/app/src/timeline/session-stream-reducers.ts @@ -2,14 +2,16 @@ import type { AgentStreamEventPayload } from "@getpaseo/protocol/messages"; import type { AgentLifecycleStatus } from "@getpaseo/protocol/agent-lifecycle"; import type { Agent } from "@/stores/session-store"; import { useSessionStore } from "@/stores/session-store"; -import type { AssistantMessageItem, StreamItem, UserMessageItem } from "@/types/stream"; +import type { AssistantMessageItem, StreamItem } from "@/types/stream"; import { applyStreamEvent, flushHeadToTail, hydrateStreamState, isAgentToolCallItem, mergeAgentToolCallItem, + replaceWithCanonicalStream, reduceStreamUpdate, + upsertUserMessageAcrossStream, } from "@/types/stream"; const AGENT_STREAM_REDUCER_FLUSH_DELAY_MS = 16 * 3; @@ -88,6 +90,7 @@ export interface ProcessTimelineResponseInput { isInitializing: boolean; hasActiveInitDeferred: boolean; initRequestDirection: InitRequestDirection; + pendingClientMessageId: string | null; } export interface ProcessTimelineResponseOutput { @@ -201,7 +204,7 @@ function shouldResolveTimelineInit({ return responseDirection === initRequestDirection; } -function deriveOptimisticLifecycleStatus( +function deriveProjectedLifecycleStatus( currentStatus: AgentLifecycleStatus, event: AgentStreamEventPayload, ): AgentLifecycleStatus | null { @@ -222,55 +225,35 @@ function deriveOptimisticLifecycleStatus( } } -function preserveReplacePathAssistantHead(params: { - tail: StreamItem[]; - currentHead: StreamItem[]; -}): { - tail: StreamItem[]; - head: StreamItem[]; -} { - const { tail, currentHead } = params; - const liveAssistant = currentHead.findLast( - (item): item is Extract => - item.kind === "assistant_message", - ); - if (!liveAssistant) { - return { tail, head: [] }; - } - const tailAssistant = tail.at(-1); - if (!tailAssistant || tailAssistant.kind !== "assistant_message") { - return { tail, head: currentHead }; - } - if (!liveAssistant.text.startsWith(tailAssistant.text)) { - return { tail, head: [] }; - } - return { - tail: tail.slice(0, -1), - head: [{ ...liveAssistant, text: tailAssistant.text }], - }; -} - function applyTimelineReplacePath(args: { timelineUnits: TimelineUnit[]; payload: ProcessTimelineResponseInput["payload"]; bootstrapPolicy: ReturnType; currentTail: StreamItem[]; currentHead: StreamItem[]; + pendingClientMessageId: string | null; + replacementPolicy: "preserve-local-submissions" | "replace-history"; toHydratedEvents: ( units: TimelineUnit[], ) => Array<{ event: AgentStreamEventPayload; timestamp: Date }>; }): TimelinePathResult { - const { timelineUnits, payload, bootstrapPolicy, currentTail, currentHead, toHydratedEvents } = - args; + const { + timelineUnits, + payload, + bootstrapPolicy, + currentTail, + currentHead, + pendingClientMessageId, + replacementPolicy, + toHydratedEvents, + } = args; const hydratedTail = hydrateStreamState(toHydratedEvents(timelineUnits), { source: "canonical" }); - const reconciledTail = reconcileLocalUserPresentationAfterReplace({ - canonicalTail: hydratedTail, + const { tail, head } = replaceWithCanonicalStream({ + canonical: hydratedTail, previousTail: currentTail, previousHead: currentHead, - }); - const { tail, head } = preserveReplacePathAssistantHead({ - tail: reconciledTail, - currentHead, + pendingClientMessageId, + policy: replacementPolicy, }); const cursor: TimelineCursor | null = payload.startCursor && payload.endCursor @@ -287,133 +270,6 @@ function applyTimelineReplacePath(args: { return { tail, head, cursor, cursorChanged: true, sideEffects }; } -function collectLocallyPresentedUserMessages(items: StreamItem[]): Array<{ - ordinal: number; - item: UserMessageItem; -}> { - const localUsers: Array<{ ordinal: number; item: UserMessageItem }> = []; - let ordinal = 0; - for (const item of items) { - if (item.kind !== "user_message") { - continue; - } - if (item.optimistic || item.images?.length || item.attachments?.length) { - localUsers.push({ ordinal, item }); - } - ordinal += 1; - } - return localUsers; -} - -function mergeCanonicalUserWithLocalPresentation( - canonical: UserMessageItem, - local: UserMessageItem, -): UserMessageItem { - return { - kind: "user_message", - id: canonical.id, - ...(canonical.clientMessageId ? { clientMessageId: canonical.clientMessageId } : {}), - text: local.text, - timestamp: local.timestamp, - ...(local.images && local.images.length > 0 ? { images: local.images } : {}), - ...(local.attachments && local.attachments.length > 0 - ? { attachments: local.attachments } - : {}), - }; -} - -interface CanonicalUserMessageIdentity { - messageId?: string; - clientMessageId?: string; - text: string; -} - -function matchesLocalUserMessageIdentity( - canonical: CanonicalUserMessageIdentity, - optimistic: UserMessageItem, -): boolean { - if (canonical.clientMessageId !== undefined) { - return canonical.clientMessageId === optimistic.id; - } - if (canonical.messageId === optimistic.id) { - return true; - } - // COMPAT(userMessageClientId): added in v0.2.0, remove after 2027-01-20 once - // the supported daemon floor emits clientMessageId on submitted user messages. - return canonical.text.length > 0 && canonical.text === optimistic.text; -} - -function reconcileLocalUserPresentationAfterReplace(params: { - canonicalTail: StreamItem[]; - previousTail: StreamItem[]; - previousHead: StreamItem[]; -}): StreamItem[] { - const localUsers = collectLocallyPresentedUserMessages([ - ...params.previousTail, - ...params.previousHead, - ]); - if (localUsers.length === 0) { - return params.canonicalTail; - } - - const canonicalUserIndexes: number[] = []; - params.canonicalTail.forEach((item, index) => { - if (item.kind === "user_message") { - canonicalUserIndexes.push(index); - } - }); - - const nextTail = [...params.canonicalTail]; - const claimedCanonicalIndexes = new Set(); - const unmatched: UserMessageItem[] = []; - - for (const local of localUsers) { - const exactIndex = canonicalUserIndexes.find((index) => { - if (claimedCanonicalIndexes.has(index)) return false; - const canonical = params.canonicalTail[index]; - return ( - canonical?.kind === "user_message" && - matchesLocalUserMessageIdentity( - { - messageId: canonical.id, - clientMessageId: canonical.clientMessageId, - text: canonical.text, - }, - local.item, - ) - ); - }); - const ordinalIndex = canonicalUserIndexes[local.ordinal]; - const ordinalItem = ordinalIndex === undefined ? undefined : params.canonicalTail[ordinalIndex]; - const canonicalIndex = - exactIndex ?? - (ordinalIndex !== undefined && - !claimedCanonicalIndexes.has(ordinalIndex) && - ordinalItem?.kind === "user_message" && - ordinalItem.clientMessageId === undefined - ? ordinalIndex - : undefined); - const canonicalItem = canonicalIndex === undefined ? undefined : nextTail[canonicalIndex]; - if (canonicalIndex === undefined || !canonicalItem || canonicalItem.kind !== "user_message") { - if (local.item.optimistic) { - unmatched.push(local.item); - } - continue; - } - nextTail[canonicalIndex] = mergeCanonicalUserWithLocalPresentation(canonicalItem, local.item); - claimedCanonicalIndexes.add(canonicalIndex); - } - - for (const item of unmatched) { - const insertionIndex = nextTail.findIndex( - (canonical) => canonical.timestamp.getTime() > item.timestamp.getTime(), - ); - nextTail.splice(insertionIndex < 0 ? nextTail.length : insertionIndex, 0, item); - } - - return nextTail; -} - interface IncrementalAcceptResult { acceptedUnits: TimelineUnit[]; cursor: TimelineCursor | undefined; @@ -515,6 +371,34 @@ function mergePrependedCanonicalTail(olderTail: StreamItem[], currentTail: Strea return olderTail; } + const remainingOlder: StreamItem[] = []; + let reconciledCurrent = currentTail; + for (const item of olderTail) { + if (item.kind !== "user_message") { + remainingOlder.push(item); + continue; + } + const result = upsertUserMessageAcrossStream({ + tail: reconciledCurrent, + head: [], + message: item, + insert: "prepend-tail", + presentation: "existing", + }); + if (result.location?.matched) { + remainingOlder.push(result.location.message); + reconciledCurrent = [ + ...result.tail.slice(0, result.location.index), + ...result.tail.slice(result.location.index + 1), + ]; + } else { + remainingOlder.push(item); + } + } + olderTail = remainingOlder; + currentTail = reconciledCurrent; + if (olderTail.length === 0) return currentTail; + const olderLast = olderTail.at(-1); const currentFirst = currentTail[0]; @@ -906,6 +790,7 @@ export function processTimelineResponse( isInitializing, hasActiveInitDeferred, initRequestDirection, + pendingClientMessageId, } = input; // ------------------------------------------------------------------ @@ -967,6 +852,10 @@ export function processTimelineResponse( hasActiveInitDeferred, }); const replace = bootstrapPolicy.replace; + const replacementPolicy = + currentCursor && currentCursor.epoch !== payload.epoch + ? "replace-history" + : "preserve-local-submissions"; const sideEffects: TimelineReducerSideEffect[] = []; const timelineResult = replace @@ -976,6 +865,8 @@ export function processTimelineResponse( bootstrapPolicy, currentTail, currentHead, + pendingClientMessageId, + replacementPolicy, toHydratedEvents, }) : applyTimelineIncrementalPath({ @@ -1230,7 +1121,7 @@ export function processAgentStreamEvent( }; // ------------------------------------------------------------------ - // Optimistic lifecycle status + // Projected lifecycle status // ------------------------------------------------------------------ let agentPatch: AgentPatch | null = null; let agentChanged = false; @@ -1241,15 +1132,15 @@ export function processAgentStreamEvent( event.type === "turn_canceled" || event.type === "turn_failed") ) { - const optimisticStatus = deriveOptimisticLifecycleStatus(currentAgent.status, event); - if (optimisticStatus) { + const projectedStatus = deriveProjectedLifecycleStatus(currentAgent.status, event); + if (projectedStatus) { const nextUpdatedAtMs = Math.max(currentAgent.updatedAt.getTime(), timestamp.getTime()); const nextLastActivityAtMs = Math.max( currentAgent.lastActivityAt.getTime(), timestamp.getTime(), ); agentPatch = { - status: optimisticStatus, + status: projectedStatus, updatedAt: new Date(nextUpdatedAtMs), lastActivityAt: new Date(nextLastActivityAtMs), }; diff --git a/packages/app/src/timeline/turn-time.test.ts b/packages/app/src/timeline/turn-time.test.ts index d63d2de7d5..3945edbcb7 100644 --- a/packages/app/src/timeline/turn-time.test.ts +++ b/packages/app/src/timeline/turn-time.test.ts @@ -22,34 +22,16 @@ function assistant(id: string, timestamp: Date): StreamItem { } describe("deriveStreamTurnTiming", () => { - it("reserves a running footer for an optimistic prompt before the host starts the turn", () => { - const optimisticPrompt = { - ...user("optimistic", new Date("2026-05-15T00:00:00.000Z")), - optimistic: true as const, - }; - - const timing = deriveStreamTurnTiming({ - agentStatus: "idle", - tail: [], - head: [optimisticPrompt], - }); - - assert.equal(timing.isActive, true); - }); - - it("does not start elapsed time from an optimistic prompt", () => { - const optimisticPrompt = { - ...user("optimistic", new Date("2026-05-15T00:00:00.000Z")), - optimistic: true as const, - }; + it("starts elapsed time from the submitted prompt", () => { + const submittedAt = new Date("2026-05-15T00:00:00.000Z"); const timing = deriveStreamTurnTiming({ agentStatus: "running", tail: [], - head: [optimisticPrompt], + head: [user("submitted", submittedAt)], }); - assert.equal(timing.runningStartedAt, null); + assert.equal(timing.runningStartedAt, submittedAt); }); it("uses the last user message as the running turn start", () => { diff --git a/packages/app/src/timeline/turn-time.ts b/packages/app/src/timeline/turn-time.ts index 4356945243..7ebd77c214 100644 --- a/packages/app/src/timeline/turn-time.ts +++ b/packages/app/src/timeline/turn-time.ts @@ -9,7 +9,6 @@ export interface TurnTiming { export interface StreamTurnTiming { byAssistantId: Map; runningStartedAt: Date | null; - isActive: boolean; } export function deriveStreamTurnTiming(params: { @@ -19,8 +18,6 @@ export function deriveStreamTurnTiming(params: { }): StreamTurnTiming { const byAssistantId = new Map(); let currentUserAt: Date | null = null; - let currentAuthoritativeUserAt: Date | null = null; - let currentUserIsOptimistic = false; let currentLastItemAt: Date | null = null; let currentAssistantIds: string[] = []; @@ -42,8 +39,6 @@ export function deriveStreamTurnTiming(params: { if (item.kind === "user_message") { flushCompletedTurn(); currentUserAt = item.timestamp; - currentAuthoritativeUserAt = item.optimistic ? null : item.timestamp; - currentUserIsOptimistic = item.optimistic === true; currentLastItemAt = null; currentAssistantIds = []; return; @@ -65,7 +60,7 @@ export function deriveStreamTurnTiming(params: { } const isRunning = params.agentStatus === "running"; - const runningStartedAt = isRunning ? currentAuthoritativeUserAt : null; + const runningStartedAt = isRunning ? currentUserAt : null; if (params.agentStatus !== "running") { flushCompletedTurn(); } @@ -73,6 +68,5 @@ export function deriveStreamTurnTiming(params: { return { byAssistantId, runningStartedAt, - isActive: isRunning || currentUserIsOptimistic, }; } diff --git a/packages/app/src/types/stream.test.ts b/packages/app/src/types/stream.test.ts index fb68d7a020..36c277b52f 100644 --- a/packages/app/src/types/stream.test.ts +++ b/packages/app/src/types/stream.test.ts @@ -4,9 +4,7 @@ import { describe, expect, it } from "vitest"; import { applyStreamEvent, - appendOptimisticUserMessageToStream, - buildOptimisticUserMessage, - clearOptimisticUserMessages, + createUserMessage, handoffCreatedAgentUserMessageToStream, hydrateStreamState, mergeToolCallDetail, @@ -14,6 +12,8 @@ import { type AgentToolCallItem, type StreamItem, isAgentToolCallItem, + upsertUserMessage, + upsertUserMessageAcrossStream, } from "./stream"; import type { AgentProvider, ToolCallDetail } from "@getpaseo/protocol/agent-types"; import type { AgentStreamEventPayload } from "@getpaseo/protocol/messages"; @@ -21,6 +21,109 @@ import { buildToolCallDisplayModel } from "@getpaseo/protocol/tool-call-display" type CanonicalToolStatus = "running" | "completed" | "failed" | "canceled"; +describe("user message identity", () => { + it("adds provider identity without replacing local presentation", () => { + const timestamp = new Date("2026-07-26T10:00:00.000Z"); + const local = createUserMessage({ + clientMessageId: "client-1", + text: "local text", + timestamp, + images: [ + { + id: "image-1", + mimeType: "image/png", + storageType: "web-indexeddb", + storageKey: "image-1.png", + createdAt: timestamp.getTime(), + }, + ], + attachments: [{ type: "text", mimeType: "text/plain", text: "attachment" }], + }); + const canonical = createUserMessage({ + id: "provider-1", + messageId: "provider-1", + clientMessageId: "client-1", + text: "provider text", + timestamp: new Date("2026-07-26T10:00:01.000Z"), + }); + + const first = upsertUserMessage([local], canonical); + const second = upsertUserMessage(first, canonical); + + expect(first).toEqual([ + { + ...local, + messageId: "provider-1", + clientMessageId: "client-1", + }, + ]); + expect(first[0]).toBe(second[0]); + }); + + it("keeps local presentation when a later canonical row omits provider identity", () => { + const timestamp = new Date("2026-07-27T10:00:00.000Z"); + const local = createUserMessage({ + clientMessageId: "client-1", + messageId: "provider-1", + text: "local text", + timestamp, + images: [ + { + id: "image-1", + mimeType: "image/png", + storageType: "web-indexeddb", + storageKey: "image-1.png", + createdAt: timestamp.getTime(), + }, + ], + attachments: [{ type: "text", mimeType: "text/plain", text: "local attachment" }], + }); + const canonicalWithoutProviderIdentity = createUserMessage({ + id: "canonical-page-row", + clientMessageId: "client-1", + text: "provider-shaped text", + timestamp: new Date("2026-07-27T10:00:01.000Z"), + }); + + const result = upsertUserMessage([local], canonicalWithoutProviderIdentity); + + expect(result).toEqual([local]); + }); + + it("matches a submitted message against a legacy canonical row that has no client identity", () => { + // Daemons before v0.2.0 do not echo clientMessageId. During agent creation the + // legacy canonical row can land before the local submission is handed off, so the + // submitted row arrives as `incoming` and must still match by text. + const timestamp = new Date("2026-07-27T11:00:00.000Z"); + const legacyCanonical = createUserMessage({ + id: "provider-1", + messageId: "provider-1", + text: "review this", + timestamp, + }); + const submitted = createUserMessage({ + clientMessageId: "client-1", + text: "review this", + timestamp: new Date("2026-07-27T11:00:01.000Z"), + attachments: [{ type: "text", mimeType: "text/plain", text: "attachment" }], + }); + + const result = handoffCreatedAgentUserMessageToStream({ + tail: [legacyCanonical], + head: [], + message: submitted, + }); + + expect(result.tail).toEqual([ + { + ...submitted, + id: "client-1", + messageId: "provider-1", + }, + ]); + }); +}); + function assistantTimeline( text: string, provider: AgentProvider = "claude", @@ -895,15 +998,15 @@ describe("stream reducer canonical tool calls", () => { assert.strictEqual(todos.items[0]?.text, "Task 1"); }); - it("preserves optimistic user message images when authoritative user message arrives", () => { + it("preserves submitted user message images when authoritative user message arrives", () => { const messageId = "msg-user-images"; - const optimisticTimestamp = new Date("2025-01-01T11:10:00Z"); - const optimisticImages = [ + const submittedTimestamp = new Date("2025-01-01T11:10:00Z"); + const submittedImages = [ { - id: "att-optimistic", + id: "att-submitted", mimeType: "image/jpeg", storageType: "native-file" as const, - storageKey: "/tmp/optimistic.jpg", + storageKey: "/tmp/submitted.jpg", createdAt: Date.now(), }, ]; @@ -911,10 +1014,10 @@ describe("stream reducer canonical tool calls", () => { { kind: "user_message", id: messageId, + clientMessageId: messageId, text: "Analyze this image", - timestamp: optimisticTimestamp, - optimistic: true, - images: optimisticImages, + timestamp: submittedTimestamp, + images: submittedImages, }, ]; const event: AgentStreamEventPayload = { @@ -933,9 +1036,9 @@ describe("stream reducer canonical tool calls", () => { assert.ok(message); assert.strictEqual(message.id, messageId); - assert.deepStrictEqual(message.images, optimisticImages); + assert.deepStrictEqual(message.images, submittedImages); assert.strictEqual(message.text, "Analyze this image"); - assert.strictEqual(message.timestamp.getTime(), optimisticTimestamp.getTime()); + assert.strictEqual(message.timestamp.getTime(), submittedTimestamp.getTime()); }); it("keeps canonical assistant/user/assistant order during replay", () => { @@ -981,7 +1084,7 @@ describe("stream reducer canonical tool calls", () => { ); }); - it("keeps live optimistic assistant merge behavior", () => { + it("keeps live submitted assistant merge behavior", () => { const state: StreamItem[] = [ { kind: "assistant_message", @@ -1117,23 +1220,23 @@ describe("turn lifecycle events", () => { }); it.each(["codex", "opencode", "pi"] satisfies AgentProvider[])( - "replaces an optimistic user message when a live %s provider-owned id echo arrives without text matching", + "replaces a submitted user message when a live %s provider-owned id echo arrives without text matching", (provider) => { - const optimisticTimestamp = new Date("2025-01-01T15:02:00Z"); + const submittedTimestamp = new Date("2025-01-01T15:02:00Z"); const serverTimestamp = new Date("2025-01-01T15:02:01Z"); - const optimistic: StreamItem = { + const submitted: StreamItem = { kind: "user_message", - id: "msg_optimistic", + id: "msg_submitted", + clientMessageId: "msg_submitted", text: "same user text", - timestamp: optimisticTimestamp, - optimistic: true, + timestamp: submittedTimestamp, images: [ { id: "image-1", mimeType: "image/png", storageType: "web-indexeddb", storageKey: "image-1", - createdAt: optimisticTimestamp.getTime(), + createdAt: submittedTimestamp.getTime(), }, ], attachments: [ @@ -1147,7 +1250,7 @@ describe("turn lifecycle events", () => { }; const state = reduceStreamUpdate( - [optimistic], + [submitted], { type: "timeline", provider, @@ -1155,6 +1258,7 @@ describe("turn lifecycle events", () => { type: "user_message", text: "server-owned rendered text", messageId: "provider-owned-id", + clientMessageId: "msg_submitted", }, }, serverTimestamp, @@ -1165,28 +1269,28 @@ describe("turn lifecycle events", () => { assert.strictEqual(userMessages.length, 1); const userMessage = userMessages[0]; invariant(userMessage?.kind === "user_message"); - assert.strictEqual(userMessage.id, "provider-owned-id"); - assert.strictEqual(userMessage.text, optimistic.text); - assert.strictEqual(userMessage.timestamp.getTime(), optimistic.timestamp.getTime()); - assert.strictEqual(userMessage.optimistic, undefined); - assert.deepStrictEqual(userMessage.images, optimistic.images); - assert.deepStrictEqual(userMessage.attachments, optimistic.attachments); + assert.strictEqual(userMessage.id, "msg_submitted"); + assert.strictEqual(userMessage.messageId, "provider-owned-id"); + assert.strictEqual(userMessage.text, submitted.text); + assert.strictEqual(userMessage.timestamp.getTime(), submitted.timestamp.getTime()); + assert.deepStrictEqual(userMessage.images, submitted.images); + assert.deepStrictEqual(userMessage.attachments, submitted.attachments); }, ); - it("replaces one optimistic plain-text user message with the next live server user message", () => { - const optimisticTimestamp = new Date("2025-01-01T15:03:00Z"); + it("replaces one submitted plain-text user message with the next live server user message", () => { + const submittedTimestamp = new Date("2025-01-01T15:03:00Z"); const serverTimestamp = new Date("2025-01-01T15:03:01Z"); - const optimistic: StreamItem = { + const submitted: StreamItem = { kind: "user_message", - id: "msg_optimistic", + id: "msg_submitted", + clientMessageId: "msg_submitted", text: "typed plain text", - timestamp: optimisticTimestamp, - optimistic: true, + timestamp: submittedTimestamp, }; const state = reduceStreamUpdate( - [optimistic], + [submitted], { type: "timeline", provider: "opencode", @@ -1204,20 +1308,20 @@ describe("turn lifecycle events", () => { assert.strictEqual(userMessages.length, 1); const userMessage = userMessages[0]; invariant(userMessage?.kind === "user_message"); - assert.strictEqual(userMessage.id, "msg_opencode_provider_owned"); + assert.strictEqual(userMessage.id, "msg_submitted"); + assert.strictEqual(userMessage.messageId, "msg_opencode_provider_owned"); assert.strictEqual(userMessage.text, "typed plain text"); - assert.strictEqual(userMessage.timestamp.getTime(), optimisticTimestamp.getTime()); - assert.strictEqual(userMessage.optimistic, undefined); + assert.strictEqual(userMessage.timestamp.getTime(), submittedTimestamp.getTime()); }); - it("replaces an optimistic image user message with the next canonical server user message", () => { - const optimisticTimestamp = new Date("2025-01-01T15:03:10Z"); + it("replaces a submitted image user message with the next canonical server user message", () => { + const submittedTimestamp = new Date("2025-01-01T15:03:10Z"); const image = { id: "image-canonical", mimeType: "image/png", storageType: "web-indexeddb" as const, storageKey: "image-canonical", - createdAt: optimisticTimestamp.getTime(), + createdAt: submittedTimestamp.getTime(), }; const attachment = { type: "text" as const, @@ -1225,16 +1329,16 @@ describe("turn lifecycle events", () => { text: "context", title: "context.txt", }; - const optimistic = buildOptimisticUserMessage({ - id: "msg_optimistic_canonical", + const submitted = createUserMessage({ + clientMessageId: "msg_submitted_canonical", text: "Analyze this", - timestamp: optimisticTimestamp, + timestamp: submittedTimestamp, images: [image], attachments: [attachment], }); const state = reduceStreamUpdate( - [optimistic], + [submitted], { type: "timeline", provider: "claude", @@ -1242,7 +1346,7 @@ describe("turn lifecycle events", () => { type: "user_message", text: "server-rendered attachment text", messageId: "provider-owned-canonical", - clientMessageId: optimistic.id, + clientMessageId: submitted.id, }, }, new Date("2025-01-01T15:03:11Z"), @@ -1253,17 +1357,17 @@ describe("turn lifecycle events", () => { assert.strictEqual(userMessages.length, 1); const userMessage = userMessages[0]; invariant(userMessage?.kind === "user_message"); - assert.strictEqual(userMessage.id, "provider-owned-canonical"); + assert.strictEqual(userMessage.id, "msg_submitted_canonical"); + assert.strictEqual(userMessage.messageId, "provider-owned-canonical"); assert.strictEqual(userMessage.text, "Analyze this"); - assert.strictEqual(userMessage.timestamp.getTime(), optimisticTimestamp.getTime()); - assert.strictEqual(userMessage.optimistic, undefined); + assert.strictEqual(userMessage.timestamp.getTime(), submittedTimestamp.getTime()); assert.deepStrictEqual(userMessage.images, [image]); assert.deepStrictEqual(userMessage.attachments, [attachment]); }); - it("places optimistic user messages through one append helper", () => { - const optimistic = buildOptimisticUserMessage({ - id: "msg_append_once", + it("places submitted user messages through the identity producer", () => { + const submitted = createUserMessage({ + clientMessageId: "msg_append_once", text: "append once", timestamp: new Date("2025-01-01T15:03:20Z"), }); @@ -1274,28 +1378,30 @@ describe("turn lifecycle events", () => { timestamp: new Date("2025-01-01T15:03:19Z"), }; - const first = appendOptimisticUserMessageToStream({ + const first = upsertUserMessageAcrossStream({ tail: [], head: [headItem], - message: optimistic, - placement: "active-head", + message: submitted, + insert: "head", + presentation: "existing", }); - const second = appendOptimisticUserMessageToStream({ + const second = upsertUserMessageAcrossStream({ tail: first.tail, head: first.head, - message: optimistic, - placement: "active-head", + message: submitted, + insert: "head", + presentation: "existing", }); assert.deepStrictEqual(first.tail, []); - assert.deepStrictEqual(first.head, [headItem, optimistic]); + assert.deepStrictEqual(first.head, [headItem, submitted]); assert.strictEqual(second.changedHead, false); assert.strictEqual(second.head, first.head); }); - it("hands rich optimistic content to an authoritative create message without duplicating it", () => { + it("hands rich submitted content to its create message without overwriting an earlier user row", () => { const timestamp = new Date("2025-01-01T15:03:20Z"); - const optimistic = buildOptimisticUserMessage({ - id: "client-user", + const submitted = createUserMessage({ + clientMessageId: "client-user", text: "", timestamp, images: [ @@ -1317,32 +1423,44 @@ describe("turn lifecycle events", () => { }, ], }); + const precedingProviderRow: StreamItem = { + kind: "user_message", + id: "provider-system-user", + messageId: "provider-system-user", + text: "provider setup prompt", + timestamp: new Date("2025-01-01T15:03:20.500Z"), + }; const canonical: StreamItem = { kind: "user_message", id: "provider-user", + messageId: "provider-user", + clientMessageId: "client-user", text: "server-rendered attachment text", timestamp: new Date("2025-01-01T15:03:21Z"), }; const handedOff = handoffCreatedAgentUserMessageToStream({ - tail: [canonical], + tail: [precedingProviderRow, canonical], head: [], - message: optimistic, + message: submitted, }); const repeated = handoffCreatedAgentUserMessageToStream({ tail: handedOff.tail, head: handedOff.head, - message: optimistic, + message: submitted, }); assert.deepStrictEqual(handedOff.tail, [ + precedingProviderRow, { kind: "user_message", - id: "provider-user", - text: optimistic.text, - timestamp: optimistic.timestamp, - images: optimistic.images, - attachments: optimistic.attachments, + id: "client-user", + clientMessageId: "client-user", + messageId: "provider-user", + text: submitted.text, + timestamp: submitted.timestamp, + images: submitted.images, + attachments: submitted.attachments, }, ]); assert.deepStrictEqual(handedOff.head, []); @@ -1364,22 +1482,22 @@ describe("turn lifecycle events", () => { ); assert.deepStrictEqual( afterNextUser.filter((item) => item.kind === "user_message").map((item) => item.id), - ["provider-user", "provider-next-user"], + ["provider-system-user", "client-user", "provider-next-user"], ); }); - it("reconciles an optimistic user message that was pending in the streaming head", () => { - const optimistic: StreamItem = { + it("reconciles a submitted user message that was pending in the streaming head", () => { + const submitted: StreamItem = { kind: "user_message", - id: "msg_head_optimistic", + id: "msg_head_submitted", + clientMessageId: "msg_head_submitted", text: "plain text in head", timestamp: new Date("2025-01-01T15:03:02Z"), - optimistic: true, }; const result = applyStreamEvent({ tail: [], - head: [optimistic], + head: [submitted], event: { type: "timeline", provider: "opencode", @@ -1393,33 +1511,33 @@ describe("turn lifecycle events", () => { source: "live", }); - assert.strictEqual(result.head.length, 0); - const userMessages = result.tail.filter((item) => item.kind === "user_message"); + assert.strictEqual(result.tail.length, 0); + const userMessages = result.head.filter((item) => item.kind === "user_message"); assert.strictEqual(userMessages.length, 1); - assert.strictEqual(userMessages[0]?.id, "provider-owned-head"); - assert.strictEqual(userMessages[0]?.optimistic, undefined); + assert.strictEqual(userMessages[0]?.id, "msg_head_submitted"); + assert.strictEqual(userMessages[0]?.messageId, "provider-owned-head"); }); - it("replaces multiple optimistic user messages in FIFO order", () => { - const optimisticTimestamp = new Date("2025-01-01T15:04:00Z"); + it("replaces multiple submitted user messages in FIFO order", () => { + const submittedTimestamp = new Date("2025-01-01T15:04:00Z"); const serverTimestamp = new Date("2025-01-01T15:04:01Z"); - const firstOptimistic: StreamItem = { + const firstSubmitted: StreamItem = { kind: "user_message", - id: "msg_optimistic_1", + id: "msg_submitted_1", + clientMessageId: "msg_submitted_1", text: "first typed text", - timestamp: optimisticTimestamp, - optimistic: true, + timestamp: submittedTimestamp, }; - const secondOptimistic: StreamItem = { + const secondSubmitted: StreamItem = { kind: "user_message", - id: "msg_optimistic_2", + id: "msg_submitted_2", + clientMessageId: "msg_submitted_2", text: "second typed text", timestamp: new Date("2025-01-01T15:04:00.500Z"), - optimistic: true, }; const afterFirstEcho = reduceStreamUpdate( - [firstOptimistic, secondOptimistic], + [firstSubmitted, secondSubmitted], { type: "timeline", provider: "opencode", @@ -1427,6 +1545,7 @@ describe("turn lifecycle events", () => { type: "user_message", text: "first server text", messageId: "provider-owned-first", + clientMessageId: "msg_submitted_1", }, }, serverTimestamp, @@ -1441,6 +1560,7 @@ describe("turn lifecycle events", () => { type: "user_message", text: "second server text", messageId: "provider-owned-second", + clientMessageId: "msg_submitted_2", }, }, new Date("2025-01-01T15:04:02Z"), @@ -1450,30 +1570,30 @@ describe("turn lifecycle events", () => { const userMessages = state.filter((item) => item.kind === "user_message"); assert.strictEqual(userMessages.length, 2); assert.deepStrictEqual( - userMessages.map((item) => [item.id, item.text, item.optimistic]), + userMessages.map((item) => [item.id, item.text, item.messageId]), [ - ["provider-owned-first", "first typed text", undefined], - ["provider-owned-second", "second typed text", undefined], + ["msg_submitted_1", "first typed text", "provider-owned-first"], + ["msg_submitted_2", "second typed text", "provider-owned-second"], ], ); }); - it("does not shift later prompts when an earlier optimistic prompt has no canonical echo", () => { + it("does not shift later prompts when an earlier submitted prompt has no canonical echo", () => { const staleTimestamp = new Date("2025-01-01T15:04:00Z"); const submittedTimestamp = new Date("2025-01-01T15:04:01Z"); const stalePrompt: StreamItem = { kind: "user_message", id: "msg_stale", + clientMessageId: "msg_stale", text: "first prompt without an echo", timestamp: staleTimestamp, - optimistic: true, }; const submittedPrompt: StreamItem = { kind: "user_message", id: "msg_submitted", + clientMessageId: "msg_submitted", text: "later submitted prompt", timestamp: submittedTimestamp, - optimistic: true, }; const state = reduceStreamUpdate( @@ -1496,15 +1616,16 @@ describe("turn lifecycle events", () => { stalePrompt, { kind: "user_message", - id: "provider-owned-submitted", + id: "msg_submitted", clientMessageId: submittedPrompt.id, + messageId: "provider-owned-submitted", text: submittedPrompt.text, timestamp: submittedPrompt.timestamp, }, ]); }); - it("appends a live server user message when no optimistic user message is pending", () => { + it("appends a live server user message when no submitted user message is pending", () => { const state = reduceStreamUpdate( [], { @@ -1523,21 +1644,11 @@ describe("turn lifecycle events", () => { const userMessages = state.filter((item) => item.kind === "user_message"); assert.strictEqual(userMessages.length, 1); assert.strictEqual(userMessages[0]?.id, "provider-owned-resume"); - assert.strictEqual(userMessages[0]?.optimistic, undefined); }); - it("does not match a server user message to an optimistic from a rewound turn after pending optimistics are cleared", () => { - const optimistic: StreamItem = { - kind: "user_message", - id: "msg_rewound_optimistic", - text: "rewound text", - timestamp: new Date("2025-01-01T15:04:04Z"), - optimistic: true, - }; - const cleared = clearOptimisticUserMessages([optimistic]); - + it("appends a server user message after a rewound local row was removed", () => { const state = reduceStreamUpdate( - cleared, + [], { type: "timeline", provider: "opencode", @@ -1555,7 +1666,6 @@ describe("turn lifecycle events", () => { assert.strictEqual(userMessages.length, 1); assert.strictEqual(userMessages[0]?.id, "provider-owned-after-rewind"); assert.strictEqual(userMessages[0]?.text, "future server echo"); - assert.strictEqual(userMessages[0]?.optimistic, undefined); }); it("keeps canonical repeated user messages distinct during hydration", () => { diff --git a/packages/app/src/types/stream.ts b/packages/app/src/types/stream.ts index 76407d84cb..7e3bfcd218 100644 --- a/packages/app/src/types/stream.ts +++ b/packages/app/src/types/stream.ts @@ -87,22 +87,345 @@ export interface UserMessageItem { kind: "user_message"; id: string; clientMessageId?: string; - text: string; - timestamp: Date; - optimistic?: true; - images?: UserMessageImageAttachment[]; - attachments?: AgentAttachment[]; -} - -export interface OptimisticUserMessageInput { - id: string; + messageId?: string; text: string; timestamp: Date; images?: UserMessageImageAttachment[]; attachments?: AgentAttachment[]; } -export type OptimisticUserMessagePlacement = "tail" | "active-head"; +export interface UserMessageInput { + id?: string; + clientMessageId?: string; + messageId?: string; + text: string; + timestamp: Date; + images?: UserMessageImageAttachment[]; + attachments?: AgentAttachment[]; +} + +export function createUserMessage(input: UserMessageInput): UserMessageItem { + const id = input.id ?? input.clientMessageId ?? input.messageId; + if (!id) { + throw new Error("User message identity is required"); + } + return { + kind: "user_message", + id, + ...(input.clientMessageId ? { clientMessageId: input.clientMessageId } : {}), + ...(input.messageId ? { messageId: input.messageId } : {}), + text: input.text, + timestamp: input.timestamp, + ...(input.images && input.images.length > 0 ? { images: input.images } : {}), + ...(input.attachments && input.attachments.length > 0 + ? { attachments: input.attachments } + : {}), + }; +} + +// COMPAT(userMessageClientId): added in v0.2.0, remove after 2027-01-20 once the +// supported daemon floor emits clientMessageId on submitted user messages. Until then a +// locally submitted row (clientMessageId, no messageId) and its canonical twin from an +// old daemon (messageId, no clientMessageId) share no identifier, so they are matched by +// the id the client supplied over the wire or by text. Direction-independent: either row +// can already be in the timeline when the other arrives. +function matchesLegacyCanonicalUserMessage( + submitted: UserMessageItem, + canonical: UserMessageItem, +): boolean { + if (submitted.clientMessageId === undefined || submitted.messageId !== undefined) return false; + if (canonical.messageId === undefined) return false; + return canonical.messageId === submitted.clientMessageId || canonical.text === submitted.text; +} + +function matchesUserMessage(existing: UserMessageItem, incoming: UserMessageItem): boolean { + if (existing.clientMessageId && incoming.clientMessageId) { + return existing.clientMessageId === incoming.clientMessageId; + } + if (existing.messageId && incoming.messageId) { + return existing.messageId === incoming.messageId; + } + return ( + matchesLegacyCanonicalUserMessage(existing, incoming) || + matchesLegacyCanonicalUserMessage(incoming, existing) + ); +} + +export function upsertUserMessage( + items: StreamItem[], + incoming: UserMessageItem, + insertAt = items.length, +): StreamItem[] { + return produceUserMessage(items, incoming, insertAt, "existing").items; +} + +type UserMessagePresentationPolicy = "existing" | "incoming"; + +interface UserMessageProductionResult { + items: StreamItem[]; + index: number; + message: UserMessageItem; + matched: boolean; +} + +function produceUserMessage( + items: StreamItem[], + incoming: UserMessageItem, + insertAt: number | null, + presentationPolicy: UserMessagePresentationPolicy, +): UserMessageProductionResult { + const index = items.findIndex( + (item) => item.kind === "user_message" && matchesUserMessage(item, incoming), + ); + if (index < 0) { + if (insertAt === null) { + return { items, index: -1, message: incoming, matched: false }; + } + return { + items: [...items.slice(0, insertAt), incoming, ...items.slice(insertAt)], + index: insertAt, + message: incoming, + matched: false, + }; + } + + const existing = items[index]; + if (!existing || existing.kind !== "user_message") { + throw new Error("User message upsert matched a non-user row"); + } + const presentation = presentationPolicy === "incoming" ? incoming : existing; + const merged = createUserMessage({ + ...presentation, + clientMessageId: incoming.clientMessageId ?? existing.clientMessageId, + messageId: incoming.messageId ?? existing.messageId, + }); + if ( + existing.id === merged.id && + existing.clientMessageId === merged.clientMessageId && + existing.messageId === merged.messageId && + existing.text === merged.text && + existing.timestamp === merged.timestamp && + existing.images === merged.images && + existing.attachments === merged.attachments + ) { + return { items, index, message: existing, matched: true }; + } + const next = [...items]; + next[index] = merged; + return { items: next, index, message: merged, matched: true }; +} + +export interface UserMessageStreamUpsertInput { + tail: StreamItem[]; + head: StreamItem[]; + message: UserMessageItem; + insert: "tail" | "head" | "prepend-tail" | "none"; + presentation: UserMessagePresentationPolicy; +} + +export interface UserMessageStreamUpsertResult extends ApplyStreamEventResult { + location: { + lane: "tail" | "head"; + index: number; + message: UserMessageItem; + matched: boolean; + } | null; +} + +export function upsertUserMessageAcrossStream( + input: UserMessageStreamUpsertInput, +): UserMessageStreamUpsertResult { + const tailResult = produceUserMessage(input.tail, input.message, null, input.presentation); + if (tailResult.matched) { + return { + tail: tailResult.items, + head: input.head, + changedTail: tailResult.items !== input.tail, + changedHead: false, + location: { + lane: "tail", + index: tailResult.index, + message: tailResult.message, + matched: true, + }, + }; + } + const headResult = produceUserMessage(input.head, input.message, null, input.presentation); + if (headResult.matched) { + return { + tail: input.tail, + head: headResult.items, + changedTail: false, + changedHead: headResult.items !== input.head, + location: { + lane: "head", + index: headResult.index, + message: headResult.message, + matched: true, + }, + }; + } + if (input.insert === "none") { + return { + tail: input.tail, + head: input.head, + changedTail: false, + changedHead: false, + location: null, + }; + } + if (input.insert === "head") { + const inserted = produceUserMessage( + input.head, + input.message, + input.head.length, + input.presentation, + ); + return { + tail: input.tail, + head: inserted.items, + changedTail: false, + changedHead: true, + location: { + lane: "head", + index: inserted.index, + message: inserted.message, + matched: false, + }, + }; + } + const inserted = produceUserMessage( + input.tail, + input.message, + input.insert === "prepend-tail" ? 0 : input.tail.length, + input.presentation, + ); + return { + tail: inserted.items, + head: input.head, + changedTail: true, + changedHead: false, + location: { + lane: "tail", + index: inserted.index, + message: inserted.message, + matched: false, + }, + }; +} + +export interface CanonicalStreamReplacementInput { + canonical: StreamItem[]; + previousTail: StreamItem[]; + previousHead: StreamItem[]; + pendingClientMessageId: string | null; + policy: "preserve-local-submissions" | "replace-history"; +} + +export interface CanonicalStreamReplacementResult { + tail: StreamItem[]; + head: StreamItem[]; +} + +function removeUserMessageAt(items: UserMessageItem[], index: number): UserMessageItem[] { + return [...items.slice(0, index), ...items.slice(index + 1)]; +} + +function preserveReplacementHead( + tail: StreamItem[], + currentHead: StreamItem[], +): CanonicalStreamReplacementResult { + const liveAssistantIndex = currentHead.findLastIndex((item) => item.kind === "assistant_message"); + const submittedRows = currentHead.filter( + (item): item is UserMessageItem => + item.kind === "user_message" && item.clientMessageId !== undefined, + ); + if (liveAssistantIndex < 0) return { tail, head: submittedRows }; + + const liveAssistant = currentHead[liveAssistantIndex]; + const tailAssistant = tail.at(-1); + if ( + liveAssistant.kind !== "assistant_message" || + !tailAssistant || + tailAssistant.kind !== "assistant_message" || + !liveAssistant.text.startsWith(tailAssistant.text) + ) { + return { tail, head: submittedRows }; + } + + const head: StreamItem[] = []; + currentHead.forEach((item, index) => { + if (index === liveAssistantIndex) { + head.push({ ...liveAssistant, text: tailAssistant.text }); + return; + } + if (item.kind === "user_message" && item.clientMessageId) { + head.push(item); + } + }); + return { tail: tail.slice(0, -1), head }; +} + +export function replaceWithCanonicalStream( + input: CanonicalStreamReplacementInput, +): CanonicalStreamReplacementResult { + let unmatchedTailMessages = input.previousTail.filter( + (item): item is UserMessageItem => + item.kind === "user_message" && item.clientMessageId !== undefined, + ); + let nextHead = input.previousHead; + const matchedHeadClientMessageIds = new Set(); + const nextTail: StreamItem[] = []; + + for (const item of input.canonical) { + if (item.kind !== "user_message") { + nextTail.push(item); + continue; + } + + const tailResult = produceUserMessage(unmatchedTailMessages, item, null, "existing"); + if (tailResult.matched) { + unmatchedTailMessages = removeUserMessageAt(unmatchedTailMessages, tailResult.index); + nextTail.push(tailResult.message); + continue; + } + + const headResult = produceUserMessage(nextHead, item, null, "existing"); + if (headResult.matched) { + nextHead = headResult.items; + if (headResult.message.clientMessageId) { + matchedHeadClientMessageIds.add(headResult.message.clientMessageId); + } + continue; + } + + nextTail.push(item); + } + + for (const local of unmatchedTailMessages) { + if ( + input.policy === "replace-history" && + local.clientMessageId !== input.pendingClientMessageId + ) { + continue; + } + const insertionIndex = nextTail.findIndex( + (item) => item.timestamp.getTime() > local.timestamp.getTime(), + ); + nextTail.splice(insertionIndex < 0 ? nextTail.length : insertionIndex, 0, local); + } + + nextHead = nextHead.filter((item) => { + if (item.kind !== "user_message" || !item.clientMessageId) return true; + if (matchedHeadClientMessageIds.has(item.clientMessageId)) return true; + return ( + input.policy === "preserve-local-submissions" || + item.clientMessageId === input.pendingClientMessageId + ); + }); + + return preserveReplacementHead(nextTail, nextHead); +} export interface AssistantMessageItem { kind: "assistant_message"; @@ -237,124 +560,23 @@ function markThoughtReady(item: ThoughtItem): ThoughtItem { }; } -function buildUserMessageItem(input: { - id: string; - clientMessageId?: string; - text: string; - timestamp: Date; - optimistic?: UserMessageItem | null; -}): UserMessageItem { - if (input.optimistic) { - return { - kind: "user_message", - id: input.id, - ...(input.clientMessageId ? { clientMessageId: input.clientMessageId } : {}), - text: input.optimistic.text, - timestamp: input.optimistic.timestamp, - ...(input.optimistic.images && input.optimistic.images.length > 0 - ? { images: input.optimistic.images } - : {}), - ...(input.optimistic.attachments && input.optimistic.attachments.length > 0 - ? { attachments: input.optimistic.attachments } - : {}), - }; - } - - return { - kind: "user_message", - id: input.id, - ...(input.clientMessageId ? { clientMessageId: input.clientMessageId } : {}), - text: input.text, - timestamp: input.timestamp, - }; -} - -export function buildOptimisticUserMessage(input: OptimisticUserMessageInput): UserMessageItem { - return { - kind: "user_message", - id: input.id, - text: input.text, - timestamp: input.timestamp, - optimistic: true, - ...(input.images && input.images.length > 0 ? { images: input.images } : {}), - ...(input.attachments && input.attachments.length > 0 - ? { attachments: input.attachments } - : {}), - }; -} - -export function appendOptimisticUserMessageToStream(params: { - tail: StreamItem[]; - head: StreamItem[]; - message: UserMessageItem; - placement: OptimisticUserMessagePlacement; -}): ApplyStreamEventResult { - const { tail, head, message, placement } = params; - if (tail.some((item) => item.id === message.id) || head.some((item) => item.id === message.id)) { - return { tail, head, changedTail: false, changedHead: false }; - } - - if (placement === "active-head" && head.length > 0) { - return { - tail, - head: [...head, message], - changedTail: false, - changedHead: true, - }; - } - - return { - tail: [...tail, message], - head, - changedTail: true, - changedHead: false, - }; -} - export function handoffCreatedAgentUserMessageToStream(params: { tail: StreamItem[]; head: StreamItem[]; message: UserMessageItem; }): ApplyStreamEventResult { - const { tail, head, message } = params; - const items = [...tail, ...head]; - const userIndex = items.findIndex((item) => item.kind === "user_message"); - if (userIndex < 0) { - return appendOptimisticUserMessageToStream({ - tail, - head, - message, - placement: "tail", - }); - } - - const userMessage = items[userIndex]; - if (!userMessage || userMessage.kind !== "user_message" || userMessage.optimistic) { - return { tail, head, changedTail: false, changedHead: false }; - } - - const handedOffMessage = buildUserMessageItem({ - id: userMessage.id, - text: message.text, - timestamp: message.timestamp, - optimistic: message, + return upsertUserMessageAcrossStream({ + ...params, + insert: "tail", + presentation: "incoming", }); - if (userIndex < tail.length) { - const nextTail = [...tail]; - nextTail[userIndex] = handedOffMessage; - return { tail: nextTail, head, changedTail: true, changedHead: false }; - } - - const nextHead = [...head]; - nextHead[userIndex - tail.length] = handedOffMessage; - return { tail, head: nextHead, changedTail: false, changedHead: true }; } function appendUserMessage( state: StreamItem[], text: string, timestamp: Date, - source: StreamUpdateSource, + _source: StreamUpdateSource, messageId?: string, clientMessageId?: string, ): StreamItem[] { @@ -364,37 +586,14 @@ function appendUserMessage( } const chunkSeed = chunk.trim() || chunk; - const entryId = messageId ?? createUniqueTimelineId(state, "user", chunkSeed, timestamp); - const optimisticIndex = state.findIndex( - (entry) => - entry.kind === "user_message" && - entry.optimistic && - (clientMessageId !== undefined - ? entry.id === clientMessageId - : source === "live" || entry.id === messageId || entry.text === chunk), - ); - const optimistic = optimisticIndex >= 0 ? (state[optimisticIndex] as UserMessageItem) : null; - - const nextItem = buildUserMessageItem({ - id: entryId, + const nextItem = createUserMessage({ + id: messageId ?? createUniqueTimelineId(state, "user", chunkSeed, timestamp), clientMessageId, + messageId, text: chunk, timestamp, - optimistic, }); - - if (optimisticIndex >= 0) { - const next = [...state]; - next[optimisticIndex] = nextItem; - return next; - } - - return [...state, nextItem]; -} - -export function clearOptimisticUserMessages(state: StreamItem[]): StreamItem[] { - const next = state.filter((item) => item.kind !== "user_message" || !item.optimistic); - return next.length === state.length ? state : next; + return upsertUserMessage(state, nextItem); } function appendAssistantMessage( @@ -426,8 +625,8 @@ function appendAssistantMessage( return [...state.slice(0, -1), updated]; } - // If the last item is a user_message (optimistic append to head during - // interrupt), look one further back for the streaming assistant_message. + // A submitted user row can follow the streaming assistant during interrupt. + // In that case, look one row further back for the assistant to extend. const secondLast = state[state.length - 2]; if ( source === "live" && @@ -1149,7 +1348,6 @@ export function flushHeadToTail(tail: StreamItem[], head: StreamItem[]): StreamI if (newItems.length === 0) { return tail; } - return [...tail, ...newItems]; } @@ -1177,8 +1375,7 @@ function shouldFlushHead(input: { return true; } - // Find the last streamable item in head (skip trailing non-streamable - // items like an optimistic user_message appended during interrupt). + // Find the last streamable item in head (skip trailing non-streamable items). let lastStreamable: StreamItem | undefined; for (let i = head.length - 1; i >= 0; i--) { if (isStreamableKind(head[i].kind)) { @@ -1237,6 +1434,32 @@ export function applyStreamEvent(params: { let changedTail = false; let changedHead = false; + if (event.type === "timeline" && event.item.type === "user_message") { + const normalized = normalizeChunk(event.item.text); + if (!normalized.hasContent) { + return { tail, head, changedTail: false, changedHead: false }; + } + const canonical = createUserMessage({ + id: + event.item.messageId ?? + createUniqueTimelineId([...tail, ...head], "user", normalized.chunk.trim(), timestamp), + messageId: event.item.messageId, + clientMessageId: event.item.clientMessageId, + text: normalized.chunk, + timestamp, + }); + const reconciled = upsertUserMessageAcrossStream({ + tail, + head, + message: canonical, + insert: "none", + presentation: "existing", + }); + if (reconciled.location) { + return reconciled; + } + } + const flushHead = () => { if (nextHead.length === 0) { return; diff --git a/packages/app/src/utils/agent-directory-sync.test.ts b/packages/app/src/utils/agent-directory-sync.test.ts index 1ec16bde26..49afa03a70 100644 --- a/packages/app/src/utils/agent-directory-sync.test.ts +++ b/packages/app/src/utils/agent-directory-sync.test.ts @@ -7,6 +7,7 @@ import { useSessionStore } from "@/stores/session-store"; import { normalizeAgentSnapshot } from "@/utils/agent-snapshots"; import { isAgentArchiving, setAgentArchiving } from "@/hooks/use-archive-agent"; import { queryClient } from "@/data/query-client"; +import { createUserMessage } from "@/types/stream"; import { applyAgentDirectoryDelta, replaceFetchedAgentDirectory } from "./agent-directory-sync"; function createAgentPayload( @@ -64,6 +65,75 @@ function permission(id: string): AgentPermissionRequest { return { id, provider: "codex", name: id, kind: "tool", title: id }; } +function beginPendingSubmission(serverId: string, agentId: string): string { + const clientMessageId = `client-${agentId}`; + useSessionStore.getState().beginAgentMessageSubmission( + serverId, + agentId, + createUserMessage({ + clientMessageId, + text: "Run this", + timestamp: new Date("2026-07-27T10:00:00.000Z"), + }), + ); + return clientMessageId; +} + +function applyAgentStatus(input: { + serverId: string; + agentId: string; + status: AgentSnapshotPayload["status"]; + updatedAt: string; +}): void { + const agent = createAgentPayload({ + id: input.agentId, + status: input.status, + updatedAt: input.updatedAt, + }); + applyAgentDirectoryDelta({ + serverId: input.serverId, + delta: { kind: "upsert", agent, project: createEntry(agent).project }, + }); +} + +describe("message submission lifecycle ordering", () => { + it("keeps idle and pending when acceptance precedes authoritative running", () => { + const serverId = "server-accept-before-running"; + const agentId = "agent-1"; + const store = useSessionStore.getState(); + store.initializeSession(serverId, null as unknown as DaemonClient); + applyAgentStatus({ + serverId, + agentId, + status: "idle", + updatedAt: "2026-07-27T10:00:00.000Z", + }); + const clientMessageId = beginPendingSubmission(serverId, agentId); + + store.acceptAgentMessageSubmission(serverId, agentId, clientMessageId); + + const session = useSessionStore.getState().sessions[serverId]; + expect(session?.agents.get(agentId)?.status).toBe("idle"); + expect(session?.messageSubmissions.get(agentId)).toEqual({ + clientMessageId, + submittedAt: new Date("2026-07-27T10:00:00.000Z"), + phase: "waiting-for-running", + }); + + applyAgentStatus({ + serverId, + agentId, + status: "running", + updatedAt: "2026-07-27T10:00:01.000Z", + }); + + expect( + useSessionStore.getState().sessions[serverId]?.messageSubmissions.get(agentId), + ).toBeUndefined(); + store.clearSession(serverId); + }); +}); + describe("replaceFetchedAgentDirectory", () => { it("preserves timeline initialization while replacing directory state", () => { const serverId = "server-initializing"; diff --git a/packages/server/src/server/agent/providers/mock-load-test-agent.test.ts b/packages/server/src/server/agent/providers/mock-load-test-agent.test.ts index 98d904076b..cacdda38e1 100644 --- a/packages/server/src/server/agent/providers/mock-load-test-agent.test.ts +++ b/packages/server/src/server/agent/providers/mock-load-test-agent.test.ts @@ -49,6 +49,25 @@ describe("MockLoadTestAgentClient", () => { }); }); + test("rejects the configured number of prompts before starting a retry", async () => { + const client = new MockLoadTestAgentClient(); + const session = await client.createSession({ + provider: "mock", + cwd: process.cwd(), + model: "ten-second-stream", + featureValues: { mockPromptRejections: 1 }, + }); + + await expect(session.startTurn("Reject this prompt.")).rejects.toThrow( + "Requested mock prompt rejection", + ); + + await expect(session.startTurn("Accept this retry.")).resolves.toEqual({ + turnId: expect.any(String), + }); + await session.interrupt(); + }); + test("returns schema-shaped JSON for structured branch-name generation", async () => { vi.useFakeTimers(); const client = new MockLoadTestAgentClient(); diff --git a/packages/server/src/server/agent/providers/mock-load-test-agent.ts b/packages/server/src/server/agent/providers/mock-load-test-agent.ts index 33b9ddface..ffab16c79f 100644 --- a/packages/server/src/server/agent/providers/mock-load-test-agent.ts +++ b/packages/server/src/server/agent/providers/mock-load-test-agent.ts @@ -583,6 +583,7 @@ export class MockLoadTestAgentSession implements AgentSession { private modeId: string | null; private modelId: string | null; private readonly rewindError: string | null; + private remainingPromptRejections: number; constructor(options: { config: AgentSessionConfig; sessionId: string; logger?: Logger }) { this.id = options.sessionId; @@ -593,6 +594,13 @@ export class MockLoadTestAgentSession implements AgentSession { typeof options.config.featureValues?.mockRewindError === "string" ? options.config.featureValues.mockRewindError : null; + const requestedPromptRejections = options.config.featureValues?.mockPromptRejections; + this.remainingPromptRejections = + typeof requestedPromptRejections === "number" && + Number.isSafeInteger(requestedPromptRejections) && + requestedPromptRejections > 0 + ? requestedPromptRejections + : 0; } async run(prompt: AgentPromptInput, options?: AgentRunOptions): Promise { @@ -611,6 +619,10 @@ export class MockLoadTestAgentSession implements AgentSession { if (this.activeTurn) { throw new Error("Mock load-test provider already has an active turn"); } + if (this.remainingPromptRejections > 0) { + this.remainingPromptRejections -= 1; + throw new Error("Requested mock prompt rejection"); + } const profile = resolveModelProfile(this.modelId); const turnId = randomUUID();