feat: migrate frontend to Vercel AI SDK with robust WebSocket transport

Replace custom useAgentWebSocket + agentStore message management with
Vercel AI SDK's useChat hook and a custom WebSocketChatTransport.

Key changes:
- Add ws-chat-transport.ts: bridges backend WS events to UIMessageChunk
- Add useAgentChat.ts: central hook wiring useChat + transport + side-channel
- Add chat-message-store.ts: localStorage persistence for messages
- Restore message persistence and timestamps after SDK migration
- Add stop button (abort signal wiring + backend interrupt)
- Fix stream race conditions (generation counter, awaitingProcessing flag)
- Fix WS cycling on rapid session changes (connect timeout guard)
- Robust undo: syncs useChat state + localStorage on undo_complete
- Replace inline session-expired banner with debounced Snackbar toast
- Refactor ToolCallGroup to use AI SDK DynamicToolPart types
- Simplify agentStore (remove message/trace management, keep UI state)

Co-authored-by: Cursor <cursoragent@cursor.com>

agent/core/agent_loop.py

@@ -115,6 +115,42 @@ def _needs_approval(
 class Handlers:
     """Handler functions for each operation type"""
 
+    @staticmethod
+    async def _abandon_pending_approval(session: Session) -> None:
+        """Cancel pending approval tools when the user continues the conversation.
+
+        Injects rejection tool-result messages into the LLM context (so the
+        history stays valid) and notifies the frontend that those tools were
+        abandoned.
+        """
+        tool_calls = session.pending_approval.get("tool_calls", [])
+        for tc in tool_calls:
+            tool_name = tc.function.name
+            abandon_msg = "Task abandoned — user continued the conversation without approving."
+
+            # Keep LLM context valid: every tool_call needs a tool result
+            tool_msg = Message(
+                role="tool",
+                content=abandon_msg,
+                tool_call_id=tc.id,
+                name=tool_name,
+            )
+            session.context_manager.add_message(tool_msg)
+
+            await session.send_event(
+                Event(
+                    event_type="tool_state_change",
+                    data={
+                        "tool_call_id": tc.id,
+                        "tool": tool_name,
+                        "state": "abandoned",
+                    },
+                )
+            )
+
+        session.pending_approval = None
+        logger.info("Abandoned %d pending approval tool(s)", len(tool_calls))
+
     @staticmethod
     @observe(name="run_agent")
     async def run_agent(
@@ -130,6 +166,11 @@ class Handlers:
 
             Laminar.set_trace_session_id(session_id=session.session_id)
 
+        # If there's a pending approval and the user sent a new message,
+        # abandon the pending tools so the LLM context stays valid.
+        if text and session.pending_approval:
+            await Handlers._abandon_pending_approval(session)
+
         # Add user message to history only if there's actual content
         if text:
             user_msg = Message(role="user", content=text)
@@ -553,14 +594,19 @@ class Handlers:
 
         # Execute all approved tools concurrently
         async def execute_tool(tc, tool_name, tool_args):
-            """Execute a single tool and return its result"""
+            """Execute a single tool and return its result.
+
+            The TraceLog already exists on the frontend (created by
+            approval_required), so we send tool_state_change instead of
+            tool_call to avoid creating a duplicate.
+            """
             await session.send_event(
                 Event(
-                    event_type="tool_call",
+                    event_type="tool_state_change",
                     data={
-                        "tool": tool_name,
-                        "arguments": tool_args,
                         "tool_call_id": tc.id,
+                        "tool": tool_name,
+                        "state": "running",
                     },
                 )
             )

configs/main_agent_config.json

@@ -3,7 +3,7 @@
   "save_sessions": true,
   "session_dataset_repo": "akseljoonas/hf-agent-sessions",
   "yolo_mode": false,
-  "confirm_cpu_jobs": false,
+  "confirm_cpu_jobs": true,
   "auto_file_upload": true,
   "mcpServers": {
     "hf-mcp-server": {

frontend/package-lock.json

@@ -8,10 +8,12 @@
       "name": "hf-agent-frontend",
       "version": "1.0.0",
       "dependencies": {
+        "@ai-sdk/react": "^3.0.93",
         "@emotion/react": "^11.13.0",
         "@emotion/styled": "^11.13.0",
         "@mui/icons-material": "^6.1.0",
         "@mui/material": "^6.1.0",
+        "ai": "^6.0.91",
         "react": "^18.3.1",
         "react-dom": "^18.3.1",
         "react-markdown": "^9.0.1",
@@ -34,6 +36,70 @@
         "vite": "^5.4.10"
       }
     },
+    "node_modules/@ai-sdk/gateway": {
+      "version": "3.0.50",
+      "resolved": "https://registry.npmjs.org/@ai-sdk/gateway/-/gateway-3.0.50.tgz",
+      "integrity": "sha512-Jdd1a8VgbD7l7r+COj0h5SuaYRfPvOJ/AO6l0OrmTPEcI2MUQPr3C4JttfpNkcheEN+gOdy0CtZWuG17bW2fjw==",
+      "license": "Apache-2.0",
+      "dependencies": {
+        "@ai-sdk/provider": "3.0.8",
+        "@ai-sdk/provider-utils": "4.0.15",
+        "@vercel/oidc": "3.1.0"
+      },
+      "engines": {
+        "node": ">=18"
+      },
+      "peerDependencies": {
+        "zod": "^3.25.76 || ^4.1.8"
+      }
+    },
+    "node_modules/@ai-sdk/provider": {
+      "version": "3.0.8",
+      "resolved": "https://registry.npmjs.org/@ai-sdk/provider/-/provider-3.0.8.tgz",
+      "integrity": "sha512-oGMAgGoQdBXbZqNG0Ze56CHjDZ1IDYOwGYxYjO5KLSlz5HiNQ9udIXsPZ61VWaHGZ5XW/jyjmr6t2xz2jGVwbQ==",
+      "license": "Apache-2.0",
+      "dependencies": {
+        "json-schema": "^0.4.0"
+      },
+      "engines": {
+        "node": ">=18"
+      }
+    },
+    "node_modules/@ai-sdk/provider-utils": {
+      "version": "4.0.15",
+      "resolved": "https://registry.npmjs.org/@ai-sdk/provider-utils/-/provider-utils-4.0.15.tgz",
+      "integrity": "sha512-8XiKWbemmCbvNN0CLR9u3PQiet4gtEVIrX4zzLxnCj06AwsEDJwJVBbKrEI4t6qE8XRSIvU2irka0dcpziKW6w==",
+      "license": "Apache-2.0",
+      "dependencies": {
+        "@ai-sdk/provider": "3.0.8",
+        "@standard-schema/spec": "^1.1.0",
+        "eventsource-parser": "^3.0.6"
+      },
+      "engines": {
+        "node": ">=18"
+      },
+      "peerDependencies": {
+        "zod": "^3.25.76 || ^4.1.8"
+      }
+    },
+    "node_modules/@ai-sdk/react": {
+      "version": "3.0.93",
+      "resolved": "https://registry.npmjs.org/@ai-sdk/react/-/react-3.0.93.tgz",
+      "integrity": "sha512-FY1HmeAfCpiAGLhIZh2QR8QFzHFZfhjMmkA9D5KC/O3eGqPeY7CwBABLkzRH+5Gkf+MfxXnEm4VF0MpmvDMjpg==",
+      "license": "Apache-2.0",
+      "dependencies": {
+        "@ai-sdk/provider-utils": "4.0.15",
+        "ai": "6.0.91",
+        "swr": "^2.2.5",
+        "throttleit": "2.1.0"
+      },
+      "engines": {
+        "node": ">=18"
+      },
+      "peerDependencies": {
+        "react": "^18 || ~19.0.1 || ~19.1.2 || ^19.2.1"
+      }
+    },
     "node_modules/@babel/code-frame": {
       "version": "7.28.6",
       "resolved": "https://registry.npmjs.org/@babel/code-frame/-/code-frame-7.28.6.tgz",
@@ -1348,6 +1414,15 @@
         }
       }
     },
+    "node_modules/@opentelemetry/api": {
+      "version": "1.9.0",
+      "resolved": "https://registry.npmjs.org/@opentelemetry/api/-/api-1.9.0.tgz",
+      "integrity": "sha512-3giAOQvZiH5F9bMlMiv8+GSPMeqg0dbaeo58/0SlA9sxSqZhnUtxzX9/2FzyhS9sWQf5S0GJE0AKBrFqjpeYcg==",
+      "license": "Apache-2.0",
+      "engines": {
+        "node": ">=8.0.0"
+      }
+    },
     "node_modules/@popperjs/core": {
       "version": "2.11.8",
       "resolved": "https://registry.npmjs.org/@popperjs/core/-/core-2.11.8.tgz",
@@ -1715,6 +1790,12 @@
         "win32"
       ]
     },
+    "node_modules/@standard-schema/spec": {
+      "version": "1.1.0",
+      "resolved": "https://registry.npmjs.org/@standard-schema/spec/-/spec-1.1.0.tgz",
+      "integrity": "sha512-l2aFy5jALhniG5HgqrD6jXLi/rUWrKvqN/qJx6yoJsgKhblVd+iqqU4RCXavm/jPityDo5TCvKMnpjKnOriy0w==",
+      "license": "MIT"
+    },
     "node_modules/@types/babel__core": {
       "version": "7.20.5",
       "resolved": "https://registry.npmjs.org/@types/babel__core/-/babel__core-7.20.5.tgz",
@@ -2155,6 +2236,15 @@
       "integrity": "sha512-WmoN8qaIAo7WTYWbAZuG8PYEhn5fkz7dZrqTBZ7dtt//lL2Gwms1IcnQ5yHqjDfX8Ft5j4YzDM23f87zBfDe9g==",
       "license": "ISC"
     },
+    "node_modules/@vercel/oidc": {
+      "version": "3.1.0",
+      "resolved": "https://registry.npmjs.org/@vercel/oidc/-/oidc-3.1.0.tgz",
+      "integrity": "sha512-Fw28YZpRnA3cAHHDlkt7xQHiJ0fcL+NRcIqsocZQUSmbzeIKRpwttJjik5ZGanXP+vlA4SbTg+AbA3bP363l+w==",
+      "license": "Apache-2.0",
+      "engines": {
+        "node": ">= 20"
+      }
+    },
     "node_modules/@vitejs/plugin-react": {
       "version": "4.7.0",
       "resolved": "https://registry.npmjs.org/@vitejs/plugin-react/-/plugin-react-4.7.0.tgz",
@@ -2200,6 +2290,24 @@
         "acorn": "^6.0.0 || ^7.0.0 || ^8.0.0"
       }
     },
+    "node_modules/ai": {
+      "version": "6.0.91",
+      "resolved": "https://registry.npmjs.org/ai/-/ai-6.0.91.tgz",
+      "integrity": "sha512-k1/8BusZMhYVxxLZt0BUZzm9HVDCCh117nyWfWUx5xjR2+tWisJbXgysL7EBMq2lgyHwgpA1jDR3tVjWSdWZXw==",
+      "license": "Apache-2.0",
+      "dependencies": {
+        "@ai-sdk/gateway": "3.0.50",
+        "@ai-sdk/provider": "3.0.8",
+        "@ai-sdk/provider-utils": "4.0.15",
+        "@opentelemetry/api": "1.9.0"
+      },
+      "engines": {
+        "node": ">=18"
+      },
+      "peerDependencies": {
+        "zod": "^3.25.76 || ^4.1.8"
+      }
+    },
     "node_modules/ajv": {
       "version": "6.12.6",
       "resolved": "https://registry.npmjs.org/ajv/-/ajv-6.12.6.tgz",
@@ -2848,6 +2956,15 @@
         "node": ">=0.10.0"
       }
     },
+    "node_modules/eventsource-parser": {
+      "version": "3.0.6",
+      "resolved": "https://registry.npmjs.org/eventsource-parser/-/eventsource-parser-3.0.6.tgz",
+      "integrity": "sha512-Vo1ab+QXPzZ4tCa8SwIHJFaSzy4R6SHf7BY79rFBDf0idraZWAkYrDjDj8uWaSm3S2TK+hJ7/t1CEmZ7jXw+pg==",
+      "license": "MIT",
+      "engines": {
+        "node": ">=18.0.0"
+      }
+    },
     "node_modules/extend": {
       "version": "3.0.2",
       "resolved": "https://registry.npmjs.org/extend/-/extend-3.0.2.tgz",
@@ -3356,6 +3473,12 @@
       "integrity": "sha512-xyFwyhro/JEof6Ghe2iz2NcXoj2sloNsWr/XsERDK/oiPCfaNhl5ONfp+jQdAZRQQ0IJWNzH9zIZF7li91kh2w==",
       "license": "MIT"
     },
+    "node_modules/json-schema": {
+      "version": "0.4.0",
+      "resolved": "https://registry.npmjs.org/json-schema/-/json-schema-0.4.0.tgz",
+      "integrity": "sha512-es94M3nTIfsEPisRafak+HDLfHXnKBhV3vU5eqPcS3flIWqcxJWgXHXiey3YrpaNsanY5ei1VoYEbOzijuq9BA==",
+      "license": "(AFL-2.1 OR BSD-3-Clause)"
+    },
     "node_modules/json-schema-traverse": {
       "version": "0.4.1",
       "resolved": "https://registry.npmjs.org/json-schema-traverse/-/json-schema-traverse-0.4.1.tgz",
@@ -5052,6 +5175,31 @@
         "url": "https://github.com/sponsors/ljharb"
       }
     },
+    "node_modules/swr": {
+      "version": "2.4.0",
+      "resolved": "https://registry.npmjs.org/swr/-/swr-2.4.0.tgz",
+      "integrity": "sha512-sUlC20T8EOt1pHmDiqueUWMmRRX03W7w5YxovWX7VR2KHEPCTMly85x05vpkP5i6Bu4h44ePSMD9Tc+G2MItFw==",
+      "license": "MIT",
+      "dependencies": {
+        "dequal": "^2.0.3",
+        "use-sync-external-store": "^1.6.0"
+      },
+      "peerDependencies": {
+        "react": "^16.11.0 || ^17.0.0 || ^18.0.0 || ^19.0.0"
+      }
+    },
+    "node_modules/throttleit": {
+      "version": "2.1.0",
+      "resolved": "https://registry.npmjs.org/throttleit/-/throttleit-2.1.0.tgz",
+      "integrity": "sha512-nt6AMGKW1p/70DF/hGBdJB57B8Tspmbp5gfJ8ilhLnt7kkr2ye7hzD6NVG8GGErk2HWF34igrL2CXmNIkzKqKw==",
+      "license": "MIT",
+      "engines": {
+        "node": ">=18"
+      },
+      "funding": {
+        "url": "https://github.com/sponsors/sindresorhus"
+      }
+    },
     "node_modules/tinyglobby": {
       "version": "0.2.15",
       "resolved": "https://registry.npmjs.org/tinyglobby/-/tinyglobby-0.2.15.tgz",
@@ -5282,6 +5430,16 @@
         "punycode": "^2.1.0"
       }
     },
+    "node_modules/use-sync-external-store": {
+      "version": "1.6.0",
+      "resolved": "https://registry.npmjs.org/use-sync-external-store/-/use-sync-external-store-1.6.0.tgz",
+      "integrity": "sha512-Pp6GSwGP/NrPIrxVFAIkOQeyw8lFenOHijQWkUTrDvrF4ALqylP2C/KCkeS9dpUM3KvYRQhna5vt7IL95+ZQ9w==",
+      "license": "MIT",
+      "peer": true,
+      "peerDependencies": {
+        "react": "^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0"
+      }
+    },
     "node_modules/vfile": {
       "version": "6.0.3",
       "resolved": "https://registry.npmjs.org/vfile/-/vfile-6.0.3.tgz",
@@ -5426,6 +5584,16 @@
         "url": "https://github.com/sponsors/sindresorhus"
       }
     },
+    "node_modules/zod": {
+      "version": "4.3.6",
+      "resolved": "https://registry.npmjs.org/zod/-/zod-4.3.6.tgz",
+      "integrity": "sha512-rftlrkhHZOcjDwkGlnUtZZkvaPHCsDATp4pGpuOOMDaTdDDXF91wuVDJoWoPsKX/3YPQ5fHuF3STjcYyKr+Qhg==",
+      "license": "MIT",
+      "peer": true,
+      "funding": {
+        "url": "https://github.com/sponsors/colinhacks"
+      }
+    },
     "node_modules/zustand": {
       "version": "5.0.10",
       "resolved": "https://registry.npmjs.org/zustand/-/zustand-5.0.10.tgz",

frontend/package.json

@@ -10,10 +10,12 @@
     "preview": "vite preview"
   },
   "dependencies": {
+    "@ai-sdk/react": "^3.0.93",
     "@emotion/react": "^11.13.0",
     "@emotion/styled": "^11.13.0",
     "@mui/icons-material": "^6.1.0",
     "@mui/material": "^6.1.0",
+    "ai": "^6.0.91",
     "react": "^18.3.1",
     "react-dom": "^18.3.1",
     "react-markdown": "^9.0.1",

frontend/src/components/Chat/AssistantMessage.tsx

@@ -1,54 +1,66 @@
+import { useMemo } from 'react';
 import { Box, Stack, Typography } from '@mui/material';
 import MarkdownContent from './MarkdownContent';
 import ToolCallGroup from './ToolCallGroup';
-import type { Message } from '@/types/agent';
+import type { UIMessage } from 'ai';
+import type { MessageMeta } from '@/types/agent';
 
 interface AssistantMessageProps {
-  message: Message;
-  /** True when this message is actively receiving streaming chunks. */
+  message: UIMessage;
   isStreaming?: boolean;
+  approveTools: (approvals: Array<{ tool_call_id: string; approved: boolean; feedback?: string | null }>) => Promise<boolean>;
 }
 
-export default function AssistantMessage({ message, isStreaming = false }: AssistantMessageProps) {
-  const renderSegments = () => {
-    if (message.segments && message.segments.length > 0) {
-      // Find the index of the last text segment (that's the one being streamed)
-      let lastTextIdx = -1;
-      for (let i = message.segments.length - 1; i >= 0; i--) {
-        if (message.segments[i].type === 'text') {
-          lastTextIdx = i;
-          break;
-        }
-      }
+/**
+ * Groups consecutive tool parts together so they render as a single
+ * ToolCallGroup (visually identical to the old segments approach).
+ */
+type DynamicToolPart = Extract<UIMessage['parts'][number], { type: 'dynamic-tool' }>;
 
-      return message.segments.map((segment, idx) => {
-        if (segment.type === 'text' && segment.content) {
-          return (
-            <MarkdownContent
-              key={idx}
-              content={segment.content}
-              isStreaming={isStreaming && idx === lastTextIdx}
-            />
-          );
-        }
-        if (segment.type === 'tools' && segment.tools && segment.tools.length > 0) {
-          return <ToolCallGroup key={idx} tools={segment.tools} />;
-        }
-        return null;
-      });
-    }
+function groupParts(parts: UIMessage['parts']) {
+  const groups: Array<
+    | { kind: 'text'; text: string; idx: number }
+    | { kind: 'tools'; tools: DynamicToolPart[]; idx: number }
+  > = [];
 
-    // Fallback: render raw content
-    if (message.content) {
-      return <MarkdownContent content={message.content} isStreaming={isStreaming} />;
+  for (let i = 0; i < parts.length; i++) {
+    const part = parts[i];
+
+    if (part.type === 'text') {
+      groups.push({ kind: 'text', text: part.text, idx: i });
+    } else if (part.type === 'dynamic-tool') {
+      const toolPart = part as DynamicToolPart;
+      const last = groups[groups.length - 1];
+      if (last?.kind === 'tools') {
+        last.tools.push(toolPart);
+      } else {
+        groups.push({ kind: 'tools', tools: [toolPart], idx: i });
+      }
     }
+    // step-start, step-end, etc. are ignored visually
+  }
+
+  return groups;
+}
+
+export default function AssistantMessage({ message, isStreaming = false, approveTools }: AssistantMessageProps) {
+  const groups = useMemo(() => groupParts(message.parts), [message.parts]);
+
+  // Find the last text group index for streaming cursor
+  let lastTextIdx = -1;
+  for (let i = groups.length - 1; i >= 0; i--) {
+    if (groups[i].kind === 'text') { lastTextIdx = i; break; }
+  }
 
-    return null;
-  };
+  const meta = message.metadata as MessageMeta | undefined;
+  const timeStr = meta?.createdAt
+    ? new Date(meta.createdAt).toLocaleTimeString([], { hour: '2-digit', minute: '2-digit' })
+    : null;
+
+  if (groups.length === 0) return null;
 
   return (
     <Box sx={{ minWidth: 0 }}>
-      {/* Role label + timestamp */}
       <Stack direction="row" alignItems="baseline" spacing={1} sx={{ mb: 0.5 }}>
         <Typography
           variant="caption"
@@ -62,19 +74,13 @@ export default function AssistantMessage({ message, isStreaming = false }: Assis
         >
           Assistant
         </Typography>
-        <Typography
-          variant="caption"
-          sx={{
-            fontSize: '0.66rem',
-            color: 'var(--muted-text)',
-            opacity: 0.6,
-          }}
-        >
-          {new Date(message.timestamp).toLocaleTimeString([], { hour: '2-digit', minute: '2-digit' })}
-        </Typography>
+        {timeStr && (
+          <Typography variant="caption" sx={{ color: 'var(--muted-text)', fontSize: '0.7rem' }}>
+            {timeStr}
+          </Typography>
+        )}
       </Stack>
 
-      {/* Message bubble */}
       <Box
         sx={{
           maxWidth: { xs: '95%', md: '85%' },
@@ -86,7 +92,27 @@ export default function AssistantMessage({ message, isStreaming = false }: Assis
           border: '1px solid var(--border)',
         }}
       >
-        {renderSegments()}
+        {groups.map((group, i) => {
+          if (group.kind === 'text' && group.text) {
+            return (
+              <MarkdownContent
+                key={group.idx}
+                content={group.text}
+                isStreaming={isStreaming && i === lastTextIdx}
+              />
+            );
+          }
+          if (group.kind === 'tools' && group.tools.length > 0) {
+            return (
+              <ToolCallGroup
+                key={group.idx}
+                tools={group.tools}
+                approveTools={approveTools}
+              />
+            );
+          }
+          return null;
+        })}
       </Box>
     </Box>
   );

frontend/src/components/Chat/ChatInput.tsx

@@ -1,6 +1,7 @@
 import { useState, useCallback, useEffect, useRef, KeyboardEvent } from 'react';
 import { Box, TextField, IconButton, CircularProgress, Typography, Menu, MenuItem, ListItemIcon, ListItemText, Chip } from '@mui/material';
 import ArrowUpwardIcon from '@mui/icons-material/ArrowUpward';
+import StopIcon from '@mui/icons-material/Stop';
 import ArrowDropDownIcon from '@mui/icons-material/ArrowDropDown';
 import { apiFetch } from '@/utils/api';
 
@@ -58,23 +59,34 @@ const findModelByPath = (path: string): ModelOption | undefined => {
 
 interface ChatInputProps {
   onSend: (text: string) => void;
+  onStop?: () => void;
   disabled?: boolean;
+  isStreaming?: boolean;
 }
 
-export default function ChatInput({ onSend, disabled = false }: ChatInputProps) {
+export default function ChatInput({ onSend, onStop, disabled = false, isStreaming = false }: ChatInputProps) {
   const [input, setInput] = useState('');
   const inputRef = useRef<HTMLTextAreaElement>(null);
-  const [selectedModelId, setSelectedModelId] = useState<string>(MODEL_OPTIONS[0].id);
+  const [selectedModelId, setSelectedModelId] = useState<string>(() => {
+    try {
+      const stored = localStorage.getItem('hf-agent-model');
+      if (stored && MODEL_OPTIONS.some(m => m.id === stored)) return stored;
+    } catch { /* localStorage unavailable */ }
+    return MODEL_OPTIONS[0].id;
+  });
   const [modelAnchorEl, setModelAnchorEl] = useState<null | HTMLElement>(null);
 
-  // Sync with backend on mount
+  // Sync with backend on mount (backend is source of truth, localStorage is just a cache)
   useEffect(() => {
     fetch('/api/config/model')
       .then((res) => (res.ok ? res.json() : null))
       .then((data) => {
         if (data?.current) {
           const model = findModelByPath(data.current);
-          if (model) setSelectedModelId(model.id);
+          if (model) {
+            setSelectedModelId(model.id);
+            try { localStorage.setItem('hf-agent-model', model.id); } catch { /* ignore */ }
+          }
         }
       })
       .catch(() => { /* ignore */ });
@@ -123,6 +135,7 @@ export default function ChatInput({ onSend, disabled = false }: ChatInputProps)
       });
       if (res.ok) {
         setSelectedModelId(model.id);
+        try { localStorage.setItem('hf-agent-model', model.id); } catch { /* ignore */ }
       }
     } catch { /* ignore */ }
   };
@@ -189,26 +202,44 @@ export default function ChatInput({ onSend, disabled = false }: ChatInputProps)
                 }
             }}
           />
-          <IconButton
-            onClick={handleSend}
-            disabled={disabled || !input.trim()}
-            sx={{
-              mt: 1,
-              p: 1,
-              borderRadius: '10px',
-              color: 'var(--muted-text)',
-              transition: 'all 0.2s',
-              '&:hover': {
+          {isStreaming ? (
+            <IconButton
+              onClick={onStop}
+              sx={{
+                mt: 1,
+                p: 1,
+                borderRadius: '10px',
                 color: 'var(--accent-yellow)',
-                bgcolor: 'var(--hover-bg)',
-              },
-              '&.Mui-disabled': {
-                opacity: 0.3,
-              },
-            }}
-          >
-            {disabled ? <CircularProgress size={20} color="inherit" /> : <ArrowUpwardIcon fontSize="small" />}
-          </IconButton>
+                transition: 'all 0.2s',
+                '&:hover': {
+                  bgcolor: 'var(--hover-bg)',
+                },
+              }}
+            >
+              <StopIcon fontSize="small" />
+            </IconButton>
+          ) : (
+            <IconButton
+              onClick={handleSend}
+              disabled={disabled || !input.trim()}
+              sx={{
+                mt: 1,
+                p: 1,
+                borderRadius: '10px',
+                color: 'var(--muted-text)',
+                transition: 'all 0.2s',
+                '&:hover': {
+                  color: 'var(--accent-yellow)',
+                  bgcolor: 'var(--hover-bg)',
+                },
+                '&.Mui-disabled': {
+                  opacity: 0.3,
+                },
+              }}
+            >
+              {disabled ? <CircularProgress size={20} color="inherit" /> : <ArrowUpwardIcon fontSize="small" />}
+            </IconButton>
+          )}
         </Box>
 
         {/* Powered By Badge */}

frontend/src/components/Chat/MessageBubble.tsx

@@ -1,36 +1,24 @@
 import UserMessage from './UserMessage';
 import AssistantMessage from './AssistantMessage';
-import type { Message } from '@/types/agent';
+import type { UIMessage } from 'ai';
 
 interface MessageBubbleProps {
-  message: Message;
-  /** True if this is the user message that starts the last turn. */
+  message: UIMessage;
   isLastTurn?: boolean;
-  /** Callback to undo (remove) the last turn. */
   onUndoTurn?: () => void;
-  /** Whether the agent is currently processing. */
   isProcessing?: boolean;
-  /** True when this message is actively receiving streaming chunks. */
   isStreaming?: boolean;
+  approveTools: (approvals: Array<{ tool_call_id: string; approved: boolean; feedback?: string | null }>) => Promise<boolean>;
 }
 
-/**
- * Thin dispatcher — routes each message to the correct
- * specialised component based on its role / content.
- */
 export default function MessageBubble({
   message,
   isLastTurn = false,
   onUndoTurn,
   isProcessing = false,
   isStreaming = false,
+  approveTools,
 }: MessageBubbleProps) {
-  // Legacy approval-only messages (from old localStorage data) — skip them.
-  // Approvals are now rendered inline within ToolCallGroup.
-  if (message.approval && !message.content && !message.segments?.length) {
-    return null;
-  }
-
   if (message.role === 'user') {
     return (
       <UserMessage
@@ -43,9 +31,14 @@ export default function MessageBubble({
   }
 
   if (message.role === 'assistant') {
-    return <AssistantMessage message={message} isStreaming={isStreaming} />;
+    return (
+      <AssistantMessage
+        message={message}
+        isStreaming={isStreaming}
+        approveTools={approveTools}
+      />
+    );
   }
 
-  // Fallback (tool messages, etc.)
   return null;
 }

frontend/src/components/Chat/MessageList.tsx

@@ -1,16 +1,15 @@
-import { useEffect, useRef, useMemo, useCallback } from 'react';
+import { useCallback, useEffect, useRef, useMemo } from 'react';
 import { Box, Stack, Typography } from '@mui/material';
 import MessageBubble from './MessageBubble';
 import ThinkingIndicator from './ThinkingIndicator';
 import { useAgentStore } from '@/store/agentStore';
-import { useSessionStore } from '@/store/sessionStore';
-import { apiFetch } from '@/utils/api';
-import { logger } from '@/utils/logger';
-import type { Message } from '@/types/agent';
+import type { UIMessage } from 'ai';
 
 interface MessageListProps {
-  messages: Message[];
+  messages: UIMessage[];
   isProcessing: boolean;
+  approveTools: (approvals: Array<{ tool_call_id: string; approved: boolean; feedback?: string | null }>) => Promise<boolean>;
+  onUndoLastTurn: () => void | Promise<void>;
 }
 
 function getGreeting(): string {
@@ -20,7 +19,6 @@ function getGreeting(): string {
   return 'Evening';
 }
 
-/** Minimal greeting shown when the conversation is empty. */
 function WelcomeGreeting() {
   const { user } = useAgentStore();
   const firstName = user?.name?.split(' ')[0] || user?.username;
@@ -58,58 +56,40 @@ function WelcomeGreeting() {
   );
 }
 
-export default function MessageList({ messages, isProcessing }: MessageListProps) {
+export default function MessageList({ messages, isProcessing, approveTools, onUndoLastTurn }: MessageListProps) {
   const scrollContainerRef = useRef<HTMLDivElement>(null);
   const stickToBottom = useRef(true);
-  const { activeSessionId } = useSessionStore();
-  const { removeLastTurn, currentTurnMessageId } = useAgentStore();
 
-  // ── Scroll-to-bottom helper ─────────────────────────────────────
   const scrollToBottom = useCallback(() => {
     const el = scrollContainerRef.current;
     if (el) el.scrollTop = el.scrollHeight;
   }, []);
 
-  // ── Track user scroll intent ────────────────────────────────────
   useEffect(() => {
     const el = scrollContainerRef.current;
     if (!el) return;
-
     const onScroll = () => {
       const distFromBottom = el.scrollHeight - el.scrollTop - el.clientHeight;
       stickToBottom.current = distFromBottom < 80;
     };
-
     el.addEventListener('scroll', onScroll, { passive: true });
     return () => el.removeEventListener('scroll', onScroll);
   }, []);
 
-  // ── Auto-scroll on new messages / state changes ─────────────────
   useEffect(() => {
     if (stickToBottom.current) scrollToBottom();
   }, [messages, isProcessing, scrollToBottom]);
 
-  // ── Auto-scroll on DOM mutations (streaming content growth) ─────
   useEffect(() => {
     const el = scrollContainerRef.current;
     if (!el) return;
-
     const observer = new MutationObserver(() => {
-      if (stickToBottom.current) {
-        el.scrollTop = el.scrollHeight;
-      }
-    });
-
-    observer.observe(el, {
-      childList: true,
-      subtree: true,
-      characterData: true,
+      if (stickToBottom.current) el.scrollTop = el.scrollHeight;
     });
-
+    observer.observe(el, { childList: true, subtree: true, characterData: true });
     return () => observer.disconnect();
   }, []);
 
-  // Find the index of the last user message (start of the last turn)
   const lastUserMsgId = useMemo(() => {
     for (let i = messages.length - 1; i >= 0; i--) {
       if (messages[i].role === 'user') return messages[i].id;
@@ -117,15 +97,13 @@ export default function MessageList({ messages, isProcessing }: MessageListProps
     return null;
   }, [messages]);
 
-  const handleUndoLastTurn = useCallback(async () => {
-    if (!activeSessionId) return;
-    try {
-      await apiFetch(`/api/undo/${activeSessionId}`, { method: 'POST' });
-      removeLastTurn(activeSessionId);
-    } catch (e) {
-      logger.error('Undo failed:', e);
+  // The last assistant message is "streaming" when we're processing
+  const lastAssistantId = useMemo(() => {
+    for (let i = messages.length - 1; i >= 0; i--) {
+      if (messages[i].role === 'assistant') return messages[i].id;
     }
-  }, [activeSessionId, removeLastTurn]);
+    return null;
+  }, [messages]);
 
   return (
     <Box
@@ -156,17 +134,18 @@ export default function MessageList({ messages, isProcessing }: MessageListProps
               key={msg.id}
               message={msg}
               isLastTurn={msg.id === lastUserMsgId}
-              onUndoTurn={handleUndoLastTurn}
+              onUndoTurn={onUndoLastTurn}
               isProcessing={isProcessing}
-              isStreaming={isProcessing && msg.id === currentTurnMessageId}
+              isStreaming={isProcessing && msg.id === lastAssistantId}
+              approveTools={approveTools}
             />
           ))
         )}
 
-        {/* Show thinking dots only when processing but no streaming message yet */}
-        {isProcessing && !currentTurnMessageId && <ThinkingIndicator />}
+        {isProcessing && (!lastAssistantId || messages.every(m => m.id !== lastAssistantId || m.parts.length === 0)) && (
+          <ThinkingIndicator />
+        )}
 
-        {/* Sentinel — keeps scroll anchor at the bottom */}
         <div />
       </Stack>
     </Box>

frontend/src/components/Chat/ToolCallGroup.tsx

@@ -1,4 +1,4 @@
-import { useCallback, useState } from 'react';
+import { useCallback, useMemo, useRef, useState } from 'react';
 import { Box, Stack, Typography, Chip, Button, TextField, IconButton, Link } from '@mui/material';
 import CheckCircleOutlineIcon from '@mui/icons-material/CheckCircleOutline';
 import ErrorOutlineIcon from '@mui/icons-material/ErrorOutline';
@@ -7,54 +7,40 @@ import OpenInNewIcon from '@mui/icons-material/OpenInNew';
 import HourglassEmptyIcon from '@mui/icons-material/HourglassEmpty';
 import LaunchIcon from '@mui/icons-material/Launch';
 import SendIcon from '@mui/icons-material/Send';
+import BlockIcon from '@mui/icons-material/Block';
 import { useAgentStore } from '@/store/agentStore';
 import { useLayoutStore } from '@/store/layoutStore';
-import { useSessionStore } from '@/store/sessionStore';
-import { apiFetch } from '@/utils/api';
 import { logger } from '@/utils/logger';
-import type { TraceLog, ToolState } from '@/types/agent';
+import type { UIMessage } from 'ai';
+
+// ---------------------------------------------------------------------------
+// Type helpers — extract the dynamic-tool part type from UIMessage
+// ---------------------------------------------------------------------------
+type DynamicToolPart = Extract<UIMessage['parts'][number], { type: 'dynamic-tool' }>;
+
+type ToolPartState = DynamicToolPart['state'];
 
 interface ToolCallGroupProps {
-  tools: TraceLog[];
+  tools: DynamicToolPart[];
+  approveTools: (approvals: Array<{ tool_call_id: string; approved: boolean; feedback?: string | null }>) => Promise<boolean>;
 }
 
-const TOOL_TIMEOUT_MS = 5 * 60 * 1000;
-
-/**
- * Resolve the effective state of a TraceLog.
- * Uses `state` field if present, otherwise infers from legacy fields
- * (backward compat with data persisted before the state refactor).
- */
-function resolveState(log: TraceLog): ToolState {
-  if (log.state) return log.state;
-  // Legacy inference
-  if (log.approvalStatus === 'pending') return 'pending_approval';
-  if (log.approvalStatus === 'rejected') return 'rejected';
-  if (log.completed && log.success === false) return 'failed';
-  if (log.completed) return 'completed';
-  // Check timeout
-  const elapsed = Date.now() - new Date(log.timestamp).getTime();
-  if (elapsed > TOOL_TIMEOUT_MS) return 'timed_out';
-  if (log.approvalStatus === 'approved') return 'running';
-  return 'calling';
-}
+// ---------------------------------------------------------------------------
+// Visual helpers
+// ---------------------------------------------------------------------------
 
-// ── Status icon based on resolved state ──────────────────────────────
-function StatusIcon({ state }: { state: ToolState }) {
+function StatusIcon({ state }: { state: ToolPartState }) {
   switch (state) {
-    case 'pending_approval':
+    case 'approval-requested':
       return <HourglassEmptyIcon sx={{ fontSize: 16, color: 'var(--accent-yellow)' }} />;
-    case 'approved':
-      return <HourglassEmptyIcon sx={{ fontSize: 16, color: 'var(--accent-green)', opacity: 0.7 }} />;
-    case 'rejected':
-    case 'failed':
-      return <ErrorOutlineIcon sx={{ fontSize: 16, color: 'error.main' }} />;
-    case 'timed_out':
-      return <ErrorOutlineIcon sx={{ fontSize: 16, color: 'var(--muted-text)' }} />;
-    case 'completed':
+    case 'output-available':
       return <CheckCircleOutlineIcon sx={{ fontSize: 16, color: 'success.main' }} />;
-    case 'calling':
-    case 'running':
+    case 'output-error':
+      return <ErrorOutlineIcon sx={{ fontSize: 16, color: 'error.main' }} />;
+    case 'output-denied':
+      return <BlockIcon sx={{ fontSize: 16, color: 'var(--muted-text)' }} />;
+    case 'input-streaming':
+    case 'input-available':
     default:
       return (
         <MoreHorizIcon
@@ -72,58 +58,61 @@ function StatusIcon({ state }: { state: ToolState }) {
   }
 }
 
-// ── Status chip label ───────────────────────────────────────────────
-function statusLabel(state: ToolState): string | null {
+function statusLabel(state: ToolPartState): string | null {
   switch (state) {
-    case 'pending_approval': return 'awaiting approval';
-    case 'approved': return 'approved';
-    case 'rejected': return 'rejected';
-    case 'timed_out': return 'timed out';
-    case 'calling':
-    case 'running': return 'running';
+    case 'approval-requested': return 'awaiting approval';
+    case 'input-streaming':
+    case 'input-available': return 'running';
+    case 'output-denied': return 'denied';
+    case 'output-error': return 'error';
     default: return null;
   }
 }
 
-function statusColor(state: ToolState): string {
+function statusColor(state: ToolPartState): string {
   switch (state) {
-    case 'pending_approval': return 'var(--accent-yellow)';
-    case 'approved': return 'var(--accent-green)';
-    case 'rejected':
-    case 'failed': return 'var(--accent-red)';
-    case 'timed_out': return 'var(--muted-text)';
+    case 'approval-requested': return 'var(--accent-yellow)';
+    case 'output-available': return 'var(--accent-green)';
+    case 'output-error': return 'var(--accent-red)';
+    case 'output-denied': return 'var(--muted-text)';
     default: return 'var(--accent-yellow)';
   }
 }
 
-// ── Inline approval UI ──────────────────────────────────────────────
+// ---------------------------------------------------------------------------
+// Inline approval UI (per-tool)
+// ---------------------------------------------------------------------------
+
 function InlineApproval({
-  log,
+  toolCallId,
+  toolName,
+  input,
   onResolve,
 }: {
-  log: TraceLog;
+  toolCallId: string;
+  toolName: string;
+  input: unknown;
   onResolve: (toolCallId: string, approved: boolean, feedback?: string) => void;
 }) {
   const [feedback, setFeedback] = useState('');
+  const args = input as Record<string, unknown> | undefined;
 
   return (
     <Box sx={{ px: 1.5, py: 1.5, borderTop: '1px solid var(--tool-border)' }}>
-      {/* Tool description */}
-      {log.tool === 'hf_jobs' && log.args && (
+      {toolName === 'hf_jobs' && args && (
         <Typography variant="body2" sx={{ color: 'var(--muted-text)', fontSize: '0.75rem', mb: 1.5 }}>
-          Execute <Box component="span" sx={{ color: 'var(--accent-yellow)', fontWeight: 500 }}>{log.tool}</Box> on{' '}
+          Execute <Box component="span" sx={{ color: 'var(--accent-yellow)', fontWeight: 500 }}>{toolName}</Box> on{' '}
           <Box component="span" sx={{ fontWeight: 500, color: 'var(--text)' }}>
-            {String(log.args.hardware_flavor || 'default')}
+            {String(args.hardware_flavor || 'default')}
           </Box>
-          {!!log.args.timeout && (
+          {!!args.timeout && (
             <> with timeout <Box component="span" sx={{ fontWeight: 500, color: 'var(--text)' }}>
-              {String(log.args.timeout)}
+              {String(args.timeout)}
             </Box></>
           )}
         </Typography>
       )}
 
-      {/* Feedback + buttons */}
       <Box sx={{ display: 'flex', gap: 1, mb: 1 }}>
         <TextField
           fullWidth
@@ -141,7 +130,7 @@ function InlineApproval({
           }}
         />
         <IconButton
-          onClick={() => onResolve(log.toolCallId || '', false, feedback || 'Rejected by user')}
+          onClick={() => onResolve(toolCallId, false, feedback || 'Rejected by user')}
           disabled={!feedback}
           size="small"
           sx={{
@@ -159,7 +148,7 @@ function InlineApproval({
       <Box sx={{ display: 'flex', gap: 1 }}>
         <Button
           size="small"
-          onClick={() => onResolve(log.toolCallId || '', false, feedback || 'Rejected by user')}
+          onClick={() => onResolve(toolCallId, false, feedback || 'Rejected by user')}
           sx={{
             flex: 1,
             textTransform: 'none',
@@ -175,7 +164,7 @@ function InlineApproval({
         </Button>
         <Button
           size="small"
-          onClick={() => onResolve(log.toolCallId || '', true)}
+          onClick={() => onResolve(toolCallId, true)}
           sx={{
             flex: 1,
             textTransform: 'none',
@@ -194,94 +183,129 @@ function InlineApproval({
   );
 }
 
-// ── Main component ──────────────────────────────────────────────────
-export default function ToolCallGroup({ tools }: ToolCallGroupProps) {
-  const { showToolOutput, setPanelTab, setActivePanelTab, clearPanelTabs } = useAgentStore();
+// ---------------------------------------------------------------------------
+// Main component
+// ---------------------------------------------------------------------------
+
+export default function ToolCallGroup({ tools, approveTools }: ToolCallGroupProps) {
+  const { setPanelTab, setActivePanelTab, clearPanelTabs } = useAgentStore();
   const { setRightPanelOpen, setLeftSidebarOpen } = useLayoutStore();
-  const { activeSessionId } = useSessionStore();
 
+  // ── Batch approval state ──────────────────────────────────────────
+  const pendingTools = useMemo(
+    () => tools.filter(t => t.state === 'approval-requested'),
+    [tools],
+  );
+  const hasPendingApprovals = pendingTools.length > 0;
+
+  const [decisions, setDecisions] = useState<Record<string, { approved: boolean; feedback?: string }>>({});
+  const [isSubmitting, setIsSubmitting] = useState(false);
+  const submittingRef = useRef(false);
+
+  // ── Send all decisions as a single batch ──────────────────────────
+  const sendBatch = useCallback(
+    async (batch: Record<string, { approved: boolean; feedback?: string }>) => {
+      if (submittingRef.current) return;
+      submittingRef.current = true;
+      setIsSubmitting(true);
+
+      const approvals = Object.entries(batch).map(([toolCallId, d]) => ({
+        tool_call_id: toolCallId,
+        approved: d.approved,
+        feedback: d.approved ? null : (d.feedback || 'Rejected by user'),
+      }));
+
+      const ok = await approveTools(approvals);
+      if (!ok) {
+        logger.error('Batch approval failed');
+        submittingRef.current = false;
+        setIsSubmitting(false);
+      }
+    },
+    [approveTools],
+  );
+
+  const handleApproveAll = useCallback(() => {
+    const batch: Record<string, { approved: boolean }> = {};
+    for (const t of pendingTools) batch[t.toolCallId] = { approved: true };
+    sendBatch(batch);
+  }, [pendingTools, sendBatch]);
+
+  const handleRejectAll = useCallback(() => {
+    const batch: Record<string, { approved: boolean }> = {};
+    for (const t of pendingTools) batch[t.toolCallId] = { approved: false };
+    sendBatch(batch);
+  }, [pendingTools, sendBatch]);
+
+  const handleIndividualDecision = useCallback(
+    (toolCallId: string, approved: boolean, feedback?: string) => {
+      const next = { ...decisions, [toolCallId]: { approved, feedback } };
+      setDecisions(next);
+      if (pendingTools.every(t => next[t.toolCallId])) {
+        sendBatch(next);
+      }
+    },
+    [decisions, pendingTools, sendBatch],
+  );
+
+  const undoDecision = useCallback((toolCallId: string) => {
+    setDecisions(prev => {
+      const next = { ...prev };
+      delete next[toolCallId];
+      return next;
+    });
+  }, []);
+
+  // ── Panel click handler ───────────────────────────────────────────
   const handleClick = useCallback(
-    (log: TraceLog) => {
-      // For hf_jobs with scripts, use tab system
-      if (log.tool === 'hf_jobs' && log.args?.script) {
+    (tool: DynamicToolPart) => {
+      const args = tool.input as Record<string, unknown> | undefined;
+
+      if (tool.toolName === 'hf_jobs' && args?.script) {
         clearPanelTabs();
-        setPanelTab({
-          id: 'script',
-          title: 'Script',
-          content: String(log.args.script),
-          language: 'python',
-        });
-        if (log.output) {
-          setPanelTab({
-            id: 'output',
-            title: 'Output',
-            content: log.output,
-            language: 'markdown',
-          });
+        setPanelTab({ id: 'script', title: 'Script', content: String(args.script), language: 'python' });
+        if (tool.state === 'output-available' && tool.output) {
+          setPanelTab({ id: 'output', title: 'Output', content: String(tool.output), language: 'markdown' });
         }
-        if (log.jobLogs) {
-          setPanelTab({
-            id: 'logs',
-            title: 'Logs',
-            content: log.jobLogs,
-            language: 'text',
-          });
-        }
-        // Default to output if it exists (most useful), otherwise script
-        setActivePanelTab(log.output ? 'output' : 'script');
+        setActivePanelTab(tool.state === 'output-available' ? 'output' : 'script');
         setRightPanelOpen(true);
         setLeftSidebarOpen(false);
         return;
       }
 
-      // Show output if completed/failed, or args if still running
-      const s = resolveState(log);
-      if ((s === 'completed' || s === 'failed') && log.output) {
-        showToolOutput(log);
-      } else if (log.args) {
-        const content = JSON.stringify(log.args, null, 2);
-        showToolOutput({ ...log, output: content });
-      } else {
-        return;
-      }
-      setRightPanelOpen(true);
-    },
-    [showToolOutput, setRightPanelOpen, setLeftSidebarOpen, clearPanelTabs, setPanelTab, setActivePanelTab],
-  );
+      if ((tool.state === 'output-available' || tool.state === 'output-error') && tool.output) {
+        let language = 'text';
+        const content = String(tool.output);
+        if (content.trim().startsWith('{') || content.trim().startsWith('[')) language = 'json';
+        else if (content.includes('```')) language = 'markdown';
 
-  const handleApprovalResolve = useCallback(
-    async (toolCallId: string, approved: boolean, feedback?: string) => {
-      if (!activeSessionId) return;
-      try {
-        const res = await apiFetch('/api/approve', {
-          method: 'POST',
-          body: JSON.stringify({
-            session_id: activeSessionId,
-            approvals: [{
-              tool_call_id: toolCallId,
-              approved,
-              feedback: approved ? null : feedback || 'Rejected by user',
-            }],
-          }),
-        });
-
-        if (res.ok) {
-          // Optimistic update with proper state transitions
-          const { updateTraceLog, updateCurrentTurnTrace, setProcessing } = useAgentStore.getState();
-          updateTraceLog(toolCallId, '', {
-            state: approved ? 'approved' : 'rejected',
-            approvalStatus: approved ? 'approved' : 'rejected', // legacy compat
-          });
-          updateCurrentTurnTrace(activeSessionId);
-          if (approved) setProcessing(true);
-        }
-      } catch (e) {
-        logger.error('Approval failed:', e);
+        useAgentStore.getState().setPanelTab({ id: 'tool_output', title: tool.toolName, content, language });
+        useAgentStore.getState().setActivePanelTab('tool_output');
+        setRightPanelOpen(true);
+      } else if (args) {
+        const content = JSON.stringify(args, null, 2);
+        useAgentStore.getState().setPanelTab({ id: 'tool_output', title: tool.toolName, content, language: 'json' });
+        useAgentStore.getState().setActivePanelTab('tool_output');
+        setRightPanelOpen(true);
       }
     },
-    [activeSessionId],
+    [clearPanelTabs, setPanelTab, setActivePanelTab, setRightPanelOpen, setLeftSidebarOpen],
   );
 
+  // ── Parse hf_jobs metadata from output ────────────────────────────
+  function parseJobMeta(output: unknown): { jobUrl?: string; jobStatus?: string } {
+    if (typeof output !== 'string') return {};
+    const urlMatch = output.match(/\*\*View at:\*\*\s*(https:\/\/[^\s\n]+)/);
+    const statusMatch = output.match(/\*\*Final Status:\*\*\s*([^\n]+)/);
+    return {
+      jobUrl: urlMatch?.[1],
+      jobStatus: statusMatch?.[1]?.trim(),
+    };
+  }
+
+  // ── Render ────────────────────────────────────────────────────────
+  const decidedCount = pendingTools.filter(t => decisions[t.toolCallId]).length;
+
   return (
     <Box
       sx={{
@@ -292,30 +316,100 @@ export default function ToolCallGroup({ tools }: ToolCallGroupProps) {
         my: 1,
       }}
     >
+      {/* Batch approval header */}
+      {hasPendingApprovals && !isSubmitting && (
+        <Box
+          sx={{
+            display: 'flex',
+            alignItems: 'center',
+            gap: 1,
+            px: 1.5,
+            py: 1,
+            borderBottom: '1px solid var(--tool-border)',
+          }}
+        >
+          <Typography
+            variant="body2"
+            sx={{ fontSize: '0.72rem', color: 'var(--muted-text)', mr: 'auto', whiteSpace: 'nowrap' }}
+          >
+            {decidedCount > 0
+              ? `${decidedCount}/${pendingTools.length} decided`
+              : `${pendingTools.length} tool${pendingTools.length > 1 ? 's' : ''} pending`}
+          </Typography>
+          <Button
+            size="small"
+            onClick={handleRejectAll}
+            sx={{
+              textTransform: 'none',
+              color: 'var(--accent-red)',
+              border: '1px solid rgba(255,255,255,0.05)',
+              fontSize: '0.72rem',
+              py: 0.5,
+              px: 1.5,
+              borderRadius: '8px',
+              '&:hover': { bgcolor: 'rgba(224,90,79,0.05)', borderColor: 'var(--accent-red)' },
+            }}
+          >
+            Reject all
+          </Button>
+          <Button
+            size="small"
+            onClick={handleApproveAll}
+            sx={{
+              textTransform: 'none',
+              color: 'var(--accent-green)',
+              border: '1px solid var(--accent-green)',
+              fontSize: '0.72rem',
+              fontWeight: 600,
+              py: 0.5,
+              px: 1.5,
+              borderRadius: '8px',
+              '&:hover': { bgcolor: 'rgba(47,204,113,0.1)' },
+            }}
+          >
+            Approve all{pendingTools.length > 1 ? ` (${pendingTools.length})` : ''}
+          </Button>
+        </Box>
+      )}
+
+      {/* Tool list */}
       <Stack divider={<Box sx={{ borderBottom: '1px solid var(--tool-border)' }} />}>
-        {tools.map((log) => {
-          const state = resolveState(log);
-          const clickable = state === 'completed' || state === 'failed' || !!log.args;
-          const label = statusLabel(state);
-          const isPendingApproval = state === 'pending_approval';
+        {tools.map((tool) => {
+          const state = tool.state;
+          const isPending = state === 'approval-requested';
+          const clickable =
+            state === 'output-available' ||
+            state === 'output-error' ||
+            !!tool.input;
+          const localDecision = decisions[tool.toolCallId];
+
+          const displayState = isPending && localDecision
+            ? (localDecision.approved ? 'input-available' : 'output-denied')
+            : state;
+          const label = statusLabel(displayState as ToolPartState);
+
+          // Parse job metadata from hf_jobs output
+          const jobMeta = tool.toolName === 'hf_jobs' && tool.state === 'output-available'
+            ? parseJobMeta(tool.output)
+            : {};
 
           return (
-            <Box key={log.id}>
+            <Box key={tool.toolCallId}>
               {/* Main tool row */}
               <Stack
                 direction="row"
                 alignItems="center"
                 spacing={1}
-                onClick={() => !isPendingApproval && handleClick(log)}
+                onClick={() => !isPending && handleClick(tool)}
                 sx={{
                   px: 1.5,
                   py: 1,
-                  cursor: isPendingApproval ? 'default' : clickable ? 'pointer' : 'default',
+                  cursor: isPending ? 'default' : clickable ? 'pointer' : 'default',
                   transition: 'background-color 0.15s',
-                  '&:hover': clickable && !isPendingApproval ? { bgcolor: 'var(--hover-bg)' } : {},
+                  '&:hover': clickable && !isPending ? { bgcolor: 'var(--hover-bg)' } : {},
                 }}
               >
-                <StatusIcon state={state} />
+                <StatusIcon state={displayState as ToolPartState} />
 
                 <Typography
                   variant="body2"
@@ -331,7 +425,7 @@ export default function ToolCallGroup({ tools }: ToolCallGroupProps) {
                     whiteSpace: 'nowrap',
                   }}
                 >
-                  {log.tool}
+                  {tool.toolName}
                 </Typography>
 
                 {label && (
@@ -343,19 +437,19 @@ export default function ToolCallGroup({ tools }: ToolCallGroupProps) {
                       fontSize: '0.65rem',
                       fontWeight: 600,
                       bgcolor: 'var(--accent-yellow-weak)',
-                      color: statusColor(state),
+                      color: statusColor(displayState as ToolPartState),
                       letterSpacing: '0.03em',
                     }}
                   />
                 )}
 
-                {clickable && !isPendingApproval && (
+                {clickable && !isPending && (
                   <OpenInNewIcon sx={{ fontSize: 14, color: 'var(--muted-text)', opacity: 0.6 }} />
                 )}
               </Stack>
 
-              {/* Job status + link row */}
-              {(log.jobUrl || log.jobStatus) && (
+              {/* Job status + link */}
+              {(jobMeta.jobUrl || jobMeta.jobStatus) && (
                 <Box
                   sx={{
                     display: 'flex',
@@ -366,21 +460,21 @@ export default function ToolCallGroup({ tools }: ToolCallGroupProps) {
                     borderTop: '1px solid var(--tool-border)',
                   }}
                 >
-                  {log.jobStatus && (
+                  {jobMeta.jobStatus && (
                     <Typography
                       variant="caption"
                       sx={{
-                        color: log.success === false ? 'var(--accent-red)' : 'var(--accent-green)',
+                        color: tool.state === 'output-error' ? 'var(--accent-red)' : 'var(--accent-green)',
                         fontSize: '0.7rem',
                         fontWeight: 600,
                       }}
                     >
-                      {log.jobStatus}
+                      {jobMeta.jobStatus}
                     </Typography>
                   )}
-                  {log.jobUrl && (
+                  {jobMeta.jobUrl && (
                     <Link
-                      href={log.jobUrl}
+                      href={jobMeta.jobUrl}
                       target="_blank"
                       rel="noopener noreferrer"
                       onClick={(e) => e.stopPropagation()}
@@ -401,9 +495,48 @@ export default function ToolCallGroup({ tools }: ToolCallGroupProps) {
                 </Box>
               )}
 
-              {/* Inline approval UI (only when pending) */}
-              {isPendingApproval && (
-                <InlineApproval log={log} onResolve={handleApprovalResolve} />
+              {/* Per-tool approval: undecided */}
+              {isPending && !localDecision && !isSubmitting && (
+                <InlineApproval
+                  toolCallId={tool.toolCallId}
+                  toolName={tool.toolName}
+                  input={tool.input}
+                  onResolve={handleIndividualDecision}
+                />
+              )}
+
+              {/* Per-tool approval: locally decided (undo available) */}
+              {isPending && localDecision && !isSubmitting && (
+                <Box
+                  sx={{
+                    display: 'flex',
+                    alignItems: 'center',
+                    justifyContent: 'space-between',
+                    px: 1.5,
+                    py: 0.75,
+                    borderTop: '1px solid var(--tool-border)',
+                  }}
+                >
+                  <Typography variant="body2" sx={{ fontSize: '0.72rem', color: 'var(--muted-text)' }}>
+                    {localDecision.approved
+                      ? 'Marked for approval'
+                      : `Marked for rejection${localDecision.feedback ? `: ${localDecision.feedback}` : ''}`}
+                  </Typography>
+                  <Button
+                    size="small"
+                    onClick={() => undoDecision(tool.toolCallId)}
+                    sx={{
+                      textTransform: 'none',
+                      fontSize: '0.7rem',
+                      color: 'var(--muted-text)',
+                      minWidth: 'auto',
+                      px: 1,
+                      '&:hover': { color: 'var(--text)' },
+                    }}
+                  >
+                    Undo
+                  </Button>
+                </Box>
               )}
             </Box>
           );

frontend/src/components/Chat/UserMessage.tsx

@@ -1,17 +1,22 @@
 import { Box, Stack, Typography, IconButton, Tooltip } from '@mui/material';
 import CloseIcon from '@mui/icons-material/Close';
-import type { Message } from '@/types/agent';
+import type { UIMessage } from 'ai';
+import type { MessageMeta } from '@/types/agent';
 
 interface UserMessageProps {
-  message: Message;
-  /** True if this message starts the last turn. */
+  message: UIMessage;
   isLastTurn?: boolean;
-  /** Callback to remove the last turn. */
   onUndoTurn?: () => void;
-  /** Whether the agent is currently processing (disables undo). */
   isProcessing?: boolean;
 }
 
+function extractText(message: UIMessage): string {
+  return message.parts
+    .filter((p): p is Extract<typeof p, { type: 'text' }> => p.type === 'text')
+    .map(p => p.text)
+    .join('');
+}
+
 export default function UserMessage({
   message,
   isLastTurn = false,
@@ -19,7 +24,11 @@ export default function UserMessage({
   isProcessing = false,
 }: UserMessageProps) {
   const showUndo = isLastTurn && !isProcessing && !!onUndoTurn;
-
+  const text = extractText(message);
+  const meta = message.metadata as MessageMeta | undefined;
+  const timeStr = meta?.createdAt
+    ? new Date(meta.createdAt).toLocaleTimeString([], { hour: '2-digit', minute: '2-digit' })
+    : null;
   return (
     <Stack
       direction="row"
@@ -27,7 +36,6 @@ export default function UserMessage({
       justifyContent="flex-end"
       alignItems="flex-start"
       sx={{
-        // Show the undo button when hovering the entire row
         '& .undo-btn': {
           opacity: 0,
           transition: 'opacity 0.15s ease',
@@ -37,7 +45,6 @@ export default function UserMessage({
         },
       }}
     >
-      {/* Undo button — visible on hover, left of the bubble */}
       {showUndo && (
         <Box className="undo-btn" sx={{ display: 'flex', alignItems: 'center', mt: 0.75 }}>
           <Tooltip title="Remove this turn" placement="left">
@@ -81,22 +88,17 @@ export default function UserMessage({
             wordBreak: 'break-word',
           }}
         >
-          {message.content}
+          {text}
         </Typography>
 
-        <Typography
-          variant="caption"
-          sx={{
-            display: 'block',
-            textAlign: 'right',
-            mt: 1,
-            fontSize: '0.68rem',
-            color: 'var(--muted-text)',
-            opacity: 0.7,
-          }}
-        >
-          {new Date(message.timestamp).toLocaleTimeString([], { hour: '2-digit', minute: '2-digit' })}
-        </Typography>
+        {timeStr && (
+          <Typography
+            variant="caption"
+            sx={{ color: 'var(--muted-text)', mt: 0.5, display: 'block', textAlign: 'right', fontSize: '0.7rem' }}
+          >
+            {timeStr}
+          </Typography>
+        )}
       </Box>
     </Stack>
   );

frontend/src/components/Layout/AppLayout.tsx

@@ -1,4 +1,4 @@
-import { useCallback, useRef, useEffect } from 'react';
+import { useCallback, useRef, useEffect, useState } from 'react';
 import {
   Avatar,
   Box,
@@ -7,6 +7,7 @@ import {
   IconButton,
   Alert,
   AlertTitle,
+  Snackbar,
   useMediaQuery,
   useTheme,
 } from '@mui/material';
@@ -20,20 +21,19 @@ import { logger } from '@/utils/logger';
 import { useSessionStore } from '@/store/sessionStore';
 import { useAgentStore } from '@/store/agentStore';
 import { useLayoutStore } from '@/store/layoutStore';
-import { useAgentWebSocket } from '@/hooks/useAgentWebSocket';
+import { useAgentChat } from '@/hooks/useAgentChat';
 import SessionSidebar from '@/components/SessionSidebar/SessionSidebar';
 import CodePanel from '@/components/CodePanel/CodePanel';
 import ChatInput from '@/components/Chat/ChatInput';
 import MessageList from '@/components/Chat/MessageList';
 import WelcomeScreen from '@/components/WelcomeScreen/WelcomeScreen';
 import { apiFetch } from '@/utils/api';
-import type { Message } from '@/types/agent';
 
 const DRAWER_WIDTH = 260;
 
 export default function AppLayout() {
   const { sessions, activeSessionId, deleteSession, updateSessionTitle } = useSessionStore();
-  const { isConnected, isProcessing, getMessages, addMessage, setProcessing, llmHealthError, setLlmHealthError, user } = useAgentStore();
+  const { isConnected, isProcessing, setProcessing, llmHealthError, setLlmHealthError, user } = useAgentStore();
   const { 
     isLeftSidebarOpen, 
     isRightPanelOpen, 
@@ -48,6 +48,9 @@ export default function AppLayout() {
   const theme = useTheme();
   const isMobile = useMediaQuery(theme.breakpoints.down('md'));
 
+  const [showExpiredToast, setShowExpiredToast] = useState(false);
+  const disconnectTimer = useRef<ReturnType<typeof setTimeout> | null>(null);
+
   const isResizing = useRef(false);
 
   const handleMouseMove = useCallback((e: MouseEvent) => {
@@ -105,10 +108,9 @@ export default function AppLayout() {
     return () => { cancelled = true; };
   }, []); // eslint-disable-line react-hooks/exhaustive-deps
 
-  const messages = activeSessionId ? getMessages(activeSessionId) : [];
   const hasAnySessions = sessions.length > 0;
 
-  useAgentWebSocket({
+  const { messages, sendMessage, stop, undoLastTurn, approveTools } = useAgentChat({
     sessionId: activeSessionId,
     onReady: () => logger.log('Agent ready'),
     onError: (error) => logger.error('Agent error:', error),
@@ -118,24 +120,29 @@ export default function AppLayout() {
     },
   });
 
+  // Debounced "session expired" toast — only fires after 2s of sustained disconnect
+  useEffect(() => {
+    if (!isConnected && messages.length > 0 && activeSessionId) {
+      disconnectTimer.current = setTimeout(() => setShowExpiredToast(true), 2000);
+    } else {
+      if (disconnectTimer.current) clearTimeout(disconnectTimer.current);
+      disconnectTimer.current = null;
+      setShowExpiredToast(false);
+    }
+    return () => {
+      if (disconnectTimer.current) clearTimeout(disconnectTimer.current);
+    };
+  }, [isConnected, messages.length, activeSessionId]);
+
   const handleSendMessage = useCallback(
     async (text: string) => {
       if (!activeSessionId || !text.trim() || isProcessing) return;
-      
-      // Lock input immediately to prevent double-sends
-      setProcessing(true);
 
-      const userMsg: Message = {
-        id: `user_${Date.now()}`,
-        role: 'user',
-        content: text.trim(),
-        timestamp: new Date().toISOString(),
-      };
-      addMessage(activeSessionId, userMsg);
+      setProcessing(true);
+      sendMessage({ text: text.trim(), metadata: { createdAt: new Date().toISOString() } });
 
       // Auto-title the session from the first user message (async, non-blocking)
-      const currentMessages = getMessages(activeSessionId);
-      const isFirstMessage = currentMessages.filter((m) => m.role === 'user').length <= 1;
+      const isFirstMessage = messages.filter((m) => m.role === 'user').length <= 1;
       if (isFirstMessage) {
         const sessionId = activeSessionId;
         apiFetch('/api/title', {
@@ -151,20 +158,8 @@ export default function AppLayout() {
             updateSessionTitle(sessionId, raw.length > 40 ? raw.slice(0, 40) + '…' : raw);
           });
       }
-
-      try {
-        await apiFetch('/api/submit', {
-          method: 'POST',
-          body: JSON.stringify({
-            session_id: activeSessionId,
-            text: text.trim(),
-          }),
-        });
-      } catch (e) {
-        logger.error('Send failed:', e);
-      }
     },
-    [activeSessionId, addMessage, getMessages, updateSessionTitle, isProcessing, setProcessing]
+    [activeSessionId, sendMessage, messages, updateSessionTitle, isProcessing, setProcessing],
   );
 
   // Close sidebar on mobile after selecting a session
@@ -363,28 +358,12 @@ export default function AppLayout() {
           >
             {activeSessionId ? (
               <>
-                <MessageList messages={messages} isProcessing={isProcessing} />
-                {!isConnected && messages.length > 0 && (
-                  <Box sx={{
-                    display: 'flex',
-                    alignItems: 'center',
-                    justifyContent: 'center',
-                    gap: 1,
-                    py: 1,
-                    px: { xs: 1, md: 2 },
-                    mb: 1,
-                    borderRadius: 'var(--radius-md)',
-                    bgcolor: 'rgba(255, 171, 0, 0.08)',
-                    border: '1px solid rgba(255, 171, 0, 0.2)',
-                  }}>
-                    <Typography variant="body2" sx={{ color: 'var(--accent-yellow)', fontFamily: 'monospace', fontSize: { xs: '0.7rem', md: '0.8rem' } }}>
-                      Session expired — create a new session to continue.
-                    </Typography>
-                  </Box>
-                )}
+                <MessageList messages={messages} isProcessing={isProcessing} approveTools={approveTools} onUndoLastTurn={undoLastTurn} />
                 <ChatInput
                   onSend={handleSendMessage}
+                  onStop={stop}
                   disabled={isProcessing || !isConnected}
+                  isStreaming={isProcessing}
                 />
               </>
             ) : (
@@ -466,6 +445,20 @@ export default function AppLayout() {
           <CodePanel />
         </Drawer>
       )}
+      <Snackbar
+        open={showExpiredToast}
+        anchorOrigin={{ vertical: 'bottom', horizontal: 'center' }}
+        onClose={() => setShowExpiredToast(false)}
+      >
+        <Alert
+          severity="warning"
+          variant="filled"
+          onClose={() => setShowExpiredToast(false)}
+          sx={{ fontFamily: 'monospace', fontSize: '0.8rem' }}
+        >
+          Session expired — create a new session to continue.
+        </Alert>
+      </Snackbar>
     </Box>
   );
 }

frontend/src/hooks/useAgentChat.ts

@@ -0,0 +1,284 @@
+/**
+ * Central hook wiring the Vercel AI SDK's useChat with our custom
+ * WebSocketChatTransport. Replaces the old useAgentWebSocket + agentStore
+ * message management.
+ */
+import { useCallback, useEffect, useMemo, useRef } from 'react';
+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 { apiFetch } from '@/utils/api';
+import { useAgentStore } from '@/store/agentStore';
+import { useSessionStore } from '@/store/sessionStore';
+import { useLayoutStore } from '@/store/layoutStore';
+import { logger } from '@/utils/logger';
+
+interface UseAgentChatOptions {
+  sessionId: string | null;
+  onReady?: () => void;
+  onError?: (error: string) => void;
+  onSessionDead?: (sessionId: string) => void;
+}
+
+export function useAgentChat({ sessionId, onReady, onError, onSessionDead }: UseAgentChatOptions) {
+  const callbacksRef = useRef({ onReady, onError, onSessionDead });
+  callbacksRef.current = { onReady, onError, onSessionDead };
+
+  const {
+    setProcessing,
+    setConnected,
+    setError,
+    setPanelTab,
+    setActivePanelTab,
+    clearPanelTabs,
+  } = useAgentStore();
+
+  const { setRightPanelOpen, setLeftSidebarOpen } = useLayoutStore();
+  const { setSessionActive } = useSessionStore();
+
+  // ── Build side-channel callbacks (stable ref) ────────────────────
+  const sideChannel = useMemo<SideChannelCallbacks>(
+    () => ({
+      onReady: () => {
+        setConnected(true);
+        setProcessing(false);
+        if (sessionId) setSessionActive(sessionId, true);
+        callbacksRef.current.onReady?.();
+      },
+      onShutdown: () => {
+        setConnected(false);
+        setProcessing(false);
+      },
+      onError: (error: string) => {
+        setError(error);
+        setProcessing(false);
+        callbacksRef.current.onError?.(error);
+      },
+      onProcessing: () => {
+        setProcessing(true);
+      },
+      onProcessingDone: () => {
+        setProcessing(false);
+      },
+      onUndoComplete: () => {
+        setProcessing(false);
+        // Remove the last turn (user msg + assistant response) from useChat state
+        const setMsgs = chatActionsRef.current.setMessages;
+        const msgs = chatActionsRef.current.messages;
+        if (setMsgs && msgs.length > 0) {
+          let lastUserIdx = -1;
+          for (let i = msgs.length - 1; i >= 0; i--) {
+            if (msgs[i].role === 'user') { lastUserIdx = i; break; }
+          }
+          const updated = lastUserIdx > 0 ? msgs.slice(0, lastUserIdx) : [];
+          setMsgs(updated);
+          if (sessionId) saveMessages(sessionId, updated);
+        }
+      },
+      onCompacted: (oldTokens: number, newTokens: number) => {
+        logger.log(`Context compacted: ${oldTokens} → ${newTokens} tokens`);
+      },
+      onPlanUpdate: (plan) => {
+        useAgentStore.getState().setPlan(plan as Array<{ id: string; content: string; status: 'pending' | 'in_progress' | 'completed' }>);
+        if (!useLayoutStore.getState().isRightPanelOpen) {
+          setRightPanelOpen(true);
+        }
+      },
+      onToolLog: (tool: string, log: string) => {
+        if (tool === 'hf_jobs') {
+          const currentTabs = useAgentStore.getState().panelTabs;
+          const logsTab = currentTabs.find(t => t.id === 'logs');
+          const newContent = logsTab
+            ? logsTab.content + '\n' + log
+            : '--- Job execution started ---\n' + log;
+
+          setPanelTab({ id: 'logs', title: 'Logs', content: newContent, language: 'text' });
+          setActivePanelTab('logs');
+
+          if (!useLayoutStore.getState().isRightPanelOpen) {
+            setRightPanelOpen(true);
+          }
+        }
+      },
+      onConnectionChange: (connected: boolean) => {
+        setConnected(connected);
+      },
+      onSessionDead: (deadSessionId: string) => {
+        logger.warn(`Session ${deadSessionId} dead, removing`);
+        callbacksRef.current.onSessionDead?.(deadSessionId);
+      },
+      onApprovalRequired: (tools) => {
+        if (!tools.length) return;
+        const firstTool = tools[0];
+        const args = firstTool.arguments as Record<string, string | undefined>;
+
+        clearPanelTabs();
+
+        if (firstTool.tool === 'hf_jobs' && args.script) {
+          setPanelTab({ id: 'script', title: 'Script', content: args.script, language: 'python', parameters: firstTool.arguments });
+          setActivePanelTab('script');
+        } else if (firstTool.tool === 'hf_repo_files' && args.content) {
+          const filename = args.path || 'file';
+          setPanelTab({
+            id: 'content',
+            title: filename.split('/').pop() || 'Content',
+            content: args.content,
+            language: filename.endsWith('.py') ? 'python' : 'text',
+            parameters: firstTool.arguments,
+          });
+          setActivePanelTab('content');
+        } else {
+          setPanelTab({
+            id: 'args',
+            title: firstTool.tool,
+            content: JSON.stringify(firstTool.arguments, null, 2),
+            language: 'json',
+            parameters: firstTool.arguments,
+          });
+          setActivePanelTab('args');
+        }
+
+        setRightPanelOpen(true);
+        setLeftSidebarOpen(false);
+      },
+      onToolCallPanel: (toolName: string, args: Record<string, unknown>) => {
+        if (toolName === 'hf_jobs' && args.operation && args.script) {
+          clearPanelTabs();
+          setPanelTab({
+            id: 'script',
+            title: 'Script',
+            content: String(args.script),
+            language: 'python',
+            parameters: args,
+          });
+          setActivePanelTab('script');
+          setRightPanelOpen(true);
+          setLeftSidebarOpen(false);
+        } else if (toolName === 'hf_repo_files' && args.operation === 'upload' && args.content) {
+          useAgentStore.getState().setPanelContent({
+            title: `File Upload: ${String(args.path || 'unnamed')}`,
+            content: String(args.content),
+            parameters: args,
+            language: String(args.path || '').endsWith('.py') ? 'python' : undefined,
+          });
+          setRightPanelOpen(true);
+          setLeftSidebarOpen(false);
+        }
+      },
+      onToolOutputPanel: (toolName: string, _toolCallId: string, output: string, success: boolean) => {
+        if (toolName === 'hf_jobs' && output) {
+          setPanelTab({ id: 'output', title: 'Output', content: output, language: 'markdown' });
+          if (!success) setActivePanelTab('output');
+        }
+      },
+    }),
+    // Zustand setters are stable
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+    [sessionId],
+  );
+
+  // ── Create transport (single stable instance for the lifetime of this hook) ──
+  const transportRef = useRef<WebSocketChatTransport | null>(null);
+  if (!transportRef.current) {
+    transportRef.current = new WebSocketChatTransport({ sideChannel });
+  }
+
+  // Keep side-channel callbacks in sync (they capture sessionId)
+  useEffect(() => {
+    transportRef.current?.updateSideChannel(sideChannel);
+  }, [sideChannel]);
+
+  // Connect / disconnect WebSocket when session changes
+  useEffect(() => {
+    transportRef.current?.connectToSession(sessionId);
+    return () => {
+      transportRef.current?.connectToSession(null);
+    };
+  }, [sessionId]);
+
+  // ── Restore persisted messages for this session ─────────────────
+  const initialMessages = useMemo(
+    () => (sessionId ? loadMessages(sessionId) : []),
+    [sessionId],
+  );
+
+  // ── Ref for chat actions (used by sideChannel callbacks created before chat) ──
+  const chatActionsRef = useRef<{
+    setMessages: ((msgs: UIMessage[]) => void) | null;
+    messages: UIMessage[];
+  }>({ setMessages: null, messages: [] });
+
+  // ── useChat from Vercel AI SDK ───────────────────────────────────
+  const chat = useChat({
+    id: sessionId || '__no_session__',
+    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);
+      setProcessing(false);
+    },
+  });
+
+  // Keep chatActionsRef in sync every render
+  chatActionsRef.current.setMessages = chat.setMessages;
+  chatActionsRef.current.messages = chat.messages;
+
+  // ── Persist messages on every user send (onFinish covers assistant turns) ──
+  const prevLenRef = useRef(initialMessages.length);
+  useEffect(() => {
+    if (!sessionId || chat.messages.length === 0) return;
+    if (chat.messages.length !== prevLenRef.current) {
+      prevLenRef.current = chat.messages.length;
+      saveMessages(sessionId, chat.messages);
+    }
+  }, [sessionId, chat.messages]);
+
+  // ── Undo last turn (calls backend + syncs useChat + localStorage) ──
+  const undoLastTurn = useCallback(async () => {
+    if (!sessionId) return;
+    try {
+      const res = await apiFetch(`/api/undo/${sessionId}`, { method: 'POST' });
+      if (!res.ok) {
+        logger.error('Undo API returned', res.status);
+        return;
+      }
+    } catch (e) {
+      logger.error('Undo failed:', e);
+    }
+    // Backend will also send undo_complete, but we apply optimistically
+    // so the UI updates immediately.
+  }, [sessionId]);
+
+  // ── Convenience: approve tools via transport ─────────────────────
+  const approveTools = useCallback(
+    async (approvals: Array<{ tool_call_id: string; approved: boolean; feedback?: string | null }>) => {
+      if (!sessionId || !transportRef.current) return false;
+      const ok = await transportRef.current.approveTools(sessionId, approvals);
+      if (ok) {
+        const hasApproved = approvals.some(a => a.approved);
+        if (hasApproved) setProcessing(true);
+      }
+      return ok;
+    },
+    [sessionId, setProcessing],
+  );
+
+  return {
+    messages: chat.messages,
+    sendMessage: chat.sendMessage,
+    stop: chat.stop,
+    status: chat.status,
+    undoLastTurn,
+    approveTools,
+    transport: transportRef.current,
+  };
+}

frontend/src/hooks/useAgentWebSocket.ts

@@ -1,636 +0,0 @@
-import { useCallback, useEffect, useRef } from 'react';
-import { useAgentStore, type PlanItem } from '@/store/agentStore';
-import { useSessionStore } from '@/store/sessionStore';
-import { useLayoutStore } from '@/store/layoutStore';
-import { getWebSocketUrl } from '@/utils/api';
-import { logger } from '@/utils/logger';
-import type { AgentEvent } from '@/types/events';
-import type { Message, TraceLog } from '@/types/agent';
-
-const WS_RECONNECT_DELAY = 1000;
-const WS_MAX_RECONNECT_DELAY = 30000;
-const WS_MAX_RETRIES = 5;
-
-interface UseAgentWebSocketOptions {
-  sessionId: string | null;
-  onReady?: () => void;
-  onError?: (error: string) => void;
-  onSessionDead?: (sessionId: string) => void;
-}
-
-export function useAgentWebSocket({
-  sessionId,
-  onReady,
-  onError,
-  onSessionDead,
-}: UseAgentWebSocketOptions) {
-  const wsRef = useRef<WebSocket | null>(null);
-  const reconnectTimeoutRef = useRef<number | null>(null);
-  const reconnectDelayRef = useRef(WS_RECONNECT_DELAY);
-  const retriesRef = useRef(0);
-
-  const {
-    addMessage,
-    updateMessage,
-    appendToMessage,
-    setProcessing,
-    setConnected,
-    setError,
-    addTraceLog,
-    updateTraceLog,
-    clearTraceLogs,
-    setPanelContent,
-    setPanelTab,
-    setActivePanelTab,
-    clearPanelTabs,
-    setPlan,
-    setCurrentTurnMessageId,
-    updateCurrentTurnTrace,
-    removeLastTurn,
-  } = useAgentStore();
-
-  const { setRightPanelOpen, setLeftSidebarOpen } = useLayoutStore();
-
-  const { setSessionActive } = useSessionStore();
-
-  const handleEvent = useCallback(
-    (event: AgentEvent) => {
-      if (!sessionId) return;
-
-      switch (event.event_type) {
-        case 'ready':
-          setConnected(true);
-          setProcessing(false);
-          setSessionActive(sessionId, true);
-          onReady?.();
-          break;
-
-        case 'processing':
-          setProcessing(true);
-          clearTraceLogs();
-          // Don't clear panel tabs here - they should persist during approval flow
-          // Tabs will be cleared when a new tool_call sets up new content
-          setCurrentTurnMessageId(null); // Start a new turn
-          break;
-
-        // ── Streaming: individual token chunks ──────────────────
-        case 'assistant_chunk': {
-          const delta = (event.data?.content as string) || '';
-          if (!delta) break;
-
-          const currentTurnMsgId = useAgentStore.getState().currentTurnMessageId;
-
-          if (currentTurnMsgId) {
-            // Append delta to the existing streaming message
-            appendToMessage(sessionId, currentTurnMsgId, delta);
-          } else {
-            // First chunk — create the message (with pending traces if any)
-            const currentTrace = useAgentStore.getState().traceLogs;
-            const messageId = `msg_${Date.now()}`;
-            const segments: Array<{ type: 'text' | 'tools'; content?: string; tools?: typeof currentTrace }> = [];
-
-            if (currentTrace.length > 0) {
-              segments.push({ type: 'tools', tools: [...currentTrace] });
-              clearTraceLogs();
-            }
-            segments.push({ type: 'text', content: delta });
-
-            const message: Message = {
-              id: messageId,
-              role: 'assistant',
-              content: delta,
-              timestamp: new Date().toISOString(),
-              segments,
-            };
-            addMessage(sessionId, message);
-            setCurrentTurnMessageId(messageId);
-          }
-          break;
-        }
-
-        // ── Streaming ended (text is already rendered via chunks) ─
-        case 'assistant_stream_end':
-          // Nothing to do — chunks already built the message.
-          // This event is just a signal that the stream is complete.
-          break;
-
-        // ── Legacy non-streaming full message (kept for backwards compat)
-        case 'assistant_message': {
-          const content = (event.data?.content as string) || '';
-          const currentTrace = useAgentStore.getState().traceLogs;
-          const currentTurnMsgId = useAgentStore.getState().currentTurnMessageId;
-
-          if (currentTurnMsgId) {
-            // Update existing message - add segments chronologically
-            const messages = useAgentStore.getState().getMessages(sessionId);
-            const existingMsg = messages.find(m => m.id === currentTurnMsgId);
-
-            if (existingMsg) {
-              const segments = existingMsg.segments ? [...existingMsg.segments] : [];
-
-              // If there are pending traces, add them as a tools segment first
-              if (currentTrace.length > 0) {
-                segments.push({ type: 'tools', tools: [...currentTrace] });
-                clearTraceLogs();
-              }
-
-              // Add the new text segment
-              if (content) {
-                segments.push({ type: 'text', content });
-              }
-
-              updateMessage(sessionId, currentTurnMsgId, {
-                content: existingMsg.content + '\n\n' + content,
-                segments,
-              });
-            }
-          } else {
-            // Create new message
-            const messageId = `msg_${Date.now()}`;
-            const segments: Array<{ type: 'text' | 'tools'; content?: string; tools?: typeof currentTrace }> = [];
-
-            // Add any pending traces first
-            if (currentTrace.length > 0) {
-              segments.push({ type: 'tools', tools: [...currentTrace] });
-              clearTraceLogs();
-            }
-
-            // Add the text
-            if (content) {
-              segments.push({ type: 'text', content });
-            }
-
-            const message: Message = {
-              id: messageId,
-              role: 'assistant',
-              content,
-              timestamp: new Date().toISOString(),
-              segments,
-            };
-            addMessage(sessionId, message);
-            setCurrentTurnMessageId(messageId);
-          }
-          break;
-        }
-
-        case 'tool_call': {
-          const toolName = (event.data?.tool as string) || 'unknown';
-          const toolCallId = (event.data?.tool_call_id as string) || '';
-          const args = (event.data?.arguments as Record<string, string | undefined>) || {};
-
-          // Don't display plan_tool in trace logs (it shows up elsewhere in the UI)
-          if (toolName !== 'plan_tool') {
-            const log: TraceLog = {
-              id: `tool_${Date.now()}_${toolCallId}`,
-              toolCallId,
-              type: 'call',
-              text: `Agent is executing ${toolName}...`,
-              tool: toolName,
-              timestamp: new Date().toISOString(),
-              state: 'running',
-              completed: false, // legacy compat
-              args,
-            };
-            addTraceLog(log);
-
-            // If no assistant message exists for this turn, create one now
-            // so the ToolCallGroup renders immediately in the chat flow.
-            const currentTurnMsgId = useAgentStore.getState().currentTurnMessageId;
-            if (!currentTurnMsgId) {
-              const messageId = `msg_${Date.now()}`;
-              const currentTrace = useAgentStore.getState().traceLogs;
-              addMessage(sessionId, {
-                id: messageId,
-                role: 'assistant',
-                content: '',
-                timestamp: new Date().toISOString(),
-                segments: [{ type: 'tools', tools: [...currentTrace] }],
-              });
-              setCurrentTurnMessageId(messageId);
-              clearTraceLogs();
-            } else {
-              updateCurrentTurnTrace(sessionId);
-            }
-          }
-
-          // Auto-expand Right Panel for specific tools
-          if (toolName === 'hf_jobs' && (args.operation === 'run' || args.operation === 'scheduled run') && args.script) {
-            // Clear any existing tabs from previous jobs before setting new script
-            clearPanelTabs();
-            // Use tab system for jobs - add script tab immediately
-            setPanelTab({
-              id: 'script',
-              title: 'Script',
-              content: args.script,
-              language: 'python',
-              parameters: args
-            });
-            setActivePanelTab('script');
-            setRightPanelOpen(true);
-            setLeftSidebarOpen(false);
-          } else if (toolName === 'hf_repo_files' && args.operation === 'upload' && args.content) {
-            setPanelContent({
-              title: `File Upload: ${args.path || 'unnamed'}`,
-              content: args.content,
-              parameters: args,
-              language: args.path?.endsWith('.py') ? 'python' : undefined
-            });
-            setRightPanelOpen(true);
-            setLeftSidebarOpen(false);
-          }
-
-          logger.log('Tool call:', toolName, args);
-          break;
-        }
-
-        case 'tool_output': {
-          const toolName = (event.data?.tool as string) || 'unknown';
-          const toolCallId = (event.data?.tool_call_id as string) || '';
-          const output = (event.data?.output as string) || '';
-          const success = event.data?.success as boolean;
-
-          // Mark the tool as completed/failed with its output
-          updateTraceLog(toolCallId, toolName, {
-            state: success ? 'completed' : 'failed',
-            completed: true, // legacy
-            output,
-            success,
-            approvalStatus: 'approved', // legacy: if we got output, it was approved/auto
-          });
-          updateCurrentTurnTrace(sessionId);
-
-          // For hf_jobs: parse job output and enrich the TraceLog with job info
-          if (toolName === 'hf_jobs' && output) {
-            const updates: Partial<TraceLog> = { approvalStatus: 'approved' as const };
-
-            // Parse job URL
-            const urlMatch = output.match(/\*\*View at:\*\*\s*(https:\/\/[^\s\n]+)/);
-            if (urlMatch) updates.jobUrl = urlMatch[1];
-
-            // Parse job status
-            const statusMatch = output.match(/\*\*Final Status:\*\*\s*([^\n]+)/);
-            if (statusMatch) updates.jobStatus = statusMatch[1].trim();
-
-            // Parse logs
-            if (output.includes('**Logs:**')) {
-              const parts = output.split('**Logs:**');
-              if (parts.length > 1) {
-                const codeBlockMatch = parts[1].trim().match(/```([\s\S]*?)```/);
-                if (codeBlockMatch) updates.jobLogs = codeBlockMatch[1].trim();
-              }
-            }
-
-            updateTraceLog(toolCallId, toolName, updates);
-            updateCurrentTurnTrace(sessionId);
-
-            // Add output tab so the user can see results (especially errors)
-            setPanelTab({
-              id: 'output',
-              title: 'Output',
-              content: output,
-              language: 'markdown',
-            });
-            // Auto-switch to output tab on failure so errors are immediately visible
-            if (!success) {
-              setActivePanelTab('output');
-            }
-          }
-
-          // Don't create message bubbles for tool outputs - they only show in trace logs
-          logger.log('Tool output:', toolName, success);
-          break;
-        }
-
-        case 'tool_log': {
-          const toolName = (event.data?.tool as string) || 'unknown';
-          const log = (event.data?.log as string) || '';
-
-          if (toolName === 'hf_jobs') {
-            const currentTabs = useAgentStore.getState().panelTabs;
-            const logsTab = currentTabs.find(t => t.id === 'logs');
-
-            // Append to existing logs tab or create new one
-            const newContent = logsTab
-              ? logsTab.content + '\n' + log
-              : '--- Job execution started ---\n' + log;
-
-            setPanelTab({
-              id: 'logs',
-              title: 'Logs',
-              content: newContent,
-              language: 'text'
-            });
-
-            // Auto-switch to logs tab when logs start streaming
-            setActivePanelTab('logs');
-
-            if (!useLayoutStore.getState().isRightPanelOpen) {
-              setRightPanelOpen(true);
-            }
-          }
-          break;
-        }
-
-        case 'plan_update': {
-          const plan = (event.data?.plan as PlanItem[]) || [];
-          setPlan(plan);
-          if (!useLayoutStore.getState().isRightPanelOpen) {
-            setRightPanelOpen(true);
-          }
-          break;
-        }
-
-        case 'approval_required': {
-          const tools = event.data?.tools as Array<{
-            tool: string;
-            arguments: Record<string, unknown>;
-            tool_call_id: string;
-          }>;
-
-          // Create or update trace logs for approval tools.
-          // The backend only sends tool_call events for non-approval tools,
-          // so we must create TraceLogs here for approval-requiring tools.
-          if (tools) {
-            for (const t of tools) {
-              // Check if a TraceLog already exists (shouldn't, but be safe)
-              const existing = useAgentStore.getState().traceLogs.find(
-                (log) => log.toolCallId === t.tool_call_id
-              );
-              if (!existing) {
-                addTraceLog({
-                  id: `tool_${Date.now()}_${t.tool_call_id}`,
-                  toolCallId: t.tool_call_id,
-                  type: 'call',
-                  text: `Approval required for ${t.tool}`,
-                  tool: t.tool,
-                  timestamp: new Date().toISOString(),
-                  state: 'pending_approval',
-                  completed: false, // legacy
-                  args: t.arguments as Record<string, unknown>,
-                  approvalStatus: 'pending', // legacy
-                });
-              } else {
-                updateTraceLog(t.tool_call_id, t.tool, {
-                  state: 'pending_approval',
-                  approvalStatus: 'pending', // legacy
-                  args: t.arguments as Record<string, unknown>,
-                });
-              }
-            }
-
-            // Ensure there's a message to render the approval UI in
-            const currentTurnMsgId = useAgentStore.getState().currentTurnMessageId;
-            if (!currentTurnMsgId) {
-              const messageId = `msg_${Date.now()}`;
-              const currentTrace = useAgentStore.getState().traceLogs;
-              addMessage(sessionId, {
-                id: messageId,
-                role: 'assistant',
-                content: '',
-                timestamp: new Date().toISOString(),
-                segments: [{ type: 'tools', tools: [...currentTrace] }],
-              });
-              setCurrentTurnMessageId(messageId);
-              clearTraceLogs();
-            } else {
-              updateCurrentTurnTrace(sessionId);
-            }
-          }
-
-          // Show the first tool's content in the panel
-          if (tools && tools.length > 0) {
-            const firstTool = tools[0];
-            const args = firstTool.arguments as Record<string, string | undefined>;
-
-            clearPanelTabs();
-
-            if (firstTool.tool === 'hf_jobs' && args.script) {
-              setPanelTab({
-                id: 'script',
-                title: 'Script',
-                content: args.script,
-                language: 'python',
-                parameters: args
-              });
-              setActivePanelTab('script');
-            } else if (firstTool.tool === 'hf_repo_files' && args.content) {
-              const filename = args.path || 'file';
-              const isPython = filename.endsWith('.py');
-              setPanelTab({
-                id: 'content',
-                title: filename.split('/').pop() || 'Content',
-                content: args.content,
-                language: isPython ? 'python' : 'text',
-                parameters: args
-              });
-              setActivePanelTab('content');
-            } else {
-              setPanelTab({
-                id: 'args',
-                title: firstTool.tool,
-                content: JSON.stringify(args, null, 2),
-                language: 'json',
-                parameters: args
-              });
-              setActivePanelTab('args');
-            }
-
-            setRightPanelOpen(true);
-            setLeftSidebarOpen(false);
-          }
-
-          setProcessing(false);
-          break;
-        }
-
-        // ── Tool state change (sent by backend after approval decisions) ──
-        case 'tool_state_change': {
-          const tcId = (event.data?.tool_call_id as string) || '';
-          const tcTool = (event.data?.tool as string) || '';
-          const newState = (event.data?.state as string) || '';
-
-          if (tcId && newState) {
-            updateTraceLog(tcId, tcTool, {
-              state: newState as import('@/types/agent').ToolState,
-              // Legacy compat
-              ...(newState === 'approved' ? { approvalStatus: 'approved' as const } : {}),
-              ...(newState === 'rejected' ? { approvalStatus: 'rejected' as const, completed: true } : {}),
-            });
-            if (sessionId) updateCurrentTurnTrace(sessionId);
-          }
-          break;
-        }
-
-        case 'turn_complete':
-          setProcessing(false);
-          setCurrentTurnMessageId(null);
-          break;
-
-        case 'compacted': {
-          const oldTokens = event.data?.old_tokens as number;
-          const newTokens = event.data?.new_tokens as number;
-          logger.log(`Context compacted: ${oldTokens} -> ${newTokens} tokens`);
-          break;
-        }
-
-        case 'error': {
-          const errorMsg = (event.data?.error as string) || 'Unknown error';
-          setError(errorMsg);
-          setProcessing(false);
-          onError?.(errorMsg);
-          break;
-        }
-
-        case 'shutdown':
-          setConnected(false);
-          setProcessing(false);
-          break;
-
-        case 'interrupted':
-          setProcessing(false);
-          break;
-
-        case 'undo_complete':
-          if (sessionId) {
-            removeLastTurn(sessionId);
-          }
-          setProcessing(false);
-          break;
-
-        default:
-          logger.log('Unknown event:', event);
-      }
-    },
-    // Zustand setters are stable, so we don't need them in deps
-    // eslint-disable-next-line react-hooks/exhaustive-deps
-    [sessionId, onReady, onError, onSessionDead]
-  );
-
-  const connect = useCallback(() => {
-    if (!sessionId) return;
-    
-    // Don't connect if already connected or connecting
-    if (wsRef.current?.readyState === WebSocket.OPEN || 
-        wsRef.current?.readyState === WebSocket.CONNECTING) {
-      return;
-    }
-
-    // Build WebSocket URL (centralized in utils/api.ts)
-    const wsUrl = getWebSocketUrl(sessionId);
-
-    logger.log('Connecting to WebSocket:', wsUrl);
-    const ws = new WebSocket(wsUrl);
-
-    ws.onopen = () => {
-      logger.log('WebSocket connected');
-      setConnected(true);
-      reconnectDelayRef.current = WS_RECONNECT_DELAY;
-      retriesRef.current = 0; // Reset retry counter on successful connect
-    };
-
-    ws.onmessage = (event) => {
-      try {
-        const data = JSON.parse(event.data) as AgentEvent;
-        handleEvent(data);
-      } catch (e) {
-        logger.error('Failed to parse WebSocket message:', e);
-      }
-    };
-
-    ws.onerror = (error) => {
-      logger.error('WebSocket error:', error);
-    };
-
-    ws.onclose = (event) => {
-      logger.log('WebSocket closed', event.code, event.reason);
-      setConnected(false);
-
-      // Don't reconnect if:
-      // - Normal closure (1000)
-      // - Session not found (4004) — session was deleted or backend restarted
-      // - Auth failed (4001) or access denied (4003) — won't succeed on retry
-      // - No session ID
-      const noRetryCodes = [1000, 4001, 4003, 4004];
-      if (!noRetryCodes.includes(event.code) && sessionId) {
-        retriesRef.current += 1;
-        if (retriesRef.current > WS_MAX_RETRIES) {
-          logger.warn(`WebSocket: max retries (${WS_MAX_RETRIES}) reached, giving up.`);
-          onSessionDead?.(sessionId);
-          return;
-        }
-        // Attempt to reconnect with exponential backoff
-        if (reconnectTimeoutRef.current) {
-          clearTimeout(reconnectTimeoutRef.current);
-        }
-        reconnectTimeoutRef.current = window.setTimeout(() => {
-          reconnectDelayRef.current = Math.min(
-            reconnectDelayRef.current * 2,
-            WS_MAX_RECONNECT_DELAY
-          );
-          connect();
-        }, reconnectDelayRef.current);
-      } else if (event.code === 4004 && sessionId) {
-        // Session not found — remove it from the store (lazy cleanup)
-        logger.warn(`Session ${sessionId} no longer exists on backend, removing.`);
-        onSessionDead?.(sessionId);
-      } else if (noRetryCodes.includes(event.code) && event.code !== 1000) {
-        logger.warn(`WebSocket permanently closed: ${event.code} ${event.reason}`);
-      }
-    };
-
-    wsRef.current = ws;
-  }, [sessionId, handleEvent]);
-
-  const disconnect = useCallback(() => {
-    if (reconnectTimeoutRef.current) {
-      clearTimeout(reconnectTimeoutRef.current);
-      reconnectTimeoutRef.current = null;
-    }
-    if (wsRef.current) {
-      wsRef.current.close();
-      wsRef.current = null;
-    }
-    setConnected(false);
-  }, []);
-
-  const sendPing = useCallback(() => {
-    if (wsRef.current?.readyState === WebSocket.OPEN) {
-      wsRef.current.send(JSON.stringify({ type: 'ping' }));
-    }
-  }, []);
-
-  // Connect when sessionId changes (with a small delay to ensure session is ready)
-  useEffect(() => {
-    if (!sessionId) {
-      disconnect();
-      return;
-    }
-
-    // Reset retry state for new session
-    retriesRef.current = 0;
-    reconnectDelayRef.current = WS_RECONNECT_DELAY;
-
-    // Small delay to ensure session is fully created on backend
-    const timeoutId = setTimeout(() => {
-      connect();
-    }, 100);
-
-    return () => {
-      clearTimeout(timeoutId);
-      disconnect();
-    };
-    // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [sessionId]);
-
-  // Heartbeat
-  useEffect(() => {
-    const interval = setInterval(sendPing, 30000);
-    return () => clearInterval(interval);
-  }, [sendPing]);
-
-  return {
-    isConnected: wsRef.current?.readyState === WebSocket.OPEN,
-    connect,
-    disconnect,
-  };
-}

frontend/src/lib/chat-message-store.ts

@@ -0,0 +1,62 @@
+/**
+ * Lightweight localStorage persistence for UIMessage arrays,
+ * keyed by session ID.
+ *
+ * Uses the same storage namespace (`hf-agent-messages`) that the
+ * old Zustand-based store used, so existing data is compatible.
+ */
+import type { UIMessage } from 'ai';
+import { logger } from '@/utils/logger';
+
+const STORAGE_KEY = 'hf-agent-messages';
+const MAX_SESSIONS = 50;
+
+type MessagesMap = Record<string, UIMessage[]>;
+
+function readAll(): MessagesMap {
+  try {
+    const raw = localStorage.getItem(STORAGE_KEY);
+    if (!raw) return {};
+    const parsed = JSON.parse(raw);
+    // Legacy format was { messagesBySession: {...} }
+    if (parsed.messagesBySession) return parsed.messagesBySession;
+    // New flat format
+    if (typeof parsed === 'object' && !Array.isArray(parsed)) return parsed;
+    return {};
+  } catch {
+    return {};
+  }
+}
+
+function writeAll(map: MessagesMap): void {
+  try {
+    localStorage.setItem(STORAGE_KEY, JSON.stringify(map));
+  } catch (e) {
+    logger.warn('Failed to persist messages:', e);
+  }
+}
+
+export function loadMessages(sessionId: string): UIMessage[] {
+  const map = readAll();
+  return map[sessionId] ?? [];
+}
+
+export function saveMessages(sessionId: string, messages: UIMessage[]): void {
+  const map = readAll();
+  map[sessionId] = messages;
+
+  // Evict oldest sessions if we exceed the cap
+  const keys = Object.keys(map);
+  if (keys.length > MAX_SESSIONS) {
+    const toRemove = keys.slice(0, keys.length - MAX_SESSIONS);
+    for (const k of toRemove) delete map[k];
+  }
+
+  writeAll(map);
+}
+
+export function deleteMessages(sessionId: string): void {
+  const map = readAll();
+  delete map[sessionId];
+  writeAll(map);
+}

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

@@ -0,0 +1,526 @@
+/**
+ * Custom ChatTransport that bridges our WebSocket-based backend protocol
+ * to the Vercel AI SDK's UIMessageChunk streaming interface.
+ *
+ * The backend stays unchanged — this adapter translates WebSocket events
+ * into the chunk types that useChat() expects.
+ */
+import type { ChatTransport, UIMessage, UIMessageChunk, ChatRequestOptions } from 'ai';
+import { apiFetch, getWebSocketUrl } from '@/utils/api';
+import { logger } from '@/utils/logger';
+import type { AgentEvent } from '@/types/events';
+
+// ---------------------------------------------------------------------------
+// Side-channel callback interface (non-chat events forwarded to the store)
+// ---------------------------------------------------------------------------
+export interface SideChannelCallbacks {
+  onReady: () => void;
+  onShutdown: () => void;
+  onError: (error: string) => void;
+  onProcessing: () => void;
+  onProcessingDone: () => void;
+  onUndoComplete: () => void;
+  onCompacted: (oldTokens: number, newTokens: number) => void;
+  onPlanUpdate: (plan: Array<{ id: string; content: string; status: string }>) => void;
+  onToolLog: (tool: string, log: string) => void;
+  onConnectionChange: (connected: boolean) => void;
+  onSessionDead: (sessionId: string) => void;
+  /** Called when approval_required arrives — lets the store manage panels */
+  onApprovalRequired: (tools: Array<{ tool: string; arguments: Record<string, unknown>; tool_call_id: string }>) => void;
+  /** Called when a tool_call arrives with panel-relevant args */
+  onToolCallPanel: (tool: string, args: Record<string, unknown>) => void;
+  /** Called when tool_output arrives with panel-relevant data */
+  onToolOutputPanel: (tool: string, toolCallId: string, output: string, success: boolean) => void;
+}
+
+// ---------------------------------------------------------------------------
+// Transport options
+// ---------------------------------------------------------------------------
+export interface WebSocketChatTransportOptions {
+  sideChannel: SideChannelCallbacks;
+}
+
+// ---------------------------------------------------------------------------
+// Constants
+// ---------------------------------------------------------------------------
+const WS_RECONNECT_DELAY = 1000;
+const WS_MAX_RECONNECT_DELAY = 30000;
+const WS_MAX_RETRIES = 5;
+const WS_PING_INTERVAL = 30000;
+
+let partIdCounter = 0;
+function nextPartId(prefix: string): string {
+  return `${prefix}-${Date.now()}-${++partIdCounter}`;
+}
+
+// ---------------------------------------------------------------------------
+// Transport implementation
+// ---------------------------------------------------------------------------
+export class WebSocketChatTransport implements ChatTransport<UIMessage> {
+  private ws: WebSocket | null = null;
+  private currentSessionId: string | null = null;
+  private sideChannel: SideChannelCallbacks;
+
+  private streamController: ReadableStreamDefaultController<UIMessageChunk> | null = null;
+  private streamGeneration = 0;
+  private textPartId: string | null = null;
+  private awaitingProcessing = false;
+
+  private connectTimeout: ReturnType<typeof setTimeout> | null = null;
+  private reconnectTimeout: ReturnType<typeof setTimeout> | null = null;
+  private reconnectDelay = WS_RECONNECT_DELAY;
+  private retries = 0;
+  private pingInterval: ReturnType<typeof setInterval> | null = null;
+
+  constructor({ sideChannel }: WebSocketChatTransportOptions) {
+    this.sideChannel = sideChannel;
+  }
+
+  /** Update side-channel callbacks (e.g. when sessionId changes). */
+  updateSideChannel(sideChannel: SideChannelCallbacks): void {
+    this.sideChannel = sideChannel;
+  }
+
+  // ── Public API ──────────────────────────────────────────────────────
+
+  /** Connect (or reconnect) to a session's WebSocket. */
+  connectToSession(sessionId: string | null): void {
+    if (this.connectTimeout) {
+      clearTimeout(this.connectTimeout);
+      this.connectTimeout = null;
+    }
+    this.disconnectWebSocket();
+    this.currentSessionId = sessionId;
+    if (sessionId) {
+      this.retries = 0;
+      this.reconnectDelay = WS_RECONNECT_DELAY;
+      this.connectTimeout = setTimeout(() => {
+        this.connectTimeout = null;
+        if (this.currentSessionId === sessionId) {
+          this.createWebSocket(sessionId);
+        }
+      }, 100);
+    }
+  }
+
+  /** Approve / reject tools. Called directly from the UI. */
+  async approveTools(
+    sessionId: string,
+    approvals: Array<{ tool_call_id: string; approved: boolean; feedback?: string | null }>,
+  ): Promise<boolean> {
+    try {
+      const res = await apiFetch('/api/approve', {
+        method: 'POST',
+        body: JSON.stringify({ session_id: sessionId, approvals }),
+      });
+      return res.ok;
+    } catch (e) {
+      logger.error('Approval request failed:', e);
+      return false;
+    }
+  }
+
+  /** Clean up everything. */
+  destroy(): void {
+    if (this.connectTimeout) {
+      clearTimeout(this.connectTimeout);
+      this.connectTimeout = null;
+    }
+    this.disconnectWebSocket();
+    this.closeActiveStream();
+  }
+
+  // ── ChatTransport interface ─────────────────────────────────────────
+
+  async sendMessages(
+    options: {
+      trigger: 'submit-message' | 'regenerate-message';
+      chatId: string;
+      messageId: string | undefined;
+      messages: UIMessage[];
+      abortSignal: AbortSignal | undefined;
+    } & ChatRequestOptions,
+  ): Promise<ReadableStream<UIMessageChunk>> {
+    const sessionId = options.chatId;
+
+    // Close any previously active stream (e.g. user sent new msg during approval)
+    this.closeActiveStream();
+    this.awaitingProcessing = false;
+
+    // Track generation to protect against late cancel from a stale stream
+    const gen = ++this.streamGeneration;
+
+    // Wire up abort signal to interrupt the backend and close the stream
+    if (options.abortSignal) {
+      const onAbort = () => {
+        if (this.streamGeneration !== gen) return;
+        logger.log('Stream aborted by user');
+        this.interruptBackend(sessionId);
+        this.endTextPart();
+        if (this.streamController) {
+          this.enqueue({ type: 'finish-step' });
+          this.enqueue({ type: 'finish', finishReason: 'stop' });
+          this.closeActiveStream();
+        }
+        this.awaitingProcessing = true;
+        this.sideChannel.onProcessingDone();
+      };
+      if (options.abortSignal.aborted) {
+        onAbort();
+      } else {
+        options.abortSignal.addEventListener('abort', onAbort, { once: true });
+      }
+    }
+
+    // Create the stream BEFORE the POST so WebSocket events arriving
+    // while the HTTP request is in-flight are captured immediately.
+    const stream = new ReadableStream<UIMessageChunk>({
+      start: (controller) => {
+        this.streamController = controller;
+        this.textPartId = null;
+      },
+      cancel: () => {
+        if (this.streamGeneration === gen) {
+          this.streamController = null;
+          this.textPartId = null;
+        }
+      },
+    });
+
+    // Extract the latest user text from the messages array
+    const lastUserMsg = [...options.messages].reverse().find(m => m.role === 'user');
+    const text = lastUserMsg
+      ? lastUserMsg.parts
+          .filter((p): p is Extract<typeof p, { type: 'text' }> => p.type === 'text')
+          .map(p => p.text)
+          .join('')
+      : '';
+
+    // POST to the existing backend endpoint
+    try {
+      await apiFetch('/api/submit', {
+        method: 'POST',
+        body: JSON.stringify({ session_id: sessionId, text }),
+      });
+    } catch (e) {
+      logger.error('Submit failed:', e);
+      this.enqueue({ type: 'error', errorText: 'Failed to send message' });
+      this.closeActiveStream();
+    }
+
+    return stream;
+  }
+
+  async reconnectToStream(): Promise<ReadableStream<UIMessageChunk> | null> {
+    return null;
+  }
+
+  /** Ask the backend to interrupt the current generation. Fire-and-forget. */
+  private interruptBackend(sessionId: string): void {
+    apiFetch(`/api/interrupt/${sessionId}`, { method: 'POST' }).catch((e) =>
+      logger.warn('Interrupt request failed:', e),
+    );
+  }
+
+  // ── WebSocket lifecycle ─────────────────────────────────────────────
+
+  private createWebSocket(sessionId: string): void {
+    if (this.ws?.readyState === WebSocket.OPEN || this.ws?.readyState === WebSocket.CONNECTING) {
+      return;
+    }
+
+    const wsUrl = getWebSocketUrl(sessionId);
+    logger.log('WS transport connecting:', wsUrl);
+    const ws = new WebSocket(wsUrl);
+
+    ws.onopen = () => {
+      logger.log('WS transport connected');
+      this.sideChannel.onConnectionChange(true);
+      this.reconnectDelay = WS_RECONNECT_DELAY;
+      this.retries = 0;
+      this.startPing();
+    };
+
+    ws.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);
+      }
+    };
+
+    ws.onerror = (err) => logger.error('WS error:', err);
+
+    ws.onclose = (evt) => {
+      logger.log('WS closed', evt.code, evt.reason);
+      this.sideChannel.onConnectionChange(false);
+      this.stopPing();
+
+      const noRetry = [1000, 4001, 4003, 4004];
+      if (evt.code === 4004 && sessionId) {
+        this.sideChannel.onSessionDead(sessionId);
+        return;
+      }
+      if (!noRetry.includes(evt.code) && this.currentSessionId === sessionId) {
+        this.retries += 1;
+        if (this.retries > WS_MAX_RETRIES) {
+          logger.warn('WS max retries reached');
+          this.sideChannel.onSessionDead(sessionId);
+          return;
+        }
+        this.reconnectTimeout = setTimeout(() => {
+          this.reconnectDelay = Math.min(this.reconnectDelay * 2, WS_MAX_RECONNECT_DELAY);
+          this.createWebSocket(sessionId);
+        }, this.reconnectDelay);
+      }
+    };
+
+    this.ws = ws;
+  }
+
+  private disconnectWebSocket(): void {
+    if (this.reconnectTimeout) {
+      clearTimeout(this.reconnectTimeout);
+      this.reconnectTimeout = null;
+    }
+    this.stopPing();
+    if (this.ws) {
+      this.ws.close();
+      this.ws = null;
+    }
+    this.sideChannel.onConnectionChange(false);
+  }
+
+  private startPing(): void {
+    this.stopPing();
+    this.pingInterval = setInterval(() => {
+      if (this.ws?.readyState === WebSocket.OPEN) {
+        this.ws.send(JSON.stringify({ type: 'ping' }));
+      }
+    }, WS_PING_INTERVAL);
+  }
+
+  private stopPing(): void {
+    if (this.pingInterval) {
+      clearInterval(this.pingInterval);
+      this.pingInterval = null;
+    }
+  }
+
+  // ── Stream helpers ──────────────────────────────────────────────────
+
+  private closeActiveStream(): void {
+    if (this.streamController) {
+      try {
+        this.streamController.close();
+      } catch {
+        // already closed
+      }
+      this.streamController = null;
+      this.textPartId = null;
+    }
+  }
+
+  private enqueue(chunk: UIMessageChunk): void {
+    try {
+      this.streamController?.enqueue(chunk);
+    } catch {
+      // stream already closed
+    }
+  }
+
+  private endTextPart(): void {
+    if (this.textPartId) {
+      this.enqueue({ type: 'text-end', id: this.textPartId });
+      this.textPartId = null;
+    }
+  }
+
+  // ── Event → UIMessageChunk mapping ──────────────────────────────────
+
+  private static readonly STREAM_EVENTS = new Set([
+    'assistant_chunk', 'assistant_stream_end', 'assistant_message',
+    'tool_call', 'tool_output', 'approval_required', 'tool_state_change',
+    'turn_complete', 'error',
+  ]);
+
+  private handleEvent(event: AgentEvent): void {
+    // After an abort, ignore stale stream events until the next 'processing'
+    if (this.awaitingProcessing && WebSocketChatTransport.STREAM_EVENTS.has(event.event_type)) {
+      return;
+    }
+
+    switch (event.event_type) {
+      // ── Side-channel only events ────────────────────────────────
+      case 'ready':
+        this.sideChannel.onReady();
+        break;
+
+      case 'shutdown':
+        this.sideChannel.onShutdown();
+        this.closeActiveStream();
+        break;
+
+      case 'interrupted':
+        this.sideChannel.onProcessingDone();
+        this.closeActiveStream();
+        break;
+
+      case 'undo_complete':
+        this.endTextPart();
+        this.closeActiveStream();
+        this.sideChannel.onUndoComplete();
+        break;
+
+      case 'compacted':
+        this.sideChannel.onCompacted(
+          (event.data?.old_tokens as number) || 0,
+          (event.data?.new_tokens as number) || 0,
+        );
+        break;
+
+      case 'plan_update':
+        this.sideChannel.onPlanUpdate(
+          (event.data?.plan as Array<{ id: string; content: string; status: string }>) || [],
+        );
+        break;
+
+      case 'tool_log':
+        this.sideChannel.onToolLog(
+          (event.data?.tool as string) || '',
+          (event.data?.log as string) || '',
+        );
+        break;
+
+      // ── Chat stream events ──────────────────────────────────────
+      case 'processing':
+        this.awaitingProcessing = false;
+        this.sideChannel.onProcessing();
+        if (this.streamController) {
+          this.enqueue({
+            type: 'start',
+            messageMetadata: { createdAt: new Date().toISOString() },
+          });
+          this.enqueue({ type: 'start-step' });
+        }
+        break;
+
+      case 'assistant_chunk': {
+        const delta = (event.data?.content as string) || '';
+        if (!delta || !this.streamController) break;
+
+        if (!this.textPartId) {
+          this.textPartId = nextPartId('text');
+          this.enqueue({ type: 'text-start', id: this.textPartId });
+        }
+        this.enqueue({ type: 'text-delta', id: this.textPartId, delta });
+        break;
+      }
+
+      case 'assistant_stream_end':
+        this.endTextPart();
+        break;
+
+      case 'assistant_message': {
+        const content = (event.data?.content as string) || '';
+        if (!content || !this.streamController) break;
+        const id = nextPartId('text');
+        this.enqueue({ type: 'text-start', id });
+        this.enqueue({ type: 'text-delta', id, delta: content });
+        this.enqueue({ type: 'text-end', id });
+        break;
+      }
+
+      case 'tool_call': {
+        if (!this.streamController) break;
+        const toolName = (event.data?.tool as string) || 'unknown';
+        const toolCallId = (event.data?.tool_call_id as string) || '';
+        const args = (event.data?.arguments as Record<string, unknown>) || {};
+
+        if (toolName === 'plan_tool') break;
+
+        this.endTextPart();
+        this.enqueue({ type: 'tool-input-start', toolCallId, toolName, dynamic: true });
+        this.enqueue({ type: 'tool-input-available', toolCallId, toolName, input: args, dynamic: true });
+
+        this.sideChannel.onToolCallPanel(toolName, args as Record<string, unknown>);
+        break;
+      }
+
+      case 'tool_output': {
+        if (!this.streamController) break;
+        const toolCallId = (event.data?.tool_call_id as string) || '';
+        const output = (event.data?.output as string) || '';
+        const success = event.data?.success as boolean;
+        const toolName = (event.data?.tool as string) || '';
+
+        if (success) {
+          this.enqueue({ type: 'tool-output-available', toolCallId, output, dynamic: true });
+        } else {
+          this.enqueue({ type: 'tool-output-error', toolCallId, errorText: output, dynamic: true });
+        }
+
+        this.sideChannel.onToolOutputPanel(toolName, toolCallId, output, success);
+        break;
+      }
+
+      case 'approval_required': {
+        const tools = event.data?.tools as Array<{
+          tool: string;
+          arguments: Record<string, unknown>;
+          tool_call_id: string;
+        }>;
+        if (!tools || !this.streamController) break;
+
+        this.endTextPart();
+
+        for (const t of tools) {
+          this.enqueue({ type: 'tool-input-start', toolCallId: t.tool_call_id, toolName: t.tool, dynamic: true });
+          this.enqueue({ type: 'tool-input-available', toolCallId: t.tool_call_id, toolName: t.tool, input: t.arguments, dynamic: true });
+          this.enqueue({ type: 'tool-approval-request', approvalId: `approval-${t.tool_call_id}`, toolCallId: t.tool_call_id });
+        }
+
+        this.sideChannel.onApprovalRequired(tools);
+        this.sideChannel.onProcessingDone();
+        break;
+      }
+
+      case 'tool_state_change': {
+        if (!this.streamController) break;
+        const tcId = (event.data?.tool_call_id as string) || '';
+        const state = (event.data?.state as string) || '';
+
+        if (state === 'rejected' || state === 'abandoned') {
+          this.enqueue({ type: 'tool-output-denied', toolCallId: tcId });
+        }
+        // 'approved' and 'running' are transient — results will follow via tool_output
+        break;
+      }
+
+      case 'turn_complete':
+        this.endTextPart();
+        if (this.streamController) {
+          this.enqueue({ type: 'finish-step' });
+          this.enqueue({ type: 'finish', finishReason: 'stop' });
+          this.closeActiveStream();
+        }
+        this.sideChannel.onProcessingDone();
+        break;
+
+      case 'error': {
+        const errorMsg = (event.data?.error as string) || 'Unknown error';
+        this.sideChannel.onError(errorMsg);
+        if (this.streamController) {
+          this.enqueue({ type: 'error', errorText: errorMsg });
+        }
+        this.sideChannel.onProcessingDone();
+        break;
+      }
+
+      default:
+        logger.log('WS transport: unknown event', event);
+    }
+  }
+}

frontend/src/store/agentStore.ts

@@ -1,6 +1,16 @@
+/**
+ * Agent store — manages UI state that is NOT handled by the Vercel AI SDK.
+ *
+ * Message state (messages, streaming, tool calls) is now managed by useChat().
+ * This store only handles:
+ *  - Connection / processing flags
+ *  - Panel state (right panel tabs)
+ *  - Plan state
+ *  - User info / error banners
+ *  - Edited scripts (for hf_jobs code editing)
+ */
 import { create } from 'zustand';
-import { persist } from 'zustand/middleware';
-import type { Message, User, TraceLog } from '@/types/agent';
+import type { User } from '@/types/agent';
 
 export interface PlanItem {
   id: string;
@@ -23,408 +33,140 @@ export interface LLMHealthError {
 }
 
 interface AgentStore {
-  // State per session (keyed by session ID)
-  messagesBySession: Record<string, Message[]>;
+  // Global UI flags
   isProcessing: boolean;
   isConnected: boolean;
   user: User | null;
   error: string | null;
   llmHealthError: LLMHealthError | null;
-  traceLogs: TraceLog[];
+
+  // Right panel
   panelContent: { title: string; content: string; language?: string; parameters?: Record<string, unknown> } | null;
   panelTabs: PanelTab[];
   activePanelTab: string | null;
+
+  // Plan
   plan: PlanItem[];
-  currentTurnMessageId: string | null; // Track the current turn's assistant message
-  editedScripts: Record<string, string>; // tool_call_id -> edited content
+
+  // Edited scripts (tool_call_id -> edited content)
+  editedScripts: Record<string, string>;
 
   // Actions
-  addMessage: (sessionId: string, message: Message) => void;
-  updateMessage: (sessionId: string, messageId: string, updates: Partial<Message>) => void;
-  clearMessages: (sessionId: string) => void;
   setProcessing: (isProcessing: boolean) => void;
   setConnected: (isConnected: boolean) => void;
   setUser: (user: User | null) => void;
   setError: (error: string | null) => void;
-  getMessages: (sessionId: string) => Message[];
-  addTraceLog: (log: TraceLog) => void;
-  updateTraceLog: (toolCallId: string, toolName: string, updates: Partial<TraceLog>) => void;
-  clearTraceLogs: () => void;
+  setLlmHealthError: (error: LLMHealthError | null) => void;
+
   setPanelContent: (content: { title: string; content: string; language?: string; parameters?: Record<string, unknown> } | null) => void;
   setPanelTab: (tab: PanelTab) => void;
   updatePanelTabContent: (tabId: string, content: string) => void;
   setActivePanelTab: (tabId: string) => void;
   clearPanelTabs: () => void;
   removePanelTab: (tabId: string) => void;
+  showToolOutput: (output: { tool: string; output?: string; args?: Record<string, unknown> }) => void;
+
   setPlan: (plan: PlanItem[]) => void;
-  setCurrentTurnMessageId: (id: string | null) => void;
-  updateCurrentTurnTrace: (sessionId: string) => void;
-  showToolOutput: (log: TraceLog) => void;
+
   setEditedScript: (toolCallId: string, content: string) => void;
   getEditedScript: (toolCallId: string) => string | undefined;
   clearEditedScripts: () => void;
-  /** Append a streaming delta to an existing message. */
-  appendToMessage: (sessionId: string, messageId: string, delta: string) => void;
-  /** Remove all messages for a session (also clears from localStorage). */
-  deleteSessionMessages: (sessionId: string) => void;
-  /** Remove the last turn (last user msg + all following assistant/tool msgs). */
-  removeLastTurn: (sessionId: string) => void;
-  setLlmHealthError: (error: LLMHealthError | null) => void;
 }
 
-export const useAgentStore = create<AgentStore>()(
-  persist(
-  (set, get) => ({
-  messagesBySession: {},
+export const useAgentStore = create<AgentStore>()((set, get) => ({
   isProcessing: false,
   isConnected: false,
   user: null,
   error: null,
   llmHealthError: null,
-  traceLogs: [],
+
   panelContent: null,
   panelTabs: [],
   activePanelTab: null,
-  plan: [],
-  currentTurnMessageId: null,
-  editedScripts: {},
 
-  addMessage: (sessionId: string, message: Message) => {
-    set((state) => {
-      const currentMessages = state.messagesBySession[sessionId] || [];
-      return {
-        messagesBySession: {
-          ...state.messagesBySession,
-          [sessionId]: [...currentMessages, message],
-        },
-      };
-    });
-  },
-
-  updateMessage: (sessionId: string, messageId: string, updates: Partial<Message>) => {
-    set((state) => {
-      const currentMessages = state.messagesBySession[sessionId] || [];
-      const updatedMessages = currentMessages.map((msg) =>
-        msg.id === messageId ? { ...msg, ...updates } : msg
-      );
-      return {
-        messagesBySession: {
-          ...state.messagesBySession,
-          [sessionId]: updatedMessages,
-        },
-      };
-    });
-  },
-
-  clearMessages: (sessionId: string) => {
-    set((state) => ({
-      messagesBySession: {
-        ...state.messagesBySession,
-        [sessionId]: [],
-      },
-    }));
-  },
-
-  setProcessing: (isProcessing: boolean) => {
-    set({ isProcessing });
-  },
-
-  setConnected: (isConnected: boolean) => {
-    set({ isConnected });
-  },
-
-  setUser: (user: User | null) => {
-    set({ user });
-  },
-
-  setError: (error: string | null) => {
-    set({ error });
-  },
+  plan: [],
 
-  getMessages: (sessionId: string) => {
-    return get().messagesBySession[sessionId] || [];
-  },
+  editedScripts: {},
 
-  addTraceLog: (log: TraceLog) => {
-    set((state) => ({
-      traceLogs: [...state.traceLogs, log],
-    }));
-  },
+  // ── Global flags ──────────────────────────────────────────────────
 
-  updateTraceLog: (toolCallId: string, toolName: string, updates: Partial<TraceLog>) => {
-    set((state) => {
-      const traceLogs = [...state.traceLogs];
-      // Prefer matching by tool_call_id (reliable), fall back to tool name (legacy)
-      let matched = false;
-      if (toolCallId) {
-        for (let i = traceLogs.length - 1; i >= 0; i--) {
-          if (traceLogs[i].toolCallId === toolCallId) {
-            traceLogs[i] = { ...traceLogs[i], ...updates };
-            matched = true;
-            break;
-          }
-        }
-      }
-      if (!matched) {
-        // Fallback: match by tool name (last uncompleted call)
-        for (let i = traceLogs.length - 1; i >= 0; i--) {
-          if (traceLogs[i].tool === toolName && traceLogs[i].type === 'call' && !traceLogs[i].completed) {
-            traceLogs[i] = { ...traceLogs[i], ...updates };
-            break;
-          }
-        }
-      }
-      return { traceLogs };
-    });
-  },
+  setProcessing: (isProcessing) => set({ isProcessing }),
+  setConnected: (isConnected) => set({ isConnected }),
+  setUser: (user) => set({ user }),
+  setError: (error) => set({ error }),
+  setLlmHealthError: (error) => set({ llmHealthError: error }),
 
-  clearTraceLogs: () => {
-    set({ traceLogs: [] });
-  },
+  // ── Panel ─────────────────────────────────────────────────────────
 
-  setPanelContent: (content) => {
-    set({ panelContent: content });
-  },
+  setPanelContent: (content) => set({ panelContent: content }),
 
   setPanelTab: (tab: PanelTab) => {
     set((state) => {
-      const existingIndex = state.panelTabs.findIndex(t => t.id === tab.id);
+      const idx = state.panelTabs.findIndex(t => t.id === tab.id);
       let newTabs: PanelTab[];
-      if (existingIndex >= 0) {
-        // Update existing tab
+      if (idx >= 0) {
         newTabs = [...state.panelTabs];
-        newTabs[existingIndex] = tab;
+        newTabs[idx] = tab;
       } else {
-        // Add new tab
         newTabs = [...state.panelTabs, tab];
       }
       return {
         panelTabs: newTabs,
-        activePanelTab: state.activePanelTab || tab.id, // Auto-select first tab
+        activePanelTab: state.activePanelTab || tab.id,
       };
     });
   },
 
-  updatePanelTabContent: (tabId: string, content: string) => {
-    set((state) => {
-      const newTabs = state.panelTabs.map(tab =>
+  updatePanelTabContent: (tabId, content) => {
+    set((state) => ({
+      panelTabs: state.panelTabs.map(tab =>
         tab.id === tabId ? { ...tab, content } : tab
-      );
-      return { panelTabs: newTabs };
-    });
+      ),
+    }));
   },
 
-  setActivePanelTab: (tabId: string) => {
-    set({ activePanelTab: tabId });
-  },
+  setActivePanelTab: (tabId) => set({ activePanelTab: tabId }),
 
-  clearPanelTabs: () => {
-    set({ panelTabs: [], activePanelTab: null });
-  },
+  clearPanelTabs: () => set({ panelTabs: [], activePanelTab: null }),
 
-  removePanelTab: (tabId: string) => {
+  removePanelTab: (tabId) => {
     set((state) => {
       const newTabs = state.panelTabs.filter(t => t.id !== tabId);
-      // If we removed the active tab, switch to another tab or null
       let newActiveTab = state.activePanelTab;
       if (state.activePanelTab === tabId) {
         newActiveTab = newTabs.length > 0 ? newTabs[newTabs.length - 1].id : null;
       }
-      return {
-        panelTabs: newTabs,
-        activePanelTab: newActiveTab,
-      };
+      return { panelTabs: newTabs, activePanelTab: newActiveTab };
     });
   },
 
-  setPlan: (plan: PlanItem[]) => {
-    set({ plan });
-  },
-
-  setCurrentTurnMessageId: (id: string | null) => {
-    set({ currentTurnMessageId: id });
-  },
+  showToolOutput: (info) => {
+    const content = info.output || (info.args ? JSON.stringify(info.args, null, 2) : 'No output available');
+    let language = 'text';
+    if (content.trim().startsWith('{') || content.trim().startsWith('[')) language = 'json';
+    else if (content.includes('```')) language = 'markdown';
 
-  updateCurrentTurnTrace: (sessionId: string) => {
     const state = get();
-    if (!state.currentTurnMessageId) return;
-
-    const currentMessages = state.messagesBySession[sessionId] || [];
-    const latestTools = state.traceLogs.length > 0 ? [...state.traceLogs] : undefined;
-    if (!latestTools) return;
-
-    // Build a lookup of the latest state for each tool by id
-    const toolById = new Map(latestTools.map(t => [t.id, t]));
-
-    const updatedMessages = currentMessages.map((msg) => {
-      if (msg.id !== state.currentTurnMessageId) return msg;
-
-      const segments = msg.segments ? [...msg.segments] : [];
-
-      // First pass: update existing tools in their original segments
-      const placedToolIds = new Set<string>();
-      for (let i = 0; i < segments.length; i++) {
-        if (segments[i].type === 'tools' && segments[i].tools) {
-          segments[i] = {
-            ...segments[i],
-            tools: segments[i].tools!.map(t => {
-              placedToolIds.add(t.id);
-              return toolById.get(t.id) || t;
-            }),
-          };
-        }
-      }
-
-      // Collect only genuinely new tools (not yet in any segment)
-      const newTools = latestTools.filter(t => !placedToolIds.has(t.id));
-
-      if (newTools.length > 0) {
-        const lastToolsIdx = segments.map((s) => s.type).lastIndexOf('tools');
-
-        if (lastToolsIdx >= 0 && lastToolsIdx === segments.length - 1) {
-          // Last segment is tools — append new tools to it
-          segments[lastToolsIdx] = {
-            ...segments[lastToolsIdx],
-            tools: [...(segments[lastToolsIdx].tools || []), ...newTools],
-          };
-        } else {
-          // Text came after previous tools — create a new segment with only new tools
-          segments.push({ type: 'tools', tools: newTools });
-        }
-      }
-
-      return { ...msg, segments };
-    });
-
+    const otherTabs = state.panelTabs.filter(t => t.id !== 'tool_output');
     set({
-      messagesBySession: {
-        ...state.messagesBySession,
-        [sessionId]: updatedMessages,
-      },
+      panelTabs: [...otherTabs, { id: 'tool_output', title: info.tool, content, language }],
+      activePanelTab: 'tool_output',
     });
   },
 
-  showToolOutput: (log: TraceLog) => {
-    // Show tool output in the right panel - only ONE tool output tab at a time
-    const state = get();
+  // ── Plan ──────────────────────────────────────────────────────────
 
-    // Determine language based on content
-    let language = 'text';
-    const content = log.output || '';
-
-    // Check if content looks like JSON
-    if (content.trim().startsWith('{') || content.trim().startsWith('[') || content.includes('```json')) {
-      language = 'json';
-    }
-    // Check if content has markdown tables or formatting
-    else if (content.includes('|') && content.includes('---') || content.includes('```')) {
-      language = 'markdown';
-    }
-
-    // Remove any existing tool output tab (only keep one)
-    const otherTabs = state.panelTabs.filter(t => t.id !== 'tool_output');
+  setPlan: (plan) => set({ plan }),
 
-    // Create/replace the single tool output tab
-    const newTab = {
-      id: 'tool_output',
-      title: log.tool,
-      content: content || 'No output available',
-      language,
-    };
+  // ── Edited scripts ────────────────────────────────────────────────
 
-    set({
-      panelTabs: [...otherTabs, newTab],
-      activePanelTab: 'tool_output',
-    });
-  },
-
-  setEditedScript: (toolCallId: string, content: string) => {
+  setEditedScript: (toolCallId, content) => {
     set((state) => ({
       editedScripts: { ...state.editedScripts, [toolCallId]: content },
     }));
   },
 
-  getEditedScript: (toolCallId: string) => {
-    return get().editedScripts[toolCallId];
-  },
-
-  clearEditedScripts: () => {
-    set({ editedScripts: {} });
-  },
-
-  appendToMessage: (sessionId: string, messageId: string, delta: string) => {
-    set((state) => {
-      const messages = state.messagesBySession[sessionId] || [];
-      return {
-        messagesBySession: {
-          ...state.messagesBySession,
-          [sessionId]: messages.map((msg) => {
-            if (msg.id !== messageId) return msg;
-            const newContent = msg.content + delta;
-            const segments = msg.segments ? [...msg.segments] : [];
-            const lastSeg = segments[segments.length - 1];
-
-            if (lastSeg && lastSeg.type === 'text') {
-              // Append to the existing text segment
-              segments[segments.length - 1] = {
-                ...lastSeg,
-                content: (lastSeg.content || '') + delta,
-              };
-            } else {
-              // Last segment is 'tools' (or empty) — start a NEW text segment
-              // so that tools and text remain visually separated.
-              segments.push({ type: 'text', content: delta });
-            }
-
-            return { ...msg, content: newContent, segments };
-          }),
-        },
-      };
-    });
-  },
-
-  deleteSessionMessages: (sessionId: string) => {
-    set((state) => {
-      const { [sessionId]: _, ...rest } = state.messagesBySession;
-      return { messagesBySession: rest };
-    });
-  },
-
-  removeLastTurn: (sessionId: string) => {
-    set((state) => {
-      const msgs = state.messagesBySession[sessionId];
-      if (!msgs || msgs.length === 0) return state;
-
-      // Find the index of the last user message
-      let lastUserIdx = -1;
-      for (let i = msgs.length - 1; i >= 0; i--) {
-        if (msgs[i].role === 'user') {
-          lastUserIdx = i;
-          break;
-        }
-      }
-      if (lastUserIdx === -1) return state; // no user message to remove
+  getEditedScript: (toolCallId) => get().editedScripts[toolCallId],
 
-      // Remove everything from that user message onward
-      return {
-        messagesBySession: {
-          ...state.messagesBySession,
-          [sessionId]: msgs.slice(0, lastUserIdx),
-        },
-      };
-    });
-  },
-
-  setLlmHealthError: (error: LLMHealthError | null) => {
-    set({ llmHealthError: error });
-  },
-}),
-    {
-      name: 'hf-agent-messages',
-      // Only persist messages — all transient UI state stays in-memory
-      partialize: (state) => ({
-        messagesBySession: state.messagesBySession,
-      }),
-    }
-  )
-);
+  clearEditedScripts: () => set({ editedScripts: {} }),
+}));

frontend/src/store/sessionStore.ts

@@ -1,7 +1,7 @@
 import { create } from 'zustand';
 import { persist } from 'zustand/middleware';
 import type { SessionMeta } from '@/types/agent';
-import { useAgentStore } from './agentStore';
+import { deleteMessages } from '@/lib/chat-message-store';
 
 interface SessionStore {
   sessions: SessionMeta[];
@@ -35,9 +35,7 @@ export const useSessionStore = create<SessionStore>()(
       },
 
       deleteSession: (id: string) => {
-        // Clean up persisted messages for this session
-        useAgentStore.getState().deleteSessionMessages(id);
-
+        deleteMessages(id);
         set((state) => {
           const newSessions = state.sessions.filter((s) => s.id !== id);
           const newActiveId =

frontend/src/types/agent.ts

@@ -1,7 +1,15 @@
 /**
- * Agent-related types
+ * Agent-related types.
+ *
+ * Message and tool-call types are now provided by the Vercel AI SDK
+ * (UIMessage, UIMessagePart, etc.). Only non-SDK types remain here.
  */
 
+/** Custom metadata attached to every UIMessage via the `metadata` field. */
+export interface MessageMeta {
+  createdAt?: string;
+}
+
 export interface SessionMeta {
   id: string;
   title: string;
@@ -9,87 +17,12 @@ export interface SessionMeta {
   isActive: boolean;
 }
 
-export interface MessageSegment {
-  type: 'text' | 'tools';
-  content?: string;
-  tools?: TraceLog[];
-}
-
-export interface Message {
-  id: string;
-  role: 'user' | 'assistant' | 'tool';
-  content: string;
-  timestamp: string;
-  segments?: MessageSegment[];
-  approval?: {
-    status: 'pending' | 'approved' | 'rejected';
-    batch: ApprovalBatch;
-    decisions?: ToolApproval[];
-  };
-  toolOutput?: string;
-}
-
-export interface ToolCall {
-  id: string;
-  tool: string;
-  arguments: Record<string, unknown>;
-  status: 'pending' | 'running' | 'completed' | 'failed';
-  output?: string;
-}
-
 export interface ToolApproval {
   tool_call_id: string;
   approved: boolean;
   feedback?: string | null;
 }
 
-export interface ApprovalBatch {
-  tools: Array<{
-    tool: string;
-    arguments: Record<string, unknown>;
-    tool_call_id: string;
-  }>;
-  count: number;
-}
-
-/**
- * Single state field for each tool call lifecycle.
- * Follows the Vercel AI SDK pattern: clear, non-overlapping states.
- */
-export type ToolState =
-  | 'calling'           // tool_call received, execution starting
-  | 'pending_approval'  // waiting for user to approve/reject
-  | 'approved'          // user approved, waiting for execution to start
-  | 'running'           // execution in progress
-  | 'completed'         // execution finished successfully
-  | 'failed'            // execution finished with error
-  | 'rejected'          // user rejected the tool call
-  | 'timed_out';        // no response after timeout
-
-// Keep backward compat alias
-export type ApprovalStatus = 'none' | 'pending' | 'approved' | 'rejected';
-
-export interface TraceLog {
-  id: string;
-  toolCallId?: string;
-  type: 'call' | 'output';
-  text: string;
-  tool: string;
-  timestamp: string;
-  /** Single source of truth for tool lifecycle state */
-  state: ToolState;
-  args?: Record<string, unknown>;
-  output?: string;
-  success?: boolean;
-  // Legacy fields — kept for backward compat with persisted data
-  completed?: boolean;
-  approvalStatus?: ApprovalStatus;
-  /** Parsed job info (URL, status, logs) for hf_jobs */
-  jobUrl?: string;
-  jobStatus?: string;
-  jobLogs?: string;
-}
-
 export interface User {
   authenticated: boolean;
   username?: string;