diff --git a/packages/ui/src/components/chat/message/parts/ToolPart.tsx b/packages/ui/src/components/chat/message/parts/ToolPart.tsx index 04bead90..67f5c659 100644 --- a/packages/ui/src/components/chat/message/parts/ToolPart.tsx +++ b/packages/ui/src/components/chat/message/parts/ToolPart.tsx @@ -12,7 +12,7 @@ import { Prism as SyntaxHighlighter } from 'react-syntax-highlighter'; import { useOptionalThemeSystem } from '@/contexts/useThemeSystem'; import { useDirectoryStore } from '@/stores/useDirectoryStore'; import { useSessionUIStore } from '@/sync/session-ui-store'; -import { useDirectorySync, useSessionMessageRecords } from '@/sync/sync-context'; +import { useDirectorySync, useSessionMessageRecords, useEnsureSessionMessages } from '@/sync/sync-context'; import { getSyncChildStores } from '@/sync/sync-refs'; import { useUIStore } from '@/stores/useUIStore'; import { useSessionActivity } from '@/hooks/useSessionActivity'; @@ -1928,6 +1928,7 @@ const ToolPart: React.FC = ({ const taskSessionId = explicitTaskSessionId ?? fallbackTaskSessionId; const childSessionMessages = useSessionMessageRecords(taskSessionId ?? '', currentDirectory); + useEnsureSessionMessages(taskSessionId ?? '', currentDirectory); const metadataTaskSummaryEntries = React.useMemo(() => { if (!isTaskTool) { diff --git a/packages/ui/src/sync/event-pipeline.ts b/packages/ui/src/sync/event-pipeline.ts index 2ce3d99b..8b056003 100644 --- a/packages/ui/src/sync/event-pipeline.ts +++ b/packages/ui/src/sync/event-pipeline.ts @@ -35,6 +35,8 @@ export type EventPipelineInput = { onReconnect?: () => void /** Called when the stream disconnects (heartbeat timeout, network error, or transport failure). */ onDisconnect?: (reason: string) => void + /** Called when transport switches (e.g. WS timeout → SSE fallback) without actual disconnection. */ + onTransportSwitch?: () => void transport?: "auto" | "ws" | "sse" } @@ -144,7 +146,7 @@ type DirectoryQueue = { } export function createEventPipeline(input: EventPipelineInput) { - const { sdk, onEvent, onReconnect, onDisconnect, routeDirectory, transport = "auto" } = input + const { sdk, onEvent, onReconnect, onDisconnect, onTransportSwitch, routeDirectory, transport = "auto" } = input const abort = new AbortController() let disconnected = false let lastEventId: string | undefined @@ -514,6 +516,12 @@ export function createEventPipeline(input: EventPipelineInput) { const code = typeof error === "object" && error !== null ? (error as { code?: unknown }).code : undefined if (currentTransport === "ws" && code === "WS_FALLBACK") { retryDelayMs = 0 + // Transport switch (WS → SSE fallback), not a real disconnection. + // No events were lost — the next attempt will use SSE and carry + // lastEventId for gapless replay. Notify consumer so it can set + // isConnected, but do NOT treat this as a disconnection requiring + // a full directory resync. + onTransportSwitch?.() } else if (!isAbortError(error)) { if (!streamErrorLogged) { streamErrorLogged = true diff --git a/packages/ui/src/sync/index.ts b/packages/ui/src/sync/index.ts index 36caf3e6..a186ae22 100644 --- a/packages/ui/src/sync/index.ts +++ b/packages/ui/src/sync/index.ts @@ -70,6 +70,7 @@ export { useSyncDirectory, useChildStoreManager, useSessionMessageRecords, + useEnsureSessionMessages, useSessionTextMessages, useUserMessageHistory, buildSessionMessageRecordsSnapshot, diff --git a/packages/ui/src/sync/sync-context.tsx b/packages/ui/src/sync/sync-context.tsx index ef15d374..49254f46 100644 --- a/packages/ui/src/sync/sync-context.tsx +++ b/packages/ui/src/sync/sync-context.tsx @@ -291,6 +291,21 @@ function getReconnectCandidateSessionIds(state: State) { } } + // Ensure parent sessions of any child sessions are also resynced. + // A child may have completed during the disconnect gap without the + // store knowing (SSE events lost). The parent's task tool part needs + // to reflect the child's final state. + const parentIds = new Set() + for (const session of state.session) { + const parentId = (session as Session & { parentID?: string | null }).parentID + if (parentId) { + parentIds.add(parentId) + } + } + for (const pid of parentIds) { + ids.add(pid) + } + return Array.from(ids) } @@ -735,7 +750,6 @@ async function resyncDirectoryAfterReconnect( .filter((record) => !!record?.info?.id) .map((record) => stripMessageDiffSnapshots(record.info)) .sort((a, b) => cmp(a.id, b.id)) - const nextMessageIds = new Set(nextMessages.map((message) => message.id)) store.setState((state: DirectoryStore) => { const sessions = [...state.session] @@ -755,13 +769,10 @@ async function resyncDirectoryAfterReconnect( sessionChanged = true } + // Merge parts: overwrite only messages present in the fetch snapshot. + // Do NOT delete parts for messages that may have been added by SSE + // events arriving between the fetch and the setState — those are more recent. const nextPartState = { ...state.part } - const previousMessages = state.message[sessionId] ?? [] - for (const message of previousMessages) { - if (!nextMessageIds.has(message.id)) { - delete nextPartState[message.id] - } - } for (const record of records) { const messageId = record?.info?.id if (!messageId) continue @@ -1104,7 +1115,25 @@ function handleEvent( } } - // Read live state, create targeted draft cloning ONLY fields the event + // Sync-layer parent resync: when a child session goes idle, schedule + // a targeted parts repair for the parent session. This ensures the + // parent's task tool part reflects the child's completion even when + // no ToolPart component is mounted. + if (payload.type === "session.idle") { + const idleSessionId = getSessionIdFromPayload(payload) + if (idleSessionId && resolvedDirectory && resolvedDirectory !== "global") { + const sessionState = store.getState() + const idleSession = sessionState.session.find((s) => s.id === idleSessionId) + const parentID = idleSession + ? (idleSession as Session & { parentID?: string | null }).parentID + : null + if (parentID) { + enqueuePartsRepair(resolvedDirectory, parentID, childStores) + } + } + } + + // Read live state, create targeted draft cloning ONLY fields that event // type will mutate. This preserves reference identity for untouched slices // so Zustand selectors skip re-renders for unrelated subscribers. const current = store.getState() @@ -1365,6 +1394,12 @@ export function SyncProvider(props: { onDisconnect: (reason) => { useConfigStore.setState({ isConnected: false, lastDisconnectReason: reason }) }, + onTransportSwitch: () => { + // Transport switched (e.g. WS timeout → SSE fallback) without + // actual disconnection. No events lost — just update connection + // state without triggering a full directory resync. + useConfigStore.setState({ isConnected: true }) + }, }) return cleanup }, [props.sdk, childStores, routingIndex, messageStreamTransport]) @@ -1848,6 +1883,79 @@ export function useSessionMessageRecords( return React.useSyncExternalStore(store.subscribe, getSnapshot, getSnapshot) } +/** + * Ensures a session's messages are loaded into the sync store. + * If the session exists in state.session but messages haven't been fetched + * (state.message[sessionID] is absent), triggers a background API fetch. + * + * This covers the case where a user navigates to an old parent session + * whose child session messages were never loaded — bootstrap only loads + * session metadata, not messages. + */ + +// Module-level in-flight tracking for useEnsureSessionMessages. +// Prevents redundant parallel fetches when multiple component instances +// (e.g. multiple ToolParts) request the same session's messages. +const _ensureMessagesLoading = new Set() + +export function useEnsureSessionMessages(sessionID: string, directory?: string) { + const store = useDirectoryStore(directory) + + React.useEffect(() => { + if (!sessionID) return + + const state = store.getState() + // Already loaded — nothing to do + if (Object.prototype.hasOwnProperty.call(state.message, sessionID)) return + // Session doesn't exist — nothing to load + if (!state.session.some((s) => s.id === sessionID)) return + + const dir = directory ?? opencodeClient.getDirectory() + const loadingKey = `${dir ?? ""}:${sessionID}` + // Already loading this session for this directory + if (_ensureMessagesLoading.has(loadingKey)) return + + _ensureMessagesLoading.add(loadingKey) + + void (async () => { + try { + const scopedClient = opencodeClient.getScopedSdkClient(dir ?? "") + const response = await scopedClient.session.messages({ + sessionID: sessionID, + limit: RECONNECT_MESSAGE_LIMIT, + }) + const records = (response.data ?? []).filter( + (record: { info?: { id?: string } }) => !!record?.info?.id, + ) + if (records.length === 0) return + + const nextMessages = records + .map((record: { info: Message }) => stripMessageDiffSnapshots(record.info)) + .filter((m: Message | null): m is Message => m !== null) + .sort((a: Message, b: Message) => cmp(a.id, b.id)) + + const nextPartState: Record = {} + for (const record of records) { + const messageId = record?.info?.id + if (!messageId) continue + nextPartState[messageId] = (record.parts ?? []) + .filter((part: Part) => !!part?.id && !RECONNECT_SKIP_PARTS.has(part.type)) + .sort((a: Part, b: Part) => cmp(a.id, b.id)) + } + + store.setState((state: DirectoryStore) => ({ + message: { ...state.message, [sessionID]: nextMessages }, + part: { ...state.part, ...nextPartState }, + })) + } catch { + // Transient failure — next navigation or reconnect will retry + } finally { + _ensureMessagesLoading.delete(loadingKey) + } + })() + }, [sessionID, store, directory]) +} + /** * Determines if a session is actively working. * Checks session_status and only falls back to incomplete assistant messages