my-pi/pi-hermes-memory/src/index.ts

/**
 * Pi Hermes Memory Extension
 *
 * Brings Hermes-style persistent memory and a learning loop to any Pi user.
 * After `pi install`, users get:
 *
 * 1. Persistent Memory — MEMORY.md + USER.md that survive across sessions
 * 2. Background Learning Loop — auto-saves notable facts every N turns
 * 3. Session-End Flush — saves memories before compaction/shutdown
 * 4. Auto-Consolidation — merges memory when full instead of erroring
 * 5. Correction Detection — immediate save on user corrections
 * 6. Procedural Skills — SKILL.md files for reusable procedures
 * 7. Tool-Call-Aware Nudge — review triggers on tool call count too
 * 8. /memory-insights — shows what's stored
 * 9. /memory-skills — lists procedural skills
 * 10. /memory-consolidate — manual consolidation trigger
 * 11. /memory-interview — onboarding interview to pre-fill user profile
 * 12. /memory-switch-project — list project memories
 * 13. Context Fencing — <memory-context> tags prevent injection through stored memory
 * 14. Memory Aging — entry timestamps guide consolidation
 *
 * See docs/ROADMAP.md for full roadmap and Hermes competitive analysis.
 */

import * as path from "node:path";
import * as fs from "node:fs";
import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
import { MemoryStore } from "./store/memory-store.js";
import { SkillStore } from "./store/skill-store.js";
import { DatabaseManager } from "./store/db.js";
import { indexSession, upsertSessionFileMetadata, pruneEphemeralReviewSessions, pruneOldSessions, retentionCutoffMs } from "./store/session-indexer.js";
import { runRecoveryMaintenance } from "./store/recovery-maintenance.js";
import { scheduleSessionBackfill, joinSessionBackfill, waitForSessionBackfill, SESSION_BACKFILL_SHUTDOWN_TIMEOUT_MS, type SessionBackfillState } from "./handlers/session-backfill.js";
import { scheduleLiveSessionIndex, waitForLiveSessionIndex, SESSION_LIVE_INDEX_SHUTDOWN_TIMEOUT_MS } from "./handlers/session-live-index.js";
import { parseSessionFile } from "./store/session-parser.js";
import { registerMemoryTool } from "./tools/memory-tool.js";
import { registerSkillTool } from "./tools/skill-tool.js";
import { registerSessionSearchTool } from "./tools/session-search-tool.js";
import { registerMemorySearchTool } from "./tools/memory-search-tool.js";
import { setupBackgroundReview } from "./handlers/background-review.js";
import { setupSessionFlush } from "./handlers/session-flush.js";
import { registerInsightsCommand } from "./handlers/insights.js";
import { triggerConsolidation, registerConsolidateCommand } from "./handlers/auto-consolidate.js";
import { setupCorrectionDetector } from "./handlers/correction-detector.js";
import { registerSkillsCommand } from "./handlers/skills-command.js";
import { registerInterviewCommand } from "./handlers/interview.js";
import { registerSwitchProjectCommand } from "./handlers/switch-project.js";
import { registerIndexSessionsCommand } from "./handlers/index-sessions.js";
import { registerLearnMemoryCommand } from "./handlers/learn-memory.js";
import { migrateThenSyncMarkdownMemories, registerSyncMarkdownMemoriesCommand } from "./handlers/sync-markdown-memories.js";
import { registerPreviewContextCommand } from "./handlers/preview-context.js";
import { registerStandingPinCommand } from "./handlers/standing-pin.js";
import { StandingInstructions } from "./store/standing-instructions.js";
import { STANDING_FILE } from "./constants.js";
import { loadConfig } from "./config.js";
import { shouldWarnAutoConsolidationFailure } from "./auto-consolidation-warning.js";
import { detectProject, detectProjectSkills } from "./project.js";
import { buildPromptContext } from "./prompt-context.js";
import { migrateLegacyProjectMemoryDirs } from "./project-memory-migration.js";
import { AGENT_ROOT } from "./paths.js";
import { isDatabaseMigrationPending } from "./extension-root-migration.js";
import { measureLifecycle, measureLifecycleSync } from "./lifecycle-timing.js";
import { createMemoryInitializer, withMemoryInitialization, type EnsureMemoryReady } from "./memory-initialization.js";
import { isSubagentSession, markSubagentSession, markSubagentSessionFromPrompt } from "./subagent-mode.js";

export function resolveProjectSkillDiscovery(
  skillStore: SkillStore,
  projectsMemoryDir: string | undefined,
  cwd?: string,
): { skillPaths: string[] } {
  const detected = detectProjectSkills(projectsMemoryDir, cwd);
  skillStore.setProjectContext(detected.name, detected.skillsDir);

  // Pi auto-discovers its own `~/.pi/agent/skills/`, but this extension keeps
  // its generated skills in a directory of its own so users can audit, wipe, or
  // ignore them without touching skills they installed themselves (#126). Both
  // of ours must therefore be contributed here.
  const skillPaths = [skillStore.getGlobalSkillsDir()];
  if (detected.skillsDir) skillPaths.push(detected.skillsDir);
  return { skillPaths };
}

export function registerProjectSkillDiscoveryHandler(
  pi: Pick<ExtensionAPI, "on">,
  skillStore: SkillStore,
  projectsMemoryDir: string | undefined,
): void {
  pi.on("resources_discover", async (event, _ctx) => {
    return resolveProjectSkillDiscovery(skillStore, projectsMemoryDir, (event as { cwd?: string }).cwd);
  });
}

export default function (pi: ExtensionAPI) {
  const config = loadConfig();
  const lazy = config.lazyInitialization === true && config.memoryMode === "policy-only";
  let sessionContext: ExtensionContext | undefined;

  const agentRoot = AGENT_ROOT;
  const legacyGlobalDir = path.join(agentRoot, "memory");
  const defaultGlobalDir = path.join(agentRoot, "pi-hermes-memory");

  const configuredMemoryDir = config.memoryDir?.trim();
  const pointsToLegacyMemoryDir = configuredMemoryDir
    ? path.resolve(configuredMemoryDir) === path.resolve(legacyGlobalDir)
    : false;

  const globalDir = !configuredMemoryDir || pointsToLegacyMemoryDir
    ? defaultGlobalDir
    : configuredMemoryDir;

  const shouldMigrateExtensionRoot = !configuredMemoryDir || pointsToLegacyMemoryDir;
  let persistenceInitialized = false;

  const store = new MemoryStore({ ...config, memoryDir: globalDir });
  // Factory may run with no session (Pi public contract). Do not snapshot
  // project identity from process.cwd() here — bind from session_start ctx.cwd
  // and from tool execute ctx.cwd.
  let projectName = "";
  const skillStore = new SkillStore({
    globalSkillsDir: path.join(globalDir, "skills"),
    piGlobalSkillsDir: path.join(agentRoot, "skills"),
    projectSkillsDir: null,
    projectName: null,
    legacySkillsDir: path.join(legacyGlobalDir, "skills"),
    migrationSentinelPath: path.join(globalDir, ".skills-migrated-to-extension-storage"),
  });
  const dbManager = new DatabaseManager(globalDir);
  let databaseClosed = false;
  const backfillState: SessionBackfillState = { inProgress: false, promise: null };
  dbManager.setQuickCheckOnOpen(config.quickCheckOnOpen ?? true);
  // No database may open before first-use migration has completed.
  let databaseMigrationPending = (lazy && shouldMigrateExtensionRoot) || (shouldMigrateExtensionRoot
    && isDatabaseMigrationPending(legacyGlobalDir, globalDir));
  dbManager.setOpenGuard(() => {
    if (databaseClosed) throw new Error("Memory session has shut down");
    if (databaseMigrationPending) throw new Error("Legacy sessions.db migration is pending");
  });
  const sessionsDir = path.join(agentRoot, "sessions");

  const refreshSkillProjectContext = (cwd?: string) => {
    const resource = resolveProjectSkillDiscovery(skillStore, config.projectsMemoryDir, cwd);
    return {
      name: skillStore.getProjectName(),
      skillsDir: skillStore.getProjectSkillsDir(),
      resource,
    };
  };

  // Keep project memory available for users upgrading from the old
  // ~/.pi/agent/<project>/ layout. This is non-destructive: legacy folders
  // remain in place while entries are copied/merged into projects-memory/.
  if (!lazy) migrateLegacyProjectMemoryDirs(agentRoot, config.projectsMemoryDir);
  // Project-scoped store: ~/.pi/agent/<projectsMemoryDir>/<project_name>/
  // Bound from session/tool ctx.cwd, never from factory process.cwd().
  const createProjectStore = (projectInfo: ReturnType<typeof detectProject>): MemoryStore | null => {
    if (!projectInfo.memoryDir) return null;
    return new MemoryStore({
      ...config,
      memoryCharLimit: config.projectCharLimit,
      memoryDir: projectInfo.memoryDir,
    });
  };
  let projectMemoryDir: string | null = null;
  let projectStore: MemoryStore | null = null;
  let projectLoad: Promise<void> | undefined;
  const projectStoreRef = () => projectStore;
  const projectNameRef = () => projectName;
  let configureProjectStore: (candidate: MemoryStore | null) => void = () => {};
  let configureMemoryToolProjectStore: (candidate: MemoryStore | null) => void = () => {};
  const bindProjectFromCwd = async (cwd?: string): Promise<void> => {
    if (!cwd) return;
    const nextProject = detectProject(config.projectsMemoryDir, cwd);
    const nextProjectMemoryDir = nextProject.memoryDir ?? null;
    if (nextProjectMemoryDir !== projectMemoryDir) {
      projectMemoryDir = nextProjectMemoryDir;
      projectStore = createProjectStore(nextProject);
      configureProjectStore(projectStore);
      configureMemoryToolProjectStore(projectStore);
      projectLoad = projectStore?.loadFromDisk().catch((error) => {
        projectMemoryDir = null;
        projectStore = null;
        projectName = "";
        projectLoad = undefined;
        throw error;
      });
    }
    await projectLoad;
    projectName = nextProject.name ?? "";
  };
  // Never written by review, consolidation or the correction detector — see
  // store/standing-instructions.ts for why provenance has to be structural.
  const standingStore = config.standingInstructionsEnabled !== false
    ? new StandingInstructions(path.join(globalDir, STANDING_FILE), undefined, undefined,
        shouldMigrateExtensionRoot ? path.join(legacyGlobalDir, STANDING_FILE) : undefined)
    : null;

  const initialization = createMemoryInitializer(async (ctx) => {
    const timingPrefix = lazy ? "memory-init" : "session-start";
    // Subagent sessions are read-only: skip the write side of initialization
    // (legacy-dir migration, markdown↔SQLite sync, pruning, backfill) while
    // still loading the store so memory_search/session_search keep working.
    const subagent = isSubagentSession(ctx);
    if (lazy && !subagent) migrateLegacyProjectMemoryDirs(agentRoot, config.projectsMemoryDir);
    if (!persistenceInitialized && !subagent) {
      try {
        await measureLifecycle(`${timingPrefix}.persistence-sync`, async () => {
          // Startup reconciles the global files plus the current project only:
          // a full sweep over every projects-memory folder serialized
          // concurrent startups on the shared mutation-lock database.
          // /memory-sync-markdown remains the repair path for other scopes.
          const startupProject = ctx?.cwd
            ? detectProject(config.projectsMemoryDir, ctx.cwd).name
            : null;
          await migrateThenSyncMarkdownMemories(
            dbManager,
            shouldMigrateExtensionRoot ? legacyGlobalDir : null,
            globalDir,
            config.projectsMemoryDir,
            agentRoot,
            {
              onlyProjects: startupProject ? [startupProject] : [],
              onMigrationSucceeded: () => {
                databaseMigrationPending = false;
              },
            },
          );
        });
        persistenceInitialized = true;
      } catch (error) {
        if (lazy) throw error;
        // Best-effort only: migration or SQLite backfill must not block startup.
      }
    }

    await measureLifecycle(`${timingPrefix}.load`, async () => {
      await store.loadFromDisk();
    });

    if (persistenceInitialized) {
      try {
        pruneEphemeralReviewSessions(dbManager);
      } catch (err) {
        console.warn(`⚠️ Ephemeral session cleanup failed: ${err instanceof Error ? err.message : String(err)}`);
      }
      // Prune sessions older than the configured retention window to bound the
      // size of the session index database (see #183). Runs before
      // scheduleSessionBackfill so pruned sessions are never re-indexed by a
      // backfill started in the same startup.
      if (config.sessionRetentionDays && config.sessionRetentionDays > 0) {
        try {
          const pruneResult = pruneOldSessions(dbManager, config.sessionRetentionDays);
          if (pruneResult.sessionsRemoved > 0) {
            console.info(
              `🧠 Pruned ${pruneResult.sessionsRemoved} old session(s) (> ${config.sessionRetentionDays} days old)`,
            );
          }
        } catch (err) {
          console.warn(`⚠️ Session pruning failed: ${err instanceof Error ? err.message : String(err)}`);
        }
      }
      try {
        await runRecoveryMaintenance({ config, globalDir });
      } catch (err) {
        console.warn(`⚠️ Snapshot retention sweep failed: ${err instanceof Error ? err.message : String(err)}`);
      }
      scheduleSessionBackfill(dbManager, sessionsDir, {
        state: backfillState,
        notify: (message, level) => {
          const ui = sessionContext?.ui;
          if (ui?.notify) {
            ui.notify(message, level);
          } else if (level === "error" || level === "warning") {
            console.warn(message);
          } else {
            console.info(message);
          }
        },
        // Exclude files older than the retention cutoff so sessions pruned by
        // retention are never re-indexed and never schedule another backfill.
        retentionCutoffMs: retentionCutoffMs(config.sessionRetentionDays),
      });
    }
  }, async (ctx) => {
    try {
      await bindProjectFromCwd(ctx.cwd);
    } finally {
      // Even a failed project bind must not leave backfill running past close.
      if (lazy) await joinSessionBackfill(backfillState);
    }
  });

  const ensureMemoryReady: EnsureMemoryReady = (ctx, signal) => initialization.ensure(ctx, signal);
  const memoryPi = withMemoryInitialization(pi, initialization);

  // Skills and pinned instructions must be available even without a lookup.
  pi.on("session_start", async (_event, ctx) => {
    sessionContext = ctx;
    // session_start runs before before_agent_start; a child's rendered prompt
    // leads with the <active_agent> tag, so detect it here early enough to
    // skip the session_start writes below.
    markSubagentSessionFromPrompt(ctx.getSystemPrompt?.() ?? "", ctx);
    // Pinned directives must not depend on migration/SQLite being healthy.
    if (standingStore) await standingStore.load();
    if (!lazy) await ensureMemoryReady(ctx);
    refreshSkillProjectContext(ctx.cwd);
    if (!isSubagentSession(ctx)) {
      await skillStore.migrateLegacySkills();
      await skillStore.ensureDiscoveredRoots();
    }
  });

  registerProjectSkillDiscoveryHandler(pi, skillStore, config.projectsMemoryDir);

  // ── 2. Inject memory policy by default; legacy mode keeps full frozen memory blocks ──
  pi.on("before_agent_start", async (event, ctx) => {
    markSubagentSession(event, ctx);
    const promptContext = await buildPromptContext(config, store, projectStoreRef(), projectNameRef(), standingStore);
    if (!promptContext) return;

    // pi 0.86+ hands handlers a normalized clone of these options and re-renders
    // the prompt from it, so writing a named `memory_policy` section keeps the
    // prompt a set of named sections; returning `systemPrompt` instead forces the
    // rendered text and flattens the sections for every later handler and for
    // section-aware providers (#251). The bundle's single system-prompt owner
    // (extensions/system-prompt.ts) orders the section via
    // config/prompt/sections.jsonc.
    //
    // Branch on that mechanism, not on the field: `systemPromptOptions` is on the
    // event at the 0.80.6 floor too, but 0.80.x-0.85.x pass the raw long-lived
    // base options and only apply a returned `systemPrompt`, so an append there
    // reaches no prompt and does not stay private either — it accumulates in the
    // object `ctx.getSystemPromptOptions()` hands to other extensions. Every
    // 0.86+ runner normalizes `sections` onto the clone it passes.
    const promptOptions = event.systemPromptOptions;
    if (promptOptions && "sections" in promptOptions) {
      (promptOptions.sections as Record<string, string>).memory_policy = promptContext;
      return;
    }

    return {
      systemPrompt: event.systemPrompt + "\n\n" + promptContext,
    };
  });

  // ── 3. Register action-specific memory write tools with SQLite sync ──
  configureMemoryToolProjectStore = registerMemoryTool(memoryPi, store, projectStoreRef, dbManager, projectNameRef, bindProjectFromCwd);

  // ── 4. Register the skill tool ──
  registerSkillTool(pi, skillStore);

  // ── 5. Setup background learning loop (with tool-call-aware nudge) ──
  setupBackgroundReview(pi, store, projectStoreRef, config, {
    dbManager,
    projectName: projectNameRef,
    ensureMemoryReady,
  });

  // ── 6. Setup session-end flush ──
  setupSessionFlush(pi, store, projectStoreRef, config, dbManager, projectNameRef, { ensureMemoryReady });

  // ── 7. Setup auto-consolidation (inject consolidator into stores) ──
  // Keep the failure in the tool result regardless; session-console logging is
  // separately configurable for users who already monitor tool results (#135).
  const runAutoConsolidation = async (
    target: "memory" | "user" | "failure",
    targetStore: MemoryStore,
    toolTarget: "memory" | "user" | "failure" | "project",
    signal?: AbortSignal,
  ) => {
    const result = await triggerConsolidation(
      pi,
      targetStore,
      target,
      signal,
      config.consolidationTimeoutMs,
      toolTarget,
      config,
      null,
      dbManager,
      toolTarget === "project" ? projectNameRef() : null,
    );
    if (result.deferred) {
      console.info(`⏳ Auto-consolidation for '${toolTarget}' deferred: ${result.error ?? "another session holds the consolidation lock"}`);
    } else if (shouldWarnAutoConsolidationFailure(config.autoConsolidationWarnOnFailure, result.consolidated, result.partial === true)) {
      console.warn(`⚠️ Auto-consolidation ${result.partial ? "partially " : ""}failed for '${toolTarget}': ${result.error ?? "no reason reported"}`);
    }
    return result;
  };

  store.setConsolidator((target, signal) => runAutoConsolidation(target, store, target, signal));
  configureProjectStore = (candidate) => {
    if (!candidate) return;
    candidate.setConsolidator((target, signal) =>
      runAutoConsolidation(target, candidate, target === "memory" ? "project" : target, signal),
    );
  };
  configureProjectStore(projectStore);
  registerConsolidateCommand(memoryPi, store, config.consolidationTimeoutMs, projectStoreRef, projectNameRef, config, dbManager);

  // ── 8. Setup correction detection ──
  setupCorrectionDetector(pi, store, projectStoreRef, config, dbManager, projectNameRef, { ensureMemoryReady });

  // ── 9. Register commands ──
  registerInsightsCommand(memoryPi, store, projectStoreRef, projectNameRef, dbManager);
  registerSkillsCommand(pi, skillStore);
  registerInterviewCommand(memoryPi, store);
  registerSwitchProjectCommand(pi, config);
  registerLearnMemoryCommand(pi);
  registerSyncMarkdownMemoriesCommand(memoryPi, dbManager, globalDir, config.projectsMemoryDir, agentRoot);
  registerPreviewContextCommand(lazy ? pi : memoryPi, store, projectStoreRef, projectNameRef, config, standingStore);
  if (standingStore) registerStandingPinCommand(pi, standingStore);

  // ── 10. Live session indexing ──
  pi.on("message_end", async (_event, ctx) => {
    if (isSubagentSession(ctx)) return;
    // Cold lazy sessions remain in Pi's JSONL files and are backfilled on use.
    if (lazy && !initialization.isReady()) return;
    scheduleLiveSessionIndex(dbManager, ctx.sessionManager, {
      onError: (err) => console.warn(`⚠️ Live session indexing failed: ${err instanceof Error ? err.message : String(err)}`),
    });
  });

  // ── 11. SQLite session search + extended memory ──
  registerSessionSearchTool(config.sessionSearch?.variant === "anchors" ? pi : memoryPi,
    dbManager, config.sessionSearch ?? { variant: "legacy" });
  registerMemorySearchTool(memoryPi, dbManager, { usageTrackingEnabled: config.usageHitTrackingEnabled });
  registerIndexSessionsCommand(memoryPi, config);

  // ── 12. Auto-index session on shutdown ──
  // Registered last, so this runs after the session-flush shutdown handler and
  // is the final DB activity. Closing here truncates the WAL via
  // PRAGMA wal_checkpoint(TRUNCATE); without it the WAL only grows to its
  // high-water mark and is never reclaimed across sessions.
  //
  // Ordering is safe: Pi's ExtensionRunner.emit() runs same-extension handlers
  // sequentially in registration order and awaits each one, so the flush above
  // fully completes before close() runs. WARNING: do not register another
  // DB-writing session_shutdown handler after this block — it would run after
  // close() and silently no-op.
  pi.on("session_shutdown", async (_event, ctx) => {
    await initialization.close();
    if (lazy && !initialization.isReady()) {
      databaseClosed = true;
      dbManager.close();
      return;
    }
    try {
      if (isSubagentSession(ctx)) return;
      measureLifecycleSync("shutdown.active-index", () => {
        const sessionFile = ctx.sessionManager.getSessionFile();
        if (sessionFile && fs.existsSync(sessionFile)) {
          const sessionData = parseSessionFile(sessionFile);
          if (sessionData) {
            dbManager.withCorruptionRecovery(() => {
              indexSession(dbManager, sessionData);
              // Keep session_files metadata in sync with the final on-disk state.
              // Pi appends the closing session entry on shutdown after the last
              // message_end, so without this upsert the stored size/mtime would be
              // stale and the next startup would re-parse this file unnecessarily.
              upsertSessionFileMetadata(dbManager, sessionFile, sessionData.id);
            });
          }
        }
      });
    } catch {
      // Silent fail — don't block shutdown
    } finally {
      try {
        await measureLifecycle("shutdown.index-waits", () => Promise.all([
          waitForSessionBackfill(SESSION_BACKFILL_SHUTDOWN_TIMEOUT_MS, backfillState),
          waitForLiveSessionIndex(SESSION_LIVE_INDEX_SHUTDOWN_TIMEOUT_MS),
        ]));
      } catch {
        // Best effort only — shutdown should not be held up by indexing errors.
      }
      try {
        databaseClosed = true;
        measureLifecycleSync("shutdown.database-close", () => dbManager.close());
      } catch { /* best effort — never block shutdown */ }
    }
  });
}