pi-memory-stone 0.1.0

package/src/retrieval/index.ts

@@ -0,0 +1,219 @@
+/**
+ * Retrieval and ranking module.
+ * Hybrid ranking: FTS score + same-project boost + recency decay + kind boost.
+ */
+
+import type { RecordRow } from "../db/index.js";
+import type { RecordKind, RecordScope } from "../db/schema.js";
+import { searchRecordsFts } from "../db/index.js";
+import { getConfig } from "../config/index.js";
+
+// ─── Kind boosts ────────────────────────────────────────────────────
+
+const KIND_BOOST: Record<string, number> = {
+  decision: 1.5,
+  preference: 1.3,
+  error_resolution: 1.4,
+  task: 1.1,
+  turn_summary: 0.9,
+  session_summary: 0.8,
+  file_activity: 0.5,
+};
+
+// ─── Recency decay ──────────────────────────────────────────────────
+
+const RECENCY_HALF_LIFE_MS = 7 * 24 * 60 * 60 * 1000; // 7 days
+
+function recencyDecay(createdAt: number): number {
+  const age = Date.now() - createdAt;
+  return Math.exp(-Math.log(2) * (age / RECENCY_HALF_LIFE_MS));
+}
+
+// ─── Query builder ──────────────────────────────────────────────────
+
+export function buildSearchQuery(
+  userPrompt: string,
+  recentFiles: string[] = [],
+): string {
+  const parts: string[] = [];
+
+  // User prompt terms (take first ~200 chars)
+  parts.push(userPrompt.slice(0, 200));
+
+  // Recent files (just filenames, not full paths)
+  for (const f of recentFiles.slice(0, 5)) {
+    const basename = f.split("/").pop() || f;
+    parts.push(basename);
+  }
+
+  return parts.join(" ");
+}
+
+// ─── Ranking ────────────────────────────────────────────────────────
+
+export interface RankedResult {
+  record: RecordRow;
+  score: number;
+  reasons: string[];
+}
+
+export function rankAndFilter(
+  records: (RecordRow & { rank: number })[],
+  currentProjectId: string | null,
+  crossProjectEnabled: boolean,
+): RankedResult[] {
+  const results: RankedResult[] = [];
+
+  for (const rec of records) {
+    // Skip non-active records
+    if (rec.status !== "active") continue;
+
+    // Cross-project/global filter. Global records are cross-project by definition
+    // and require explicit cross-project retrieval.
+    if (rec.scope === "global") {
+      if (!crossProjectEnabled) continue;
+    } else if (rec.project_id && currentProjectId && rec.project_id !== currentProjectId) {
+      continue;
+    }
+
+    // Compute hybrid score
+    let score = 1.0 / (1.0 + (rec.rank ?? 0)); // Normalize FTS rank: lower rank = better
+
+    // Same project boost
+    const reasons: string[] = [];
+    if (rec.project_id && currentProjectId && rec.project_id === currentProjectId) {
+      score *= 1.5;
+      reasons.push("same-project");
+    }
+
+    // Global preference boost
+    if (rec.scope === "global") {
+      score *= 1.2;
+      reasons.push("global-preference");
+    }
+
+    // Kind boost
+    const kindBoost = KIND_BOOST[rec.kind] ?? 1.0;
+    score *= kindBoost;
+    if (kindBoost !== 1.0) reasons.push(`kind:${rec.kind}`);
+
+    // Recency decay
+    const decay = recencyDecay(rec.created_at);
+    score *= decay;
+
+    // Confidence multiplier
+    score *= rec.confidence;
+
+    // Importance multiplier
+    score *= 0.5 + rec.importance; // Scale: 0.5-1.5
+
+    results.push({ record: rec, score, reasons });
+  }
+
+  // Sort by score descending, then recency
+  results.sort((a, b) => {
+    if (b.score !== a.score) return b.score - a.score;
+    return b.record.created_at - a.record.created_at;
+  });
+
+  return results;
+}
+
+// ─── Full retrieval pipeline ────────────────────────────────────────
+
+export function retrieve(
+  userPrompt: string,
+  currentProjectId: string | null,
+  recentFiles: string[] = [],
+  opts?: {
+    limit?: number;
+    crossProjectEnabled?: boolean;
+    kindFilter?: RecordKind[];
+    scopeFilter?: RecordScope[];
+  },
+): RankedResult[] {
+  const config = getConfig();
+  const limit = opts?.limit ?? config.maxInjectedRecords;
+  const crossProject = opts?.crossProjectEnabled ?? config.crossProjectEnabled;
+
+  const query = buildSearchQuery(userPrompt, recentFiles);
+
+  // Get more candidates than needed (ranking will filter)
+  const candidates = searchRecordsFts(query, limit * 10, opts?.kindFilter, opts?.scopeFilter);
+
+  const ranked = rankAndFilter(candidates, currentProjectId, crossProject);
+
+  // Return top results
+  return ranked.slice(0, limit);
+}
+
+// ─── Injection packet builder ────────────────────────────────────────
+
+export interface InjectionPacket {
+  header: string;
+  items: Array<{
+    ref: string;
+    kind: string;
+    text: string;
+    project?: string;
+    staleHint?: string;
+  }>;
+  footer: string;
+  recordCount: number;
+}
+
+export function buildInjectionPacket(results: RankedResult[]): InjectionPacket {
+  const items: InjectionPacket["items"] = [];
+
+  for (const r of results) {
+    const age = Date.now() - r.record.created_at;
+    const ageDays = Math.floor(age / (24 * 60 * 60 * 1000));
+    let staleHint: string | undefined;
+
+    if (ageDays > 30) {
+      staleHint = `⚠️ This memory is ${ageDays} days old and may be stale.`;
+    } else if (ageDays > 7) {
+      staleHint = `📅 From ${ageDays} days ago.`;
+    }
+
+    items.push({
+      ref: r.record.id,
+      kind: r.record.kind,
+      text: r.record.text,
+      project: r.record.project_id === null ? "global" : undefined,
+      staleHint,
+    });
+  }
+
+  return {
+    header: `Memory: loaded ${items.length} relevant items from past sessions. These may help provide context.`,
+    items,
+    footer:
+      "Use memory_open with a ref for full text. If any memory is irrelevant or conflicting, mention it and I can /memory-forget <ref> it.",
+    recordCount: items.length,
+  };
+}
+
+export function formatInjectionForLlm(packet: InjectionPacket, maxTokens = 1000): string {
+  const lines: string[] = [packet.header, ""];
+
+  // Rough token estimate: ~4 chars per token
+  let charBudget = maxTokens * 4;
+  let usedChars = lines.join("\n").length;
+
+  for (const item of packet.items) {
+    const itemHeader = `[${item.kind} ref=${item.ref}]`;
+    const itemText = item.text.slice(0, 300); // Truncate per item
+    const itemFooter = item.staleHint ? `  ${item.staleHint}` : "";
+    const itemLine = `${itemHeader} ${itemText}${itemFooter}`;
+
+    if (usedChars + itemLine.length > charBudget) break;
+    lines.push(itemLine);
+    usedChars += itemLine.length + 1;
+  }
+
+  lines.push("");
+  lines.push(packet.footer);
+
+  return lines.join("\n");
+}

package/src/tools/index.ts

@@ -0,0 +1,257 @@
+/**
+ * Tools: memory_search, memory_open, memory_remember, memory_forget
+ */
+
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import { Type } from "typebox";
+import { StringEnum } from "@earendil-works/pi-ai";
+import { getRecord, softForgetRecord, upsertRecord } from "../db/index.js";
+import { retrieve, buildInjectionPacket, formatInjectionForLlm } from "../retrieval/index.js";
+import { getProjectId, getConfig } from "../config/index.js";
+import type { RecordKind, RecordScope } from "../db/schema.js";
+
+export function registerTools(pi: ExtensionAPI): void {
+  // ── memory_search ───────────────────────────────────────────────
+
+  pi.registerTool({
+    name: "memory_search",
+    label: "Search Memory",
+    description:
+      "Search your memory stone for relevant records from past pi sessions. Use this to recall past decisions, preferences, tasks, or error resolutions.",
+    promptSnippet: "Search memory stone by query, with optional kind/scope/limit filters",
+    promptGuidelines: [
+      "Use memory_search to recall relevant context from past sessions before making decisions.",
+      "Set kind to filter by record type: decision, preference, task, error_resolution, turn_summary.",
+      "Set scope to 'global' for cross-project memories, or omit for current project only.",
+    ],
+    parameters: Type.Object({
+      query: Type.String({ description: "Search query text" }),
+      kind: Type.Optional(
+        StringEnum([
+          "decision",
+          "preference",
+          "task",
+          "error_resolution",
+          "turn_summary",
+          "session_summary",
+        ] as const),
+      ),
+      scope: Type.Optional(StringEnum(["project", "global"] as const)),
+      limit: Type.Optional(Type.Number({ description: "Max results (default 5)" })),
+    }),
+    async execute(toolCallId, params, _signal, _onUpdate, ctx) {
+      const projectId = getProjectId(ctx.cwd);
+      const config = getConfig(ctx.cwd);
+      const limit = params.limit ?? 5;
+
+      const results = retrieve(params.query, projectId, [], {
+        limit,
+        crossProjectEnabled: params.scope === "global" || config.crossProjectEnabled,
+        kindFilter: params.kind ? [params.kind as RecordKind] : undefined,
+        scopeFilter: params.scope ? [params.scope as RecordScope] : undefined,
+      });
+
+      if (results.length === 0) {
+        return {
+          content: [{ type: "text", text: "No matching memories found." }],
+          details: { query: params.query, results: [] },
+        };
+      }
+
+      const packet = buildInjectionPacket(results);
+      const formatted = formatInjectionForLlm(packet, 2000);
+
+      return {
+        content: [{ type: "text", text: formatted }],
+        details: {
+          query: params.query,
+          results: results.map((r) => ({
+            id: r.record.id,
+            kind: r.record.kind,
+            score: r.score,
+            text: r.record.text.slice(0, 200),
+            reasons: r.reasons,
+          })),
+        },
+      };
+    },
+  });
+
+  // ── memory_open ─────────────────────────────────────────────────
+
+  pi.registerTool({
+    name: "memory_open",
+    label: "Open Memory",
+    description:
+      "Open a specific memory record by its reference ID. Returns the full record text and metadata. The 'ref' can be obtained from memory_search results or injection packets.",
+    promptSnippet: "Open specific memory record by ref ID to see full content",
+    promptGuidelines: [
+      "Use memory_open when you need the full text of a specific memory record referenced in an injection packet or search result.",
+    ],
+    parameters: Type.Object({
+      ref: Type.String({ description: "Memory record reference ID" }),
+    }),
+    async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
+      const record = getRecord(params.ref);
+
+      if (!record) {
+        return {
+          content: [{ type: "text", text: `Memory record ${params.ref} not found.` }],
+          details: { ref: params.ref, found: false },
+        };
+      }
+
+      const currentProjectId = getProjectId(ctx.cwd);
+      const visibleInCurrentProject =
+        record.scope === "global" || record.project_id === null || record.project_id === currentProjectId;
+
+      if (record.status !== "active" || !visibleInCurrentProject) {
+        return {
+          content: [{ type: "text", text: `Memory record ${params.ref} is not available.` }],
+          details: { ref: params.ref, found: false, unavailable: true },
+        };
+      }
+
+      // Redacted excerpt (already redacted at storage time)
+      const lines: string[] = [];
+      lines.push(`Memory Record: ${record.id}`);
+      lines.push(`Kind: ${record.kind}`);
+      lines.push(`Scope: ${record.scope}`);
+      lines.push(`Project: ${record.project_id ?? "global"}`);
+      lines.push(`Created: ${new Date(record.created_at).toISOString()}`);
+      lines.push(`Status: ${record.status}`);
+      lines.push("");
+      lines.push(record.text);
+
+      return {
+        content: [{ type: "text", text: lines.join("\n") }],
+        details: {
+          ref: params.ref,
+          found: true,
+          kind: record.kind,
+          scope: record.scope,
+          status: record.status,
+          created_at: record.created_at,
+        },
+      };
+    },
+  });
+
+  // ── memory_remember ─────────────────────────────────────────────
+
+  pi.registerTool({
+    name: "memory_remember",
+    label: "Remember",
+    description:
+      "Explicitly store a memory record. Only use when the user explicitly asks you to remember something. Default scope is project. Use scope='global' only when the user explicitly says to remember globally.",
+    promptSnippet: "Store an explicit memory record when user asks to remember something",
+    promptGuidelines: [
+      "Use memory_remember only when the user explicitly asks you to remember a decision, preference, or fact.",
+      "Default to scope='project'. Only use scope='global' when the user explicitly says 'remember this globally' or 'remember for all projects'.",
+    ],
+    parameters: Type.Object({
+      kind: StringEnum([
+        "decision",
+        "preference",
+        "task",
+        "error_resolution",
+        "turn_summary",
+        "session_summary",
+      ] as const),
+      text: Type.String({ description: "Memory text to store" }),
+      scope: Type.Optional(StringEnum(["project", "global"] as const)),
+      tags: Type.Optional(Type.String({ description: "Comma-separated tags" })),
+      importance: Type.Optional(Type.Number({ description: "Importance 0-1 (default 0.5)" })),
+    }),
+    async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
+      const projectId = getProjectId(ctx.cwd);
+      let scope = params.scope ?? "project";
+
+      // Safety: never allow global for implementation details, paths, etc.
+      const isSensitiveForGlobal =
+        /\b(?:password|secret|token|key|\.env|localhost|127\.0\.0\.1|internal|private)\b/i.test(
+          params.text,
+        );
+
+      const downgradedToProject = isSensitiveForGlobal && scope === "global";
+      if (downgradedToProject) {
+        scope = "project";
+      }
+
+      const recordId = upsertRecord({
+        kind: params.kind as RecordKind,
+        scope: scope as RecordScope,
+        project_id: scope === "global" ? null : projectId,
+        text: params.text,
+        tags: params.tags,
+        importance: params.importance ?? 0.5,
+        confidence: 1.0,
+      });
+
+      return {
+        content: [
+          {
+            type: "text",
+            text: downgradedToProject
+              ? `Cannot store as global: text appears to contain sensitive data (secrets, paths, hostnames). Stored as project-scoped instead: [${params.kind}] ${recordId}`
+              : `Memory stored: [${params.kind}] ${recordId} (scope: ${scope})`,
+          },
+        ],
+        details: { id: recordId, kind: params.kind, scope, downgradedToProject },
+      };
+    },
+  });
+
+  // ── memory_forget ───────────────────────────────────────────────
+
+  pi.registerTool({
+    name: "memory_forget",
+    label: "Forget Memory",
+    description:
+      "Soft-forget a memory record by its reference ID. The record is hidden from future searches but can be restored. Use --hard for permanent deletion (requires explicit user confirmation).",
+    promptSnippet: "Forget a memory record by ID (soft or hard)",
+    promptGuidelines: [
+      "Use memory_forget when the user asks to remove or forget a memory reference.",
+      "Default to soft forget. Only use hard=true when the user explicitly asks to permanently delete.",
+    ],
+    parameters: Type.Object({
+      ref: Type.String({ description: "Memory record reference ID" }),
+      hard: Type.Optional(Type.Boolean({ description: "Permanently delete (default: soft forget)" })),
+    }),
+    async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
+      const record = getRecord(params.ref);
+
+      if (!record) {
+        return {
+          content: [{ type: "text", text: `Memory record ${params.ref} not found.` }],
+          details: { ref: params.ref, found: false },
+        };
+      }
+
+      if (params.hard) {
+        // For hard delete via tool, we require the user to explicitly confirm
+        // The tool should note this requires user interaction
+        return {
+          content: [
+            {
+              type: "text",
+              text: `Permanent deletion requires explicit confirmation. Please use /memory-forget ${params.ref} --hard to permanently delete this record.\n\nRecord: [${record.kind}] ${record.text.slice(0, 200)}`,
+            },
+          ],
+          details: { ref: params.ref, requiresConfirmation: true },
+        };
+      }
+
+      softForgetRecord(params.ref);
+      return {
+        content: [
+          {
+            type: "text",
+            text: `Memory record ${params.ref} has been soft-forgotten. It will no longer appear in searches.`,
+          },
+        ],
+        details: { ref: params.ref, forgotten: true, hard: false },
+      };
+    },
+  });
+}

package/test/indexing.test.ts

@@ -0,0 +1,97 @@
+/**
+ * Tests for incremental session indexing.
+ */
+
+import { describe, it, beforeEach, after } from "node:test";
+import assert from "node:assert/strict";
+import { existsSync, unlinkSync, mkdtempSync, rmSync, writeFileSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { indexSessionOnAgentEnd } from "../src/indexing/index.js";
+import { closeDb, getDb, getDbPath } from "../src/db/index.js";
+
+const testMemoryDir = mkdtempSync(join(tmpdir(), "pi-memory-stone-indexing-"));
+process.env.PI_MEMORY_STONE_DB_PATH = join(testMemoryDir, "memory.db");
+
+function cleanDb() {
+  const dbPath = getDbPath();
+  closeDb();
+  for (const f of [dbPath, dbPath + "-wal", dbPath + "-shm"]) {
+    try { if (existsSync(f)) unlinkSync(f); } catch {}
+  }
+}
+
+function makeUserEntry(id: string, text: string) {
+  return {
+    type: "message",
+    id,
+    parentId: null,
+    message: {
+      role: "user",
+      content: text,
+      timestamp: Date.now(),
+    },
+  };
+}
+
+function makeAssistantEntry(id: string, text: string) {
+  return {
+    type: "message",
+    id,
+    parentId: null,
+    message: {
+      role: "assistant",
+      content: [{ type: "text", text }],
+      timestamp: Date.now(),
+    },
+  };
+}
+
+function makeContext(sessionFile: string, branch: unknown[]) {
+  return {
+    cwd: testMemoryDir,
+    sessionManager: {
+      getSessionFile: () => sessionFile,
+      getSessionId: () => "session-1",
+      getBranch: () => branch,
+      getLeafId: () => (branch.at(-1) as { id?: string } | undefined)?.id,
+    },
+  } as any;
+}
+
+describe("indexSessionOnAgentEnd", () => {
+  beforeEach(() => cleanDb());
+
+  after(() => {
+    cleanDb();
+    rmSync(testMemoryDir, { recursive: true, force: true });
+  });
+
+  it("indexes only entries after the last indexed entry when timestamps are absent", async () => {
+    const sessionFile = join(testMemoryDir, "session.jsonl");
+    writeFileSync(sessionFile, "");
+
+    const firstBranch = [
+      makeUserEntry("001", "First question"),
+      makeAssistantEntry("002", "First answer"),
+    ];
+    const first = await indexSessionOnAgentEnd(makeContext(sessionFile, firstBranch), {});
+    assert.equal(first.errors.length, 0);
+
+    const fullBranch = [
+      ...firstBranch,
+      makeUserEntry("003", "Second question"),
+      makeAssistantEntry("004", "Second answer"),
+    ];
+    const second = await indexSessionOnAgentEnd(makeContext(sessionFile, fullBranch), {});
+    assert.equal(second.errors.length, 0);
+
+    const rows = getDb()
+      .prepare("SELECT text FROM records WHERE kind = 'turn_summary' ORDER BY created_at")
+      .all() as Array<{ text: string }>;
+
+    assert.equal(rows.length, 2);
+    assert.ok(rows.some((r) => r.text.includes("First question")));
+    assert.ok(rows.some((r) => r.text.includes("Second question")));
+  });
+});