Frontend: useAgentChat hook — WebSocket connection and message streaming

[frankbria]

Parent Issue

Part of #500 — depends on #501 and #502 (backend must be running)

Context

The web-ui currently has useEventSource (SSE, one-way) and useExecutionMonitor (aggregates SSE events into UI state). For interactive sessions, we need a bidirectional WebSocket hook that connects to /ws/sessions/{id}/chat, sends user messages, and accumulates streamed token events into structured chat messages.

Existing Code to Build On

• web-ui/src/hooks/useEventSource.ts — reconnect pattern to adapt for WS
• web-ui/src/hooks/useExecutionMonitor.ts — state accumulation pattern to follow
• Reference: Optio apps/web/src/hooks/use-logs.ts + use-websocket.ts

What to Build

New file: web-ui/src/hooks/useAgentChat.ts

type MessageRole =
  | "user"
  | "assistant"
  | "tool_use"
  | "tool_result"
  | "thinking"
  | "system"
  | "error";

interface ChatMessage {
  id: string;           // client-generated UUID per turn
  role: MessageRole;
  content: string;
  toolName?: string;    // for tool_use / tool_result
  toolInput?: unknown;  // for tool_use
  createdAt: string;
}

interface AgentChatState {
  messages: ChatMessage[];
  status: "idle" | "connecting" | "thinking" | "streaming" | "error" | "disconnected";
  costUsd: number;
  inputTokens: number;
  outputTokens: number;
  error: string | null;
  connected: boolean;
}

interface UseAgentChat {
  state: AgentChatState;
  sendMessage: (content: string) => void;
  interrupt: () => void;
  clearMessages: () => void;
}

export function useAgentChat(sessionId: string | null): UseAgentChat

Behavior

1. On sessionId change: open WebSocket to /ws/sessions/{sessionId}/chat?token=<JWT>
2. Send { type: "ping" } every 30s to keep alive
3. On text_delta events: append content to the current in-progress assistant message
4. On tool_use_start: push a new tool_use message
5. On tool_result: push a tool_result message
6. On thinking: push a thinking message
7. On cost_update: update costUsd, inputTokens, outputTokens in state
8. On done: finalize the in-progress assistant message, set status to idle
9. On error: set error field, set status to error
10. On WS close: set connected = false, status to disconnected, auto-reconnect with exponential backoff (max 5 attempts)
11. sendMessage(content): send { type: "message", content } over WS, push optimistic user message, set status to thinking
12. interrupt(): send { type: "interrupt" } over WS
13. Use requestAnimationFrame batching for text_delta updates to avoid render thrash

Auth

Get the JWT from the existing auth context/local storage — look at how other hooks pass credentials.

Acceptance Criteria

• Hook connects to WebSocket on mount, disconnects on unmount
• Streaming text deltas update the in-progress message character-by-character
• Tool use and tool result messages appear inline between assistant turns
• Thinking messages render separately from text content
• sendMessage pushes optimistic user message immediately, then sends over WS
• interrupt() sends interrupt message
• Reconnects automatically on disconnect (up to 5 attempts, exponential backoff)
• status transitions correctly: idle → thinking → streaming → idle
• No memory leaks: cleanup on unmount

Out of Scope

• The visual chat panel (tracked in the next sub-issue)
• Terminal (separate sub-issue)

[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

The current web-ui has no auth context — the v2 API client (file:web-ui/src/lib/api.ts) uses withCredentials: true (cookie-based auth), not a JWT Bearer token. The legacy web-ui stored a JWT in localStorage under auth_token. The existing hooks (useEventSource, useExecutionMonitor, useTaskStream) establish the patterns to follow: callback refs for stable closures, useRef + requestAnimationFrame batching for high-frequency updates, and exponential-backoff reconnect. No third-party WebSocket library is present in file:web-ui/package.json.

Approach

Build useAgentChat as a self-contained hook in file:web-ui/src/hooks/useAgentChat.ts, mirroring the layered pattern of useEventSource → useTaskStream → useExecutionMonitor. The hook manages the raw WebSocket directly (no extra library needed), uses useRef + requestAnimationFrame batching for text_delta events (matching useExecutionMonitor), and reads the JWT from localStorage under auth_token (matching the legacy pattern, which is the only token storage in the codebase). The WS base URL is derived from a NEXT_PUBLIC_WS_URL env var (matching the legacy websocket.ts pattern).

---

Implementation Steps

1. Add types to file:web-ui/src/types/index.ts

Append the following type definitions (do not modify existing exports):

• MessageRole — union of "user" | "assistant" | "tool_use" | "tool_result" | "thinking" | "system" | "error"
• ChatMessage — interface with id: string, role: MessageRole, content: string, toolName?: string, toolInput?: unknown, createdAt: string
• AgentChatStatus — union of "idle" | "connecting" | "thinking" | "streaming" | "error" | "disconnected"
• AgentChatState — interface with messages: ChatMessage[], status: AgentChatStatus, costUsd: number, inputTokens: number, outputTokens: number, error: string | null, connected: boolean

---

2. Create file:web-ui/src/hooks/useAgentChat.ts

Structure the hook in the following logical sections:

2a. Constants and helpers

• Define WS_BASE_URL from process.env.NEXT_PUBLIC_WS_URL || 'ws://localhost:8000'
• Define MAX_RETRIES = 5 and BASE_RETRY_DELAY_MS = 1000
• Define a getToken() helper that reads localStorage.getItem('auth_token') safely (guard typeof window !== 'undefined')
• Define a generateId() helper using crypto.randomUUID() (available in all modern browsers and Node 19+; no library needed)

2b. Refs (no re-render on mutation)

Declare the following refs inside the hook:

Ref	Type	Purpose
wsRef	WebSocket | null	The live socket instance
retriesRef	number	Current reconnect attempt count
retryTimerRef	ReturnType<typeof setTimeout> | null	Pending reconnect timer
pingTimerRef	ReturnType<typeof setInterval> | null	30-second keepalive interval
rafRef	number | null	Pending requestAnimationFrame handle
pendingDeltaRef	string	Accumulated text_delta content not yet flushed
inProgressIdRef	string | null	id of the current in-progress assistant message
stateRef	AgentChatState	Mirror of React state for use inside callbacks without stale closures

2c. React state

Single useState<AgentChatState> initialized to:

{ messages: [], status: 'idle', costUsd: 0, inputTokens: 0, outputTokens: 0, error: null, connected: false }

Keep stateRef.current in sync with every setState call via a wrapper updateState(patch) that merges the patch, updates stateRef.current, and calls setState.

2d. requestAnimationFrame flush for text_delta

Define scheduleDeltaFlush():

• If rafRef.current !== null, return early (already scheduled)
• Schedule requestAnimationFrame(() => { ... }):
  • Drain pendingDeltaRef.current into the in-progress assistant message in state
  • Set status to 'streaming' if not already
  • Reset rafRef.current = null

2e. WebSocket event handlers

Define these as stable functions using useCallback with empty deps (they read from refs, not props):

• handleOpen() — set connected: true, status: 'connecting' → 'idle', reset retriesRef.current = 0, start ping interval (setInterval every 30 000 ms sending { type: "ping" })
• handleMessage(event: MessageEvent) — parse JSON.parse(event.data) and dispatch by type:
  • text_delta: append event.delta to pendingDeltaRef.current; if inProgressIdRef.current is null, push a new assistant ChatMessage (with a fresh generateId()) and set inProgressIdRef.current; call scheduleDeltaFlush()
  • tool_use_start: flush any pending delta first; push a new tool_use ChatMessage with toolName and toolInput; set status: 'streaming'
  • tool_result: push a tool_result ChatMessage
  • thinking: push a thinking ChatMessage
  • cost_update: update costUsd, inputTokens, outputTokens
  • done: flush pending delta; set inProgressIdRef.current = null; set status: 'idle'
  • error: set error field and status: 'error'
• handleClose() — clear ping interval; set connected: false, status: 'disconnected'; trigger reconnect logic
• handleError() — log the error; handleClose() will fire immediately after for cleanup

2f. connect() function

Define connect() (not a hook, just a plain function called from the effect):

1. Guard: if sessionId is null, return
2. Set status: 'connecting'
3. Build URL: `${WS_BASE_URL}/ws/sessions/${sessionId}/chat?token=${encodeURIComponent(getToken() ?? '')}`
4. Instantiate new WebSocket(url), assign to wsRef.current
5. Attach onopen, onmessage, onclose, onerror handlers

2g. Reconnect logic (inside handleClose)

Mirror the pattern from useEventSource:

• If retriesRef.current < MAX_RETRIES: schedule setTimeout(connect, BASE_RETRY_DELAY_MS * 2 ** retriesRef.current), increment retriesRef.current
• Else: set status: 'error', error: 'Max reconnect attempts reached'

2h. disconnect() cleanup helper

• Cancel rafRef.current via cancelAnimationFrame
• Clear pingTimerRef.current via clearInterval
• Clear retryTimerRef.current via clearTimeout
• Close wsRef.current if open
• Null all refs

2i. Main useEffect

useEffect(() => {
  if (!sessionId) return;
  connect();
  return () => disconnect();
}, [sessionId]);

Also add a separate cleanup-only effect for unmount that calls disconnect() (handles the case where sessionId never changes but the component unmounts).

2j. Public API functions

• sendMessage(content: string):

  1. Guard: WS must be open (wsRef.current?.readyState === WebSocket.OPEN)
  2. Push an optimistic user ChatMessage immediately via updateState
  3. Set status: 'thinking'
  4. Send JSON.stringify({ type: "message", content }) over the socket

• interrupt():

  1. Guard: WS must be open
  2. Send JSON.stringify({ type: "interrupt" })

• clearMessages():

  1. Call updateState({ messages: [], error: null })

2k. Return value

return { state, sendMessage, interrupt, clearMessages };

---

3. Export from file:web-ui/src/hooks/index.ts

Add the following export line:

export { useAgentChat, type AgentChatState, type ChatMessage, type MessageRole } from './useAgentChat';

---

4. Status transition diagram

stateDiagram
  [*] --> idle
  idle --> connecting : sessionId set / connect()
  connecting --> idle : WS open (handleOpen)
  idle --> thinking : sendMessage()
  thinking --> streaming : first text_delta received
  streaming --> idle : done event
  idle --> disconnected : WS close (no retries left)
  connecting --> disconnected : WS close (no retries left)
  streaming --> disconnected : WS close (no retries left)
  disconnected --> connecting : auto-reconnect attempt
  idle --> error : error event
  disconnected --> error : max retries exceeded

Loading

---

5. Environment variable

Ensure NEXT_PUBLIC_WS_URL is documented in file:.env.example alongside the existing NEXT_PUBLIC_API_URL and NEXT_PUBLIC_SSE_URL entries, pointing to ws://localhost:8000 as the default.

Import In IDE

🤖 Prompt for AI Agents

## Observations

The current `web-ui` has no auth context — the v2 API client (`file:web-ui/src/lib/api.ts`) uses `withCredentials: true` (cookie-based auth), not a JWT Bearer token. The legacy `web-ui` stored a JWT in `localStorage` under `auth_token`. The existing hooks (`useEventSource`, `useExecutionMonitor`, `useTaskStream`) establish the patterns to follow: callback refs for stable closures, `useRef` + `requestAnimationFrame` batching for high-frequency updates, and exponential-backoff reconnect. No third-party WebSocket library is present in `file:web-ui/package.json`.

## Approach

Build `useAgentChat` as a self-contained hook in `file:web-ui/src/hooks/useAgentChat.ts`, mirroring the layered pattern of `useEventSource` → `useTaskStream` → `useExecutionMonitor`. The hook manages the raw `WebSocket` directly (no extra library needed), uses `useRef` + `requestAnimationFrame` batching for `text_delta` events (matching `useExecutionMonitor`), and reads the JWT from `localStorage` under `auth_token` (matching the legacy pattern, which is the only token storage in the codebase). The WS base URL is derived from a `NEXT_PUBLIC_WS_URL` env var (matching the legacy `websocket.ts` pattern).

---

## Implementation Steps

### 1. Add types to `file:web-ui/src/types/index.ts`

Append the following type definitions (do not modify existing exports):

- `MessageRole` — union of `"user" | "assistant" | "tool_use" | "tool_result" | "thinking" | "system" | "error"`
- `ChatMessage` — interface with `id: string`, `role: MessageRole`, `content: string`, `toolName?: string`, `toolInput?: unknown`, `createdAt: string`
- `AgentChatStatus` — union of `"idle" | "connecting" | "thinking" | "streaming" | "error" | "disconnected"`
- `AgentChatState` — interface with `messages: ChatMessage[]`, `status: AgentChatStatus`, `costUsd: number`, `inputTokens: number`, `outputTokens: number`, `error: string | null`, `connected: boolean`

---

### 2. Create `file:web-ui/src/hooks/useAgentChat.ts`

Structure the hook in the following logical sections:

#### 2a. Constants and helpers

- Define `WS_BASE_URL` from `process.env.NEXT_PUBLIC_WS_URL || 'ws://localhost:8000'`
- Define `MAX_RETRIES = 5` and `BASE_RETRY_DELAY_MS = 1000`
- Define a `getToken()` helper that reads `localStorage.getItem('auth_token')` safely (guard `typeof window !== 'undefined'`)
- Define a `generateId()` helper using `crypto.randomUUID()` (available in all modern browsers and Node 19+; no library needed)

#### 2b. Refs (no re-render on mutation)

Declare the following refs inside the hook:

| Ref | Type | Purpose |
|---|---|---|
| `wsRef` | `WebSocket \| null` | The live socket instance |
| `retriesRef` | `number` | Current reconnect attempt count |
| `retryTimerRef` | `ReturnType<typeof setTimeout> \| null` | Pending reconnect timer |
| `pingTimerRef` | `ReturnType<typeof setInterval> \| null` | 30-second keepalive interval |
| `rafRef` | `number \| null` | Pending `requestAnimationFrame` handle |
| `pendingDeltaRef` | `string` | Accumulated `text_delta` content not yet flushed |
| `inProgressIdRef` | `string \| null` | `id` of the current in-progress assistant message |
| `stateRef` | `AgentChatState` | Mirror of React state for use inside callbacks without stale closures |

#### 2c. React state

Single `useState<AgentChatState>` initialized to:
```
{ messages: [], status: 'idle', costUsd: 0, inputTokens: 0, outputTokens: 0, error: null, connected: false }
```

Keep `stateRef.current` in sync with every `setState` call via a wrapper `updateState(patch)` that merges the patch, updates `stateRef.current`, and calls `setState`.

#### 2d. `requestAnimationFrame` flush for `text_delta`

Define `scheduleDeltaFlush()`:
- If `rafRef.current !== null`, return early (already scheduled)
- Schedule `requestAnimationFrame(() => { ... })`:
  - Drain `pendingDeltaRef.current` into the in-progress assistant message in state
  - Set `status` to `'streaming'` if not already
  - Reset `rafRef.current = null`

#### 2e. WebSocket event handlers

Define these as stable functions using `useCallback` with empty deps (they read from refs, not props):

- **`handleOpen()`** — set `connected: true`, `status: 'connecting' → 'idle'`, reset `retriesRef.current = 0`, start ping interval (`setInterval` every 30 000 ms sending `{ type: "ping" }`)
- **`handleMessage(event: MessageEvent)`** — parse `JSON.parse(event.data)` and dispatch by `type`:
  - `text_delta`: append `event.delta` to `pendingDeltaRef.current`; if `inProgressIdRef.current` is null, push a new assistant `ChatMessage` (with a fresh `generateId()`) and set `inProgressIdRef.current`; call `scheduleDeltaFlush()`
  - `tool_use_start`: flush any pending delta first; push a new `tool_use` `ChatMessage` with `toolName` and `toolInput`; set `status: 'streaming'`
  - `tool_result`: push a `tool_result` `ChatMessage`
  - `thinking`: push a `thinking` `ChatMessage`
  - `cost_update`: update `costUsd`, `inputTokens`, `outputTokens`
  - `done`: flush pending delta; set `inProgressIdRef.current = null`; set `status: 'idle'`
  - `error`: set `error` field and `status: 'error'`
- **`handleClose()`** — clear ping interval; set `connected: false`, `status: 'disconnected'`; trigger reconnect logic
- **`handleError()`** — log the error; `handleClose()` will fire immediately after for cleanup

#### 2f. `connect()` function

Define `connect()` (not a hook, just a plain function called from the effect):
1. Guard: if `sessionId` is null, return
2. Set `status: 'connecting'`
3. Build URL: `` `${WS_BASE_URL}/ws/sessions/${sessionId}/chat?token=${encodeURIComponent(getToken() ?? '')}` ``
4. Instantiate `new WebSocket(url)`, assign to `wsRef.current`
5. Attach `onopen`, `onmessage`, `onclose`, `onerror` handlers

#### 2g. Reconnect logic (inside `handleClose`)

Mirror the pattern from `useEventSource`:
- If `retriesRef.current < MAX_RETRIES`: schedule `setTimeout(connect, BASE_RETRY_DELAY_MS * 2 ** retriesRef.current)`, increment `retriesRef.current`
- Else: set `status: 'error'`, `error: 'Max reconnect attempts reached'`

#### 2h. `disconnect()` cleanup helper

- Cancel `rafRef.current` via `cancelAnimationFrame`
- Clear `pingTimerRef.current` via `clearInterval`
- Clear `retryTimerRef.current` via `clearTimeout`
- Close `wsRef.current` if open
- Null all refs

#### 2i. Main `useEffect`

```
useEffect(() => {
  if (!sessionId) return;
  connect();
  return () => disconnect();
}, [sessionId]);
```

Also add a separate cleanup-only effect for unmount that calls `disconnect()` (handles the case where `sessionId` never changes but the component unmounts).

#### 2j. Public API functions

- **`sendMessage(content: string)`**:
  1. Guard: WS must be open (`wsRef.current?.readyState === WebSocket.OPEN`)
  2. Push an optimistic `user` `ChatMessage` immediately via `updateState`
  3. Set `status: 'thinking'`
  4. Send `JSON.stringify({ type: "message", content })` over the socket

- **`interrupt()`**:
  1. Guard: WS must be open
  2. Send `JSON.stringify({ type: "interrupt" })`

- **`clearMessages()`**:
  1. Call `updateState({ messages: [], error: null })`

#### 2k. Return value

```ts
return { state, sendMessage, interrupt, clearMessages };
```

---

### 3. Export from `file:web-ui/src/hooks/index.ts`

Add the following export line:

```
export { useAgentChat, type AgentChatState, type ChatMessage, type MessageRole } from './useAgentChat';
```

---

### 4. Status transition diagram

```mermaid
stateDiagram
  [*] --> idle
  idle --> connecting : sessionId set / connect()
  connecting --> idle : WS open (handleOpen)
  idle --> thinking : sendMessage()
  thinking --> streaming : first text_delta received
  streaming --> idle : done event
  idle --> disconnected : WS close (no retries left)
  connecting --> disconnected : WS close (no retries left)
  streaming --> disconnected : WS close (no retries left)
  disconnected --> connecting : auto-reconnect attempt
  idle --> error : error event
  disconnected --> error : max retries exceeded
```

---

### 5. Environment variable

Ensure `NEXT_PUBLIC_WS_URL` is documented in `file:.env.example` alongside the existing `NEXT_PUBLIC_API_URL` and `NEXT_PUBLIC_SSE_URL` entries, pointing to `ws://localhost:8000` as the default.

---

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

#228 - fix: prevent duplicate task generation button for late-joining users [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!