@chendpoc/pi-memory 0.1.12 → 0.2.0

package/src/pi-extension.ts

@@ -1,243 +1,183 @@
 import type { AgentMessage } from "@earendil-works/pi-agent-core";
-import type { UserMessage } from "@earendil-works/pi-ai";
-import { StringEnum } from "@earendil-works/pi-ai";
 import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
-import { truncateHead } from "@earendil-works/pi-coding-agent";
-import { Type } from "typebox";
-
-import {
-  createOllamaLLMClient,
-  createOllamaMemoryHelper,
-  ollamaHealthCheck,
-  type OllamaConfig,
-} from "./adapters/ollamaClient.js";
-import {
-  createOpenAICompatLLMClient,
-  createOpenAICompatMemoryHelper,
-  openaiCompatHealthCheck,
-  type OpenAICompatConfig,
-} from "./adapters/openaiCompatClient.js";
-import {
-  createPiLLMClient,
-  DEFAULT_HELPER_MODEL,
-  DEFAULT_HELPER_PROVIDER,
-  resolveMemoryHelperLLM,
-} from "./adapters/piComplete.js";
-import type { MemoryConfig } from "./config.js";
-import {
-  loadMemorySettings,
-  resolveHelperModelSpec,
-} from "./settings.js";
-import { createFallbackQuery } from "./fallback/index.js";
-import type { RerankOptions } from "./fallback/llmRerank.js";
-import type { MemoryHelperLLM } from "./preflight/detectIntents.js";
-import { runMemoryPreflight } from "./preflight/hook.js";
+import findLastIndex from "lodash/findLastIndex.js";
+
+import { createLlmClient, type LlmClient } from "./adapters/llm/index.js";
+import { loadEnv, readPiMemoryEnv, resolveMemoryAgentDir } from "./config/index.js";
+import { createConsolidateScheduler, startConsolidateInterval, type ConsolidateScheduler } from "./consolidate/scheduler.js";
+import { registerCommands } from "./commands/index.js";
+import { registerCompactHandlers } from "./compact/register.js";
+import { runEpisodicPreflight } from "./preflight/episodic.js";
+import { mergePrivateMemoryBlocks, renderMemoryCapPrivateMemory } from "./preflight/render.js";
+import { isSubagentSession } from "./preflight/session.js";
 import { injectPrivateMemoryContext } from "./preflight/strip.js";
-import { MemoryService } from "./service.js";
-import type { LLMClient } from "./trainer/llmExtractor.js";
-import {
-  createMemoryAppendTool,
-  MEMORY_APPEND_DESCRIPTION,
-  MEMORY_APPEND_NAME,
-  MEMORY_APPEND_PROMPT_GUIDELINES,
-  MEMORY_APPEND_PROMPT_SNIPPET,
-} from "./tools/memoryAppend.js";
-import {
-  createMemoryRecallTool,
-  MEMORY_RECALL_DESCRIPTION,
-  MEMORY_RECALL_NAME,
-  MEMORY_RECALL_PROMPT_GUIDELINES,
-  MEMORY_RECALL_PROMPT_SNIPPET,
-} from "./tools/memoryRecall.js";
-
-const MemoryRecallParams = Type.Object({
-  mode: Type.Optional(
-    StringEnum(["direct_relation", "path_query", "typed_neighborhood"] as const),
-  ),
-  anchor_mentions: Type.Array(Type.String()),
-  relation_constraints: Type.Optional(Type.Array(Type.String())),
-  candidate_type: Type.Optional(Type.String()),
-  scope_filter: Type.Optional(Type.Array(Type.String())),
-  target_slot: Type.Optional(StringEnum(["head", "tail"] as const)),
-  time_window: Type.Optional(Type.String()),
-  evidence_budget: Type.Optional(Type.Number()),
-  result_limit: Type.Optional(Type.Number()),
-});
-
-const MemoryAppendParams = Type.Object({
-  content: Type.String({ description: "New entries to append (markdown bullet points)." }),
-});
-
-const RECALL_MAX_LINES = 200;
-const RECALL_MAX_BYTES = 32_000;
-
-let sharedService: MemoryService | null = null;
-let sessionCfg: MemoryConfig | null = null;
-let settingsHelperModel: string | undefined;
-let settingsOllama: OllamaConfig | null = null;
-let settingsVllm: OpenAICompatConfig | null = null;
-let sharedHelper: MemoryHelperLLM | null = null;
-let sharedLLMClient: LLMClient | null = null;
-/**
- * Per-turn preflight result, set by before_agent_start and consumed by context.
- * userPayload = raw event.prompt (no host scaffolding).
- * privateContext = <private_memory> block to inject.
- */
-let turnPreflight: { userPayload: string; privateContext: string } | null = null;
-
-/**
- * True on the very first turn of a session; used to force the intent helper
- * (bypassing the lexical gate) so the first message always checks memory.
- */
+import { enqueueShutdownMetadata, readParentSession } from "./shutdown/enqueue.js";
+import { resolveSidecarPaths } from "./sidecar/paths.js";
+import { formatTimestamp } from "./utils/time.js";
+import { createReindexScheduler, type ReindexScheduler } from "./sidecar/reindexBridge.js";
+import { ensureSidecarRunning, stopSidecar } from "./sidecar/sidecarManager.js";
+import { MemoryStore } from "./store/memoryStore.js";
+
+loadEnv();
+
+type TurnPreflight = {
+  userPayload: string;
+  privateContext: string;
+};
+
+let memoryStore: MemoryStore | null = null;
+let sidecarPaths: ReturnType<typeof resolveSidecarPaths> | null = null;
+let reindexScheduler: ReindexScheduler | null = null;
+let consolidateScheduler: ConsolidateScheduler | null = null;
+let stopConsolidateInterval: (() => void) | null = null;
+let llmClient: LlmClient | null = null;
+let sessionMemoryCap: string | null = null;
+let turnPreflight: TurnPreflight | null = null;
 let isFirstTurn = false;
+let isSubagent = false;
 
-export function getSharedMemoryService(): MemoryService | null {
-  return sharedService;
+function getUserMessageText(message: AgentMessage): string | null {
+  if (message.role !== "user") return null;
+  if (typeof message.content === "string") return message.content;
+  return message.content
+    .filter((block) => block.type === "text")
+    .map((block) => block.text)
+    .join("\n");
 }
 
-function getHelperModelSpec(pi: ExtensionAPI): string | undefined {
-  return resolveHelperModelSpec(pi.getFlag("memory-helper-model"), settingsHelperModel);
+function setUserMessageText(message: AgentMessage, text: string): AgentMessage {
+  if (message.role !== "user") return message;
+  return { ...message, content: text } as AgentMessage;
 }
 
-function isOllamaModel(spec: string | undefined): boolean {
-  return !!spec?.startsWith("ollama/");
+function findLastUserMessageIndex(messages: AgentMessage[]): number {
+  return findLastIndex(messages, (message) => message.role === "user");
 }
 
-async function refreshMemoryHelper(ctx: ExtensionContext, pi: ExtensionAPI): Promise<void> {
-  const helperModel = getHelperModelSpec(pi);
+async function refreshLlm(ctx: ExtensionContext, pi: ExtensionAPI): Promise<void> {
+  const env = readPiMemoryEnv();
+  const flagModel = pi.getFlag("memory-helper-model");
+  const modelSpec = typeof flagModel === "string" ? flagModel : env.helperModel;
+  llmClient = await createLlmClient({ ctx, modelSpec, env });
+}
 
-  if (settingsVllm) {
-    const vllmOk = await openaiCompatHealthCheck(settingsVllm.baseUrl);
-    if (vllmOk) {
-      sharedHelper = createOpenAICompatMemoryHelper(settingsVllm);
-      sharedLLMClient = createOpenAICompatLLMClient(settingsVllm);
-      return;
-    }
-  }
+async function bootstrapSidecar(): Promise<void> {
+  if (!memoryStore || !sidecarPaths) return;
 
-  if ((isOllamaModel(helperModel) || settingsOllama) && settingsOllama) {
-    const ollamaOk = await ollamaHealthCheck(settingsOllama.baseUrl);
-    if (ollamaOk) {
-      sharedHelper = createOllamaMemoryHelper(settingsOllama);
-      sharedLLMClient = createOllamaLLMClient(settingsOllama);
-      return;
-    }
-  }
+  await memoryStore.ensureInitialized();
+  await ensureSidecarRunning({
+    socketPath: sidecarPaths.socketPath,
+    dbPath: sidecarPaths.dbPath,
+  });
 
-  sharedHelper = await resolveMemoryHelperLLM(ctx, helperModel);
-  sharedLLMClient = createPiLLMClient(ctx, helperModel);
-}
+  reindexScheduler ??= createReindexScheduler({
+    sidecar: sidecarPaths,
+    agentDir: memoryStore.agentDir,
+    getDocuments: () => memoryStore!.exportForIndex(),
+    debounceMs: readPiMemoryEnv().reindexDebounceMs,
+  });
 
-function getRerankOpts(): RerankOptions | null {
-  if (!sharedLLMClient) return null;
-  return { client: sharedLLMClient };
-}
+  memoryStore.onSyncToSidecar(() => reindexScheduler?.schedule());
 
-function getUserMessageText(message: AgentMessage): string | null {
-  if (message.role !== "user") return null;
-  const user = message as UserMessage;
-  if (typeof user.content === "string") return user.content;
-  return user.content
-    .filter((block): block is { type: "text"; text: string } => block.type === "text")
-    .map((block) => block.text)
-    .join("\n");
+  await reindexScheduler.runNow();
 }
 
-function setUserMessageText(message: UserMessage, text: string): UserMessage {
-  return { ...message, content: text };
-}
+function bootstrapConsolidate(): void {
+  if (!memoryStore) return;
 
-function findLastUserMessageIndex(messages: AgentMessage[]): number {
-  for (let i = messages.length - 1; i >= 0; i--) {
-    if (messages[i]?.role === "user") return i;
-  }
-  return -1;
+  consolidateScheduler = createConsolidateScheduler({
+    getStore: () => memoryStore,
+    getAgentDir: () => memoryStore?.agentDir ?? null,
+    getLlm: () => llmClient,
+    debounceMs: readPiMemoryEnv().consolidateDebounceMs,
+    onComplete: preloadSessionMemoryCap,
+  });
+
+  memoryStore.onConsolidateCheck(() => consolidateScheduler?.schedule());
+
+  stopConsolidateInterval?.();
+  stopConsolidateInterval = startConsolidateInterval(() => {
+    void consolidateScheduler?.runNow();
+  });
+
+  void consolidateScheduler.runNow();
 }
 
-function formatMemoryStatus(service: MemoryService): string {
-  const snap = service.getStatus();
-  const lines = [
-    `status: ${snap.status}`,
-    snap.mode ? `mode: ${snap.mode}` : null,
-    snap.reason ? `reason: ${snap.reason}` : null,
-    sharedHelper
-      ? `helper: ${settingsVllm ? `vllm/${settingsVllm.model}` : settingsOllama ? `ollama/${settingsOllama.model}` : (settingsHelperModel ?? "pi-ai")}`
-      : "helper: none (regex only)",
-    settingsVllm ? `vllm: ${settingsVllm.baseUrl} (${settingsVllm.model})` : null,
-    settingsOllama ? `ollama: ${settingsOllama.baseUrl} (${settingsOllama.model})` : null,
-    snap.health ? `health: ${JSON.stringify(snap.health)}` : null,
-  ].filter(Boolean);
-  return lines.join("\n");
+async function preloadSessionMemoryCap(): Promise<void> {
+  if (!memoryStore) return;
+  const fallback = await memoryStore.readForFallback();
+  sessionMemoryCap = renderMemoryCapPrivateMemory(fallback) || null;
 }
 
 export default function piMemoryExtension(pi: ExtensionAPI): void {
-  pi.registerFlag("memory-helper-model", {
-    description: `Model for memory intent helper (default: ${DEFAULT_HELPER_PROVIDER}/${DEFAULT_HELPER_MODEL})`,
-    type: "string",
-  });
-
   pi.on("session_start", async (_event, ctx) => {
-    const loaded = loadMemorySettings();
-    const cfg = loaded.config;
-    settingsHelperModel = loaded.helperModel;
-    settingsOllama = loaded.ollama;
-    settingsVllm = loaded.vllm;
-    sessionCfg = cfg;
-    turnPreflight = null;
+    const agentDir = resolveMemoryAgentDir();
+    memoryStore = new MemoryStore({ agentDir });
+    sidecarPaths = resolveSidecarPaths(agentDir);
+    isSubagent = isSubagentSession(ctx);
     isFirstTurn = true;
-    sharedHelper = null;
-
-    if (cfg.provider === "disabled") return;
+    turnPreflight = null;
 
-    const service = new MemoryService(cfg);
-    sharedService = service;
+    await refreshLlm(ctx, pi);
 
-    await service.start();
-    service.startSessionIndex();
-    if (cfg.trainer.auto_interval) {
-      service.startAutoTrainer();
-    }
     try {
-      await refreshMemoryHelper(ctx, pi);
+      await bootstrapSidecar();
+      bootstrapConsolidate();
+      await preloadSessionMemoryCap();
     } catch {
-      /* helper unavailable — regex-only preflight */
+      // sidecar optional; preflight falls back to MEMORY.md
     }
   });
 
-  pi.on("session_shutdown", async () => {
-    if (sharedService) {
-      await sharedService.stop();
-      sharedService = null;
+  pi.on("session_shutdown", async (event, ctx) => {
+    if (memoryStore) {
+      const header = ctx.sessionManager.getHeader() as unknown as Record<string, unknown> | null;
+      void enqueueShutdownMetadata(memoryStore.agentDir, {
+        sessionFile: ctx.sessionManager.getSessionFile() ?? "",
+        parentSession: readParentSession(header),
+        reason: event.reason,
+        isSubagent,
+        enqueuedAt: formatTimestamp(),
+      }).catch(() => {});
     }
-    sessionCfg = null;
-    settingsHelperModel = undefined;
-    settingsOllama = null;
-    settingsVllm = null;
-    sharedHelper = null;
-    sharedLLMClient = null;
+
+    stopConsolidateInterval?.();
+    stopConsolidateInterval = null;
+    consolidateScheduler = null;
+    reindexScheduler = null;
+    memoryStore = null;
+    sidecarPaths = null;
+    llmClient = null;
+    sessionMemoryCap = null;
     turnPreflight = null;
     isFirstTurn = false;
-  });
+    isSubagent = false;
 
-  pi.on("model_select", async (_event, ctx) => {
-    await refreshMemoryHelper(ctx, pi);
+    try {
+      await stopSidecar();
+    } catch {
+      // ignore shutdown errors
+    }
   });
 
-  pi.on("agent_start", () => {
-    turnPreflight = null;
+  pi.on("model_select", async (_event, ctx) => {
+    await refreshLlm(ctx, pi);
   });
 
   pi.on("before_agent_start", async (event, ctx) => {
-    const service = sharedService;
-    const cfg = sessionCfg;
-    if (!service || !cfg || cfg.provider === "disabled") return;
+    if (!memoryStore || !sidecarPaths) return;
 
-    const userPayload = String(((event as unknown) as { prompt?: string }).prompt ?? "").trim();
+    const userPayload = String(event.prompt ?? "").trim();
     if (!userPayload) return;
 
     const forceHelper = isFirstTurn;
     isFirstTurn = false;
 
+    if (isSubagent) {
+      turnPreflight = { userPayload, privateContext: sessionMemoryCap ?? "" };
+      return;
+    }
+
     const workingUi = ctx.hasUI
       ? {
           show: (msg: string) => ctx.ui.setWorkingMessage(msg),
@@ -248,110 +188,28 @@ export default function piMemoryExtension(pi: ExtensionAPI): void {
 
     try {
       workingUi?.show("Recalling memory…");
-      const result = await runMemoryPreflight(userPayload, service, {
-        helper: sharedHelper,
-        forceHelper,
-        fallback,
+      const env = readPiMemoryEnv();
+      const result = await runEpisodicPreflight(userPayload, {
+        socketPath: sidecarPaths.socketPath,
+        agentDir: memoryStore.agentDir,
+        store: memoryStore,
+        llm: llmClient,
+        force: forceHelper,
+        budgetMs: env.preflightBudgetMs,
         signal: ctx.signal,
-        rerankOpts: getRerankOpts(),
         onProgress: workingUi?.update,
       });
-      turnPreflight = result?.privateContext
-        ? { userPayload, privateContext: result.privateContext }
-        : null;
+
+      const episodicContext = result?.privateContext ?? null;
+      const combined = mergePrivateMemoryBlocks(sessionMemoryCap, episodicContext);
+      turnPreflight = { userPayload, privateContext: combined ?? "" };
     } finally {
       workingUi?.clear();
     }
   });
 
-  const initialSettings = loadMemorySettings();
-  settingsHelperModel = initialSettings.helperModel;
-  settingsOllama = initialSettings.ollama;
-  settingsVllm = initialSettings.vllm;
-
-  const fallback = createFallbackQuery({
-    sessionsDir: initialSettings.config.sessionsDir,
-    memoryMdPaths: initialSettings.config.memoryMdPaths,
-  });
-
-  pi.registerCommand("memory", {
-    description: "Show pi-memory sidecar status and bundle info",
-    handler: async (_args, ctx) => {
-      const service = sharedService;
-      if (!service) {
-        ctx.ui.notify("pi-memory: service not started", "warning");
-        return;
-      }
-      ctx.ui.notify(formatMemoryStatus(service), "info");
-    },
-  });
-
-  pi.registerTool({
-    name: MEMORY_RECALL_NAME,
-    label: "Memory Recall",
-    description: MEMORY_RECALL_DESCRIPTION,
-    promptSnippet: MEMORY_RECALL_PROMPT_SNIPPET,
-    promptGuidelines: [...MEMORY_RECALL_PROMPT_GUIDELINES],
-    parameters: MemoryRecallParams,
-    async execute(_toolCallId, params, signal, onUpdate) {
-      const service = sharedService;
-      if (!service) {
-        return {
-          content: [{ type: "text", text: "Memory service not started." }],
-          details: { truncated: false },
-          isError: true,
-        };
-      }
-      const onProgress = onUpdate
-        ? (msg: string) =>
-            (onUpdate as (u: { content: unknown[]; details: Record<string, unknown> }) => void)?.({
-              content: [{ type: "text", text: msg }],
-              details: { phase: "querying" },
-            })
-        : undefined;
-      const tool = createMemoryRecallTool(service, fallback, getRerankOpts());
-      const result = await tool.run(JSON.stringify(params), signal, onProgress);
-      const truncated = truncateHead(result.content, {
-        maxLines: RECALL_MAX_LINES,
-        maxBytes: RECALL_MAX_BYTES,
-      });
-      let text = truncated.content;
-      if (truncated.truncated) {
-        text += `\n\n[truncated: ${truncated.totalLines} lines, ${truncated.totalBytes} bytes]`;
-      }
-      return {
-        content: [{ type: "text", text }],
-        details: { truncated: truncated.truncated },
-        isError: result.isError ?? false,
-      };
-    },
-  });
-
-  const memoryMdPath = initialSettings.config.memoryMdPaths[0];
-  if (memoryMdPath) {
-    const appendTool = createMemoryAppendTool(memoryMdPath);
-    pi.registerTool({
-      name: MEMORY_APPEND_NAME,
-      label: "Memory Append",
-      description: MEMORY_APPEND_DESCRIPTION,
-      promptSnippet: MEMORY_APPEND_PROMPT_SNIPPET,
-      promptGuidelines: [...MEMORY_APPEND_PROMPT_GUIDELINES],
-      parameters: MemoryAppendParams,
-      async execute(_toolCallId, params) {
-        const result = await appendTool.run(JSON.stringify(params));
-        return {
-          content: [{ type: "text", text: result.content }],
-          details: {},
-          isError: result.isError ?? false,
-        };
-      },
-    });
-  }
-
   pi.on("context", async (event, ctx) => {
-    const service = sharedService;
-    const cfg = sessionCfg;
-    if (!service || !cfg || cfg.provider === "disabled") return;
+    if (!memoryStore || !sidecarPaths) return;
 
     const userIndex = findLastUserMessageIndex(event.messages);
     if (userIndex < 0) return;
@@ -362,43 +220,55 @@ export default function piMemoryExtension(pi: ExtensionAPI): void {
     let privateContext: string | undefined;
     let userPayload: string;
 
-    if (turnPreflight?.privateContext) {
-      // Happy path: before_agent_start already ran preflight.
-      // userPayload is the raw prompt; scaffolded may have host-added prefix.
-      privateContext = turnPreflight.privateContext;
+    if (turnPreflight && turnPreflight.userPayload === scaffolded.trim()) {
+      privateContext = turnPreflight.privateContext || undefined;
       userPayload = turnPreflight.userPayload;
     } else {
-      // Fallback: before_agent_start didn't fire (e.g. first context after
-      // session restore). Run preflight inline; no scaffold separation.
-      const userTurnCount = event.messages.filter((m) => m.role === "user").length;
-      const preflight = await runMemoryPreflight(scaffolded, service, {
-        helper: sharedHelper,
-        forceHelper: userTurnCount === 1,
-        fallback,
-        signal: ctx.signal,
-        rerankOpts: getRerankOpts(),
-      });
-      if (!preflight?.privateContext) return;
-      privateContext = preflight.privateContext;
-      userPayload = scaffolded;
-      // Cache so tool-call follow-up context calls reuse without re-running.
-      turnPreflight = { userPayload: scaffolded, privateContext };
+      const userTurnCount = event.messages.filter((message) => message.role === "user").length;
+      const env = readPiMemoryEnv();
+
+      if (isSubagent) {
+        privateContext = sessionMemoryCap ?? undefined;
+        userPayload = scaffolded;
+        turnPreflight = { userPayload: scaffolded, privateContext: privateContext ?? "" };
+      } else {
+        const preflight = await runEpisodicPreflight(scaffolded, {
+          socketPath: sidecarPaths.socketPath,
+          agentDir: memoryStore.agentDir,
+          store: memoryStore,
+          llm: llmClient,
+          force: userTurnCount === 1,
+          budgetMs: env.preflightBudgetMs,
+          signal: ctx.signal,
+        });
+        privateContext = mergePrivateMemoryBlocks(
+          sessionMemoryCap,
+          preflight?.privateContext,
+        ) || undefined;
+        userPayload = scaffolded;
+        turnPreflight = { userPayload: scaffolded, privateContext: privateContext ?? "" };
+      }
     }
 
-    const injectedText = injectPrivateMemoryContext(
-      scaffolded,
-      userPayload,
-      privateContext,
-    );
+    if (!privateContext) return;
 
+    const injectedText = injectPrivateMemoryContext(scaffolded, userPayload, privateContext);
     const messages = [...event.messages];
-    messages[userIndex] = setUserMessageText(
-      messages[userIndex] as UserMessage,
-      injectedText,
-    );
-
+    messages[userIndex] = setUserMessageText(messages[userIndex]!, injectedText);
     return { messages };
   });
+
+  registerCommands(pi, {
+    getMemoryStore: () => memoryStore,
+    onRemembered: preloadSessionMemoryCap,
+    getAgentDir: () => memoryStore?.agentDir ?? resolveMemoryAgentDir(),
+  });
+
+  registerCompactHandlers(pi, {
+    getMemoryStore: () => memoryStore,
+    getLlmClient: () => llmClient,
+    onCompactionIngested: preloadSessionMemoryCap,
+  });
 }
 
 export { piMemoryExtension };