# Pi Agent Architecture - Top-Down Overview ## Executive Summary The Pi Agent is a **stateful, event-driven agent framework** built in TypeScript. It provides: 1. **Core Agent** - Low-level agent loop with message/tool streaming 2. **Agent Harness** - High-level session management with persistence, branching, and compaction Both layers follow the **same core pattern**: stream LLM response → execute tools → emit events → repeat. --- ## Architecture Layers ``` ┌─────────────────────────────────────────────────────────────────────────────────────┐ │ APPLICATION LAYER │ │ • Creates Agent/AgentHarness instances │ │ • Subscribes to events for UI updates │ │ • Provides tools and model configuration │ └─────────────────────────────────────────────────────────────────────────────────────┘ │ ┌─────────────────────────────┼─────────────────────────────┐ │ │ │ ▼ ▼ ▼ ┌────────────────────────┐ ┌────────────────────────────────┐ ┌─────────────────┐ │ Agent (Core) │ │ Agent Harness (High-Level) │ │ Agent-Loop │ │ │ │ │ │ │ │ • State management │ │ • Session persistence │ │ • Turn │ │ • Event streaming │ │ • Branching/compaction │ │ • Tool exec │ │ • Steering/follow-up │ │ • Skills/templates │ │ • Message │ │ queues │ │ • Tool context binding │ │ streaming │ │ • Hook system │ │ • State snapshots │ │ │ └────────────────────────┘ └────────────────────────────────┘ └─────────────────┘ │ │ ▼ ▼ ┌──────────────────────────┐ ┌─────────────────┐ │ LLM Provider API │ │ Session Repo │ │ (via @earendil-works) │ │ (JSONL/ │ └──────────────────────────┘ │ Memory) │ └─────────────────┘ ``` --- ## Core Concepts ### 1. AgentMessage ```typescript type AgentMessage = Message | CustomAgentMessages[keyof CustomAgentMessages] ``` The unified message type that combines: - **LLM messages**: `user`, `assistant`, `toolResult` (from pi-ai) - **Custom messages**: Application-specific types (via declaration merging) ### 2. AgentEvent ```typescript type AgentEvent = | { type: "agent_start" } | { type: "agent_end"; messages: AgentMessage[] } | { type: "turn_start" } | { type: "turn_end"; message: AgentMessage; toolResults: ToolResultMessage[] } | { type: "message_start"; message: AgentMessage } | { type: "message_update"; message: AgentMessage; assistantMessageEvent: AssistantMessageEvent } | { type: "message_end"; message: AgentMessage } | { type: "tool_execution_start"; toolCallId: string; toolName: string; args: any } | { type: "tool_execution_update"; toolCallId: string; toolName: string; args: any; partialResult: any } | { type: "tool_execution_end"; toolCallId: string; toolName: string; result: any; isError: boolean } ``` **Event Flow per Turn:** ``` turn_start message_start (user prompt) message_end message_start (assistant streaming) message_update (multiple - as chunks arrive) message_end tool_execution_start (if tool calls present) tool_execution_update (if tool streams partial results) tool_execution_end turn_end ``` ### 3. AgentTool ```typescript interface AgentTool { name: string; label: string; description: string; parameters: TSchema; execute( toolCallId: string, params: Static, signal?: AbortSignal, onUpdate?: AgentToolUpdateCallback ): Promise>; } ``` --- ## Process Flow Diagrams ### Prompt Flow (High-Level) ``` User Input │ ▼ Agent.prompt("Hello") │ ├─► normalizePromptInput() → AgentMessage[] │ ├─► runWithLifecycle() │ ├─► Set isStreaming=true │ └─► Create abort controller │ ▼ runAgentLoop() │ ├─► Emit: agent_start ├─► Emit: turn_start ├─► Emit: message_start/end (prompts) │ ▼ runLoop() - Main Loop │ ├─► Check steering queue (drain if any) ├─► Check follow-up queue (skip if not first turn) │ ▼ streamAssistantResponse() │ ├─► transformContext() [optional] ├─► convertToLlm() → Message[] ├─► Build Context {systemPrompt, messages, tools} ├─► Resolve API key ├─► Call streamFn(model, context, options) │ ▼ Assistant Message Stream │ ├─► message_start (assistant) ├─► message_update (text chunks) ├─► message_update (toolCall blocks) ├─► message_end │ ▼ executeToolCalls() │ ├─► Check if sequential/parallel execution ├─► For each tool call: │ ├─► prepareToolCall() │ │ ├─► Find tool by name │ │ ├─► Validate arguments │ │ └─► beforeToolCall() hook │ │ │ ├─► executePreparedToolCall() │ │ └─► tool.execute() with onUpdate callback │ │ │ └─► finalizeExecutedToolCall() │ └─► afterToolCall() hook │ ├─► Emit: tool_execution_start/update/end └─► Emit: message_start/end (toolResult) │ ▼ turn_end │ ├─► Check prepareNextTurn hook ├─► Check shouldStopAfterTurn hook ├─► Drain steering queue └─► Drain follow-up queue │ ├─► If steering/follow-up exists → repeat loop └─► If no more messages → agent_end ``` ### Tool Execution Flow (Detailed) ``` Tool Call from LLM │ ▼ prepareToolCall() │ ├─► Find tool in currentContext.tools │ └─► If not found → immediate error │ ├─► prepareToolCallArguments() [optional] │ ├─► validateToolArguments() │ └─► If invalid → immediate error │ └─► beforeToolCall() hook ├─► Return {block: true} → error └─► Continue │ ▼ executePreparedToolCall() │ ├─► Call tool.execute() with onUpdate callback │ └─► tool calls onUpdate(partialResult) during execution │ ├─► onUpdate() → emit tool_execution_update └─► Return {result, isError} │ ▼ finalizeExecutedToolCall() │ └─► afterToolCall() hook ├─► Override content/details/usage/terminate └─► Return {toolCall, result, isError} │ ▼ emitToolExecutionEnd() │ └─► Emit: tool_execution_end │ ▼ createToolResultMessage() │ └─► Create ToolResultMessage with: ├─► toolCallId ├─► toolName ├─► content ├─► details ├─► usage └─► isError │ ▼ emitToolResultMessage() │ ├─► Emit: message_start └─► Emit: message_end ``` ### Session Persistence Flow ``` AgentHarness.handleAgentEvent() │ ├─► message_end → session.appendMessage() │ └─► Storage: write entry to JSONL file │ ├─► turn_end → flushPendingSessionWrites() │ ├─► Write all pending entries │ ├─► Emit: save_point │ └─► session.getStorage().setLeafId() │ └─► agent_end → flushPendingSessionWrites() ├─► Write leaf entry pointing to last message └─► Emit: settled │ ▼ Session Tree Structure: root ├─► message (user prompt #1) ├─► message (assistant #1) ├─► tool_result (result #1) ├─► turn_end ├─► message (user prompt #2) ├─► message (assistant #2) ├─► compaction (summary of history) ├─► message (assistant continues) └─► leaf → points to current head ``` --- ## Hook System ### Agent-Level Hooks (agent-loop.ts) ```typescript interface AgentLoopConfig { // Message transformation convertToLlm: (messages: AgentMessage[]) => Message[] transformContext?: (messages: AgentMessage[]) => AgentMessage[] // Lifecycle hooks beforeToolCall?: (context: BeforeToolCallContext) => BeforeToolCallResult afterToolCall?: (context: AfterToolCallContext) => AfterToolCallResult shouldStopAfterTurn?: (context: ShouldStopAfterTurnContext) => boolean prepareNextTurn?: (context: PrepareNextTurnContext) => AgentLoopTurnUpdate // Queue draining getSteeringMessages?: () => AgentMessage[] getFollowUpMessages?: () => AgentMessage[] } ``` ### Harness-Level Hooks (agent-harness.ts) ```typescript // Hook types in AgentHarnessEventResultMap: type HookName = | "before_agent_start" | "context" | "tool_call" | "tool_result" | "session_before_compact" | "session_before_tree" | "before_provider_request" | "before_provider_payload" ``` **Hook Execution Order per Turn:** ``` 1. before_agent_start (harness) 2. context (harness) → transformContext 3. streamAssistantResponse ├─► Before provider request (harness) ├─► convertToLlm (agent) └─► LLM call 4. For each tool call: ├─► tool_call (harness) → beforeToolCall ├─► Execute tool └─► tool_result (harness) → afterToolCall 5. turn_end 6. shouldStopAfterTurn (agent) 7. prepareNextTurn (agent) 8. Drain steering/follow-up queues ``` --- ## Data Flow Summary ``` ┌────────────────────────────────────────────────────────────────────────────────┐ │ AGENT LIFECYCLE - DATA FLOW │ ├────────────────────────────────────────────────────────────────────────────────┤ │ 1. INPUT │ │ • prompt("Hello") → normalizePromptInput() │ │ → AgentMessage[] │ │ 2. INITIATE │ │ • createMutableAgentState() │ │ • runWithLifecycle() │ │ 3. LOOP CONTROL │ │ • runLoop() │ │ ├─► Steering queue? → drain and inject │ │ └─► Follow-up queue? (after first turn) │ │ 4. LLM STREAMING │ │ • transformContext() [optional] │ │ • convertToLlm() │ │ • streamFn() │ │ → AssistantMessage stream (text + toolCalls) │ │ 5. TOOL EXECUTION │ │ • executeToolCalls() │ │ ├─► prepareToolCall() │ │ │ ├─► beforeToolCall() hook │ │ │ └─► Validate args │ │ ├─► executePreparedToolCall() │ │ │ └─► tool.execute() │ │ └─► finalizeExecutedToolCall() │ │ └─► afterToolCall() hook │ │ 6. UPDATE STATE │ │ • Push assistant message to state.messages │ │ • Push toolResult messages to state.messages │ │ 7. TERMINATION CHECK │ │ • shouldStopAfterTurn? → exit │ │ • prepareNextTurn? → update context/model │ │ • Drain steering/follow-up → continue │ │ 8. FINISH │ │ • emit agent_end │ │ • finishRun() → reset isStreaming │ └────────────────────────────────────────────────────────────────────────────────┘ ``` --- ## Key Design Patterns ### 1. Event-Driven Architecture - All external communication via `AgentEvent` stream - Hooks can be async and are awaited in order - Abort signal propagated through all operations ### 2. State Isolation - `AgentState` is read-only externally - `AgentHarness` snapshots state per turn - Context transforms return new arrays (immutability) ### 3. Layered Abstraction ``` Low-level (agent-loop.ts) • Pure async iteration • No session management • No tool context binding High-level (agent-harness.ts) • Session persistence • Branching/compaction • Hook system for customization ``` ### 4. Extensibility Points - **Custom messages**: Extend `CustomAgentMessages` interface - **Custom hooks**: Add handlers via `subscribe()`/`on()` - **Tool context**: Pass `toolContext` to harness constructor - **Storage**: Implement `SessionStorage` interface --- ## Learning Path 1. **Start with types.ts** - Understand `AgentMessage`, `AgentEvent`, `AgentTool` 2. **Read agent-loop.ts** - See how messages flow through the loop 3. **Study agent.ts** - See how Agent wraps the loop with state management 4. **Read agent-harness.ts** - See how session management hooks into the loop 5. **Explore session/* files** - Understand persistence and branching 6. **Study tools/* files** - See concrete tool implementations --- ## Next Steps See individual markdown files in this folder for: - `AGENT-LOOP-DETAILED.md` - Deep dive into the agent loop - `HOOK-SYSTEM.md` - Complete hook documentation - `SESSION-ARCHITECTURE.md` - Session persistence details - `TOOL-EXECUTION.md` - Tool execution mechanics