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

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] },

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 messageID = (part as { messageID?: string }).messageID
-      const sessionID = (part as { sessionID?: string }).sessionID
+      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<string, unknown>

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<string, PermissionRequest[]>,
   state?: Partial<DirectoryStore>,
@@ -183,6 +186,56 @@ function createChildStores(entries: Array<[string, StoreApi<DirectoryStore>]>) {
   } 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

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
+type OptimisticAddInput = { sessionID: string; directory?: string | null; message: Message; parts: Part[] }
-let _optimisticRemove: ((input: { sessionID: string; messageID: string }) => void) | null = null
+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,
+  add: (input: OptimisticAddInput) => void,
-  remove: (input: { sessionID: string; messageID: string }) => 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<void>
@@ -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()

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,

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
+    const partMessageID = (part as { messageID?: unknown }).messageID
-    return typeof messageID === "string" && messageID.length > 0 ? messageID : null
+    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
+      const props = payload.properties as { sessionID?: string; part?: Part }
-      if (part?.messageID && part.sessionID) {
+      const part = props.part as (Part & { sessionID?: string; messageID?: string }) | undefined
-        setIndexedMessage(routingIndex, part.sessionID, part.messageID, directory)
+      const sessionID = part?.sessionID ?? props.sessionID
+      const messageID = part?.messageID
+      if (messageID && sessionID) {
+        setIndexedMessage(routingIndex, sessionID, messageID, directory)
       }
       return
     }

packages/ui/src/sync/use-sync.ts

@@ -257,16 +257,16 @@ export function useSync() {
 
   // Optimistic operations
   const getOptimistic = useCallback(
-    (sessionID: string): OptimisticItem[] => {
+    (sessionID: string, directoryOverride?: string | null): OptimisticItem[] => {
-      const key = `${directory}\n${sessionID}`
+      const key = `${directoryOverride || directory}\n${sessionID}`
       return [...(optimistic.current.get(key)?.values() ?? [])]
     },
     [directory],
   )
 
   const setOptimistic = useCallback(
-    (sessionID: string, item: OptimisticItem) => {
+    (sessionID: string, item: OptimisticItem, directoryOverride?: string | null) => {
-      const key = `${directory}\n${sessionID}`
+      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) => {
+    (sessionID: string, messageID?: string, directoryOverride?: string | null) => {
-      const key = `${directory}\n${sessionID}`
+      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[] }) => {
+    (input: { sessionID: string; directory?: string | null; message: Message; parts: Part[] }) => {
-      setOptimistic(input.sessionID, { message: input.message, parts: input.parts })
+      setOptimistic(input.sessionID, { message: input.message, parts: input.parts }, input.directory)
-      const current = store.getState()
+      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 }) => {
+    (input: { sessionID: string; directory?: string | null; messageID: string }) => {
-      clearOptimistic(input.sessionID, input.messageID)
+      clearOptimistic(input.sessionID, input.messageID, input.directory)
-      const current = store.getState()
+      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(