1  import type { ExtensionAPI, ExtensionCommandContext } from "@mariozechner/pi-coding-agent";
  2  import { extractSnippets, extractText } from "./snippets";
  3  import type { Snippet } from "./snippets";
  4  import { pickAction, pickSnippet } from "./ui";
  5  import type { PickResult } from "./ui";
  6  import { copyToClipboard, insertIntoEditor, runSnippet } from "./actions";
  7  
  8  type ParsedArgs = {
  9  	scope: "last" | "all";
 10  	action?: "copy" | "insert" | "run";
 11  	index?: number;
 12  	includeInline: boolean;
 13  	limit: number;
 14  };
 15  
 16  function parseArgs(args?: string): ParsedArgs {
 17  	const tokens = args?.trim().split(/\s+/).filter(Boolean) ?? [];
 18  	const parsed: ParsedArgs = { scope: "all", includeInline: true, limit: 200 };
 19  
 20  	for (const token of tokens) {
 21  		if (token === "all" || token === "last") {
 22  			parsed.scope = token;
 23  			continue;
 24  		}
 25  		if (token === "inline") {
 26  			parsed.includeInline = true;
 27  			continue;
 28  		}
 29  		if (token === "blocks") {
 30  			parsed.includeInline = false;
 31  			continue;
 32  		}
 33  		if (token === "copy" || token === "insert" || token === "run") {
 34  			parsed.action = token;
 35  			continue;
 36  		}
 37  		if (token.startsWith("limit=")) {
 38  			const value = Number.parseInt(token.slice("limit=".length), 10);
 39  			if (!Number.isNaN(value) && value > 0) parsed.limit = value;
 40  			continue;
 41  		}
 42  		if (/^\d+$/.test(token)) {
 43  			parsed.index = Math.max(0, Number.parseInt(token, 10) - 1);
 44  		}
 45  	}
 46  
 47  	return parsed;
 48  }
 49  
 50  function collectSnippets(
 51  	ctx: ExtensionCommandContext,
 52  	scope: "last" | "all",
 53  	includeInline: boolean,
 54  	limit: number,
 55  ): Snippet[] {
 56  	const branchEntries = ctx.sessionManager.getBranch();
 57  	const candidateEntries = branchEntries.filter((entry) => {
 58  		if (entry.type !== "message") return false;
 59  		const role = entry.message?.role;
 60  		return role === "assistant" || role === "toolResult";
 61  	});
 62  
 63  	if (candidateEntries.length === 0) return [];
 64  
 65  	const sorted = candidateEntries.slice().sort((a, b) => {
 66  		const aTime = Date.parse(a.timestamp);
 67  		const bTime = Date.parse(b.timestamp);
 68  		return bTime - aTime;
 69  	});
 70  
 71  	const entriesToScan = scope === "all" ? sorted : [sorted[0]!];
 72  
 73  	let snippets: Snippet[] = [];
 74  	let nextId = 0;
 75  	for (const entry of entriesToScan) {
 76  		if (snippets.length >= limit) break;
 77  		const label = new Date(entry.timestamp).toLocaleTimeString();
 78  
 79  		// Extract bash tool calls directly (these often render as `$ ...` in the UI but are stored structurally)
 80  		if (entry.type === "message" && entry.message?.role === "assistant" && Array.isArray(entry.message.content)) {
 81  			for (const block of entry.message.content as any[]) {
 82  				if (snippets.length >= limit) break;
 83  				if (!block || typeof block !== "object") continue;
 84  				if (block.type !== "toolCall") continue;
 85  				if (block.name !== "bash") continue;
 86  				const cmd = typeof block.arguments?.command === "string" ? block.arguments.command.trim() : "";
 87  				if (!cmd) continue;
 88  
 89  				snippets.push({
 90  					id: nextId,
 91  					type: "block",
 92  					language: "bash",
 93  					content: cmd,
 94  					messageId: entry.id,
 95  					sourceLabel: label,
 96  				});
 97  				nextId += 1;
 98  			}
 99  		}
100  
101  		const text = extractText(entry.message.content);
102  		if (!text) continue;
103  		const extracted = extractSnippets(text, entry.id, label, nextId, includeInline, limit - snippets.length);
104  		snippets = snippets.concat(extracted);
105  		nextId = snippets.length;
106  	}
107  
108  	return snippets;
109  }
110  
111  export default function codeActionsExtension(pi: ExtensionAPI) {
112  	pi.registerCommand("code", {
113  		description: "Pick code from assistant messages and copy/insert/run it",
114  		handler: async (args, ctx) => {
115  			if (!ctx.hasUI && (!args || args.trim().length === 0)) {
116  				return;
117  			}
118  
119  			const parsed = parseArgs(args);
120  			const snippets = collectSnippets(ctx, parsed.scope, parsed.includeInline, parsed.limit);
121  
122  			if (snippets.length === 0) {
123  				if (ctx.hasUI) ctx.ui.notify("No code snippets found. /code tracks code blocks and filepaths only.", "warning");
124  				return;
125  			}
126  
127  			let snippet: Snippet | undefined;
128  			let pickedAction: PickResult["action"] | undefined;
129  			if (parsed.index !== undefined) {
130  				snippet = snippets[parsed.index];
131  				if (!snippet) {
132  					if (ctx.hasUI) ctx.ui.notify("Snippet index out of range.", "warning");
133  					return;
134  				}
135  			} else {
136  				if (!ctx.hasUI) return;
137  				const result = await pickSnippet(ctx, snippets);
138  				if (!result) return;
139  				snippet = result.snippet;
140  				pickedAction = result.action;
141  			}
142  
143  			let action = parsed.action ?? pickedAction;
144  			if (!action) {
145  				if (!ctx.hasUI) return;
146  				action = await pickAction(ctx);
147  			}
148  			if (!action) return;
149  
150  			if (action === "copy") {
151  				const ok = await copyToClipboard(pi, snippet.content);
152  				if (ctx.hasUI) {
153  					ctx.ui.notify(ok ? "Copied to clipboard." : "Failed to copy to clipboard.", ok ? "info" : "error");
154  				}
155  				return;
156  			}
157  
158  			if (action === "insert") {
159  				if (!ctx.hasUI) return;
160  				insertIntoEditor(ctx, snippet.content);
161  				ctx.ui.notify("Inserted snippet into editor.", "info");
162  				return;
163  			}
164  
165  			if (action === "run") {
166  				if (!ctx.hasUI) return;
167  				const ok = await ctx.ui.confirm("Run snippet?", "This will execute the selected snippet in your shell.");
168  				if (!ok) return;
169  				await runSnippet(pi, ctx, snippet.content);
170  			}
171  		},
172  	});
173  }