diff --git a/apps/web/app/s/[sessionId]/page.tsx b/apps/web/app/s/[sessionId]/page.tsx index 862db6ab9f..ffb8ce9dec 100644 --- a/apps/web/app/s/[sessionId]/page.tsx +++ b/apps/web/app/s/[sessionId]/page.tsx @@ -154,7 +154,7 @@ export default async function SessionPage({ return ( <> - {initialState ? : null} + {initialState ? : null} ; + /** Session ID to force-merge even if it's active (for navigation refresh) */ + sessionId?: string; }; -export function StateHydrator({ initialState }: StateHydratorProps) { +export function StateHydrator({ initialState, sessionId }: StateHydratorProps) { const store = useAppStoreApi(); // Use useLayoutEffect to hydrate state synchronously before child effects run. @@ -16,9 +18,11 @@ export function StateHydrator({ initialState }: StateHydratorProps) { // decide whether to fetch data. useLayoutEffect(() => { if (Object.keys(initialState).length) { - store.getState().hydrate(initialState); + store.getState().hydrate(initialState, { + forceMergeSessionId: sessionId, + }); } - }, [initialState, store]); + }, [initialState, sessionId, store]); return null; } diff --git a/apps/web/lib/state/hydration/hydrator.ts b/apps/web/lib/state/hydration/hydrator.ts index 573d49b486..2438fa7385 100644 --- a/apps/web/lib/state/hydration/hydrator.ts +++ b/apps/web/lib/state/hydration/hydrator.ts @@ -10,6 +10,8 @@ export type HydrationOptions = { activeSessionId?: string | null; /** Whether to skip hydrating session runtime state (shell, processes, git) */ skipSessionRuntime?: boolean; + /** Force merge this session even if it's active (for navigation refresh) */ + forceMergeSessionId?: string | null; }; /** @@ -26,7 +28,7 @@ export function hydrateState( state: Partial, options: HydrationOptions = {} ): void { - const { activeSessionId = null, skipSessionRuntime = false } = options; + const { activeSessionId = null, skipSessionRuntime = false, forceMergeSessionId = null } = options; // Kanban slice - always safe to hydrate if (state.kanban) deepMerge(draft.kanban, state.kanban); @@ -66,24 +68,24 @@ export function hydrateState( // Session slice - careful with active sessions if (state.messages) { if (state.messages.bySession) { - mergeSessionMap(draft.messages.bySession, state.messages.bySession, activeSessionId); + mergeSessionMap(draft.messages.bySession, state.messages.bySession, activeSessionId, forceMergeSessionId); } if (state.messages.metaBySession) { - mergeSessionMap(draft.messages.metaBySession, state.messages.metaBySession, activeSessionId); + mergeSessionMap(draft.messages.metaBySession, state.messages.metaBySession, activeSessionId, forceMergeSessionId); } } if (state.turns) { if (state.turns.bySession) { - mergeSessionMap(draft.turns.bySession, state.turns.bySession, activeSessionId); + mergeSessionMap(draft.turns.bySession, state.turns.bySession, activeSessionId, forceMergeSessionId); } if (state.turns.activeBySession) { - mergeSessionMap(draft.turns.activeBySession, state.turns.activeBySession, activeSessionId); + mergeSessionMap(draft.turns.activeBySession, state.turns.activeBySession, activeSessionId, forceMergeSessionId); } } if (state.taskSessions) deepMerge(draft.taskSessions, state.taskSessions); if (state.taskSessionsByTask) deepMerge(draft.taskSessionsByTask, state.taskSessionsByTask); if (state.sessionAgentctl) { - mergeSessionMap(draft.sessionAgentctl.itemsBySessionId, state.sessionAgentctl?.itemsBySessionId, activeSessionId); + mergeSessionMap(draft.sessionAgentctl.itemsBySessionId, state.sessionAgentctl?.itemsBySessionId, activeSessionId, forceMergeSessionId); } if (state.worktrees) deepMerge(draft.worktrees, state.worktrees); if (state.sessionWorktreesBySessionId) { @@ -96,15 +98,15 @@ export function hydrateState( if (!skipSessionRuntime) { if (state.terminal) deepMerge(draft.terminal, state.terminal); if (state.shell) { - mergeSessionMap(draft.shell.outputs, state.shell?.outputs, activeSessionId); - mergeSessionMap(draft.shell.statuses, state.shell?.statuses, activeSessionId); + mergeSessionMap(draft.shell.outputs, state.shell?.outputs, activeSessionId, forceMergeSessionId); + mergeSessionMap(draft.shell.statuses, state.shell?.statuses, activeSessionId, forceMergeSessionId); } if (state.processes) deepMerge(draft.processes, state.processes); if (state.gitStatus) { - mergeSessionMap(draft.gitStatus.bySessionId, state.gitStatus?.bySessionId, activeSessionId); + mergeSessionMap(draft.gitStatus.bySessionId, state.gitStatus?.bySessionId, activeSessionId, forceMergeSessionId); } if (state.contextWindow) { - mergeSessionMap(draft.contextWindow.bySessionId, state.contextWindow?.bySessionId, activeSessionId); + mergeSessionMap(draft.contextWindow.bySessionId, state.contextWindow?.bySessionId, activeSessionId, forceMergeSessionId); } if (state.agents) deepMerge(draft.agents, state.agents); } diff --git a/apps/web/lib/state/hydration/merge-strategies.ts b/apps/web/lib/state/hydration/merge-strategies.ts index d662c416a6..0ff2308434 100644 --- a/apps/web/lib/state/hydration/merge-strategies.ts +++ b/apps/web/lib/state/hydration/merge-strategies.ts @@ -37,20 +37,25 @@ export function deepMerge>(target: Draft, sourc /** * Merge strategy for sessionId-keyed maps * Only merges sessions that don't exist or are not currently active + * @param forceMergeSessionId - If provided, this session will be merged even if it's active */ export function mergeSessionMap( target: Draft>, source: Record | undefined, - activeSessionId: string | null + activeSessionId: string | null, + forceMergeSessionId?: string | null ): void { if (!source) return; for (const sessionId in source) { - // Skip if this is the active session to avoid overwriting live data - if (sessionId === activeSessionId) continue; + // Force merge if this is the forceMergeSessionId (for navigation refresh) + const shouldForceMerge = forceMergeSessionId && sessionId === forceMergeSessionId; - // Only set if not present - if (!(sessionId in target)) { + // Skip if this is the active session to avoid overwriting live data (unless force merge) + if (!shouldForceMerge && sessionId === activeSessionId) continue; + + // Merge the session data + if (shouldForceMerge || !(sessionId in target)) { // eslint-disable-next-line @typescript-eslint/no-explicit-any (target as any)[sessionId] = source[sessionId]; } diff --git a/apps/web/lib/state/store.ts b/apps/web/lib/state/store.ts index 1a187ab998..a05bbb0d94 100644 --- a/apps/web/lib/state/store.ts +++ b/apps/web/lib/state/store.ts @@ -1,6 +1,6 @@ import { createStore } from 'zustand/vanilla'; import { immer } from 'zustand/middleware/immer'; -import { hydrateState } from './hydration/hydrator'; +import { hydrateState, type HydrationOptions } from './hydration/hydrator'; import type { Repository, Branch, Message, Turn, TaskSession } from '@/lib/types/http'; import { createKanbanSlice, @@ -144,7 +144,7 @@ export type AppState = { connection: typeof defaultUIState['connection']; // Actions from all slices - hydrate: (state: Partial) => void; + hydrate: (state: Partial, options?: HydrationOptions) => void; setActiveWorkspace: (workspaceId: string | null) => void; setWorkspaces: (workspaces: WorkspaceState['items']) => void; setActiveBoard: (boardId: string | null) => void; @@ -382,7 +382,7 @@ export function createAppStore(initialState?: Partial) { connection: merged.connection, // Add hydrate method // eslint-disable-next-line @typescript-eslint/no-explicit-any - hydrate: (state) => set((draft) => hydrateState(draft as any, state)), + hydrate: (state, options) => set((draft) => hydrateState(draft as any, state, options)), })) ); }