pi-extensions/memory/index.ts

/**
 * memory — cross-agent memory access for pi.
 *
 * READS span this extension's own canonical store (project + user scope) AND
 * other agent systems' memory files (Claude, Copilot, AGENTS.md, Gemini,
 * Aider), surfaced verbatim and source-tagged. WRITES only ever touch the
 * canonical store (see ./store.ts) — foreign files are never modified.
 *
 * Surfaces
 * ────────
 *   Tools (LLM-callable):  memory_search · memory_read · memory_write · memory_forget
 *   Commands (human):      /memory-list · /memory-search · /memory-show
 *                          /memory-write · /memory-forget
 *   Context injection:     on before_agent_start, the canonical index (one line
 *                          per memory) is appended to the system prompt so the
 *                          model always knows what it can pull. Only the
 *                          canonical store is indexed — foreign systems already
 *                          load their own files, so we don't double-inject them.
 *
 * Settings (~/.pi/agent/settings.json, key "memory")
 *   enabled       boolean  master switch (default true; read at load)
 *   injectIndex   boolean  append the index to the system prompt (default true)
 *   readForeign   boolean  include foreign agent memory in reads (default true)
 */

import { homedir } from "node:os";
import type {
	ExtensionAPI,
	ExtensionCommandContext,
	ExtensionContext,
} from "@mariozechner/pi-coding-agent";
import { Type } from "typebox";
import { readExtensionBooleanSetting } from "../shared/pi-settings.js";
import {
	buildIndex,
	type CanonicalMemory,
	findByName,
	isMemoryScope,
	isMemoryType,
	listMemories,
	type MemoryScope,
	MEMORY_SCOPES,
	type MemoryType,
	MEMORY_TYPES,
	deleteMemory,
	normalizeEnumArg,
	searchMemories,
	storeDirForScope,
	writeMemory,
} from "./store.js";
import {
	discoverForeignSources,
	findForeignByName,
	type ForeignSource,
	FOREIGN_SYSTEMS,
	formatForeignDigest,
	formatForeignFull,
	isForeignSystem,
} from "./sources.js";

const EXTENSION_NAME = "memory";

// ── settings ──────────────────────────────────────────────────────────────

const isEnabled = () => readExtensionBooleanSetting(EXTENSION_NAME, "enabled", true);
const shouldInjectIndex = () => readExtensionBooleanSetting(EXTENSION_NAME, "injectIndex", true);
const shouldReadForeign = () => readExtensionBooleanSetting(EXTENSION_NAME, "readForeign", true);

// ── parameter schemas ───────────────────────────────────────────────────────
// These are closed string vocabularies, but they are declared as plain strings
// (allowed values spelled out in the description) rather than anyOf/const
// unions on purpose: some local models emit const-union values JSON-quoted
// (e.g. `"\"feedback\""`), which strict schema validation then rejects BEFORE
// the tool runs. Keeping them lenient lets the call reach `execute`, where
// normalizeEnumArg strips the stray quoting and we validate against the
// vocabulary ourselves — with a helpful error instead of a hard schema reject.

const memoryTypeSchema = Type.String({
	description: `Memory type — one of: ${MEMORY_TYPES.join(" | ")}.`,
});
const writeScopeSchema = Type.String({
	description: `Store scope — one of: ${MEMORY_SCOPES.join(" | ")} (default project).`,
});
const scopeFilterSchema = Type.String({
	description: `Scope filter — one of: ${MEMORY_SCOPES.join(" | ")} | all (default all).`,
});
const foreignSystemSchema = Type.String({
	description: `Foreign agent system — one of: ${FOREIGN_SYSTEMS.join(" | ")}.`,
});

// ── small helpers ───────────────────────────────────────────────────────

function text(body: string) {
	return { content: [{ type: "text" as const, text: body }], details: undefined };
}

function dirs(ctx: ExtensionContext | ExtensionCommandContext) {
	const cwd = ctx.cwd;
	const home = homedir();
	return { cwd, home };
}

/** Canonical memories across scopes, optionally narrowed by a scope filter. */
function loadCanonical(
	ctx: ExtensionContext | ExtensionCommandContext,
	scopeFilter: MemoryScope | "all" = "all",
): CanonicalMemory[] {
	const { cwd, home } = dirs(ctx);
	const out: CanonicalMemory[] = [];
	for (const scope of MEMORY_SCOPES) {
		if (scopeFilter !== "all" && scopeFilter !== scope) continue;
		out.push(...listMemories(storeDirForScope(scope, cwd, home), scope));
	}
	return out;
}

function formatMemory(m: CanonicalMemory): string {
	const meta = `[${m.scope}/${m.type}] ${m.name}`;
	const stamps = [m.created && `created ${m.created}`, m.updated && `updated ${m.updated}`]
		.filter(Boolean)
		.join(" · ");
	const desc = m.description ? `\ndescription: ${m.description}` : "";
	return `${meta}${desc}${stamps ? `\n${stamps}` : ""}\n\n${m.body}`;
}

function foreignSnippet(content: string, lowerQuery: string): string {
	const idx = content.toLowerCase().indexOf(lowerQuery);
	if (idx === -1) return "";
	const start = Math.max(0, idx - 40);
	const end = Math.min(content.length, idx + lowerQuery.length + 60);
	const prefix = start > 0 ? "…" : "";
	const suffix = end < content.length ? "…" : "";
	return `${prefix}${content.slice(start, end).replace(/\s+/g, " ").trim()}${suffix}`;
}

function searchForeign(sources: readonly ForeignSource[], query: string) {
	const lowerQuery = query.trim().toLowerCase();
	if (!lowerQuery) return [];
	return sources
		.filter((s) => s.content.toLowerCase().includes(lowerQuery))
		.map((s) => ({ source: s, snippet: foreignSnippet(s.content, lowerQuery) }));
}

/**
 * Argument completer for /memory-show and /memory-forget. Completers don't get a
 * ctx, so we best-effort use process.cwd() (the session cwd in practice) to list
 * canonical memory names across both scopes.
 */
function completeMemoryName(prefix: string) {
	try {
		const cwd = process.cwd();
		const home = homedir();
		const names = new Set<string>();
		for (const scope of MEMORY_SCOPES) {
			for (const m of listMemories(storeDirForScope(scope, cwd, home), scope)) {
				names.add(m.name);
			}
		}
		const hits = [...names]
			.filter((n) => n.toLowerCase().startsWith(prefix.toLowerCase()))
			.map((n) => ({ value: n, label: n }));
		return hits.length > 0 ? hits : null;
	} catch {
		return null;
	}
}

// ── extension ───────────────────────────────────────────────────────────

export default function (pi: ExtensionAPI) {
	if (!isEnabled()) return;

	// ── Context injection ──────────────────────────────────────────────────
	// Append the canonical index to the system prompt on every agent loop start.
	// before_agent_start lets us chain onto the assembled prompt; we add a clearly
	// delimited block and skip entirely when the store is empty (no token cost).
	pi.on("before_agent_start", async (event, ctx) => {
		try {
			if (!shouldInjectIndex()) return undefined;
			const index = buildIndex(loadCanonical(ctx));
			if (!index) return undefined;
			const block = [
				"<persistent-memory>",
				"You have a persistent memory store (project + user scope). The entries below",
				"are available — call `memory_read` with a `name` to load a full body,",
				"`memory_search` to search by keyword, `memory_write` to save a new/updated",
				"memory, and `memory_forget` to remove one. Memory from other agent tools",
				"(CLAUDE.md, AGENTS.md, etc.) is also readable via `memory_read`/`memory_search`.",
				"",
				index,
				"</persistent-memory>",
			].join("\n");
			return { systemPrompt: `${event.systemPrompt}\n\n${block}` };
		} catch {
			return undefined;
		}
	});

	// ── Tools ────────────────────────────────────────────────────────────

	pi.registerTool({
		name: "memory_search",
		label: "Memory Search",
		description:
			"Keyword/substring search across persistent memory — this agent's canonical " +
			"store (project + user) and, unless disabled, other agent systems' memory " +
			"(CLAUDE.md, Copilot, AGENTS.md, Gemini, Aider). Case-insensitive.",
		promptSnippet: "Search persistent memory by keyword",
		parameters: Type.Object({
			query: Type.String({ description: "Case-insensitive substring to search for." }),
			scope: Type.Optional(
				scopeFilterSchema,
			),
			type: Type.Optional(memoryTypeSchema),
			includeForeign: Type.Optional(
				Type.Boolean({
					description:
						"Also search other agent systems' memory files. Defaults to true.",
				}),
			),
		}),
		async execute(_id, params, _signal, _onUpdate, ctx) {
			const scopeArg = normalizeEnumArg(params.scope);
			const scope: MemoryScope | "all" = isMemoryScope(scopeArg) ? scopeArg : "all";
			let canonical = loadCanonical(ctx, scope);
			const typeArg = normalizeEnumArg(params.type);
			if (isMemoryType(typeArg)) {
				canonical = canonical.filter((m) => m.type === typeArg);
			}
			const hits = searchMemories(canonical, params.query);

			const includeForeign = params.includeForeign !== false && shouldReadForeign();
			const { cwd, home } = dirs(ctx);
			const foreignHits = includeForeign
				? searchForeign(discoverForeignSources(cwd, home), params.query)
				: [];

			if (hits.length === 0 && foreignHits.length === 0) {
				return text(`No memory matches "${params.query}".`);
			}
			const lines: string[] = [];
			if (hits.length > 0) {
				lines.push(`Canonical matches (${hits.length}):`);
				for (const h of hits) {
					lines.push(`- [${h.memory.scope}/${h.memory.type}] ${h.memory.name} — ${h.snippet}`);
				}
			}
			if (foreignHits.length > 0) {
				if (lines.length) lines.push("");
				lines.push(`Foreign matches (${foreignHits.length}):`);
				for (const f of foreignHits) {
					lines.push(`- [${f.source.system}/${f.source.scope}] ${f.source.path} — ${f.snippet}`);
				}
			}
			return text(lines.join("\n"));
		},
	});

	pi.registerTool({
		name: "memory_read",
		label: "Memory Read",
		description:
			"Read persistent memory. With `name`: return that canonical memory's full body — " +
			"or, if the name matches a foreign file (e.g. CLAUDE.md, MEMORY.md, project_x), that " +
			"file verbatim. With `system`: dump that agent system's memory files verbatim. With " +
			"neither: return a map of all available memory (canonical index + foreign source list).",
		promptSnippet: "Read a stored memory or list what is available",
		parameters: Type.Object({
			name: Type.Optional(
				Type.String({
					description:
						"Name of a canonical memory, or a foreign memory file name like " +
						"'MEMORY.md' / 'CLAUDE.md', to read in full.",
				}),
			),
			system: Type.Optional(foreignSystemSchema),
			scope: Type.Optional(scopeFilterSchema),
		}),
		async execute(_id, params, _signal, _onUpdate, ctx) {
			const { cwd, home } = dirs(ctx);
			const scopeArg = normalizeEnumArg(params.scope);
			const scope: MemoryScope | "all" = isMemoryScope(scopeArg) ? scopeArg : "all";

			if (params.name) {
				const found = findByName(loadCanonical(ctx, scope), params.name);
				if (found) return text(formatMemory(found));
				// Fall back to a foreign file of the same name (CLAUDE.md, MEMORY.md, …).
				if (shouldReadForeign()) {
					const foreign = findForeignByName(discoverForeignSources(cwd, home), params.name);
					if (foreign) return text(formatForeignFull([foreign]));
				}
				return text(`No memory named "${params.name}" (canonical or foreign).`);
			}

			const systemArg = normalizeEnumArg(params.system);
			if (systemArg) {
				if (!isForeignSystem(systemArg)) {
					return text(
						`Unknown system "${params.system}". Valid systems: ${FOREIGN_SYSTEMS.join(", ")}.`,
					);
				}
				if (!shouldReadForeign()) {
					return text("Foreign memory reading is disabled (memory.readForeign = false).");
				}
				const sources = discoverForeignSources(cwd, home, [systemArg]);
				if (sources.length === 0) {
					return text(`No ${systemArg} memory found for this project.`);
				}
				return text(formatForeignFull(sources));
			}

			// The map: canonical index + foreign digest.
			const index = buildIndex(loadCanonical(ctx, scope));
			const parts: string[] = [];
			parts.push(index ? `Canonical memory:\n${index}` : "Canonical memory: (empty)");
			if (shouldReadForeign()) {
				const digest = formatForeignDigest(discoverForeignSources(cwd, home));
				parts.push(digest ? `\nForeign memory sources:\n${digest}` : "\nForeign memory sources: (none found)");
			}
			return text(parts.join("\n"));
		},
	});

	pi.registerTool({
		name: "memory_write",
		label: "Memory Write",
		description:
			"Create or update a canonical memory. Writing a name that matches an existing " +
			"memory updates it. Saved to the project store by default; pass scope='user' for " +
			"a machine-global memory. Use the typed taxonomy: user, feedback, project, reference.",
		promptSnippet: "Save a durable memory to the canonical store",
		promptGuidelines: [
			"Save a memory when you learn something durable and reusable across sessions: a " +
				"user preference or correction (type 'feedback'), facts about the user (type " +
				"'user'), ongoing project context (type 'project'), or a pointer to an external " +
				"system (type 'reference'). Keep `description` a one-line hook; put detail in `content`.",
			"Prefer updating an existing memory (same name) over creating a near-duplicate.",
		],
		parameters: Type.Object({
			name: Type.String({ description: "Short human-readable title; also the memory's identity." }),
			description: Type.String({ description: "One-line summary used in the index and search." }),
			type: memoryTypeSchema,
			content: Type.String({ description: "The memory body (markdown)." }),
			scope: Type.Optional(writeScopeSchema),
		}),
		async execute(_id, params, _signal, _onUpdate, ctx) {
			const { cwd, home } = dirs(ctx);
			const scope: MemoryScope = isMemoryScope(normalizeEnumArg(params.scope))
				? (normalizeEnumArg(params.scope) as MemoryScope)
				: "project";
			const typeArg = normalizeEnumArg(params.type);
			if (!isMemoryType(typeArg)) {
				return text(`Invalid type "${params.type}". Valid types: ${MEMORY_TYPES.join(", ")}.`);
			}
			const type: MemoryType = typeArg;
			const dir = storeDirForScope(scope, cwd, home);
			const result = writeMemory(dir, {
				name: params.name,
				type,
				description: params.description,
				body: params.content,
			});
			const verb = result.created ? "Created" : "Updated";
			return {
				...text(`${verb} ${scope} memory "${result.memory.name}" (${type}) at ${result.memory.path}`),
				details: { scope, ...result },
			};
		},
	});

	pi.registerTool({
		name: "memory_forget",
		label: "Memory Forget",
		description:
			"Delete a canonical memory by name. Only the canonical store is affected — foreign " +
			"agent memory is never modified. Defaults to the project scope.",
		promptSnippet: "Delete a stored memory",
		parameters: Type.Object({
			name: Type.String({ description: "Name of the canonical memory to delete." }),
			scope: Type.Optional(writeScopeSchema),
		}),
		async execute(_id, params, _signal, _onUpdate, ctx) {
			const { cwd, home } = dirs(ctx);
			const scopeArg = normalizeEnumArg(params.scope);
			const scope: MemoryScope = isMemoryScope(scopeArg) ? scopeArg : "project";
			const removed = deleteMemory(storeDirForScope(scope, cwd, home), params.name);
			return text(
				removed
					? `Forgot ${scope} memory "${params.name}".`
					: `No ${scope} memory named "${params.name}" to forget.`,
			);
		},
	});

	// ── Slash commands ─────────────────────────────────────────────────────

	pi.registerCommand("memory-list", {
		description: "List canonical memories and discovered foreign memory sources",
		handler: async (_args, ctx) => {
			const { cwd, home } = dirs(ctx);
			const index = buildIndex(loadCanonical(ctx));
			const lines = [index ? `Canonical memory:\n${index}` : "Canonical memory: (empty)"];
			if (shouldReadForeign()) {
				const digest = formatForeignDigest(discoverForeignSources(cwd, home));
				lines.push(digest ? `\nForeign sources:\n${digest}` : "\nForeign sources: (none found)");
			}
			ctx.ui.notify(lines.join("\n"), "info");
		},
	});

	pi.registerCommand("memory-search", {
		description: "Search memory by keyword (usage: /memory-search <query>)",
		handler: async (args, ctx) => {
			const query = args.trim();
			if (!query) {
				ctx.ui.notify("Usage: /memory-search <query>", "error");
				return;
			}
			const { cwd, home } = dirs(ctx);
			const hits = searchMemories(loadCanonical(ctx), query);
			const foreignHits = shouldReadForeign()
				? searchForeign(discoverForeignSources(cwd, home), query)
				: [];
			if (hits.length === 0 && foreignHits.length === 0) {
				ctx.ui.notify(`No memory matches "${query}".`, "info");
				return;
			}
			const lines: string[] = [];
			for (const h of hits) {
				lines.push(`[${h.memory.scope}/${h.memory.type}] ${h.memory.name} — ${h.snippet}`);
			}
			for (const f of foreignHits) {
				lines.push(`[${f.source.system}/${f.source.scope}] ${f.source.path} — ${f.snippet}`);
			}
			ctx.ui.notify(lines.join("\n"), "info");
		},
	});

	pi.registerCommand("memory-show", {
		description: "Show a canonical memory in full (usage: /memory-show <name>)",
		getArgumentCompletions: completeMemoryName,
		handler: async (args, ctx) => {
			const name = args.trim();
			if (!name) {
				ctx.ui.notify("Usage: /memory-show <name>", "error");
				return;
			}
			const found = findByName(loadCanonical(ctx), name);
			ctx.ui.notify(found ? formatMemory(found) : `No canonical memory named "${name}".`, found ? "info" : "error");
		},
	});

	pi.registerCommand("memory-write", {
		description: "Interactively create or update a canonical memory",
		handler: async (_args, ctx) => {
			const scope = await ctx.ui.select("Memory scope", [...MEMORY_SCOPES]);
			if (!scope || !isMemoryScope(scope)) return;
			const type = await ctx.ui.select("Memory type", [...MEMORY_TYPES]);
			if (!type || !isMemoryType(type)) return;
			const name = await ctx.ui.input("Memory name", "short title");
			if (!name || !name.trim()) return;
			const description = await ctx.ui.input("One-line description");
			if (description === undefined) return;
			const body = await ctx.ui.editor(`Memory body — ${name.trim()}`);
			if (body === undefined) return;
			const { cwd, home } = dirs(ctx);
			try {
				const result = writeMemory(storeDirForScope(scope, cwd, home), {
					name,
					type,
					description: description ?? "",
					body,
				});
				ctx.ui.notify(
					`${result.created ? "Created" : "Updated"} ${scope} memory "${result.memory.name}"`,
					"info",
				);
			} catch (err) {
				ctx.ui.notify(`Write failed: ${(err as Error).message}`, "error");
			}
		},
	});

	pi.registerCommand("memory-forget", {
		description: "Delete a canonical memory (usage: /memory-forget <name> [project|user])",
		getArgumentCompletions: completeMemoryName,
		handler: async (args, ctx) => {
			const parts = args.trim().split(/\s+/).filter(Boolean);
			const name = parts[0];
			if (!name) {
				ctx.ui.notify("Usage: /memory-forget <name> [project|user]", "error");
				return;
			}
			const scopeArg = normalizeEnumArg(parts[1]);
			const scope: MemoryScope = isMemoryScope(scopeArg) ? scopeArg : "project";
			const ok = await ctx.ui.confirm(
				"Forget memory?",
				`Delete ${scope} memory "${name}"? This removes the file.`,
			);
			if (!ok) return;
			const { cwd, home } = dirs(ctx);
			const removed = deleteMemory(storeDirForScope(scope, cwd, home), name);
			ctx.ui.notify(
				removed ? `Forgot ${scope} memory "${name}".` : `No ${scope} memory named "${name}".`,
				removed ? "info" : "warning",
			);
		},
	});
}