16 KiB
16 KiB
Pi Agent Architecture - Top-Down Overview
Executive Summary
The Pi Agent is a stateful, event-driven agent framework built in TypeScript. It provides:
- Core Agent - Low-level agent loop with message/tool streaming
- 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
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
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
interface AgentTool<TParameters extends TSchema, TDetails> {
name: string;
label: string;
description: string;
parameters: TSchema;
execute(
toolCallId: string,
params: Static<TParameters>,
signal?: AbortSignal,
onUpdate?: AgentToolUpdateCallback<TDetails>
): Promise<AgentToolResult<TDetails>>;
}
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)
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)
// 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
AgentEventstream - Hooks can be async and are awaited in order
- Abort signal propagated through all operations
2. State Isolation
AgentStateis read-only externallyAgentHarnesssnapshots 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
CustomAgentMessagesinterface - Custom hooks: Add handlers via
subscribe()/on() - Tool context: Pass
toolContextto harness constructor - Storage: Implement
SessionStorageinterface
Learning Path
- Start with types.ts - Understand
AgentMessage,AgentEvent,AgentTool - Read agent-loop.ts - See how messages flow through the loop
- Study agent.ts - See how Agent wraps the loop with state management
- Read agent-harness.ts - See how session management hooks into the loop
- Explore session/ files* - Understand persistence and branching
- 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 loopHOOK-SYSTEM.md- Complete hook documentationSESSION-ARCHITECTURE.md- Session persistence detailsTOOL-EXECUTION.md- Tool execution mechanics