feat(api): WebSocket endpoint for streaming agent chat (#502)

Adds /ws/sessions/{session_id}/chat WebSocket endpoint for bidirectional streaming agent chat.

- JWT auth (close 1008 on failure), session validation (close 4008 for ended/unknown)
- Concurrent receive + relay tasks for non-blocking interrupt handling
- Cost/token DB update persisted before sending "done" (accurate client reads)
- Cleanup guard wraps all post-accept logic in try/finally
- SessionChatManager tracks connections, interrupt events, and token queues per session;
  unregister() checks websocket identity to prevent late disconnects from tearing newer state
- Stub _run_streaming_adapter() ready for real Anthropic integration (#503)
- 12 integration tests covering auth, session validation, protocol, interrupt, cost, cleanup

Closes #502

Author: frankbria

codeframe/ui/routers/session_chat_ws.py

@@ -0,0 +1,325 @@
+"""WebSocket router for per-session streaming agent chat.
+
+Endpoint:
+    WS /ws/sessions/{session_id}/chat?token=<JWT>
+
+Client → Server message types:
+    {"type": "message", "content": "..."}
+    {"type": "interrupt"}
+    {"type": "ping"}
+
+Server → Client message types:
+    {"type": "text_delta", "content": "..."}
+    {"type": "tool_use_start", "tool_name": "...", "tool_input": {...}}
+    {"type": "tool_result", "tool_name": "...", "content": "..."}
+    {"type": "thinking", "content": "..."}
+    {"type": "cost_update", "cost_usd": 0.003, "input_tokens": 100, "output_tokens": 50}
+    {"type": "done"}
+    {"type": "error", "message": "..."}
+    {"type": "pong"}
+
+Architecture note: This router intentionally handles session orchestration
+(activation, interrupt coordination, cost aggregation) rather than delegating
+to a core service. A future refactoring (TODO: extract to
+core.session_chat_service) would make this a thin transport adapter. See
+GitHub issue #502 for context.
+
+Workspace scoping: This endpoint is intentionally exempt from workspace_path
+query param validation. Auth is already scoped to a user via JWT and the
+session_id identifies the resource — workspace scoping on a per-session WS
+would be redundant. Revisit if multi-tenant workspace isolation is required.
+"""
+
+import asyncio
+import json
+import logging
+from typing import Optional
+
+import jwt as pyjwt
+from fastapi import APIRouter, WebSocket, WebSocketDisconnect
+from sqlalchemy import select
+
+from codeframe.auth.manager import SECRET, JWT_ALGORITHM, JWT_AUDIENCE, get_async_session_maker
+from codeframe.auth.models import User
+from codeframe.ui.shared import session_chat_manager
+
+logger = logging.getLogger(__name__)
+
+router = APIRouter(tags=["websocket"])
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+async def _authenticate_websocket(websocket: WebSocket) -> Optional[int]:
+    """Validate JWT from query param. Returns user_id or closes with 1008."""
+    token = websocket.query_params.get("token")
+    if not token:
+        await websocket.close(code=1008, reason="Authentication required: missing token")
+        return None
+
+    try:
+        payload = pyjwt.decode(token, SECRET, algorithms=[JWT_ALGORITHM], audience=JWT_AUDIENCE)
+        user_id_str = payload.get("sub")
+        if not user_id_str:
+            await websocket.close(code=1008, reason="Invalid token: missing subject")
+            return None
+        user_id = int(user_id_str)
+    except pyjwt.ExpiredSignatureError:
+        await websocket.close(code=1008, reason="Token expired")
+        return None
+    except (pyjwt.InvalidTokenError, ValueError) as exc:
+        logger.debug("WebSocket JWT decode error: %s", exc)
+        await websocket.close(code=1008, reason="Invalid authentication token")
+        return None
+
+    try:
+        async_session_maker = get_async_session_maker()
+        async with async_session_maker() as session:
+            result = await session.execute(select(User).where(User.id == user_id))
+            user = result.scalar_one_or_none()
+            if user is None:
+                await websocket.close(code=1008, reason="User not found")
+                return None
+            if not user.is_active:
+                await websocket.close(code=1008, reason="User is inactive")
+                return None
+    except Exception as exc:
+        logger.error("WebSocket user lookup error: %s", exc)
+        await websocket.close(code=1008, reason="Authentication failed")
+        return None
+
+    return user_id
+
+
+async def _run_streaming_adapter(
+    session_id: str,
+    user_message: str,
+    token_queue: asyncio.Queue,
+    interrupt_event: asyncio.Event,
+) -> None:
+    """Stub for the real Anthropic streaming adapter (tracked in #503).
+
+    Signature is the integration point — replace the body when the real adapter
+    is ready. Must push dicts matching the server→client protocol into
+    token_queue, and check interrupt_event.is_set() periodically.
+    """
+    words = user_message.split() or ["..."]
+    for word in words:
+        if interrupt_event.is_set():
+            await token_queue.put({"type": "done"})
+            return
+        await token_queue.put({"type": "text_delta", "content": word + " "})
+        await asyncio.sleep(0.01)
+
+    await token_queue.put(
+        {"type": "cost_update", "cost_usd": 0.0, "input_tokens": 0, "output_tokens": 0}
+    )
+    await token_queue.put({"type": "done"})
+
+
+# ---------------------------------------------------------------------------
+# Endpoint
+# ---------------------------------------------------------------------------
+
+
+@router.websocket("/ws/sessions/{session_id}/chat")
+async def session_chat_ws(session_id: str, websocket: WebSocket) -> None:
+    """Bidirectional WebSocket for streaming agent chat on an interactive session."""
+    # --- Auth ---
+    user_id = await _authenticate_websocket(websocket)
+    if user_id is None:
+        return
+
+    # --- Session validation ---
+    db = getattr(websocket.app.state, "db", None)
+    if db is None:
+        await websocket.close(code=1011, reason="Database unavailable")
+        return
+
+    session = await asyncio.to_thread(db.interactive_sessions.get, session_id)
+    if session is None or session.get("state") == "ended":
+        await websocket.close(code=4008, reason="Session not found or ended")
+        return
+
+    # --- Accept connection; everything after this point must run inside the
+    #     try/finally so unregister() and close() always execute even if
+    #     update_state, register, or get_token_queue raises. ---
+    await websocket.accept()
+
+    relay: Optional[asyncio.Task] = None
+    adapter_task: list[Optional[asyncio.Task]] = [None]
+
+    try:
+        await asyncio.to_thread(db.interactive_sessions.update_state, session_id, "active")
+        await session_chat_manager.register(session_id, websocket)
+
+        token_queue = await session_chat_manager.get_token_queue(session_id)
+
+        # ---- Relay task: token_queue → WebSocket ----------------------------
+        async def _relay() -> None:
+            """Forward adapter events to the client; persist cost on 'done'."""
+            turn_cost = {"cost_usd": 0.0, "input_tokens": 0, "output_tokens": 0}
+            try:
+                while True:
+                    event = await token_queue.get()
+                    event_type = event.get("type")
+
+                    if event_type == "cost_update":
+                        turn_cost["cost_usd"] += event.get("cost_usd", 0.0)
+                        turn_cost["input_tokens"] += event.get("input_tokens", 0)
+                        turn_cost["output_tokens"] += event.get("output_tokens", 0)
+                        try:
+                            await websocket.send_json(event)
+                        except Exception as exc:
+                            logger.warning(
+                                "session_id=%s send_json(cost_update) failed: %s", session_id, exc
+                            )
+                            return
+
+                    elif event_type == "done":
+                        # Persist cost BEFORE sending "done" so clients that
+                        # immediately fetch session stats observe accurate totals.
+                        if (
+                            turn_cost["cost_usd"]
+                            or turn_cost["input_tokens"]
+                            or turn_cost["output_tokens"]
+                        ):
+                            try:
+                                await asyncio.to_thread(
+                                    db.interactive_sessions.update_cost,
+                                    session_id,
+                                    turn_cost["cost_usd"],
+                                    turn_cost["input_tokens"],
+                                    turn_cost["output_tokens"],
+                                )
+                            except Exception as exc:
+                                logger.error(
+                                    "session_id=%s update_cost failed: %s turn_cost=%s",
+                                    session_id,
+                                    exc,
+                                    turn_cost,
+                                )
+                        turn_cost = {"cost_usd": 0.0, "input_tokens": 0, "output_tokens": 0}
+                        try:
+                            await websocket.send_json(event)
+                        except Exception as exc:
+                            logger.warning(
+                                "session_id=%s send_json(done) failed: %s", session_id, exc
+                            )
+                            return
+
+                    else:
+                        try:
+                            await websocket.send_json(event)
+                        except Exception as exc:
+                            logger.warning(
+                                "session_id=%s send_json(%s) failed: %s",
+                                session_id,
+                                event_type,
+                                exc,
+                            )
+                            return
+            finally:
+                # Flush any cost accumulated during a cancelled/aborted turn
+                if (
+                    turn_cost["cost_usd"]
+                    or turn_cost["input_tokens"]
+                    or turn_cost["output_tokens"]
+                ):
+                    try:
+                        await asyncio.to_thread(
+                            db.interactive_sessions.update_cost,
+                            session_id,
+                            turn_cost["cost_usd"],
+                            turn_cost["input_tokens"],
+                            turn_cost["output_tokens"],
+                        )
+                    except Exception as exc:
+                        logger.error(
+                            "session_id=%s relay finally update_cost failed: %s", session_id, exc
+                        )
+
+        # ---- Receive task: WebSocket → dispatch -----------------------------
+        async def _receive() -> None:
+            """Read client messages and dispatch to ping/interrupt/message handlers."""
+            while True:
+                try:
+                    raw = await websocket.receive_text()
+                except WebSocketDisconnect:
+                    raise
+
+                try:
+                    msg = json.loads(raw)
+                except json.JSONDecodeError:
+                    try:
+                        await websocket.send_json({"type": "error", "message": "Invalid JSON"})
+                    except Exception:
+                        pass
+                    continue
+
+                msg_type = msg.get("type")
+
+                if msg_type == "ping":
+                    await websocket.send_json({"type": "pong"})
+
+                elif msg_type == "interrupt":
+                    await session_chat_manager.signal_interrupt(session_id)
+
+                elif msg_type == "message":
+                    content = msg.get("content", "")
+
+                    # Cancel any in-flight adapter
+                    if adapter_task[0] and not adapter_task[0].done():
+                        adapter_task[0].cancel()
+                        try:
+                            await adapter_task[0]
+                        except (asyncio.CancelledError, Exception) as exc:
+                            logger.debug(
+                                "session_id=%s adapter cancelled: %s", session_id, exc
+                            )
+
+                    # Reset interrupt and drain stale queue items
+                    await session_chat_manager.reset_interrupt(session_id)
+                    while not token_queue.empty():
+                        try:
+                            token_queue.get_nowait()
+                        except asyncio.QueueEmpty:
+                            break
+
+                    interrupt_event = await session_chat_manager.get_interrupt_event(session_id)
+                    adapter_task[0] = asyncio.create_task(
+                        _run_streaming_adapter(session_id, content, token_queue, interrupt_event)
+                    )
+
+        relay = asyncio.create_task(_relay())
+        await _receive()
+
+    except WebSocketDisconnect:
+        logger.debug("Session chat WebSocket disconnected: session_id=%s", session_id)
+    except Exception as exc:
+        logger.error("Session chat WebSocket error: %s", exc, exc_info=True)
+        try:
+            await websocket.send_json({"type": "error", "message": str(exc)})
+        except Exception:
+            pass
+    finally:
+        if relay is not None:
+            relay.cancel()
+            try:
+                await relay
+            except (asyncio.CancelledError, Exception):
+                pass
+        if adapter_task[0] and not adapter_task[0].done():
+            adapter_task[0].cancel()
+            try:
+                await adapter_task[0]
+            except (asyncio.CancelledError, Exception) as exc:
+                logger.debug("session_id=%s adapter cleanup: %s", session_id, exc)
+        await session_chat_manager.unregister(session_id, websocket)
+        try:
+            await websocket.close()
+        except Exception:
+            pass

codeframe/ui/server.py

@@ -34,6 +34,7 @@
     proof_v2,
     review_v2,
     schedule_v2,
+    session_chat_ws,
     streaming_v2,
     tasks_v2,
     templates_v2,
@@ -487,6 +488,7 @@ async def test_broadcast(message: dict, project_id: int = None):
 app.include_router(gates_v2.router)         # /api/v2/gates
 app.include_router(git_v2.router)           # /api/v2/git
 app.include_router(interactive_sessions_v2.router)  # /api/v2/sessions
+app.include_router(session_chat_ws.router)          # /ws/sessions/{id}/chat
 app.include_router(pr_v2.router)            # /api/v2/pr
 app.include_router(prd_v2.router)           # /api/v2/prd
 app.include_router(proof_v2.router)         # /api/v2/proof