feat: cooperative cancellation, session persistence, background generation

backend/routes/agent.py

@@ -298,6 +298,18 @@ async def interrupt_session(
     return {"status": "interrupted", "session_id": session_id}
 
 
+@router.get("/session/{session_id}/messages")
+async def get_session_messages(
+    session_id: str, user: dict = Depends(get_current_user)
+) -> list[dict]:
+    """Return the session's message history from memory."""
+    _check_session_access(session_id, user)
+    agent_session = session_manager.sessions.get(session_id)
+    if not agent_session or not agent_session.is_active:
+        raise HTTPException(status_code=404, detail="Session not found or inactive")
+    return [msg.model_dump() for msg in agent_session.session.context_manager.items]
+
+
 @router.post("/undo/{session_id}")
 async def undo_session(session_id: str, user: dict = Depends(get_current_user)) -> dict:
     """Undo the last turn in a session."""

frontend/src/components/Layout/AppLayout.tsx

@@ -110,7 +110,7 @@ export default function AppLayout() {
 
   const hasAnySessions = sessions.length > 0;
 
-  const { messages, sendMessage, stop, undoLastTurn, approveTools } = useAgentChat({
+  const { messages, sendMessage, stop, undoLastTurn, approveTools, flushMessages } = useAgentChat({
     sessionId: activeSessionId,
     onReady: () => logger.log('Agent ready'),
     onError: (error) => logger.error('Agent error:', error),
@@ -209,7 +209,7 @@ export default function AppLayout() {
         },
       }}
     >
-      <SessionSidebar onClose={handleSidebarClose} />
+      <SessionSidebar onClose={handleSidebarClose} onBeforeSwitch={flushMessages} />
     </Drawer>
   );
 

frontend/src/components/SessionSidebar/SessionSidebar.tsx

@@ -16,6 +16,7 @@ import { apiFetch } from '@/utils/api';
 
 interface SessionSidebarProps {
   onClose?: () => void;
+  onBeforeSwitch?: () => void;
 }
 
 /** Small coloured dot for connection status */
@@ -32,7 +33,7 @@ const StatusDot = ({ connected }: { connected: boolean }) => (
   />
 );
 
-export default function SessionSidebar({ onClose }: SessionSidebarProps) {
+export default function SessionSidebar({ onClose, onBeforeSwitch }: SessionSidebarProps) {
   const { sessions, activeSessionId, createSession, deleteSession, switchSession } =
     useSessionStore();
   const { isConnected, setPlan, clearPanel } =
@@ -44,6 +45,7 @@ export default function SessionSidebar({ onClose }: SessionSidebarProps) {
 
   const handleNewSession = useCallback(async () => {
     if (isCreatingSession) return;
+    onBeforeSwitch?.();
     setIsCreatingSession(true);
     setCapacityError(null);
     try {
@@ -63,7 +65,7 @@ export default function SessionSidebar({ onClose }: SessionSidebarProps) {
     } finally {
       setIsCreatingSession(false);
     }
-  }, [isCreatingSession, createSession, setPlan, clearPanel, onClose]);
+  }, [isCreatingSession, onBeforeSwitch, createSession, setPlan, clearPanel, onClose]);
 
   const handleDelete = useCallback(
     async (sessionId: string, e: React.MouseEvent) => {
@@ -81,12 +83,13 @@ export default function SessionSidebar({ onClose }: SessionSidebarProps) {
 
   const handleSelect = useCallback(
     (sessionId: string) => {
+      onBeforeSwitch?.();
       switchSession(sessionId);
       setPlan([]);
       clearPanel();
       onClose?.();
     },
-    [switchSession, setPlan, clearPanel, onClose],
+    [onBeforeSwitch, switchSession, setPlan, clearPanel, onClose],
   );
 
   const formatTime = (d: string) =>
@@ -270,7 +273,7 @@ export default function SessionSidebar({ onClose }: SessionSidebarProps) {
         )}
       </Box>
 
-      {/* ── Footer: New Session + status ──────────────────────────── */}
+      {/* ── Footer: New Task + status ──────────────────────────── */}
       <Divider sx={{ opacity: 0.5 }} />
       <Box
         sx={{
@@ -319,7 +322,7 @@ export default function SessionSidebar({ onClose }: SessionSidebarProps) {
           ) : (
             <>
               <AddIcon sx={{ fontSize: 16 }} />
-              New Session
+              New Task
             </>
           )}
         </Box>

frontend/src/hooks/useAgentChat.ts

@@ -8,6 +8,7 @@ import { useChat } from '@ai-sdk/react';
 import type { UIMessage } from 'ai';
 import { WebSocketChatTransport, type SideChannelCallbacks } from '@/lib/ws-chat-transport';
 import { loadMessages, saveMessages } from '@/lib/chat-message-store';
+import { llmMessagesToUIMessages } from '@/lib/convert-llm-messages';
 import { apiFetch } from '@/utils/api';
 import { useAgentStore } from '@/store/agentStore';
 import { useSessionStore } from '@/store/sessionStore';
@@ -210,12 +211,6 @@ export function useAgentChat({ sessionId, onReady, onError, onSessionDead }: Use
     messages: initialMessages,
     transport: transportRef.current!,
     experimental_throttle: 80,
-    onFinish: ({ messages, isAbort, isError }) => {
-      if (isAbort || isError) return;
-      if (sessionId && messages.length > 0) {
-        saveMessages(sessionId, messages);
-      }
-    },
     onError: (error) => {
       logger.error('useChat error:', error);
       setError(error.message);
@@ -227,7 +222,30 @@ export function useAgentChat({ sessionId, onReady, onError, onSessionDead }: Use
   chatActionsRef.current.setMessages = chat.setMessages;
   chatActionsRef.current.messages = chat.messages;
 
-  // ── Persist messages on every user send (onFinish covers assistant turns) ──
+  // ── Hydrate from backend when switching to a session ──────────────
+  useEffect(() => {
+    if (!sessionId) return;
+    let cancelled = false;
+    apiFetch(`/api/session/${sessionId}/messages`)
+      .then((res) => (res.ok ? res.json() : null))
+      .then((data) => {
+        if (cancelled || !data || !Array.isArray(data) || data.length === 0) return;
+        const uiMsgs = llmMessagesToUIMessages(data);
+        if (uiMsgs.length > 0) {
+          chat.setMessages(uiMsgs);
+          saveMessages(sessionId, uiMsgs);
+        }
+      })
+      .catch(() => { /* backend unreachable — localStorage fallback is fine */ });
+    return () => { cancelled = true; };
+  }, [sessionId]); // eslint-disable-line react-hooks/exhaustive-deps
+
+  // ── Persist messages ──────────────────────────────────────────────
+  const flushRef = useRef<{ sid: string | null; msgs: UIMessage[] }>({ sid: null, msgs: [] });
+  flushRef.current.sid = sessionId;
+  flushRef.current.msgs = chat.messages;
+
+  // Save whenever message count changes (covers user sends + new assistant msgs)
   const prevLenRef = useRef(initialMessages.length);
   useEffect(() => {
     if (!sessionId || chat.messages.length === 0) return;
@@ -267,6 +285,12 @@ export function useAgentChat({ sessionId, onReady, onError, onSessionDead }: Use
     [sessionId, setProcessing],
   );
 
+  // ── Flush current messages to localStorage (call before switching sessions) ──
+  const flushMessages = useCallback(() => {
+    const { sid, msgs } = flushRef.current;
+    if (sid && msgs.length > 0) saveMessages(sid, msgs);
+  }, []);
+
   return {
     messages: chat.messages,
     sendMessage: chat.sendMessage,
@@ -274,6 +298,7 @@ export function useAgentChat({ sessionId, onReady, onError, onSessionDead }: Use
     status: chat.status,
     undoLastTurn,
     approveTools,
+    flushMessages,
     transport: transportRef.current,
   };
 }

frontend/src/lib/convert-llm-messages.ts

@@ -0,0 +1,96 @@
+/**
+ * Convert backend LLM messages (litellm format) to Vercel AI SDK UIMessage format.
+ */
+import type { UIMessage } from 'ai';
+
+interface LLMToolCall {
+  id: string;
+  function: { name: string; arguments: string };
+}
+
+interface LLMMessage {
+  role: 'user' | 'assistant' | 'tool' | 'system';
+  content: string | null;
+  tool_calls?: LLMToolCall[] | null;
+  tool_call_id?: string | null;
+  name?: string | null;
+}
+
+let idCounter = 0;
+function nextId(): string {
+  return `msg-${Date.now()}-${++idCounter}`;
+}
+
+export function llmMessagesToUIMessages(messages: LLMMessage[]): UIMessage[] {
+  // Build a map of tool_call_id -> tool result for pairing
+  const toolResults = new Map<string, { output: string; isError: boolean }>();
+  for (const msg of messages) {
+    if (msg.role === 'tool' && msg.tool_call_id) {
+      toolResults.set(msg.tool_call_id, {
+        output: msg.content || '',
+        isError: false,
+      });
+    }
+  }
+
+  const uiMessages: UIMessage[] = [];
+
+  for (const msg of messages) {
+    if (msg.role === 'system') continue;
+    if (msg.role === 'tool') continue; // handled via tool_calls pairing
+
+    if (msg.role === 'user') {
+      uiMessages.push({
+        id: nextId(),
+        role: 'user',
+        parts: [{ type: 'text', text: msg.content || '' }],
+      });
+      continue;
+    }
+
+    if (msg.role === 'assistant') {
+      const parts: UIMessage['parts'] = [];
+
+      if (msg.content) {
+        parts.push({ type: 'text', text: msg.content });
+      }
+
+      if (msg.tool_calls) {
+        for (const tc of msg.tool_calls) {
+          let input: Record<string, unknown> = {};
+          try {
+            input = JSON.parse(tc.function.arguments);
+          } catch { /* malformed */ }
+
+          const result = toolResults.get(tc.id);
+          if (result) {
+            parts.push({
+              type: 'dynamic-tool',
+              toolCallId: tc.id,
+              toolName: tc.function.name,
+              state: 'output-available',
+              input,
+              output: result.output,
+            });
+          } else {
+            parts.push({
+              type: 'dynamic-tool',
+              toolCallId: tc.id,
+              toolName: tc.function.name,
+              state: 'input-available',
+              input,
+            });
+          }
+        }
+      }
+
+      uiMessages.push({
+        id: nextId(),
+        role: 'assistant',
+        parts,
+      });
+    }
+  }
+
+  return uiMessages;
+}

frontend/src/lib/ws-chat-transport.ts

@@ -66,6 +66,9 @@ export class WebSocketChatTransport implements ChatTransport<UIMessage> {
   private currentSessionId: string | null = null;
   private sideChannel: SideChannelCallbacks;
 
+  /** Background WebSockets kept alive so the backend agent keeps running. */
+  private backgroundSockets: Map<string, WebSocket> = new Map();
+
   private streamController: ReadableStreamDefaultController<UIMessageChunk> | null = null;
   private streamGeneration = 0;
   private abortedGeneration = 0;
@@ -130,9 +133,56 @@ export class WebSocketChatTransport implements ChatTransport<UIMessage> {
       clearTimeout(this.connectTimeout);
       this.connectTimeout = null;
     }
-    this.disconnectWebSocket();
+
+    // Move current WS to background instead of closing it
+    if (this.ws && this.currentSessionId && this.currentSessionId !== sessionId) {
+      const oldId = this.currentSessionId;
+      const oldWs = this.ws;
+      this.backgroundSockets.set(oldId, oldWs);
+      // Replace handler: background sockets only need ping/pong
+      oldWs.onmessage = (evt) => {
+        try {
+          const raw = JSON.parse(evt.data);
+          if (raw.type === 'pong') return;
+          // Silently discard — backend keeps running, we'll load results from localStorage
+        } catch { /* ignore */ }
+      };
+      oldWs.onclose = () => {
+        this.backgroundSockets.delete(oldId);
+      };
+      this.ws = null;
+      this.stopPing();
+    } else {
+      this.disconnectWebSocket();
+    }
+
     this.currentSessionId = sessionId;
     if (sessionId) {
+      // Promote background socket if one exists for this session
+      const bg = this.backgroundSockets.get(sessionId);
+      if (bg && (bg.readyState === WebSocket.OPEN || bg.readyState === WebSocket.CONNECTING)) {
+        this.backgroundSockets.delete(sessionId);
+        this.ws = bg;
+        // Restore full event handling
+        bg.onmessage = (evt) => {
+          try {
+            const raw = JSON.parse(evt.data);
+            if (raw.type === 'pong') return;
+            this.handleEvent(raw as AgentEvent);
+          } catch (e) {
+            logger.error('WS parse error:', e);
+          }
+        };
+        bg.onclose = (evt) => {
+          logger.log('WS closed', evt.code, evt.reason);
+          this.sideChannel.onConnectionChange(false);
+          this.stopPing();
+        };
+        this.sideChannel.onConnectionChange(true);
+        this.startPing();
+        return;
+      }
+
       this.retries = 0;
       this.reconnectDelay = WS_RECONNECT_DELAY;
       this.connectTimeout = setTimeout(() => {
@@ -171,6 +221,9 @@ export class WebSocketChatTransport implements ChatTransport<UIMessage> {
       document.removeEventListener('visibilitychange', this.boundVisibilityHandler);
       this.boundVisibilityHandler = null;
     }
+    // Close all background sockets
+    for (const ws of this.backgroundSockets.values()) ws.close();
+    this.backgroundSockets.clear();
     this.disconnectWebSocket();
     this.closeActiveStream();
   }