17 KiB
Session Tree Implementation Plan
Reference: session-tree.md
Phase 1: SessionManager Core ✅
- Update entry types with
id,parentIdfields (using SessionEntryBase) - Add
versionfield toSessionHeader - Change
CompactionEntry.firstKeptEntryIndex→firstKeptEntryId - Add
BranchSummaryEntrytype - Add
CustomEntrytype for hooks - Add
byId: Map<string, SessionEntry>index - Add
leafId: stringtracking - Implement
getPath(fromId?)tree traversal - Implement
getTree()returningSessionTreeNode[] - Implement
getEntry(id)lookup - Implement
getLeafUuid()andgetLeafEntry()helpers - Update
_buildIndex()to populatebyIdmap - Rename
saveXXX()toappendXXX()(returns id, advances leaf) - Add
appendCustomEntry(customType, data)for hooks - Update
buildSessionContext()to usegetPath()traversal
Phase 2: Migration ✅
- Add
CURRENT_SESSION_VERSION = 2constant - Implement
migrateV1ToV2()with extensible migration chain - Update
setSessionFile()to detect version and migrate - Implement
_rewriteFile()for post-migration persistence - Handle
firstKeptEntryIndex→firstKeptEntryIdconversion in migration
Phase 3: Branching ✅
- Implement
branch(id)- switch leaf pointer - Implement
branchWithSummary(id, summary)- create summary entry - Implement
createBranchedSession(leafId)- extract path to new file - Update
AgentSession.branch()to use new API
Phase 4: Compaction Integration ✅
- Update
compaction.tsto work with IDs - Update
prepareCompaction()to returnfirstKeptEntryId - Update
compact()to returnCompactionResultwithfirstKeptEntryId - Update
AgentSessioncompaction methods - Add
firstKeptEntryIdtobefore_compacthook event
Phase 5: Testing ✅
migration.test.ts- v1 to v2 migration, idempotencybuild-context.test.ts- context building with tree structure, compaction, branchestree-traversal.test.ts- append operations, getPath, getTree, branchingfile-operations.test.ts- loadEntriesFromFile, findMostRecentSessionsave-entry.test.ts- custom entry integration- Update existing compaction tests for new types
Remaining Work
Compaction Refactor
- Use
CompactionResulttype for hook return value - Make
CompactionEntry<T>generic with optionaldetails?: Tfield for hook-specific data - Make
CompactionResult<T>generic to match - Update
SessionEventBaseto passsessionManagerandmodelRegistryinstead of derived fields - Update
before_compactevent:- Pass
preparation: CompactionPreparationinstead of individual fields - Pass
previousCompactions: CompactionEntry[](newest first) instead ofpreviousSummary?: string - Keep:
customInstructions,model,signal - Drop:
resolveApiKey(usemodelRegistry.getApiKey()),cutPoint,entries
- Pass
- Update hook example
custom-compaction.tsto use new API - Update
getSessionFile()to returnstring | undefinedfor in-memory sessions - Update
before_switchto havetargetSessionFile,switchto havepreviousSessionFile
Reference: #314 - Structured compaction with anchored iterative summarization needs details field to store ArtifactIndex and version markers.
Branch Summary Design ✅
Current type:
export interface BranchSummaryEntry extends SessionEntryBase {
type: "branch_summary";
summary: string;
fromId: string; // References the abandoned leaf
fromHook?: boolean; // Whether summary was generated by a hook
details?: unknown; // File tracking: { readFiles, modifiedFiles }
}
fromIdfield references the abandoned leaffromHookfield distinguishes pi-generated vs hook-generated summariesdetailsfield for file tracking- Branch summarizer implemented with structured output format
- Uses serialization approach (same as compaction) to prevent model confusion
- Tests for
branchWithSummary()flow
Entry Labels ✅
- Add
LabelEntrytype withtargetIdandlabelfields - Add
labelsById: Map<string, string>private field - Build labels map in
_buildIndex()via linear scan - Add
getLabel(id)method - Add
appendLabelChange(targetId, label)method (undefined clears) - Update
createBranchedSession()to filter out LabelEntry and recreate from resolved map buildSessionContext()already ignores LabelEntry (only handles message types)- Add
label?: stringtoSessionTreeNode, populated bygetTree() - Display labels in UI (tree-selector shows labels)
/labelcommand (implemented in tree-selector)
CustomMessageEntry
Hook-injected messages that participate in LLM context. Unlike CustomEntry<T> (for hook state only), these are sent to the model.
export interface CustomMessageEntry<T = unknown> extends SessionEntryBase {
type: "custom_message";
customType: string; // Hook identifier
content: string | (TextContent | ImageContent)[]; // Message content (same as UserMessage)
details?: T; // Hook-specific data for state reconstruction on reload
display: boolean; // Whether to display in TUI
}
Behavior:
- Type definition matching plan
appendCustomMessageEntry(customType, content, display, details?)in SessionManagerbuildSessionContext()includes custom_message entries as user messages- Exported from main index
- TUI rendering:
display: false- hidden entirelydisplay: true- rendered with purple styling (customMessageBg, customMessageText, customMessageLabel theme colors)registerCustomMessageRenderer(customType, renderer)in HookAPI for custom renderers- Renderer returns inner Component, TUI wraps in styled Box
Hook API Changes ✅
Renamed:
renderCustomMessage()→registerCustomMessageRenderer()
New: sendMessage() ✅
Replaces send(). Always creates CustomMessageEntry, never user messages.
type HookMessage<T = unknown> = Pick<CustomMessageEntry<T>, 'customType' | 'content' | 'display' | 'details'>;
sendMessage(message: HookMessage, triggerTurn?: boolean): void;
Implementation:
- Uses agent's queue mechanism with
_hookDatamarker on AppMessage message_endhandler routes based on marker presenceAgentSession.sendHookMessage()handles three cases:- Streaming: queues via
agent.queueMessage(), loop processes and emitsmessage_end - Not streaming + triggerTurn: direct append +
agent.continue() - Not streaming + no trigger: direct append only
- Streaming: queues via
- TUI updates via event (streaming) or explicit rebuild (non-streaming)
New: appendEntry() ✅
For hook state persistence (NOT in LLM context):
appendEntry(customType: string, data?: unknown): void;
Calls sessionManager.appendCustomEntry() directly.
New: registerCommand() (types ✅, wiring TODO)
// HookAPI (the `pi` object) - utilities available to all hooks:
interface HookAPI {
sendMessage(message: HookMessage, triggerTurn?: boolean): void;
appendEntry(customType: string, data?: unknown): void;
registerCommand(name: string, options: RegisteredCommand): void;
registerCustomMessageRenderer(customType: string, renderer: CustomMessageRenderer): void;
exec(command: string, args: string[], options?: ExecOptions): Promise<ExecResult>;
}
// HookEventContext - passed to event handlers, has stable context:
interface HookEventContext {
ui: HookUIContext;
hasUI: boolean;
cwd: string;
sessionManager: SessionManager;
modelRegistry: ModelRegistry;
}
// Note: exec moved to HookAPI, sessionManager/modelRegistry moved from SessionEventBase
// HookCommandContext - passed to command handlers:
interface HookCommandContext {
args: string; // Everything after /commandname
ui: HookUIContext;
hasUI: boolean;
cwd: string;
sessionManager: SessionManager;
modelRegistry: ModelRegistry;
}
// Note: exec and sendMessage accessed via `pi` closure
registerCommand(name: string, options: {
description?: string;
handler: (ctx: HookCommandContext) => Promise<void>;
}): void;
Handler return:
void- command completed (usesendMessage()withtriggerTurn: trueto prompt LLM)
Wiring (all in AgentSession.prompt()):
- Add hook commands to autocomplete in interactive-mode
_tryExecuteHookCommand()in AgentSession handles command execution- Build HookCommandContext with ui (from hookRunner), exec, sessionManager, etc.
- If handler returns string, use as prompt text
- If handler returns undefined, return early (no LLM call)
- Works for all modes (interactive, RPC, print) via shared AgentSession
New: ui.custom() ✅
For arbitrary hook UI with keyboard focus:
interface HookUIContext {
// ... existing: select, confirm, input, notify
/** Show custom component with keyboard focus. Call done() when finished. */
custom(component: Component, done: () => void): void;
}
See also: CustomEntry<T> for storing hook state that does NOT participate in context.
New: context event ✅
Fires before messages are sent to the LLM, allowing hooks to modify context non-destructively.
interface ContextEvent {
type: "context";
/** Messages that will be sent to the LLM */
messages: Message[];
}
interface ContextEventResult {
/** Modified messages to send instead */
messages?: Message[];
}
// In HookAPI:
on(event: "context", handler: HookHandler<ContextEvent, ContextEventResult | void>): void;
Example use case: Dynamic Context Pruning (discussion #330)
Non-destructive pruning of tool results to reduce context size:
export default function(pi: HookAPI) {
// Register /prune command
pi.registerCommand("prune", {
description: "Mark tool results for pruning",
handler: async (ctx) => {
// Show UI to select which tool results to prune
// Append custom entry recording pruning decisions:
// { toolResultId, strategy: "summary" | "truncate" | "remove" }
pi.appendEntry("tool-result-pruning", { ... });
}
});
// Intercept context before LLM call
pi.on("context", async (event, ctx) => {
// Find all pruning entries in session
const entries = ctx.sessionManager.getEntries();
const pruningRules = entries
.filter(e => e.type === "custom" && e.customType === "tool-result-pruning")
.map(e => e.data);
// Apply pruning rules to messages
const prunedMessages = applyPruning(event.messages, pruningRules);
return { messages: prunedMessages };
});
}
Benefits:
- Original tool results stay intact in session
- Pruning is stored as custom entries, survives session reload
- Works with branching (pruning entries are part of the tree)
- Trade-off: cache busting on first submission after pruning
Investigate: context event vs before_agent_start ✅
References:
Current context event:
- Fires before each LLM call within the agent loop
- Receives
AgentMessage[](deep copy, safe to modify) - Returns
Message[](inconsistent with input type) - Modifications are transient (not persisted to session)
- No TUI visibility of what was changed
- Use case: non-destructive pruning, dynamic context manipulation
Type inconsistency: Event receives AgentMessage[] but result returns Message[]:
interface ContextEvent {
messages: AgentMessage[]; // Input
}
interface ContextEventResult {
messages?: Message[]; // Output - different type!
}
Questions:
- Should input/output both be
Message[](LLM format)? - Or both be
AgentMessage[]with conversion happening after? - Where does
AgentMessage[]→Message[]conversion currently happen?
Proposed before_agent_start event:
- Fires once when user submits a prompt, before
agent_start - Allows hooks to inject additional content that gets persisted to session
- Injected content is visible in TUI (observability)
- Does not bust prompt cache (appended after user message, not modifying system prompt)
Key difference:
| Aspect | context |
before_agent_start |
|---|---|---|
| When | Before each LLM call | Once per user prompt |
| Persisted | No | Yes (as SystemMessage) |
| TUI visible | No | Yes (collapsible) |
| Cache impact | Can bust cache | Append-only, cache-safe |
| Use case | Transient manipulation | Persistent context injection |
Implementation (completed):
- Reuses
HookMessagetype (no new message type needed) - Handler returns
{ message: Pick<HookMessage, "customType" | "content" | "display" | "details"> } - Message is appended to agent state AND persisted to session before
agent.prompt()is called - Renders using existing
HookMessageComponent(or custom renderer if registered) - How does it interact with compaction? (treated like user messages?)
- Can hook return multiple messages or just one?
Implementation sketch:
interface BeforeAgentStartEvent {
type: "before_agent_start";
userMessage: UserMessage; // The prompt user just submitted
}
interface BeforeAgentStartResult {
/** Additional context to inject (persisted as SystemMessage) */
inject?: {
label: string; // Shown in collapsed TUI state
content: string | (TextContent | ImageContent)[];
};
}
HTML Export
- Add collapsible sidebar showing full tree structure
- Allow selecting any node in tree to view that path
- Add "reset to session leaf" button
- Render full path (no compaction resolution needed)
- Responsive: collapse sidebar on mobile
UI Commands ✅
/branch- Creates new session file from current path (usescreateBranchedSession())/tree- In-session tree navigation via tree-selector component- Shows full tree structure with labels
- Navigate between branches (moves leaf pointer)
- Shows current position
- Generates branch summaries when switching branches
Tree Selector Improvements ✅
- Active line highlight using
selectedBgtheme color - Filter modes via
^O(forward) /Shift+^O(backward):default: hides label/custom entriesno-tools: default minus tool resultsuser-only: just user messageslabeled-only: just labeled entriesall: everything
Documentation
Review and update all docs:
docs/hooks.md- Major update for hook API:pi.send()→pi.sendMessage()with new signature- New
pi.appendEntry()for state persistence - New
pi.registerCommand()for custom slash commands - New
pi.registerCustomMessageRenderer()for custom TUI rendering HookCommandContextinterface and handler patternsHookMessage<T>type- Updated event signatures (
SessionEventBase,before_compact, etc.)
docs/hooks-v2.md- Review/merge or remove if obsoletedocs/sdk.md- Update for:HookMessageandisHookMessage()Agent.prompt(AppMessage)overload- Session v2 tree structure
- SessionManager API changes
docs/session.md- Update for v2 tree structure, new entry typesdocs/custom-tools.md- Check if hook changes affect custom toolsdocs/rpc.md- Check if hook commands work in RPC modedocs/skills.md- Review for any hook-related updatesdocs/extension-loading.md- Reviewdocs/theme.md- Added selectedBg, customMessageBg/Text/Label color tokens (50 total)README.md- Update hook examples if any
Examples
Review and update examples:
examples/hooks/- Update existing, add new examples:- Review
custom-compaction.tsfor new API - Add
registerCommand()example - Add
sendMessage()example - Add
registerCustomMessageRenderer()example
- Review
examples/sdk/- Update for new session/hook APIsexamples/custom-tools/- Review for compatibility
Before Release
- Run full automated test suite:
npm test - Manual testing of tree navigation and branch summarization
- Verify compaction with file tracking works correctly
Notes
- All append methods return the new entry's ID
- Migration rewrites file on first load if version < CURRENT_VERSION
- Existing sessions become linear chains after migration (parentId = previous entry)
- Tree features available immediately after migration
- SessionHeader does NOT have id/parentId (it's metadata, not part of tree)
- Session is append-only: entries cannot be modified or deleted, only branching changes the leaf pointer