How useAgentSession Manages Loading and Streaming State in pi-web

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, 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).

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

// 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 manages the server-side AgentSession objects that generate these events.

Summary

• useAgentSession in 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, 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.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}