Files
pi_harness/packages/agent/learning/01-ARCHITECTURE-OVERVIEW.md
2026-07-29 10:59:18 +07:00

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:

  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

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 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