pi-yank 0.1.0 → 0.1.2

package/README.md

@@ -3,11 +3,13 @@
 Lightweight `/yank` extension for [pi](https://github.com/badlogic/pi-mono).
 
 `/yank` copies the last assistant response to the clipboard, or lets you choose an individual fenced code block from that response.
+You can also target earlier messages with `/yank N` (e.g. `/yank 2` operates on the 2nd-to-last assistant message).
 
 ## Features
 
 - Copy the last assistant message to the clipboard
-- Detect fenced Markdown code blocks in the last assistant message
+- Use `/yank N` to target the N-th latest assistant message (`/yank 2` = 2nd-to-last)
+- Detect fenced Markdown code blocks in the selected assistant message
 - Use arrow keys + Enter to choose a specific code block
 - Optionally disable the menu for the current session so `/yank` behaves like `/copy`
 - Reuse pi's official clipboard pipeline via `copyToClipboard()` for `/copy`-compatible behavior
@@ -53,17 +55,19 @@ cp /path/to/pi-yank/extensions/yank.ts ~/.pi/agent/extensions/yank.ts
 
 ## Usage
 
-Run:
-
 ```text
-/yank
+/yank       # copy from the latest assistant message
+/yank 1     # same as /yank
+/yank 2     # copy from the 2nd-to-last assistant message
+/yank N     # copy from the N-th latest assistant message
 ```
 
 Behavior:
 
 - No assistant message yet → shows `No agent messages to copy yet.`
-- Last assistant message has no fenced code blocks → copies the full message
-- Last assistant message has code blocks → opens a selection menu
+- N-th latest message not found → shows `No N-th latest agent message to copy yet.`
+- Selected assistant message has no fenced code blocks → copies the full message
+- Selected assistant message has code blocks → opens a selection menu
 - Selecting a code block copies only that code block's content
 - Selecting the session option disables the menu for later `/yank` calls in the current session
 - Pressing `Escape` in the selector cancels without copying
@@ -78,12 +82,14 @@ This makes it more robust for LLM output such as:
 - `~~~` fenced blocks
 - unterminated fences, which are ignored
 
-Menu previews now include more context when available, for example:
+Menu previews include more context when available, for example:
 
-- `1. Full last message (24 lines • 812 chars)`
+- `1. Full latest message (24 lines • 812 chars)`
 - `2. Code block — markdown: # Title (18 lines • 640 chars)`
 - `3. Code block — js: console.log("hi") (1 line • 18 chars)`
-- `4. Always copy full last message in this session`
+- `4. Always copy full selected message in this session`
+
+When using `/yank N` with N > 1, the menu title shows the target (e.g. "Yank from 2nd latest response").
 
 ## Development
 

package/extensions/yank-core.js

@@ -14,6 +14,23 @@
  * @returns {CodeBlock[]}
  */
 export function extractCodeBlocks(markdown) {
+	const topLevelBlocks = extractTopLevelCodeBlocks(markdown);
+	if (topLevelBlocks.length === 1) {
+		const [wrapper] = topLevelBlocks;
+		const innerBlocks = extractTopLevelCodeBlocks(wrapper.content);
+		if (innerBlocks.length > 1) {
+			return withPreviews(innerBlocks);
+		}
+	}
+
+	return withPreviews(topLevelBlocks);
+}
+
+/**
+ * @param {string} markdown
+ * @returns {CodeBlock[]}
+ */
+function extractTopLevelCodeBlocks(markdown) {
 	/** @type {CodeBlock[]} */
 	const blocks = [];
 	/** @type {OpenCodeBlock[]} */
@@ -59,9 +76,8 @@ export function extractCodeBlocks(markdown) {
 
 			const finishedBlock = openBlocks.pop();
 			if (finishedBlock && openBlocks.length === 0) {
-				const content = finishedBlock.lines.join("\n");
 				blocks.push({
-					content,
+					content: finishedBlock.lines.join("\n"),
 					preview: "",
 					startLine: finishedBlock.startLine,
 					infoString: finishedBlock.infoString,
@@ -82,16 +98,18 @@ export function extractCodeBlocks(markdown) {
 		});
 	}
 
-	return blocks
-		.sort((left, right) => left.startLine - right.startLine)
-		.map((block, index) => ({
-			...block,
-			preview: buildCodeBlockPreview(
-				block.content,
-				index + 1,
-				block.infoString,
-			),
-		}));
+	return blocks.sort((left, right) => left.startLine - right.startLine);
+}
+
+/**
+ * @param {CodeBlock[]} blocks
+ * @returns {CodeBlock[]}
+ */
+function withPreviews(blocks) {
+	return blocks.map((block, index) => ({
+		...block,
+		preview: buildCodeBlockPreview(block.content, index + 1, block.infoString),
+	}));
 }
 
 /**

package/extensions/yank-helpers.js

@@ -0,0 +1,78 @@
+/**
+ * @typedef {{ type?: unknown; message?: { role?: unknown; stopReason?: unknown; content?: unknown } }} SessionEntryLike
+ */
+
+/**
+ * @param {string} args
+ * @returns {{ ok: true, nth: number } | { ok: false, message: string }}
+ */
+export function parseYankArgs(args) {
+	const trimmed = args.trim();
+	if (trimmed.length === 0) {
+		return { ok: true, nth: 1 };
+	}
+
+	if (!/^\d+$/.test(trimmed)) {
+		return {
+			ok: false,
+			message: 'Usage: /yank [N] where N is a positive integer',
+		};
+	}
+
+	const nth = Number.parseInt(trimmed, 10);
+	if (!Number.isSafeInteger(nth) || nth < 1) {
+		return {
+			ok: false,
+			message: 'Usage: /yank [N] where N is a positive integer',
+		};
+	}
+
+	return { ok: true, nth };
+}
+
+/**
+ * @param {SessionEntryLike['message']} message
+ * @returns {string | undefined}
+ */
+export function extractAssistantText(message) {
+	const content = Array.isArray(message?.content) ? message.content : [];
+	if (message?.stopReason === 'aborted' && content.length === 0) {
+		return undefined;
+	}
+
+	let text = '';
+	for (const item of content) {
+		if (!item || typeof item !== 'object') continue;
+		const part = /** @type {{ type?: unknown; text?: unknown }} */ (item);
+		if (part.type === 'text' && typeof part.text === 'string') {
+			text += part.text;
+		}
+	}
+
+	const trimmed = text.trim();
+	return trimmed || undefined;
+}
+
+/**
+ * @param {SessionEntryLike[]} entries
+ * @param {number} nth
+ * @returns {string | undefined}
+ */
+export function getNthLatestAssistantText(entries, nth = 1) {
+	let seen = 0;
+	for (let i = entries.length - 1; i >= 0; i--) {
+		const entry = entries[i];
+		if (entry?.type !== 'message') continue;
+
+		const message = entry.message;
+		if (message?.role !== 'assistant') continue;
+
+		const text = extractAssistantText(message);
+		if (!text) continue;
+
+		seen += 1;
+		if (seen === nth) return text;
+	}
+
+	return undefined;
+}

package/extensions/yank.ts

@@ -2,57 +2,15 @@ import {
 	copyToClipboard,
 	type ExtensionAPI,
 	type ExtensionCommandContext,
-	type SessionEntry,
 } from "@mariozechner/pi-coding-agent";
 import { extractCodeBlocks } from "./yank-core.js";
+import { getNthLatestAssistantText, parseYankArgs } from "./yank-helpers.js";
 
-interface AssistantMessageLike {
-	role?: unknown;
-	stopReason?: unknown;
-	content?: unknown;
-}
-
-const COPY_SUCCESS_MESSAGE = "Copied last agent message to clipboard";
+const COPY_SUCCESS_MESSAGE = "Copied agent message to clipboard";
 const COPY_EMPTY_MESSAGE = "No agent messages to copy yet.";
 const COPY_CODE_BLOCK_SUCCESS_MESSAGE = "Copied code block to clipboard";
 const MENU_DISABLED_MESSAGE =
-	"Future /yank commands will copy the full last message in this session";
-
-function getLastAssistantText(entries: SessionEntry[]): string | undefined {
-	for (let i = entries.length - 1; i >= 0; i--) {
-		const entry = entries[i];
-		if (entry.type !== "message") continue;
-
-		const message = entry.message as AssistantMessageLike;
-		if (message.role !== "assistant") continue;
-
-		const text = extractAssistantText(message);
-		if (text) return text;
-	}
-
-	return undefined;
-}
-
-function extractAssistantText(
-	message: AssistantMessageLike,
-): string | undefined {
-	const content = Array.isArray(message.content) ? message.content : [];
-	if (message.stopReason === "aborted" && content.length === 0) {
-		return undefined;
-	}
-
-	let text = "";
-	for (const item of content) {
-		if (!item || typeof item !== "object") continue;
-		const part = item as { type?: unknown; text?: unknown };
-		if (part.type === "text" && typeof part.text === "string") {
-			text += part.text;
-		}
-	}
-
-	const trimmed = text.trim();
-	return trimmed || undefined;
-}
+	"Future /yank commands will copy the full selected message in this session";
 
 function countNonEmptyLines(text: string): number {
 	return text.split(/\r?\n/).filter((line) => line.trim().length > 0).length;
@@ -96,11 +54,26 @@ export default function yankExtension(pi: ExtensionAPI) {
 	let skipMenuForSession = false;
 
 	pi.registerCommand("yank", {
-		description: "Copy the last assistant message or a code block from it",
-		handler: async (_args, ctx) => {
-			const text = getLastAssistantText(ctx.sessionManager.getBranch());
+		description:
+			"Copy the latest assistant message, or use /yank N for the N-th latest one",
+		handler: async (args, ctx) => {
+			const parsed = parseYankArgs(args);
+			if (!parsed.ok) {
+				ctx.ui.notify(parsed.message, "error");
+				return;
+			}
+
+			const text = getNthLatestAssistantText(
+				ctx.sessionManager.getBranch(),
+				parsed.nth,
+			);
 			if (!text) {
-				ctx.ui.notify(COPY_EMPTY_MESSAGE, "error");
+				ctx.ui.notify(
+					parsed.nth === 1
+						? COPY_EMPTY_MESSAGE
+						: `No ${parsed.nth}-th latest agent message to copy yet.`,
+					"error",
+				);
 				return;
 			}
 
@@ -118,7 +91,12 @@ export default function yankExtension(pi: ExtensionAPI) {
 			}
 
 			const menu = buildMenuOptions(text, codeBlocks);
-			const selected = await ctx.ui.select("Yank from last response", menu.all);
+			const selected = await ctx.ui.select(
+				parsed.nth === 1
+					? "Yank from latest response"
+					: `Yank from ${parsed.nth}-th latest response`,
+				menu.all,
+			);
 			if (!selected) return;
 
 			try {

package/package.json

@@ -1,7 +1,7 @@
 {
 	"name": "pi-yank",
 	"private": false,
-	"version": "0.1.0",
+	"version": "0.1.2",
 	"type": "module",
 	"description": "Lightweight /yank extension for pi that copies the last assistant message or a selected code block.",
 	"keywords": [

package/tests/yank-command.test.mjs

@@ -0,0 +1,48 @@
+import assert from "node:assert/strict";
+import test from "node:test";
+import {
+	extractAssistantText,
+	getNthLatestAssistantText,
+	parseYankArgs,
+} from "../extensions/yank-helpers.js";
+
+test("parseYankArgs defaults to latest message", () => {
+	assert.deepEqual(parseYankArgs(""), { ok: true, nth: 1 });
+	assert.deepEqual(parseYankArgs("   "), { ok: true, nth: 1 });
+});
+
+test("parseYankArgs accepts a positive integer", () => {
+	assert.deepEqual(parseYankArgs("1"), { ok: true, nth: 1 });
+	assert.deepEqual(parseYankArgs("2"), { ok: true, nth: 2 });
+	assert.deepEqual(parseYankArgs("  12  "), { ok: true, nth: 12 });
+});
+
+test("parseYankArgs rejects invalid values", () => {
+	assert.equal(parseYankArgs("0").ok, false);
+	assert.equal(parseYankArgs("-1").ok, false);
+	assert.equal(parseYankArgs("1 2").ok, false);
+	assert.equal(parseYankArgs("abc").ok, false);
+});
+
+test("extractAssistantText ignores aborted empty assistant messages", () => {
+	assert.equal(
+		extractAssistantText({ role: "assistant", stopReason: "aborted", content: [] }),
+		undefined,
+	);
+});
+
+test("getNthLatestAssistantText returns the requested assistant message", () => {
+	const entries = [
+		{ type: "message", message: { role: "user", content: [{ type: "text", text: "u1" }] } },
+		{ type: "message", message: { role: "assistant", content: [{ type: "text", text: "a1" }] } },
+		{ type: "message", message: { role: "assistant", stopReason: "aborted", content: [] } },
+		{ type: "message", message: { role: "assistant", content: [{ type: "text", text: "a2" }] } },
+		{ type: "message", message: { role: "user", content: [{ type: "text", text: "u2" }] } },
+		{ type: "message", message: { role: "assistant", content: [{ type: "text", text: "a3" }] } },
+	];
+
+	assert.equal(getNthLatestAssistantText(entries, 1), "a3");
+	assert.equal(getNthLatestAssistantText(entries, 2), "a2");
+	assert.equal(getNthLatestAssistantText(entries, 3), "a1");
+	assert.equal(getNthLatestAssistantText(entries, 4), undefined);
+});

package/tests/yank-core.test.mjs

@@ -5,7 +5,7 @@ import {
 	extractCodeBlocks,
 } from "../extensions/yank-core.js";
 
-test("extracts a single outer markdown block when nested fences appear as literal examples", () => {
+test("prefers inner sibling code blocks when a single outer markdown wrapper encloses them", () => {
 	const text = `我来为你创建一个 Markdown 格式的测试文本：
 
 \`\`\`markdown
@@ -32,10 +32,11 @@ def hello():
 说明文字`;
 
 	const blocks = extractCodeBlocks(text);
-	assert.equal(blocks.length, 1);
+	assert.equal(blocks.length, 2);
 	assert.match(blocks[0].content, /function hello\(\)/);
-	assert.match(blocks[0].content, /def hello\(\):/);
-	assert.equal(blocks[0].preview, "markdown: # Markdown 格式测试文档");
+	assert.match(blocks[1].content, /def hello\(\):/);
+	assert.equal(blocks[0].preview, "javascript: function hello() {");
+	assert.equal(blocks[1].preview, "python: def hello():");
 });
 
 test("extracts sibling top-level code blocks independently", () => {
@@ -59,6 +60,35 @@ after`;
 	assert.equal(blocks[1].preview, 'python: print("b")');
 });
 
+test("unwraps a single outer markdown wrapper around filename and script blocks", () => {
+	const text = `下面给你一套 pi / pi-mono 迁移的一键备份 +恢复脚本（macOS/Linux）。
+
+保存为：
+
+\`\`\`bash
+backup-pi.sh
+\`\`\`
+
+内容如下：
+
+\`\`\`bash
+#!/usr/bin/env bash
+set -euo pipefail
+TIMESTAMP="$(date +%Y%m%d-%H%M%S)"
+\`\`\``;
+
+	const wrapped = `\`\`\`markdown
+${text}
+\`\`\``;
+	const blocks = extractCodeBlocks(wrapped);
+
+	assert.equal(blocks.length, 2);
+	assert.equal(blocks[0].content, "backup-pi.sh");
+	assert.match(blocks[1].content, /set -euo pipefail/);
+	assert.equal(blocks[0].preview, "bash: backup-pi.sh");
+	assert.equal(blocks[1].preview, "bash: #!/usr/bin/env bash");
+});
+
 test("supports tilde fences", () => {
 	const text = `~~~sql
 select 1;