replace websocket transport with SSE

backend/dependencies.py

@@ -1,6 +1,5 @@
 """Authentication dependencies for FastAPI routes.
 
-Provides auth validation for both REST and WebSocket endpoints.
 - In dev mode (OAUTH_CLIENT_ID not set): auth is bypassed, returns a default "dev" user.
 - In production: validates Bearer tokens or cookies against HF OAuth.
 """
@@ -11,7 +10,7 @@ import time
 from typing import Any
 
 import httpx
-from fastapi import HTTPException, Request, WebSocket, status
+from fastapi import HTTPException, Request, status
 
 logger = logging.getLogger(__name__)
 
@@ -142,31 +141,3 @@ async def get_current_user(request: Request) -> dict[str, Any]:
     )
 
 
-async def get_ws_user(websocket: WebSocket) -> dict[str, Any] | None:
-    """Extract and validate user from WebSocket connection.
-
-    WebSocket doesn't support custom headers from browser, so we check:
-    1. ?token= query parameter
-    2. hf_access_token cookie (sent automatically for same-origin)
-
-    Returns user dict or None if not authenticated.
-    In dev mode, returns the default dev user.
-    """
-    if not AUTH_ENABLED:
-        return DEV_USER
-
-    # Try query param
-    token = websocket.query_params.get("token")
-    if token:
-        user = await _extract_user_from_token(token)
-        if user:
-            return user
-
-    # Try cookie (works for same-origin WebSocket)
-    token = websocket.cookies.get("hf_access_token")
-    if token:
-        user = await _extract_user_from_token(token)
-        if user:
-            return user
-
-    return None

backend/routes/agent.py

@@ -1,25 +1,22 @@
-"""Agent API routes - WebSocket and REST endpoints.
+"""Agent API routes — REST + SSE endpoints.
 
 All routes (except /health) require authentication via the get_current_user
 dependency. In dev mode (no OAUTH_CLIENT_ID), auth is bypassed automatically.
 """
 
+import json
 import logging
-import os
 from typing import Any
 
-from dependencies import get_current_user, get_ws_user
+from dependencies import get_current_user
 from fastapi import (
     APIRouter,
     Depends,
     HTTPException,
     Request,
-    WebSocket,
-    WebSocketDisconnect,
 )
+from fastapi.responses import StreamingResponse
 from litellm import acompletion
-
-from agent.core.agent_loop import _resolve_hf_router_params
 from models import (
     ApprovalRequest,
     HealthResponse,
@@ -29,7 +26,8 @@ from models import (
     SubmitRequest,
 )
 from session_manager import MAX_SESSIONS, SessionCapacityError, session_manager
-from websocket import manager as ws_manager
+
+from agent.core.agent_loop import _resolve_hf_router_params
 
 logger = logging.getLogger(__name__)
 
@@ -286,6 +284,84 @@ async def submit_approval(
     return {"status": "submitted", "session_id": request.session_id}
 
 
+@router.post("/chat/{session_id}")
+async def chat_sse(
+    session_id: str,
+    request: Request,
+    user: dict = Depends(get_current_user),
+) -> StreamingResponse:
+    """SSE endpoint: submit input or approval, then stream events until turn ends."""
+    _check_session_access(session_id, user)
+
+    agent_session = session_manager.sessions.get(session_id)
+    if not agent_session or not agent_session.is_active:
+        raise HTTPException(status_code=404, detail="Session not found or inactive")
+
+    # Parse body
+    body = await request.json()
+
+    # Subscribe BEFORE submitting so we never miss events — even if the
+    # agent loop processes the submission before this coroutine continues.
+    broadcaster = agent_session.broadcaster
+    sub_id, event_queue = broadcaster.subscribe()
+
+    # Submit the operation
+    text = body.get("text")
+    approvals = body.get("approvals")
+
+    try:
+        if approvals:
+            formatted = [
+                {
+                    "tool_call_id": a["tool_call_id"],
+                    "approved": a["approved"],
+                    "feedback": a.get("feedback"),
+                    "edited_script": a.get("edited_script"),
+                }
+                for a in approvals
+            ]
+            success = await session_manager.submit_approval(session_id, formatted)
+        elif text is not None:
+            success = await session_manager.submit_user_input(session_id, text)
+        else:
+            broadcaster.unsubscribe(sub_id)
+            raise HTTPException(status_code=400, detail="Must provide 'text' or 'approvals'")
+
+        if not success:
+            broadcaster.unsubscribe(sub_id)
+            raise HTTPException(status_code=404, detail="Session not found or inactive")
+    except HTTPException:
+        raise
+    except Exception:
+        broadcaster.unsubscribe(sub_id)
+        raise
+
+    # Terminal events that end the SSE stream
+    TERMINAL_EVENTS = {"turn_complete", "approval_required", "error", "interrupted", "shutdown"}
+
+    async def event_generator():
+        try:
+            while True:
+                msg = await event_queue.get()
+                event_type = msg.get("event_type", "")
+                # Format as SSE
+                yield f"data: {json.dumps(msg)}\n\n"
+                if event_type in TERMINAL_EVENTS:
+                    break
+        finally:
+            broadcaster.unsubscribe(sub_id)
+
+    return StreamingResponse(
+        event_generator(),
+        media_type="text/event-stream",
+        headers={
+            "Cache-Control": "no-cache",
+            "Connection": "keep-alive",
+            "X-Accel-Buffering": "no",
+        },
+    )
+
+
 @router.post("/interrupt/{session_id}")
 async def interrupt_session(
     session_id: str, user: dict = Depends(get_current_user)
@@ -344,77 +420,3 @@ async def shutdown_session(
     return {"status": "shutdown_requested", "session_id": session_id}
 
 
-@router.websocket("/ws/{session_id}")
-async def websocket_endpoint(websocket: WebSocket, session_id: str) -> None:
-    """WebSocket endpoint for real-time events.
-
-    Authentication is done via:
-    - ?token= query parameter (for browsers that can't send WS headers)
-    - Cookie (automatic for same-origin connections)
-    - Dev mode bypass (when OAUTH_CLIENT_ID is not set)
-
-    NOTE: We must accept() before close() so the browser receives our custom
-    close codes (4001, 4003, 4004).  If we close() before accept(), Starlette
-    sends HTTP 403 and the browser only sees code 1006 (abnormal closure).
-    """
-    logger.info(f"WebSocket connection request for session {session_id}")
-
-    # Authenticate the WebSocket connection
-    user = await get_ws_user(websocket)
-    if not user:
-        logger.warning(
-            f"WebSocket rejected: authentication failed for session {session_id}"
-        )
-        await websocket.accept()
-        await websocket.close(code=4001, reason="Authentication required")
-        return
-
-    # Verify session exists
-    info = session_manager.get_session_info(session_id)
-    if not info:
-        logger.warning(f"WebSocket rejected: session {session_id} not found")
-        await websocket.accept()
-        await websocket.close(code=4004, reason="Session not found")
-        return
-
-    # Verify user owns the session
-    if not session_manager.verify_session_access(session_id, user["user_id"]):
-        logger.warning(
-            f"WebSocket rejected: user {user['user_id']} denied access to session {session_id}"
-        )
-        await websocket.accept()
-        await websocket.close(code=4003, reason="Access denied")
-        return
-
-    had_buffered = await ws_manager.connect(websocket, session_id)
-
-    # Send "ready" on fresh connections so the frontend knows the session
-    # is alive.  Skip it when buffered events were flushed — those already
-    # contain the correct state and a ready would incorrectly reset
-    # isProcessing on the frontend.
-    if not had_buffered:
-        try:
-            await websocket.send_json(
-                {
-                    "event_type": "ready",
-                    "data": {"message": "Agent initialized"},
-                }
-            )
-        except Exception as e:
-            logger.error(f"Failed to send ready event for session {session_id}: {e}")
-
-    try:
-        while True:
-            # Keep connection alive, handle ping/pong
-            data = await websocket.receive_json()
-
-            # Handle client messages (e.g., ping)
-            if data.get("type") == "ping":
-                await websocket.send_json({"type": "pong"})
-
-    except WebSocketDisconnect:
-        logger.info(f"WebSocket disconnected for session {session_id}")
-    except Exception as e:
-        logger.error(f"WebSocket error for session {session_id}: {e}")
-    finally:
-        ws_manager.disconnect(session_id)

backend/session_manager.py

@@ -8,8 +8,6 @@ from datetime import datetime
 from pathlib import Path
 from typing import Any, Optional
 
-from websocket import manager as ws_manager
-
 from agent.config import load_config
 from agent.core.agent_loop import process_submission
 from agent.core.session import Event, OpType, Session
@@ -40,6 +38,44 @@ class Submission:
 logger = logging.getLogger(__name__)
 
 
+class EventBroadcaster:
+    """Reads from the agent's event queue and fans out to SSE subscribers.
+
+    Events that arrive when no subscribers are listening are discarded.
+    With SSE each turn is a separate request, so there is no reconnect
+    scenario that would need buffered replay.
+    """
+
+    def __init__(self, event_queue: asyncio.Queue):
+        self._source = event_queue
+        self._subscribers: dict[int, asyncio.Queue] = {}
+        self._counter = 0
+
+    def subscribe(self) -> tuple[int, asyncio.Queue]:
+        """Create a new subscriber. Returns (id, queue)."""
+        self._counter += 1
+        sub_id = self._counter
+        q: asyncio.Queue = asyncio.Queue()
+        self._subscribers[sub_id] = q
+        return sub_id, q
+
+    def unsubscribe(self, sub_id: int) -> None:
+        self._subscribers.pop(sub_id, None)
+
+    async def run(self) -> None:
+        """Main loop — reads from source queue and broadcasts."""
+        while True:
+            try:
+                event: Event = await self._source.get()
+                msg = {"event_type": event.event_type, "data": event.data}
+                for q in self._subscribers.values():
+                    await q.put(msg)
+            except asyncio.CancelledError:
+                break
+            except Exception as e:
+                logger.error(f"EventBroadcaster error: {e}")
+
+
 @dataclass
 class AgentSession:
     """Wrapper for an agent session with its associated resources."""
@@ -53,6 +89,7 @@ class AgentSession:
     task: asyncio.Task | None = None
     created_at: datetime = field(default_factory=datetime.utcnow)
     is_active: bool = True
+    broadcaster: Any = None
 
 
 class SessionCapacityError(Exception):
@@ -126,7 +163,7 @@ class SessionManager:
 
         # Run blocking constructors in a thread to keep the event loop responsive.
         # Without this, Session.__init__ → ContextManager → litellm.get_max_tokens()
-        # blocks all HTTP/WebSocket handling.
+        # blocks all HTTP/SSE handling.
         import time as _time
 
         def _create_session_sync():
@@ -182,7 +219,7 @@ class SessionManager:
         event_queue: asyncio.Queue,
         tool_router: ToolRouter,
     ) -> None:
-        """Run the agent loop for a session and forward events to WebSocket."""
+        """Run the agent loop for a session and broadcast events via EventBroadcaster."""
         agent_session = self.sessions.get(session_id)
         if not agent_session:
             logger.error(f"Session {session_id} not found")
@@ -190,10 +227,10 @@ class SessionManager:
 
         session = agent_session.session
 
-        # Start event forwarder task
-        event_forwarder = asyncio.create_task(
-            self._forward_events(session_id, event_queue)
-        )
+        # Start event broadcaster task
+        broadcaster = EventBroadcaster(event_queue)
+        agent_session.broadcaster = broadcaster
+        broadcast_task = asyncio.create_task(broadcaster.run())
 
         try:
             async with tool_router:
@@ -223,9 +260,9 @@ class SessionManager:
                         )
 
         finally:
-            event_forwarder.cancel()
+            broadcast_task.cancel()
             try:
-                await event_forwarder
+                await broadcast_task
             except asyncio.CancelledError:
                 pass
 
@@ -237,19 +274,6 @@ class SessionManager:
 
             logger.info(f"Session {session_id} ended")
 
-    async def _forward_events(
-        self, session_id: str, event_queue: asyncio.Queue
-    ) -> None:
-        """Forward events from the agent to the WebSocket."""
-        while True:
-            try:
-                event: Event = await event_queue.get()
-                await ws_manager.send_event(session_id, event.event_type, event.data)
-            except asyncio.CancelledError:
-                break
-            except Exception as e:
-                logger.error(f"Error forwarding event for {session_id}: {e}")
-
     async def submit(self, session_id: str, operation: Operation) -> bool:
         """Submit an operation to a session."""
         async with self._lock:
@@ -320,8 +344,6 @@ class SessionManager:
         if not agent_session:
             return False
 
-        ws_manager.clear_buffer(session_id)
-
         # Clean up sandbox Space before cancelling the task
         await self._cleanup_sandbox(agent_session.session)
 

backend/websocket.py

@@ -1,105 +0,0 @@
-"""WebSocket connection manager for real-time communication."""
-
-import logging
-from collections import deque
-from typing import Any
-
-from fastapi import WebSocket
-
-logger = logging.getLogger(__name__)
-
-# Max events buffered per session while WS is disconnected.
-MAX_EVENT_BUFFER = 500
-
-
-class ConnectionManager:
-    """Manages WebSocket connections for multiple sessions."""
-
-    def __init__(self) -> None:
-        # session_id -> WebSocket
-        self.active_connections: dict[str, WebSocket] = {}
-        # session_id -> events buffered while WS was disconnected
-        self._event_buffers: dict[str, deque[dict[str, Any]]] = {}
-
-    async def connect(self, websocket: WebSocket, session_id: str) -> bool:
-        """Accept a WebSocket connection, register it, and flush buffered events.
-
-        Returns True if buffered events were flushed (i.e. this is a reconnect
-        that had missed events).
-        """
-        logger.info(f"Attempting to accept WebSocket for session {session_id}")
-        await websocket.accept()
-        self.active_connections[session_id] = websocket
-
-        # Flush events that were buffered while the WS was disconnected
-        buffered = self._event_buffers.pop(session_id, None)
-        if buffered:
-            logger.info(
-                f"Flushing {len(buffered)} buffered events for session {session_id}"
-            )
-            for message in buffered:
-                try:
-                    await websocket.send_json(message)
-                except Exception:
-                    logger.error(
-                        f"Error flushing buffered event for session {session_id}"
-                    )
-                    break
-
-        logger.info(f"WebSocket connected and registered for session {session_id}")
-        return bool(buffered)
-
-    def disconnect(self, session_id: str) -> None:
-        """Remove a WebSocket connection."""
-        if session_id in self.active_connections:
-            del self.active_connections[session_id]
-        logger.info(f"WebSocket disconnected for session {session_id}")
-
-    def clear_buffer(self, session_id: str) -> None:
-        """Clear the event buffer for a session (e.g. on session delete)."""
-        self._event_buffers.pop(session_id, None)
-
-    async def send_event(
-        self, session_id: str, event_type: str, data: dict[str, Any] | None = None
-    ) -> None:
-        """Send an event to a specific session's WebSocket.
-
-        If no WebSocket is connected, the event is buffered so it can be
-        replayed when the client reconnects.
-        """
-        message: dict[str, Any] = {"event_type": event_type}
-        if data is not None:
-            message["data"] = data
-
-        if session_id not in self.active_connections:
-            buf = self._event_buffers.setdefault(
-                session_id, deque(maxlen=MAX_EVENT_BUFFER)
-            )
-            buf.append(message)
-            return
-
-        try:
-            await self.active_connections[session_id].send_json(message)
-        except Exception as e:
-            logger.error(f"Error sending to session {session_id}: {e}")
-            self.disconnect(session_id)
-            # Buffer the event that failed to send
-            buf = self._event_buffers.setdefault(
-                session_id, deque(maxlen=MAX_EVENT_BUFFER)
-            )
-            buf.append(message)
-
-    async def broadcast(
-        self, event_type: str, data: dict[str, Any] | None = None
-    ) -> None:
-        """Broadcast an event to all connected sessions."""
-        for session_id in list(self.active_connections.keys()):
-            await self.send_event(session_id, event_type, data)
-
-    def is_connected(self, session_id: str) -> bool:
-        """Check if a session has an active WebSocket connection."""
-        return session_id in self.active_connections
-
-
-# Global connection manager instance
-manager = ConnectionManager()

frontend/src/components/SessionChat.tsx

@@ -2,8 +2,8 @@
  * Per-session chat component.
  *
  * Each session renders its own SessionChat. The hook (useAgentChat) always
- * runs — keeping the WebSocket alive and processing events — but only the
- * active session renders visible UI (MessageList + ChatInput).
+ * runs — processing events — but only the active session renders visible
+ * UI (MessageList + ChatInput).
  */
 import { useCallback, useEffect, useRef } from 'react';
 import { useAgentChat } from '@/hooks/useAgentChat';
@@ -25,7 +25,7 @@ export default function SessionChat({ sessionId, isActive, onSessionDead }: Sess
   const { isConnected, isProcessing, setProcessing, activityStatus } = useAgentStore();
   const { updateSessionTitle } = useSessionStore();
 
-  const { messages, sendMessage, stop, undoLastTurn, approveTools, transport } = useAgentChat({
+  const { messages, sendMessage, stop, undoLastTurn, approveTools } = useAgentChat({
     sessionId,
     isActive,
     onReady: () => logger.log(`Session ${sessionId} ready`),
@@ -38,16 +38,10 @@ export default function SessionChat({ sessionId, isActive, onSessionDead }: Sess
   const prevActiveRef = useRef(isActive);
   useEffect(() => {
     if (isActive && !prevActiveRef.current) {
-      // Force reconnect if WS is dead (e.g. retries exhausted while on another session)
-      if (transport && !transport.isWebSocketConnected()) {
-        transport.connectToSession(sessionId);
-      }
-
       const store = useAgentStore.getState();
 
-      // Sync WebSocket connection state
-      const wsConnected = transport?.isWebSocketConnected() ?? false;
-      store.setConnected(wsConnected);
+      // SSE transport has no persistent connection — always connected
+      store.setConnected(true);
 
       // Check if this session has pending approvals in its messages
       const lastAssistant = [...messages].reverse().find(m => m.role === 'assistant');
@@ -109,7 +103,7 @@ export default function SessionChat({ sessionId, isActive, onSessionDead }: Sess
       }
     }
     prevActiveRef.current = isActive;
-  }, [isActive, messages, transport]);
+  }, [isActive, messages]);
 
   const handleSendMessage = useCallback(
     async (text: string) => {

frontend/src/hooks/useAgentChat.ts

@@ -1,6 +1,6 @@
 /**
- * Central hook wiring the Vercel AI SDK's useChat with our custom
- * WebSocketChatTransport.
+ * Central hook wiring the Vercel AI SDK's useChat with our SSE-based
+ * ChatTransport.
  *
  * In the per-session architecture, each session mounts its own instance
  * of this hook. The `isActive` flag controls whether side-channel
@@ -9,8 +9,8 @@
  */
 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 { type UIMessage, lastAssistantMessageIsCompleteWithApprovalResponses } from 'ai';
+import { SSEChatTransport, type SideChannelCallbacks } from '@/lib/sse-chat-transport';
 import { loadMessages, saveMessages } from '@/lib/chat-message-store';
 import { llmMessagesToUIMessages } from '@/lib/convert-llm-messages';
 import { apiFetch } from '@/utils/api';
@@ -47,7 +47,6 @@ export function useAgentChat({ sessionId, isActive, onReady, onError, onSessionD
   const { setSessionActive, setNeedsAttention } = useSessionStore();
 
   // -- Build side-channel callbacks (stable ref) --------------------------
-  // These check isActiveRef to decide whether to update global UI state.
   const sideChannel = useMemo<SideChannelCallbacks>(
     () => ({
       onReady: () => {
@@ -83,19 +82,9 @@ export function useAgentChat({ sessionId, isActive, onReady, onError, onSessionD
         }
       },
       onUndoComplete: () => {
+        // Undo is handled client-side in undoLastTurn(). With SSE, undo_complete
+        // events are discarded (no subscriber listening between turns).
         if (isActiveRef.current) setProcessing(false);
-        // 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);
-          saveMessages(sessionId, updated);
-        }
       },
       onCompacted: (oldTokens: number, newTokens: number) => {
         logger.log(`Context compacted: ${oldTokens} -> ${newTokens} tokens`);
@@ -132,11 +121,7 @@ export function useAgentChat({ sessionId, isActive, onReady, onError, onSessionD
       },
       onApprovalRequired: (tools) => {
         if (!tools.length) return;
-
-        // Always mark the session as needing attention
         setNeedsAttention(sessionId, true);
-
-        // Only update global UI if this is the active session
         if (!isActiveRef.current) return;
 
         setActivityStatus({ type: 'waiting-approval' });
@@ -199,30 +184,21 @@ export function useAgentChat({ sessionId, isActive, onReady, onError, onSessionD
         if (isActiveRef.current) setActivityStatus({ type: 'tool', toolName, description });
       },
     }),
-    // sessionId is the only real dependency — Zustand setters are stable
     // eslint-disable-next-line react-hooks/exhaustive-deps
     [sessionId],
   );
 
   // -- Create transport (one per session, stable for lifetime) ------------
-  const transportRef = useRef<WebSocketChatTransport | null>(null);
+  const transportRef = useRef<SSEChatTransport | null>(null);
   if (!transportRef.current) {
-    transportRef.current = new WebSocketChatTransport({ sideChannel });
+    transportRef.current = new SSEChatTransport(sessionId, sideChannel);
   }
 
-  // Keep side-channel callbacks in sync (they capture isActiveRef)
+  // Keep side-channel callbacks in sync
   useEffect(() => {
     transportRef.current?.updateSideChannel(sideChannel);
   }, [sideChannel]);
 
-  // Connect WebSocket on mount, disconnect on unmount
-  useEffect(() => {
-    transportRef.current?.connectToSession(sessionId);
-    return () => {
-      transportRef.current?.connectToSession(null);
-    };
-  }, [sessionId]);
-
   // Destroy transport on unmount
   useEffect(() => {
     return () => {
@@ -249,6 +225,10 @@ export function useAgentChat({ sessionId, isActive, onReady, onError, onSessionD
     messages: initialMessages,
     transport: transportRef.current!,
     experimental_throttle: 80,
+    // After all approval responses are set, auto-send to continue the agent loop.
+    // Without this, addToolApprovalResponse only updates the UI — it won't trigger
+    // sendMessages on the transport.
+    sendAutomaticallyWhen: lastAssistantMessageIsCompleteWithApprovalResponses,
     onError: (error) => {
       logger.error('useChat error:', error);
       if (isActiveRef.current) {
@@ -267,15 +247,12 @@ export function useAgentChat({ sessionId, isActive, onReady, onError, onSessionD
     let cancelled = false;
     (async () => {
       try {
-        // Fetch messages and session info (for pending approval) in parallel
         const [msgsRes, infoRes] = await Promise.all([
           apiFetch(`/api/session/${sessionId}/messages`),
           apiFetch(`/api/session/${sessionId}`),
         ]);
-
         if (cancelled) return;
 
-        // Extract pending approval tool IDs from session info
         let pendingIds: Set<string> | undefined;
         if (infoRes.ok) {
           const info = await infoRes.json();
@@ -315,25 +292,46 @@ export function useAgentChat({ sessionId, isActive, onReady, onError, onSessionD
     }
   }, [sessionId, chat.messages]);
 
-  // -- Undo last turn -----------------------------------------------------
+  // -- Undo last turn (REST call + client-side message removal) -----------
+  // With SSE there's no persistent connection to receive the undo_complete
+  // event, so we handle message removal on the frontend after a successful
+  // REST call to the backend.
   const undoLastTurn = useCallback(async () => {
     try {
       const res = await apiFetch(`/api/undo/${sessionId}`, { method: 'POST' });
       if (!res.ok) {
         logger.error('Undo API returned', res.status);
+        return;
+      }
+      // Remove the last user turn + assistant response from the UI
+      const msgs = chatActionsRef.current.messages;
+      const setMsgs = chatActionsRef.current.setMessages;
+      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);
+        saveMessages(sessionId, updated);
       }
+      if (isActiveRef.current) setProcessing(false);
     } catch (e) {
       logger.error('Undo failed:', e);
     }
-  }, [sessionId]);
+  }, [sessionId, setProcessing]);
 
-  // -- Approve tools via transport ----------------------------------------
+  // -- Approve tools ------------------------------------------------------
   const approveTools = useCallback(
     async (approvals: Array<{ tool_call_id: string; approved: boolean; feedback?: string | null; edited_script?: string | null }>) => {
-      if (!transportRef.current) return false;
+      // Store edited scripts so the transport can read them when sendMessages is called
+      for (const a of approvals) {
+        if (a.edited_script) {
+          useAgentStore.getState().setEditedScript(a.tool_call_id, a.edited_script);
+        }
+      }
 
-      // Transition SDK tool state from approval-requested → approval-responded
-      // so the approval UI disappears immediately (survives session switches).
+      // Update SDK tool state — this triggers sendMessages() via the transport
       for (const a of approvals) {
         chat.addToolApprovalResponse({
           id: `approval-${a.tool_call_id}`,
@@ -342,25 +340,26 @@ export function useAgentChat({ sessionId, isActive, onReady, onError, onSessionD
         });
       }
 
-      const ok = await transportRef.current.approveTools(sessionId, approvals);
-      if (ok) {
-        // Clear needsAttention since user has responded
-        setNeedsAttention(sessionId, false);
-        const hasApproved = approvals.some(a => a.approved);
-        if (hasApproved && isActiveRef.current) setProcessing(true);
-      }
-      return ok;
+      setNeedsAttention(sessionId, false);
+      const hasApproved = approvals.some(a => a.approved);
+      if (hasApproved && isActiveRef.current) setProcessing(true);
+      return true;
     },
     [sessionId, chat, setProcessing, setNeedsAttention],
   );
 
+  // -- Stop (abort SSE stream + interrupt backend agent loop) ---------------
+  const stop = useCallback(() => {
+    chat.stop();
+    apiFetch(`/api/interrupt/${sessionId}`, { method: 'POST' }).catch(() => {});
+  }, [sessionId, chat]);
+
   return {
     messages: chat.messages,
     sendMessage: chat.sendMessage,
-    stop: chat.stop,
+    stop,
     status: chat.status,
     undoLastTurn,
     approveTools,
-    transport: transportRef.current,
   };
 }

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

@@ -0,0 +1,361 @@
+/**
+ * SSE-based ChatTransport that bridges our backend event protocol
+ * to the Vercel AI SDK's UIMessageChunk streaming interface.
+ *
+ * Each sendMessages() call does a POST → SSE response.
+ * One request per turn phase (initial message, or approval continuation).
+ */
+import type { ChatTransport, UIMessage, UIMessageChunk, ChatRequestOptions } from 'ai';
+import { apiFetch } from '@/utils/api';
+import { logger } from '@/utils/logger';
+import type { AgentEvent } from '@/types/events';
+import { useAgentStore } from '@/store/agentStore';
+
+// ---------------------------------------------------------------------------
+// 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;
+  onApprovalRequired: (tools: Array<{ tool: string; arguments: Record<string, unknown>; tool_call_id: string }>) => void;
+  onToolCallPanel: (tool: string, args: Record<string, unknown>) => void;
+  onToolOutputPanel: (tool: string, toolCallId: string, output: string, success: boolean) => void;
+  onStreaming: () => void;
+  onToolRunning: (toolName: string, description?: string) => void;
+}
+
+// ---------------------------------------------------------------------------
+// Helpers
+// ---------------------------------------------------------------------------
+let partIdCounter = 0;
+function nextPartId(prefix: string): string {
+  return `${prefix}-${Date.now()}-${++partIdCounter}`;
+}
+
+/** Parse an SSE text stream into AgentEvent objects. */
+function createSSEParserStream(): TransformStream<string, AgentEvent> {
+  let buffer = '';
+  return new TransformStream<string, AgentEvent>({
+    transform(chunk, controller) {
+      buffer += chunk;
+      const lines = buffer.split('\n');
+      // Keep the last (possibly incomplete) line in the buffer
+      buffer = lines.pop() || '';
+      for (const line of lines) {
+        const trimmed = line.trim();
+        if (trimmed.startsWith('data: ')) {
+          try {
+            const json = JSON.parse(trimmed.slice(6));
+            controller.enqueue(json as AgentEvent);
+          } catch {
+            logger.warn('SSE parse error:', trimmed);
+          }
+        }
+      }
+    },
+    flush(controller) {
+      // Process any remaining data in buffer
+      if (buffer.trim().startsWith('data: ')) {
+        try {
+          const json = JSON.parse(buffer.trim().slice(6));
+          controller.enqueue(json as AgentEvent);
+        } catch { /* ignore incomplete */ }
+      }
+    },
+  });
+}
+
+/** Transform AgentEvent objects into UIMessageChunk objects for the Vercel AI SDK. */
+function createEventToChunkStream(sideChannel: SideChannelCallbacks): TransformStream<AgentEvent, UIMessageChunk> {
+  let textPartId: string | null = null;
+
+  function endTextPart(controller: TransformStreamDefaultController<UIMessageChunk>) {
+    if (textPartId) {
+      controller.enqueue({ type: 'text-end', id: textPartId });
+      textPartId = null;
+    }
+  }
+
+  return new TransformStream<AgentEvent, UIMessageChunk>({
+    transform(event, controller) {
+      switch (event.event_type) {
+        // -- Side-channel only events ----------------------------------------
+        case 'ready':
+          sideChannel.onReady();
+          break;
+
+        case 'shutdown':
+          sideChannel.onShutdown();
+          break;
+
+        case 'interrupted':
+          endTextPart(controller);
+          controller.enqueue({ type: 'finish-step' });
+          controller.enqueue({ type: 'finish', finishReason: 'stop' });
+          sideChannel.onProcessingDone();
+          break;
+
+        case 'undo_complete':
+          endTextPart(controller);
+          sideChannel.onUndoComplete();
+          break;
+
+        case 'compacted':
+          sideChannel.onCompacted(
+            (event.data?.old_tokens as number) || 0,
+            (event.data?.new_tokens as number) || 0,
+          );
+          break;
+
+        case 'plan_update':
+          sideChannel.onPlanUpdate(
+            (event.data?.plan as Array<{ id: string; content: string; status: string }>) || [],
+          );
+          break;
+
+        case 'tool_log':
+          sideChannel.onToolLog(
+            (event.data?.tool as string) || '',
+            (event.data?.log as string) || '',
+          );
+          break;
+
+        // -- Chat stream events ----------------------------------------------
+        case 'processing':
+          sideChannel.onProcessing();
+          controller.enqueue({ type: 'start', messageMetadata: { createdAt: new Date().toISOString() } });
+          controller.enqueue({ type: 'start-step' });
+          break;
+
+        case 'assistant_chunk': {
+          const delta = (event.data?.content as string) || '';
+          if (!delta) break;
+          if (!textPartId) {
+            textPartId = nextPartId('text');
+            controller.enqueue({ type: 'text-start', id: textPartId });
+            sideChannel.onStreaming();
+          }
+          controller.enqueue({ type: 'text-delta', id: textPartId, delta });
+          break;
+        }
+
+        case 'assistant_stream_end':
+          endTextPart(controller);
+          break;
+
+        case 'assistant_message': {
+          const content = (event.data?.content as string) || '';
+          if (!content) break;
+          const id = nextPartId('text');
+          controller.enqueue({ type: 'text-start', id });
+          controller.enqueue({ type: 'text-delta', id, delta: content });
+          controller.enqueue({ type: 'text-end', id });
+          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, unknown>) || {};
+          if (toolName === 'plan_tool') break;
+
+          endTextPart(controller);
+          controller.enqueue({ type: 'tool-input-start', toolCallId, toolName, dynamic: true });
+          controller.enqueue({ type: 'tool-input-available', toolCallId, toolName, input: args, dynamic: true });
+
+          sideChannel.onToolRunning(toolName, (args as Record<string, unknown>)?.description as string | undefined);
+          sideChannel.onToolCallPanel(toolName, args as Record<string, unknown>);
+          break;
+        }
+
+        case 'tool_output': {
+          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 (toolName === 'plan_tool' || toolCallId.startsWith('plan_tool')) break;
+
+          if (success) {
+            controller.enqueue({ type: 'tool-output-available', toolCallId, output, dynamic: true });
+          } else {
+            controller.enqueue({ type: 'tool-output-error', toolCallId, errorText: output, dynamic: true });
+          }
+          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) break;
+
+          endTextPart(controller);
+          for (const t of tools) {
+            controller.enqueue({ type: 'tool-input-start', toolCallId: t.tool_call_id, toolName: t.tool, dynamic: true });
+            controller.enqueue({ type: 'tool-input-available', toolCallId: t.tool_call_id, toolName: t.tool, input: t.arguments, dynamic: true });
+            controller.enqueue({ type: 'tool-approval-request', approvalId: `approval-${t.tool_call_id}`, toolCallId: t.tool_call_id });
+          }
+          sideChannel.onApprovalRequired(tools);
+          // DON'T emit finish here — the stream will close naturally and the SDK
+          // will see there's a pending approval. The SDK calls sendMessages again
+          // after addToolApprovalResponse.
+          break;
+        }
+
+        case 'tool_state_change': {
+          const tcId = (event.data?.tool_call_id as string) || '';
+          const state = (event.data?.state as string) || '';
+          const toolName = (event.data?.tool as string) || '';
+          const jobUrl = (event.data?.jobUrl as string) || undefined;
+
+          if (tcId.startsWith('plan_tool')) break;
+
+          if (jobUrl && tcId) {
+            useAgentStore.getState().setJobUrl(tcId, jobUrl);
+          }
+          if (state === 'running' && toolName) {
+            sideChannel.onToolRunning(toolName);
+          }
+          if (state === 'rejected' || state === 'abandoned') {
+            controller.enqueue({ type: 'tool-output-denied', toolCallId: tcId });
+          }
+          break;
+        }
+
+        case 'turn_complete':
+          endTextPart(controller);
+          controller.enqueue({ type: 'finish-step' });
+          controller.enqueue({ type: 'finish', finishReason: 'stop' });
+          sideChannel.onProcessingDone();
+          break;
+
+        case 'error': {
+          const errorMsg = (event.data?.error as string) || 'Unknown error';
+          sideChannel.onError(errorMsg);
+          controller.enqueue({ type: 'error', errorText: errorMsg });
+          sideChannel.onProcessingDone();
+          break;
+        }
+
+        default:
+          logger.log('SSE transport: unknown event', event);
+      }
+    },
+  });
+}
+
+// ---------------------------------------------------------------------------
+// Transport implementation
+// ---------------------------------------------------------------------------
+export class SSEChatTransport implements ChatTransport<UIMessage> {
+  private sessionId: string;
+  private sideChannel: SideChannelCallbacks;
+
+  constructor(sessionId: string, sideChannel: SideChannelCallbacks) {
+    this.sessionId = sessionId;
+    this.sideChannel = sideChannel;
+    // Mark as connected immediately — no persistent connection to establish
+    sideChannel.onConnectionChange(true);
+  }
+
+  updateSideChannel(sideChannel: SideChannelCallbacks): void {
+    this.sideChannel = sideChannel;
+  }
+
+  destroy(): void {
+    // Nothing to clean up — no persistent connections
+  }
+
+  // -- 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 = this.sessionId;
+
+    // Detect: is this an approval continuation or a new user message?
+    // After addToolApprovalResponse, the SDK calls sendMessages again.
+    // The last assistant message will have tool parts in 'approval-responded' state.
+    const lastAssistant = [...options.messages].reverse().find(m => m.role === 'assistant');
+    const approvedParts = lastAssistant?.parts.filter(
+      (p) => p.type === 'dynamic-tool' && p.state === 'approval-responded'
+    ) || [];
+
+    let body: Record<string, unknown>;
+    if (approvedParts.length > 0) {
+      // Approval continuation — extract approval decisions
+      const approvals = approvedParts.map((p) => {
+        if (p.type !== 'dynamic-tool') return null;
+        const approved = p.approval?.approved ?? true;
+        // Get edited script from agentStore if available
+        const editedScript = useAgentStore.getState().getEditedScript(p.toolCallId);
+        return {
+          tool_call_id: p.toolCallId,
+          approved,
+          feedback: approved ? null : (p.approval?.reason || 'Rejected by user'),
+          edited_script: editedScript ?? null,
+        };
+      }).filter(Boolean);
+      body = { approvals };
+    } else {
+      // Normal user message
+      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('')
+        : '';
+      body = { text };
+    }
+
+    // POST to SSE endpoint
+    const response = await apiFetch(`/api/chat/${sessionId}`, {
+      method: 'POST',
+      body: JSON.stringify(body),
+      signal: options.abortSignal,
+      headers: {
+        'Content-Type': 'application/json',
+        'Accept': 'text/event-stream',
+      },
+    });
+
+    if (!response.ok) {
+      const errorText = await response.text().catch(() => 'Request failed');
+      throw new Error(`Chat request failed: ${response.status} ${errorText}`);
+    }
+
+    if (!response.body) {
+      throw new Error('No response body');
+    }
+
+    // Pipe: response bytes → text → SSE events → UIMessageChunks
+    return response.body
+      .pipeThrough(new TextDecoderStream())
+      .pipeThrough(createSSEParserStream())
+      .pipeThrough(createEventToChunkStream(this.sideChannel));
+  }
+
+  async reconnectToStream(): Promise<ReadableStream<UIMessageChunk> | null> {
+    // Not needed — each turn is a separate request
+    return null;
+  }
+}

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

@@ -1,599 +0,0 @@
-/**
- * Custom ChatTransport that bridges our WebSocket-based backend protocol
- * to the Vercel AI SDK's UIMessageChunk streaming interface.
- *
- * Each instance manages a single session's WebSocket connection.
- * In the per-session architecture, every session owns its own transport.
- */
-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';
-import { useAgentStore } from '@/store/agentStore';
-
-// ---------------------------------------------------------------------------
-// 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;
-  /** Called when assistant text starts streaming */
-  onStreaming: () => void;
-  /** Called when a tool starts running (non-plan) */
-  onToolRunning: (toolName: string, description?: string) => 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 abortedGeneration = 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;
-  private boundVisibilityHandler: (() => void) | null = null;
-
-  constructor({ sideChannel }: WebSocketChatTransportOptions) {
-    this.sideChannel = sideChannel;
-    this.setupVisibilityHandler();
-  }
-
-  private setupVisibilityHandler(): void {
-    this.boundVisibilityHandler = () => {
-      if (document.visibilityState === 'hidden') {
-        return;
-      }
-
-      if (document.visibilityState === 'visible' && this.currentSessionId) {
-        const wsState = this.ws?.readyState;
-        if (wsState !== WebSocket.OPEN && wsState !== WebSocket.CONNECTING) {
-          logger.log('Tab visible: WS is dead, reconnecting immediately');
-          this.retries = 0;
-          this.reconnectDelay = WS_RECONNECT_DELAY;
-          this.createWebSocket(this.currentSessionId);
-        } else if (wsState === WebSocket.OPEN) {
-          this.ws!.send(JSON.stringify({ type: 'ping' }));
-        }
-      }
-    };
-    document.addEventListener('visibilitychange', this.boundVisibilityHandler);
-  }
-
-  /** Update side-channel callbacks (e.g. when isActive changes). */
-  updateSideChannel(sideChannel: SideChannelCallbacks): void {
-    this.sideChannel = sideChannel;
-  }
-
-  /** Check if the WebSocket is currently connected. */
-  isWebSocketConnected(): boolean {
-    return this.ws?.readyState === WebSocket.OPEN;
-  }
-
-  // -- Public API ----------------------------------------------------------
-
-  /** Connect (or reconnect) to a session's WebSocket. */
-  connectToSession(sessionId: string | null): void {
-    if (this.connectTimeout) {
-      clearTimeout(this.connectTimeout);
-      this.connectTimeout = null;
-    }
-
-    // Same session — no-op
-    if (sessionId === this.currentSessionId && this.ws?.readyState === WebSocket.OPEN) {
-      return;
-    }
-
-    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; edited_script?: 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;
-    }
-    if (this.boundVisibilityHandler) {
-      document.removeEventListener('visibilitychange', this.boundVisibilityHandler);
-      this.boundVisibilityHandler = 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();
-
-    // Track generation to protect against late cancel from a stale stream
-    const gen = ++this.streamGeneration;
-    logger.log(`sendMessages: gen=${gen}, awaitingProcessing=${this.awaitingProcessing}, abortedGen=${this.abortedGeneration}`);
-
-    // 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 (gen=${gen})`);
-        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.abortedGeneration = this.streamGeneration;
-        logger.log(`Abort complete: awaitingProcessing=true, abortedGen=${this.abortedGeneration}`);
-        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, will reconnect on session switch');
-          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)) {
-      logger.log(`Filtering stale "${event.event_type}" (gen=${this.streamGeneration}, aborted=${this.abortedGeneration})`);
-      return;
-    }
-
-    switch (event.event_type) {
-      // -- Side-channel only events ----------------------------------------
-      case 'ready':
-        if (this.streamController) {
-          // Reconnect during an active turn — don't reset processing state,
-          // just signal that the connection is back.
-          this.sideChannel.onConnectionChange(true);
-        } else {
-          this.sideChannel.onReady();
-        }
-        break;
-
-      case 'shutdown':
-        this.sideChannel.onShutdown();
-        this.closeActiveStream();
-        break;
-
-      case 'interrupted':
-        this.sideChannel.onProcessingDone();
-        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':
-        if (this.awaitingProcessing) {
-          if (this.streamGeneration <= this.abortedGeneration) {
-            logger.log(`Ignoring stale "processing" (gen=${this.streamGeneration} <= aborted=${this.abortedGeneration})`);
-            break;
-          }
-          logger.log(`Accepting "processing" for new generation (gen=${this.streamGeneration}, aborted=${this.abortedGeneration})`);
-          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.sideChannel.onStreaming();
-        }
-        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.onToolRunning(toolName, (args as Record<string, unknown>)?.description as string | undefined);
-        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 (toolName === 'plan_tool' || toolCallId.startsWith('plan_tool')) break;
-
-        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': {
-        const tcId = (event.data?.tool_call_id as string) || '';
-        const state = (event.data?.state as string) || '';
-        const toolName = (event.data?.tool as string) || '';
-        const jobUrl = (event.data?.jobUrl as string) || undefined;
-
-        if (tcId.startsWith('plan_tool')) break;
-
-        if (jobUrl && tcId) {
-          useAgentStore.getState().setJobUrl(tcId, jobUrl);
-        }
-
-        if (state === 'running' && toolName) {
-          this.sideChannel.onToolRunning(toolName);
-        }
-
-        if (this.streamController && (state === 'rejected' || state === 'abandoned')) {
-          this.enqueue({ type: 'tool-output-denied', toolCallId: tcId });
-        }
-        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/utils/api.ts

@@ -38,10 +38,4 @@ export async function apiFetch(
   }
 
   return response;
-}
-
-/** Build the WebSocket URL for a session. */
-export function getWebSocketUrl(sessionId: string): string {
-  const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
-  return `${protocol}//${window.location.host}/api/ws/${sessionId}`;
-}
+}