# How useAgentSession Manages Loading and Streaming State in pi-web

> Discover how useAgentSession orchestrates loading and streaming state in pi-web. Learn about its finite-state reducer for real-time SSE message management.

- Repository: [Alex Yang/pi-web](https://github.com/agegr/pi-web)
- Tags: internals
- Published: 2026-07-31

---

**useAgentSession** is the core React hook that coordinates session hydration, Server-Sent Events (SSE) consumption, and a finite-state reducer to manage real-time message streaming in the pi-web chat application.

In the **pi-web** repository, `useAgentSession` serves as the central nervous system for interactive chat sessions. Located at [`hooks/useAgentSession.ts`](https://github.com/agegr/pi-web/blob/main/hooks/useAgentSession.ts), this hook orchestrates the retrieval of conversation history, maintains resilient connections to the agent backend, and manages complex UI state transitions during streaming responses. Its architecture ensures the chat interface remains synchronized even under flaky network conditions.

## Loading Static and Live Session Data

When a component mounts, `useAgentSession` initiates hydration by calling `loadSession` (lines 40-78). This function performs parallel fetches to reconstruct the conversation state from two sources:

*   **Static session file** (`/api/sessions/:id`): Populates `messages`, `entryIds`, `tree`, and `leafId` (lines 44-63).
*   **Live agent state** (`/api/sessions/:id/state`): Updates `contextUsage`, `systemPrompt`, `thinkingLevel`, and queued messages (lines 73-89).

If the session file is missing and returns a **404**, the hook clears the UI state to prevent stale data (lines 47-53). For background synchronization, `loadContext` (lines 102-113) provides a lightweight refresh mechanism that updates only the volatile agent configuration without reloading the entire message history.

## Managing Server-Sent Event Connections

The hook maintains a persistent connection to the backend via `EventSource` to receive real-time updates.

### Connection Initialization

`ensureEventsConnected` (lines 125-133) creates an `EventSource` pointing to `/api/agent/[id]/events`. The first `connected` event resolves the internal connection promise, signaling that the channel is ready to receive events.

### Reconnection and Error Handling

When the connection drops, the hook distinguishes between fatal closures and recoverable errors (lines 150-166). If the agent is still running and the closure is accidental, `ensureEventsConnected` automatically re-establishes the stream. The hook stores the active `EventSource` in `eventSourceRef` and cleans it up on unmount via `closeEvents` (lines 120-124).

## The Streaming State Machine

At the heart of the streaming logic sits `streamReducer` (lines 31-49), a finite-state machine that tracks whether the agent is actively generating content and holds the interim `streamingMessage`.

### Reducer States and Transitions

The reducer handles four distinct action types:

*   **`start`** → Sets `isStreaming: true` and initializes `streamingMessage` to null.
*   **`update`** → Preserves `isStreaming: true` while merging the latest partial message payload.
*   **`end`** / **`reset`** → Clears the streaming flag and finalizes the message buffer.

### Event Processing Pipeline

`handleAgentEvent` (lines 210-262) parses incoming SSE payloads and routes them to the appropriate state updates:

| Event | State Transition |
|-------|------------------|
| `agent_start` | Sets `agentRunning` true, pushes `waiting_model` phase, dispatches `streamReducer` start. |
| `message_start` / `message_update` | Dispatches `update` actions to keep `streamingMessage` current (lines 263-308). |
| `message_end` | Inserts the completed message into `messages`, resets streaming via `dispatch({type:"reset"})`, and returns to `waiting_model` phase. |
| `agent_end` | Marks the run finished, reloads the session file, and clears streaming flags. |
| `tool_execution_*` | Adjusts `agentPhase` to `running_tools` and tracks active tool calls. |

This architecture ensures the UI reflects the exact state of the backend generation process, from initial token to final completion.

## Prompt Lifecycle and Network Resilience

### Initiating a Prompt Run

When a user submits a message, `handleSend` (lines 665-724) generates a unique **run ID** (`promptRunIdRef`) to disambiguate concurrent runs. It sets `agentRunning` to true, starts the streaming state machine (`dispatch({type:"start"})`), and transmits the prompt via `sendAgentCommand` (defined in [`lib/agent-client.ts`](https://github.com/agegr/pi-web/blob/main/lib/agent-client.ts)).

### State Reconciliation and Polling

If the SSE connection drops during generation, `reconcileAgentState` (lines 851-874) implements a fallback polling mechanism. It queries `/api/agent/:id` to check if the server reports idle (`!state.isStreaming && !state.isPromptRunning`). When the server confirms completion, `finishPromptWithoutStream` cleans up the UI state, ensuring the chat never permanently stalls.

### Background Synchronization

A periodic `setInterval` (every `AGENT_STATE_RECONCILE_MS` = 15 seconds) combined with `visibilitychange` and `online` event listeners (lines 882-904) triggers automatic reconciliation when users return to the tab or restore connectivity. Additionally, `waitForPromptSettlement` (lines 94-115) ensures the UI settles cleanly after the server finishes processing a run, even if the final SSE event was lost.

## Session Control Operations

Beyond streaming and loading, `useAgentSession` exposes imperative controls for runtime session management.

### Aborting Active Runs

`handleAbort` (lines 202-218) transmits `abort` or `abort_bash` commands to the backend and immediately clears the local streaming state, allowing users to terminate long-running operations.

### Forking and Navigation

`handleFork` (lines 220-237) issues a `fork` command to the server, receives a new session ID, and notifies the parent component to transition to the branched conversation. `handleNavigate` (lines 239-247) updates the `leafId` and reloads the context when users navigate through message history branches.

### Implementation Example

```typescript
// hooks/useAgentSession.ts
import { useAgentSession } from '@/hooks/useAgentSession';

export default function ChatTab({ session }) {
  const { 
    messages, 
    agentRunning, 
    handleSend, 
    handleAbort, 
    handleFork,
    reconcileAgentState 
  } = useAgentSession({
    session,
    newSessionCwd: null,
    onAgentEnd: () => console.log('run finished'),
  });

  return (
    <div>
      <ChatWindow
        messages={messages}
        isRunning={agentRunning}
        onSend={handleSend}
        onAbort={handleAbort}
        onFork={handleFork}
      />
      <button onClick={() => reconcileAgentState(session.id)}>
        Force Refresh
      </button>
    </div>
  );
}

```

The SSE endpoint is implemented in `app/api/agent/[id]/events/route.ts`, while [`lib/rpc-manager.ts`](https://github.com/agegr/pi-web/blob/main/lib/rpc-manager.ts) manages the server-side `AgentSession` objects that generate these events.

## Summary

*   **useAgentSession** in [`hooks/useAgentSession.ts`](https://github.com/agegr/pi-web/blob/main/hooks/useAgentSession.ts) centralizes session hydration, SSE handling, and streaming state management for the pi-web chat interface.
*   **Dual loading strategy**: `loadSession` (lines 40-78) fetches static message history while `loadContext` (lines 102-113) refreshes live agent configuration.
*   **Resilient event streaming**: `ensureEventsConnected` (lines 125-166) maintains the `EventSource` connection with automatic reconnection logic.
*   **Finite-state reducer**: `streamReducer` (lines 31-49) provides deterministic state transitions for `isStreaming` and `streamingMessage` during content generation.
*   **Graceful degradation**: `reconcileAgentState` (lines 851-874) and periodic polling (lines 882-904) ensure UI synchronization even when the SSE connection fails.

## Frequently Asked Questions

### How does useAgentSession handle network interruptions during streaming?

If the `EventSource` connection drops while the agent is running, `useAgentSession` falls back to polling via `reconcileAgentState` (lines 851-874). This function queries the agent status endpoint to detect when the server becomes idle, then calls `finishPromptWithoutStream` to synchronize the UI. A periodic timer (every 15 seconds) and browser event listeners for `visibilitychange` and `online` automatically trigger this reconciliation.

### What is the purpose of the streamReducer in useAgentSession?

`streamReducer` (lines 31-49) implements a finite-state machine for tracking streaming message progress. It manages three core states: `start` (initializing the stream), `update` (appending partial content), and `reset` (finalizing the message). This ensures that `isStreaming` and `streamingMessage` update atomically and predictably as SSE events arrive via `handleAgentEvent`.

### How does useAgentSession distinguish between loading historical messages and live updates?

The hook uses `loadSession` (lines 40-78) for initial hydration, which fetches both the static JSONL message file and the live agent state in parallel. For subsequent lightweight updates, `loadContext` (lines 102-113) refreshes only the volatile configuration (context usage, thinking level) without reloading the entire message tree.

### Where is the SSE endpoint defined that useAgentSession connects to?

The client-side connection logic resides in [`hooks/useAgentSession.ts`](https://github.com/agegr/pi-web/blob/main/hooks/useAgentSession.ts), specifically in `ensureEventsConnected` (lines 125-133). The corresponding server-side endpoint is implemented at `app/api/agent/[id]/events/route.ts`, which streams events from the `AgentSession` objects managed by [`lib/rpc-manager.ts`](https://github.com/agegr/pi-web/blob/main/lib/rpc-manager.ts).