pi-code 0.1.0

package/LICENSE

@@ -0,0 +1,22 @@
+MIT License
+
+Copyright (c) 2025 Mario Zechner
+Copyright (c) 2026 ilovepixelart
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

package/README.md

@@ -0,0 +1,51 @@
+# pi-code
+
+Claude Code experience for the [pi](https://pi.dev) coding agent, in one package. Point pi at a project that already has a `.claude/` directory and it reads your existing config: rules, commands, skills, hooks, output styles, MCP servers, and agents. It also adds the Claude Code features pi lacks: a todo overlay, checkpoints, memory, web search, and subagents.
+
+![pi-code demo](demos/hero.gif)
+
+## Install
+
+```bash
+pi install pi-code                       # from npm (when published)
+pi install git:github.com/ilovepixelart/pi-code
+pi install ~/Documents/pi-code           # local path, then /reload after edits
+```
+
+One `pi install`, everything below loads. Each feature is an extension under [`extensions/`](extensions).
+
+## What it does
+
+| Feature | Reads / provides | Extension |
+|---|---|---|
+| Global + project rules | `~/.claude/rules`, `.claude/rules` (+ `paths:` frontmatter scoping) | `claude-rules.ts` |
+| Custom slash commands | `.claude/commands/*.md` → pi prompt templates | `commands.ts` |
+| Skills | `.claude/skills` → pi skill discovery | `skills.ts` |
+| Hooks | `.claude/settings.json` hooks on pi lifecycle events | `hooks.ts` |
+| Output styles | `.claude/output-styles` + active `outputStyle`, `/output-style` switcher | `output-styles.ts` |
+| CLAUDE.md `@imports` | resolves `@path` imports pi's native loader skips | `context-imports.ts` |
+| MCP servers | `.mcp.json`, `~/.claude.json`, `.pi/mcp.json`; stdio + HTTP | `mcp.ts` |
+| Subagents / Task | `~/.claude/agents` + project `.claude/agents`, background runs | `subagent/` |
+| Plan mode | `plan_mode_complete` tool, exact tool snapshot/restore | `plan-mode/` |
+| Todo list | persistent overlay, status machine, compaction-safe | `todo.ts` |
+| Checkpoints / rewind | shadow-repo snapshots, hard-reset restore | `git-checkpoint.ts` |
+| Persistent memory | per-project memories, index injected each session | `memory.ts` |
+| WebSearch / WebFetch | key-free DuckDuckGo search, SSRF-guarded fetch | `web.ts` |
+| AskUserQuestion | vendored example | `question.ts` |
+| Statusline | turn state + session cost | `status-line.ts` |
+| Notifications | vendored example | `notify.ts` |
+
+`CLAUDE.md` itself needs no extension: pi loads `CLAUDE.md` / `AGENTS.md` context files natively (global + walking cwd to root). `context-imports.ts` only adds the `@import` resolution pi's loader lacks, appending the imported files without re-injecting the base.
+
+Vendored bases (`question`, `notify`, `status-line`) come from pi's MIT example extensions (see [LICENSE](LICENSE)). Personal config lives in [dot-pi](https://github.com/ilovepixelart/dot-pi) (`~/.pi/agent`).
+
+## Development
+
+```bash
+npm install
+npm run check           # biome + strict tsc + vitest, the whole gate
+scripts/e2e.sh          # drives the real pi TUI via tmux (needs a working model)
+scripts/record-demos.sh # re-records demos/*.tape with vhs at low thinking
+```
+
+Extensions live in `extensions/`, tests in `tests/`. Install locally with `pi install ~/Documents/pi-code`, then `/reload` after edits.

package/extensions/claude-rules.ts

@@ -0,0 +1,136 @@
+/**
+ * Claude Rules Extension
+ *
+ * Replicates Claude Code's rules loading:
+ * - Global rules (~/.claude/rules/*.md) are inlined in full into the system prompt.
+ * - Project rules (.claude/rules/*.md) are listed as pointers the agent can read on demand.
+ *
+ * Path-scoped rules: a rule file may declare `paths:` frontmatter (a glob or
+ * list of globs). Project-rule pointers surface that scope so the agent knows
+ * to read the rule when working on matching files. Frontmatter is stripped
+ * from inlined global rules.
+ *
+ * Adapted from the pi v0.74.2 claude-rules example.
+ */
+
+import * as fs from 'node:fs'
+import * as os from 'node:os'
+import * as path from 'node:path'
+import type { ExtensionAPI } from '@earendil-works/pi-coding-agent'
+
+export interface Frontmatter {
+  paths: string[]
+  body: string
+}
+
+function unquote(value: string): string {
+  return value.replace(/^["']|["']$/g, '')
+}
+
+function splitInline(value: string): string[] {
+  return value
+    .replace(/^\[|\]$/g, '')
+    .split(',')
+    .map((entry) => unquote(entry.trim()))
+    .filter(Boolean)
+}
+
+function parsePaths(frontmatter: string): string[] {
+  const lines = frontmatter.split('\n')
+  const index = lines.findIndex((line) => /^\s*paths\s*:/.test(line))
+  if (index === -1) return []
+  const inline = lines[index].replace(/^\s*paths\s*:/, '').trim()
+  if (inline) return splitInline(inline)
+  const items: string[] = []
+  for (let i = index + 1; i < lines.length; i++) {
+    const match = /^\s*-\s*(.+)$/.exec(lines[i])
+    if (!match) break
+    items.push(unquote(match[1].trim()))
+  }
+  return items
+}
+
+/** Split YAML-ish frontmatter off the front of a rule file, extracting `paths`. */
+export function parseFrontmatter(content: string): Frontmatter {
+  const match = /^---\r?\n([\s\S]*?)\r?\n---\r?\n?/.exec(content)
+  if (!match) return { paths: [], body: content }
+  return { paths: parsePaths(match[1]), body: content.slice(match[0].length) }
+}
+
+/** A project-rule pointer line, annotated with its path scope when present. */
+export function formatRulePointer(rel: string, paths: string[]): string {
+  const ref = `- .claude/rules/${rel}`
+  return paths.length > 0 ? `${ref} — applies when working on: ${paths.join(', ')}` : ref
+}
+
+/** Recursively find all .md files in a directory. */
+function findMarkdownFiles(dir: string, basePath = ''): string[] {
+  if (!fs.existsSync(dir)) return []
+  const results: string[] = []
+  for (const entry of fs.readdirSync(dir, { withFileTypes: true })) {
+    const relativePath = basePath ? `${basePath}/${entry.name}` : entry.name
+    if (entry.isDirectory()) {
+      results.push(...findMarkdownFiles(path.join(dir, entry.name), relativePath))
+    } else if (entry.isFile() && entry.name.endsWith('.md')) {
+      results.push(relativePath)
+    }
+  }
+  return results
+}
+
+function readGlobalRules(globalRulesDir: string): string {
+  return findMarkdownFiles(globalRulesDir)
+    .map((file) => parseFrontmatter(fs.readFileSync(path.join(globalRulesDir, file), 'utf-8')).body.trim())
+    .filter((content) => content.length > 0)
+    .join('\n\n')
+}
+
+interface ProjectRule {
+  rel: string
+  paths: string[]
+}
+
+function readProjectRules(projectRulesDir: string): ProjectRule[] {
+  return findMarkdownFiles(projectRulesDir).map((rel) => {
+    try {
+      return { rel, paths: parseFrontmatter(fs.readFileSync(path.join(projectRulesDir, rel), 'utf-8')).paths }
+    } catch {
+      return { rel, paths: [] }
+    }
+  })
+}
+
+export default function claudeRulesExtension(pi: ExtensionAPI) {
+  const globalRulesDir = path.join(os.homedir(), '.claude', 'rules')
+  let globalRules = ''
+  let projectRules: ProjectRule[] = []
+
+  pi.on('session_start', async (_event, ctx) => {
+    globalRules = readGlobalRules(globalRulesDir)
+    // Project rule filenames and their paths: frontmatter are surfaced in the system prompt,
+    // so read them only for a trusted project.
+    const trusted = ctx.isProjectTrusted?.() ?? false
+    projectRules = trusted ? readProjectRules(path.join(ctx.cwd, '.claude', 'rules')) : []
+
+    if (globalRules.length > 0 || projectRules.length > 0) {
+      ctx.ui.notify(`Rules loaded: global ${globalRules.length > 0 ? 'yes' : 'no'}, project ${projectRules.length}`, 'info')
+    }
+  })
+
+  pi.on('before_agent_start', async (event) => {
+    let addition = ''
+
+    if (globalRules.length > 0) {
+      addition += `\n\n## Global Rules\n\nThese rules always apply:\n\n${globalRules}`
+    }
+
+    if (projectRules.length > 0) {
+      const rulesList = projectRules.map((rule) => formatRulePointer(rule.rel, rule.paths)).join('\n')
+      addition += `\n\n## Project Rules\n\nThe following project rules are available in .claude/rules/:\n\n${rulesList}\n\nRead the relevant rule file with the read tool before working on the files it covers; rules with an "applies when" scope are path-scoped.`
+    }
+
+    if (addition.length === 0) return
+
+    return { systemPrompt: event.systemPrompt + addition }
+  })
+}

package/extensions/commands.ts

@@ -0,0 +1,45 @@
+/**
+ * Claude Commands Extension
+ *
+ * Bridges Claude Code's custom slash commands into pi. On resources_discover
+ * it hands pi the existing `.claude/commands` directories (user then project)
+ * as prompt-template paths, so `/name` invokes `.claude/commands/name.md` the
+ * same way pi loads its own `.pi/prompts`. pi's `$ARGUMENTS` / `$1` / `${1:-x}`
+ * substitution overlaps Claude Code's, so most command files work unchanged.
+ *
+ * Not bridged (pi's prompt engine ignores them): `!` bash execution, `@` file
+ * refs, `allowed-tools`/`model` frontmatter, and namespaced subdirectories
+ * (discovery is non-recursive).
+ *
+ * Docs: https://code.claude.com/docs/en/slash-commands.md
+ */
+
+import * as fs from 'node:fs'
+import * as os from 'node:os'
+import * as path from 'node:path'
+import type { ExtensionAPI } from '@earendil-works/pi-coding-agent'
+
+function isDirectory(target: string): boolean {
+  try {
+    return fs.statSync(target).isDirectory()
+  } catch {
+    return false
+  }
+}
+
+/** Existing `.claude/commands` directories, user first then project. */
+export function commandDirs(cwd: string, home: string): string[] {
+  const candidates = [path.join(home, '.claude', 'commands'), path.join(cwd, '.claude', 'commands')]
+  const dirs: string[] = []
+  for (const dir of candidates) {
+    if (!dirs.includes(dir) && isDirectory(dir)) dirs.push(dir)
+  }
+  return dirs
+}
+
+export default function commandsExtension(pi: ExtensionAPI) {
+  pi.on('resources_discover', async (_event, ctx) => {
+    const promptPaths = commandDirs(ctx.cwd, os.homedir())
+    return promptPaths.length > 0 ? { promptPaths } : undefined
+  })
+}

package/extensions/context-imports.ts

@@ -0,0 +1,118 @@
+/**
+ * Context Imports Extension
+ *
+ * pi loads CLAUDE.md / AGENTS.md context files natively but does not resolve
+ * Claude Code's `@path` imports inside them. This fills that one gap: on
+ * before_agent_start it reads the already-loaded context files from
+ * systemPromptOptions, resolves any `@path` imports (recursive, depth-capped,
+ * cycle-safe; ~ expands to home, relative paths resolve against the importing
+ * file), and appends ONLY the imported content. pi already injected the base
+ * files, so nothing is duplicated.
+ *
+ * Security: context files can come from an untrusted project, so imports are
+ * confined (after resolving symlinks) to the working directory and the user's
+ * own ~/.claude and ~/.pi config roots. An import that escapes those roots
+ * (absolute paths, ~/.ssh, ../.. traversal, symlinks) is ignored, so a hostile
+ * CLAUDE.md cannot read arbitrary files into the prompt. Imports inside fenced
+ * code blocks are also skipped.
+ *
+ * Docs: https://code.claude.com/docs/en/memory.md (imports)
+ */
+
+import * as fs from 'node:fs'
+import * as os from 'node:os'
+import * as path from 'node:path'
+import type { ExtensionAPI } from '@earendil-works/pi-coding-agent'
+
+const MAX_IMPORT_DEPTH = 5
+
+export function expandHome(target: string, home: string): string {
+  if (target === '~') return home
+  if (target.startsWith('~/')) return path.join(home, target.slice(2))
+  return target
+}
+
+function isUnder(target: string, roots: string[]): boolean {
+  return roots.some((root) => target === root || target.startsWith(root + path.sep))
+}
+
+/** Realpath the roots that exist; used both to seed and to bound the import search. */
+export function realRoots(candidates: string[]): string[] {
+  const roots: string[] = []
+  for (const candidate of candidates) {
+    try {
+      roots.push(fs.realpathSync(candidate))
+    } catch {
+      // a root that does not exist has nothing under it to allow
+    }
+  }
+  return roots
+}
+
+export interface ImportedFile {
+  path: string
+  body: string
+}
+
+/**
+ * Collect the contents of every file transitively imported via `@path`, in
+ * discovery order. Imports are resolved through symlinks and kept within
+ * `allowedRoots` (which must already be realpath'd).
+ */
+export function collectImports(content: string, fromDir: string, home: string, allowedRoots: string[], seen: Set<string>, depth = 0): ImportedFile[] {
+  if (depth >= MAX_IMPORT_DEPTH) return []
+  const out: ImportedFile[] = []
+  let inFence = false
+  for (const line of content.split('\n')) {
+    if (line.trimStart().startsWith('```')) {
+      inFence = !inFence
+      continue
+    }
+    if (inFence) continue
+    for (const match of line.matchAll(/(^|\s)@(\S+)/g)) {
+      const resolved = path.resolve(fromDir, expandHome(match[2], home))
+      let real: string
+      try {
+        real = fs.realpathSync(resolved)
+      } catch {
+        continue
+      }
+      if (seen.has(real)) continue
+      seen.add(real)
+      if (!isUnder(real, allowedRoots)) continue
+      let body: string
+      try {
+        // real may be a directory (EISDIR) or vanish after the realpath (ENOENT/EACCES).
+        body = fs.readFileSync(real, 'utf-8')
+      } catch {
+        continue
+      }
+      out.push({ path: real, body: body.trim() })
+      out.push(...collectImports(body, path.dirname(real), home, allowedRoots, seen, depth + 1))
+    }
+  }
+  return out
+}
+
+export default function contextImportsExtension(pi: ExtensionAPI) {
+  pi.on('before_agent_start', async (event) => {
+    const contextFiles: Array<{ path: string; content: string }> = event.systemPromptOptions?.contextFiles ?? []
+    if (contextFiles.length === 0) return
+
+    const home = os.homedir()
+    const cwd = event.systemPromptOptions?.cwd ?? process.cwd()
+    const allowedRoots = realRoots([cwd, path.join(home, '.claude'), path.join(home, '.pi')])
+    // Seed with the loaded context file paths so pi's own files are never re-imported.
+    const seen = realRoots(contextFiles.map((file) => file.path))
+    const seenSet = new Set(seen)
+
+    const imported: ImportedFile[] = []
+    for (const file of contextFiles) {
+      imported.push(...collectImports(file.content, path.dirname(file.path), home, allowedRoots, seenSet))
+    }
+    if (imported.length === 0) return
+
+    const section = imported.map((entry) => `### ${entry.path}\n\n${entry.body}`).join('\n\n')
+    return { systemPrompt: `${event.systemPrompt}\n\n## Imported context (@)\n\n${section}` }
+  })
+}

package/extensions/git-checkpoint.ts

@@ -0,0 +1,199 @@
+/**
+ * Git Checkpoint Extension
+ *
+ * Claude Code style /rewind built on a per-session shadow git repo.
+ *
+ * Snapshots commit the entire working tree (untracked files included, the
+ * project's .gitignore is honored) into a bare repo under
+ * ~/.pi/agent/checkpoints/<session>, using --git-dir/--work-tree so the
+ * project's own git state is never touched. Each user prompt gets one
+ * checkpoint persisted as {entryId, ref, prompt, createdAt} in the session
+ * file, so /rewind works across restarts, resumes, and forks. Code restore
+ * checks the snapshot out over the working tree, resetting file contents to
+ * the checkpoint (files created after the checkpoint are left in place).
+ */
+
+import * as os from 'node:os'
+import * as path from 'node:path'
+import type { ExtensionAPI, ExtensionCommandContext, ExtensionContext } from '@earendil-works/pi-coding-agent'
+
+const CUSTOM_TYPE = 'git-checkpoint'
+const PROMPT_SNIPPET_LENGTH = 60
+const RESTORE_MODES = ['Code and conversation', 'Conversation only', 'Code only']
+
+interface Checkpoint {
+  entryId: string
+  ref: string
+  prompt: string
+  createdAt: string
+}
+
+export function sessionSlug(sessionFile: string | undefined): string {
+  if (!sessionFile) return `ephemeral-${process.pid}`
+  return path.basename(sessionFile).replace(/[^\w.-]+/g, '_')
+}
+
+function extractText(content: unknown): string {
+  if (typeof content === 'string') return content
+  if (!Array.isArray(content)) return ''
+  return content
+    .filter((part) => part?.type === 'text' && typeof part.text === 'string')
+    .map((part) => part.text)
+    .join(' ')
+}
+
+function promptSnippet(content: unknown): string {
+  const text = extractText(content).replace(/\s+/g, ' ').trim()
+  if (text.length <= PROMPT_SNIPPET_LENGTH) return text
+  return `${text.slice(0, PROMPT_SNIPPET_LENGTH)}…`
+}
+
+function findLastUserMessage(ctx: ExtensionContext): { entryId: string; prompt: string } | undefined {
+  const branch = ctx.sessionManager.getBranch()
+  for (let i = branch.length - 1; i >= 0; i--) {
+    const entry = branch[i]
+    if (entry?.type === 'message' && entry.message.role === 'user') {
+      return { entryId: entry.id, prompt: promptSnippet(entry.message.content) }
+    }
+  }
+  return undefined
+}
+
+function checkpointLabel(checkpoint: Checkpoint, index: number): string {
+  const time = new Date(checkpoint.createdAt).toLocaleTimeString()
+  const marker = checkpoint.ref ? '' : ' [no code snapshot]'
+  return `${index + 1}. ${time}  ${checkpoint.prompt || '(empty prompt)'}${marker}`
+}
+
+export default function (pi: ExtensionAPI) {
+  const checkpoints = new Map<string, Checkpoint>()
+  let pending: { ref: string; createdAt: string } | undefined
+  let shadowDir: string | undefined
+  let workTree: string | undefined
+
+  function gitShadow(args: string[]): ReturnType<ExtensionAPI['exec']> {
+    if (!shadowDir || !workTree) return Promise.resolve({ stdout: '', stderr: 'shadow repo not initialized', code: 1, killed: false })
+    return pi.exec('git', ['--git-dir', shadowDir, '--work-tree', workTree, ...args], { cwd: workTree })
+  }
+
+  async function ensureShadow(ctx: ExtensionContext): Promise<void> {
+    workTree = ctx.cwd
+    const sessionFile = (ctx.sessionManager as { getSessionFile?: () => string | undefined }).getSessionFile?.()
+    shadowDir = path.join(os.homedir(), '.pi', 'agent', 'checkpoints', sessionSlug(sessionFile))
+    const check = await pi.exec('git', ['--git-dir', shadowDir, 'rev-parse', '--git-dir'], { cwd: ctx.cwd })
+    if (check.code !== 0) {
+      await pi.exec('git', ['init', '--bare', '-b', 'main', shadowDir], { cwd: ctx.cwd })
+      await pi.exec('git', ['--git-dir', shadowDir, 'config', 'user.email', 'checkpoint@pi-code'], { cwd: ctx.cwd })
+      await pi.exec('git', ['--git-dir', shadowDir, 'config', 'user.name', 'pi-code-checkpoint'], { cwd: ctx.cwd })
+    }
+  }
+
+  async function snapshot(): Promise<{ ref: string; createdAt: string } | undefined> {
+    const createdAt = new Date().toISOString()
+    const add = await gitShadow(['add', '-A'])
+    if (add.code !== 0) return undefined
+    const commit = await gitShadow(['commit', '-m', 'checkpoint'])
+    if (commit.code !== 0) {
+      // nothing changed since the last snapshot: reuse HEAD, or create the first empty commit
+      const head = await gitShadow(['rev-parse', 'HEAD'])
+      if (head.code === 0) return { ref: head.stdout.trim(), createdAt }
+      const empty = await gitShadow(['commit', '--allow-empty', '-m', 'checkpoint'])
+      if (empty.code !== 0) return undefined
+    }
+    const sha = await gitShadow(['rev-parse', 'HEAD'])
+    return sha.code === 0 ? { ref: sha.stdout.trim(), createdAt } : undefined
+  }
+
+  async function restoreCode(ctx: ExtensionCommandContext, checkpoint: Checkpoint): Promise<boolean> {
+    if (!checkpoint.ref) {
+      ctx.ui.notify('Checkpoint has no code snapshot; code left untouched', 'warning')
+      return true
+    }
+    const result = await gitShadow(['checkout', '-f', checkpoint.ref, '--', '.'])
+    if (result.code !== 0) {
+      ctx.ui.notify(`Code restore failed: ${result.stderr.trim()}`, 'warning')
+      return false
+    }
+    return true
+  }
+
+  async function restoreConversation(ctx: ExtensionCommandContext, entryId: string): Promise<boolean> {
+    try {
+      // navigateTree's published type omits editorText, but it is present at runtime (docs/extensions.md)
+      const result = (await ctx.navigateTree(entryId, { summarize: false })) as { cancelled: boolean; editorText?: string }
+      if (result.cancelled) return false
+      if (typeof result.editorText === 'string') ctx.ui.setEditorText(result.editorText)
+      return true
+    } catch (error) {
+      const message = error instanceof Error ? error.message : String(error)
+      ctx.ui.notify(`Conversation restore failed: ${message}`, 'error')
+      return false
+    }
+  }
+
+  async function runRestoreMode(ctx: ExtensionCommandContext, checkpoint: Checkpoint): Promise<void> {
+    const mode = await ctx.ui.select('Restore mode:', [...RESTORE_MODES])
+    if (!mode) return
+    if (mode !== 'Conversation only' && !(await restoreCode(ctx, checkpoint))) return
+    if (mode !== 'Code only' && !(await restoreConversation(ctx, checkpoint.entryId))) return
+    ctx.ui.notify('Rewind complete', 'info')
+  }
+
+  pi.on('session_start', async (_event, ctx) => {
+    await ensureShadow(ctx)
+    checkpoints.clear()
+    for (const entry of ctx.sessionManager.getEntries()) {
+      if (entry.type !== 'custom' || entry.customType !== CUSTOM_TYPE) continue
+      const checkpoint = entry.data as Checkpoint | undefined
+      if (checkpoint?.entryId) checkpoints.set(checkpoint.entryId, checkpoint)
+    }
+  })
+
+  // Snapshot code state before the LLM acts in this turn. The user message
+  // that started the turn is not persisted yet at turn_start (it lands on
+  // message_end), so the checkpoint is only keyed and saved at turn_end.
+  pi.on('turn_start', async () => {
+    pending = await snapshot()
+  })
+
+  pi.on('turn_end', async (_event, ctx) => {
+    const snap = pending
+    pending = undefined
+    if (!snap) return
+
+    const target = findLastUserMessage(ctx)
+    if (!target || checkpoints.has(target.entryId)) return
+
+    const checkpoint: Checkpoint = { entryId: target.entryId, ref: snap.ref, prompt: target.prompt, createdAt: snap.createdAt }
+    checkpoints.set(checkpoint.entryId, checkpoint)
+    pi.appendEntry(CUSTOM_TYPE, checkpoint)
+  })
+
+  pi.registerCommand('rewind', {
+    description: 'Rewind code and/or conversation to a previous checkpoint',
+    handler: async (_args, ctx) => {
+      if (!ctx.hasUI) return
+      const ordered = [...checkpoints.values()].reverse()
+      if (ordered.length === 0) {
+        ctx.ui.notify('No checkpoints recorded yet', 'info')
+        return
+      }
+      const labels = ordered.map(checkpointLabel)
+      const choice = await ctx.ui.select('Rewind to checkpoint:', labels)
+      if (!choice) return
+      const checkpoint = ordered[labels.indexOf(choice)]
+      if (checkpoint) await runRestoreMode(ctx, checkpoint)
+    },
+  })
+
+  pi.on('session_before_fork', async (event, ctx) => {
+    const checkpoint = checkpoints.get(event.entryId)
+    if (!checkpoint?.ref || !ctx.hasUI) return
+
+    const choice = await ctx.ui.select('Restore code state?', ['Yes, restore code to that point', 'No, keep current code'])
+    if (choice?.startsWith('Yes')) {
+      const result = await gitShadow(['checkout', '-f', checkpoint.ref, '--', '.'])
+      ctx.ui.notify(result.code === 0 ? 'Code restored to checkpoint' : `Restore failed: ${result.stderr.trim()}`, result.code === 0 ? 'info' : 'warning')
+    }
+  })
+}