Forráskód Böngészése

feat: checkpoint/rollback functionality for session recovery

Adds checkpoint/rollback feature inspired by pi-rollback:

User commands:
- /checkpoint save <name> [description] - save current state
- /checkpoint list - show all checkpoints
- /checkpoint delete <name> - remove checkpoint
- /rollback <name> - rollback to checkpoint

Agent tools:
- checkpoint tool - agent can checkpoint before risky ops
- rollback tool - agent can self-recover from bad paths

Implementation:
- Checkpoints anchor to last user message
- Rollback uses OpenCode's session.revert API
- Stored in .slim/checkpoints.json
- Auto-cleanup on session deletion

Closes #374
Alvin Real 3 hónapja
szülő
commit
c7a3e4cd4a

+ 179 - 0
src/checkpoints/commands.ts

@@ -0,0 +1,179 @@
+/**
+ * Checkpoint slash commands - user-facing interface.
+ */
+
+import type { CheckpointManager } from './manager';
+
+export interface CommandContext {
+  sessionID: string;
+  arguments: string;
+}
+
+export interface CommandResult {
+  parts: Array<{ type: string; text?: string }>;
+}
+
+export function createCheckpointCommands(manager: CheckpointManager) {
+  async function handleCheckpointSave(
+    ctx: CommandContext,
+  ): Promise<CommandResult> {
+    const args = ctx.arguments.trim();
+    const match = args.match(/^([^\s]+)(?:\s+(.+))?$/);
+
+    if (!match) {
+      return {
+        parts: [
+          {
+            type: 'text',
+            text: 'Usage: /checkpoint save <name> [description]\nExample: /checkpoint save before-refactor Starting auth refactor',
+          },
+        ],
+      };
+    }
+
+    const [, name, description] = match;
+    const result = await manager.saveCheckpoint(ctx.sessionID, name, description);
+
+    if (result.success) {
+      return {
+        parts: [
+          {
+            type: 'text',
+            text: `✓ Checkpoint "${name}" saved.${description ? `\n  Description: ${description}` : ''}`,
+          },
+        ],
+      };
+    }
+
+    return {
+      parts: [
+        {
+          type: 'text',
+          text: `✗ Failed to save checkpoint: ${result.error}`,
+        },
+      ],
+    };
+  }
+
+  async function handleCheckpointList(
+    ctx: CommandContext,
+  ): Promise<CommandResult> {
+    const checkpoints = await manager.listCheckpoints(ctx.sessionID);
+
+    if (checkpoints.length === 0) {
+      return {
+        parts: [
+          {
+            type: 'text',
+            text: 'No checkpoints saved for this session.\nUse `/checkpoint save <name>` to create one.',
+          },
+        ],
+      };
+    }
+
+    const lines = checkpoints.map((cp) => {
+      const date = new Date(cp.createdAt).toLocaleString();
+      const desc = cp.description ? `\n    ${cp.description}` : '';
+      return `  • ${cp.name} (${date})${desc}`;
+    });
+
+    return {
+      parts: [
+        {
+          type: 'text',
+          text: `Checkpoints for this session:\n${lines.join('\n')}\n\nUse /rollback <name> to restore a checkpoint.`,
+        },
+      ],
+    };
+  }
+
+  async function handleCheckpointDelete(
+    ctx: CommandContext,
+  ): Promise<CommandResult> {
+    const name = ctx.arguments.trim();
+
+    if (!name) {
+      return {
+        parts: [
+          {
+            type: 'text',
+            text: 'Usage: /checkpoint delete <name>\nExample: /checkpoint delete before-refactor',
+          },
+        ],
+      };
+    }
+
+    const result = await manager.deleteCheckpoint(ctx.sessionID, name);
+
+    if (result.success) {
+      return {
+        parts: [
+          {
+            type: 'text',
+            text: `✓ Checkpoint "${name}" deleted.`,
+          },
+        ],
+      };
+    }
+
+    return {
+      parts: [
+        {
+          type: 'text',
+          text: `✗ Failed to delete checkpoint: ${result.error}`,
+        },
+      ],
+    };
+  }
+
+  async function handleRollback(
+    ctx: CommandContext,
+  ): Promise<CommandResult> {
+    const name = ctx.arguments.trim();
+
+    if (!name) {
+      return {
+        parts: [
+          {
+            type: 'text',
+            text: 'Usage: /rollback <name>\nExample: /rollback before-refactor\n\nAvailable checkpoints:\n' +
+              (await manager.listCheckpoints(ctx.sessionID))
+                .map((cp) => `  • ${cp.name}`)
+                .join('\n') ||
+              '  (none)',
+          },
+        ],
+      };
+    }
+
+    const result = await manager.rollback(ctx.sessionID, name);
+
+    if (result.success) {
+      return {
+        parts: [
+          {
+            type: 'text',
+            text: `✓ ${result.message || `Rolled back to checkpoint "${name}".`}`,
+          },
+        ],
+      };
+    }
+
+    return {
+      parts: [
+        {
+            type: 'text',
+            text: `✗ Rollback failed: ${result.error}`,
+          },
+        ],
+      };
+    }
+  }
+
+  return {
+    handleCheckpointSave,
+    handleCheckpointList,
+    handleCheckpointDelete,
+    handleRollback,
+  };
+}

+ 17 - 0
src/checkpoints/index.ts

@@ -0,0 +1,17 @@
+/**
+ * Checkpoint module - session rollback functionality.
+ *
+ * Provides:
+ * - /checkpoint save <name> [description] - Save current state
+ * - /checkpoint list - Show all checkpoints
+ * - /checkpoint delete <name> - Remove a checkpoint
+ * - /rollback <name> - Rollback to checkpoint
+ * - checkpoint tool - Agent-driven checkpointing
+ * - rollback tool - Agent-driven rollback
+ */
+
+export { createCheckpointManager, type CheckpointManager } from './manager';
+export { createCheckpointTools, type CheckpointTools } from './tool';
+export { createCheckpointCommands, type CommandContext, type CommandResult } from './commands';
+export { CheckpointStorage } from './store';
+export type { Checkpoint, CheckpointStore } from './types';

+ 218 - 0
src/checkpoints/manager.ts

@@ -0,0 +1,218 @@
+/**
+ * Checkpoint manager - core logic for saving and restoring checkpoints.
+ */
+
+import type { PluginInput } from '@opencode-ai/plugin';
+import { log } from '../utils/logger';
+import { CheckpointStorage } from './store';
+import type { Checkpoint } from './types';
+
+export interface CheckpointManager {
+  /** Save a checkpoint for the current session */
+  saveCheckpoint(
+    sessionID: string,
+    name: string,
+    description?: string,
+  ): Promise<{ success: boolean; checkpoint?: Checkpoint; error?: string }>;
+
+  /** List all checkpoints for a session */
+  listCheckpoints(sessionID: string): Promise<Checkpoint[]>;
+
+  /** Delete a checkpoint by name */
+  deleteCheckpoint(
+    sessionID: string,
+    name: string,
+  ): Promise<{ success: boolean; error?: string }>;
+
+  /** Rollback to a checkpoint */
+  rollback(
+    sessionID: string,
+    name: string,
+  ): Promise<{ success: boolean; message?: string; error?: string }>;
+
+  /** Cleanup all checkpoints for a session */
+  cleanupSession(sessionID: string): Promise<number>;
+}
+
+interface MessageInfo {
+  id: string;
+  info?: { role: string; id?: string };
+}
+
+export function createCheckpointManager(ctx: PluginInput): CheckpointManager {
+  const storage = new CheckpointStorage(ctx.directory);
+  const client = ctx.client;
+
+  async function getLastUserMessage(sessionID: string): Promise<string | null> {
+    try {
+      const result = await client.session.messages({ path: { id: sessionID } });
+      const messages = (result.data || []) as MessageInfo[];
+
+      // Find the last user message
+      for (let i = messages.length - 1; i >= 0; i--) {
+        if (messages[i]?.info?.role === 'user') {
+          return messages[i]?.info?.id || messages[i]?.id || null;
+        }
+      }
+      return null;
+    } catch (err) {
+      log('[checkpoint] Failed to get messages', { error: String(err) });
+      return null;
+    }
+  }
+
+  async function getFirstUserMessageAfter(
+    sessionID: string,
+    anchorMessageID: string,
+  ): Promise<string | null> {
+    try {
+      const result = await client.session.messages({ path: { id: sessionID } });
+      const messages = (result.data || []) as MessageInfo[];
+
+      let foundAnchor = false;
+      for (const msg of messages) {
+        if (msg.id === anchorMessageID || msg.info?.id === anchorMessageID) {
+          foundAnchor = true;
+          continue;
+        }
+        if (foundAnchor && msg.info?.role === 'user') {
+          return msg.info?.id || msg.id;
+        }
+      }
+      return null;
+    } catch (err) {
+      log('[checkpoint] Failed to find message after anchor', { error: String(err) });
+      return null;
+    }
+  }
+
+  async function isSessionBusy(sessionID: string): Promise<boolean> {
+    try {
+      const status = await client.session.status();
+      const sessionStatus = (status.data as Record<string, { type?: string }> | undefined)?.[sessionID];
+      return sessionStatus?.type === 'busy';
+    } catch {
+      return false;
+    }
+  }
+
+  return {
+    async saveCheckpoint(sessionID, name, description) {
+      try {
+        // Check if checkpoint with this name already exists
+        const existing = await storage.getByName(sessionID, name);
+        if (existing) {
+          return {
+            success: false,
+            error: `Checkpoint "${name}" already exists. Delete it first or use a different name.`,
+          };
+        }
+
+        // Get the last user message as anchor
+        const anchorMessageID = await getLastUserMessage(sessionID);
+        if (!anchorMessageID) {
+          return {
+            success: false,
+            error: 'No user message found to anchor checkpoint. Send a message first.',
+          };
+        }
+
+        const checkpoint: Checkpoint = {
+          id: `cp-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`,
+          name,
+          description,
+          sessionID,
+          anchorMessageID,
+          directory: ctx.directory,
+          createdAt: Date.now(),
+        };
+
+        await storage.add(checkpoint);
+        log('[checkpoint] Saved', { name, sessionID, anchorMessageID });
+
+        return { success: true, checkpoint };
+      } catch (err) {
+        const error = String(err);
+        log('[checkpoint] Failed to save', { error, name, sessionID });
+        return { success: false, error };
+      }
+    },
+
+    async listCheckpoints(sessionID) {
+      return storage.listForSession(sessionID);
+    },
+
+    async deleteCheckpoint(sessionID, name) {
+      try {
+        const checkpoint = await storage.getByName(sessionID, name);
+        if (!checkpoint) {
+          return { success: false, error: `Checkpoint "${name}" not found.` };
+        }
+
+        await storage.delete(checkpoint.id);
+        log('[checkpoint] Deleted', { name, sessionID });
+        return { success: true };
+      } catch (err) {
+        const error = String(err);
+        log('[checkpoint] Failed to delete', { error, name, sessionID });
+        return { success: false, error };
+      }
+    },
+
+    async rollback(sessionID, name) {
+      try {
+        const checkpoint = await storage.getByName(sessionID, name);
+        if (!checkpoint) {
+          return { success: false, error: `Checkpoint "${name}" not found.` };
+        }
+
+        // Abort if session is busy
+        if (await isSessionBusy(sessionID)) {
+          try {
+            await client.session.abort({ path: { id: sessionID } });
+            log('[checkpoint] Aborted busy session before rollback', { sessionID });
+          } catch (err) {
+            log('[checkpoint] Failed to abort session', { error: String(err) });
+          }
+        }
+
+        // Find the first user message after the anchor
+        const targetMessageID = await getFirstUserMessageAfter(
+          sessionID,
+          checkpoint.anchorMessageID,
+        );
+
+        if (!targetMessageID) {
+          return {
+            success: false,
+            error: 'No messages found after checkpoint anchor. Nothing to rollback.',
+          };
+        }
+
+        // Call OpenCode's revert API
+        await client.session.revert({
+          path: { id: sessionID },
+          body: { messageID: targetMessageID },
+        });
+
+        log('[checkpoint] Rollback complete', { name, sessionID, targetMessageID });
+        return {
+          success: true,
+          message: `Rolled back to checkpoint "${name}". Context pruned to before message ${targetMessageID.slice(0, 8)}...`,
+        };
+      } catch (err) {
+        const error = String(err);
+        log('[checkpoint] Rollback failed', { error, name, sessionID });
+        return { success: false, error };
+      }
+    },
+
+    async cleanupSession(sessionID) {
+      const count = await storage.deleteForSession(sessionID);
+      if (count > 0) {
+        log('[checkpoint] Cleaned up session', { sessionID, count });
+      }
+      return count;
+    },
+  };
+}

+ 154 - 0
src/checkpoints/store.ts

@@ -0,0 +1,154 @@
+/**
+ * Checkpoint storage - persists checkpoints to disk.
+ */
+
+import * as fs from 'node:fs/promises';
+import * as path from 'node:path';
+import type { Checkpoint, CheckpointStore } from './types';
+
+const STORE_FILENAME = '.slim/checkpoints.json';
+
+export class CheckpointStorage {
+  private storePath: string;
+  private cache: CheckpointStore | null = null;
+
+  constructor(directory: string) {
+    this.storePath = path.join(directory, STORE_FILENAME);
+  }
+
+  /**
+   * Ensure the .slim directory exists.
+   */
+  private async ensureDir(): Promise<void> {
+    const dir = path.dirname(this.storePath);
+    try {
+      await fs.mkdir(dir, { recursive: true });
+    } catch {
+      // Directory may already exist
+    }
+  }
+
+  /**
+   * Load the checkpoint store from disk.
+   */
+  async load(): Promise<CheckpointStore> {
+    if (this.cache) return this.cache;
+
+    try {
+      const data = await fs.readFile(this.storePath, 'utf8');
+      const parsed = JSON.parse(data) as CheckpointStore;
+      this.cache = {
+        checkpoints: parsed.checkpoints || {},
+        bySession: parsed.bySession || {},
+      };
+      return this.cache;
+    } catch {
+      // File doesn't exist or is corrupted - start fresh
+      this.cache = { checkpoints: {}, bySession: {} };
+      return this.cache;
+    }
+  }
+
+  /**
+   * Save the checkpoint store to disk.
+   */
+  private async save(store: CheckpointStore): Promise<void> {
+    await this.ensureDir();
+    await fs.writeFile(this.storePath, JSON.stringify(store, null, 2), 'utf8');
+    this.cache = store;
+  }
+
+  /**
+   * Add a new checkpoint.
+   */
+  async add(checkpoint: Checkpoint): Promise<void> {
+    const store = await this.load();
+    store.checkpoints[checkpoint.id] = checkpoint;
+
+    if (!store.bySession[checkpoint.sessionID]) {
+      store.bySession[checkpoint.sessionID] = [];
+    }
+    if (!store.bySession[checkpoint.sessionID].includes(checkpoint.id)) {
+      store.bySession[checkpoint.sessionID].push(checkpoint.id);
+    }
+
+    await this.save(store);
+  }
+
+  /**
+   * Get a checkpoint by ID.
+   */
+  async get(id: string): Promise<Checkpoint | null> {
+    const store = await this.load();
+    return store.checkpoints[id] || null;
+  }
+
+  /**
+   * Get a checkpoint by name (for a specific session).
+   */
+  async getByName(sessionID: string, name: string): Promise<Checkpoint | null> {
+    const store = await this.load();
+    const ids = store.bySession[sessionID] || [];
+    for (const id of ids) {
+      const cp = store.checkpoints[id];
+      if (cp?.name === name) return cp;
+    }
+    return null;
+  }
+
+  /**
+   * List all checkpoints for a session.
+   */
+  async listForSession(sessionID: string): Promise<Checkpoint[]> {
+    const store = await this.load();
+    const ids = store.bySession[sessionID] || [];
+    return ids
+      .map((id) => store.checkpoints[id])
+      .filter((cp): cp is Checkpoint => !!cp)
+      .sort((a, b) => b.createdAt - a.createdAt);
+  }
+
+  /**
+   * Delete a checkpoint.
+   */
+  async delete(id: string): Promise<boolean> {
+    const store = await this.load();
+    const cp = store.checkpoints[id];
+    if (!cp) return false;
+
+    delete store.checkpoints[id];
+
+    const sessionCheckpoints = store.bySession[cp.sessionID] || [];
+    store.bySession[cp.sessionID] = sessionCheckpoints.filter((cid) => cid !== id);
+
+    await this.save(store);
+    return true;
+  }
+
+  /**
+   * Delete all checkpoints for a session (cleanup).
+   */
+  async deleteForSession(sessionID: string): Promise<number> {
+    const store = await this.load();
+    const ids = store.bySession[sessionID] || [];
+    let count = 0;
+
+    for (const id of ids) {
+      if (store.checkpoints[id]) {
+        delete store.checkpoints[id];
+        count++;
+      }
+    }
+
+    delete store.bySession[sessionID];
+    await this.save(store);
+    return count;
+  }
+
+  /**
+   * Clear the in-memory cache.
+   */
+  clearCache(): void {
+    this.cache = null;
+  }
+}

+ 94 - 0
src/checkpoints/tool.ts

@@ -0,0 +1,94 @@
+/**
+ * Checkpoint tools - exposed to agents for self-checkpointing.
+ */
+
+import type { PluginInput } from '@opencode-ai/plugin';
+import type { CheckpointManager } from './manager';
+
+export interface CheckpointTools {
+  checkpoint: (args: {
+    name: string;
+    description?: string;
+  }) => Promise<{ output: string; metadata?: Record<string, unknown> }>;
+  rollback: (args: {
+    name: string;
+  }) => Promise<{ output: string; metadata?: Record<string, unknown> }>;
+}
+
+export function createCheckpointTools(
+  ctx: PluginInput,
+  manager: CheckpointManager,
+  getSessionID: () => string | undefined,
+): CheckpointTools {
+  return {
+    async checkpoint(args) {
+      const sessionID = getSessionID();
+      if (!sessionID) {
+        return {
+          output: 'Error: No active session to checkpoint.',
+          metadata: { error: 'no_session' },
+        };
+      }
+
+      if (!args.name || args.name.trim() === '') {
+        return {
+          output: 'Error: Checkpoint name is required.',
+          metadata: { error: 'missing_name' },
+        };
+      }
+
+      const result = await manager.saveCheckpoint(
+        sessionID,
+        args.name.trim(),
+        args.description,
+      );
+
+      if (result.success && result.checkpoint) {
+        return {
+          output: `Checkpoint "${args.name}" saved. You can rollback to this state later if needed.`,
+          metadata: {
+            checkpointId: result.checkpoint.id,
+            anchorMessageID: result.checkpoint.anchorMessageID,
+            createdAt: result.checkpoint.createdAt,
+          },
+        };
+      }
+
+      return {
+        output: `Failed to save checkpoint: ${result.error}`,
+        metadata: { error: result.error },
+      };
+    },
+
+    async rollback(args) {
+      const sessionID = getSessionID();
+      if (!sessionID) {
+        return {
+          output: 'Error: No active session to rollback.',
+          metadata: { error: 'no_session' },
+        };
+      }
+
+      if (!args.name || args.name.trim() === '') {
+        return {
+          output: 'Error: Checkpoint name is required.',
+          metadata: { error: 'missing_name' },
+        };
+      }
+
+      const result = await manager.rollback(sessionID, args.name.trim());
+
+      if (result.success) {
+        return {
+          output: result.message || `Rolled back to checkpoint "${args.name}".`,
+          metadata: { success: true },
+        };
+      }
+
+      return {
+        output: `Rollback failed: ${result.error}`,
+        metadata: { error: result.error },
+      };
+    },
+  };
+}

+ 27 - 0
src/checkpoints/types.ts

@@ -0,0 +1,27 @@
+/**
+ * Checkpoint types for session rollback functionality.
+ */
+
+export interface Checkpoint {
+  /** Unique identifier for the checkpoint */
+  id: string;
+  /** User-defined name for the checkpoint */
+  name: string;
+  /** Optional description of why checkpoint was created */
+  description?: string;
+  /** Session ID this checkpoint belongs to */
+  sessionID: string;
+  /** Message ID of the last user message at checkpoint time (the anchor) */
+  anchorMessageID: string;
+  /** Workspace directory */
+  directory: string;
+  /** Timestamp when checkpoint was created */
+  createdAt: number;
+}
+
+export interface CheckpointStore {
+  /** Map of checkpoint ID to checkpoint */
+  checkpoints: Record<string, Checkpoint>;
+  /** Index: sessionID -> checkpoint IDs */
+  bySession: Record<string, string[]>;
+}

+ 111 - 16
src/index.ts

@@ -19,6 +19,11 @@ import {
 import { processImageAttachments } from './hooks/image-hook';
 import { createInterviewManager } from './interview';
 import { createBuiltinMcps } from './mcp';
+import {
+  createCheckpointCommands,
+  createCheckpointManager,
+  createCheckpointTools,
+} from './checkpoints';
 import {
   getMultiplexer,
   MultiplexerSessionManager,
@@ -30,6 +35,7 @@ import {
   createCouncilTool,
   createWebfetchTool,
 } from './tools';
+import type { CheckpointManager } from './checkpoints';
 import { resolveRuntimeAgentName, rewriteDisplayNameMentions } from './utils';
 import { initLogger, log } from './utils/logger';
 import { SubagentDepthTracker } from './utils/subagent-depth';
@@ -250,9 +256,18 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
     });
     interviewManager = createInterviewManager(ctx, config);
 
+    // Initialize checkpoint manager for session rollback
+    const checkpointManager = createCheckpointManager(ctx);
+    const checkpointCommands = createCheckpointCommands(checkpointManager);
+
+    // Get current session ID for tools
+    let currentSessionID: string | undefined;
+    const checkpointTools = createCheckpointTools(ctx, checkpointManager, () => currentSessionID);
+
     toolCount =
       Object.keys(councilTools).length +
       Object.keys(todoContinuationHook.tool).length +
+      Object.keys(checkpointTools).length +
       1 + // webfetch
       2; // ast_grep_search, ast_grep_replace
   } catch (err) {
@@ -320,6 +335,7 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
       ...councilTools,
       webfetch,
       ...todoContinuationHook.tool,
+      ...checkpointTools,
       ast_grep_search,
       ast_grep_replace,
     },
@@ -520,6 +536,30 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
       }
 
       interviewManager.registerCommand(opencodeConfig);
+
+      // Register checkpoint commands
+      const configCommand = opencodeConfig.command as
+        | Record<string, unknown>
+        | undefined;
+      if (!configCommand?.['checkpoint']) {
+        if (!opencodeConfig.command) {
+          opencodeConfig.command = {};
+        }
+        (opencodeConfig.command as Record<string, unknown>)['checkpoint'] = {
+          template: 'checkpoint <subcommand> [args]',
+          description:
+            'Save and restore session checkpoints — /checkpoint save <name>, /checkpoint list, /checkpoint delete <name>',
+        };
+      }
+      if (!configCommand?.['rollback']) {
+        if (!opencodeConfig.command) {
+          opencodeConfig.command = {};
+        }
+        (opencodeConfig.command as Record<string, unknown>)['rollback'] = {
+          template: 'rollback <name>',
+          description: 'Rollback to a saved checkpoint',
+        };
+      }
     },
 
     event: async (input) => {
@@ -588,6 +628,19 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
         if (sessionID) {
           sessionAgentMap.delete(sessionID);
         }
+
+        // Cleanup checkpoints for deleted session
+        if (sessionID && checkpointManager) {
+          await checkpointManager.cleanupSession(sessionID);
+        }
+      }
+
+      // Track current session ID for checkpoint tools
+      if (event.type === 'session.status' || event.type === 'session.created') {
+        const sid = event.properties?.sessionID ?? event.properties?.info?.id;
+        if (sid) {
+          currentSessionID = sid;
+        }
       }
     },
 
@@ -608,23 +661,65 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
     // Direct interception of /auto-continue command — bypasses LLM
     // round-trip
     'command.execute.before': async (input, output) => {
-      await todoContinuationHook.handleCommandExecuteBefore(
-        input as {
-          command: string;
-          sessionID: string;
-          arguments: string;
-        },
-        output as { parts: Array<{ type: string; text?: string }> },
-      );
+      const cmdInput = input as {
+        command: string;
+        sessionID: string;
+        arguments: string;
+      };
+      const cmdOutput = output as { parts: Array<{ type: string; text?: string }> };
+
+      await todoContinuationHook.handleCommandExecuteBefore(cmdInput, cmdOutput);
+
+      await interviewManager.handleCommandExecuteBefore(cmdInput, cmdOutput);
+
+      // Handle checkpoint commands
+      if (cmdInput.command === 'checkpoint') {
+        const args = cmdInput.arguments.trim();
+        const subcommandMatch = args.match(/^(\S+)(?:\s+(.*))?$/);
+        if (subcommandMatch) {
+          const [, subcommand, rest] = subcommandMatch;
+          switch (subcommand) {
+            case 'save': {
+              const result = await checkpointCommands.handleCheckpointSave({
+                sessionID: cmdInput.sessionID,
+                arguments: rest || '',
+              });
+              cmdOutput.parts.push(...result.parts);
+              break;
+            }
+            case 'list': {
+              const result = await checkpointCommands.handleCheckpointList({
+                sessionID: cmdInput.sessionID,
+                arguments: '',
+              });
+              cmdOutput.parts.push(...result.parts);
+              break;
+            }
+            case 'delete': {
+              const result = await checkpointCommands.handleCheckpointDelete({
+                sessionID: cmdInput.sessionID,
+                arguments: rest || '',
+              });
+              cmdOutput.parts.push(...result.parts);
+              break;
+            }
+            default: {
+              cmdOutput.parts.push({
+                type: 'text',
+                text: 'Unknown checkpoint subcommand. Use: save, list, or delete',
+              });
+            }
+          }
+        }
+      }
 
-      await interviewManager.handleCommandExecuteBefore(
-        input as {
-          command: string;
-          sessionID: string;
-          arguments: string;
-        },
-        output as { parts: Array<{ type: string; text?: string }> },
-      );
+      if (cmdInput.command === 'rollback') {
+        const result = await checkpointCommands.handleRollback({
+          sessionID: cmdInput.sessionID,
+          arguments: cmdInput.arguments,
+        });
+        cmdOutput.parts.push(...result.parts);
+      }
     },
 
     'chat.headers': chatHeadersHook['chat.headers'],