Backend: WebSocket endpoint for streaming agent chat (/ws/sessions/{id}/chat)

[frankbria]

Parent Issue

Part of #500 — depends on #501

Context

The existing /ws WebSocket (in codeframe/ui/routers/websocket.py) handles project-level status broadcasts (subscribe/unsubscribe). This issue adds a dedicated per-session WebSocket endpoint that handles bidirectional chat with a coding agent — receiving user messages and relaying streamed agent tokens back to the browser.

What to Build

New file: codeframe/ui/routers/session_chat_ws.py

New WebSocket endpoint:

WS /ws/sessions/{session_id}/chat?token=<JWT>

Message protocol (client → server):

{ "type": "message", "content": "Add error handling to the auth module" }
{ "type": "interrupt" }
{ "type": "ping" }

Message protocol (server → client):

{ "type": "text_delta", "content": "Sure, I'll " }
{ "type": "text_delta", "content": "add try/except..." }
{ "type": "tool_use_start", "tool_name": "Read", "tool_input": {"file_path": "..."} }
{ "type": "tool_result", "tool_name": "Read", "content": "..." }
{ "type": "thinking", "content": "I need to check..." }
{ "type": "cost_update", "cost_usd": 0.0031, "input_tokens": 1200, "output_tokens": 340 }
{ "type": "done" }
{ "type": "error", "message": "..." }
{ "type": "pong" }

Endpoint logic

1. Authenticate JWT (same pattern as websocket.py)
2. Load session from DB — return 4008 if session not found or ended
3. Accept connection, set session state to active in DB
4. Start a background task that calls the streaming chat adapter (see issue Backend: WebSocket endpoint for streaming agent chat (/ws/sessions/{id}/chat) #502) and relays token events to the WebSocket
5. Accept incoming message events and route them to the running adapter
6. Accept interrupt events and signal the adapter to stop mid-generation
7. On adapter completion: write cost/token update to DB, send done event
8. On disconnect or session end: clean up gracefully

File to modify: codeframe/ui/server.py

Register the new router.

Connection manager

Reuse or extend codeframe/ui/shared.py manager pattern to track active session WebSockets so the streaming adapter can push to them.

Acceptance Criteria

• WebSocket endpoint accepts connections with valid JWT
• Rejects invalid/expired JWT with close code 1008
• Rejects connections to non-existent or ended sessions
• Relays streamed text deltas as individual text_delta events (not buffered to complete)
• Relays tool use start and result events
• Accepts interrupt message and stops generation
• Updates session cost/token counts in DB after each turn
• Handles client disconnect cleanly (no leaked coroutines)
• Integration test in tests/api/test_session_chat_ws.py using FastAPI WebSocket test client

Out of Scope

• The actual streaming call to Anthropic (tracked in Backend: WebSocket endpoint for streaming agent chat (/ws/sessions/{id}/chat) #502)
• Frontend (tracked in separate issues)

[traycerai]

Warning

Ticket Assist will be discontinued on May 1, 2026. After that date, the feature will no longer be available. We suggest installing the Traycer IDE Extension for continued AI-powered, spec-driven development. Thank you for choosing Traycer!

Plan

Observations

• file:codeframe/ui/routers/websocket.py is the existing project-level WS router (not registered in server.py — it's a v1 legacy router). Its JWT auth pattern (pyjwt.decode + get_async_session_maker() user lookup) is the exact pattern to replicate.
• file:codeframe/ui/shared.py has ConnectionManager (connect/disconnect/broadcast) and SharedState (tracks LeadAgent by project_id). The new endpoint needs a parallel per-session connection registry.
• file:codeframe/ui/server.py registers only v2 routers; the new router must be added here.
• file:codeframe/core/streaming.py's EventPublisher provides a proven async queue-based pub/sub pattern that can be reused for routing streamed tokens to the WebSocket.
• The issue notes the actual streaming adapter call to Anthropic is out of scope (Backend: WebSocket endpoint for streaming agent chat (/ws/sessions/{id}/chat) #502); a stub/placeholder is sufficient.

Approach

Create file:codeframe/ui/routers/session_chat_ws.py as a self-contained FastAPI APIRouter with a single @router.websocket endpoint. Extend file:codeframe/ui/shared.py with a SessionChatManager to track active session WebSocket connections and interrupt signals. Register the router in file:codeframe/ui/server.py. The streaming adapter is stubbed with an asyncio.Queue-based interface so the real adapter (issue #502) can be plugged in later.

---

Implementation Steps

1. Extend file:codeframe/ui/shared.py — Add SessionChatManager

Add a new class SessionChatManager below ConnectionManager. It must be thread-safe using asyncio.Lock.

Responsibilities:

• Track active WebSocket connections keyed by session_id: str → WebSocket
• Track per-session interrupt signals keyed by session_id: str → asyncio.Event
• Track per-session token queues keyed by session_id: str → asyncio.Queue (for the streaming adapter to push events into)

Methods to add:

• async register(session_id, websocket) — stores the connection and creates an interrupt Event and a token Queue for the session
• async unregister(session_id) — removes all tracking state for the session
• async get_interrupt_event(session_id) -> asyncio.Event | None
• async get_token_queue(session_id) -> asyncio.Queue | None
• async signal_interrupt(session_id) — sets the interrupt event

Add a module-level singleton: session_chat_manager = SessionChatManager()

---

2. Create file:codeframe/ui/routers/session_chat_ws.py

Create a new APIRouter(tags=["websocket"]) with no prefix (path is defined on the endpoint itself).

2a. JWT Authentication helper

Extract the JWT decode + user-existence check into a private async helper _authenticate_websocket(websocket) -> int | None that mirrors the pattern in file:codeframe/ui/routers/websocket.py:

• Read token from websocket.query_params
• Decode with pyjwt.decode(token, SECRET, algorithms=[JWT_ALGORITHM], audience=JWT_AUDIENCE)
• Extract user_id from sub claim
• Verify user is active via get_async_session_maker()
• On any failure, call await websocket.close(code=1008, reason="...") and return None

2b. Streaming adapter stub

Define a private async function _run_streaming_adapter(session_id, user_message, token_queue, interrupt_event) that acts as the placeholder for the real Anthropic streaming call (issue #502). It should:

• Accept a user_message: str, an asyncio.Queue to push event dicts into, and an asyncio.Event for interrupt signalling
• Push typed event dicts matching the server→client protocol (e.g., {"type": "text_delta", "content": "..."}, {"type": "done"}, {"type": "cost_update", ...})
• Periodically check interrupt_event.is_set() and exit early if set, pushing {"type": "done"}
• When the real adapter is ready, this function signature is the integration point

2c. Main WebSocket endpoint

@router.websocket("/ws/sessions/{session_id}/chat")
async def session_chat_ws(session_id: str, websocket: WebSocket)

Flow:

sequenceDiagram
    participant C as Client
    participant WS as session_chat_ws
    participant DB as Database/SessionManager
    participant SM as SessionChatManager
    participant Adapter as _run_streaming_adapter

    C->>WS: WS connect ?token=JWT
    WS->>WS: _authenticate_websocket() → user_id or close(1008)
    WS->>DB: load session by session_id
    alt session not found or ended
        WS-->>C: close(4008, "Session not found or ended")
    end
    WS->>C: websocket.accept()
    WS->>DB: update session state → "active"
    WS->>SM: register(session_id, websocket)

    loop Message receive loop
        C->>WS: {"type": "message", "content": "..."}
        WS->>WS: cancel previous adapter task if running
        WS->>SM: reset interrupt event
        WS->>Adapter: asyncio.create_task(_run_streaming_adapter(...))
        Adapter-->>SM: push events to token_queue
        WS->>C: relay events from token_queue (text_delta, tool_use_start, etc.)
        Adapter-->>SM: push {"type": "done"}
        WS->>C: send {"type": "done"}
        WS->>DB: write cost/token update

        C->>WS: {"type": "interrupt"}
        WS->>SM: signal_interrupt(session_id)

        C->>WS: {"type": "ping"}
        WS->>C: {"type": "pong"}
    end

    C-->>WS: disconnect (WebSocketDisconnect)
    WS->>SM: unregister(session_id)
    WS->>DB: update session state → "idle" (or leave as-is)

Loading

Implementation details:

• Use two concurrent asyncio.Tasks inside the endpoint:
  1. Receive task — loops on websocket.receive_text(), parses JSON, dispatches message/interrupt/ping
  2. Relay task — loops on token_queue.get(), sends each event dict via websocket.send_json()
• Use asyncio.wait({receive_task, relay_task}, return_when=asyncio.FIRST_EXCEPTION) or a shared asyncio.Event to coordinate shutdown
• On WebSocketDisconnect: cancel the adapter task and relay task, call session_chat_manager.unregister(session_id)
• On adapter completion ({"type": "done"}): send done to client, write cost/token update to DB (use asyncio.to_thread if DB is sync)
• On any unhandled exception: send {"type": "error", "message": "..."} before closing

Session DB interaction:

• Load session: use the existing SessionManager from file:codeframe/core/session_manager.py or a direct DB lookup — check that session_id exists and is not in ended state
• Update session state: call the appropriate DB update method (e.g., db.update_session(session_id, {"state": "active"})) — use asyncio.to_thread since Database is synchronous
• Write cost update after each turn: store cost_usd, input_tokens, output_tokens via the DB

Note: Since the Database class is synchronous (SQLite), wrap all DB calls in asyncio.to_thread(...) to avoid blocking the event loop, consistent with the pattern in file:codeframe/ui/shared.py's start_agent.

---

3. Register the router in file:codeframe/ui/server.py

In the "Router Mounting" section, import and include the new router:

• Import session_chat_ws from codeframe.ui.routers
• Add app.include_router(session_chat_ws.router) alongside the other router registrations
• Add a "session_chat" or extend the existing "websocket" tag entry in OPENAPI_TAGS if desired

---

4. Create file:tests/api/test_session_chat_ws.py

Follow the class-scoped fixture pattern from file:tests/api/conftest.py. Use FastAPI's TestClient with with client.websocket_connect(...) context manager (available via starlette.testclient).

Test classes and cases:

Test Class	Cases
TestSessionChatWSAuth	Rejects missing token (close 1008), rejects invalid/expired token (close 1008), accepts valid JWT
TestSessionChatWSSession	Rejects non-existent session (close 4008), rejects ended session (close 4008)
TestSessionChatWSProtocol	ping → pong, message → streams text_delta events → done, interrupt stops generation mid-stream
TestSessionChatWSCleanup	Client disconnect cleans up SessionChatManager state (no leaked tasks)

Use create_test_jwt_token(user_id=1) from file:tests/api/conftest.py for auth tokens. Mock _run_streaming_adapter to return controlled event sequences without real Anthropic calls.

Import In IDE

🤖 Prompt for AI Agents

## Observations

- `file:codeframe/ui/routers/websocket.py` is the existing project-level WS router (not registered in `server.py` — it's a v1 legacy router). Its JWT auth pattern (`pyjwt.decode` + `get_async_session_maker()` user lookup) is the exact pattern to replicate.
- `file:codeframe/ui/shared.py` has `ConnectionManager` (connect/disconnect/broadcast) and `SharedState` (tracks `LeadAgent` by `project_id`). The new endpoint needs a parallel per-session connection registry.
- `file:codeframe/ui/server.py` registers only v2 routers; the new router must be added here.
- `file:codeframe/core/streaming.py`'s `EventPublisher` provides a proven async queue-based pub/sub pattern that can be reused for routing streamed tokens to the WebSocket.
- The issue notes the actual streaming adapter call to Anthropic is out of scope (#502); a stub/placeholder is sufficient.

## Approach

Create `file:codeframe/ui/routers/session_chat_ws.py` as a self-contained FastAPI `APIRouter` with a single `@router.websocket` endpoint. Extend `file:codeframe/ui/shared.py` with a `SessionChatManager` to track active session WebSocket connections and interrupt signals. Register the router in `file:codeframe/ui/server.py`. The streaming adapter is stubbed with an `asyncio.Queue`-based interface so the real adapter (issue #502) can be plugged in later.

---

## Implementation Steps

### 1. Extend `file:codeframe/ui/shared.py` — Add `SessionChatManager`

Add a new class `SessionChatManager` below `ConnectionManager`. It must be thread-safe using `asyncio.Lock`.

**Responsibilities:**
- Track active WebSocket connections keyed by `session_id: str` → `WebSocket`
- Track per-session interrupt signals keyed by `session_id: str` → `asyncio.Event`
- Track per-session token queues keyed by `session_id: str` → `asyncio.Queue` (for the streaming adapter to push events into)

**Methods to add:**
- `async register(session_id, websocket)` — stores the connection and creates an interrupt `Event` and a token `Queue` for the session
- `async unregister(session_id)` — removes all tracking state for the session
- `async get_interrupt_event(session_id) -> asyncio.Event | None`
- `async get_token_queue(session_id) -> asyncio.Queue | None`
- `async signal_interrupt(session_id)` — sets the interrupt event

Add a module-level singleton: `session_chat_manager = SessionChatManager()`

---

### 2. Create `file:codeframe/ui/routers/session_chat_ws.py`

Create a new `APIRouter(tags=["websocket"])` with no prefix (path is defined on the endpoint itself).

#### 2a. JWT Authentication helper

Extract the JWT decode + user-existence check into a private async helper `_authenticate_websocket(websocket) -> int | None` that mirrors the pattern in `file:codeframe/ui/routers/websocket.py`:
- Read `token` from `websocket.query_params`
- Decode with `pyjwt.decode(token, SECRET, algorithms=[JWT_ALGORITHM], audience=JWT_AUDIENCE)`
- Extract `user_id` from `sub` claim
- Verify user is active via `get_async_session_maker()`
- On any failure, call `await websocket.close(code=1008, reason="...")` and return `None`

#### 2b. Streaming adapter stub

Define a private async function `_run_streaming_adapter(session_id, user_message, token_queue, interrupt_event)` that acts as the placeholder for the real Anthropic streaming call (issue #502). It should:
- Accept a `user_message: str`, an `asyncio.Queue` to push event dicts into, and an `asyncio.Event` for interrupt signalling
- Push typed event dicts matching the server→client protocol (e.g., `{"type": "text_delta", "content": "..."}`, `{"type": "done"}`, `{"type": "cost_update", ...}`)
- Periodically check `interrupt_event.is_set()` and exit early if set, pushing `{"type": "done"}`
- When the real adapter is ready, this function signature is the integration point

#### 2c. Main WebSocket endpoint

```
@router.websocket("/ws/sessions/{session_id}/chat")
async def session_chat_ws(session_id: str, websocket: WebSocket)
```

**Flow:**

```mermaid
sequenceDiagram
    participant C as Client
    participant WS as session_chat_ws
    participant DB as Database/SessionManager
    participant SM as SessionChatManager
    participant Adapter as _run_streaming_adapter

    C->>WS: WS connect ?token=JWT
    WS->>WS: _authenticate_websocket() → user_id or close(1008)
    WS->>DB: load session by session_id
    alt session not found or ended
        WS-->>C: close(4008, "Session not found or ended")
    end
    WS->>C: websocket.accept()
    WS->>DB: update session state → "active"
    WS->>SM: register(session_id, websocket)

    loop Message receive loop
        C->>WS: {"type": "message", "content": "..."}
        WS->>WS: cancel previous adapter task if running
        WS->>SM: reset interrupt event
        WS->>Adapter: asyncio.create_task(_run_streaming_adapter(...))
        Adapter-->>SM: push events to token_queue
        WS->>C: relay events from token_queue (text_delta, tool_use_start, etc.)
        Adapter-->>SM: push {"type": "done"}
        WS->>C: send {"type": "done"}
        WS->>DB: write cost/token update

        C->>WS: {"type": "interrupt"}
        WS->>SM: signal_interrupt(session_id)

        C->>WS: {"type": "ping"}
        WS->>C: {"type": "pong"}
    end

    C-->>WS: disconnect (WebSocketDisconnect)
    WS->>SM: unregister(session_id)
    WS->>DB: update session state → "idle" (or leave as-is)
```

**Implementation details:**

- Use two concurrent `asyncio.Task`s inside the endpoint:
  1. **Receive task** — loops on `websocket.receive_text()`, parses JSON, dispatches `message`/`interrupt`/`ping`
  2. **Relay task** — loops on `token_queue.get()`, sends each event dict via `websocket.send_json()`
- Use `asyncio.wait({receive_task, relay_task}, return_when=asyncio.FIRST_EXCEPTION)` or a shared `asyncio.Event` to coordinate shutdown
- On `WebSocketDisconnect`: cancel the adapter task and relay task, call `session_chat_manager.unregister(session_id)`
- On adapter completion (`{"type": "done"}`): send `done` to client, write cost/token update to DB (use `asyncio.to_thread` if DB is sync)
- On any unhandled exception: send `{"type": "error", "message": "..."}` before closing

**Session DB interaction:**
- Load session: use the existing `SessionManager` from `file:codeframe/core/session_manager.py` or a direct DB lookup — check that `session_id` exists and is not in `ended` state
- Update session state: call the appropriate DB update method (e.g., `db.update_session(session_id, {"state": "active"})`) — use `asyncio.to_thread` since `Database` is synchronous
- Write cost update after each turn: store `cost_usd`, `input_tokens`, `output_tokens` via the DB

> **Note:** Since the `Database` class is synchronous (SQLite), wrap all DB calls in `asyncio.to_thread(...)` to avoid blocking the event loop, consistent with the pattern in `file:codeframe/ui/shared.py`'s `start_agent`.

---

### 3. Register the router in `file:codeframe/ui/server.py`

In the "Router Mounting" section, import and include the new router:

- Import `session_chat_ws` from `codeframe.ui.routers`
- Add `app.include_router(session_chat_ws.router)` alongside the other router registrations
- Add a `"session_chat"` or extend the existing `"websocket"` tag entry in `OPENAPI_TAGS` if desired

---

### 4. Create `file:tests/api/test_session_chat_ws.py`

Follow the class-scoped fixture pattern from `file:tests/api/conftest.py`. Use FastAPI's `TestClient` with `with client.websocket_connect(...)` context manager (available via `starlette.testclient`).

**Test classes and cases:**

| Test Class | Cases |
|---|---|
| `TestSessionChatWSAuth` | Rejects missing token (close 1008), rejects invalid/expired token (close 1008), accepts valid JWT |
| `TestSessionChatWSSession` | Rejects non-existent session (close 4008), rejects `ended` session (close 4008) |
| `TestSessionChatWSProtocol` | `ping` → `pong`, `message` → streams `text_delta` events → `done`, `interrupt` stops generation mid-stream |
| `TestSessionChatWSCleanup` | Client disconnect cleans up `SessionChatManager` state (no leaked tasks) |

Use `create_test_jwt_token(user_id=1)` from `file:tests/api/conftest.py` for auth tokens. Mock `_run_streaming_adapter` to return controlled event sequences without real Anthropic calls.

---

Execution Information

Branch: main
Commit: cb074e3

---

💡 Tips

Supported Commands (Inside Comments)

• Use @traycerai generate to iterate on the previous version of the implementation plan.

Supported Commands (Inside Description)

• Add @traycerai ignore anywhere in the ticket description to prevent this ticket from being processed.
• Add @traycerai branch:<branch-name> anywhere in the ticket description to specify the target branch for the implementation plan.

Community

• Join our Discord Community to get help, request features, and share feedback.
• Follow us on X/Twitter for updates and announcements.

[coderabbitai]

🔗 Related PRs

#133 - WebSocket Subscription Filtering for Project-Scoped Broadcasts [merged]
#134 - WebSocket Subscription Filtering for Project-Scoped Broadcasts [merged]
#148 - fix: Add WebSocket health check and E2E test authentication [merged]
#224 - feat(websocket): implement proactive messaging system [merged]

---

📝 Issue Planner

_{Check the box below or use the @coderabbitai plan command to generate an implementation plan and prompts that you can use with your favorite coding assistant.}

• Create Plan

---

🧪 Issue enrichment is currently in open beta.

You can configure auto-planning by selecting labels in the issue_enrichment configuration.

To disable automatic issue enrichment, add the following to your .coderabbit.yaml:

issue_enrichment:
  auto_enrich:
    enabled: false

💬 Have feedback or questions? Drop into our discord!