diff --git a/packages/app/src/components/composer-actions.test.ts b/packages/app/src/components/composer-actions.test.ts index 3067c52cf..9c385cee2 100644 --- a/packages/app/src/components/composer-actions.test.ts +++ b/packages/app/src/components/composer-actions.test.ts @@ -189,6 +189,7 @@ function createFakeStream(initialHead: Map = new Map()): F 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); diff --git a/packages/app/src/components/composer-actions.ts b/packages/app/src/components/composer-actions.ts index 82e27fc9b..539717419 100644 --- a/packages/app/src/components/composer-actions.ts +++ b/packages/app/src/components/composer-actions.ts @@ -9,7 +9,13 @@ import { userAttachmentsOnly, } from "@/attachments/workspace-attachment-utils"; import { splitComposerAttachmentsForSubmit } from "@/components/composer-attachments"; -import { generateMessageId, type StreamItem } from "@/types/stream"; +import { + appendOptimisticUserMessageToStream, + buildOptimisticUserMessage, + generateMessageId, + type StreamItem, + type UserMessageItem, +} from "@/types/stream"; import type { PickedImageAttachmentInput } from "@/hooks/image-attachment-picker"; export interface QueuedComposerMessage { @@ -49,6 +55,7 @@ export interface ComposerCancelClient { } 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; @@ -130,15 +137,13 @@ export async function dispatchComposerAgentMessage( ): Promise { const wirePayload = splitComposerAttachmentsForSubmit(input.attachments); const messageId = generateMessageId(); - const userMessage: StreamItem = { - kind: "user_message", + const userMessage = buildOptimisticUserMessage({ id: messageId, text: input.text, timestamp: new Date(), - optimistic: true, - ...(wirePayload.images.length > 0 ? { images: wirePayload.images } : {}), - ...(wirePayload.attachments.length > 0 ? { attachments: wirePayload.attachments } : {}), - }; + images: wirePayload.images, + attachments: wirePayload.attachments, + }); appendUserMessageToStream(input.agentId, userMessage, input.stream); const imagesData = await input.encodeImages(wirePayload.images); await input.client.sendAgentMessage(input.agentId, input.text, { @@ -150,23 +155,29 @@ export async function dispatchComposerAgentMessage( function appendUserMessageToStream( agentId: string, - userMessage: StreamItem, + userMessage: UserMessageItem, stream: AgentStreamWriter, ): void { - const head = stream.getHead(agentId); - if (head && head.length > 0) { + const result = appendOptimisticUserMessageToStream({ + tail: stream.getTail(agentId) ?? [], + head: stream.getHead(agentId) ?? [], + message: userMessage, + placement: "active-head", + }); + if (result.changedHead) { stream.setHead((prev) => { const next = new Map(prev); - next.set(agentId, [...(prev.get(agentId) ?? []), userMessage]); + next.set(agentId, result.head); + return next; + }); + } + if (result.changedTail) { + stream.setTail((prev) => { + const next = new Map(prev); + next.set(agentId, result.tail); return next; }); - return; } - stream.setTail((prev) => { - const next = new Map(prev); - next.set(agentId, [...(prev.get(agentId) ?? []), userMessage]); - return next; - }); } export interface QueueComposerMessageInput { diff --git a/packages/app/src/components/composer.tsx b/packages/app/src/components/composer.tsx index c52330b36..8ec2e2af5 100644 --- a/packages/app/src/components/composer.tsx +++ b/packages/app/src/components/composer.tsx @@ -1029,6 +1029,7 @@ export function Composer({ throw new Error("Host is not connected"); } 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), diff --git a/packages/app/src/contexts/session-context.tsx b/packages/app/src/contexts/session-context.tsx index 3544a54a3..6d54f6fb0 100644 --- a/packages/app/src/contexts/session-context.tsx +++ b/packages/app/src/contexts/session-context.tsx @@ -13,6 +13,7 @@ import { type ProcessTimelineResponseOutput, type TimelineReducerSideEffect, } from "@/timeline/session-stream-reducers"; +import { useCreateFlowStore } from "@/stores/create-flow-store"; import { TIMELINE_FETCH_PAGE_SIZE } from "@/timeline/timeline-fetch-policy"; import type { AgentAttachment, SessionOutboundMessage } from "@server/shared/messages"; import { parseServerInfoStatusPayload } from "@server/shared/messages"; @@ -375,6 +376,7 @@ function finalizeTimelineApplication(input: { } if (shouldMarkAuthoritativeHistoryApplied) { setAgentAuthoritativeHistoryApplied(serverId, agentId, true); + useCreateFlowStore.getState().clearByAgent({ serverId, agentId }); } if (result.initResolution === "resolve") { resolveInitDeferred(initKey); diff --git a/packages/app/src/hooks/use-agent-screen-state-machine.test.ts b/packages/app/src/hooks/use-agent-screen-state-machine.test.ts index f4148f49c..b4e304e9b 100644 --- a/packages/app/src/hooks/use-agent-screen-state-machine.test.ts +++ b/packages/app/src/hooks/use-agent-screen-state-machine.test.ts @@ -57,13 +57,12 @@ function createAgentWithStatus({ id, status }: { id: string; status: Agent["stat function createBaseInput(): AgentScreenMachineInput { return { agent: null, - placeholderAgent: null, + continuity: { kind: "none" }, missingAgentState: { kind: "idle" }, isConnected: true, isArchivingCurrentAgent: false, isHistorySyncing: false, needsAuthoritativeSync: false, - shouldUseOptimisticStream: false, hasHydratedHistoryBefore: false, }; } @@ -298,8 +297,7 @@ describe("deriveAgentScreenViewState", () => { const memory = createBaseMemory(); const input: AgentScreenMachineInput = { ...createBaseInput(), - placeholderAgent: createAgent("draft-agent"), - shouldUseOptimisticStream: true, + continuity: { kind: "optimistic-create", agent: createAgent("draft-agent") }, }; const result = deriveAgentScreenViewState({ input, memory }); @@ -335,8 +333,7 @@ describe("deriveAgentScreenViewState", () => { const input: AgentScreenMachineInput = { ...createBaseInput(), agent: createAgentWithStatus({ id: "agent-1", status: "idle" }), - placeholderAgent: createAgent("agent-1"), - shouldUseOptimisticStream: true, + continuity: { kind: "optimistic-create", agent: createAgent("agent-1") }, needsAuthoritativeSync: true, isHistorySyncing: true, hasHydratedHistoryBefore: false, @@ -347,14 +344,14 @@ describe("deriveAgentScreenViewState", () => { expect(ready.source).toBe("optimistic"); expect(ready.agent.status).toBe("running"); + expect(ready.sync).toEqual({ status: "catching_up", ui: "silent" }); }); it("keeps optimistic flow non-blocking while transitioning to authoritative stream", () => { const initialMemory = createBaseMemory(); const optimisticInput: AgentScreenMachineInput = { ...createBaseInput(), - placeholderAgent: createAgent("draft-agent"), - shouldUseOptimisticStream: true, + continuity: { kind: "optimistic-create", agent: createAgent("draft-agent") }, }; const optimistic = deriveAgentScreenViewState({ @@ -382,8 +379,7 @@ describe("deriveAgentScreenViewState", () => { const input: AgentScreenMachineInput = { ...createBaseInput(), agent: createAgentWithStatus({ id: "agent-1", status: "idle" }), - placeholderAgent: createAgent("agent-1"), - shouldUseOptimisticStream: true, + continuity: { kind: "optimistic-create", agent: createAgent("agent-1") }, }; const result = deriveAgentScreenViewState({ input, memory }); @@ -402,14 +398,12 @@ describe("deriveAgentScreenViewState", () => { { ...createBaseInput(), agent: createAgentWithStatus({ id: "agent-1", status: "idle" }), - placeholderAgent: createAgent("agent-1"), - shouldUseOptimisticStream: true, + continuity: { kind: "optimistic-create", agent: createAgent("agent-1") }, }, { ...createBaseInput(), agent: createAgentWithStatus({ id: "agent-1", status: "running" }), - placeholderAgent: createAgent("agent-1"), - shouldUseOptimisticStream: true, + continuity: { kind: "optimistic-create", agent: createAgent("agent-1") }, }, { ...createBaseInput(), @@ -431,8 +425,7 @@ describe("deriveAgentScreenViewState", () => { const input: AgentScreenMachineInput = { ...createBaseInput(), agent: createAgentWithStatus({ id: "agent-1", status: "initializing" }), - placeholderAgent: createAgent("agent-1"), - shouldUseOptimisticStream: true, + continuity: { kind: "optimistic-create", agent: createAgent("agent-1") }, }; const result = deriveAgentScreenViewState({ input, memory }); @@ -447,8 +440,7 @@ describe("deriveAgentScreenViewState", () => { const input: AgentScreenMachineInput = { ...createBaseInput(), agent: createAgentWithStatus({ id: "agent-1", status: "running" }), - placeholderAgent: createAgent("agent-1"), - shouldUseOptimisticStream: true, + continuity: { kind: "optimistic-create", agent: createAgent("agent-1") }, }; const result = deriveAgentScreenViewState({ input, memory }); @@ -463,8 +455,7 @@ describe("deriveAgentScreenViewState", () => { const input: AgentScreenMachineInput = { ...createBaseInput(), agent: createAgentWithStatus({ id: "agent-1", status: "error" }), - placeholderAgent: createAgent("agent-1"), - shouldUseOptimisticStream: true, + continuity: { kind: "optimistic-create", agent: createAgent("agent-1") }, }; const result = deriveAgentScreenViewState({ input, memory }); @@ -479,8 +470,7 @@ describe("deriveAgentScreenViewState", () => { const input: AgentScreenMachineInput = { ...createBaseInput(), agent: createAgentWithStatus({ id: "agent-1", status: "closed" }), - placeholderAgent: createAgent("agent-1"), - shouldUseOptimisticStream: true, + continuity: { kind: "optimistic-create", agent: createAgent("agent-1") }, }; const result = deriveAgentScreenViewState({ input, memory }); diff --git a/packages/app/src/hooks/use-agent-screen-state-machine.ts b/packages/app/src/hooks/use-agent-screen-state-machine.ts index 5a9875e2f..be254a6a3 100644 --- a/packages/app/src/hooks/use-agent-screen-state-machine.ts +++ b/packages/app/src/hooks/use-agent-screen-state-machine.ts @@ -24,19 +24,26 @@ export type AgentScreenMissingState = export interface AgentScreenMachineInput { agent: AgentScreenAgent | null; - placeholderAgent: AgentScreenAgent | null; missingAgentState: AgentScreenMissingState; isConnected: boolean; isArchivingCurrentAgent: boolean; isHistorySyncing: boolean; needsAuthoritativeSync: boolean; - shouldUseOptimisticStream: boolean; + continuity: AgentScreenContinuity; hasHydratedHistoryBefore: boolean; } +export type AgentScreenContinuity = + | { kind: "none" } + | { kind: "optimistic-create"; agent: AgentScreenAgent }; + +function hasOptimisticCreateContinuity(input: AgentScreenMachineInput): boolean { + return input.continuity.kind === "optimistic-create"; +} + function shouldBlockInitialAuthoritativeReadyState(input: AgentScreenMachineInput): boolean { return ( - !input.shouldUseOptimisticStream && + !hasOptimisticCreateContinuity(input) && !input.hasHydratedHistoryBefore && (input.needsAuthoritativeSync || input.isHistorySyncing) ); @@ -93,8 +100,7 @@ function updateInitialSyncFailureMemory(args: { function shouldUseOptimisticCreateFlowAgent(input: AgentScreenMachineInput): boolean { return ( - input.shouldUseOptimisticStream && - Boolean(input.placeholderAgent) && + input.continuity.kind === "optimistic-create" && (!input.agent || input.agent.status === "initializing" || input.agent.status === "idle") ); } @@ -104,29 +110,31 @@ function resolveCandidateAgent(args: { useOptimisticCreateFlowAgent: boolean; }): AgentScreenAgent | null { const { input, useOptimisticCreateFlowAgent } = args; - if (input.agent && useOptimisticCreateFlowAgent && input.placeholderAgent) { - return { ...input.agent, status: input.placeholderAgent.status }; + const continuityAgent = + input.continuity.kind === "optimistic-create" ? input.continuity.agent : null; + if (input.agent && useOptimisticCreateFlowAgent && continuityAgent) { + return { ...input.agent, status: continuityAgent.status }; } - return input.agent ?? input.placeholderAgent; + return input.agent ?? continuityAgent; } function resolveAgentScreenSource(args: { useOptimisticCreateFlowAgent: boolean; hasAgent: boolean; - shouldUseOptimisticStream: boolean; + hasOptimisticCreateContinuity: boolean; }): "authoritative" | "optimistic" | "stale" { if (args.useOptimisticCreateFlowAgent) return "optimistic"; if (args.hasAgent) return "authoritative"; - if (args.shouldUseOptimisticStream) return "optimistic"; + if (args.hasOptimisticCreateContinuity) return "optimistic"; return "stale"; } function resolveCatchingUpUi(args: { - shouldUseOptimisticStream: boolean; + hasOptimisticCreateContinuity: boolean; hasHydratedHistoryBefore: boolean; hadInitialSyncFailure: boolean; }): "overlay" | "silent" { - if (args.shouldUseOptimisticStream) return "silent"; + if (args.hasOptimisticCreateContinuity) return "silent"; if (args.hasHydratedHistoryBefore) return "silent"; if (args.hadInitialSyncFailure) return "silent"; return "overlay"; @@ -147,7 +155,7 @@ function resolveAgentScreenSync(args: { return { status: "catching_up", ui: resolveCatchingUpUi({ - shouldUseOptimisticStream: input.shouldUseOptimisticStream, + hasOptimisticCreateContinuity: hasOptimisticCreateContinuity(input), hasHydratedHistoryBefore: input.hasHydratedHistoryBefore, hadInitialSyncFailure, }), @@ -227,7 +235,7 @@ export function deriveAgentScreenViewState({ const source = resolveAgentScreenSource({ useOptimisticCreateFlowAgent, hasAgent: Boolean(input.agent), - shouldUseOptimisticStream: input.shouldUseOptimisticStream, + hasOptimisticCreateContinuity: hasOptimisticCreateContinuity(input), }); const sync = resolveAgentScreenSync({ diff --git a/packages/app/src/hooks/use-draft-agent-create-flow.ts b/packages/app/src/hooks/use-draft-agent-create-flow.ts index 3286b1240..1c92c9454 100644 --- a/packages/app/src/hooks/use-draft-agent-create-flow.ts +++ b/packages/app/src/hooks/use-draft-agent-create-flow.ts @@ -2,7 +2,9 @@ import { useCallback, useMemo, useReducer } from "react"; import type { ComposerAttachment } from "@/attachments/types"; import { splitComposerAttachmentsForSubmit } from "@/components/composer-attachments"; import { useCreateFlowStore } from "@/stores/create-flow-store"; +import { useSessionStore } from "@/stores/session-store"; import { + buildOptimisticUserMessage, generateMessageId, type StreamItem, type UserMessageImageAttachment, @@ -110,6 +112,9 @@ 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 formErrorMessage = machine.tag === "draft" ? machine.errorMessage : ""; const isSubmitting = machine.tag === "creating"; @@ -128,19 +133,13 @@ export function useDraftAgentCreateFlow({ } return [ - { - kind: "user_message", + buildOptimisticUserMessage({ id: machine.attempt.clientMessageId, text: machine.attempt.text, timestamp: machine.attempt.timestamp, - optimistic: true, - ...(machine.attempt.images && machine.attempt.images.length > 0 - ? { images: machine.attempt.images } - : {}), - ...(machine.attempt.attachments && machine.attempt.attachments.length > 0 - ? { attachments: machine.attempt.attachments } - : {}), - }, + images: machine.attempt.images, + attachments: machine.attempt.attachments, + }), ]; }, [machine]); @@ -228,6 +227,19 @@ export function useDraftAgentCreateFlow({ if (createResult.agentId) { updatePendingAgentId({ draftId, agentId: createResult.agentId }); + appendOptimisticUserMessageToAgentStream( + pendingServerId, + createResult.agentId, + buildOptimisticUserMessage({ + id: attempt.clientMessageId, + text: attempt.text, + timestamp: attempt.timestamp, + images: attempt.images, + attachments: attempt.attachments, + }), + { placement: "tail", skipIfUserMessageExists: true }, + ); + markPendingCreateLifecycle({ draftId, lifecycle: "sent" }); } await onCreateSuccess({ result: createResult.result, attempt }); @@ -241,6 +253,7 @@ export function useDraftAgentCreateFlow({ } }, [ + appendOptimisticUserMessageToAgentStream, clearPendingCreateAttempt, createRequest, draftId, diff --git a/packages/app/src/panels/agent-panel.tsx b/packages/app/src/panels/agent-panel.tsx index 57f0409ef..a8604eae0 100644 --- a/packages/app/src/panels/agent-panel.tsx +++ b/packages/app/src/panels/agent-panel.tsx @@ -27,6 +27,7 @@ import { useAgentInitialization } from "@/hooks/use-agent-initialization"; import { useAgentInputDraft, type AgentInputDraft } from "@/hooks/use-agent-input-draft"; import { type AgentScreenAgent, + type AgentScreenContinuity, type AgentScreenMissingState, type AgentScreenViewState, useAgentScreenStateMachine, @@ -60,7 +61,6 @@ import type { StreamItem } from "@/types/stream"; import { getInitDeferred, getInitKey } from "@/utils/agent-initialization"; import { derivePendingPermissionKey, normalizeAgentSnapshot } from "@/utils/agent-snapshots"; import type { WorkspaceFileOpenRequest } from "@/workspace/file-open"; -import { mergePendingCreateImages } from "@/utils/pending-create-images"; import { navigateToAgent } from "@/utils/navigate-to-agent"; import { deriveSidebarStateBucket } from "@/utils/sidebar-agent-state"; import { buildDraftAgentSetup, type ClientSlashCommand } from "@/client-slash-commands"; @@ -313,8 +313,22 @@ const EMPTY_PENDING_PERMISSIONS = new Map(); const EMPTY_PENDING_PERMISSION_LIST: PendingPermission[] = []; type RouteBottomAnchorRequest = ReturnType; -type PendingCreateByDraftId = ReturnType["pendingByDraftId"]; -type PendingCreateAttempt = PendingCreateByDraftId[string]; + +function findActiveCreateHandoff(input: { + pendingByDraftId: ReturnType["pendingByDraftId"]; + serverId: string; + agentId?: string; +}): boolean { + if (!input.agentId) { + return false; + } + return Object.values(input.pendingByDraftId).some( + (pending) => + pending.lifecycle === "sent" && + pending.serverId === input.serverId && + pending.agentId === input.agentId, + ); +} function toErrorMessage(error: unknown): string { if (error instanceof Error) { @@ -327,27 +341,6 @@ function isNotFoundErrorMessage(message: string): boolean { return /agent not found|not found/i.test(message); } -function findPendingCreateForPanel(input: { - pendingByDraftId: PendingCreateByDraftId; - serverId: string; - agentId?: string; -}): PendingCreateAttempt | null { - if (!input.agentId) { - return null; - } - const values = Object.values(input.pendingByDraftId); - for (const entry of values) { - if ( - entry.lifecycle === "active" && - entry.serverId === input.serverId && - entry.agentId === input.agentId - ) { - return entry; - } - } - return null; -} - type AgentLookupState = | { tag: "idle" } | { tag: "loading" } @@ -630,7 +623,6 @@ function ChatAgentContent({ }, (a, b) => a === b || JSON.stringify(a) === JSON.stringify(b), ); - const pendingByDraftId = useCreateFlowStore((state) => state.pendingByDraftId); const isInitializingFromMap = useSessionStore((state) => agentId ? (state.sessions[serverId]?.initializingAgents?.get(agentId) ?? false) : false, ); @@ -645,6 +637,9 @@ function ChatAgentContent({ const agentHistorySyncGeneration = useSessionStore((state) => agentId ? (state.sessions[serverId]?.agentHistorySyncGeneration?.get(agentId) ?? -1) : -1, ); + const hasActiveCreateHandoff = useCreateFlowStore((state) => + findActiveCreateHandoff({ pendingByDraftId: state.pendingByDraftId, serverId, agentId }), + ); const hasSession = useSessionStore((state) => Boolean(state.sessions[serverId])); const { ensureAgentIsInitialized } = useAgentInitialization({ serverId, @@ -654,11 +649,6 @@ function ChatAgentContent({ kind: "idle", }); - const pendingCreate = useMemo( - () => findPendingCreateForPanel({ pendingByDraftId, serverId, agentId }), - [agentId, pendingByDraftId, serverId], - ); - const isPendingCreateForPanel = Boolean(pendingCreate); const hasHydratedHistoryBefore = hasAppliedAuthoritativeHistory; const attentionController = useAgentAttentionClear({ @@ -768,43 +758,36 @@ function ChatAgentContent({ return agentHistorySyncGeneration < historySyncGeneration; }, [agentHistorySyncGeneration, agentId, historySyncGeneration]); - const shouldUseOptimisticStream = isPendingCreateForPanel; - const authoritativeStatus = agentState.status; - const isAuthoritativeBootstrapping = - authoritativeStatus === "initializing" || authoritativeStatus === "idle"; - const showPendingCreateSubmitLoading = - isPendingCreateForPanel && (!authoritativeStatus || isAuthoritativeBootstrapping); - const canFinalizePendingCreate = Boolean(authoritativeStatus) && !isAuthoritativeBootstrapping; - const agent = useMemo( () => buildChatAgentFromState(agentState, projectPlacement), [agentState, projectPlacement], ); - - const placeholderAgent: AgentScreenAgent | null = useMemo(() => { - if (!shouldUseOptimisticStream || !agentId) { - return null; + const continuity = useMemo(() => { + if (!hasActiveCreateHandoff || !agentId) { + return { kind: "none" }; } return { - serverId, - id: agentId, - status: "running", - cwd: ".", - projectPlacement: null, + kind: "optimistic-create", + agent: { + serverId, + id: agentId, + status: "running", + cwd: agent?.cwd ?? ".", + projectPlacement: agent?.projectPlacement ?? null, + }, }; - }, [agentId, serverId, shouldUseOptimisticStream]); + }, [agent, agentId, hasActiveCreateHandoff, serverId]); const viewState = useAgentScreenStateMachine({ routeKey: `${serverId}:${agentId ?? ""}`, input: { agent: agent ?? null, - placeholderAgent, missingAgentState, isConnected, isArchivingCurrentAgent, isHistorySyncing, needsAuthoritativeSync, - shouldUseOptimisticStream, + continuity, hasHydratedHistoryBefore, }, }); @@ -872,7 +855,7 @@ function ChatAgentContent({ if (!agentId) { return; } - if (agentState.id || shouldUseOptimisticStream) { + if (agentState.id) { if (missingAgentState.kind !== "idle") { setMissingAgentState({ kind: "idle" }); } @@ -936,7 +919,6 @@ function ChatAgentContent({ isConnected, missingAgentState.kind, serverId, - shouldUseOptimisticStream, ]); const animatedContentStyle = useMemo( @@ -966,9 +948,6 @@ function ChatAgentContent({ isArchivingCurrentAgent={isArchivingCurrentAgent} agentState={agentState} effectiveAgent={effectiveAgent} - pendingCreate={pendingCreate} - shouldUseOptimisticStream={shouldUseOptimisticStream} - canFinalizePendingCreate={canFinalizePendingCreate} routeBottomAnchorRequest={routeBottomAnchorRequest} hasAppliedAuthoritativeHistory={hasAppliedAuthoritativeHistory} panelToast={panelToast} @@ -978,7 +957,6 @@ function ChatAgentContent({ handleAddImagesCallback={handleAddImagesCallback} handleComposerHeightChange={handleComposerHeightChange} handleMessageSent={handleMessageSent} - showPendingCreateSubmitLoading={showPendingCreateSubmitLoading} showHistorySyncOverlay={showHistorySyncOverlay} cwd={agentCwd} attentionController={attentionController} @@ -994,9 +972,6 @@ function ChatAgentReadyContent({ isArchivingCurrentAgent, agentState, effectiveAgent, - pendingCreate, - shouldUseOptimisticStream, - canFinalizePendingCreate, routeBottomAnchorRequest, hasAppliedAuthoritativeHistory, panelToast, @@ -1006,7 +981,6 @@ function ChatAgentReadyContent({ handleAddImagesCallback, handleComposerHeightChange, handleMessageSent, - showPendingCreateSubmitLoading, showHistorySyncOverlay, cwd, attentionController, @@ -1018,9 +992,6 @@ function ChatAgentReadyContent({ isArchivingCurrentAgent: boolean; agentState: ChatAgentSelectedState; effectiveAgent: AgentScreenAgent; - pendingCreate: PendingCreateAttempt | null; - shouldUseOptimisticStream: boolean; - canFinalizePendingCreate: boolean; routeBottomAnchorRequest: RouteBottomAnchorRequest; hasAppliedAuthoritativeHistory: boolean; panelToast: ReturnType; @@ -1030,7 +1001,6 @@ function ChatAgentReadyContent({ handleAddImagesCallback: (addImages: (images: ImageAttachment[]) => void) => void; handleComposerHeightChange: (height: number) => void; handleMessageSent: () => void; - showPendingCreateSubmitLoading: boolean; showHistorySyncOverlay: boolean; cwd: string; attentionController: ReturnType; @@ -1055,9 +1025,6 @@ function ChatAgentReadyContent({ serverId={serverId} agentId={agentId} agent={effectiveAgent} - pendingCreate={pendingCreate} - shouldUseOptimisticStream={shouldUseOptimisticStream} - canFinalizePendingCreate={canFinalizePendingCreate} routeBottomAnchorRequest={routeBottomAnchorRequest} hasAppliedAuthoritativeHistory={hasAppliedAuthoritativeHistory} toast={panelToast.api} @@ -1073,7 +1040,7 @@ function ChatAgentReadyContent({ isArchivingCurrentAgent={isArchivingCurrentAgent} archivedAt={agentState.archivedAt} cwd={cwd} - isSubmitLoading={showPendingCreateSubmitLoading} + isSubmitLoading={false} agentInputDraft={agentInputDraft} onAttentionInputFocus={attentionController.clearOnInputFocus} onAttentionPromptSend={attentionController.clearOnPromptSend} @@ -1113,9 +1080,6 @@ function AgentStreamSection({ serverId, agentId, agent, - pendingCreate, - shouldUseOptimisticStream, - canFinalizePendingCreate, routeBottomAnchorRequest, hasAppliedAuthoritativeHistory, toast, @@ -1125,9 +1089,6 @@ function AgentStreamSection({ serverId: string; agentId?: string; agent: AgentScreenAgent; - pendingCreate: PendingCreateAttempt | null; - shouldUseOptimisticStream: boolean; - canFinalizePendingCreate: boolean; routeBottomAnchorRequest: RouteBottomAnchorRequest; hasAppliedAuthoritativeHistory: boolean; toast: ReturnType["api"]; @@ -1163,98 +1124,6 @@ function AgentStreamSection({ } return new Map(pendingPermissionList.map((permission) => [permission.key, permission])); }, [pendingPermissionList]); - const setAgentStreamTail = useSessionStore((state) => state.setAgentStreamTail); - const markPendingCreateLifecycle = useCreateFlowStore((state) => state.markLifecycle); - const clearPendingCreate = useCreateFlowStore((state) => state.clear); - - const optimisticStreamItems = useMemo(() => { - if (!shouldUseOptimisticStream || !pendingCreate) { - return EMPTY_STREAM_ITEMS; - } - return [ - { - kind: "user_message", - id: pendingCreate.clientMessageId, - text: pendingCreate.text, - timestamp: new Date(pendingCreate.timestamp), - optimistic: true, - ...(pendingCreate.images && pendingCreate.images.length > 0 - ? { images: pendingCreate.images } - : {}), - ...(pendingCreate.attachments && pendingCreate.attachments.length > 0 - ? { attachments: pendingCreate.attachments } - : {}), - }, - ]; - }, [pendingCreate, shouldUseOptimisticStream]); - - const mergedStreamItems = useMemo(() => { - if (optimisticStreamItems.length === 0) { - return streamItems; - } - const optimistic = optimisticStreamItems[0]; - if (!optimistic) { - return streamItems; - } - const alreadyHasOptimistic = streamItems.some( - (item) => item.kind === "user_message" && item.id === optimistic.id, - ); - return alreadyHasOptimistic ? streamItems : [...optimisticStreamItems, ...streamItems]; - }, [optimisticStreamItems, streamItems]); - - useEffect(() => { - if (!shouldUseOptimisticStream || !pendingCreate) { - return; - } - const hasUserMessage = streamItems.some( - (item) => item.kind === "user_message" && item.id === pendingCreate.clientMessageId, - ); - if (!hasUserMessage || !canFinalizePendingCreate) { - return; - } - - const pendingImages = pendingCreate.images; - const pendingAttachments = pendingCreate.attachments; - const hasPendingImages = Boolean(pendingImages && pendingImages.length > 0); - const hasPendingAttachments = Boolean(pendingAttachments && pendingAttachments.length > 0); - if (agentId && (hasPendingImages || hasPendingAttachments)) { - setAgentStreamTail(serverId, (previous) => { - const current = previous.get(agentId); - if (!current) { - return previous; - } - - const merged = mergePendingCreateImages({ - streamItems: current, - clientMessageId: pendingCreate.clientMessageId, - images: pendingImages, - attachments: pendingAttachments, - }); - if (merged === current) { - return previous; - } - - const next = new Map(previous); - next.set(agentId, merged); - return next; - }); - } - markPendingCreateLifecycle({ - draftId: pendingCreate.draftId, - lifecycle: "sent", - }); - clearPendingCreate({ draftId: pendingCreate.draftId }); - }, [ - agentId, - canFinalizePendingCreate, - clearPendingCreate, - markPendingCreateLifecycle, - pendingCreate, - serverId, - setAgentStreamTail, - shouldUseOptimisticStream, - streamItems, - ]); return ( [0]) => { const normalized = normalizeAgentSnapshot(agentSnapshot, serverId); - retargetCurrentTab({ kind: "agent", agentId: agentSnapshot.id }); useSessionStore.getState().setAgents(serverId, (prev) => { const next = new Map(prev); next.set(agentSnapshot.id, normalized); return next; }); + retargetCurrentTab({ kind: "agent", agentId: agentSnapshot.id }); }, [retargetCurrentTab, serverId], ); diff --git a/packages/app/src/screens/workspace/workspace-pane-content.test.tsx b/packages/app/src/screens/workspace/workspace-pane-content.test.tsx index d700da58c..fa353542f 100644 --- a/packages/app/src/screens/workspace/workspace-pane-content.test.tsx +++ b/packages/app/src/screens/workspace/workspace-pane-content.test.tsx @@ -48,14 +48,14 @@ function ProbePanel() { return null; } -const tab: WorkspaceTabDescriptor = { +const agentTab: WorkspaceTabDescriptor = { key: "agent_agent-a", tabId: "agent_agent-a", kind: "agent", target: { kind: "agent", agentId: "agent-a" }, }; -function buildContent() { +function buildContent(tab: WorkspaceTabDescriptor = agentTab) { return buildWorkspacePaneContentModel({ tab, normalizedServerId: "server-a", @@ -120,4 +120,43 @@ describe("WorkspacePaneContent", () => { focusPane: expect.any(Function), }); }); + + it("keeps pane content mounted when a draft tab is retargeted in place", () => { + container = document.createElement("div"); + document.body.appendChild(container); + root = createRoot(container); + + const draftContent = buildContent({ + key: "draft-a", + tabId: "draft-a", + kind: "draft", + target: { kind: "draft", draftId: "draft-a" }, + }); + const agentContent = buildContent({ + key: "draft-a", + tabId: "draft-a", + kind: "agent", + target: { kind: "agent", agentId: "agent-a" }, + }); + + act(() => { + root?.render( + , + ); + }); + act(() => { + root?.render( + , + ); + }); + + expect(mountCount).toHaveBeenCalledTimes(1); + expect(unmountCount).not.toHaveBeenCalled(); + expect(snapshots).toHaveLength(2); + expect(snapshots[1]?.paneContextValue.tabId).toBe("draft-a"); + expect(snapshots[1]?.paneContextValue.target).toEqual({ + kind: "agent", + agentId: "agent-a", + }); + }); }); diff --git a/packages/app/src/screens/workspace/workspace-pane-content.tsx b/packages/app/src/screens/workspace/workspace-pane-content.tsx index 27fe2ba6c..3987980fd 100644 --- a/packages/app/src/screens/workspace/workspace-pane-content.tsx +++ b/packages/app/src/screens/workspace/workspace-pane-content.tsx @@ -42,7 +42,7 @@ export function buildWorkspacePaneContentModel({ const registration = getPanelRegistration(tab.kind); invariant(registration, `No panel registration for kind: ${tab.kind}`); return { - key: `${normalizedServerId}:${normalizedWorkspaceId}:${tab.tabId}:${tab.kind}`, + key: `${normalizedServerId}:${normalizedWorkspaceId}:${tab.tabId}`, Component: registration.component, paneContextValue: { serverId: normalizedServerId, diff --git a/packages/app/src/screens/workspace/workspace-screen.tsx b/packages/app/src/screens/workspace/workspace-screen.tsx index d11229f8f..4c2088da5 100644 --- a/packages/app/src/screens/workspace/workspace-screen.tsx +++ b/packages/app/src/screens/workspace/workspace-screen.tsx @@ -1746,9 +1746,6 @@ function WorkspaceScreenContent({ const unpinWorkspaceAgent = useWorkspaceLayoutStore((state) => state.unpinAgent); const hideWorkspaceAgent = useWorkspaceLayoutStore((state) => state.hideAgent); const retargetWorkspaceTab = useWorkspaceLayoutStore((state) => state.retargetTab); - const convertWorkspaceDraftToAgent = useWorkspaceLayoutStore( - (state) => state.convertDraftToAgent, - ); const reconcileWorkspaceTabs = useWorkspaceLayoutStore((state) => state.reconcileTabs); const splitWorkspacePane = useWorkspaceLayoutStore((state) => state.splitPane); const splitWorkspacePaneEmpty = useWorkspaceLayoutStore((state) => state.splitPaneEmpty); @@ -2842,10 +2839,6 @@ function WorkspaceScreenContent({ if (!persistenceKey) { return; } - if (input.tab.kind === "draft" && target.kind === "agent") { - convertWorkspaceDraftToAgent(persistenceKey, input.tab.tabId, target.agentId); - return; - } retargetWorkspaceTab(persistenceKey, input.tab.tabId, target); }, onOpenWorkspaceFile: (request: WorkspaceFileOpenRequest) => { @@ -2875,7 +2868,6 @@ function WorkspaceScreenContent({ openImportSheet, openWorkspaceChildTabFocused, persistenceKey, - convertWorkspaceDraftToAgent, retargetWorkspaceTab, ], ); diff --git a/packages/app/src/stores/create-flow-store.test.ts b/packages/app/src/stores/create-flow-store.test.ts index 78f0d6ec4..17844e4bd 100644 --- a/packages/app/src/stores/create-flow-store.test.ts +++ b/packages/app/src/stores/create-flow-store.test.ts @@ -71,4 +71,21 @@ describe("create-flow-store", () => { useCreateFlowStore.getState().clear({ draftId: "draft-1" }); expect(useCreateFlowStore.getState().pendingByDraftId["draft-1"]).toBeUndefined(); }); + + it("clears pending handoff state by agent", () => { + useCreateFlowStore.getState().setPending({ + draftId: "draft-1", + serverId: "server-1", + agentId: null, + clientMessageId: "msg-1", + text: "hello", + timestamp: Date.now(), + }); + useCreateFlowStore.getState().updateAgentId({ draftId: "draft-1", agentId: "agent-1" }); + useCreateFlowStore.getState().markLifecycle({ draftId: "draft-1", lifecycle: "sent" }); + + useCreateFlowStore.getState().clearByAgent({ serverId: "server-1", agentId: "agent-1" }); + + expect(useCreateFlowStore.getState().pendingByDraftId["draft-1"]).toBeUndefined(); + }); }); diff --git a/packages/app/src/stores/create-flow-store.ts b/packages/app/src/stores/create-flow-store.ts index 84cad6e49..fd6af5311 100644 --- a/packages/app/src/stores/create-flow-store.ts +++ b/packages/app/src/stores/create-flow-store.ts @@ -4,7 +4,7 @@ import type { AgentAttachment } from "@server/shared/messages"; export type CreateFlowLifecycleState = "active" | "abandoned" | "sent"; -interface PendingCreateAttempt { +export interface PendingCreateAttempt { draftId: string; serverId: string; agentId: string | null; @@ -23,6 +23,7 @@ interface CreateFlowState { markLifecycle: (input: { draftId: string; lifecycle: CreateFlowLifecycleState }) => void; rekeyDraft: (input: { fromDraftId: string; toDraftId: string }) => void; clear: (input: { draftId: string }) => void; + clearByAgent: (input: { serverId: string; agentId: string }) => void; clearAll: () => void; } @@ -89,5 +90,17 @@ export const useCreateFlowStore = create((set) => ({ const { [draftId]: _removed, ...rest } = state.pendingByDraftId; return { pendingByDraftId: rest }; }), + clearByAgent: ({ serverId, agentId }) => + set((state) => { + const next = Object.fromEntries( + Object.entries(state.pendingByDraftId).filter( + ([, pending]) => pending.serverId !== serverId || pending.agentId !== agentId, + ), + ); + if (Object.keys(next).length === Object.keys(state.pendingByDraftId).length) { + return state; + } + return { pendingByDraftId: next }; + }), clearAll: () => set({ pendingByDraftId: {} }), })); diff --git a/packages/app/src/stores/session-store.ts b/packages/app/src/stores/session-store.ts index 3f5f3ec95..61044035e 100644 --- a/packages/app/src/stores/session-store.ts +++ b/packages/app/src/stores/session-store.ts @@ -3,7 +3,12 @@ import { create } from "zustand"; import { subscribeWithSelector } from "zustand/middleware"; import type { DaemonClient } from "@server/client/daemon-client"; import type { AgentDirectoryEntry } from "@/types/agent-directory"; -import type { StreamItem } from "@/types/stream"; +import { + appendOptimisticUserMessageToStream, + type OptimisticUserMessagePlacement, + type StreamItem, + type UserMessageItem, +} from "@/types/stream"; import type { PendingPermission } from "@/types/shared"; import type { ComposerAttachment } from "@/attachments/types"; import type { AgentLifecycleStatus } from "@server/shared/agent-lifecycle"; @@ -351,6 +356,15 @@ interface SessionStoreActions { agentId: string, state: { tail?: StreamItem[]; head?: StreamItem[] }, ) => void; + appendOptimisticUserMessageToAgentStream: ( + serverId: string, + agentId: string, + message: UserMessageItem, + options: { + placement: OptimisticUserMessagePlacement; + skipIfUserMessageExists?: boolean; + }, + ) => boolean; clearAgentStreamHead: (serverId: string, agentId: string) => void; setAgentTimelineCursor: ( serverId: string, @@ -836,6 +850,50 @@ export const useSessionStore = create()( }); }, + appendOptimisticUserMessageToAgentStream: (serverId, agentId, message, options) => { + let didAppend = false; + set((prev) => { + const session = prev.sessions[serverId]; + if (!session) { + return prev; + } + + const currentTail = session.agentStreamTail.get(agentId) ?? []; + const currentHead = session.agentStreamHead.get(agentId) ?? []; + const result = appendOptimisticUserMessageToStream({ + tail: currentTail, + head: currentHead, + message, + placement: options.placement, + skipIfUserMessageExists: options.skipIfUserMessageExists, + }); + if (!result.changedTail && !result.changedHead) { + return prev; + } + + const nextTail = result.changedTail + ? new Map(session.agentStreamTail).set(agentId, result.tail) + : session.agentStreamTail; + const nextHead = result.changedHead + ? new Map(session.agentStreamHead).set(agentId, result.head) + : session.agentStreamHead; + didAppend = true; + + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { + ...session, + agentStreamTail: nextTail, + agentStreamHead: nextHead, + }, + }, + }; + }); + return didAppend; + }, + clearAgentStreamHead: (serverId, agentId) => { set((prev) => { const session = prev.sessions[serverId]; diff --git a/packages/app/src/timeline/session-stream-reducers.test.ts b/packages/app/src/timeline/session-stream-reducers.test.ts index 13bdec08e..496ac6c96 100644 --- a/packages/app/src/timeline/session-stream-reducers.test.ts +++ b/packages/app/src/timeline/session-stream-reducers.test.ts @@ -1,6 +1,11 @@ import { describe, expect, it } from "vitest"; import type { AgentStreamEventPayload } from "@server/shared/messages"; -import { hydrateStreamState, type AgentToolCallItem, type StreamItem } from "@/types/stream"; +import { + buildOptimisticUserMessage, + hydrateStreamState, + type AgentToolCallItem, + type StreamItem, +} from "@/types/stream"; import { createAgentStreamReducerQueue, processTimelineResponse, @@ -103,6 +108,17 @@ function makeAssistantItem(text: string, id = `assistant-${text.length}`): Strea }; } +function makeOptimisticUserMessage( + text: string, + id = `optimistic-${text.length}`, +): Extract { + return buildOptimisticUserMessage({ + id, + text, + timestamp: new Date(1000), + }); +} + function getAssistantTexts(items: StreamItem[]): string[] { return items .filter((item): item is Extract => { @@ -275,6 +291,76 @@ describe("processTimelineResponse", () => { expect(assistant?.timestamp.toISOString()).toBe("2025-01-01T12:00:04.000Z"); }); + it("reconciles an optimistic user message during tail replacement", () => { + const image = { + id: "optimistic-image", + mimeType: "image/png", + storageType: "web-indexeddb" as const, + storageKey: "optimistic-image", + createdAt: 1000, + }; + const attachment = { + type: "text" as const, + mimeType: "text/plain" as const, + text: "attached context", + title: "context.txt", + }; + const optimistic = buildOptimisticUserMessage({ + id: "optimistic-create-user", + text: "Analyze this", + timestamp: new Date(1000), + images: [image], + attachments: [attachment], + }); + + const result = processTimelineResponse({ + ...baseTimelineInput, + currentTail: [optimistic], + payload: { + ...baseTimelineInput.payload, + reset: true, + startCursor: { seq: 1 }, + endCursor: { seq: 1 }, + entries: [ + { + ...makeTimelineEntry(1, "Analyze this", "user_message"), + item: { + type: "user_message", + text: "Analyze this", + messageId: "canonical-create-user", + }, + }, + ], + }, + }); + + const userMessages = result.tail.filter((item) => item.kind === "user_message"); + expect(userMessages).toHaveLength(1); + expect(userMessages[0]).toMatchObject({ + id: "canonical-create-user", + text: "Analyze this", + images: [image], + attachments: [attachment], + }); + expect(userMessages[0]?.optimistic).toBeUndefined(); + }); + + it("keeps an unmatched optimistic user message during tail replacement", () => { + const optimistic = makeOptimisticUserMessage("still sending", "optimistic-unmatched"); + + const result = processTimelineResponse({ + ...baseTimelineInput, + currentTail: [optimistic], + payload: { + ...baseTimelineInput.payload, + reset: true, + entries: [], + }, + }); + + expect(result.tail).toEqual([optimistic]); + }); + it("sets cursor to null when reset=true but no cursors in payload", () => { const existingTail: StreamItem[] = [ { @@ -363,6 +449,40 @@ describe("processTimelineResponse", () => { expect(result.error).toBe(null); }); + it("reconciles an optimistic user message during an after-page response", () => { + const existingCursor: TimelineCursor = { + epoch: "epoch-1", + startSeq: 1, + endSeq: 1, + }; + const optimistic = makeOptimisticUserMessage("sent while catching up", "optimistic-after"); + + const result = processTimelineResponse({ + ...baseTimelineInput, + currentTail: [optimistic], + currentCursor: existingCursor, + payload: { + ...baseTimelineInput.payload, + epoch: "epoch-1", + entries: [ + { + ...makeTimelineEntry(2, "sent while catching up", "user_message"), + item: { + type: "user_message", + text: "sent while catching up", + messageId: "canonical-after", + }, + }, + ], + }, + }); + + 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(); + }); + it("keeps an active assistant head live when an incremental fetch accepts same-turn assistant text", () => { const existingCursor: TimelineCursor = { epoch: "epoch-1", @@ -483,6 +603,43 @@ describe("processTimelineResponse", () => { }); }); + it("does not reconcile an active optimistic user message from a before-page response", () => { + const optimistic = makeOptimisticUserMessage("active prompt", "optimistic-active"); + const existingCursor: TimelineCursor = { + epoch: "epoch-1", + startSeq: 3, + endSeq: 5, + }; + + const result = processTimelineResponse({ + ...baseTimelineInput, + currentTail: [optimistic], + currentCursor: existingCursor, + payload: { + ...baseTimelineInput.payload, + direction: "before", + epoch: "epoch-1", + startCursor: { seq: 1 }, + endCursor: { seq: 2 }, + entries: [ + { + ...makeTimelineEntry(1, "older prompt", "user_message"), + item: { + type: "user_message", + text: "older prompt", + messageId: "canonical-before", + }, + }, + ], + }, + }); + + 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); + }); + it("leaves the cursor alone when a before page makes no progress", () => { const currentTail: StreamItem[] = [ { diff --git a/packages/app/src/timeline/session-stream-reducers.ts b/packages/app/src/timeline/session-stream-reducers.ts index e8a0ee738..0f1c07f38 100644 --- a/packages/app/src/timeline/session-stream-reducers.ts +++ b/packages/app/src/timeline/session-stream-reducers.ts @@ -2,7 +2,7 @@ import type { AgentStreamEventPayload } from "@server/shared/messages"; import type { AgentLifecycleStatus } from "@server/shared/agent-lifecycle"; import type { Agent } from "@/stores/session-store"; import { useSessionStore } from "@/stores/session-store"; -import type { StreamItem } from "@/types/stream"; +import type { StreamItem, UserMessageItem } from "@/types/stream"; import { applyStreamEvent, hydrateStreamState, @@ -240,15 +240,22 @@ function applyTimelineReplacePath(args: { timelineUnits: TimelineUnit[]; payload: ProcessTimelineResponseInput["payload"]; bootstrapPolicy: ReturnType; + currentTail: StreamItem[]; currentHead: StreamItem[]; toHydratedEvents: ( units: TimelineUnit[], ) => Array<{ event: AgentStreamEventPayload; timestamp: Date }>; }): TimelinePathResult { - const { timelineUnits, payload, bootstrapPolicy, currentHead, toHydratedEvents } = args; + const { timelineUnits, payload, bootstrapPolicy, currentTail, currentHead, toHydratedEvents } = + args; const hydratedTail = hydrateStreamState(toHydratedEvents(timelineUnits), { source: "canonical" }); + const reconciledTail = reconcileOptimisticUsersAfterReplace({ + canonicalTail: hydratedTail, + previousTail: currentTail, + previousHead: currentHead, + }); const { tail, head } = preserveReplacePathAssistantHead({ - tail: hydratedTail, + tail: reconciledTail, currentHead, }); const cursor: TimelineCursor | null = @@ -266,6 +273,93 @@ function applyTimelineReplacePath(args: { return { tail, head, cursor, cursorChanged: true, sideEffects }; } +function collectOptimisticUserMessages(items: StreamItem[]): Array<{ + ordinal: number; + item: UserMessageItem; +}> { + const optimistic: 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 }); + } + ordinal += 1; + } + return optimistic; +} + +function mergeCanonicalUserWithOptimistic( + canonical: UserMessageItem, + optimistic: UserMessageItem, +): UserMessageItem { + const shouldPreserveImages = + (!canonical.images || canonical.images.length === 0) && + Boolean(optimistic.images && optimistic.images.length > 0); + const shouldPreserveAttachments = + (!canonical.attachments || canonical.attachments.length === 0) && + Boolean(optimistic.attachments && optimistic.attachments.length > 0); + return { + ...canonical, + ...(shouldPreserveImages ? { images: optimistic.images } : {}), + ...(shouldPreserveAttachments ? { attachments: optimistic.attachments } : {}), + }; +} + +function reconcileOptimisticUsersAfterReplace(params: { + canonicalTail: StreamItem[]; + previousTail: StreamItem[]; + previousHead: StreamItem[]; +}): StreamItem[] { + const optimisticUsers = collectOptimisticUserMessages([ + ...params.previousTail, + ...params.previousHead, + ]); + if (optimisticUsers.length === 0) { + return params.canonicalTail; + } + + const canonicalUserIndexes: number[] = []; + params.canonicalTail.forEach((item, index) => { + if (item.kind === "user_message") { + canonicalUserIndexes.push(index); + } + }); + + let changed = false; + const nextTail = [...params.canonicalTail]; + let searchFromOrdinal = 0; + const unmatched: UserMessageItem[] = []; + + for (const optimistic of optimisticUsers) { + const canonicalOrdinal = canonicalUserIndexes.findIndex( + (_index, ordinal) => ordinal >= Math.max(optimistic.ordinal, searchFromOrdinal), + ); + if (canonicalOrdinal < 0) { + unmatched.push(optimistic.item); + continue; + } + + const canonicalIndex = canonicalUserIndexes[canonicalOrdinal]; + const canonicalItem = canonicalIndex !== undefined ? nextTail[canonicalIndex] : undefined; + if (!canonicalItem || canonicalItem.kind !== "user_message") { + unmatched.push(optimistic.item); + continue; + } + nextTail[canonicalIndex] = mergeCanonicalUserWithOptimistic(canonicalItem, optimistic.item); + searchFromOrdinal = canonicalOrdinal + 1; + changed = true; + } + + if (unmatched.length === 0) { + return changed ? nextTail : params.canonicalTail; + } + + return [...nextTail, ...unmatched]; +} + interface IncrementalAcceptResult { acceptedUnits: TimelineUnit[]; cursor: TimelineCursor | undefined; @@ -531,6 +625,7 @@ export function processTimelineResponse( timelineUnits, payload, bootstrapPolicy, + currentTail, currentHead, toHydratedEvents, }) diff --git a/packages/app/src/types/stream.test.ts b/packages/app/src/types/stream.test.ts index 65d9e2881..abd7b2512 100644 --- a/packages/app/src/types/stream.test.ts +++ b/packages/app/src/types/stream.test.ts @@ -4,6 +4,8 @@ import { describe, expect, it } from "vitest"; import { applyStreamEvent, + appendOptimisticUserMessageToStream, + buildOptimisticUserMessage, clearOptimisticUserMessages, hydrateStreamState, mergeToolCallDetail, @@ -1030,6 +1032,102 @@ describe("turn lifecycle events", () => { assert.strictEqual(userMessage.optimistic, undefined); }); + it("replaces an optimistic image user message with the next canonical server user message", () => { + const optimisticTimestamp = 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(), + }; + const attachment = { + type: "text" as const, + mimeType: "text/plain" as const, + text: "context", + title: "context.txt", + }; + const optimistic = buildOptimisticUserMessage({ + id: "msg_optimistic_canonical", + text: "Analyze this", + timestamp: optimisticTimestamp, + images: [image], + attachments: [attachment], + }); + + const state = reduceStreamUpdate( + [optimistic], + { + type: "timeline", + provider: "claude", + item: { + type: "user_message", + text: "Analyze this", + messageId: "provider-owned-canonical", + }, + }, + new Date("2025-01-01T15:03:11Z"), + { source: "canonical" }, + ); + + const userMessages = state.filter((item) => item.kind === "user_message"); + assert.strictEqual(userMessages.length, 1); + const userMessage = userMessages[0]; + invariant(userMessage?.kind === "user_message"); + assert.strictEqual(userMessage.id, "provider-owned-canonical"); + assert.strictEqual(userMessage.optimistic, undefined); + 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", + text: "append once", + timestamp: new Date("2025-01-01T15:03:20Z"), + }); + const headItem: StreamItem = { + kind: "assistant_message", + id: "assistant-head", + text: "streaming", + timestamp: new Date("2025-01-01T15:03:19Z"), + }; + + const first = appendOptimisticUserMessageToStream({ + tail: [], + head: [headItem], + message: optimistic, + placement: "active-head", + }); + const second = appendOptimisticUserMessageToStream({ + tail: first.tail, + head: first.head, + 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("reconciles an optimistic user message that was pending in the streaming head", () => { const optimistic: StreamItem = { kind: "user_message", diff --git a/packages/app/src/types/stream.ts b/packages/app/src/types/stream.ts index 96edfdbaf..d6c0b36c9 100644 --- a/packages/app/src/types/stream.ts +++ b/packages/app/src/types/stream.ts @@ -64,6 +64,16 @@ export interface UserMessageItem { attachments?: AgentAttachment[]; } +export interface OptimisticUserMessageInput { + id: string; + text: string; + timestamp: Date; + images?: UserMessageImageAttachment[]; + attachments?: AgentAttachment[]; +} + +export type OptimisticUserMessagePlacement = "tail" | "active-head"; + export interface AssistantMessageItem { kind: "assistant_message"; id: string; @@ -206,11 +216,61 @@ function buildUserMessageItem(input: { }; } +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 } + : {}), + }; +} + +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))) + ) { + 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, + }; +} + function appendUserMessage( state: StreamItem[], text: string, timestamp: Date, - source: StreamUpdateSource, messageId?: string, ): StreamItem[] { const { chunk, hasContent } = normalizeChunk(text); @@ -220,10 +280,9 @@ function appendUserMessage( const chunkSeed = chunk.trim() || chunk; const entryId = messageId ?? createUniqueTimelineId(state, "user", chunkSeed, timestamp); - const optimisticIndex = - source === "live" - ? state.findIndex((entry) => entry.kind === "user_message" && entry.optimistic) - : -1; + const optimisticIndex = state.findIndex( + (entry) => entry.kind === "user_message" && entry.optimistic, + ); const optimistic = optimisticIndex >= 0 ? (state[optimisticIndex] as UserMessageItem) : null; const nextItem = buildUserMessageItem({ @@ -688,9 +747,7 @@ function reduceTimelineEvent( const item = event.item; switch (item.type) { case "user_message": - return finalizeActiveThoughts( - appendUserMessage(state, item.text, timestamp, source, item.messageId), - ); + return finalizeActiveThoughts(appendUserMessage(state, item.text, timestamp, item.messageId)); case "assistant_message": return finalizeActiveThoughts( appendAssistantMessage(state, item.text, timestamp, source, item.messageId), diff --git a/packages/app/src/utils/pending-create-images.test.ts b/packages/app/src/utils/pending-create-images.test.ts deleted file mode 100644 index ba84c1da2..000000000 --- a/packages/app/src/utils/pending-create-images.test.ts +++ /dev/null @@ -1,154 +0,0 @@ -import { describe, expect, it } from "vitest"; -import type { StreamItem } from "@/types/stream"; -import { mergePendingCreateImages } from "./pending-create-images"; -import type { AgentAttachment } from "@server/shared/messages"; - -function userMessage(params: { - id: string; - text: string; - images?: Array<{ - id: string; - storageType: "native-file"; - storageKey: string; - mimeType: string; - createdAt: number; - }>; - attachments?: AgentAttachment[]; -}): StreamItem { - return { - kind: "user_message", - id: params.id, - text: params.text, - timestamp: new Date("2026-01-01T00:00:00Z"), - ...(params.images ? { images: params.images } : {}), - ...(params.attachments ? { attachments: params.attachments } : {}), - }; -} - -function buildImage(id: string) { - return [ - { - id, - storageType: "native-file" as const, - storageKey: `/tmp/${id}.jpg`, - mimeType: "image/jpeg", - createdAt: Date.now(), - }, - ]; -} - -function buildReviewAttachment(): AgentAttachment { - return { - type: "review", - mimeType: "application/paseo-review", - cwd: "/repo", - mode: "base", - comments: [], - }; -} - -describe("mergePendingCreateImages", () => { - it("returns same reference when pending images are absent", () => { - const streamItems = [userMessage({ id: "msg-1", text: "hello" })]; - const result = mergePendingCreateImages({ - streamItems, - clientMessageId: "msg-1", - images: [], - }); - expect(result).toBe(streamItems); - }); - - it("merges images by clientMessageId when matched message has none", () => { - const streamItems = [userMessage({ id: "msg-1", text: "hello" })]; - const images = buildImage("image-1"); - const result = mergePendingCreateImages({ - streamItems, - clientMessageId: "msg-1", - images, - }); - - expect(result).not.toBe(streamItems); - const updated = result[0]; - expect(updated?.kind).toBe("user_message"); - if (updated?.kind !== "user_message") { - throw new Error("Expected user_message item"); - } - expect(updated.images).toEqual(images); - }); - - it("merges attachments by clientMessageId when matched message has none", () => { - const streamItems = [userMessage({ id: "msg-1", text: "hello" })]; - const attachments = [buildReviewAttachment()]; - const result = mergePendingCreateImages({ - streamItems, - clientMessageId: "msg-1", - attachments, - }); - - expect(result).not.toBe(streamItems); - const updated = result[0]; - expect(updated?.kind).toBe("user_message"); - if (updated?.kind !== "user_message") { - throw new Error("Expected user_message item"); - } - expect(updated.attachments).toEqual(attachments); - }); - - it("does not merge when clientMessageId does not match", () => { - const streamItems = [userMessage({ id: "msg-1", text: "same text" })]; - const images = buildImage("image-2"); - const result = mergePendingCreateImages({ - streamItems, - clientMessageId: "missing-id", - images, - }); - - expect(result).toBe(streamItems); - }); - - it("does not overwrite existing user message images", () => { - const existingImages = buildImage("existing"); - const streamItems = [userMessage({ id: "msg-1", text: "hello", images: existingImages })]; - const result = mergePendingCreateImages({ - streamItems, - clientMessageId: "msg-1", - images: buildImage("new"), - }); - - expect(result).toBe(streamItems); - const unchanged = result[0]; - expect(unchanged?.kind).toBe("user_message"); - if (unchanged?.kind !== "user_message") { - throw new Error("Expected user_message item"); - } - expect(unchanged.images).toEqual(existingImages); - }); - - it("does not overwrite existing user message attachments", () => { - const existingAttachments = [buildReviewAttachment()]; - const streamItems = [ - userMessage({ id: "msg-1", text: "hello", attachments: existingAttachments }), - ]; - const result = mergePendingCreateImages({ - streamItems, - clientMessageId: "msg-1", - attachments: [ - { - type: "github_issue", - mimeType: "application/github-issue", - number: 7, - title: "Issue", - url: "https://github.com/getpaseo/paseo/issues/7", - }, - ], - }); - - expect(result).toBe(streamItems); - const unchanged = result[0]; - expect(unchanged?.kind).toBe("user_message"); - if (unchanged?.kind !== "user_message") { - throw new Error("Expected user_message item"); - } - expect(unchanged.attachments).toEqual(existingAttachments); - }); -}); diff --git a/packages/app/src/utils/pending-create-images.ts b/packages/app/src/utils/pending-create-images.ts deleted file mode 100644 index ea700f754..000000000 --- a/packages/app/src/utils/pending-create-images.ts +++ /dev/null @@ -1,48 +0,0 @@ -import type { StreamItem, UserMessageImageAttachment } from "@/types/stream"; -import type { AgentAttachment } from "@server/shared/messages"; - -interface MergePendingCreateImagesParams { - streamItems: StreamItem[]; - clientMessageId: string; - images?: UserMessageImageAttachment[]; - attachments?: AgentAttachment[]; -} - -export function mergePendingCreateImages({ - streamItems, - clientMessageId, - images, - attachments, -}: MergePendingCreateImagesParams): StreamItem[] { - const hasPendingImages = Boolean(images && images.length > 0); - const hasPendingAttachments = Boolean(attachments && attachments.length > 0); - if (!hasPendingImages && !hasPendingAttachments) { - return streamItems; - } - - const targetIndex = streamItems.findIndex( - (item) => item.kind === "user_message" && item.id === clientMessageId, - ); - if (targetIndex < 0) { - return streamItems; - } - - const target = streamItems[targetIndex]; - if (target.kind !== "user_message") { - return streamItems; - } - const shouldMergeImages = hasPendingImages && (!target.images || target.images.length === 0); - const shouldMergeAttachments = - hasPendingAttachments && (!target.attachments || target.attachments.length === 0); - if (!shouldMergeImages && !shouldMergeAttachments) { - return streamItems; - } - - const next = [...streamItems]; - next[targetIndex] = { - ...target, - ...(shouldMergeImages ? { images } : {}), - ...(shouldMergeAttachments ? { attachments } : {}), - }; - return next; -}