Joakim Persson
77 lines
2.3 KiB
TypeScript
77 lines
2.3 KiB
TypeScript
/**
|
|
* Git Checkpoint Extension
|
|
*
|
|
* Creates a git stash checkpoint at the start of each turn, keyed to the
|
|
* session entry ID. If you /fork from a past entry, you're offered the option
|
|
* to restore the code state from that point.
|
|
*
|
|
* Silently skips turns where the working directory isn't inside a git repo or
|
|
* where there are no changes to stash.
|
|
*
|
|
* Status bar shows the number of checkpoints saved in the current session.
|
|
*/
|
|
|
|
import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
|
|
|
|
export default function (pi: ExtensionAPI) {
|
|
// entryId → stash ref (e.g. "refs/stash" or a full sha)
|
|
const checkpoints = new Map<string, string>();
|
|
|
|
async function inGitRepo(): Promise<boolean> {
|
|
try {
|
|
const result = await pi.exec("git", ["rev-parse", "--git-dir"]);
|
|
return result.exitCode === 0;
|
|
} catch {
|
|
return false;
|
|
}
|
|
}
|
|
|
|
function updateStatus(ctx: { ui: { setStatus: (id: string, text: string) => void } }): void {
|
|
const n = checkpoints.size;
|
|
if (n > 0) {
|
|
ctx.ui.setStatus("git-checkpoint", `⎇ ${n} checkpoint${n === 1 ? "" : "s"}`);
|
|
} else {
|
|
ctx.ui.setStatus("git-checkpoint", "");
|
|
}
|
|
}
|
|
|
|
pi.on("turn_start", async (_event, ctx) => {
|
|
if (!(await inGitRepo())) return;
|
|
|
|
// Capture the current entry ID at turn start — this is the entry the
|
|
// user will fork from if they want to restore code to this point.
|
|
const entryId = ctx.sessionManager.getLeafId();
|
|
if (!entryId) return;
|
|
|
|
// git stash create makes a stash object without touching the working tree.
|
|
// Returns the stash ref on stdout, or empty string if nothing to stash.
|
|
const result = await pi.exec("git", ["stash", "create"]);
|
|
const ref = result.stdout.trim();
|
|
|
|
if (ref) {
|
|
checkpoints.set(entryId, ref);
|
|
updateStatus(ctx);
|
|
}
|
|
});
|
|
|
|
pi.on("session_before_fork", async (event, ctx) => {
|
|
const ref = checkpoints.get(event.entryId);
|
|
if (!ref) return;
|
|
if (!ctx.hasUI) return;
|
|
|
|
const choice = await ctx.ui.select(
|
|
"Restore code to this checkpoint?",
|
|
["Yes, restore code state", "No, keep current code"],
|
|
);
|
|
|
|
if (choice?.startsWith("Yes")) {
|
|
await pi.exec("git", ["stash", "apply", ref]);
|
|
ctx.ui.notify("Code restored to checkpoint", "success");
|
|
}
|
|
});
|
|
|
|
pi.on("session_shutdown", () => {
|
|
checkpoints.clear();
|
|
});
|
|
}
|