kimiko-agent/extensions/memory/index.ts

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);
    },
  });
}