# Pi Agent Architecture - Complete Summary

## Quick Reference for Julia Reimplementation

---

## 1. Core Architecture (Top-Down)

```
┌─────────────────────────────────────────────────────────────────────┐
│ APPLICATION LAYER                                                   │
│ • Agent (Low-level)                                                 │
│ • AgentHarness (High-level)                                         │
└─────────────────────────────────────────────────────────────────────┘
                             │
        ┌────────────────────┼────────────────────┐
        │                    │                    │
        ▼                    ▼                    ▼
┌───────────────┐  ┌──────────────────┐  ┌──────────────────┐
│ Agent Core    │  │ Session System   │  │ Tool Execution   │
│ • Async loop  │  │ • Tree storage   │  │ • Prepare        │
│ • Event       │  │ • Branching      │  │ • Execute        │
│ • Message     │  │ • Compaction     │  │ • Finalize       │
│ • Hooks       │  │ • Context        │  │ • Streaming      │
└───────────────┘  └──────────────────┘  └──────────────────┘
        │                    │                    │
        ▼                    ▼                    ▼
┌─────────────────────────────────────────────────────────────────────┐
│ LLM PROVIDER LAYER                                                  │
│ • StreamFn (streaming interface)                                    │
│ • Models (LLM catalog)                                              │
└─────────────────────────────────────────────────────────────────────┘
```

---

## 2. Key Components

### Agent Core

**Files**: `src/agent.ts`, `src/agent-loop.ts`, `src/types.ts`

**Responsibilities**:
- State management (messages, tools, isStreaming, pendingToolCalls)
- Event streaming (agent_start, turn_start, message_start, etc.)
- Queue management (steering, follow-up)
- Hook execution (beforeToolCall, afterToolCall, etc.)

**Key Types**:
```typescript
type AgentMessage = Message | CustomAgentMessages
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 }
  | { type: "message_end"; message: AgentMessage }
  | { type: "tool_execution_start"; ... }
  | { type: "tool_execution_end"; ... }

interface AgentContext {
  systemPrompt: string
  messages: AgentMessage[]
  tools?: AgentTool<any>[]
}
```

### AgentHarness

**Files**: `src/harness/agent-harness.ts`

**Responsibilities**:
- Session persistence (JSONL/Memory)
- Branching (create conversation paths)
- Compaction (summarize history)
- Tool context binding
- Hook system (before_agent_start, tool_call, tool_result, etc.)
- Queue management (steer, followUp, nextTurn)

**Key Types**:
```typescript
interface AgentHarnessEvent<TSkill, TPromptTemplate> =
  | { type: "agent_start" }  // From core
  | { type: "before_agent_start" }  // Harness-specific
  | { type: "tool_call"; ... }
  | { type: "tool_result"; ... }
  | { type: "session_before_compact"; ... }
  | { type: "session_before_tree"; ... }
  // ... more harness events

interface SessionContext {
  systemPrompt: string
  messages: AgentMessage[]
  thinkingLevel: ThinkingLevel
  model: { provider: string; modelId: string } | null
  activeToolNames: string[] | null
}
```

### Session System

**Files**: `src/harness/session/`

**Responsibilities**:
- Conversation persistence as tree
- Context building from tree
- Branching and forking
- Compaction
- Entry types (message, model_change, compaction, branch_summary, etc.)

**Key Types**:
```typescript
interface SessionTreeEntry {
  id: string
  parentId: string | null
  timestamp: string
  type: string  // "message", "compaction", "branch_summary", etc.
}

interface CompactionEntry extends SessionTreeEntry {
  type: "compaction"
  summary: string
  firstKeptEntryId?: string
  tokensBefore: number
  retainedTail?: AgentMessage[]
}
```

### Tool System

**Files**: `src/harness/tools/`

**Responsibilities**:
- Tool definition and execution
- Sequential vs parallel execution
- Streaming updates
- Error handling
- Before/after hooks

**Key Types**:
```typescript
interface AgentTool<TParameters, TDetails> extends Tool<TParameters> {
  label: string
  execute(
    toolCallId: string,
    params: Static<TParameters>,
    signal?: AbortSignal,
    onUpdate?: AgentToolUpdateCallback<TDetails>
  ): Promise<AgentToolResult<TDetails>>
}

interface AgentToolResult<T> {
  content: (TextContent | ImageContent)[]
  details: T
  usage?: Usage
  addedToolNames?: string[]
  terminate?: boolean
}
```

---

## 3. Message Flow

```
User Input
    │
    ├─► Agent.prompt("Hello")
    │   └─► normalizePromptInput() → AgentMessage[]
    │
    └─► runWithLifecycle()
         ├─► isStreaming = true
         └─► runAgentLoop()
              │
              ├─► agent_start
              ├─► turn_start
              ├─► message_start/end (user)
              │
              ├─► streamAssistantResponse()
              │    ├─► transformContext() [optional]
              │    ├─► convertToLlm()
              │    └─► streamFn() → LLM
              │
              ├─► executeToolCalls()
              │    ├─► prepareToolCall()
              │    │    ├─► Find tool
              │    │    ├─► Validate args
              │    │    └─► beforeToolCall() [hook]
              │    │
              │    ├─► executePreparedToolCall()
              │    │    └─► tool.execute() with onUpdate
              │    │
              │    └─► finalizeExecutedToolCall()
              │         └─► afterToolCall() [hook]
              │
              └─► turn_end
                   ├─► prepareNextTurn() [hook]
                   ├─► shouldStopAfterTurn() [hook]
                   ├─► Drain steering queue
                   └─► Drain follow-up queue
                   
              ┌─► Continue? → Repeat
              └─► Stop? → agent_end
```

---

## 4. Hook System

| Hook | Layer | When | Can Block? | Use Case |
|------|-------|------|------------|----------|
| `convertToLlm` | Agent | Before LLM | No | Filter messages |
| `transformContext` | Agent | Before LLM | Yes | Prune context |
| `beforeToolCall` | Agent | Before tool | Yes | Permission checks |
| `afterToolCall` | Agent | After tool | Yes | Override results |
| `shouldStopAfterTurn` | Agent | After turn | Yes | Request early stop |
| `prepareNextTurn` | Agent | Before next | Yes | Update config |
| `getSteeringMessages` | Agent | After turn | Yes | Interrupt agent |
| `getFollowUpMessages` | Agent | When stop | Yes | Queue messages |

**Harness Hooks**:
- `before_agent_start` - Modify system prompt
- `context` - Transform context
- `tool_call` - Log/before tool
- `tool_result` - Log/after tool
- `session_before_compact` - Customize compaction
- `session_before_tree` - Customize branching
- `before_provider_request` - Modify stream options
- `before_provider_payload` - Modify LLM payload

---

## 5. Session Tree

```
root (parentId: null)
  ├─ message [id: 1]  ← User prompt
  │   └─ message [id: 2]  ← Assistant
  │       └─ tool_result [id: 3]
  │           └─ message [id: 4]
  │               └─ compaction [id: 5]
  │                   ├─ summary: "..."
  │                   ├─ firstKeptEntryId: msg6.id
  │                   ├─ tokensBefore: 10000
  │                   ├─ retainedTail: [msg6, msg7]
  │                   └─ msg6 [id: 6]  ← Retained
  │                       └─ ... (rest of retained)
  │                           └─ leaf [id: 8]  ← Current head
  │
  └─ branch_summary [id: 9]  ← Branch point
      └─ message [id: 10]  ← New branch
          └─ leaf [id: 11]  ← New head
```

**Key Operations**:
- `getBranch()` → Get entries from leaf to root
- `buildContext()` → Project entries to messages
- `fork()` → Create branch at entry
- `compact()` → Summarize history

---

## 6. Tool Execution Flow

```
1. LLM sends tool call
   └─► AssistantMessage with toolCall block

2. prepareToolCall()
   ├─► Find tool
   ├─► prepareArguments() [optional]
   ├─► validateToolArguments()
   └─► beforeToolCall() [hook]
        ├─► block: true → Error
        └─► block: undefined → Continue

3. executePreparedToolCall()
   └─► tool.execute(toolCallId, params, signal, onUpdate)
        ├─► onUpdate(partialResult) [streaming]
        └─► Return: { content, details, ... }

4. finalizeExecutedToolCall()
   └─► afterToolCall() [hook]
        ├─► Override: content, details, isError, usage, terminate
        └─► Use executed result

5. Emit events
   ├─► tool_execution_start
   ├─► tool_execution_update [streaming]
   └─► tool_execution_end
   │
   └─► createToolResultMessage()
        └─► Emit: message_start/end (toolResult)
```

---

## 7. Data Types

### Messages

```typescript
interface Message {
  role: "user" | "assistant" | "toolResult"
  content: (TextContent | ImageContent)[]
  api?: string
  provider?: string
  model?: string
  usage?: Usage
  stopReason?: StopReason
  errorMessage?: string
  timestamp: number
}

interface TextContent { type: "text"; text: string }
interface ImageContent { type: "image"; mediaType: string; data: string }
```

### Events

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

### Tools

```typescript
interface AgentTool<TParameters, TDetails> extends Tool<TParameters> {
  label: string
  prepareArguments?: (args: unknown) => Static<TParameters>
  execute(
    toolCallId: string,
    params: Static<TParameters>,
    signal?: AbortSignal,
    onUpdate?: AgentToolUpdateCallback<TDetails>
  ): Promise<AgentToolResult<TDetails>>
}

interface AgentToolResult<T> {
  content: (TextContent | ImageContent)[]
  details: T
  usage?: Usage
  addedToolNames?: string[]
  terminate?: boolean
}
```

---

## 8. State Management

### Agent State

```typescript
interface AgentState {
  systemPrompt: string
  model: Model<any>
  thinkingLevel: ThinkingLevel
  tools: AgentTool<any>[]
  messages: AgentMessage[]
  isStreaming: boolean
  streamingMessage?: AgentMessage
  pendingToolCalls: Set<string>
  errorMessage?: string
}
```

### State Mutations

| Event | State Change |
|-------|-------------|
| `message_start` | `streamingMessage = message` |
| `message_update` | `streamingMessage = message` |
| `message_end` | `messages.push(message)`, `streamingMessage = undefined` |
| `tool_execution_start` | `pendingToolCalls.add(toolCallId)` |
| `tool_execution_end` | `pendingToolCalls.delete(toolCallId)` |
| `turn_end` | `errorMessage = message.errorMessage` (if error) |
| `agent_end` | `streamingMessage = undefined` |

---

## 9. Queue System

### Steering Queue

**Purpose**: Interrupt agent while working

**Mode**: `"all"` or `"one-at-a-time"`

**Flow**: After turn ends → Drain → Inject into context → Next LLM call

### Follow-up Queue

**Purpose**: Queue messages for after agent stops

**Mode**: `"all"` or `"one-at-a-time"`

**Flow**: When agent would stop → Drain → Set as pending → Continue loop

---

## 10. Entry Types

| Type | Purpose |
|------|---------|
| `message` | User/assistant/toolResult messages |
| `model_change` | Model switch (`setModel()`) |
| `thinking_level_change` | Reasoning level (`setThinkingLevel()`) |
| `active_tools_change` | Tools change (`setActiveTools()`) |
| `compaction` | History summary (`compact()`) |
| `branch_summary` | Branch point (branching) |
| `custom` | App data (not visible to model) |
| `custom_message` | Custom message |
| `label` | User-assigned label |
| `leaf` | Current session head |

---

## 11. Common Patterns

### Context Window Management

```typescript
transformContext: async (messages, signal) => {
  if (estimateTokens(messages) > MAX_TOKENS) {
    return pruneOldestMessages(messages, Math.floor(MAX_TOKENS * 0.3))
  }
  return messages
}
```

### Tool Permission Checks

```typescript
beforeToolCall: async ({ toolCall, args }, signal) => {
  if (toolCall.name === "bash" && signal?.aborted) {
    return { block: true, reason: "Operation aborted" }
  }
  if (toolCall.name === "bash" && !await canExecute(args)) {
    return { block: true, reason: "Permission denied" }
  }
  return undefined
}
```

### Streaming Updates

```typescript
execute: async (id, params, signal, onUpdate) => {
  for await (const item of longProcess()) {
    if (signal?.aborted) throw new Error("Aborted")
    onUpdate({ 
      content: [{ type: "text", text: `Progress: ${item}%` }],
      details: { progress: item }
    })
  }
  return finalResult
}
```

### Early Termination

```typescript
shouldStopAfterTurn: async ({ message, toolResults }) => {
  // Check if model indicates completion
  if (message.content.some(c => c.text?.includes("TASK_COMPLETE"))) {
    return true
  }
  // Stop if all tool calls set terminate
  return toolResults.every(r => r.terminate)
}
```

---

## 12. Learning Path

1. **Start with types** - Understand AgentMessage, AgentEvent, AgentTool
2. **Study agent-loop** - See how messages flow through the loop
3. **Read hooks** - Understand customization points
4. **Explore session** - See persistence and tree structure
5. **Study tools** - Understand tool execution
6. **Read harness** - See high-level API
7. **Design in Julia** - Implement step by step

---

## 13. Implementation Checklist

### Phase 1: Data Types (Julia)
- [ ] AgentMessage equivalent
- [ ] AgentEvent types
- [ ] AgentTool interface
- [ ] AgentContext

### Phase 2: Core Agent
- [ ] Agent class with state
- [ ] Event streaming
- [ ] Message queue (steering, follow-up)

### Phase 3: Agent Loop
- [ ] runAgentLoop()
- [ ] streamAssistantResponse()
- [ ] executeToolCalls()
- [ ] Tool preparation and execution
- [ ] Event emission

### Phase 4: Hooks
- [ ] Hook registration
- [ ] Hook execution
- [ ] Return value handling

### Phase 5: Session
- [ ] SessionTreeEntry types
- [ ] Tree structure
- [ ] Context building
- [ ] Persistence

### Phase 6: AgentHarness
- [ ] High-level API
- [ ] Queue management
- [ ] Branching
- [ ] Compaction

---

## 14. Quick Reference Cards

### Agent Core

| Function | Purpose |
|----------|---------|
| `runAgentLoop()` | Start new conversation |
| `runAgentLoopContinue()` | Continue existing |
| `streamAssistantResponse()` | Stream LLM |
| `executeToolCalls()` | Execute tools |
| `prepareToolCall()` | Prepare tool execution |
| `executePreparedToolCall()` | Execute tool |
| `finalizeExecutedToolCall()` | Finalize tool |

### AgentHarness

| Method | Purpose |
|--------|---------|
| `prompt()` | Run conversation |
| `skill()` | Execute skill |
| `promptFromTemplate()` | Run template |
| `steer()` | Interrupt agent |
| `followUp()` | Queue message |
| `nextTurn()` | Queue next turn |
| `compact()` | Compress context |
| `navigateTree()` | Branch conversation |
| `setModel()` | Change model |
| `setThinkingLevel()` | Change reasoning |
| `setTools()` | Set tools |
| `setActiveTools()` | Set active tools |

### Hooks

| Hook | Layer | Purpose |
|------|-------|---------|
| `convertToLlm` | Agent | Convert messages |
| `transformContext` | Agent | Manipulate context |
| `beforeToolCall` | Agent | Block tools |
| `afterToolCall` | Agent | Override results |
| `shouldStopAfterTurn` | Agent | Request stop |
| `prepareNextTurn` | Agent | Update config |
| `getSteeringMessages` | Agent | Interrupt |
| `getFollowUpMessages` | Agent | Queue messages |

### Session

| Method | Purpose |
|--------|---------|
| `buildContext()` | Get LLM context |
| `appendMessage()` | Add message |
| `fork()` | Create branch |
| `compact()` | Compress history |

---

**You now have a complete reference for reimplementing the Pi Agent in Julia!**

Start with the data types, implement the core loop, add hooks, then build up to the harness and session layers.