Pi extension: /review command for branching and returning with summary

review-extension-corrected.ts

import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";

export default function (pi: ExtensionAPI) {
  // Store review state
  let reviewState: {
    originalSessionFile: string;
    reviewSessionFile: string;
    branchFromId: string;
  } | undefined;

  pi.registerCommand("review", {
    description: "Start a review in a forked session",
    handler: async (args, ctx) => {
      // Store current state
      const originalSessionFile = ctx.sessionManager.getSessionFile();
      const branchFromId = ctx.sessionManager.getLeafId() || "root";
      
      // Fork to new session file
      const result = await ctx.fork(branchFromId);
      if (result.cancelled) return;
      
      // Store state for return
      reviewState = {
        originalSessionFile,
        reviewSessionFile: ctx.sessionManager.getSessionFile(),
        branchFromId,
      };
      
      // Mark this as a review session with custom entry
      pi.appendEntry("review_start", {
        originalSession: originalSessionFile,
        branchFromId,
        startedAt: new Date().toISOString(),
      });
      
      ctx.ui.notify("Review session started. Use /end-review when done.", "info");
    },
  });

  pi.registerCommand("end-review", {
    description: "Complete review and return to original session",
    handler: async (args, ctx) => {
      if (!reviewState) {
        ctx.ui.notify("Not in a review session", "error");
        return;
      }

      const summary = args || "Review completed";
      
      // Append summary to review session
      pi.appendEntry("review_summary", {
        summary,
        completedAt: new Date().toISOString(),
      });
      
      // Switch back to original session
      // Note: There's no direct API to switch sessions, so you'd need to:
      // 1. Tell user to manually switch back, OR
      // 2. Use ctx.shutdown() and restart with original session
      
      ctx.ui.notify(
        `Review complete: ${summary}\n` +
        `Review session: ${reviewState.reviewSessionFile}\n` +
        `Manually switch back to: ${reviewState.originalSessionFile}`,
        "info"
      );
      
      reviewState = undefined;
    },
  });
}

review-extension-final.ts

import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";

export default function (pi: ExtensionAPI) {
  // Store where we branched from
  let reviewOriginId: string | undefined;

  pi.registerCommand("review", {
    description: "Start a review branch from root, return when done",
    handler: async (args, ctx) => {
      // Store current position
      reviewOriginId = ctx.sessionManager.getLeafId() || "root";
      
      // Get the root entry (first user message or session start)
      const entries = ctx.sessionManager.getEntries();
      const rootEntry = entries.find(e => e.parentId === null);
      
      if (!rootEntry) {
        ctx.ui.notify("No root entry found", "error");
        return;
      }
      
      // Navigate to root to create a branch from there
      // This changes the leaf pointer to root
      const result = await ctx.navigateTree(rootEntry.id, { summarize: false });
      if (result.cancelled) return;
      
      // Add a marker entry so we know we're in review mode
      pi.appendEntry("review_start", {
        originId: reviewOriginId,
        startedAt: new Date().toISOString(),
      });
      
      ctx.ui.notify(
        "Review branch started from root. Use /end-review when done.",
        "info"
      );
    },
  });

  pi.registerCommand("end-review", {
    description: "Complete review and return to original position with summary",
    handler: async (args, ctx) => {
      if (!reviewOriginId) {
        ctx.ui.notify("Not in a review branch", "error");
        return;
      }

      const summary = args || "Review completed";
      
      // Get review path to count entries
      const reviewBranch = ctx.sessionManager.getBranch();
      
      // Navigate back to origin with summary
      const result = await ctx.navigateTree(reviewOriginId, { 
        summarize: true, // This will auto-summarize the review branch
      });
      
      if (result.cancelled) {
        return;
      }
      
      // Add custom entry with review metadata
      pi.appendEntry("review_complete", {
        summary,
        reviewEntryCount: reviewBranch.length,
        completedAt: new Date().toISOString(),
      });
      
      ctx.ui.notify(`Review complete: ${summary}`, "info");
      reviewOriginId = undefined;
    },
  });
}

review-extension.ts

import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";

export default function (pi: ExtensionAPI) {
  // Store where we branched from
  let reviewOriginId: string | undefined;

  pi.registerCommand("review", {
    description: "Start a review branch, return when done",
    handler: async (args, ctx) => {
      // Get current leaf (where we're branching from)
      const currentLeaf = ctx.sessionManager.getLeaf();
      
      // Branch from root (null parent)
      ctx.sessionManager.resetLeaf(); // Sets leaf to null
      
      // Store origin in a custom entry
      reviewOriginId = currentLeaf ?? "root";
      const entry = {
        type: "custom",
        id: generateId(),
        parentId: null, // Root of new branch
        timestamp: new Date().toISOString(),
        customType: "review_start",
        data: {
          originId: reviewOriginId,
          startedAt: new Date().toISOString(),
        },
      };
      
      ctx.sessionManager.appendEntry(entry);
      ctx.ui.notify("Review branch started. Use /end-review when done.", "info");
    },
  });

  pi.registerCommand("end-review", {
    description: "Complete review and return to original branch",
    handler: async (args, ctx) => {
      if (!reviewOriginId) {
        ctx.ui.notify("Not in a review branch", "error");
        return;
      }

      // Get review conversation
      const reviewBranch = ctx.sessionManager.getBranch();
      
      // Create summary (you could ask LLM here)
      const summary = args || "Review completed";
      
      // Navigate back to origin
      if (reviewOriginId !== "root") {
        const result = await ctx.navigateTree(reviewOriginId);
        if (result.cancelled) return;
      } else {
        ctx.sessionManager.resetLeaf();
      }
      
      // Append review summary at origin point
      const summaryEntry = {
        type: "custom",
        id: generateId(),
        parentId: reviewOriginId === "root" ? null : reviewOriginId,
        timestamp: new Date().toISOString(),
        customType: "review_summary",
        data: {
          summary,
          reviewEntryCount: reviewBranch.length,
        },
      };
      
      ctx.sessionManager.appendEntry(summaryEntry);
      ctx.ui.notify(`Review complete: ${summary}`, "info");
      reviewOriginId = undefined;
    },
  });
}

function generateId(): string {
  return Math.random().toString(36).substring(2, 15);
}