commit 44609be7c540db540a1ae32dfef5033dcc663a0f Author: wassname <1103714+wassname@users.noreply.github.com> Date: Wed Apr 8 15:03:25 2026 +0800 Initial: edit-last extension for pi - Extracts assistant text since last user message - Two commands: edit-last (built-in editor) and edit-last-external (opens $EDITOR/Zed and waits) - Quote-prefixes lines to preserve code blocks diff --git a/index.ts b/index.ts new file mode 100644 index 0000000..10e28f1 --- /dev/null +++ b/index.ts @@ -0,0 +1,169 @@ +import type { ExtensionAPI } from "@mariozechner/pi-coding-agent"; +import { Type } from "@sinclair/typebox"; +import { readFileSync, writeFileSync, unlinkSync } from "node:fs"; +import { mkdtempSync, rmdirSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { execSync } from "node:child_process"; + +type TextBlock = { type?: string; text?: string }; + +function extractAssistantText(content: unknown): string { + if (!Array.isArray(content)) return ""; + + return content + .flatMap((block) => { + if (!block || typeof block !== "object") return []; + const b = block as TextBlock; + return b.type === "text" && typeof b.text === "string" ? [b.text] : []; + }) + .join("\n") + .trim(); +} + +function quotePrefixLines(text: string): string { + return text + .split("\n") + .map((line) => `> ${line}`) + .join("\n"); +} + +function getAssistantTextSinceLastUser(ctx: any): string | null { + const branch = ctx.sessionManager.getBranch(); + + let lastUserIdx = -1; + for (let i = branch.length - 1; i >= 0; i--) { + const entry = branch[i]; + if ( + entry?.type === "message" && + entry.message && + entry.message.role === "user" + ) { + lastUserIdx = i; + break; + } + } + + const chunks: string[] = []; + + for (let i = Math.max(0, lastUserIdx + 1); i < branch.length; i++) { + const entry = branch[i]; + if (entry?.type !== "message") continue; + + const msg = entry.message; + if (!msg || msg.role !== "assistant") continue; + + const text = extractAssistantText(msg.content); + if (text) chunks.push(text); + } + + if (chunks.length === 0) return null; + + return chunks + .map((chunk) => quotePrefixLines(chunk)) + .join("\n>\n"); +} + +function buildRevisionPrompt(editedText: string): string { + return [ + "Please revise your reply to my last message using the edited/annotated assistant draft below.", + "Treat my edits as the target direction.", + "If I left inline notes, incorporate them rather than repeating them literally.", + "", + "--- edited assistant draft ---", + editedText.trim(), + ].join("\n"); +} + +// Open editor and wait for it to close, returning the edited content +function editInExternalEditor(content: string): string | null { + const editor = process.env.VISUAL || process.env.EDITOR || "vi"; + + const tmpDir = mkdtempSync(join(tmpdir(), "pi-edit-")); + const tmpFile = join(tmpDir, "edit.txt"); + + writeFileSync(tmpFile, content, "utf-8"); + + try { + execSync(`${editor} "${tmpFile}"`, { + stdio: "inherit", + env: { ...process.env, VISUAL: editor, EDITOR: editor }, + }); + + const result = readFileSync(tmpFile, "utf-8"); + return result; + } catch (e) { + return null; + } finally { + unlinkSync(tmpFile); + rmdirSync(tmpDir); + } +} + +export default function (pi: ExtensionAPI) { + pi.registerCommand("edit-last-external", { + description: + "Edit assistant text since last user message in external editor (VISUAL/EDITOR), then send for revision.", + handler: async (_args, ctx) => { + if (!ctx.hasUI) return; + + await ctx.waitForIdle(); + + const draft = getAssistantTextSinceLastUser(ctx); + if (!draft) { + ctx.ui.notify( + "No assistant text found since the last user message.", + "warning" + ); + return; + } + + const prompt = buildRevisionPrompt(draft); + + // Open external editor and wait for it to close + const edited = editInExternalEditor(prompt); + + if (edited == null || !edited.trim()) { + ctx.ui.notify("No changes or cancelled.", "info"); + return; + } + + // Send the edited content as a user message + pi.sendUserMessage(edited.trim()); + }, + }); + + // Also keep the UI-based version that pre-fills editor + pi.registerCommand("edit-last", { + description: + "Edit all assistant text since the last user message in the built-in editor.", + handler: async (_args, ctx) => { + if (!ctx.hasUI) return; + + await ctx.waitForIdle(); + + const draft = getAssistantTextSinceLastUser(ctx); + if (!draft) { + ctx.ui.notify( + "No assistant text found since the last user message.", + "warning" + ); + return; + } + + const edited = await ctx.ui.editor( + "Edit assistant text since last user message", + draft + ); + if (edited == null) return; + + const cleaned = edited.trim(); + if (!cleaned) { + ctx.ui.notify("Nothing to send.", "warning"); + return; + } + + pi.sendUserMessage(buildRevisionPrompt(cleaned)); + }, + }); +} \ No newline at end of file diff --git a/package.json b/package.json new file mode 100644 index 0000000..222a535 --- /dev/null +++ b/package.json @@ -0,0 +1,8 @@ +{ + "name": "pi-edit-last", + "version": "1.0.0", + "description": "Edit assistant replies since last user message, with external editor support", + "pi": { + "extensions": ["./index.ts"] + } +}