From 04c9aa57ddb5ce52564017abe57bf4c89ac093c1 Mon Sep 17 00:00:00 2001 From: Bohdan Triapitsyn Date: Wed, 3 Jun 2026 22:38:15 +0300 Subject: [PATCH] fix: show streaming chat messages after startup Read chat state from the active session directory Align status rows with the active session --- .../ui/src/components/chat/ChatContainer.tsx | 24 ++++++++------ packages/ui/src/hooks/useAssistantStatus.ts | 13 +++++--- packages/ui/src/hooks/useSessionActivity.ts | 3 +- packages/ui/src/sync/session-ui-store.ts | 32 ++++++++++++++----- 4 files changed, 49 insertions(+), 23 deletions(-) diff --git a/packages/ui/src/components/chat/ChatContainer.tsx b/packages/ui/src/components/chat/ChatContainer.tsx index b315c94f..dda0700a 100644 --- a/packages/ui/src/components/chat/ChatContainer.tsx +++ b/packages/ui/src/components/chat/ChatContainer.tsx @@ -355,6 +355,7 @@ export const ChatContainer: React.FC = ({ autoOpenDraft = tr const { t } = useI18n(); // Session UI state const currentSessionId = useSessionUIStore((s) => s.currentSessionId); + const currentSessionDirectory = useSessionUIStore((s) => s.currentSessionDirectory); const openNewSessionDraft = useSessionUIStore((s) => s.openNewSessionDraft); const setCurrentSession = useSessionUIStore((s) => s.setCurrentSession); const newSessionDraft = useSessionUIStore((s) => s.newSessionDraft); @@ -364,6 +365,7 @@ export const ChatContainer: React.FC = ({ autoOpenDraft = tr // Sync actions const sync = useSync(); const syncDirectory = useSyncDirectory(); + const effectiveSessionDirectory = currentSessionDirectory ?? syncDirectory; const ensureSessionRenderable = React.useCallback( (sessionId: string) => sync.ensureSessionRenderable(sessionId), [sync], @@ -396,45 +398,48 @@ export const ChatContainer: React.FC = ({ autoOpenDraft = tr [streamingMessageId], ), ); - const sessionMessageCount = useSessionMessageCount(currentSessionId ?? ''); + const sessionMessageCount = useSessionMessageCount(currentSessionId ?? '', effectiveSessionDirectory); const hasRenderableSessionSnapshot = useDirectorySync( React.useCallback( (state) => (currentSessionId ? getSessionMaterializationStatus(state, currentSessionId).renderable : false), [currentSessionId], ), + effectiveSessionDirectory, ); // Messages from sync system - const sessionMessageRecords = useSessionMessageRecords(currentSessionId ?? ''); + const sessionMessageRecords = useSessionMessageRecords(currentSessionId ?? '', effectiveSessionDirectory); const sessionMessages = currentSessionId ? sessionMessageRecords : EMPTY_MESSAGES; const sessionPrefetchInfo = React.useSyncExternalStore( React.useCallback( (notify) => currentSessionId - ? subscribeSessionPrefetch(syncDirectory, currentSessionId, notify) + ? subscribeSessionPrefetch(effectiveSessionDirectory, currentSessionId, notify) : () => undefined, - [currentSessionId, syncDirectory], + [currentSessionId, effectiveSessionDirectory], ), React.useCallback( - () => currentSessionId ? getSessionPrefetch(syncDirectory, currentSessionId) : undefined, - [currentSessionId, syncDirectory], + () => currentSessionId ? getSessionPrefetch(effectiveSessionDirectory, currentSessionId) : undefined, + [currentSessionId, effectiveSessionDirectory], ), React.useCallback(() => undefined, []), ); // Sessions from sync system - const sessions = useSessions(); + const sessions = useSessions(effectiveSessionDirectory); // Plan detection - watches messages for plan creation and signals store usePlanDetection(currentSessionId ?? '', sessionMessages); // Session status from sync system - const sessionStatusForCurrent = useSessionStatus(currentSessionId ?? '') ?? IDLE_SESSION_STATUS; + const sessionStatusForCurrent = useSessionStatus(currentSessionId ?? '', effectiveSessionDirectory) ?? IDLE_SESSION_STATUS; // Permissions & questions from sync system const allPermissions = useDirectorySync( React.useCallback((s) => s.permission ?? {}, []), + effectiveSessionDirectory, ); const allQuestions = useDirectorySync( React.useCallback((s) => s.question ?? {}, []), + effectiveSessionDirectory, ); // Convert Record → Map for blockingRequests helpers @@ -782,8 +787,9 @@ export const ChatContainer: React.FC = ({ autoOpenDraft = tr React.useEffect(() => { if (!currentSessionId) return; if (hasRenderableSessionSnapshot) return; + if (effectiveSessionDirectory !== syncDirectory) return; void ensureSessionRenderable(currentSessionId); - }, [currentSessionId, ensureSessionRenderable, hasRenderableSessionSnapshot]); + }, [currentSessionId, effectiveSessionDirectory, ensureSessionRenderable, hasRenderableSessionSnapshot, syncDirectory]); if (!currentSessionId && !draftOpen) { // With auto-open, the draft welcome opens on the next tick (effect below), diff --git a/packages/ui/src/hooks/useAssistantStatus.ts b/packages/ui/src/hooks/useAssistantStatus.ts index 516c01ea..0367a5b3 100644 --- a/packages/ui/src/hooks/useAssistantStatus.ts +++ b/packages/ui/src/hooks/useAssistantStatus.ts @@ -249,6 +249,7 @@ const getToolDisplayName = (part: ToolPart): string => { export function useAssistantStatus(): AssistantStatusSnapshot { const currentSessionId = useSessionUIStore((state) => state.currentSessionId); + const currentSessionDirectory = useSessionUIStore((state) => state.currentSessionDirectory); const rawSessionMessages = useDirectorySync( React.useCallback((state) => { @@ -256,7 +257,8 @@ export function useAssistantStatus(): AssistantStatusSnapshot { return EMPTY_MESSAGES; } return state.message[currentSessionId] ?? EMPTY_MESSAGES; - }, [currentSessionId]) + }, [currentSessionId]), + currentSessionDirectory ?? undefined, ); // Only subscribe to parts for the last assistant message — avoids re-render @@ -273,11 +275,12 @@ export function useAssistantStatus(): AssistantStatusSnapshot { const genericKey = `${currentSessionId ?? ''}:${lastAssistantId ?? ''}`; const parts = lastAssistantId ? (state.part[lastAssistantId] ?? EMPTY_PARTS) : EMPTY_PARTS; return encodeParsedStatus(createParsedStatus(parts, genericKey)); - }, [currentSessionId, lastAssistantId]) + }, [currentSessionId, lastAssistantId]), + currentSessionDirectory ?? undefined, ); - const sessionPermissionRequests = useSessionPermissions(currentSessionId ?? ''); - const sessionQuestionRequests = useSessionQuestions(currentSessionId ?? ''); + const sessionPermissionRequests = useSessionPermissions(currentSessionId ?? '', currentSessionDirectory ?? undefined); + const sessionQuestionRequests = useSessionQuestions(currentSessionId ?? '', currentSessionDirectory ?? undefined); const sessionAbortRecord = useSessionUIStore( React.useCallback((state) => { @@ -290,7 +293,7 @@ export function useAssistantStatus(): AssistantStatusSnapshot { const { phase: activityPhase, isWorking: isPhaseWorking } = useCurrentSessionActivity(); - const currentSessionStatus = useSessionStatus(currentSessionId ?? ''); + const currentSessionStatus = useSessionStatus(currentSessionId ?? '', currentSessionDirectory ?? undefined); const sessionRetryAttempt = currentSessionStatus?.type === 'retry' ? (currentSessionStatus as { type: 'retry'; attempt?: number }).attempt diff --git a/packages/ui/src/hooks/useSessionActivity.ts b/packages/ui/src/hooks/useSessionActivity.ts index f18edb5d..2c814193 100644 --- a/packages/ui/src/hooks/useSessionActivity.ts +++ b/packages/ui/src/hooks/useSessionActivity.ts @@ -66,5 +66,6 @@ export function useSessionActivity(sessionId: string | null | undefined, directo export function useCurrentSessionActivity(): SessionActivityResult { const currentSessionId = useSessionUIStore((state) => state.currentSessionId); - return useSessionActivity(currentSessionId); + const currentSessionDirectory = useSessionUIStore((state) => state.currentSessionDirectory); + return useSessionActivity(currentSessionId, currentSessionDirectory ?? undefined); } diff --git a/packages/ui/src/sync/session-ui-store.ts b/packages/ui/src/sync/session-ui-store.ts index 83c8fd13..b0377db6 100644 --- a/packages/ui/src/sync/session-ui-store.ts +++ b/packages/ui/src/sync/session-ui-store.ts @@ -199,6 +199,7 @@ export type SessionHistoryMeta = { export type SessionUIState = { currentSessionId: string | null + currentSessionDirectory: string | null newSessionDraft: NewSessionDraftState abortPromptSessionId: string | null abortPromptExpiresAt: number | null @@ -344,6 +345,10 @@ const resolveSessionDirectory = ( if (attachmentDirectory) return attachmentDirectory const metaPath = getWtMeta(sessionId)?.path if (typeof metaPath === "string" && metaPath.trim().length > 0) return normalizePath(metaPath) + const runtimeMemory = runtimeSessionMemory.get(runtimeMemoryKey()) + if (runtimeMemory?.sessionId === sessionId && runtimeMemory.directory) { + return normalizePath(runtimeMemory.directory) + } const sessions = getAllSyncSessions() const target = sessions.find((s) => s.id === sessionId) if (!target) return null @@ -391,6 +396,7 @@ const writeRuntimeSessionMemory = (key: string, patch: Partial()((set, get) => ({ currentSessionId: null, + currentSessionDirectory: null, newSessionDraft: { ...DEFAULT_DRAFT }, abortPromptSessionId: null, abortPromptExpiresAt: null, @@ -418,10 +424,6 @@ export const useSessionUIStore = create()((set, get) => ({ activeSessionByRuntime.set(key, id) const previousSessionId = get().currentSessionId - - // Set currentSessionId immediately so the skeleton renders without delay. - set({ currentSessionId: id }) - const directoryState = useDirectoryStore.getState() const sessionDir = resolveSessionDirectory( @@ -430,6 +432,10 @@ export const useSessionUIStore = create()((set, get) => ({ ) const fallbackDir = opencodeClient.getDirectory() ?? directoryState.currentDirectory ?? null const resolvedDir = (directoryHint ? normalizePath(directoryHint) : null) ?? sessionDir ?? fallbackDir + + // Set the directory together with the session id so chat hooks read the + // same child store that send/SSE events will update during startup races. + set({ currentSessionId: id, currentSessionDirectory: id ? resolvedDir ?? null : null }) writeRuntimeSessionMemory(key, { sessionId: id, directory: resolvedDir ?? null }) try { @@ -493,6 +499,7 @@ export const useSessionUIStore = create()((set, get) => ({ } set({ currentSessionId: restoredSessionId, + currentSessionDirectory: restoredSessionId ? restoredDirectory : null, newSessionDraft: restoredSessionId ? { ...DEFAULT_DRAFT } : restoredDraft, abortPromptSessionId: null, abortPromptExpiresAt: null, @@ -501,7 +508,7 @@ export const useSessionUIStore = create()((set, get) => ({ pendingChangesBarDismissed: new Map(), }) if (restoredSessionId) { - setActiveSession(opencodeClient.getDirectory() ?? "", restoredSessionId) + setActiveSession(restoredDirectory ?? opencodeClient.getDirectory() ?? "", restoredSessionId) } else { setActiveSession("", "") } @@ -575,6 +582,7 @@ export const useSessionUIStore = create()((set, get) => ({ ...nextDraft, }, currentSessionId: null, + currentSessionDirectory: null, error: null, }) @@ -1225,6 +1233,11 @@ export const useSessionUIStore = create()((set, get) => ({ }, getDirectoryForSession: (sessionId) => { + if (sessionId === get().currentSessionId && get().currentSessionDirectory) { + return get().currentSessionDirectory + } + const resolved = resolveSessionDirectory(sessionId, (sid) => get().worktreeMetadata.get(sid)) + if (resolved) return resolved const attachmentDirectory = getAttachedSessionDirectory(getAttachmentForSession(sessionId)) if (attachmentDirectory) return attachmentDirectory const sessions = getAllSyncSessions() @@ -1292,9 +1305,12 @@ export const useSessionUIStore = create()((set, get) => ({ // Handled by sync system's SSE stream }, - setSessionDirectory: () => { - // Session directory is owned by sync child stores via SSE events. - // This is now a no-op — kept for interface compatibility during migration. + setSessionDirectory: (sessionId, directory) => { + const normalized = normalizePath(directory) + if (sessionId === get().currentSessionId) { + set({ currentSessionDirectory: normalized }) + writeRuntimeSessionMemory(runtimeMemoryKey(), { sessionId, directory: normalized }) + } }, // ---------------------------------------------------------------------------