pi /code extension

copy.ts

import type { ExtensionAPI, ExtensionContext, SessionEntry, SessionMessageEntry } from "@mariozechner/pi-coding-agent";
import { copyToClipboard } from "@mariozechner/pi-coding-agent";

/**
 * Missing vs upstream branch (requires core changes):
 * 1) Inline code-block labels in assistant messages:
 *    - @mariozechner/pi-tui: add MarkdownOptions + code-block hooks/labels in
 *      packages/tui/src/components/markdown.ts and export via packages/tui/src/index.ts.
 *    - @mariozechner/pi-coding-agent: wire registry + label rendering in
 *      packages/coding-agent/src/modes/interactive/components/assistant-message.ts and
 *      packages/coding-agent/src/modes/interactive/interactive-mode.ts.
 * 2) Show labels only while /copy autocomplete is open:
 *    - Add autocomplete visibility callbacks in
 *      packages/coding-agent/src/modes/interactive/components/custom-editor.ts and
 *      wire it in packages/coding-agent/src/modes/interactive/interactive-mode.ts.
 * 3) Autocomplete select-to-submit:
 *    - Add submitOnSelect to AutocompleteItem/SelectItem in
 *      packages/tui/src/autocomplete.ts and packages/tui/src/components/select-list.ts,
 *      then honor it in packages/tui/src/components/editor.ts.
 * 4) Render-accurate block extraction (no regex drift in nested lists/fences):
 *    - Emit canonical code-block metadata from Markdown renderer in
 *      packages/tui/src/components/markdown.ts.
 * 5) Stable per-message numbering across multi-part assistant content:
 *    - Maintain a per-message code-block registry keyed by content block order in
 *      packages/coding-agent/src/modes/interactive/interactive-mode.ts.
 * 6) Cannot override built-in commands from extensions:
 *    - The loader skips extension commands that conflict with built-ins.
 *    - Upstream needs a command override mechanism (prefer extensions) or a
 *      built-in /copy hook that delegates to extensions.
 */

type CodeBlock = {
	language: string;
	code: string;
	preview: string;
};

type TextBlock = {
	type?: string;
	text?: string;
};

const CODE_BLOCK_REGEX = /```([^\n]*)\n([\s\S]*?)```/g;

function extractTextContent(content: unknown): string {
	if (typeof content === "string") {
		return content;
	}

	if (!Array.isArray(content)) {
		return "";
	}

	const parts: string[] = [];
	for (const part of content) {
		if (!part || typeof part !== "object") continue;
		const block = part as TextBlock;
		if (block.type === "text" && typeof block.text === "string") {
			parts.push(block.text);
		}
	}

	return parts.join("\n");
}

function buildPreview(code: string): string {
	const firstLine = code.split(/\r?\n/, 1)[0] ?? "";
	const trimmed = firstLine.trim();
	return trimmed.length > 60 ? `${trimmed.slice(0, 60)}…` : trimmed;
}

function extractCodeBlocks(markdown: string): CodeBlock[] {
	const blocks: CodeBlock[] = [];
	let match: RegExpExecArray | null;
	while ((match = CODE_BLOCK_REGEX.exec(markdown)) !== null) {
		const language = (match[1] ?? "").trim();
		const code = (match[2] ?? "").trimEnd();
		blocks.push({ language, code, preview: buildPreview(code) });
	}
	return blocks;
}

function findLastAssistantEntry(entries: SessionEntry[]): SessionMessageEntry | undefined {
	for (let i = entries.length - 1; i >= 0; i--) {
		const entry = entries[i];
		if (entry.type === "message" && entry.message.role === "assistant") {
			return entry;
		}
	}
	return undefined;
}

function getLastAssistantData(ctx: ExtensionContext): { text: string; blocks: CodeBlock[] } | undefined {
	const branch = ctx.sessionManager.getBranch();
	const lastAssistant = findLastAssistantEntry(branch);
	if (!lastAssistant) {
		return undefined;
	}

	const text = extractTextContent(lastAssistant.message.content);
	return { text, blocks: extractCodeBlocks(text) };
}

function parseBlockIndex(args: string): number | null {
	const trimmed = args.trim();
	if (!trimmed) return null;
	const normalized = trimmed.startsWith("#") ? trimmed.slice(1) : trimmed;
	if (!/^\d+$/.test(normalized)) {
		return Number.NaN;
	}
	return Number.parseInt(normalized, 10);
}

function notify(ctx: ExtensionContext, message: string, type: "info" | "warning" | "error"): void {
	if (ctx.hasUI) {
		ctx.ui.notify(message, type);
	}
}

export default function (pi: ExtensionAPI) {
	let lastAssistantText: string | undefined;
	let lastCodeBlocks: CodeBlock[] = [];

	const updateCacheFromSession = (ctx: ExtensionContext) => {
		const data = getLastAssistantData(ctx);
		if (!data) {
			lastAssistantText = undefined;
			lastCodeBlocks = [];
			return;
		}
		lastAssistantText = data.text;
		lastCodeBlocks = data.blocks;
	};

	const updateCacheFromMessage = (message: { role?: string; content?: unknown }) => {
		if (message.role !== "assistant") return;
		const text = extractTextContent(message.content);
		lastAssistantText = text;
		lastCodeBlocks = extractCodeBlocks(text);
	};

	pi.on("session_start", (_event, ctx) => updateCacheFromSession(ctx));
	pi.on("session_switch", (_event, ctx) => updateCacheFromSession(ctx));
	pi.on("session_tree", (_event, ctx) => updateCacheFromSession(ctx));
	pi.on("session_compact", (_event, ctx) => updateCacheFromSession(ctx));
	pi.on("session_fork", (_event, ctx) => updateCacheFromSession(ctx));
	pi.on("turn_end", (event) => updateCacheFromMessage(event.message));

	pi.registerCommand("code", {
		description: "Copy last assistant message or a code block to clipboard",
		getArgumentCompletions: (prefix: string) => {
			if (!lastAssistantText || lastCodeBlocks.length === 0) return null;
			const items = lastCodeBlocks.map((block, index) => {
				const value = String(index + 1);
				const label = `#${value}`;
				const descriptionParts = [block.language, block.preview].filter((part) => part.length > 0);
				return {
					value,
					label,
					description: descriptionParts.length > 0 ? descriptionParts.join(" · ") : undefined,
				};
			});

			const trimmedPrefix = prefix.trim();
			const filtered = trimmedPrefix
				? items.filter((item) => item.value.startsWith(trimmedPrefix) || item.label.startsWith(trimmedPrefix))
				: items;
			return filtered.length > 0 ? filtered : null;
		},
		handler: async (args, ctx) => {
			const data = getLastAssistantData(ctx);
			if (!data) {
				notify(ctx, "No assistant message found.", "warning");
				return;
			}

			if (!data.text.trim()) {
				notify(ctx, "Last assistant message has no text.", "warning");
				return;
			}

			const index = parseBlockIndex(args);
			if (index === null) {
				try {
					copyToClipboard(data.text);
					notify(ctx, "Copied last assistant message to clipboard.", "info");
				} catch (error) {
					notify(ctx, error instanceof Error ? error.message : String(error), "error");
				}
				return;
			}

			if (!Number.isFinite(index) || index < 1) {
				notify(ctx, "Usage: /copy [code block number]", "warning");
				return;
			}

			if (data.blocks.length === 0) {
				notify(ctx, "No code blocks in the last assistant message.", "warning");
				return;
			}

			const block = data.blocks[index - 1];
			if (!block) {
				notify(ctx, `Code block ${index} not found.`, "warning");
				return;
			}

			try {
				copyToClipboard(block.code);
				const suffix = block.language ? ` (${block.language})` : "";
				notify(ctx, `Copied code block ${index}${suffix}.`, "info");
			} catch (error) {
				notify(ctx, error instanceof Error ? error.message : String(error), "error");
			}
		},
	});
}