# Session Tree Implementation Plan

Reference: [session-tree.md](./session-tree.md)

## Phase 1: SessionManager Core ✅

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

## Phase 2: Migration ✅

- [x] Add `CURRENT_SESSION_VERSION = 2` constant
- [x] Implement `migrateV1ToV2()` with extensible migration chain
- [x] Update `setSessionFile()` to detect version and migrate
- [x] Implement `_rewriteFile()` for post-migration persistence
- [x] Handle `firstKeptEntryIndex` → `firstKeptEntryId` conversion in migration

## Phase 3: Branching ✅

- [x] Implement `branch(id)` - switch leaf pointer
- [x] Implement `branchWithSummary(id, summary)` - create summary entry
- [x] Implement `createBranchedSession(leafId)` - extract path to new file
- [x] Update `AgentSession.branch()` to use new API

## Phase 4: Compaction Integration ✅

- [x] Update `compaction.ts` to work with IDs
- [x] Update `prepareCompaction()` to return `firstKeptEntryId`
- [x] Update `compact()` to return `CompactionResult` with `firstKeptEntryId`
- [x] Update `AgentSession` compaction methods
- [x] Add `firstKeptEntryId` to `before_compact` hook event

## Phase 5: Testing ✅

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

---

## Remaining Work

### Compaction Refactor

- [x] Use `CompactionResult` type for hook return value
- [x] Make `CompactionEntry<T>` generic with optional `details?: T` field for hook-specific data
- [x] Make `CompactionResult<T>` generic to match
- [x] Update `SessionEventBase` to pass `sessionManager` and `modelRegistry` instead of derived fields
- [x] 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`
- [x] Update hook example `custom-compaction.ts` to use new API
- [x] Update `getSessionFile()` to return `string | undefined` for in-memory sessions
- [x] Update `before_switch` to have `targetSessionFile`, `switch` to have `previousSessionFile`

Reference: [#314](https://github.com/badlogic/pi-mono/pull/314) - Structured compaction with anchored iterative summarization needs `details` field to store `ArtifactIndex` and version markers.

### Branch Summary Design

Current type:
```typescript
export interface BranchSummaryEntry extends SessionEntryBase {
  type: "branch_summary";
  summary: string;
}
```

Questions to resolve:
- [ ] Add `abandonedLeafId` field to reference what was abandoned?
- [ ] Store metadata about why the branch happened?
- [ ] Who generates the summary - user, LLM, or both options?
- [ ] Design and implement branch summarizer
- [ ] Add tests for `branchWithSummary()` flow

### Entry Labels ✅

- [x] Add `LabelEntry` type with `targetId` and `label` fields
- [x] Add `labelsById: Map<string, string>` private field
- [x] Build labels map in `_buildIndex()` via linear scan
- [x] Add `getLabel(id)` method
- [x] Add `appendLabelChange(targetId, label)` method (undefined clears)
- [x] Update `createBranchedSession()` to filter out LabelEntry and recreate from resolved map
- [x] `buildSessionContext()` already ignores LabelEntry (only handles message types)
- [x] Add `label?: string` to `SessionTreeNode`, populated by `getTree()`
- [ ] Display labels in UI (tree view, path view) - deferred to UI phase
- [ ] `/label` command - deferred to UI phase

### CustomMessageEntry<T>

Hook-injected messages that participate in LLM context. Unlike `CustomEntry<T>` (for hook state only), these are sent to the model.

```typescript
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:
- [x] Type definition matching plan
- [x] `appendCustomMessageEntry(customType, content, display, details?)` in SessionManager
- [x] `buildSessionContext()` includes custom_message entries as user messages
- [x] Exported from main index
- [x] TUI rendering:
  - `display: false` - hidden entirely
  - `display: true` - rendered with purple styling (customMessageBg, customMessageText, customMessageLabel theme colors)
  - [x] `registerCustomMessageRenderer(customType, renderer)` in HookAPI for custom renderers
  - [x] 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.

```typescript
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):

```typescript
appendEntry(customType: string, data?: unknown): void;
```

Calls `sessionManager.appendCustomEntry()` directly.

**New: `registerCommand()` (types ✅, wiring TODO)**

```typescript
// 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()):
- [x] Add hook commands to autocomplete in interactive-mode
- [x] `_tryExecuteHookCommand()` in AgentSession handles command execution
- [x] Build HookCommandContext with ui (from hookRunner), exec, sessionManager, etc.
- [x] If handler returns string, use as prompt text
- [x] If handler returns undefined, return early (no LLM call)
- [x] Works for all modes (interactive, RPC, print) via shared AgentSession

**New: `ui.custom()` ✅**

For arbitrary hook UI with keyboard focus:

```typescript
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.

```typescript
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](https://github.com/badlogic/pi-mono/discussions/330))

Non-destructive pruning of tool results to reduce context size:

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

Reference: [#324](https://github.com/badlogic/pi-mono/issues/324)

**Current `context` event:**
- Fires before each LLM call within the agent loop
- Receives `AgentMessage[]` (deep copy, safe to modify)
- Modifications are transient (not persisted to session)
- No TUI visibility of what was changed
- Use case: non-destructive pruning, dynamic context manipulation

**Problem:** `AgentMessage` includes custom types (hookMessage, bashExecution, etc.) that need conversion to LLM `Message[]` before sending. Need to verify:
- [ ] Where does `AgentMessage[]` → `Message[]` conversion happen relative to `context` event?
- [ ] Should hooks work with `AgentMessage[]` or `Message[]`?
- [ ] Is the current abstraction level correct?

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

**Design questions:**
- [ ] Should `before_agent_start` create a new message type (`SystemMessage` with `role: "system"`)?
- [ ] How should it render in TUI? (label when collapsed, full content when expanded)
- [ ] How does it interact with compaction? (treated like user messages?)
- [ ] Can hook return multiple messages or just one?

**Implementation sketch:**
```typescript
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

Design new commands based on refactored SessionManager:

**`/branch`** - Current behavior (creates new session file from path)
- [ ] Review if this is still the right UX with tree structure
- [ ] Consider: should this use `createBranchedSession()` or `branch()`?

**`/branch-here`** - In-place branching (new)
- [ ] Use `branch(id)` to move leaf pointer without creating new file
- [ ] Subsequent messages become new branch in same file
- [ ] Design: how to select branch point? (similar to current `/branch` UI?)

**`/branches`** - List/navigate branches (new)
- [ ] Show tree structure or list of branch points
- [ ] Allow switching between branches (move leaf pointer)
- [ ] Show current position in tree

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