Fix session state refresh on navigation (#115)

This commit is contained in:
João Salavisa
2026-01-28 17:23:53 +00:00
committed by GitHub
parent 1888f8465f
commit 2ebeaaab37
5 changed files with 33 additions and 22 deletions
+1 -1
View File
@@ -154,7 +154,7 @@ export default async function SessionPage({
return (
<>
{initialState ? <StateHydrator initialState={initialState} /> : null}
{initialState ? <StateHydrator initialState={initialState} sessionId={sessionId ?? undefined} /> : null}
<TaskPageContent
task={task}
sessionId={sessionId}
+7 -3
View File
@@ -6,9 +6,11 @@ import { useAppStoreApi } from '@/components/state-provider';
type StateHydratorProps = {
initialState: Partial<AppState>;
/** 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;
}
+12 -10
View File
@@ -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<AppState>,
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);
}
@@ -37,20 +37,25 @@ export function deepMerge<T extends Record<string, any>>(target: Draft<T>, 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<T>(
target: Draft<Record<string, T>>,
source: Record<string, T> | 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];
}
+3 -3
View File
@@ -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<AppState>) => void;
hydrate: (state: Partial<AppState>, 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<AppState>) {
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)),
}))
);
}