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
This commit is contained in:
Bohdan Triapitsyn
2026-06-04 13:32:16 +03:00
parent 5d7b5c89f2
commit e9d7f913d9
7 changed files with 169 additions and 32 deletions
@@ -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] },
+8 -5
View File
@@ -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<string, unknown>
@@ -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
+12 -5
View File
@@ -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<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()
+2
View File
@@ -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,
+19 -6
View File
@@ -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
}
+26 -16
View File
@@ -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(