co-mono/packages/coding-agent/docs/session-tree-plan.md

Session Tree Implementation Plan

Reference: session-tree.md

Phase 1: SessionManager Core ✅

• Update entry types with id, parentId fields (using SessionEntryBase)
• Add version field to SessionHeader
• Change CompactionEntry.firstKeptEntryIndex → firstKeptEntryId
• Add BranchSummaryEntry type
• Add CustomEntry type for hooks
• Add byId: Map<string, SessionEntry> index
• Add leafId: string tracking
• Implement getPath(fromId?) tree traversal
• Implement getTree() returning SessionTreeNode[]
• Implement getEntry(id) lookup
• Implement getLeafUuid() and getLeafEntry() helpers
• Update _buildIndex() to populate byId map
• Rename saveXXX() to appendXXX() (returns id, advances leaf)
• Add appendCustomEntry(customType, data) for hooks
• Update buildSessionContext() to use getPath() traversal

Phase 2: Migration ✅

• Add CURRENT_SESSION_VERSION = 2 constant
• Implement migrateV1ToV2() with extensible migration chain
• Update setSessionFile() to detect version and migrate
• Implement _rewriteFile() for post-migration persistence
• Handle firstKeptEntryIndex → firstKeptEntryId conversion 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.ts to work with IDs
• Update prepareCompaction() to return firstKeptEntryId
• Update compact() to return CompactionResult with firstKeptEntryId
• Update AgentSession compaction methods
• Add firstKeptEntryId to before_compact hook event

Phase 5: Testing ✅

• migration.test.ts - v1 to v2 migration, idempotency
• build-context.test.ts - context building with tree structure, compaction, branches
• tree-traversal.test.ts - append operations, getPath, getTree, branching
• file-operations.test.ts - loadEntriesFromFile, findMostRecentSession
• save-entry.test.ts - custom entry integration
• Update existing compaction tests for new types

---

Remaining Work

Compaction Refactor

• Use CompactionResult type for hook return value
• Make CompactionEntry<T> generic with optional details?: T field for hook-specific data
• Make CompactionResult<T> generic to match
• Update SessionEventBase to pass sessionManager and modelRegistry instead of derived fields
• Update before_compact event:
  • Pass preparation: CompactionPreparation instead of individual fields
  • Pass previousCompactions: CompactionEntry[] (newest first) instead of previousSummary?: string
  • Keep: customInstructions, model, signal
  • Drop: resolveApiKey (use modelRegistry.getApiKey()), cutPoint, entries
• Update hook example custom-compaction.ts to use new API
• Update getSessionFile() to return string | undefined for in-memory sessions
• Update before_switch to have targetSessionFile, switch to have previousSessionFile

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

• fromId field references the abandoned leaf
• fromHook field distinguishes pi-generated vs hook-generated summaries
• details field 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 LabelEntry type with targetId and label fields
• 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?: string to SessionTreeNode, populated by getTree()
• Display labels in UI (tree-selector shows labels)
• /label command (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 SessionManager
• buildSessionContext() includes custom_message entries as user messages
• Exported from main index
• TUI rendering:
  • display: false - hidden entirely
  • display: 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 _hookData marker on AppMessage
• message_end handler routes based on marker presence
• AgentSession.sendHookMessage() handles three cases:
  • Streaming: queues via agent.queueMessage(), loop processes and emits message_end
  • Not streaming + triggerTurn: direct append + agent.continue()
  • Not streaming + no trigger: direct append only
• 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 (use sendMessage() with triggerTurn: true to 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:

• #324 - before_agent_start proposal
• #330 - Dynamic Context Pruning (why context was added)

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 HookMessage type (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 (uses createBranchedSession())
• /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

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
  • HookCommandContext interface and handler patterns
  • HookMessage<T> type
  • Updated event signatures (SessionEventBase, before_compact, etc.)
• docs/hooks-v2.md - Review/merge or remove if obsolete
• docs/sdk.md - Update for:
  • HookMessage and isHookMessage()
  • Agent.prompt(AppMessage) overload
  • Session v2 tree structure
  • SessionManager API changes
• docs/session.md - Update for v2 tree structure, new entry types
• docs/custom-tools.md - Check if hook changes affect custom tools
• docs/rpc.md - Check if hook commands work in RPC mode
• docs/skills.md - Review for any hook-related updates
• docs/extension-loading.md - Review
• docs/theme.md - Add customMessageBg/Text/Label color tokens
• README.md - Update hook examples if any

Examples

Review and update examples:

• examples/hooks/ - Update existing, add new examples:
  • Review custom-compaction.ts for new API
  • Add registerCommand() example
  • Add sendMessage() example
  • Add registerCustomMessageRenderer() example
• examples/sdk/ - Update for new session/hook APIs
• examples/custom-tools/ - Review for compatibility

---

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