36 KiB
36 KiB
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<string> ← 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: "<tool_call name=bash>")
message_update (assistant: "<tool_call name=bash>")
message_update (assistant: "<tool_call name=bash>")
message_end (assistant: "<tool_call name=bash>")
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:
- System Architecture - Layers and components
- Message Flow - Complete prompt flow
- Hook System - Hook execution order
- Tool Execution - Tool call lifecycle
- Session Branching - Tree navigation
- Context Compaction - History summarization
- State Mutation - Event-driven state changes
- Queue Flow - Steering and follow-up draining
- 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.