# Pi Agent Architecture - Visual Diagrams ## 1. System Architecture ``` ┌─────────────────────────────────────────────────────────────────────────────────────────────┐ │ APPLICATION LAYER │ │ │ │ ┌──────────────────┐ ┌──────────────────┐ ┌──────────────────┐ │ │ │ Agent User │ │ AgentHarness │ │ AgentHarness │ │ │ │ (Low-Level) │ │ (High-Level) │ │ (Custom App) │ │ │ └────────┬─────────┘ └────────┬─────────┘ └────────┬─────────┘ │ │ │ │ │ │ │ └────────┬───────────────┴───────────────────────┬┘ │ │ │ │ │ │ ▼ ▼ │ │ ┌──────────────────┐ ┌──────────────────┐ │ │ │ Agent Core │ │ AgentHarness │ │ │ │ │ │ │ │ │ │ • State mgmt │ │ • Session │ │ │ │ • Event stream │ │ • Compaction │ │ │ │ • Queue mgmt │ │ • Branching │ │ │ │ • Hook system │ │ • Skills │ │ │ └────────┬─────────┘ └────────┬─────────┘ │ └────────────────────┼─────────────────────────────────────┼─────────────────────────────────┘ │ │ ┌────────────┴────────────┐ ┌──────────────┴──────────────┐ │ │ │ │ ▼ ▼ ▼ ▼ ┌──────────────────┐ ┌──────────────────┐ ┌──────────────────┐ ┌──────────────────┐ │ Agent Loop │ │ Agent Context │ │ Agent State │ │ Agent Event │ │ │ │ │ │ │ │ │ │ • runAgentLoop │ │ • Messages │ │ • Tools │ │ • agent_start │ │ • runLoop │ │ • System prompt │ │ • Messages │ │ • agent_end │ │ • streamResponse │ │ • Tools │ │ • isStreaming │ │ • turn_start │ │ • executeTools │ │ │ │ • pendingCalls │ │ • turn_end │ └────────┬─────────┘ └──────────────────┘ └──────────────────┘ │ • message_start │ │ │ • message_update │ ▼ │ • message_end │ ┌───────────────────────────────────────────────────────────────────────▼───────────────────┐ │ AGENT CORE (agent.ts, agent-loop.ts) │ └─────────────────────────────────────────────────────────────────────────────────────────────┘ ┌─────────────────────────────────────────────────────────────────────────────────────────────┐ │ SESSION LAYER │ │ │ │ ┌──────────────────┐ ┌──────────────────┐ ┌──────────────────┐ │ │ │ Session │ │ SessionStorage │ │ SessionRepo │ │ │ │ │ │ │ │ │ │ │ │ • Tree structure │ │ • Memory │ │ • Create │ │ │ │ • Context build │ │ • JSONL │ │ • Open │ │ │ │ • Branching │ │ │ │ • List │ │ │ │ • Compaction │ │ │ │ • Fork │ │ │ └────────┬─────────┘ └──────────────────┘ └──────────────────┘ │ │ │ │ │ ▼ │ │ ┌───────────────────────────────────────────────────────────────────────────────────┐ │ │ │ Session Tree │ │ │ │ │ │ │ │ root (null) │ │ │ │ ├─ message [id:1] ← User prompt │ │ │ │ │ └─ message [id:2] ← Assistant response │ │ │ │ │ └─ tool_result [id:3] ← Tool call result │ │ │ │ │ └─ message [id:4] ← User continuation │ │ │ │ │ └─ compaction [id:5] ← History summarized │ │ │ │ │ ├─ retained: [msg6, msg7] ← Recent messages kept │ │ │ │ │ └─ message [id:8] ← After compaction │ │ │ │ │ └─ leaf [id:9] ← Current head (cursor) │ │ │ │ │ │ │ │ │ └─ branch_summary [id:10] ← Branch point with summary │ │ │ │ └─ message [id:11] ← New branch message │ │ │ │ └─ leaf [id:12] ← New branch head │ │ │ │ │ │ │ └────────────────────────────────────────────────────────────────────────────────────┘ │ └─────────────────────────────────────────────────────────────────────────────────────────────┘ ┌─────────────────────────────────────────────────────────────────────────────────────────────┐ │ LLM PROVIDER LAYER │ │ │ │ ┌──────────────────┐ ┌──────────────────┐ ┌──────────────────┐ │ │ │ StreamFn │ │ Models API │ │ Provider API │ │ │ │ │ │ │ │ │ │ │ │ • streamSimple │ │ • completeSimple │ │ • OpenAI │ │ │ │ • completeSimple │ │ • Models catalog │ │ • Anthropic │ │ │ │ │ │ │ │ • Custom │ │ │ └──────────────────┘ └──────────────────┘ └──────────────────┘ │ │ │ └─────────────────────────────────────────────────────────────────────────────────────────────┘ ``` --- ## 2. Message Flow Diagram ``` ┌─────────────────────────────────────────────────────────────────────────────────────────────┐ │ PROMPT FLOW │ └─────────────────────────────────────────────────────────────────────────────────────────────┘ User Input │ ├─► string: "Build a web app" │ ├─► AgentMessage: { role: "user", content: [...] } │ └─► AgentMessage[]: [{...}, {...}] │ ▼ Agent.prompt(input) │ ├─► normalizePromptInput() │ ├─► string → { role: "user", content: [{ type: "text", text: input }] } │ ├─► AgentMessage → [message] │ └─► AgentMessage[] → messages │ └─► runPromptMessages() │ └─► runWithLifecycle() │ ├─► Set isStreaming = true ├─► Create abort controller │ └─► runAgentLoop() │ ├─► emit: agent_start ├─► emit: turn_start ├─► emit: message_start (user prompt) ├─► emit: message_end (user prompt) │ └─► runLoop() │ ├─► Check steering queue (drain if any) ├─► Check follow-up queue (skip if first turn) │ └─► streamAssistantResponse() │ ├─► transformContext() [optional] │ └─► AgentMessage[] → AgentMessage[] │ ├─► convertToLlm() │ └─► AgentMessage[] → Message[] │ ├─► Build LLM Context │ └─► { systemPrompt, messages, tools } │ ├─► Resolve API key (from hook) │ └─► Call streamFn() │ ├─► LLM Provider API │ └─► AssistantMessageEventStream │ ├─► message_start (assistant) ├─► message_update (text chunk 1) ├─► message_update (text chunk 2) ├─► message_update (toolCall) └─► message_end (assistant) │ └─► executeToolCalls() │ ├─► Sequential mode: tool calls one-by-one │ └─► Parallel mode: tool calls concurrently │ ├─► prepareToolCall() │ ├─► Find tool by name │ ├─► prepareArguments() [optional] │ ├─► validateToolArguments() │ └─► beforeToolCall() hook │ ├─► Return {block: true, reason} │ └─► Return undefined │ ├─► executePreparedToolCall() │ ├─► onUpdate(partialResult) [streaming updates] │ └─► tool.execute() │ └─► finalizeExecutedToolCall() └─► afterToolCall() hook ├─► Override: content, details, isError, usage └─► Use executed result │ └─► Emit: tool_execution_start/update/end │ └─► Create ToolResultMessage │ └─► Emit: message_start/end (toolResult) │ └─► turn_end │ ├─► prepareNextTurn() hook │ └─► Return: context/model/thinkingLevel │ ├─► shouldStopAfterTurn() hook │ └─► Return: boolean │ ├─► Drain steering queue │ └─► getSteeringMessages() → inject │ └─► Drain follow-up queue └─► getFollowUpMessages() → inject │ ├─► Steering/follow-up exists? → Repeat from streamAssistantResponse() └─► No more messages → emit: agent_end │ └─► finishRun() └─► isStreaming = false ┌─────────────────────────────────────────────────────────────────────────────────────────────┐ │ CONTINUATION FLOW │ └─────────────────────────────────────────────────────────────────────────────────────────────┘ Agent.continue() │ ├─► Validate last message (must be user/toolResult) │ └─► runAgentLoopContinue() │ └─► runLoop() from current context (no new prompts) │ └─► Same flow as above, starting from current context ``` --- ## 3. Hook System Flow ``` ┌─────────────────────────────────────────────────────────────────────────────────────────────┐ │ HOOK EXECUTION ORDER │ └─────────────────────────────────────────────────────────────────────────────────────────────┘ AgentHarness.prompt() │ ├─► before_agent_start (harness hook) │ └─► Can return: messages, systemPrompt │ ├─► transformContext() (agent hook) │ └─► AgentMessage[] → AgentMessage[] │ └─► streamAssistantResponse() │ ├─► before_provider_request (harness hook) │ └─► Can modify: streamOptions │ ├─► convertToLlm() (agent hook) │ └─► AgentMessage[] → Message[] │ ├─► streamFn() │ └─► message_end (assistant) │ └─► executeToolCalls() │ ├─► For each tool call: │ │ ├─► tool_call (harness hook) │ │ └─► Can return: block, reason │ │ │ ├─► executePreparedToolCall() │ │ │ └─► tool_result (harness hook) │ └─► Can return: content, details, isError, usage, terminate │ └─► turn_end │ ├─► shouldStopAfterTurn() (agent hook) │ └─► Return: boolean │ ├─► prepareNextTurn() (agent hook) │ └─► Return: context/model/thinkingLevel │ ├─► Drain steering queue │ └─► getSteeringMessages() (agent hook) │ └─► Drain follow-up queue └─► getFollowUpMessages() (agent hook) │ ├─► Continue? → Repeat from streamAssistantResponse() └─► Stop? → agent_end (harness hook) ``` --- ## 4. Tool Execution Flow ``` ┌─────────────────────────────────────────────────────────────────────────────────────────────┐ │ TOOL EXECUTION FLOW │ └─────────────────────────────────────────────────────────────────────────────────────────────┘ Assistant Message with Tool Call │ ├─► { type: "toolCall", id: "tc_123", name: "bash", arguments: { command: "ls" } } │ ▼ prepareToolCall() │ ├─► Find tool in currentContext.tools │ └─► Not found? → immediate error result │ ├─► prepareToolCallArguments() [optional shim] │ └─► Transform arguments before validation │ ├─► validateToolArguments() │ └─► Validate against tool parameters schema │ └─► beforeToolCall() hook │ ├─► Return { block: true, reason: "..." } │ └─► Emit: tool_execution_start/update/end (error) │ └─► Tool NOT executed │ └─► Return undefined │ ▼ executePreparedToolCall() │ ├─► tool.execute(toolCallId, validatedArgs, signal, onUpdate) │ │ │ ├─► Long-running operation │ │ └─► onUpdate({ content: [...], details: {...} }) │ │ └─► Emit: tool_execution_update │ │ │ └─► Return: { content, details, usage, ... } │ └─► Return: { result, isError } │ ▼ finalizeExecutedToolCall() │ └─► afterToolCall() hook │ ├─► Return override: { content, details, isError, usage, terminate } │ └─► Merge: result = { ...result, ...override } │ └─► Return: { toolCall, result, isError } │ ▼ emitToolExecutionEnd() │ └─► Emit: tool_execution_end │ ▼ createToolResultMessage() │ └─► Create ToolResultMessage ├─► toolCallId: tc_123 ├─► toolName: bash ├─► content: result.content ├─► details: result.details ├─► usage: result.usage ├─► isError: result.isError └─► timestamp: Date.now() │ ▼ emitToolResultMessage() │ ├─► Emit: message_start (toolResult) └─► Emit: message_end (toolResult) │ ▼ Push to context.messages │ ▼ Available for next LLM call ``` --- ## 5. Session Tree Navigation ``` ┌─────────────────────────────────────────────────────────────────────────────────────────────┐ │ SESSION BRANCHING │ └─────────────────────────────────────────────────────────────────────────────────────────────┘ Original Session Tree: │ ├─ root │ └─ message [user #1] [id: 1] │ └─ message [assistant #1] [id: 2] │ └─ tool_result [id: 3] │ └─ message [user #2] [id: 4] │ └─ leaf [id: 5] ← Current head │ ▼ Navigate to entry [id: 2] with summarize=true │ ├─► Collect entries from leaf to target │ └─► [leaf, msg4, tool_result, msg2] (path) │ ├─► Common ancestor: root │ ├─► Entries to summarize: [msg4, tool_result] │ ├─► Generate branch summary via LLM │ ├─► Create branch_summary entry │ └─► { type: "branch_summary", summary: "...", fromId: 2 } │ └─► Fork session at target [id: 2] │ ├─► Clone entries up to target │ └─► [root, msg1, msg2, branch_summary] │ └─► Set new leaf to [id: 2] │ ▼ New Session Tree: │ ├─ root │ └─ message [user #1] [id: 1] │ └─ message [assistant #1] [id: 2] │ └─ branch_summary [id: 6] ← New branch point │ └─ leaf [id: 7] ← New head │ └─ Original branch (still exists) └─ message [user #2] [id: 4] └─ tool_result [id: 3] └─ leaf [id: 5] ← Old head ``` --- ## 6. Context Window Compaction ``` ┌─────────────────────────────────────────────────────────────────────────────────────────────┐ │ CONTEXT COMPACTION │ └─────────────────────────────────────────────────────────────────────────────────────────────┘ Original Context (10,000 tokens): │ ├─ message [user #1] ├─ message [assistant #1] ├─ tool_result [id: 1] ├─ message [user #2] ├─ message [assistant #2] ├─ tool_result [id: 2] ├─ message [user #3] ├─ message [assistant #3] ├─ tool_result [id: 3] ├─ message [user #4] ├─ message [assistant #4] ├─ tool_result [id: 4] ├─ message [user #5] ├─ message [assistant #5] └─ leaf [current] │ ▼ Compact (threshold: 8,000 tokens) │ ├─► prepareCompaction() │ │ │ ├─► Estimate tokens: 10,000 │ ├─► Target: 6,000 (80% of 8,000) │ ├─► Find cut point: after message [assistant #3] │ ├─► Messages to summarize: [msg1, msg2, ..., msg3] │ └─► Retained tail: [msg4, msg5, leaf] │ ├─► LLM call to generate summary │ └─► Create compaction entry │ ├─► summary: "User asked X, assistant did Y, then Z..." ├─► firstKeptEntryId: msg4.id ├─► tokensBefore: 10,000 ├─► retainedTail: [msg4, msg5, leaf] └─► details: { readFiles: [...], modifiedFiles: [...] } │ ▼ Persisted Session Tree: │ ├─ root │ └─ message [user #1] │ └─ ... (original entries) │ └─ compaction [id: new] ← New entry │ ├─ summary: "User asked X..." │ ├─ firstKeptEntryId: msg4.id │ ├─ tokensBefore: 10000 │ ├─ retainedTail: [msg4, msg5, leaf] │ └─ details: {...} │ └─ msg4 [id: msg4] │ └─ message [assistant #4] │ └─ tool_result [id: 4] │ └─ message [user #5] │ └─ message [assistant #5] │ └─ leaf [id: leaf] │ └─ Context for LLM: └─ [compaction summary, retainedTail messages] ``` --- ## 7. State Mutation Flow ``` ┌─────────────────────────────────────────────────────────────────────────────────────────────┐ │ STATE MUTATION ON EVENTS │ └─────────────────────────────────────────────────────────────────────────────────────────────┘ Agent State: │ ├─ systemPrompt: string ├─ model: Model ├─ thinkingLevel: ThinkingLevel ├─ tools: AgentTool[] ├─ messages: AgentMessage[] ├─ isStreaming: boolean ├─ streamingMessage: AgentMessage? ← Partial assistant message ├─ pendingToolCalls: Set ← Currently executing └─ errorMessage: string? │ ▼ Events and State Changes: │ ├─ agent_start │ ├─ isStreaming = true │ ├─ streamingMessage = undefined │ └─ errorMessage = undefined │ ├─ message_start (user/assistant/toolResult) │ └─ No state change (just event emission) │ ├─ message_update (assistant only) │ └─ streamingMessage = updatedMessage │ ├─ message_end │ ├─ streamingMessage = undefined │ └─ messages.push(message) │ ├─ tool_execution_start │ └─ pendingToolCalls.add(toolCallId) │ ├─ tool_execution_end │ └─ pendingToolCalls.delete(toolCallId) │ ├─ turn_end │ └─ if (message.errorMessage) errorMessage = message.errorMessage │ └─ agent_end ├─ streamingMessage = undefined └─ (run finishes, state cleared on finishRun()) ``` --- ## 8. Queue Flow ``` ┌─────────────────────────────────────────────────────────────────────────────────────────────┐ │ QUEUE DRAINING FLOW │ └─────────────────────────────────────────────────────────────────────────────────────────────┘ Steering Queue (mode: "one-at-a-time"): │ ├─ Queue: [msg1, msg2, msg3] │ ├─ After turn ends: │ ├─► getSteeringMessages() │ ├─► mode = "one-at-a-time" │ ├─► Drain: [msg1] │ └─► Queue remaining: [msg2, msg3] │ ├─► Inject msg1 into context │ └─► Next LLM call includes: [...original, msg1] │ ▼ After next turn: │ ├─► getSteeringMessages() │ ├─► Drain: [msg2] │ └─► Queue remaining: [msg3] │ └─► Inject msg2 into context │ └─► ... and so on until queue empty Follow-up Queue (mode: "all"): │ ├─ Queue: [msg1, msg2, msg3] │ ├─ Agent would stop (no more tool calls) │ ├─► getFollowUpMessages() │ ├─► mode = "all" │ ├─► Drain: [msg1, msg2, msg3] │ └─► Queue remaining: [] │ ├─► Set as pendingMessages │ └─► Inner loop continues with: [...original, msg1, msg2, msg3] │ └─► All three messages injected together ``` --- ## 9. Event Sequence Examples ### Example 1: Simple Prompt ``` agent_start turn_start message_start (user: "Hello") message_end (user: "Hello") message_start (assistant: "") message_update (assistant: "H") message_update (assistant: "He") message_update (assistant: "Hel") message_update (assistant: "Hell") message_update (assistant: "Hello") message_end (assistant: "Hello") turn_end agent_end ``` ### Example 2: Tool Execution ``` agent_start turn_start message_start (user: "List files") message_end (user: "List files") message_start (assistant: "") message_update (assistant: "") message_update (assistant: "") message_update (assistant: "") message_update (assistant: "") message_update (assistant: "") message_end (assistant: "") tool_execution_start (bash: { command: "ls -la" }) tool_execution_update (bash: { progress: 0 }) tool_execution_update (bash: { progress: 50 }) tool_execution_update (bash: { progress: 100 }) tool_execution_end (bash: { exitCode: 0 }) message_start (toolResult: "drwxr-xr-x...") message_end (toolResult: "drwxr-xr-x...") turn_end agent_end ``` ### Example 3: Steering ``` agent_start turn_start message_start (user: "Build app") message_end (user: "Build app") message_start (assistant: "") message_update (assistant: "Building...") turn_end │ ├─ User types: "Wait, use React" │ ├─► steer("Wait, use React") │ └─► steeringQueue.push(msg) │ ├─► getSteeringMessages() │ └─► Drain and return [msg] │ ├─► Inject into context │ └─► Next turn with: [...original, "Wait, use React"] turn_start message_start (user: "Wait, use React") message_end (user: "Wait, use React") message_start (assistant: "") message_update (assistant: "Using React...") turn_end agent_end ``` --- ## Summary **Key Diagrams**: 1. System Architecture - Layers and components 2. Message Flow - Complete prompt flow 3. Hook System - Hook execution order 4. Tool Execution - Tool call lifecycle 5. Session Branching - Tree navigation 6. Context Compaction - History summarization 7. State Mutation - Event-driven state changes 8. Queue Flow - Steering and follow-up draining 9. Event Sequences - Real examples These diagrams show how data flows through the agent system from user input to LLM response to tool execution and back.