diff --git a/src/components/app/AppContent.tsx b/src/components/app/AppContent.tsx index a542082a..4dd6979a 100644 --- a/src/components/app/AppContent.tsx +++ b/src/components/app/AppContent.tsx @@ -192,7 +192,9 @@ function AppContentInner() { onSessionNotProcessing={markSessionAsNotProcessing} processingSessions={processingSessions} onReplaceTemporarySession={replaceTemporarySession} - onNavigateToSession={(targetSessionId: string) => navigate(`/session/${targetSessionId}`)} + onNavigateToSession={(targetSessionId: string, options) => + navigate(`/session/${targetSessionId}`, { replace: Boolean(options?.replace) }) + } onShowSettings={() => setShowSettings(true)} externalMessageUpdate={externalMessageUpdate} newSessionTrigger={newSessionTrigger} diff --git a/src/components/chat/hooks/useChatRealtimeHandlers.ts b/src/components/chat/hooks/useChatRealtimeHandlers.ts index 855ee788..342ea117 100644 --- a/src/components/chat/hooks/useChatRealtimeHandlers.ts +++ b/src/components/chat/hooks/useChatRealtimeHandlers.ts @@ -1,7 +1,8 @@ import { useEffect, useRef } from 'react'; import type { Dispatch, MutableRefObject, SetStateAction } from 'react'; + import { usePaletteOps } from '../../../contexts/PaletteOpsContext'; -import type { PendingPermissionRequest } from '../types/types'; +import type { PendingPermissionRequest, SessionNavigationOptions } from '../types/types'; import type { Project, ProjectSession, LLMProvider } from '../../../types/app'; import type { SessionStore, NormalizedMessage } from '../../../stores/useSessionStore'; @@ -67,7 +68,7 @@ interface UseChatRealtimeHandlersArgs { onSessionProcessing?: (sessionId?: string | null) => void; onSessionNotProcessing?: (sessionId?: string | null) => void; onReplaceTemporarySession?: (sessionId?: string | null) => void; - onNavigateToSession?: (sessionId: string) => void; + onNavigateToSession?: (sessionId: string, options?: SessionNavigationOptions) => void; onWebSocketReconnect?: () => void; sessionStore: SessionStore; } @@ -273,13 +274,53 @@ export function useChatRealtimeHandlers({ break; } - // Clear pending session + const actualSessionId = + typeof msg.actualSessionId === 'string' && msg.actualSessionId.trim().length > 0 + ? msg.actualSessionId + : null; const pendingSessionId = sessionStorage.getItem('pendingSessionId'); - if (pendingSessionId && !currentSessionId && msg.exitCode === 0) { - const actualId = msg.actualSessionId || pendingSessionId; - setCurrentSessionId(actualId); - if (msg.actualSessionId) { - onNavigateToSession?.(actualId); + const completedSuccessfully = msg.exitCode === undefined || msg.exitCode === 0; + const isVisibleSession = + Boolean( + sid + && ( + sid === activeViewSessionId + || sid === pendingSessionId + || pendingViewSessionRef.current?.sessionId === sid + ), + ); + + if (actualSessionId && sid && actualSessionId !== sid) { + sessionStore.replaceSessionId(sid, actualSessionId); + + if (isVisibleSession) { + setCurrentSessionId(actualSessionId); + + if (pendingViewSessionRef.current) { + const pendingSession = pendingViewSessionRef.current.sessionId; + if (!pendingSession || pendingSession === sid) { + pendingViewSessionRef.current.sessionId = actualSessionId; + } + } + } + + if (completedSuccessfully && pendingSessionId === sid) { + sessionStorage.removeItem('pendingSessionId'); + } + + if (isVisibleSession) { + onNavigateToSession?.(actualSessionId, { replace: true }); + setTimeout(() => { void paletteOps.refreshProjects(); }, 500); + } + break; + } + + // Clear pending session + if (pendingSessionId && !currentSessionId && completedSuccessfully) { + const resolvedSessionId = actualSessionId || pendingSessionId; + setCurrentSessionId(resolvedSessionId); + if (actualSessionId) { + onNavigateToSession?.(resolvedSessionId, { replace: true }); } sessionStorage.removeItem('pendingSessionId'); setTimeout(() => { void paletteOps.refreshProjects(); }, 500); diff --git a/src/components/chat/hooks/useChatSessionState.ts b/src/components/chat/hooks/useChatSessionState.ts index a241e7bc..6bff4a88 100644 --- a/src/components/chat/hooks/useChatSessionState.ts +++ b/src/components/chat/hooks/useChatSessionState.ts @@ -1,11 +1,13 @@ import { useCallback, useEffect, useLayoutEffect, useMemo, useRef, useState } from 'react'; import type { MutableRefObject } from 'react'; + import { authenticatedFetch } from '../../../utils/api'; -import type { ChatMessage, Provider } from '../types/types'; import type { Project, ProjectSession, LLMProvider } from '../../../types/app'; -import { createCachedDiffCalculator, type DiffCalculator } from '../utils/messageTransforms'; -import { normalizedToChatMessages } from './useChatMessages'; import type { SessionStore, NormalizedMessage } from '../../../stores/useSessionStore'; +import type { ChatMessage, Provider } from '../types/types'; +import { createCachedDiffCalculator, type DiffCalculator } from '../utils/messageTransforms'; + +import { normalizedToChatMessages } from './useChatMessages'; const MESSAGES_PER_PAGE = 20; const INITIAL_VISIBLE_MESSAGES = 100; @@ -211,6 +213,7 @@ export function useChatSessionState({ const activeSessionId = selectedSession?.id || currentSessionId || null; const [pendingUserMessage, setPendingUserMessage] = useState(null); + const flushedPendingUserMessageRef = useRef(null); // Tell the store which session we're viewing so it only re-renders for this one const prevActiveForStoreRef = useRef(null); @@ -219,17 +222,29 @@ export function useChatSessionState({ sessionStore.setActiveSession(activeSessionId); } - // When a real session ID arrives and we have a pending user message, flush it to the store - const prevActiveSessionRef = useRef(null); - if (activeSessionId && activeSessionId !== prevActiveSessionRef.current && pendingUserMessage) { + useEffect(() => { + if (!pendingUserMessage) { + flushedPendingUserMessageRef.current = null; + return; + } + + if (!activeSessionId) { + return; + } + + if (flushedPendingUserMessageRef.current === pendingUserMessage) { + return; + } + const prov = (localStorage.getItem('selected-provider') as LLMProvider) || 'claude'; const normalized = chatMessageToNormalized(pendingUserMessage, activeSessionId, prov); if (normalized) { sessionStore.appendRealtime(activeSessionId, normalized); } + + flushedPendingUserMessageRef.current = pendingUserMessage; setPendingUserMessage(null); - } - prevActiveSessionRef.current = activeSessionId; + }, [activeSessionId, pendingUserMessage, sessionStore]); const storeMessages = activeSessionId ? sessionStore.getMessages(activeSessionId) : []; diff --git a/src/components/chat/types/types.ts b/src/components/chat/types/types.ts index 6cce91fc..81bd5a5b 100644 --- a/src/components/chat/types/types.ts +++ b/src/components/chat/types/types.ts @@ -91,6 +91,10 @@ export interface Question { multiSelect?: boolean; } +export type SessionNavigationOptions = { + replace?: boolean; +}; + export interface ChatInterfaceProps { selectedProject: Project | null; selectedSession: ProjectSession | null; @@ -105,7 +109,7 @@ export interface ChatInterfaceProps { onSessionNotProcessing?: (sessionId?: string | null) => void; processingSessions?: Set; onReplaceTemporarySession?: (sessionId?: string | null) => void; - onNavigateToSession?: (targetSessionId: string) => void; + onNavigateToSession?: (targetSessionId: string, options?: SessionNavigationOptions) => void; onShowSettings?: () => void; autoExpandTools?: boolean; showRawParameters?: boolean; diff --git a/src/components/main-content/types/types.ts b/src/components/main-content/types/types.ts index 17a7c9ec..d090852d 100644 --- a/src/components/main-content/types/types.ts +++ b/src/components/main-content/types/types.ts @@ -1,5 +1,7 @@ import type { Dispatch, SetStateAction } from 'react'; + import type { AppTab, Project, ProjectSession } from '../../../types/app'; +import type { SessionNavigationOptions } from '../../chat/types/types'; export type SessionLifecycleHandler = (sessionId?: string | null) => void; @@ -50,7 +52,7 @@ export type MainContentProps = { onSessionNotProcessing: SessionLifecycleHandler; processingSessions: Set; onReplaceTemporarySession: SessionLifecycleHandler; - onNavigateToSession: (targetSessionId: string) => void; + onNavigateToSession: (targetSessionId: string, options?: SessionNavigationOptions) => void; onShowSettings: () => void; externalMessageUpdate: number; newSessionTrigger: number; diff --git a/src/hooks/useProjectsState.ts b/src/hooks/useProjectsState.ts index b5cba267..d920fba2 100644 --- a/src/hooks/useProjectsState.ts +++ b/src/hooks/useProjectsState.ts @@ -5,6 +5,7 @@ import { api } from '../utils/api'; import type { AppSocketMessage, AppTab, + LLMProvider, LoadingProgress, Project, ProjectSession, @@ -557,7 +558,42 @@ export function useProjectsState({ return; } } - }, [sessionId, projects, selectedProject?.projectId, selectedSession?.id, selectedSession?.__provider]); + + // Session id is in the URL but not yet present on any project payload (common + // right after `session_created` + navigate, before the next projects refresh). + // Without a `selectedSession`, chat state clears `currentSessionId` and the + // UI stops reading the session store even though messages stream under this id. + if (selectedSession?.id === sessionId) { + return; + } + + if (!selectedProject) { + return; + } + + let providerFromStorage: string | null = null; + try { + providerFromStorage = localStorage.getItem('selected-provider'); + } catch { + providerFromStorage = null; + } + + const normalizedProvider: LLMProvider = + providerFromStorage === 'cursor' + ? 'cursor' + : providerFromStorage === 'codex' + ? 'codex' + : providerFromStorage === 'gemini' + ? 'gemini' + : 'claude'; + + setSelectedSession({ + id: sessionId, + __provider: normalizedProvider, + __projectId: selectedProject.projectId, + summary: '', + }); + }, [sessionId, projects, selectedProject, selectedSession?.id, selectedSession?.__provider]); const handleProjectSelect = useCallback( (project: Project) => { diff --git a/src/stores/useSessionStore.ts b/src/stores/useSessionStore.ts index ef581e12..86925048 100644 --- a/src/stores/useSessionStore.ts +++ b/src/stores/useSessionStore.ts @@ -104,17 +104,126 @@ function createEmptySlot(): SessionSlot { } /** - * Compute merged messages: server + realtime, deduped by id. - * Server messages take priority (they're the persisted source of truth). - * Realtime messages that aren't yet in server stay (in-flight streaming). + * Compute merged messages: server + realtime, deduped by id and adjacent + * assistant echo (same trimmed text), so finalized stream rows do not stack + * on top of the persisted copy before realtime is cleared. */ +function userTextFingerprint(m: NormalizedMessage): string | null { + if (m.kind !== 'text' || m.role !== 'user') return null; + const t = (m.content || '').trim(); + return t.length > 0 ? t : null; +} + +/** + * After `finalizeStreaming`, the client holds a synthetic assistant `text` row + * while the sessions API soon returns the same reply with a different id. + * Those sit back-to-back in merged order and look like duplicate bubbles until + * `refreshFromServer` clears realtime. Collapse same-text assistant rows and + * stream_placeholder → text when content matches. + */ +function dedupeAdjacentAssistantEchoes(merged: NormalizedMessage[]): NormalizedMessage[] { + const out: NormalizedMessage[] = []; + for (const m of merged) { + const prev = out[out.length - 1]; + if (prev) { + if (prev.kind === 'stream_delta' && m.kind === 'text' && m.role === 'assistant') { + const ps = (prev.content || '').trim(); + const ms = (m.content || '').trim(); + if (ps.length > 0 && ps === ms) { + out[out.length - 1] = m; + continue; + } + } + if ( + prev.kind === 'text' + && m.kind === 'text' + && prev.role === 'assistant' + && m.role === 'assistant' + ) { + const ms = (m.content || '').trim(); + if (ms.length > 0 && ms === (prev.content || '').trim()) { + continue; + } + } + } + out.push(m); + } + return out; +} + function computeMerged(server: NormalizedMessage[], realtime: NormalizedMessage[]): NormalizedMessage[] { if (realtime.length === 0) return server; - if (server.length === 0) return realtime; + if (server.length === 0) return dedupeAdjacentAssistantEchoes(realtime); const serverIds = new Set(server.map(m => m.id)); - const extra = realtime.filter(m => !serverIds.has(m.id)); + const serverUserTexts = new Set( + server.map(userTextFingerprint).filter((t): t is string => t !== null), + ); + const extra = realtime.filter((m) => { + if (serverIds.has(m.id)) return false; + // Optimistic user rows use `local_*` ids; once the same text exists on the + // server-backed copy, drop the realtime echo to avoid duplicate bubbles. + if (m.id.startsWith('local_')) { + const fp = userTextFingerprint(m); + if (fp && serverUserTexts.has(fp)) return false; + } + return true; + }); if (extra.length === 0) return server; - return [...server, ...extra]; + return dedupeAdjacentAssistantEchoes([...server, ...extra]); +} + +function compareMessagesByTimestamp(left: NormalizedMessage, right: NormalizedMessage): number { + const leftTime = Date.parse(left.timestamp); + const rightTime = Date.parse(right.timestamp); + + if (Number.isNaN(leftTime) || Number.isNaN(rightTime) || leftTime === rightTime) { + return 0; + } + + return leftTime - rightTime; +} + +function rewriteMessageSessionId( + msg: NormalizedMessage, + fromSessionId: string, + toSessionId: string, +): NormalizedMessage { + const streamingSourceId = `__streaming_${fromSessionId}`; + const nextId = msg.id === streamingSourceId ? `__streaming_${toSessionId}` : msg.id; + + if (msg.sessionId === toSessionId && nextId === msg.id) { + return msg; + } + + return { + ...msg, + id: nextId, + sessionId: toSessionId, + }; +} + +function mergeMessagesById( + existing: NormalizedMessage[], + incoming: NormalizedMessage[], +): NormalizedMessage[] { + if (existing.length === 0) return incoming; + if (incoming.length === 0) return existing; + + const merged = [...existing, ...incoming]; + const deduped: NormalizedMessage[] = []; + const seen = new Set(); + + for (const msg of merged) { + if (seen.has(msg.id)) { + continue; + } + + seen.add(msg.id); + deduped.push(msg); + } + + deduped.sort(compareMessagesByTimestamp); + return deduped; } /** @@ -141,28 +250,59 @@ const MAX_REALTIME_MESSAGES = 500; export function useSessionStore() { const storeRef = useRef(new Map()); + const sessionAliasesRef = useRef(new Map()); const activeSessionIdRef = useRef(null); // Bump to force re-render — only when the active session's data changes const [, setTick] = useState(0); const notify = useCallback((sessionId: string) => { - if (sessionId === activeSessionIdRef.current) { + const aliases = sessionAliasesRef.current; + let resolvedSessionId = sessionId; + const visited = new Set(); + + while (aliases.has(resolvedSessionId) && !visited.has(resolvedSessionId)) { + visited.add(resolvedSessionId); + resolvedSessionId = aliases.get(resolvedSessionId)!; + } + + if (resolvedSessionId === activeSessionIdRef.current) { setTick(n => n + 1); } }, []); - const setActiveSession = useCallback((sessionId: string | null) => { - activeSessionIdRef.current = sessionId; + const resolveSessionId = useCallback((sessionId: string | null | undefined): string | null => { + if (!sessionId) { + return null; + } + + const aliases = sessionAliasesRef.current; + let resolvedSessionId = sessionId; + const visited = new Set(); + + while (aliases.has(resolvedSessionId) && !visited.has(resolvedSessionId)) { + visited.add(resolvedSessionId); + resolvedSessionId = aliases.get(resolvedSessionId)!; + } + + return resolvedSessionId; }, []); + const setActiveSession = useCallback((sessionId: string | null) => { + activeSessionIdRef.current = resolveSessionId(sessionId); + }, [resolveSessionId]); + const getSlot = useCallback((sessionId: string): SessionSlot => { + const resolvedSessionId = resolveSessionId(sessionId) ?? sessionId; const store = storeRef.current; - if (!store.has(sessionId)) { - store.set(sessionId, createEmptySlot()); + if (!store.has(resolvedSessionId)) { + store.set(resolvedSessionId, createEmptySlot()); } - return store.get(sessionId)!; - }, []); + return store.get(resolvedSessionId)!; + }, [resolveSessionId]); - const has = useCallback((sessionId: string) => storeRef.current.has(sessionId), []); + const has = useCallback((sessionId: string) => { + const resolvedSessionId = resolveSessionId(sessionId) ?? sessionId; + return storeRef.current.has(resolvedSessionId); + }, [resolveSessionId]); /** * Fetch messages from the provider sessions endpoint and populate serverMessages. @@ -179,9 +319,10 @@ export function useSessionStore() { offset?: number; } = {}, ) => { - const slot = getSlot(sessionId); + const resolvedSessionId = resolveSessionId(sessionId) ?? sessionId; + const slot = getSlot(resolvedSessionId); slot.status = 'loading'; - notify(sessionId); + notify(resolvedSessionId); try { const params = new URLSearchParams(); @@ -191,7 +332,7 @@ export function useSessionStore() { } const qs = params.toString(); - const url = `/api/providers/sessions/${encodeURIComponent(sessionId)}/messages${qs ? `?${qs}` : ''}`; + const url = `/api/providers/sessions/${encodeURIComponent(resolvedSessionId)}/messages${qs ? `?${qs}` : ''}`; const response = await authenticatedFetch(url); if (!response.ok) { @@ -212,15 +353,15 @@ export function useSessionStore() { slot.tokenUsage = data.tokenUsage; } - notify(sessionId); + notify(resolvedSessionId); return slot; } catch (error) { - console.error(`[SessionStore] fetch failed for ${sessionId}:`, error); + console.error(`[SessionStore] fetch failed for ${resolvedSessionId}:`, error); slot.status = 'error'; - notify(sessionId); + notify(resolvedSessionId); return slot; } - }, [getSlot, notify]); + }, [getSlot, notify, resolveSessionId]); /** * Load older (paginated) messages and prepend to serverMessages. @@ -234,7 +375,8 @@ export function useSessionStore() { limit?: number; } = {}, ) => { - const slot = getSlot(sessionId); + const resolvedSessionId = resolveSessionId(sessionId) ?? sessionId; + const slot = getSlot(resolvedSessionId); if (!slot.hasMore) return slot; const params = new URLSearchParams(); @@ -243,7 +385,7 @@ export function useSessionStore() { params.append('offset', String(slot.offset)); const qs = params.toString(); - const url = `/api/providers/sessions/${encodeURIComponent(sessionId)}/messages${qs ? `?${qs}` : ''}`; + const url = `/api/providers/sessions/${encodeURIComponent(resolvedSessionId)}/messages${qs ? `?${qs}` : ''}`; try { const response = await authenticatedFetch(url); @@ -256,43 +398,54 @@ export function useSessionStore() { slot.hasMore = Boolean(data.hasMore); slot.offset = slot.offset + olderMessages.length; recomputeMergedIfNeeded(slot); - notify(sessionId); + notify(resolvedSessionId); return slot; } catch (error) { - console.error(`[SessionStore] fetchMore failed for ${sessionId}:`, error); + console.error(`[SessionStore] fetchMore failed for ${resolvedSessionId}:`, error); return slot; } - }, [getSlot, notify]); + }, [getSlot, notify, resolveSessionId]); /** * Append a realtime (WebSocket) message to the correct session slot. * This works regardless of which session is actively viewed. */ const appendRealtime = useCallback((sessionId: string, msg: NormalizedMessage) => { - const slot = getSlot(sessionId); - let updated = [...slot.realtimeMessages, msg]; + const resolvedSessionId = resolveSessionId(sessionId) ?? sessionId; + const slot = getSlot(resolvedSessionId); + const normalizedMessage = + msg.sessionId === resolvedSessionId + ? msg + : { ...msg, sessionId: resolvedSessionId }; + let updated = [...slot.realtimeMessages, normalizedMessage]; if (updated.length > MAX_REALTIME_MESSAGES) { updated = updated.slice(-MAX_REALTIME_MESSAGES); } slot.realtimeMessages = updated; recomputeMergedIfNeeded(slot); - notify(sessionId); - }, [getSlot, notify]); + notify(resolvedSessionId); + }, [getSlot, notify, resolveSessionId]); /** * Append multiple realtime messages at once (batch). */ const appendRealtimeBatch = useCallback((sessionId: string, msgs: NormalizedMessage[]) => { if (msgs.length === 0) return; - const slot = getSlot(sessionId); - let updated = [...slot.realtimeMessages, ...msgs]; + const resolvedSessionId = resolveSessionId(sessionId) ?? sessionId; + const slot = getSlot(resolvedSessionId); + const normalizedMessages = msgs.map((msg) => + msg.sessionId === resolvedSessionId + ? msg + : { ...msg, sessionId: resolvedSessionId }, + ); + let updated = [...slot.realtimeMessages, ...normalizedMessages]; if (updated.length > MAX_REALTIME_MESSAGES) { updated = updated.slice(-MAX_REALTIME_MESSAGES); } slot.realtimeMessages = updated; recomputeMergedIfNeeded(slot); - notify(sessionId); - }, [getSlot, notify]); + notify(resolvedSessionId); + }, [getSlot, notify, resolveSessionId]); /** * Re-fetch serverMessages from the provider sessions endpoint. @@ -305,12 +458,13 @@ export function useSessionStore() { projectPath?: string; } = {}, ) => { - const slot = getSlot(sessionId); + const resolvedSessionId = resolveSessionId(sessionId) ?? sessionId; + const slot = getSlot(resolvedSessionId); try { const params = new URLSearchParams(); const qs = params.toString(); - const url = `/api/providers/sessions/${encodeURIComponent(sessionId)}/messages${qs ? `?${qs}` : ''}`; + const url = `/api/providers/sessions/${encodeURIComponent(resolvedSessionId)}/messages${qs ? `?${qs}` : ''}`; const response = await authenticatedFetch(url); if (!response.ok) throw new Error(`HTTP ${response.status}`); @@ -323,40 +477,43 @@ export function useSessionStore() { // drop realtime messages that the server has caught up with to prevent unbounded growth. slot.realtimeMessages = []; recomputeMergedIfNeeded(slot); - notify(sessionId); + notify(resolvedSessionId); } catch (error) { - console.error(`[SessionStore] refresh failed for ${sessionId}:`, error); + console.error(`[SessionStore] refresh failed for ${resolvedSessionId}:`, error); } - }, [getSlot, notify]); + }, [getSlot, notify, resolveSessionId]); /** * Update session status. */ const setStatus = useCallback((sessionId: string, status: SessionStatus) => { - const slot = getSlot(sessionId); + const resolvedSessionId = resolveSessionId(sessionId) ?? sessionId; + const slot = getSlot(resolvedSessionId); slot.status = status; - notify(sessionId); - }, [getSlot, notify]); + notify(resolvedSessionId); + }, [getSlot, notify, resolveSessionId]); /** * Check if a session's data is stale (>30s old). */ const isStale = useCallback((sessionId: string) => { - const slot = storeRef.current.get(sessionId); + const resolvedSessionId = resolveSessionId(sessionId) ?? sessionId; + const slot = storeRef.current.get(resolvedSessionId); if (!slot) return true; return Date.now() - slot.fetchedAt > STALE_THRESHOLD_MS; - }, []); + }, [resolveSessionId]); /** * Update or create a streaming message (accumulated text so far). * Uses a well-known ID so subsequent calls replace the same message. */ const updateStreaming = useCallback((sessionId: string, accumulatedText: string, msgProvider: LLMProvider) => { - const slot = getSlot(sessionId); - const streamId = `__streaming_${sessionId}`; + const resolvedSessionId = resolveSessionId(sessionId) ?? sessionId; + const slot = getSlot(resolvedSessionId); + const streamId = `__streaming_${resolvedSessionId}`; const msg: NormalizedMessage = { id: streamId, - sessionId, + sessionId: resolvedSessionId, timestamp: new Date().toISOString(), provider: msgProvider, kind: 'stream_delta', @@ -370,17 +527,18 @@ export function useSessionStore() { slot.realtimeMessages = [...slot.realtimeMessages, msg]; } recomputeMergedIfNeeded(slot); - notify(sessionId); - }, [getSlot, notify]); + notify(resolvedSessionId); + }, [getSlot, notify, resolveSessionId]); /** * Finalize streaming: convert the streaming message to a regular text message. * The well-known streaming ID is replaced with a unique text message ID. */ const finalizeStreaming = useCallback((sessionId: string) => { - const slot = storeRef.current.get(sessionId); + const resolvedSessionId = resolveSessionId(sessionId) ?? sessionId; + const slot = storeRef.current.get(resolvedSessionId); if (!slot) return; - const streamId = `__streaming_${sessionId}`; + const streamId = `__streaming_${resolvedSessionId}`; const idx = slot.realtimeMessages.findIndex(m => m.id === streamId); if (idx >= 0) { const stream = slot.realtimeMessages[idx]; @@ -392,35 +550,104 @@ export function useSessionStore() { role: 'assistant', }; recomputeMergedIfNeeded(slot); - notify(sessionId); + notify(resolvedSessionId); } - }, [notify]); + }, [notify, resolveSessionId]); /** * Clear realtime messages for a session (e.g., after stream completes and server fetch catches up). */ const clearRealtime = useCallback((sessionId: string) => { - const slot = storeRef.current.get(sessionId); + const resolvedSessionId = resolveSessionId(sessionId) ?? sessionId; + const slot = storeRef.current.get(resolvedSessionId); if (slot) { slot.realtimeMessages = []; recomputeMergedIfNeeded(slot); - notify(sessionId); + notify(resolvedSessionId); } - }, [notify]); + }, [notify, resolveSessionId]); /** * Get merged messages for a session (for rendering). */ const getMessages = useCallback((sessionId: string): NormalizedMessage[] => { - return storeRef.current.get(sessionId)?.merged ?? []; - }, []); + const resolvedSessionId = resolveSessionId(sessionId) ?? sessionId; + return storeRef.current.get(resolvedSessionId)?.merged ?? []; + }, [resolveSessionId]); /** * Get session slot (for status, pagination info, etc.). */ const getSessionSlot = useCallback((sessionId: string): SessionSlot | undefined => { - return storeRef.current.get(sessionId); - }, []); + const resolvedSessionId = resolveSessionId(sessionId) ?? sessionId; + return storeRef.current.get(resolvedSessionId); + }, [resolveSessionId]); + + const replaceSessionId = useCallback((fromSessionId: string, toSessionId: string) => { + const resolvedFromSessionId = resolveSessionId(fromSessionId) ?? fromSessionId; + const resolvedToSessionId = resolveSessionId(toSessionId) ?? toSessionId; + + if (resolvedFromSessionId === resolvedToSessionId) { + sessionAliasesRef.current.set(fromSessionId, resolvedToSessionId); + return; + } + + const store = storeRef.current; + const sourceSlot = store.get(resolvedFromSessionId); + const targetSlot = store.get(resolvedToSessionId) ?? createEmptySlot(); + + if (sourceSlot) { + const migratedServerMessages = sourceSlot.serverMessages.map((msg) => + rewriteMessageSessionId(msg, resolvedFromSessionId, resolvedToSessionId), + ); + const migratedRealtimeMessages = sourceSlot.realtimeMessages.map((msg) => + rewriteMessageSessionId(msg, resolvedFromSessionId, resolvedToSessionId), + ); + + targetSlot.serverMessages = mergeMessagesById(targetSlot.serverMessages, migratedServerMessages); + targetSlot.realtimeMessages = mergeMessagesById(targetSlot.realtimeMessages, migratedRealtimeMessages); + if (targetSlot.realtimeMessages.length > MAX_REALTIME_MESSAGES) { + targetSlot.realtimeMessages = targetSlot.realtimeMessages.slice(-MAX_REALTIME_MESSAGES); + } + targetSlot.status = + sourceSlot.status === 'error' + ? 'error' + : sourceSlot.status === 'streaming' || targetSlot.status === 'streaming' + ? 'streaming' + : sourceSlot.status === 'loading' || targetSlot.status === 'loading' + ? 'loading' + : targetSlot.status; + targetSlot.fetchedAt = Math.max(targetSlot.fetchedAt, sourceSlot.fetchedAt, Date.now()); + targetSlot.total = Math.max( + targetSlot.total, + sourceSlot.total, + targetSlot.serverMessages.length, + targetSlot.realtimeMessages.length, + ); + targetSlot.hasMore = targetSlot.hasMore || sourceSlot.hasMore; + targetSlot.offset = Math.max(targetSlot.offset, sourceSlot.offset); + targetSlot.tokenUsage = targetSlot.tokenUsage ?? sourceSlot.tokenUsage; + recomputeMergedIfNeeded(targetSlot); + + store.set(resolvedToSessionId, targetSlot); + store.delete(resolvedFromSessionId); + } + + sessionAliasesRef.current.set(resolvedFromSessionId, resolvedToSessionId); + sessionAliasesRef.current.set(fromSessionId, resolvedToSessionId); + + for (const [aliasSessionId, targetSessionId] of sessionAliasesRef.current.entries()) { + if (targetSessionId === resolvedFromSessionId) { + sessionAliasesRef.current.set(aliasSessionId, resolvedToSessionId); + } + } + + if (activeSessionIdRef.current === resolvedFromSessionId) { + activeSessionIdRef.current = resolvedToSessionId; + } + + notify(resolvedToSessionId); + }, [notify, resolveSessionId]); return useMemo(() => ({ getSlot, @@ -438,11 +665,12 @@ export function useSessionStore() { clearRealtime, getMessages, getSessionSlot, + replaceSessionId, }), [ getSlot, has, fetchFromServer, fetchMore, appendRealtime, appendRealtimeBatch, refreshFromServer, setActiveSession, setStatus, isStale, updateStreaming, finalizeStreaming, - clearRealtime, getMessages, getSessionSlot, + clearRealtime, getMessages, getSessionSlot, replaceSessionId, ]); }