From e9d7f913d9f11f345c5b0f3652bbefc72d3f4f22 Mon Sep 17 00:00:00 2001 From: Bohdan Triapitsyn Date: Thu, 4 Jun 2026 13:32:16 +0300 Subject: [PATCH] fix: route new prompt sync to target directory Keeps optimistic prompt state in the session directory Routes live assistant part updates using upstream event payloads Adds regressions for startup session switch races --- .../src/sync/__tests__/event-reducer.test.ts | 49 +++++++++++++++++ packages/ui/src/sync/event-reducer.ts | 13 +++-- packages/ui/src/sync/session-actions.test.ts | 53 +++++++++++++++++++ packages/ui/src/sync/session-actions.ts | 17 ++++-- packages/ui/src/sync/session-ui-store.ts | 2 + packages/ui/src/sync/sync-context.tsx | 25 ++++++--- packages/ui/src/sync/use-sync.ts | 42 +++++++++------ 7 files changed, 169 insertions(+), 32 deletions(-) diff --git a/packages/ui/src/sync/__tests__/event-reducer.test.ts b/packages/ui/src/sync/__tests__/event-reducer.test.ts index 518a4e50..0c624563 100644 --- a/packages/ui/src/sync/__tests__/event-reducer.test.ts +++ b/packages/ui/src/sync/__tests__/event-reducer.test.ts @@ -40,6 +40,21 @@ function partUpdatedEvent(): Event { } as Event } +function topLevelSessionOnlyPartUpdatedEvent(): Event { + return { + type: "message.part.updated", + properties: { + sessionID: "ses_1", + part: { + id: "prt_1", + messageID: "msg_1", + type: "text", + text: "hello", + }, + }, + } as Event +} + describe("applyDirectoryEvent", () => { test("returns typed materialization when delta arrives before parts", () => { const result = applyDirectoryEvent(state(), deltaEvent()) @@ -78,6 +93,40 @@ describe("applyDirectoryEvent", () => { }) }) + test("uses top-level session id and part message id for part update materialization", () => { + const draft = state() + const result = applyDirectoryEvent(draft, topLevelSessionOnlyPartUpdatedEvent()) + + expect(draft.part.msg_1.map((item) => item.id)).toEqual(["prt_1"]) + expect(result).toEqual({ + changed: true, + materialization: { + type: "incomplete-session-snapshot", + sessionID: "ses_1", + messageID: "msg_1", + partID: "prt_1", + }, + }) + }) + + test("uses top-level session id for delta materialization", () => { + const result = applyDirectoryEvent(state(), { + type: "message.part.delta", + properties: { + sessionID: "ses_1", + messageID: "msg_1", + partID: "prt_1", + field: "text", + delta: "hello", + }, + } as Event) + + expect(result).toEqual({ + changed: false, + materialization: { type: "incomplete-session-snapshot", sessionID: "ses_1", messageID: "msg_1", partID: "prt_1" }, + }) + }) + test("applies part update without materialization when owning message exists", () => { const draft = state({ message: { ses_1: [{ id: "msg_1", sessionID: "ses_1", role: "assistant", time: { created: 1 } } as never] }, diff --git a/packages/ui/src/sync/event-reducer.ts b/packages/ui/src/sync/event-reducer.ts index 33de3e60..c925b3d0 100644 --- a/packages/ui/src/sync/event-reducer.ts +++ b/packages/ui/src/sync/event-reducer.ts @@ -339,13 +339,15 @@ export function applyDirectoryEvent( } case "message.part.updated": { - const part = (event.properties as { part: Part }).part + const props = event.properties as { sessionID?: string; part: Part } + const part = props.part if (SKIP_PARTS.has(part.type)) { syncDebug.reducer.partSkipped((part as { messageID: string }).messageID, part.id, part.type) return false } - const messageID = (part as { messageID: string }).messageID - const sessionID = (part as { sessionID?: string }).sessionID + const messageID = (part as { messageID?: string }).messageID + const sessionID = props.sessionID ?? (part as { sessionID?: string }).sessionID + if (!messageID) return false const missingOwningMessage = !hasMessage(draft, sessionID, messageID) const parts = draft.part[messageID] if (!parts) { @@ -413,6 +415,7 @@ export function applyDirectoryEvent( case "message.part.delta": { const props = event.properties as { + sessionID?: string messageID: string partID: string field: string @@ -423,7 +426,7 @@ export function applyDirectoryEvent( syncDebug.reducer.partDeltaNoParts(props.messageID, props.partID) return { changed: false, - materialization: { type: "incomplete-session-snapshot", messageID: props.messageID, partID: props.partID }, + materialization: { type: "incomplete-session-snapshot", sessionID: props.sessionID, messageID: props.messageID, partID: props.partID }, } } const result = Binary.search(parts, props.partID, (p) => p.id) @@ -431,7 +434,7 @@ export function applyDirectoryEvent( syncDebug.reducer.partDeltaNotFound(props.messageID, props.partID) return { changed: false, - materialization: { type: "incomplete-session-snapshot", messageID: props.messageID, partID: props.partID }, + materialization: { type: "incomplete-session-snapshot", sessionID: props.sessionID, messageID: props.messageID, partID: props.partID }, } } const existing = parts[result.index] as Record diff --git a/packages/ui/src/sync/session-actions.test.ts b/packages/ui/src/sync/session-actions.test.ts index 4e8608de..e1964473 100644 --- a/packages/ui/src/sync/session-actions.test.ts +++ b/packages/ui/src/sync/session-actions.test.ts @@ -159,6 +159,9 @@ import { INITIAL_STATE } from "./types" import type { DirectoryStore } from "./child-store" import type { Message, OpencodeClient, Part, Session } from "@opencode-ai/sdk/v2/client" +type OptimisticAddCall = { sessionID: string; directory?: string | null; message: Message; parts: Part[] } +type OptimisticRemoveCall = { sessionID: string; directory?: string | null; messageID: string } + function createStore( permissions: Record, state?: Partial, @@ -183,6 +186,56 @@ function createChildStores(entries: Array<[string, StoreApi]>) { } as unknown as import("./child-store").ChildStoreManager } +describe("optimisticSend target directory", () => { + beforeEach(() => { + replyCalls.length = 0 + scopedClientDirectories.length = 0 + }) + + test("passes the prompt directory to optimistic state during session switch races", async () => { + const currentStore = createStore({}) + const targetStore = createStore({}) + const childStores = createChildStores([ + ["/current/project", currentStore], + ["/target/project", targetStore], + ]) + let optimisticAdd: OptimisticAddCall | null = null + let optimisticRemove: OptimisticRemoveCall | null = null + let sentMessageID = "" + + const { optimisticSend, setActionRefs, setOptimisticRefs } = await import("./session-actions") + setActionRefs(mockSdk as unknown as OpencodeClient, childStores, () => "/current/project") + setOptimisticRefs( + (input) => { + optimisticAdd = input + }, + (input) => { + optimisticRemove = input + }, + ) + + await optimisticSend({ + sessionId: "session-new", + directory: "/target/project", + content: "hello", + providerID: "provider", + modelID: "model", + send: async (messageID) => { + sentMessageID = messageID + }, + }) + + expect(optimisticAdd).not.toBeNull() + const add = optimisticAdd as unknown as OptimisticAddCall + expect(add.directory).toBe("/target/project") + expect(add.sessionID).toBe("session-new") + expect(add.message.id).toBe(sentMessageID) + expect(optimisticRemove).toBe(null) + expect(targetStore.getState().session_status["session-new"]?.type).toBe("busy") + expect(currentStore.getState().session_status["session-new"]).toBe(undefined) + }) +}) + describe("respondToPermission passes directory", () => { beforeEach(() => { replyCalls.length = 0 diff --git a/packages/ui/src/sync/session-actions.ts b/packages/ui/src/sync/session-actions.ts index 2c297344..f4360c85 100644 --- a/packages/ui/src/sync/session-actions.ts +++ b/packages/ui/src/sync/session-actions.ts @@ -26,8 +26,11 @@ const UNREVERT_REFETCH_RETRY_MS = 150 let _sdk: OpencodeClient | null = null let _childStores: ChildStoreManager | null = null let _getDirectory: () => string = () => "" -let _optimisticAdd: ((input: { sessionID: string; message: Message; parts: Part[] }) => void) | null = null -let _optimisticRemove: ((input: { sessionID: string; messageID: string }) => void) | null = null +type OptimisticAddInput = { sessionID: string; directory?: string | null; message: Message; parts: Part[] } +type OptimisticRemoveInput = { sessionID: string; directory?: string | null; messageID: string } + +let _optimisticAdd: ((input: OptimisticAddInput) => void) | null = null +let _optimisticRemove: ((input: OptimisticRemoveInput) => void) | null = null const wait = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms)) @@ -84,8 +87,8 @@ export function setActionRefs( } export function setOptimisticRefs( - add: (input: { sessionID: string; message: Message; parts: Part[] }) => void, - remove: (input: { sessionID: string; messageID: string }) => void, + add: (input: OptimisticAddInput) => void, + remove: (input: OptimisticRemoveInput) => void, ) { _optimisticAdd = add _optimisticRemove = remove @@ -539,6 +542,7 @@ export async function optimisticSend(input: { providerID: string modelID: string agent?: string + directory?: string | null files?: Array<{ type: "file"; mime: string; url: string; filename: string }> /** The actual API call — receives the optimistic messageID so the server can use the same ID */ send: (messageID: string) => Promise @@ -549,7 +553,8 @@ export async function optimisticSend(input: { await waitForConnectionOrThrow() - const store = dirStore() + const targetDirectory = input.directory ?? dir() + const store = targetDirectory ? dirStoreForDirectory(targetDirectory) : dirStore() const messageID = ascendingId("msg") const textPartId = ascendingId("prt") @@ -579,6 +584,7 @@ export async function optimisticSend(input: { // Insert into store + register in shadow Map (for mergeOptimisticPage cleanup) _optimisticAdd({ sessionID: input.sessionId, + directory: targetDirectory, message: optimisticMessage, parts: optimisticParts, }) @@ -598,6 +604,7 @@ export async function optimisticSend(input: { // Rollback via optimistic infrastructure _optimisticRemove({ sessionID: input.sessionId, + directory: targetDirectory, messageID, }) const s = store.getState() diff --git a/packages/ui/src/sync/session-ui-store.ts b/packages/ui/src/sync/session-ui-store.ts index b0377db6..0c2777f1 100644 --- a/packages/ui/src/sync/session-ui-store.ts +++ b/packages/ui/src/sync/session-ui-store.ts @@ -111,6 +111,7 @@ export function routeMessage(params: { providerID: params.providerID, modelID: params.modelID, agent: params.agent, + directory: requestDirectory, files: params.files, send: (messageID) => opencodeClient.sendCommand({ id: params.sessionId, @@ -135,6 +136,7 @@ export function routeMessage(params: { providerID: params.providerID, modelID: params.modelID, agent: params.agent, + directory: requestDirectory, files: params.files, send: (messageID) => opencodeClient.sendMessage({ id: params.sessionId, diff --git a/packages/ui/src/sync/sync-context.tsx b/packages/ui/src/sync/sync-context.tsx index 4a0f832a..225af12f 100644 --- a/packages/ui/src/sync/sync-context.tsx +++ b/packages/ui/src/sync/sync-context.tsx @@ -571,11 +571,21 @@ const getSessionIdFromPayload = (event: Event): string | null => { } if (event.type === "message.part.updated") { + const sessionID = props.sessionID + if (typeof sessionID === "string" && sessionID.length > 0) { + return sessionID + } + const part = props.part if (!part || typeof part !== "object") { return null } - const sessionID = (part as { sessionID?: unknown }).sessionID + const partSessionID = (part as { sessionID?: unknown }).sessionID + return typeof partSessionID === "string" && partSessionID.length > 0 ? partSessionID : null + } + + if (event.type === "message.part.delta" || event.type === "message.part.removed") { + const sessionID = props.sessionID return typeof sessionID === "string" && sessionID.length > 0 ? sessionID : null } @@ -658,8 +668,8 @@ const getMessageIdFromPayload = (event: Event): string | null => { if (!part || typeof part !== "object") { return null } - const messageID = (part as { messageID?: unknown }).messageID - return typeof messageID === "string" && messageID.length > 0 ? messageID : null + const partMessageID = (part as { messageID?: unknown }).messageID + return typeof partMessageID === "string" && partMessageID.length > 0 ? partMessageID : null } return null @@ -957,9 +967,12 @@ const updateRoutingIndexFromEvent = ( } case "message.part.updated": { - const part = (payload.properties as { part?: Part }).part as (Part & { sessionID?: string; messageID?: string }) | undefined - if (part?.messageID && part.sessionID) { - setIndexedMessage(routingIndex, part.sessionID, part.messageID, directory) + const props = payload.properties as { sessionID?: string; part?: Part } + const part = props.part as (Part & { sessionID?: string; messageID?: string }) | undefined + const sessionID = part?.sessionID ?? props.sessionID + const messageID = part?.messageID + if (messageID && sessionID) { + setIndexedMessage(routingIndex, sessionID, messageID, directory) } return } diff --git a/packages/ui/src/sync/use-sync.ts b/packages/ui/src/sync/use-sync.ts index 91509833..a32320fe 100644 --- a/packages/ui/src/sync/use-sync.ts +++ b/packages/ui/src/sync/use-sync.ts @@ -257,16 +257,16 @@ export function useSync() { // Optimistic operations const getOptimistic = useCallback( - (sessionID: string): OptimisticItem[] => { - const key = `${directory}\n${sessionID}` + (sessionID: string, directoryOverride?: string | null): OptimisticItem[] => { + const key = `${directoryOverride || directory}\n${sessionID}` return [...(optimistic.current.get(key)?.values() ?? [])] }, [directory], ) const setOptimistic = useCallback( - (sessionID: string, item: OptimisticItem) => { - const key = `${directory}\n${sessionID}` + (sessionID: string, item: OptimisticItem, directoryOverride?: string | null) => { + const key = `${directoryOverride || directory}\n${sessionID}` const list = optimistic.current.get(key) const sorted: OptimisticItem = { message: item.message, parts: sortParts(item.parts) } if (list) { @@ -279,8 +279,8 @@ export function useSync() { ) const clearOptimistic = useCallback( - (sessionID: string, messageID?: string) => { - const key = `${directory}\n${sessionID}` + (sessionID: string, messageID?: string, directoryOverride?: string | null) => { + const key = `${directoryOverride || directory}\n${sessionID}` if (!messageID) { optimistic.current.delete(key) return @@ -293,6 +293,14 @@ export function useSync() { [directory], ) + const getOptimisticStore = useCallback( + (directoryOverride?: string | null) => { + if (!directoryOverride || directoryOverride === directory) return store + return childStores.ensureChild(directoryOverride, { bootstrap: false }) + }, + [childStores, directory, store], + ) + // Fetch messages from API const fetchMessages = useCallback( async (sessionID: string, limit: number, before?: string) => { @@ -490,9 +498,10 @@ export function useSync() { // Optimistic add (for prompt submission) const optimisticAdd = useCallback( - (input: { sessionID: string; message: Message; parts: Part[] }) => { - setOptimistic(input.sessionID, { message: input.message, parts: input.parts }) - const current = store.getState() + (input: { sessionID: string; directory?: string | null; message: Message; parts: Part[] }) => { + setOptimistic(input.sessionID, { message: input.message, parts: input.parts }, input.directory) + const targetStore = getOptimisticStore(input.directory) + const current = targetStore.getState() const message = { ...current.message } const part = { ...current.part } @@ -505,16 +514,17 @@ export function useSync() { // Insert parts part[input.message.id] = sortParts(input.parts) - store.setState({ message, part }) + targetStore.setState({ message, part }) }, - [store, setOptimistic], + [getOptimisticStore, setOptimistic], ) // Optimistic remove (for rollback on error) const optimisticRemove = useCallback( - (input: { sessionID: string; messageID: string }) => { - clearOptimistic(input.sessionID, input.messageID) - const current = store.getState() + (input: { sessionID: string; directory?: string | null; messageID: string }) => { + clearOptimistic(input.sessionID, input.messageID, input.directory) + const targetStore = getOptimisticStore(input.directory) + const current = targetStore.getState() const message = { ...current.message } const part = { ...current.part } @@ -529,9 +539,9 @@ export function useSync() { } delete part[input.messageID] - store.setState({ message, part }) + targetStore.setState({ message, part }) }, - [store, clearOptimistic], + [clearOptimistic, getOptimisticStore], ) return useMemo(