diff --git a/packages/app/e2e/assistant-fork-menu.spec.ts b/packages/app/e2e/assistant-fork-menu.spec.ts index a1ab1eafb..c09a90921 100644 --- a/packages/app/e2e/assistant-fork-menu.spec.ts +++ b/packages/app/e2e/assistant-fork-menu.spec.ts @@ -1,7 +1,7 @@ import { expect, test as base, type Page } from "./fixtures"; import { scrollAgentChatToBottom } from "./helpers/agent-bottom-anchor"; import { awaitAssistantMessage } from "./helpers/agent-stream"; -import { expectComposerVisible } from "./helpers/composer"; +import { expectComposerVisible, submitMessage } from "./helpers/composer"; import { getE2EDaemonPort } from "./helpers/daemon-port"; import { openAgentRoute, @@ -11,6 +11,7 @@ import { } from "./helpers/mock-agent"; import { getServerId } from "./helpers/server-id"; import { seedSavedSettingsHosts } from "./helpers/settings"; +import { submitNewWorkspaceEmpty } from "./helpers/new-workspace"; const test = base.extend<{ seedForkWorkspace: (options: MockAgentOptions) => Promise; @@ -75,6 +76,33 @@ test.describe("Assistant fork menu", () => { await expectChatHistoryPill(page); }); + test("keeps the fork attachment after submitting an existing-workspace draft tab", async ({ + page, + seedForkWorkspace, + }) => { + const session = await seedForkWorkspace({ + repoPrefix: "assistant-fork-tab-submit-", + title: "Assistant fork tab submit", + initialPrompt: "emit 1 coalesced agent stream updates for assistant fork tab submit.", + model: "ten-second-stream", + }); + + await openAgentRoute(page, session); + await expectComposerVisible(page); + await awaitAssistantMessage(page); + await session.client.waitForFinish(session.agentId, 45_000); + + await openAssistantForkMenu(page); + await page.getByTestId("assistant-fork-menu-new-tab").click(); + await expectChatHistoryPill(page); + + await submitMessage(page, ""); + + const userMessage = page.getByTestId("user-message").filter({ hasText: "Chat history" }).last(); + await expect(userMessage).toBeVisible({ timeout: 30_000 }); + await expect(userMessage).not.toContainText("Source agent:"); + }); + test("forks an assistant turn into New Workspace and keeps the attachment across host changes", async ({ page, seedForkWorkspace, @@ -118,4 +146,31 @@ test.describe("Assistant fork menu", () => { .click(); await expectChatHistoryPill(page); }); + + test("keeps the fork attachment after the new agent receives its user message", async ({ + page, + seedForkWorkspace, + }) => { + const session = await seedForkWorkspace({ + repoPrefix: "assistant-fork-submit-", + title: "Assistant fork submit", + initialPrompt: "emit 1 coalesced agent stream updates for assistant fork submit.", + model: "ten-second-stream", + }); + + await openAgentRoute(page, session); + await expectComposerVisible(page); + await awaitAssistantMessage(page); + await session.client.waitForFinish(session.agentId, 45_000); + + await openAssistantForkMenu(page); + await page.getByTestId("assistant-fork-menu-new-workspace").click(); + await expectChatHistoryPill(page); + + await submitNewWorkspaceEmpty(page); + + const userMessage = page.getByTestId("user-message").filter({ hasText: "Chat history" }).last(); + await expect(userMessage).toBeVisible({ timeout: 30_000 }); + await expect(userMessage).not.toContainText("Source agent:"); + }); }); diff --git a/packages/app/src/composer/draft/create-flow.ts b/packages/app/src/composer/draft/create-flow.ts index 82d063aac..f44b8ef0a 100644 --- a/packages/app/src/composer/draft/create-flow.ts +++ b/packages/app/src/composer/draft/create-flow.ts @@ -123,8 +123,8 @@ export function useDraftAgentCreateFlow({ const updatePendingAgentId = useCreateFlowStore((state) => state.updateAgentId); const markPendingCreateLifecycle = useCreateFlowStore((state) => state.markLifecycle); const clearPendingCreateAttempt = useCreateFlowStore((state) => state.clear); - const appendOptimisticUserMessageToAgentStream = useSessionStore( - (state) => state.appendOptimisticUserMessageToAgentStream, + const handoffCreatedAgentUserMessage = useSessionStore( + (state) => state.handoffCreatedAgentUserMessage, ); const formErrorMessage = machine.tag === "draft" ? machine.errorMessage : ""; @@ -189,7 +189,7 @@ export function useDraftAgentCreateFlow({ if (createResult.agentId) { updatePendingAgentId({ draftId, agentId: createResult.agentId }); - appendOptimisticUserMessageToAgentStream( + handoffCreatedAgentUserMessage( pendingServerId, createResult.agentId, buildOptimisticUserMessage({ @@ -199,7 +199,6 @@ export function useDraftAgentCreateFlow({ images: attempt.images, attachments: attempt.attachments, }), - { placement: "tail", skipIfUserMessageExists: true }, ); markPendingCreateLifecycle({ draftId, lifecycle: "sent" }); } @@ -216,7 +215,7 @@ export function useDraftAgentCreateFlow({ } }, [ - appendOptimisticUserMessageToAgentStream, + handoffCreatedAgentUserMessage, clearPendingCreateAttempt, createRequest, draftId, diff --git a/packages/app/src/stores/session-store.ts b/packages/app/src/stores/session-store.ts index 902b7ebee..bea9f1c75 100644 --- a/packages/app/src/stores/session-store.ts +++ b/packages/app/src/stores/session-store.ts @@ -4,8 +4,7 @@ import { subscribeWithSelector } from "zustand/middleware"; import type { DaemonClient } from "@getpaseo/client/internal/daemon-client"; import type { AgentDirectoryEntry } from "@/types/agent-directory"; import { - appendOptimisticUserMessageToStream, - type OptimisticUserMessagePlacement, + handoffCreatedAgentUserMessageToStream, type StreamItem, type UserMessageItem, } from "@/types/stream"; @@ -434,14 +433,10 @@ interface SessionStoreActions { agentId: string, state: { tail?: StreamItem[]; head?: StreamItem[] }, ) => void; - appendOptimisticUserMessageToAgentStream: ( + handoffCreatedAgentUserMessage: ( serverId: string, agentId: string, message: UserMessageItem, - options: { - placement: OptimisticUserMessagePlacement; - skipIfUserMessageExists?: boolean; - }, ) => boolean; clearAgentStreamHead: (serverId: string, agentId: string) => void; setAgentTimelineCursor: ( @@ -960,8 +955,8 @@ export const useSessionStore = create()( }); }, - appendOptimisticUserMessageToAgentStream: (serverId, agentId, message, options) => { - let didAppend = false; + handoffCreatedAgentUserMessage: (serverId, agentId, message) => { + let didHandoff = false; set((prev) => { const session = prev.sessions[serverId]; if (!session) { @@ -970,12 +965,10 @@ export const useSessionStore = create()( const currentTail = session.agentStreamTail.get(agentId) ?? []; const currentHead = session.agentStreamHead.get(agentId) ?? []; - const result = appendOptimisticUserMessageToStream({ + const result = handoffCreatedAgentUserMessageToStream({ tail: currentTail, head: currentHead, message, - placement: options.placement, - skipIfUserMessageExists: options.skipIfUserMessageExists, }); if (!result.changedTail && !result.changedHead) { return prev; @@ -987,7 +980,7 @@ export const useSessionStore = create()( const nextHead = result.changedHead ? new Map(session.agentStreamHead).set(agentId, result.head) : session.agentStreamHead; - didAppend = true; + didHandoff = true; return { ...prev, @@ -1001,7 +994,7 @@ export const useSessionStore = create()( }, }; }); - return didAppend; + return didHandoff; }, clearAgentStreamHead: (serverId, agentId) => { diff --git a/packages/app/src/timeline/session-stream-reducers.test.ts b/packages/app/src/timeline/session-stream-reducers.test.ts index 848c34982..a3e6acf5e 100644 --- a/packages/app/src/timeline/session-stream-reducers.test.ts +++ b/packages/app/src/timeline/session-stream-reducers.test.ts @@ -352,6 +352,29 @@ describe("processTimelineResponse", () => { attachments: [attachment], }); expect(userMessages[0]?.optimistic).toBeUndefined(); + + const repeated = processTimelineResponse({ + ...baseTimelineInput, + currentTail: result.tail, + payload: { + ...baseTimelineInput.payload, + reset: true, + startCursor: { seq: 1 }, + endCursor: { seq: 1 }, + entries: [ + { + ...makeTimelineEntry(1, "Analyze this", "user_message"), + item: { + type: "user_message", + text: "server-rendered attachment text", + messageId: "canonical-create-user", + }, + }, + ], + }, + }); + + expect(repeated.tail.filter((item) => item.kind === "user_message")).toEqual(userMessages); }); it("keeps an unmatched optimistic user message during tail replacement", () => { diff --git a/packages/app/src/timeline/session-stream-reducers.ts b/packages/app/src/timeline/session-stream-reducers.ts index d4b1a5c39..22b0062f2 100644 --- a/packages/app/src/timeline/session-stream-reducers.ts +++ b/packages/app/src/timeline/session-stream-reducers.ts @@ -262,7 +262,7 @@ function applyTimelineReplacePath(args: { const { timelineUnits, payload, bootstrapPolicy, currentTail, currentHead, toHydratedEvents } = args; const hydratedTail = hydrateStreamState(toHydratedEvents(timelineUnits), { source: "canonical" }); - const reconciledTail = reconcileOptimisticUsersAfterReplace({ + const reconciledTail = reconcileLocalUserPresentationAfterReplace({ canonicalTail: hydratedTail, previousTail: currentTail, previousHead: currentHead, @@ -286,50 +286,50 @@ function applyTimelineReplacePath(args: { return { tail, head, cursor, cursorChanged: true, sideEffects }; } -function collectOptimisticUserMessages(items: StreamItem[]): Array<{ +function collectLocallyPresentedUserMessages(items: StreamItem[]): Array<{ ordinal: number; item: UserMessageItem; }> { - const optimistic: 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) { - optimistic.push({ ordinal, item }); + if (item.optimistic || item.images?.length || item.attachments?.length) { + localUsers.push({ ordinal, item }); } ordinal += 1; } - return optimistic; + return localUsers; } -function mergeCanonicalUserWithOptimistic( +function mergeCanonicalUserWithLocalPresentation( canonical: UserMessageItem, - optimistic: UserMessageItem, + local: UserMessageItem, ): UserMessageItem { return { kind: "user_message", id: canonical.id, - text: optimistic.text, - timestamp: optimistic.timestamp, - ...(optimistic.images && optimistic.images.length > 0 ? { images: optimistic.images } : {}), - ...(optimistic.attachments && optimistic.attachments.length > 0 - ? { attachments: optimistic.attachments } + text: local.text, + timestamp: local.timestamp, + ...(local.images && local.images.length > 0 ? { images: local.images } : {}), + ...(local.attachments && local.attachments.length > 0 + ? { attachments: local.attachments } : {}), }; } -function reconcileOptimisticUsersAfterReplace(params: { +function reconcileLocalUserPresentationAfterReplace(params: { canonicalTail: StreamItem[]; previousTail: StreamItem[]; previousHead: StreamItem[]; }): StreamItem[] { - const optimisticUsers = collectOptimisticUserMessages([ + const localUsers = collectLocallyPresentedUserMessages([ ...params.previousTail, ...params.previousHead, ]); - if (optimisticUsers.length === 0) { + if (localUsers.length === 0) { return params.canonicalTail; } @@ -345,22 +345,32 @@ function reconcileOptimisticUsersAfterReplace(params: { let searchFromOrdinal = 0; const unmatched: UserMessageItem[] = []; - for (const optimistic of optimisticUsers) { - const canonicalOrdinal = canonicalUserIndexes.findIndex( - (_index, ordinal) => ordinal >= Math.max(optimistic.ordinal, searchFromOrdinal), - ); + for (const local of localUsers) { + const canonicalOrdinal = canonicalUserIndexes.findIndex((index, ordinal) => { + if (ordinal < searchFromOrdinal) { + return false; + } + if (local.item.optimistic) { + return ordinal >= local.ordinal; + } + return params.canonicalTail[index]?.id === local.item.id; + }); if (canonicalOrdinal < 0) { - unmatched.push(optimistic.item); + if (local.item.optimistic) { + unmatched.push(local.item); + } continue; } const canonicalIndex = canonicalUserIndexes[canonicalOrdinal]; const canonicalItem = canonicalIndex !== undefined ? nextTail[canonicalIndex] : undefined; if (!canonicalItem || canonicalItem.kind !== "user_message") { - unmatched.push(optimistic.item); + if (local.item.optimistic) { + unmatched.push(local.item); + } continue; } - nextTail[canonicalIndex] = mergeCanonicalUserWithOptimistic(canonicalItem, optimistic.item); + nextTail[canonicalIndex] = mergeCanonicalUserWithLocalPresentation(canonicalItem, local.item); searchFromOrdinal = canonicalOrdinal + 1; changed = true; } diff --git a/packages/app/src/types/stream.test.ts b/packages/app/src/types/stream.test.ts index 0c8979971..20d953a5d 100644 --- a/packages/app/src/types/stream.test.ts +++ b/packages/app/src/types/stream.test.ts @@ -7,6 +7,7 @@ import { appendOptimisticUserMessageToStream, buildOptimisticUserMessage, clearOptimisticUserMessages, + handoffCreatedAgentUserMessageToStream, hydrateStreamState, mergeToolCallDetail, reduceStreamUpdate, @@ -1284,27 +1285,86 @@ describe("turn lifecycle events", () => { message: optimistic, placement: "active-head", }); - const skipped = appendOptimisticUserMessageToStream({ - tail: [ - { - kind: "user_message", - id: "canonical-user", - text: "already canonical", - timestamp: new Date("2025-01-01T15:03:21Z"), - }, - ], - head: [], - message: optimistic, - placement: "tail", - skipIfUserMessageExists: true, - }); - assert.deepStrictEqual(first.tail, []); assert.deepStrictEqual(first.head, [headItem, optimistic]); assert.strictEqual(second.changedHead, false); assert.strictEqual(second.head, first.head); - assert.strictEqual(skipped.changedTail, false); - assert.strictEqual(skipped.tail.length, 1); + }); + + it("hands rich optimistic content to an authoritative create message without duplicating it", () => { + const timestamp = new Date("2025-01-01T15:03:20Z"); + const optimistic = buildOptimisticUserMessage({ + id: "client-user", + text: "", + timestamp, + images: [ + { + id: "image-1", + mimeType: "image/png", + storageType: "web-indexeddb", + storageKey: "image-1", + createdAt: timestamp.getTime(), + }, + ], + attachments: [ + { + type: "text", + mimeType: "text/plain", + text: "Previous conversation", + title: "Chat history", + contextKind: "chat_history", + }, + ], + }); + const canonical: StreamItem = { + kind: "user_message", + id: "provider-user", + text: "server-rendered attachment text", + timestamp: new Date("2025-01-01T15:03:21Z"), + }; + + const handedOff = handoffCreatedAgentUserMessageToStream({ + tail: [canonical], + head: [], + message: optimistic, + }); + const repeated = handoffCreatedAgentUserMessageToStream({ + tail: handedOff.tail, + head: handedOff.head, + message: optimistic, + }); + + assert.deepStrictEqual(handedOff.tail, [ + { + kind: "user_message", + id: "provider-user", + text: optimistic.text, + timestamp: optimistic.timestamp, + images: optimistic.images, + attachments: optimistic.attachments, + }, + ]); + assert.deepStrictEqual(handedOff.head, []); + assert.deepStrictEqual(repeated.tail, handedOff.tail); + assert.deepStrictEqual(repeated.head, handedOff.head); + + const afterNextUser = reduceStreamUpdate( + handedOff.tail, + { + type: "timeline", + provider: "claude", + item: { + type: "user_message", + text: "Next prompt", + messageId: "provider-next-user", + }, + }, + new Date("2025-01-01T15:04:00Z"), + ); + assert.deepStrictEqual( + afterNextUser.filter((item) => item.kind === "user_message").map((item) => item.id), + ["provider-user", "provider-next-user"], + ); }); it("reconciles an optimistic user message that was pending in the streaming head", () => { diff --git a/packages/app/src/types/stream.ts b/packages/app/src/types/stream.ts index 1f6f6665b..e059c0e98 100644 --- a/packages/app/src/types/stream.ts +++ b/packages/app/src/types/stream.ts @@ -272,23 +272,14 @@ export function buildOptimisticUserMessage(input: OptimisticUserMessageInput): U }; } -function hasUserMessage(state: StreamItem[]): boolean { - return state.some((item) => item.kind === "user_message"); -} - export function appendOptimisticUserMessageToStream(params: { tail: StreamItem[]; head: StreamItem[]; message: UserMessageItem; placement: OptimisticUserMessagePlacement; - skipIfUserMessageExists?: boolean; }): ApplyStreamEventResult { const { tail, head, message, placement } = params; - if ( - tail.some((item) => item.id === message.id) || - head.some((item) => item.id === message.id) || - (params.skipIfUserMessageExists && (hasUserMessage(tail) || hasUserMessage(head))) - ) { + if (tail.some((item) => item.id === message.id) || head.some((item) => item.id === message.id)) { return { tail, head, changedTail: false, changedHead: false }; } @@ -309,6 +300,45 @@ export function appendOptimisticUserMessageToStream(params: { }; } +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, + }); + 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,