--- /dev/null
+import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
+import { spawn } from "node:child_process";
+import fs from "node:fs/promises";
+import path from "node:path";
+
+const TOP_FILES = 8;
+const TOP_RANGES = 12;
+const MAX_RG_EVENTS = 600;
+const SNIPPET_CONTEXT_LINES = 2;
+const READ_WINDOW_LINES = 140;
+const RG_TIMEOUT_MS = 15_000;
+
+interface MemoryMatch {
+ file: string;
+ line: number;
+ text: string;
+}
+
+interface CandidateFile {
+ file: string;
+ matches: MemoryMatch[];
+ score: number;
+ keywords: string[];
+}
+
+interface CandidateRange {
+ file: string;
+ start: number;
+ end: number;
+ offset: number;
+ limit: number;
+ heading?: string;
+ headingLine?: number;
+ matchLines: number[];
+ keywords: string[];
+ snippet: string;
+ score: number;
+}
+
+interface ActiveMemorySearch {
+ directory: string;
+ keywords: string[];
+ allowed: Map<string, Array<{ start: number; end: number }>>;
+}
+
+function parseArgs(input: string): string[] {
+ const args: string[] = [];
+ let current = "";
+ let quote: "'" | '"' | undefined;
+ let escaping = false;
+
+ for (const char of input) {
+ if (escaping) {
+ current += char;
+ escaping = false;
+ continue;
+ }
+ if (char === "\\") {
+ escaping = true;
+ continue;
+ }
+ if (quote) {
+ if (char === quote) {
+ quote = undefined;
+ } else {
+ current += char;
+ }
+ continue;
+ }
+ if (char === "'" || char === '"') {
+ quote = char;
+ continue;
+ }
+ if (/\s/.test(char)) {
+ if (current.length > 0) {
+ args.push(current);
+ current = "";
+ }
+ continue;
+ }
+ current += char;
+ }
+
+ if (escaping) current += "\\";
+ if (quote) throw new Error(`Unclosed ${quote} quote`);
+ if (current.length > 0) args.push(current);
+ return args;
+}
+
+function normalizePath(cwd: string, inputPath: string): string {
+ const trimmed = inputPath.startsWith("@") ? inputPath.slice(1) : inputPath;
+ return path.resolve(cwd, trimmed);
+}
+
+function isMarkdownFile(filePath: string): boolean {
+ const ext = path.extname(filePath).toLowerCase();
+ return ext === ".md" || ext === ".markdown" || ext === ".mdown" || ext === ".mkd";
+}
+
+function isInside(parent: string, child: string): boolean {
+ const relative = path.relative(parent, child);
+ return relative === "" || (!relative.startsWith("..") && !path.isAbsolute(relative));
+}
+
+function displayPath(cwd: string, filePath: string): string {
+ const relative = path.relative(cwd, filePath);
+ if (relative && !relative.startsWith("..") && !path.isAbsolute(relative)) return relative;
+ return filePath;
+}
+
+function truncateLine(line: string, max = 220): string {
+ return line.length <= max ? line : `${line.slice(0, max - 1)}…`;
+}
+
+function keywordHits(text: string, keywords: string[]): string[] {
+ const lower = text.toLowerCase();
+ return keywords.filter((keyword) => lower.includes(keyword.toLowerCase()));
+}
+
+function scoreMatches(matches: MemoryMatch[], keywords: string[]): { score: number; hits: string[] } {
+ const hits = new Set<string>();
+ let headingMatches = 0;
+ for (const match of matches) {
+ for (const keyword of keywordHits(match.text, keywords)) hits.add(keyword);
+ if (/^\s*#{1,6}\s+/.test(match.text)) headingMatches += 1;
+ }
+ return {
+ score: matches.length * 3 + hits.size * 10 + headingMatches * 8,
+ hits: [...hits],
+ };
+}
+
+async function runRipgrep(directory: string, keywords: string[]): Promise<MemoryMatch[]> {
+ const args = [
+ "--json",
+ "--line-number",
+ "--ignore-case",
+ "--fixed-strings",
+ "--glob",
+ "*.md",
+ "--glob",
+ "*.markdown",
+ "--glob",
+ "*.mdown",
+ "--glob",
+ "*.mkd",
+ ];
+ for (const keyword of keywords) args.push("-e", keyword);
+ args.push("--", directory);
+
+ return new Promise((resolve, reject) => {
+ const child = spawn("rg", args, { stdio: ["ignore", "pipe", "pipe"] });
+ const timer = setTimeout(() => child.kill("SIGTERM"), RG_TIMEOUT_MS);
+ let stdout = "";
+ let stderr = "";
+
+ child.stdout.setEncoding("utf8");
+ child.stderr.setEncoding("utf8");
+ child.stdout.on("data", (chunk) => {
+ stdout += chunk;
+ // Keep memory bounded if the corpus is very noisy.
+ if (stdout.length > 5_000_000) child.kill("SIGTERM");
+ });
+ child.stderr.on("data", (chunk) => {
+ stderr += chunk;
+ });
+ child.on("error", (error) => {
+ clearTimeout(timer);
+ reject(error);
+ });
+ child.on("close", (code, signal) => {
+ clearTimeout(timer);
+ if (signal && signal !== "SIGTERM") {
+ reject(new Error(`rg stopped by signal ${signal}`));
+ return;
+ }
+ if (code !== 0 && code !== 1 && signal !== "SIGTERM") {
+ reject(new Error(stderr.trim() || `rg exited with code ${code}`));
+ return;
+ }
+
+ const matches: MemoryMatch[] = [];
+ for (const line of stdout.split("\n")) {
+ if (!line.trim()) continue;
+ let event: any;
+ try {
+ event = JSON.parse(line);
+ } catch {
+ continue;
+ }
+ if (event.type !== "match") continue;
+ const fileText = event.data?.path?.text;
+ const lineNumber = event.data?.line_number;
+ const lineText = event.data?.lines?.text;
+ if (typeof fileText !== "string" || typeof lineNumber !== "number" || typeof lineText !== "string") {
+ continue;
+ }
+ matches.push({ file: path.resolve(fileText), line: lineNumber, text: lineText.replace(/\r?\n$/, "") });
+ if (matches.length >= MAX_RG_EVENTS) break;
+ }
+ resolve(matches);
+ });
+ });
+}
+
+function findNearestHeading(lines: string[], beforeLine: number): { line: number; text: string } | undefined {
+ for (let i = Math.min(beforeLine, lines.length) - 1; i >= 0; i -= 1) {
+ const text = lines[i] ?? "";
+ if (/^\s*#{1,6}\s+/.test(text)) return { line: i + 1, text: text.trim() };
+ }
+ return undefined;
+}
+
+function findNextHeading(lines: string[], afterLine: number): number | undefined {
+ for (let i = Math.max(0, afterLine); i < lines.length; i += 1) {
+ if (/^\s*#{1,6}\s+/.test(lines[i] ?? "")) return i + 1;
+ }
+ return undefined;
+}
+
+function clusterLines(lineNumbers: number[]): number[][] {
+ const sorted = [...new Set(lineNumbers)].sort((a, b) => a - b);
+ const clusters: number[][] = [];
+ for (const line of sorted) {
+ const last = clusters[clusters.length - 1];
+ if (!last || line - last[last.length - 1] > 40) {
+ clusters.push([line]);
+ } else {
+ last.push(line);
+ }
+ }
+ return clusters;
+}
+
+async function buildRanges(cwd: string, candidates: CandidateFile[], keywords: string[]): Promise<CandidateRange[]> {
+ const ranges: CandidateRange[] = [];
+ for (const candidate of candidates) {
+ let content: string;
+ try {
+ content = await fs.readFile(candidate.file, "utf8");
+ } catch {
+ continue;
+ }
+ const lines = content.split(/\r?\n/);
+ const clusters = clusterLines(candidate.matches.map((match) => match.line));
+ for (const cluster of clusters) {
+ const firstMatch = cluster[0];
+ const lastMatch = cluster[cluster.length - 1];
+ const heading = findNearestHeading(lines, firstMatch);
+ const nextHeading = findNextHeading(lines, lastMatch + 1);
+
+ let start = Math.max(1, firstMatch - 20);
+ if (heading && firstMatch - heading.line <= 80) start = heading.line;
+
+ let end = Math.min(lines.length, Math.max(lastMatch + 60, start + 80));
+ if (nextHeading && nextHeading > lastMatch && nextHeading - 1 >= start) {
+ end = Math.min(end, nextHeading - 1);
+ }
+ end = Math.min(lines.length, start + READ_WINDOW_LINES - 1, end);
+
+ const snippetStart = Math.max(1, firstMatch - SNIPPET_CONTEXT_LINES);
+ const snippetEnd = Math.min(lines.length, lastMatch + SNIPPET_CONTEXT_LINES);
+ const snippetLines: string[] = [];
+ for (let line = snippetStart; line <= snippetEnd; line += 1) {
+ const marker = cluster.includes(line) ? ">" : " ";
+ snippetLines.push(`${marker} ${String(line).padStart(5, " ")}: ${truncateLine(lines[line - 1] ?? "")}`);
+ }
+
+ const clusterKeywords = new Set<string>();
+ for (const line of cluster) {
+ for (const keyword of keywordHits(lines[line - 1] ?? "", keywords)) clusterKeywords.add(keyword);
+ }
+
+ ranges.push({
+ file: candidate.file,
+ start,
+ end,
+ offset: start,
+ limit: end - start + 1,
+ heading: heading?.text,
+ headingLine: heading?.line,
+ matchLines: cluster,
+ keywords: [...clusterKeywords],
+ snippet: snippetLines.join("\n"),
+ score: candidate.score + cluster.length * 5 + clusterKeywords.size * 5,
+ });
+ }
+ }
+ return ranges.sort((a, b) => b.score - a.score).slice(0, TOP_RANGES);
+}
+
+function groupCandidates(matches: MemoryMatch[], keywords: string[]): CandidateFile[] {
+ const byFile = new Map<string, MemoryMatch[]>();
+ for (const match of matches) {
+ const list = byFile.get(match.file) ?? [];
+ list.push(match);
+ byFile.set(match.file, list);
+ }
+
+ const candidates: CandidateFile[] = [];
+ for (const [file, fileMatches] of byFile) {
+ const { score, hits } = scoreMatches(fileMatches, keywords);
+ candidates.push({ file, matches: fileMatches, score, keywords: hits });
+ }
+ return candidates.sort((a, b) => b.score - a.score).slice(0, TOP_FILES);
+}
+
+function buildPrompt(cwd: string, directory: string, keywords: string[], matches: MemoryMatch[], ranges: CandidateRange[]): string {
+ const lines: string[] = [];
+ lines.push(`Memory search results for ${directory}`);
+ lines.push(`Keywords: ${keywords.map((keyword) => JSON.stringify(keyword)).join(", ")}`);
+ lines.push(`Search method: ripgrep over Markdown files only (*.md, *.markdown, *.mdown, *.mkd).`);
+ lines.push(`Total matched lines sampled: ${matches.length}${matches.length >= MAX_RG_EVENTS ? ` (capped at ${MAX_RG_EVENTS})` : ""}.`);
+ lines.push("");
+ lines.push("Instructions for the agent:");
+ lines.push("1. Do not bulk-read this Markdown directory.");
+ lines.push("2. Use Pi's `read` tool only on the suggested ranges below. The memory extension will block other Markdown reads inside this directory for this turn.");
+ lines.push("3. Start with the top 1-3 ranges. If those are insufficient, read additional listed ranges or ask for a refined `/memory <directory> <keywords...>` search.");
+ lines.push("4. When answering, cite paths and line numbers from the files you read.");
+ lines.push("");
+
+ if (ranges.length === 0) {
+ lines.push("No candidate ranges were found.");
+ lines.push("Answer that no matching Markdown memory entries were found, and suggest more specific or alternate keywords.");
+ return lines.join("\n");
+ }
+
+ lines.push("Suggested candidate ranges:");
+ for (let i = 0; i < ranges.length; i += 1) {
+ const range = ranges[i];
+ const fileDisplay = displayPath(cwd, range.file);
+ lines.push("");
+ lines.push(`${i + 1}. ${fileDisplay}:${range.start}-${range.end}`);
+ lines.push(` read arguments: { "path": ${JSON.stringify(range.file)}, "offset": ${range.offset}, "limit": ${range.limit} }`);
+ if (range.heading) lines.push(` nearest heading: ${range.heading} (line ${range.headingLine})`);
+ lines.push(` matched lines: ${range.matchLines.join(", ")}`);
+ if (range.keywords.length) lines.push(` matched keywords: ${range.keywords.map((k) => JSON.stringify(k)).join(", ")}`);
+ lines.push(" snippet:");
+ lines.push("```text");
+ lines.push(range.snippet);
+ lines.push("```");
+ }
+
+ lines.push("");
+ lines.push("Now read the most relevant listed range(s), then summarize what the memory says about these keywords.");
+ return lines.join("\n");
+}
+
+function buildAllowedMap(ranges: CandidateRange[]): Map<string, Array<{ start: number; end: number }>> {
+ const allowed = new Map<string, Array<{ start: number; end: number }>>();
+ for (const range of ranges) {
+ const list = allowed.get(range.file) ?? [];
+ list.push({ start: range.start, end: range.end });
+ allowed.set(range.file, list);
+ }
+ return allowed;
+}
+
+export default function memoryExtension(pi: ExtensionAPI) {
+ let activeMemory: ActiveMemorySearch | undefined;
+
+ pi.registerCommand("memory", {
+ description: "Search Markdown memory: /memory <directory> <keyword...>",
+ handler: async (args, ctx) => {
+ let parsed: string[];
+ try {
+ parsed = parseArgs(args ?? "");
+ } catch (error: any) {
+ ctx.ui.notify(error?.message ?? "Failed to parse /memory arguments", "error");
+ return;
+ }
+
+ const [directoryArg, ...keywordArgs] = parsed;
+ const keywords = keywordArgs.map((keyword) => keyword.trim()).filter(Boolean);
+ if (!directoryArg || keywords.length === 0) {
+ ctx.ui.notify("Usage: /memory <directory> <keyword...>", "error");
+ return;
+ }
+
+ const directory = normalizePath(ctx.cwd, directoryArg);
+ let stat;
+ try {
+ stat = await fs.stat(directory);
+ } catch {
+ ctx.ui.notify(`Memory directory not found: ${directory}`, "error");
+ return;
+ }
+ if (!stat.isDirectory()) {
+ ctx.ui.notify(`Memory path is not a directory: ${directory}`, "error");
+ return;
+ }
+
+ ctx.ui.notify(`Searching Markdown memory in ${directory}...`, "info");
+
+ let matches: MemoryMatch[];
+ try {
+ matches = await runRipgrep(directory, keywords);
+ } catch (error: any) {
+ ctx.ui.notify(`memory rg failed: ${error?.message ?? String(error)}`, "error");
+ return;
+ }
+
+ const candidates = groupCandidates(matches, keywords);
+ const ranges = await buildRanges(ctx.cwd, candidates, keywords);
+ activeMemory = {
+ directory,
+ keywords,
+ allowed: buildAllowedMap(ranges),
+ };
+
+ const prompt = buildPrompt(ctx.cwd, directory, keywords, matches, ranges);
+ pi.sendUserMessage(prompt);
+ },
+ });
+
+ pi.on("before_agent_start", async (event) => {
+ if (!activeMemory) return;
+ return {
+ systemPrompt:
+ event.systemPrompt +
+ "\n\nActive /memory search policy: For this turn, do not bulk-read the Markdown memory directory. Only use the read tool on the exact path/offset/limit ranges listed in the user's /memory search results. If more information is needed, ask the user to run a refined /memory search.",
+ };
+ });
+
+ pi.on("tool_call", async (event, ctx) => {
+ if (!activeMemory || event.toolName !== "read") return;
+ const input = event.input as { path?: string; offset?: number; limit?: number };
+ if (!input?.path) return;
+
+ const requestedPath = normalizePath(ctx.cwd, input.path);
+ if (!isMarkdownFile(requestedPath) || !isInside(activeMemory.directory, requestedPath)) return;
+
+ const allowedRanges = activeMemory.allowed.get(requestedPath);
+ if (!allowedRanges) {
+ return {
+ block: true,
+ reason: `Blocked by /memory: ${requestedPath} was not one of the top returned Markdown ranges. Run /memory with refined keywords if this file is needed.`,
+ };
+ }
+
+ if (typeof input.offset !== "number" || typeof input.limit !== "number") {
+ return {
+ block: true,
+ reason: "Blocked by /memory: Markdown memory reads must include offset and limit from the suggested ranges.",
+ };
+ }
+
+ const start = input.offset;
+ const end = input.offset + input.limit - 1;
+ const allowed = allowedRanges.some((range) => start >= range.start && end <= range.end);
+ if (!allowed) {
+ const rangesText = allowedRanges.map((range) => `${range.start}-${range.end}`).join(", ");
+ return {
+ block: true,
+ reason: `Blocked by /memory: requested lines ${start}-${end} are outside the allowed ranges for this file (${rangesText}).`,
+ };
+ }
+ });
+
+ pi.on("agent_end", async () => {
+ activeMemory = undefined;
+ });
+}
projects / dotfiles / commitdiff