fix: per-session state management for smooth task switching

Background sessions now maintain their own state snapshots in agentStore
so streaming, tool calls, and panel updates are not lost when switching
tasks or backgrounding the browser tab.

- Add sessionStates map + updateSession/switchActiveSession/syncSnapshot
- Side-channel callbacks always update per-session state (no isActiveRef guard)
- Global setters (setPanel, setPanelView, etc.) sync to active snapshot
- Re-sync state on browser tab visibility change
- Fix direct state mutation in switchActiveSession

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

frontend/src/components/Layout/AppLayout.tsx

@@ -123,6 +123,7 @@ export default function AppLayout() {
 
   const handleSessionDead = useCallback(
     (deadSessionId: string) => {
+      useAgentStore.getState().clearSessionState(deadSessionId);
       deleteSession(deadSessionId);
     },
     [deleteSession],

frontend/src/components/SessionChat.tsx

@@ -5,11 +5,10 @@
  * runs — processing events — but only the active session renders visible
  * UI (MessageList + ChatInput).
  */
-import { useCallback, useEffect, useRef } from 'react';
+import { useCallback, useEffect } from 'react';
 import { useAgentChat } from '@/hooks/useAgentChat';
 import { useAgentStore } from '@/store/agentStore';
 import { useSessionStore } from '@/store/sessionStore';
-import { useLayoutStore } from '@/store/layoutStore';
 import MessageList from '@/components/Chat/MessageList';
 import ChatInput from '@/components/Chat/ChatInput';
 import { apiFetch } from '@/utils/api';
@@ -22,7 +21,7 @@ interface SessionChatProps {
 }
 
 export default function SessionChat({ sessionId, isActive, onSessionDead }: SessionChatProps) {
-  const { isConnected, isProcessing, setProcessing, activityStatus } = useAgentStore();
+  const { isConnected, isProcessing, activityStatus, updateSession } = useAgentStore();
   const { updateSessionTitle } = useSessionStore();
 
   const { messages, sendMessage, stop, status, undoLastTurn, approveTools } = useAgentChat({
@@ -33,85 +32,28 @@ export default function SessionChat({ sessionId, isActive, onSessionDead }: Sess
     onSessionDead,
   });
 
-  // When this session becomes active, sync ALL global agentStore state
-  // so the UI correctly reflects this session's current state.
-  // When it becomes inactive, clear global state so the next session starts clean.
-  const prevActiveRef = useRef(isActive);
+  // When this session becomes active, restore its per-session state to the
+  // global flat fields. The per-session state map is kept up-to-date by
+  // side-channel callbacks even while the session is in the background.
   useEffect(() => {
-    if (isActive && !prevActiveRef.current) {
-      // ── Becoming active: restore this session's state ──
-      const store = useAgentStore.getState();
-
-      // SSE transport has no persistent connection — always connected
-      store.setConnected(true);
-
-      // Check if this session has pending approvals in its messages
-      const lastAssistant = [...messages].reverse().find(m => m.role === 'assistant');
-      const hasPendingApproval = lastAssistant?.parts.some(
-        (p) => p.type === 'dynamic-tool' && p.state === 'approval-requested'
-      ) ?? false;
-      const hasApprovedRunning = lastAssistant?.parts.some(
-        (p) => p.type === 'dynamic-tool' && p.state === 'approval-responded'
-      ) ?? false;
-
-      if (hasPendingApproval) {
-        store.setActivityStatus({ type: 'waiting-approval' });
-        store.setProcessing(false);
+    if (isActive) {
+      useAgentStore.getState().switchActiveSession(sessionId);
+      useAgentStore.getState().setConnected(true);
+    }
+  }, [isActive, sessionId]);
 
-        // Restore panel for the first pending tool
-        const pendingTool = lastAssistant!.parts.find(
-          (p) => p.type === 'dynamic-tool' && p.state === 'approval-requested'
-        );
-        if (pendingTool && pendingTool.type === 'dynamic-tool') {
-          const args = pendingTool.input as Record<string, string | undefined>;
-          if (pendingTool.toolName === 'hf_jobs' && args?.script) {
-            store.setPanel(
-              { title: 'Script', script: { content: args.script, language: 'python' }, parameters: pendingTool.input as Record<string, unknown> },
-              'script',
-              true,
-            );
-          } else if (pendingTool.toolName === 'hf_repo_files' && args?.content) {
-            const filename = args.path || 'file';
-            store.setPanel({
-              title: filename.split('/').pop() || 'Content',
-              script: { content: args.content, language: filename.endsWith('.py') ? 'python' : 'text' },
-              parameters: pendingTool.input as Record<string, unknown>,
-            });
-          } else {
-            store.setPanel({
-              title: pendingTool.toolName,
-              output: { content: JSON.stringify(pendingTool.input, null, 2), language: 'json' },
-            }, 'output');
-          }
-          useLayoutStore.getState().setRightPanelOpen(true);
-        }
-      } else if (hasApprovedRunning) {
-        // Tools were approved but still executing — show processing state
-        store.setActivityStatus({ type: 'tool', toolName: 'running' });
-        store.setProcessing(true);
-      } else {
-        // Check if any tools are still running (non-approval tools like bash, read, etc.)
-        const runningTool = lastAssistant?.parts.find(
-          (p) => p.type === 'dynamic-tool' && (p.state === 'input-available' || p.state === 'input-streaming')
-        );
-        if (runningTool && runningTool.type === 'dynamic-tool') {
-          const desc = (runningTool.input as Record<string, unknown>)?.description as string | undefined;
-          store.setActivityStatus({ type: 'tool', toolName: runningTool.toolName, description: desc });
-          store.setProcessing(true);
-        } else {
-          store.setActivityStatus({ type: 'idle' });
-          store.setProcessing(false);
-        }
+  // Re-sync state when the browser tab regains focus (Chrome throttles
+  // timers in background tabs which can stall the AI SDK's update flushing).
+  useEffect(() => {
+    if (!isActive) return;
+    const onVisible = () => {
+      if (document.visibilityState === 'visible') {
+        useAgentStore.getState().switchActiveSession(sessionId);
       }
-    } else if (!isActive && prevActiveRef.current) {
-      // ── Becoming inactive: clear global state so the next session starts clean ──
-      const store = useAgentStore.getState();
-      store.setActivityStatus({ type: 'idle' });
-      store.setProcessing(false);
-      store.clearPanel();
-    }
-    prevActiveRef.current = isActive;
-  }, [isActive, messages]);
+    };
+    document.addEventListener('visibilitychange', onVisible);
+    return () => document.removeEventListener('visibilitychange', onVisible);
+  }, [isActive, sessionId]);
 
   // SDK status is the ground truth — if it's streaming/submitted, agent is busy
   const sdkBusy = status === 'streaming' || status === 'submitted';
@@ -121,7 +63,7 @@ export default function SessionChat({ sessionId, isActive, onSessionDead }: Sess
     async (text: string) => {
       if (!text.trim() || busy) return;
 
-      setProcessing(true);
+      updateSession(sessionId, { isProcessing: true });
       sendMessage({ text: text.trim(), metadata: { createdAt: new Date().toISOString() } });
 
       // Auto-title the session from the first user message
@@ -141,7 +83,7 @@ export default function SessionChat({ sessionId, isActive, onSessionDead }: Sess
           });
       }
     },
-    [sessionId, sendMessage, messages, updateSessionTitle, busy, setProcessing],
+    [sessionId, sendMessage, messages, updateSessionTitle, busy, updateSession],
   );
 
   // Don't render UI for background sessions — hooks still run

frontend/src/components/SessionSidebar/SessionSidebar.tsx

@@ -54,6 +54,7 @@ export default function SessionSidebar({ onClose }: SessionSidebarProps) {
   const handleDelete = useCallback(
     async (sessionId: string, e: React.MouseEvent) => {
       e.stopPropagation();
+      useAgentStore.getState().clearSessionState(sessionId);
       try {
         await apiFetch(`/api/session/${sessionId}`, { method: 'DELETE' });
         deleteSession(sessionId);
@@ -67,11 +68,11 @@ export default function SessionSidebar({ onClose }: SessionSidebarProps) {
   const handleSelect = useCallback(
     (sessionId: string) => {
       switchSession(sessionId);
-      setPlan([]);
-      clearPanel();
+      // Per-session state (plan, panel, activity) is restored automatically
+      // by SessionChat's useEffect when isActive flips to true.
       onClose?.();
     },
-    [switchSession, setPlan, clearPanel, onClose],
+    [switchSession, onClose],
   );
 
   const formatTime = (d: string) =>

frontend/src/hooks/useAgentChat.ts

@@ -3,9 +3,9 @@
  * ChatTransport.
  *
  * In the per-session architecture, each session mounts its own instance
- * of this hook. The `isActive` flag controls whether side-channel
- * callbacks update the global UI stores (agentStore / layoutStore) or
- * only per-session metadata (sessionStore.needsAttention).
+ * of this hook. Side-channel callbacks always update the session's own
+ * state via `updateSession()`. If the session is currently active, the
+ * store automatically mirrors updates to the flat global fields.
  */
 import { useCallback, useEffect, useMemo, useRef } from 'react';
 import { useChat } from '@ai-sdk/react';
@@ -34,87 +34,87 @@ export function useAgentChat({ sessionId, isActive, onReady, onError, onSessionD
   const isActiveRef = useRef(isActive);
   isActiveRef.current = isActive;
 
-  const {
-    setProcessing,
-    setConnected,
-    setActivityStatus,
-    setError,
-    setPanel,
-    setPanelOutput,
-  } = useAgentStore();
+  const { setNeedsAttention } = useSessionStore();
 
-  const { setRightPanelOpen, setLeftSidebarOpen } = useLayoutStore();
-  const { setSessionActive, setNeedsAttention } = useSessionStore();
+  // Helper: update this session's state (mirrors to globals if active)
+  const updateSession = useAgentStore.getState().updateSession;
 
   // -- Build side-channel callbacks (stable ref) --------------------------
   const sideChannel = useMemo<SideChannelCallbacks>(
     () => ({
       onReady: () => {
+        updateSession(sessionId, { isProcessing: false });
         if (isActiveRef.current) {
-          setConnected(true);
-          setProcessing(false);
+          useAgentStore.getState().setConnected(true);
         }
-        setSessionActive(sessionId, true);
+        useSessionStore.getState().setSessionActive(sessionId, true);
         callbacksRef.current.onReady?.();
       },
       onShutdown: () => {
+        updateSession(sessionId, { isProcessing: false });
         if (isActiveRef.current) {
-          setConnected(false);
-          setProcessing(false);
+          useAgentStore.getState().setConnected(false);
         }
       },
       onError: (error: string) => {
+        updateSession(sessionId, { isProcessing: false });
         if (isActiveRef.current) {
-          setError(error);
-          setProcessing(false);
+          useAgentStore.getState().setError(error);
         }
         callbacksRef.current.onError?.(error);
       },
       onProcessing: () => {
-        if (isActiveRef.current) {
-          setProcessing(true);
-          setActivityStatus({ type: 'thinking' });
-        }
+        updateSession(sessionId, {
+          isProcessing: true,
+          activityStatus: { type: 'thinking' },
+        });
       },
       onProcessingDone: () => {
-        if (isActiveRef.current) {
-          setProcessing(false);
-        }
+        updateSession(sessionId, { isProcessing: false });
       },
       onUndoComplete: () => {
-        // Undo is handled client-side in undoLastTurn(). With SSE, undo_complete
-        // events are discarded (no subscriber listening between turns).
-        if (isActiveRef.current) setProcessing(false);
+        updateSession(sessionId, { isProcessing: false });
       },
       onCompacted: (oldTokens: number, newTokens: number) => {
         logger.log(`Context compacted: ${oldTokens} -> ${newTokens} tokens`);
       },
       onPlanUpdate: (plan) => {
-        if (!isActiveRef.current) return;
-        useAgentStore.getState().setPlan(plan as Array<{ id: string; content: string; status: 'pending' | 'in_progress' | 'completed' }>);
-        if (!useLayoutStore.getState().isRightPanelOpen) {
-          setRightPanelOpen(true);
+        const typed = plan as Array<{ id: string; content: string; status: 'pending' | 'in_progress' | 'completed' }>;
+        updateSession(sessionId, { plan: typed });
+        if (isActiveRef.current && !useLayoutStore.getState().isRightPanelOpen) {
+          useLayoutStore.getState().setRightPanelOpen(true);
         }
       },
       onToolLog: (tool: string, log: string) => {
-        if (!isActiveRef.current) return;
-        if (tool === 'hf_jobs' || tool === 'sandbox') {
-          const state = useAgentStore.getState();
-          const existingOutput = state.panelData?.output?.content || '';
-          const header = tool === 'sandbox' ? '--- Sandbox creation ---' : '--- Job execution started ---';
-          const newContent = existingOutput
-            ? existingOutput + '\n' + log
-            : header + '\n' + log;
+        const STREAMABLE_TOOLS = new Set(['hf_jobs', 'sandbox', 'bash']);
+        if (!STREAMABLE_TOOLS.has(tool)) return;
 
-          setPanelOutput({ content: newContent, language: 'text' });
+        const sessState = useAgentStore.getState().getSessionState(sessionId);
+        const existingOutput = sessState.panelData?.output?.content || '';
 
-          if (!useLayoutStore.getState().isRightPanelOpen) {
-            setRightPanelOpen(true);
-          }
+        const newContent = existingOutput
+          ? existingOutput + '\n' + log
+          : log;
+
+        if (!sessState.panelData) {
+          const title = tool === 'bash' ? 'Sandbox' : tool === 'sandbox' ? 'Sandbox' : 'Job Output';
+          updateSession(sessionId, {
+            panelData: { title, output: { content: newContent, language: 'text' } },
+            panelView: 'output',
+          });
+        } else {
+          updateSession(sessionId, {
+            panelData: { ...sessState.panelData, output: { content: newContent, language: 'text' } },
+            panelView: 'output',
+          });
+        }
+
+        if (isActiveRef.current && !useLayoutStore.getState().isRightPanelOpen) {
+          useLayoutStore.getState().setRightPanelOpen(true);
         }
       },
       onConnectionChange: (connected: boolean) => {
-        if (isActiveRef.current) setConnected(connected);
+        if (isActiveRef.current) useAgentStore.getState().setConnected(connected);
       },
       onSessionDead: (deadSessionId: string) => {
         logger.warn(`Session ${deadSessionId} dead, removing`);
@@ -123,66 +123,105 @@ export function useAgentChat({ sessionId, isActive, onReady, onError, onSessionD
       onApprovalRequired: (tools) => {
         if (!tools.length) return;
         setNeedsAttention(sessionId, true);
-        if (!isActiveRef.current) return;
 
-        setActivityStatus({ type: 'waiting-approval' });
+        updateSession(sessionId, { activityStatus: { type: 'waiting-approval' } });
+
+        // Build panel data for this session's pending approval
         const firstTool = tools[0];
         const args = firstTool.arguments as Record<string, string | undefined>;
 
+        let panelUpdate: Partial<import('@/store/agentStore').PerSessionState> | undefined;
         if (firstTool.tool === 'hf_jobs' && args.script) {
-          setPanel(
-            { title: 'Script', script: { content: args.script, language: 'python' }, parameters: firstTool.arguments as Record<string, unknown> },
-            'script',
-            true,
-          );
+          panelUpdate = {
+            panelData: {
+              title: 'Script',
+              script: { content: args.script, language: 'python' },
+              parameters: firstTool.arguments as Record<string, unknown>,
+            },
+            panelView: 'script' as const,
+            panelEditable: true,
+          };
         } else if (firstTool.tool === 'hf_repo_files' && args.content) {
           const filename = args.path || 'file';
-          setPanel({
-            title: filename.split('/').pop() || 'Content',
-            script: { content: args.content, language: filename.endsWith('.py') ? 'python' : 'text' },
-            parameters: firstTool.arguments as Record<string, unknown>,
-          });
+          panelUpdate = {
+            panelData: {
+              title: filename.split('/').pop() || 'Content',
+              script: { content: args.content, language: filename.endsWith('.py') ? 'python' : 'text' },
+              parameters: firstTool.arguments as Record<string, unknown>,
+            },
+          };
         } else {
-          setPanel({
-            title: firstTool.tool,
-            output: { content: JSON.stringify(firstTool.arguments, null, 2), language: 'json' },
-          }, 'output');
+          panelUpdate = {
+            panelData: {
+              title: firstTool.tool,
+              output: { content: JSON.stringify(firstTool.arguments, null, 2), language: 'json' },
+            },
+            panelView: 'output' as const,
+          };
         }
+        if (panelUpdate) updateSession(sessionId, panelUpdate);
 
-        setRightPanelOpen(true);
-        setLeftSidebarOpen(false);
+        if (isActiveRef.current) {
+          useLayoutStore.getState().setRightPanelOpen(true);
+          useLayoutStore.getState().setLeftSidebarOpen(false);
+        }
       },
       onToolCallPanel: (toolName: string, args: Record<string, unknown>) => {
-        if (!isActiveRef.current) return;
         if (toolName === 'hf_jobs' && args.operation && args.script) {
-          setPanel(
-            { title: 'Script', script: { content: String(args.script), language: 'python' }, parameters: args },
-            'script',
-          );
-          setRightPanelOpen(true);
-          setLeftSidebarOpen(false);
+          updateSession(sessionId, {
+            panelData: {
+              title: 'Script',
+              script: { content: String(args.script), language: 'python' },
+              parameters: args,
+            },
+            panelView: 'script',
+          });
+          if (isActiveRef.current) {
+            useLayoutStore.getState().setRightPanelOpen(true);
+            useLayoutStore.getState().setLeftSidebarOpen(false);
+          }
         } else if (toolName === 'hf_repo_files' && args.operation === 'upload' && args.content) {
-          setPanel({
-            title: `File Upload: ${String(args.path || 'unnamed')}`,
-            script: { content: String(args.content), language: String(args.path || '').endsWith('.py') ? 'python' : 'text' },
-            parameters: args,
+          updateSession(sessionId, {
+            panelData: {
+              title: `File Upload: ${String(args.path || 'unnamed')}`,
+              script: { content: String(args.content), language: String(args.path || '').endsWith('.py') ? 'python' : 'text' },
+              parameters: args,
+            },
+          });
+          if (isActiveRef.current) {
+            useLayoutStore.getState().setRightPanelOpen(true);
+            useLayoutStore.getState().setLeftSidebarOpen(false);
+          }
+        } else if (toolName === 'bash' && args.command) {
+          updateSession(sessionId, {
+            panelData: {
+              title: 'Sandbox',
+              script: { content: String(args.command), language: 'bash' },
+            },
+            panelView: 'output',
           });
-          setRightPanelOpen(true);
-          setLeftSidebarOpen(false);
         }
       },
       onToolOutputPanel: (toolName: string, _toolCallId: string, output: string, success: boolean) => {
-        if (!isActiveRef.current) return;
+        const sessState = useAgentStore.getState().getSessionState(sessionId);
         if (toolName === 'hf_jobs' && output) {
-          setPanelOutput({ content: output, language: 'markdown' });
-          if (!success) useAgentStore.getState().setPanelView('output');
+          updateSession(sessionId, {
+            panelData: sessState.panelData
+              ? { ...sessState.panelData, output: { content: output, language: 'markdown' } }
+              : { title: 'Output', output: { content: output, language: 'markdown' } },
+            panelView: !success ? 'output' : sessState.panelView,
+          });
+        } else if (toolName === 'bash') {
+          if (!success) {
+            updateSession(sessionId, { panelView: 'output' });
+          }
         }
       },
       onStreaming: () => {
-        if (isActiveRef.current) setActivityStatus({ type: 'streaming' });
+        updateSession(sessionId, { activityStatus: { type: 'streaming' } });
       },
       onToolRunning: (toolName: string, description?: string) => {
-        if (isActiveRef.current) setActivityStatus({ type: 'tool', toolName, description });
+        updateSession(sessionId, { activityStatus: { type: 'tool', toolName, description } });
       },
     }),
     // eslint-disable-next-line react-hooks/exhaustive-deps
@@ -232,9 +271,9 @@ export function useAgentChat({ sessionId, isActive, onReady, onError, onSessionD
     sendAutomaticallyWhen: lastAssistantMessageIsCompleteWithApprovalResponses,
     onError: (error) => {
       logger.error('useChat error:', error);
+      updateSession(sessionId, { isProcessing: false });
       if (isActiveRef.current) {
-        setError(error.message);
-        setProcessing(false);
+        useAgentStore.getState().setError(error.message);
       }
     },
   });
@@ -316,11 +355,11 @@ export function useAgentChat({ sessionId, isActive, onReady, onError, onSessionD
         setMsgs(updated);
         saveMessages(sessionId, updated);
       }
-      if (isActiveRef.current) setProcessing(false);
+      updateSession(sessionId, { isProcessing: false });
     } catch (e) {
       logger.error('Undo failed:', e);
     }
-  }, [sessionId, setProcessing]);
+  }, [sessionId, updateSession]);
 
   // -- Approve tools ------------------------------------------------------
   const approveTools = useCallback(
@@ -343,10 +382,12 @@ export function useAgentChat({ sessionId, isActive, onReady, onError, onSessionD
 
       setNeedsAttention(sessionId, false);
       const hasApproved = approvals.some(a => a.approved);
-      if (hasApproved && isActiveRef.current) setProcessing(true);
+      if (hasApproved) {
+        updateSession(sessionId, { isProcessing: true });
+      }
       return true;
     },
-    [sessionId, chat, setProcessing, setNeedsAttention],
+    [sessionId, chat, updateSession, setNeedsAttention],
   );
 
   // -- Stop (abort SSE stream + interrupt backend agent loop) ---------------

frontend/src/store/agentStore.ts

@@ -8,6 +8,12 @@
  *  - Plan state
  *  - User info / error banners
  *  - Edited scripts (for hf_jobs code editing)
+ *
+ * Per-session state:
+ *  Each session maintains its own snapshot of processing/activity/panel/plan
+ *  state in `sessionStates`. Background sessions keep updating their own
+ *  snapshot via `updateSession()`. The active session's snapshot is mirrored
+ *  to the flat top-level fields so the UI reads from a single place.
  */
 import { create } from 'zustand';
 import type { User } from '@/types/agent';
@@ -46,8 +52,31 @@ export type ActivityStatus =
   | { type: 'waiting-approval' }
   | { type: 'streaming' };
 
+/** State that is tracked per-session (each session has its own copy). */
+export interface PerSessionState {
+  isProcessing: boolean;
+  activityStatus: ActivityStatus;
+  panelData: PanelData | null;
+  panelView: PanelView;
+  panelEditable: boolean;
+  plan: PlanItem[];
+}
+
+const defaultSessionState: PerSessionState = {
+  isProcessing: false,
+  activityStatus: { type: 'idle' },
+  panelData: null,
+  panelView: 'script',
+  panelEditable: false,
+  plan: [],
+};
+
 interface AgentStore {
-  // Global UI flags
+  // ── Per-session state map ───────────────────────────────────────────
+  sessionStates: Record<string, PerSessionState>;
+  activeSessionId: string | null;
+
+  // ── Flat state (mirrors active session — UI reads from here) ────────
   isProcessing: boolean;
   isConnected: boolean;
   activityStatus: ActivityStatus;
@@ -69,7 +98,21 @@ interface AgentStore {
   // Job URLs (tool_call_id -> job URL) for HF jobs
   jobUrls: Record<string, string>;
 
-  // Actions
+  // ── Per-session actions ─────────────────────────────────────────────
+
+  /** Update a session's state. If it's the active session, also update flat state. */
+  updateSession: (sessionId: string, updates: Partial<PerSessionState>) => void;
+
+  /** Get a session's current state (from map, not flat). */
+  getSessionState: (sessionId: string) => PerSessionState;
+
+  /** Switch the active session — restores its state to flat fields. */
+  switchActiveSession: (sessionId: string) => void;
+
+  /** Remove a session's state from the map. */
+  clearSessionState: (sessionId: string) => void;
+
+  // ── Global actions (not per-session) ────────────────────────────────
   setProcessing: (isProcessing: boolean) => void;
   setConnected: (isConnected: boolean) => void;
   setActivityStatus: (status: ActivityStatus) => void;
@@ -94,7 +137,29 @@ interface AgentStore {
   getJobUrl: (toolCallId: string) => string | undefined;
 }
 
+/**
+ * Helper: patch the active session's snapshot with partial per-session fields.
+ * Returns the `sessionStates` slice to spread into a `set()` call, or `{}`
+ * if there's no active session snapshot to update.
+ */
+function syncSnapshot(
+  state: AgentStore,
+  patch: Partial<PerSessionState>,
+): { sessionStates: Record<string, PerSessionState> } | Record<string, never> {
+  const { activeSessionId, sessionStates } = state;
+  if (!activeSessionId || !sessionStates[activeSessionId]) return {};
+  return {
+    sessionStates: {
+      ...sessionStates,
+      [activeSessionId]: { ...sessionStates[activeSessionId], ...patch },
+    },
+  };
+}
+
 export const useAgentStore = create<AgentStore>()((set, get) => ({
+  sessionStates: {},
+  activeSessionId: null,
+
   isProcessing: false,
   isConnected: false,
   activityStatus: { type: 'idle' },
@@ -111,6 +176,89 @@ export const useAgentStore = create<AgentStore>()((set, get) => ({
   editedScripts: {},
   jobUrls: {},
 
+  // ── Per-session state management ──────────────────────────────────
+
+  updateSession: (sessionId, updates) => {
+    const state = get();
+    const current = state.sessionStates[sessionId] || { ...defaultSessionState };
+    const updated = { ...current, ...updates };
+
+    // Apply the processing→idle side effect
+    const processingCleared = 'isProcessing' in updates && !updates.isProcessing;
+    if (processingCleared) {
+      if (updated.activityStatus.type !== 'waiting-approval') {
+        updated.activityStatus = { type: 'idle' };
+      }
+    }
+
+    const isActive = state.activeSessionId === sessionId;
+
+    // Build flat-state mirror: only the fields explicitly in `updates`
+    // (plus activityStatus when the processing→idle side-effect fires).
+    // This prevents overwriting flat fields changed by global setters
+    // (e.g. setPanelView called from CodePanel) with stale snapshot values.
+    let flatMirror: Record<string, unknown> = {};
+    if (isActive) {
+      for (const key of Object.keys(updates)) {
+        flatMirror[key] = updated[key as keyof PerSessionState];
+      }
+      // Side-effect may have changed activityStatus even if it wasn't in updates
+      if (processingCleared) {
+        flatMirror.activityStatus = updated.activityStatus;
+      }
+    }
+
+    set({
+      sessionStates: { ...state.sessionStates, [sessionId]: updated },
+      ...flatMirror,
+    });
+  },
+
+  getSessionState: (sessionId) => {
+    return get().sessionStates[sessionId] || { ...defaultSessionState };
+  },
+
+  switchActiveSession: (sessionId) => {
+    const state = get();
+
+    // Build a new sessionStates map (never mutate the existing object)
+    const updatedStates = { ...state.sessionStates };
+
+    // Save current active session's flat state back to its snapshot
+    if (state.activeSessionId && state.activeSessionId !== sessionId) {
+      updatedStates[state.activeSessionId] = {
+        isProcessing: state.isProcessing,
+        activityStatus: state.activityStatus,
+        panelData: state.panelData,
+        panelView: state.panelView,
+        panelEditable: state.panelEditable,
+        plan: state.plan,
+      };
+    }
+
+    // Restore the new session's state
+    const incoming = updatedStates[sessionId] || { ...defaultSessionState };
+    set({
+      activeSessionId: sessionId,
+      sessionStates: updatedStates,
+      isProcessing: incoming.isProcessing,
+      activityStatus: incoming.activityStatus,
+      panelData: incoming.panelData,
+      panelView: incoming.panelView,
+      panelEditable: incoming.panelEditable,
+      plan: incoming.plan,
+      // Clear transient error on switch
+      error: null,
+    });
+  },
+
+  clearSessionState: (sessionId) => {
+    set((state) => {
+      const { [sessionId]: _, ...rest } = state.sessionStates;
+      return { sessionStates: rest };
+    });
+  },
+
   // ── Global flags ──────────────────────────────────────────────────
 
   setProcessing: (isProcessing) => {
@@ -125,32 +273,56 @@ export const useAgentStore = create<AgentStore>()((set, get) => ({
   setLlmHealthError: (error) => set({ llmHealthError: error }),
 
   // ── Panel (single-artifact) ───────────────────────────────────────
-
-  setPanel: (data, view, editable) => set({
-    panelData: data,
-    panelView: view ?? (data.script ? 'script' : 'output'),
-    panelEditable: editable ?? false,
+  // Each setter also patches the active session's snapshot so that
+  // getSessionState() stays consistent with flat state.
+
+  setPanel: (data, view, editable) => set((state) => {
+    const patch: Partial<PerSessionState> = {
+      panelData: data,
+      panelView: view ?? (data.script ? 'script' : 'output'),
+      panelEditable: editable ?? false,
+    };
+    return { ...patch, ...syncSnapshot(state, patch) };
   }),
 
-  setPanelView: (view) => set({ panelView: view }),
+  setPanelView: (view) => set((state) => {
+    const patch: Partial<PerSessionState> = { panelView: view };
+    return { ...patch, ...syncSnapshot(state, patch) };
+  }),
 
-  setPanelOutput: (output) => set((state) => ({
-    panelData: state.panelData ? { ...state.panelData, output } : null,
-  })),
+  setPanelOutput: (output) => set((state) => {
+    const panelData = state.panelData
+      ? { ...state.panelData, output }
+      : { title: 'Output', output };
+    const patch: Partial<PerSessionState> = { panelData, panelView: 'output' };
+    return { ...patch, ...syncSnapshot(state, patch) };
+  }),
 
-  updatePanelScript: (content) => set((state) => ({
-    panelData: state.panelData?.script
+  updatePanelScript: (content) => set((state) => {
+    const panelData = state.panelData?.script
       ? { ...state.panelData, script: { ...state.panelData.script, content } }
-      : state.panelData,
-  })),
+      : state.panelData;
+    if (!panelData) return {};
+    const patch: Partial<PerSessionState> = { panelData };
+    return { ...patch, ...syncSnapshot(state, patch) };
+  }),
 
-  lockPanel: () => set({ panelEditable: false }),
+  lockPanel: () => set((state) => {
+    const patch: Partial<PerSessionState> = { panelEditable: false };
+    return { ...patch, ...syncSnapshot(state, patch) };
+  }),
 
-  clearPanel: () => set({ panelData: null, panelView: 'script', panelEditable: false }),
+  clearPanel: () => set((state) => {
+    const patch: Partial<PerSessionState> = { panelData: null, panelView: 'script', panelEditable: false };
+    return { ...patch, ...syncSnapshot(state, patch) };
+  }),
 
   // ── Plan ──────────────────────────────────────────────────────────
 
-  setPlan: (plan) => set({ plan }),
+  setPlan: (plan) => set((state) => {
+    const patch: Partial<PerSessionState> = { plan };
+    return { ...patch, ...syncSnapshot(state, patch) };
+  }),
 
   // ── Edited scripts ────────────────────────────────────────────────