import { fileURLToPath } from "node:url"; import { dirname, join, resolve } from "node:path"; import type { ExtensionAPI } from "@earendil-works/pi-coding-agent"; import { Type } from "typebox"; import { MemoryTurnAccess } from "./access.mjs"; import { MemoryStore } from "./store.mjs"; const DETAIL_TOOLS = ["memory_read", "memory_write", "memory_forget"]; const EXPLICIT_MEMORY_INTENT = /\b(?:forget|forgot|memories|memory|recall|remember|remembered)\b|\b(?:last time|previous conversation)\b/i; const extensionInstanceRoot = resolve(dirname(fileURLToPath(import.meta.url)), "..", ".."); const memoryRoot = join(process.env.PI_CODING_AGENT_DIR ?? extensionInstanceRoot, "memory"); const store = new MemoryStore(memoryRoot); function textResult(text: string, details: unknown = {}) { return { content: [{ type: "text" as const, text }], details }; } export default function memoryExtension(pi: ExtensionAPI) { const turnAccess = new MemoryTurnAccess(); let preparedInput: string | undefined; function resetTurnAccess() { turnAccess.reset(); const compactTools = pi.getActiveTools().filter((name) => !DETAIL_TOOLS.includes(name)); pi.setActiveTools([...new Set([...compactTools, "memory_search"])]); } function enableDetailTools() { pi.setActiveTools([...new Set([...pi.getActiveTools(), ...DETAIL_TOOLS])]); } function prepareTurn(text: string) { resetTurnAccess(); // Explicit memory requests reveal the gated tools early so the model knows the // capability exists. The workflow gate still requires search/read in order. if (EXPLICIT_MEMORY_INTENT.test(text)) enableDetailTools(); } pi.registerTool({ name: "memory_search", label: "Search Memory", description: "Required first step for every recall, remember, update, or forget request. Search durable memory by keywords; executing this tool enables detailed memory tools even when no matches are found.", promptSnippet: "First step for every recall, remember, update, or forget request; enables detailed memory tools", promptGuidelines: [ "Always call memory_search first for recall, remember, update, or forget requests; it enables the detailed memory tools even when no match exists.", ], parameters: Type.Object({ query: Type.String({ description: "Keywords for the person, project, preference, decision, or fact" }), limit: Type.Optional(Type.Integer({ minimum: 1, maximum: 20, description: "Maximum candidates (default 8)" })), }), async execute(_toolCallId, params) { enableDetailTools(); const result = await store.search(params.query, params.limit ?? 8); turnAccess.recordSearch(result.results.map((item) => item.id)); if (!result.results.length) { return textResult(`No memory candidates found for: ${params.query}\nDetailed memory tools are now enabled.`, result); } const lines = result.results.map((item) => { const tags = item.tags.length ? ` [${item.tags.join(", ")}]` : ""; const description = item.description ? ` — ${item.description}` : ""; return `- ${item.id} | ${item.type} | ${item.timestamp ?? "timestamp unavailable"}${tags}${description}`; }); if (result.warnings.length) lines.push(`\nWarnings: ${result.warnings.length} invalid concept(s); use /memory-check --repair.`); return textResult(lines.join("\n"), result); }, }); pi.registerTool({ name: "memory_read", label: "Read Memory", description: "Load full bodies for up to 10 concept IDs returned by memory_search in the current user turn. Output is capped at 45,000 characters.", parameters: Type.Object({ ids: Type.Array(Type.String({ description: "Bundle-relative concept ID without .md" }), { minItems: 1, maxItems: 10 }), }), async execute(_toolCallId, params) { const normalizedIds = params.ids.map((id) => store.normalizeId(id)); for (const id of normalizedIds) turnAccess.requireSearchedId(id); const result = await store.read(normalizedIds); turnAccess.recordRead(result.documents.map((document) => document.id)); return textResult(result.text, result); }, }); pi.registerTool({ name: "memory_write", label: "Write Memory", description: "Create or replace one atomic durable concept after searching in the current user turn; existing concepts must also be read first. Supply semantic fields and Markdown body only; storage mechanics are deterministic.", parameters: Type.Object({ id: Type.String({ description: "Stable lowercase bundle-relative ID, e.g. people/alice; no .md" }), type: Type.String({ description: "Short descriptive type, e.g. Person, Preference, Project, Decision" }), title: Type.Optional(Type.String({ description: "Human-readable title" })), description: Type.Optional(Type.String({ description: "One-sentence retrieval summary" })), tags: Type.Optional(Type.Array(Type.String(), { maxItems: 20 })), body: Type.String({ description: "Structured Markdown body without YAML frontmatter or a timestamp" }), }), async execute(_toolCallId, params) { const id = store.normalizeId(params.id); turnAccess.requireWrite(id, await store.exists(id)); const result = await store.upsert({ ...params, id }); return textResult(`${result.action}: ${result.id} (${result.timestamp})${result.attempts > 1 ? " after verified retry" : ""}`, result); }, }); pi.registerTool({ name: "memory_forget", label: "Forget Memory", description: "Permanently delete one durable concept after searching and reading it in the current user turn. Confirmation must exactly match the concept ID.", parameters: Type.Object({ id: Type.String({ description: "Concept ID to delete" }), confirm: Type.String({ description: "Repeat the exact concept ID to confirm permanent deletion" }), }), async execute(_toolCallId, params) { const id = store.normalizeId(params.id); turnAccess.requireForget(id); const result = await store.forget({ ...params, id }); return textResult(`deleted: ${result.id}`, result); }, }); pi.on("input", (event) => { // Do not revoke access while an earlier turn is still executing. Queued or // steering input is prepared when its own before_agent_start event arrives. if (event.streamingBehavior !== undefined) return; prepareTurn(event.text); preparedInput = event.text; }); pi.on("before_agent_start", async (event) => { if (preparedInput !== event.prompt) prepareTurn(event.prompt); preparedInput = undefined; try { const hints = await store.hints(event.prompt, 2); if (hints.length === 0) return; const lines = hints.map((hint) => { const summary = (hint.description || hint.title).replace(/\s+/g, " ").slice(0, 180); return `- ${hint.id} — ${summary}`; }); return { // This is rebuilt for each user turn and is never appended to session history. systemPrompt: `${event.systemPrompt}\n\nEphemeral long-term memory candidates (metadata only):\n${lines.join("\n")}\nTreat these only as retrieval hints. If relevant, call memory_search and then memory_read; otherwise ignore them.`, }; } catch (error) { console.warn(`Memory hint retrieval failed: ${error instanceof Error ? error.message : String(error)}`); } }); pi.on("session_start", async (_event, ctx) => { try { const report = await store.initialize(); resetTurnAccess(); if (report.repaired > 0 && ctx.hasUI) { ctx.ui.notify(`Memory repaired ${report.repaired} concept(s); run /memory-check for details.`, "warning"); } } catch (error) { const message = error instanceof Error ? error.message : String(error); console.warn(`Memory initialization failed: ${message}`); if (ctx.hasUI) ctx.ui.notify(`Memory initialization failed: ${message}`, "error"); } }); pi.registerCommand("memory-check", { description: "Validate memory; pass --repair for deterministic repair and regenerated indexes", handler: async (args, ctx) => { try { const report = await store.audit({ repair: args.trim() === "--repair" }); const summary = `${report.valid ? "valid" : "invalid"}: ${report.concepts} concepts, ${report.issues.length} issues, ${report.repaired} repaired`; if (ctx.hasUI) ctx.ui.notify(summary, report.valid ? "info" : "warning"); else console.log(summary); } catch (error) { const message = error instanceof Error ? error.message : String(error); if (ctx.hasUI) ctx.ui.notify(message, "error"); else console.error(message); } }, }); pi.registerCommand("memory-path", { description: "Show the OKF memory bundle path", handler: async (_args, ctx) => { if (ctx.hasUI) ctx.ui.notify(memoryRoot, "info"); else console.log(memoryRoot); }, }); }