32 KiB
32 KiB
Memory & Context Management in Pi Agent
1. Architecture Overview
The agent manages memory at two layers:
┌─────────────────────────────────────────────────────────────────┐
│ Agent Harness │
│ ┌──────────────┐ ┌──────────────┐ ┌─────────────────────┐ │
│ │ Session │ │ Compaction │ │ Branch Navigation │ │
│ │ (JSONL tree)│ │ (summarize) │ │ (reset + summarize)│ │
│ └──────┬───────┘ └──────────────┘ └─────────────────────┘ │
└─────────┼───────────────────────────────────────────────────────┘
│ builds context
┌─────────▼───────────────────────────────────────────────────────┐
│ Agent Class │
│ ┌───────────────────────────────────────────────────────────┐ │
│ │ _state.messages: AgentMessage[] (linear transcript) │ │
│ │ _state.tools, systemPrompt, model │ │
│ └───────────────────────────────────────────────────────────┘ │
│ subscribe() → events → UI updates │
└─────────────────────────────────────────────────────────────────┘
Key separation:
- In-memory (
Agent): linear transcript for the current run. Cleared onreset(). - On-disk (
Session): persistent tree of entries in JSONL files. Survives restarts. - Compaction: replaces old on-disk history with an LLM-generated summary, controlling context window usage.
2. Data Flow: From Prompt to LLM Call
agent.prompt("Read README.md")
│
▼
┌──────────────────────────┐
│ normalizePromptInput() │ → { role: "user", content: "..." }
└──────────┬───────────────┘
│
▼
┌──────────────────────────┐
│ runWithLifecycle() │ → sets isStreaming=true
│ runAgentLoop() │
└──────────┬───────────────┘
│
▼
┌─────────────────────────────────────────────────────────────┐
│ runLoop() — the main while(true) loop │
│ │
│ Inner loop: │
│ 1. Inject steering/follow-up messages │
│ 2. streamAssistantResponse() │
│ ┌─────────────────────────────────────────────┐ │
│ │ transformContext() │ │
│ │ AgentMessage[] → AgentMessage[] │ │
│ │ (prune, inject external context) │ │
│ └──────────────┬──────────────────────────────┘ │
│ ▼ │
│ ┌─────────────────────────────────────────────┐ │
│ │ convertToLlm() │ │
│ │ AgentMessage[] → Message[] │ │
│ │ (filter to user/assistant/toolResult only) │ │
│ └──────────────┬──────────────────────────────┘ │
│ ▼ │
│ ┌─────────────────────────────────────────────┐ │
│ │ streamFunction() → LLM provider │ │
│ │ { systemPrompt, messages, tools } │ │
│ └──────────────┬──────────────────────────────┘ │
│ ▼ │
│ Stream events: start → delta* → done │
│ ▼ │
│ Return AssistantMessage │
│ 3. Extract toolCall blocks │
│ 4. If toolCalls: executeToolCalls() │
│ → create toolResult messages │
│ → append to context │
│ 5. Check shouldStopAfterTurn / prepareNextTurn │
│ 6. Check steering/follow-up queues │
│ → Loop if more tool calls or queued messages │
│ │
│ Outer loop: │
│ → Check follow-up queue for messages after agent would │
│ stop │
└─────────────────────────────────────────────────────────────┘
3. Context Building (On-Disk → In-Memory)
The AgentHarness bridges on-disk session data to the in-memory agent loop.
session.buildContext()
│
▼
┌─────────────────────────────────────────────────────────────┐
│ getBranch() — walk from leaf → root via parentId │
│ │
│ Tree structure: │
│ │
│ ┌──────┐ ┌──────┐ ┌──────────┐ ┌──────────┐ │
│ │msg 1 │───▶│msg 2 │───▶│ msg 3 │───▶│ msg 4 │ │
│ │user │ │assist│ │ toolCall │ │ user │ │
│ └──────┘ └──────┘ └──────────┘ └──────────┘ │
│ │
│ Path to root: [msg1, msg2, msg3, msg4] │
└──────────┬────────────────────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────┐
│ defaultContextEntryTransform() — THE KEY STEP │
│ │
│ Finds latest "compaction" entry in path: │
│ │
│ ┌──────┐ ┌──────────┐ ┌──────┐ ┌──────┐ │
│ │msg 1 │ │compaction│ │msg 3 │ │msg 4 │ │
│ │user │ │summary X │ │msg 2 │ │assist│ │
│ └──────┘ └──────────┘ └──────┘ └──────┘ │
│ │ │ │ │ │
│ ├──────────────┤ │ │ │
│ │ SKIPPED │ │ │ │
│ │ (summarized)│ │ │ │
│ └──────────────┼───────────┘ │ │
│ ▼ ▼ │
│ Include compaction entry Include entries after │
│ + firstKeptEntryId the compaction point │
│ │
│ Result: [compaction, msg3, msg4] │
│ → compaction entry becomes a "compactionSummary" message │
└──────────┬────────────────────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────┐
│ sessionEntryToContextMessages() │
│ │
│ For each entry: │
│ message → [message] │
│ compaction → [compactionSummary, ...retainedTail] │
│ branch_summary → [branchSummaryMessage] │
│ custom_message → [customMessage] │
│ other → [] (omitted from LLM context) │
│ │
│ Flat result: [compactionSummary, msg3, msg4] │
└──────────┬────────────────────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────┐
│ deriveSessionContextState() │
│ │
│ Extracts from entries: │
│ thinkingLevel ← latest thinking_level_change or assistant │
│ model ← latest model_change or assistant │
│ activeTools ← latest active_tools_change │
└─────────────────────────────────────────────────────────────┘
4. Token Estimation
Before compaction can decide whether to trigger, it needs to know how many tokens the context uses.
estimateContextTokens(messages)
│
├── Has provider-reported usage on last assistant msg?
│ ├── YES → use actual usage + estimate tail
│ │ (accurate — avoids compounding error)
│ │
│ └── NO → estimate all messages from scratch
│
▼
estimateTokens(message) [character heuristic: chars / 4]
│
├── role === "user"
│ content.length / 4
│ (images ≈ 4800 chars each)
│
├── role === "assistant"
│ sum of all content blocks:
│ text blocks → text.length
│ thinking blocks → thinking.length
│ toolCall blocks → name.length + JSON.stringify(args).length
│
├── role === "toolResult" / "custom"
│ content.length / 4
│
├── role === "bashExecution"
│ (command.length + output.length) / 4
│
└── role === "compactionSummary" / "branchSummary"
summary.length / 4
Why chars / 4? Rough heuristic: ~4 ASCII characters ≈ 1 token. Conservative estimate to avoid under-counting.
5. Compaction — The Core Memory Management
5.1 Trigger Condition
shouldCompact(contextTokens, contextWindow, settings)
→ contextTokens > contextWindow - reserveTokens
Defaults:
reserveTokens: 16384 (~16K tokens for summary prompt + output)
keepRecentTokens: 20000 (~20K tokens of recent history to keep)
Example (Claude with 200K context window):
Triggers when: contextTokens > 200000 - 16384 = 183616
5.2 Finding the Cut Point
findCutPoint(entries, startIndex, endIndex, keepRecentTokens)
│
│ Walk BACKWARD from endIndex
│
├── Accumulate estimated tokens per message
├── Stop when accumulated ≥ keepRecentTokens
├── Snap to nearest valid cut point
│
▼
┌─────────────────────────────────────────────────────────────┐
│ Valid cut points (safe to split): │
│ - user message │
│ - assistant message │
│ - custom message │
│ - branch_summary │
│ │
│ NOT valid (tool results stay with their call): │
│ - toolResult message (skipped) │
│ │
│ Example: │
│ │
│ ┌────┐ ┌────┐ ┌────┐ ┌────┐ ┌────┐ ┌────┐ ┌────┐ │
│ │ u1 │ │ a1 │ │ tr1│ │ u2 │ │ a2 │ │ tr2│ │ u3 │ │
│ └────┘ └────┘ └────┘ └────┘ └────┘ └────┘ └────┘ │
│ ▲ ▲ ▲ │
│ │ │ │ │
│ └── kept └── cut └── discarded │
│ (~20K tokens) point history │
└─────────────────────────────────────────────────────────────┘
5.3 The Compaction Process
prepareCompaction(branchEntries, settings)
│
├── Find previous compaction (if any) → previousSummary
├── Estimate tokens of current context
├── findCutPoint() → firstKeptEntryId
│
├── Split into 3 groups:
│ │
│ ├── messagesToSummarize: entries BEFORE cut point
│ │ (these become the summary)
│ │
│ ├── retainedTail: entries AFTER cut point
│ │ (these stay verbatim)
│ │
│ └── turnPrefixMessages: if cut splits a turn
│ (the beginning of an interrupted turn)
│
└── Extract file operations from messagesToSummarize:
→ readFiles, modifiedFiles
compact(preparation, model, models)
│
├── If isSplitTurn:
│ ├── generateSummary(messagesToSummarize) → history summary
│ ├── generateTurnPrefixSummary(turnPrefixMessages) → turn context
│ └── Combine: history + "---" + turn prefix
│
├── Else (normal):
│ ├── Has previousSummary?
│ │ ├── YES → UPDATE_SUMMARIZATION_PROMPT (iterative)
│ │ └── NO → SUMMARIZATION_PROMPT (fresh)
│ └── Call LLM with conversation text + prompt
│
├── Append file operations:
│ "Files read: [...]\nFiles modified: [...]"
│
└── Return:
{
summary: "## Goal...\n## Progress...\n...",
firstKeptEntryId: "entry-uuid",
tokensBefore: 185000,
retainedTail: [msg3, msg4, ...],
details: { readFiles: [...], modifiedFiles: [...] }
}
5.4 Summary Format
The LLM generates a structured summary:
## Goal
- [What is the user trying to accomplish?]
## Constraints & Preferences
- [Any constraints, preferences, or requirements]
## Progress
### Done
- [x] [Completed tasks]
### In Progress
- [ ] [Current work]
### Blocked
- [Issues preventing progress]
## Key Decisions
- **[Decision]**: [Brief rationale]
## Next Steps
1. [Ordered list of what should happen next]
## Critical Context
- [Any data, examples, or references needed to continue]
Files read: [src/index.ts, package.json]
Files modified: [src/index.ts]
5.5 Iterative Compaction
Successive compact calls update the existing summary rather than replacing it:
Compaction 1 (at ~185K tokens):
Summary: "## Goal: Build a login page..."
firstKeptEntryId: "entry-003"
Compaction 2 (at ~185K tokens again):
previousSummary: "## Goal: Build a login page..."
→ UPDATE_SUMMARIZATION_PROMPT
→ PRESERVES existing information
→ ADDS new progress (move "In Progress" → "Done")
→ NEW summary: "## Goal: Build a login page... Add OAuth..."
firstKeptEntryId: "entry-003" (same boundary)
6. Session Storage — JSONL Format
Session file: .pi/sessions/--home-user--/2024-01-15T10-30-00_abc123.jsonl
Line 1 (header):
{"type":"session","version":3,"id":"abc123","timestamp":"2024-01-15T10:30:00.000Z",
"cwd":"/home/user/project","parentSession":"...","metadata":{}}
Line 2+ (entries, one per line):
{"type":"message","id":"e001","parentId":null,"timestamp":"...","message":{...}}
{"type":"message","id":"e002","parentId":"e001","timestamp":"...","message":{...}}
{"type":"compaction","id":"e003","parentId":"e002","timestamp":"...",
"summary":"## Goal: ...\n...","firstKeptEntryId":"e001",
"tokensBefore":185000}
{"type":"leaf","id":"e004","parentId":"e003","timestamp":"...",
"targetId":"e002"}
Entry Types
| Type | LLM Context? | Purpose |
|---|---|---|
message |
Yes | User, assistant, toolResult |
compaction |
Yes (as summary message) | Replaces compacted history |
branch_summary |
Yes (as summary message) | Summary of diverged branch |
leaf |
No | Points to current tree leaf |
thinking_level_change |
No | Tracking thinking level changes |
model_change |
No | Tracking model changes |
active_tools_change |
No | Tracking tool enable/disable |
custom |
No (unless projector configured) | App-defined data |
custom_message |
Yes | App-defined messages |
label |
No | Human-readable labels |
session_info |
No | Session name history |
7. Session Tree (Branching)
Sessions form a tree, not a linear log. This lets users "go back" and try a different approach.
Session tree:
┌───[e01]───┐
│ user: "a" │
└─────┬──────┘
▼
┌──────────┐
│ assist 1 │
└─────┬────┘
▼
┌──────────┐
│ toolCall │
└─────┬────┘
▼
┌──────────┐
│ toolRes 1│
└─────┬────┘
▼
┌──────────┐
│ user: "b"│ ← user goes back here
└─────┬────┘
│
┌─────┴─────┐
│ │
┌──────────┐ ┌──────────┐
│ user: "c" │ │ user: "d" │ ← branch point
└────┬─────┘ └────┬─────┘
│ │
┌────▼─────┐ ┌────▼─────┐
│ assist 2 │ │ assist 3 │ ← current leaf (d)
└──────────┘ └──────────┘
When user navigates to "user: b":
- Leaf moves from "d" back to "b"
- Branch summary generated for diverged work ("c" → "assist 2")
- New work branches from "b":
┌─────┐
│ user: "e" │ ← new branch
└─────┬─────┘
▼
┌──────────┐
│ assist 4 │
└──────────┘
Context sent to LLM:
[compaction summary, user:b, user:e, assist:4]
→ The old "c"/"assist 2" branch is replaced by its summary
8. Pending Writes — Batching Session Persistence
To avoid writing every message individually during a run:
Agent loop events
│
▼
handleAgentEvent(event)
│
├── message_end → pendingSessionWrites.push({ type: "message", message })
├── turn_end → flushPendingSessionWrites() (save_point)
├── agent_end → flushPendingSessionWrites()
│
▼
flushPendingSessionWrites()
│
├── Iterate pendingSessionWrites[]
│ ├── message → session.appendMessage(msg)
│ ├── model_change → session.appendModelChange(provider, id)
│ ├── thinking_level_change → session.appendThinkingLevelChange(level)
│ ├── active_tools_change → session.appendActiveToolsChange(names)
│ ├── custom → session.appendCustomEntry(type, data)
│ ├── custom_message → session.appendCustomMessageEntry(...)
│ ├── label → session.appendLabel(targetId, label)
│ ├── session_info → session.appendSessionName(name)
│ └── leaf → session.getStorage().setLeafId(targetId)
│
└── Shift all writes → empty pending list
During the run, messages are accumulated in pendingSessionWrites and only flushed to disk at save_point (end of each turn) or agent_end.
9. Hooks & Extensibility — Context Control Points
The harness exposes hooks at every memory management boundary:
┌────────────────────────────────────────────────────────────────┐
│ Hooks │
├────────────────────────┬─────────────────────────────────────┤
│ Hook │ When │
├────────────────────────┼─────────────────────────────────────┤
│ before_agent_start │ Before each prompt, can add │
│ │ messages or modify system prompt │
├────────────────────────┼─────────────────────────────────────┤
│ context │ Before each LLM call, can prune/ │
│ │ modify AgentMessage[] │
├────────────────────────┼─────────────────────────────────────┤
│ before_provider_request│ Before each provider API call, can │
│ │ modify headers, retries, timeout │
├────────────────────────┼─────────────────────────────────────┤
│ before_provider_payload│ Before sending payload to provider,│
│ │ can modify the request body │
├────────────────────────┼─────────────────────────────────────┤
│ after_provider_response│ After receiving response, for │
│ │ logging/metrics │
├────────────────────────┼─────────────────────────────────────┤
│ tool_call │ Before tool execution, can │
│ │ return { block: true } │
├────────────────────────┼─────────────────────────────────────┤
│ tool_result │ After tool execution, can override │
│ │ content, details, isError, terminate│
├────────────────────────┼─────────────────────────────────────┤
│ session_before_compact│ Before compaction, can cancel or │
│ │ provide custom compaction result │
├────────────────────────┼─────────────────────────────────────┤
│ session_before_tree │ Before branch navigation, can │
│ │ cancel or provide custom summary │
├────────────────────────┼─────────────────────────────────────┤
│ prepareNextTurn │ Between turns, can replace context, │
│ │ model, or thinkingLevel │
├────────────────────────┼─────────────────────────────────────┤
│ shouldStopAfterTurn │ After a turn, if true the loop │
│ │ exits (agent_end, no more LLM calls)│
└────────────────────────┴─────────────────────────────────────┘
10. Complete Lifecycle: Long Session
Session starts empty
│
▼
Turn 1: "Create a React component"
Context: [compaction summary (empty)]
Messages exchanged: ~2K tokens
└─ Session: [user1, assist1, toolCall, toolRes1, assist2]
│
▼
Turn 2-10: Iterative development
Context: growing with each turn
Total context: ~50K tokens
└─ Session: [user1..assist2, user2..assist20]
│
▼
Turn 15: Context approaching limit (~170K tokens)
shouldCompact() → true
└─ Compaction 1:
- Summarizes turns 1-12
- Keeps turns 13-15 verbatim
- Summary: "## Goal: React component ## Progress: built X, Y"
│
▼
Turn 20: Context ~180K tokens
shouldCompact() → true
└─ Compaction 2 (iterative update):
- Updates existing summary with new progress
- "## Goal: React component ## Done: built X,Y ## New: added auth"
│
▼
Turn 25: Context ~186K tokens → triggers compaction
└─ Compaction 3:
- Summary now covers ~22 turns of history
- Retained tail: last 20K tokens (~5 turns)
- Context window freed: ~186K → ~25K tokens
│
▼
User navigates to Turn 8:
- Branch summary generated for Turns 9-25
- Leaf moves back to Turn 8
- Context: [compaction, branch_summary, turns 1-8]
│
▼
User continues from Turn 8:
- New branch grows from Turn 8
- Old branch (9-25) replaced by branch_summary
│
▼
Session ends, JSONL file persists on disk
Next session: loads from JSONL, rebuilds context
11. Token Budget Summary
Example: Claude Sonnet (200K context window)
┌─────────────────────────────────────────────────────────────┐
│ Context Window: 200,000 tokens │
├─────────────────────────────────────────────────────────────┤
│ Reserved for summary: 16,384 tokens │
├─────────────────────────────────────────────────────────────┤
│ Keep recent: 20,000 tokens │
├─────────────────────────────────────────────────────────────┤
│ Max context before compaction: 183,616 tokens │
│ (= 200000 - 16384) │
├─────────────────────────────────────────────────────────────┤
│ After compaction: ~25,000 tokens │
│ (20,000 tail + ~5,000 summary) │
│ → ~158,616 tokens freed │
└─────────────────────────────────────────────────────────────┘
12. Key Files Reference
| File | Responsibility |
|---|---|
agent-loop.ts |
Core loop, tool execution, streaming |
agent.ts |
Stateful Agent class, event system |
agent-harness.ts |
High-level harness, hooks, session management |
compaction/compaction.ts |
Token estimation, cut point, LLM summarization |
compaction/branch-summarization.ts |
Branch divergence summarization |
session/session.ts |
Session tree, entry appending, context building |
session/jsonl-storage.ts |
JSONL file read/write |
session/jsonl-repo.ts |
Session repo: create/open/list/delete/fork |
types.ts |
All type definitions |
messages.ts |
convertToLlm(), custom message helpers |
system-prompt.ts |
System prompt building |
skills.ts |
Skill management |