@chendpoc/pi-memory 0.1.11 → 0.1.13

package/src/tools/memoryRecall.ts

@@ -55,7 +55,7 @@ export class MemoryRecallTool {
     };
   }
 
-  async run(argsJson: string, signal?: AbortSignal): Promise<ToolResult> {
+  async run(argsJson: string, signal?: AbortSignal, onProgress?: (msg: string) => void): Promise<ToolResult> {
     const args = parseArgs(argsJson);
     if ("error" in args) {
       return { content: args.error, isError: true };
@@ -67,29 +67,53 @@ export class MemoryRecallTool {
       return { content: validation, isError: true };
     }
 
-    return this.runIntent(intent, signal);
+    return this.runIntent(intent, signal, onProgress);
   }
 
-  async runIntent(intent: QueryIntent, signal?: AbortSignal): Promise<ToolResult> {
+  async runIntent(intent: QueryIntent, signal?: AbortSignal, onProgress?: (msg: string) => void): Promise<ToolResult> {
     if (this.service.status() !== "ready") {
       return this.fallbackResult(intent, "service_unavailable", "fallback");
     }
 
-    let result = await this.service.query(intent, signal);
+    // Emit a progress message if the sidecar takes longer than 500 ms.
+    let progressTimer: ReturnType<typeof setTimeout> | null = null;
+    if (onProgress) {
+      progressTimer = setTimeout(() => {
+        onProgress("Querying episodic memory…");
+      }, 500);
+    }
+
+    let result;
+    try {
+      result = await this.service.query(intent, signal);
+    } finally {
+      if (progressTimer != null) clearTimeout(progressTimer);
+    }
+
     if (result.transportError || result.errorClass === "unavailable") {
       return this.fallbackResult(intent, "service_unavailable", "fallback");
     }
 
     if (result.errorClass === "retryable") {
-      await sleep(500, signal);
-      result = await this.service.query(intent, signal);
+      let retryTimer: ReturnType<typeof setTimeout> | null = null;
+      if (onProgress) {
+        retryTimer = setTimeout(() => onProgress("Retrying memory query…"), 500);
+      }
+      let retryResult;
+      try {
+        await sleep(500, signal);
+        retryResult = await this.service.query(intent, signal);
+      } finally {
+        if (retryTimer != null) clearTimeout(retryTimer);
+      }
       if (
-        result.transportError ||
-        result.errorClass === "unavailable" ||
-        result.errorClass === "retryable"
+        retryResult.transportError ||
+        retryResult.errorClass === "unavailable" ||
+        retryResult.errorClass === "retryable"
       ) {
         return this.fallbackResult(intent, "retryable_failed", "fallback_after_retry");
       }
+      result = retryResult;
     }
 
     if (result.errorClass === "permanent") {

package/src/trainer/index.ts

@@ -129,7 +129,7 @@ export async function trainBundle(
   };
 }
 
-export { loadSessions } from "./sessionLoader.js";
+export { loadSessionFile, loadSessions } from "./sessionLoader.js";
 export type { LoadedSession, SessionTurn, SessionLoaderOptions } from "./sessionLoader.js";
 export { extractFacts, extractFactsFromSessions, RELATION_CATALOG, ALL_RELATIONS } from "./extractFacts.js";
 export type {

package/src/trainer/scheduler.ts

@@ -5,6 +5,8 @@ export interface SchedulerConfig {
   interval: string | null;
   /** Passed through to trainBundle. */
   trainConfig?: Omit<TrainBundleConfig, "full" | "dryRun">;
+  /** Called after a successful train tick (no error). Use to notify MemoryService. */
+  onSuccess?: () => void;
 }
 
 export interface TrainScheduler {
@@ -69,6 +71,7 @@ export function createTrainScheduler(
         full: false,
         dryRun: false,
       });
+      config.onSuccess?.();
     } catch (err) {
       error = err instanceof Error ? err.message : String(err);
     }

package/src/trainer/sessionLoader.ts

@@ -1,6 +1,7 @@
 import { createHash } from "node:crypto";
 import fs from "node:fs/promises";
 import path from "node:path";
+import { messageText, parseJsonlSession as parseActiveBranchJsonl } from "../session/activeBranch.js";
 
 export interface SessionTurn {
   role: string;
@@ -14,6 +15,8 @@ export interface LoadedSession {
   createdAt: string;
   filePath: string;
   modifiedAt: Date;
+  parentSessionFile?: string;
+  parentSessionId?: string;
   turns: SessionTurn[];
 }
 
@@ -29,39 +32,6 @@ interface PiSessionFile {
   messages?: PiSessionMessage[];
 }
 
-interface JsonlSessionHeader {
-  type: "session";
-  id?: string;
-  timestamp?: string;
-  title?: string;
-}
-
-interface JsonlMessageLine {
-  type: "message";
-  message?: {
-    role?: string;
-    content?: unknown;
-  };
-}
-
-function messageText(content: unknown): string {
-  if (typeof content === "string") return content;
-  if (!Array.isArray(content)) return "";
-  const parts: string[] = [];
-  for (const block of content) {
-    if (typeof block === "string") {
-      parts.push(block);
-      continue;
-    }
-    if (block && typeof block === "object") {
-      const b = block as Record<string, unknown>;
-      if (typeof b.text === "string") parts.push(b.text);
-      else if (typeof b.content === "string") parts.push(b.content);
-    }
-  }
-  return parts.join("\n");
-}
-
 export interface SessionLoaderOptions {
   sessionsDir: string;
   modifiedAfter?: Date | null;
@@ -94,49 +64,6 @@ async function collectSessionFiles(dir: string): Promise<string[]> {
   return files;
 }
 
-function parseJsonlSession(raw: string, filePath: string): LoadedSession | null {
-  const lines = raw.split("\n").filter((l) => l.trim());
-  if (lines.length === 0) return null;
-
-  let header: JsonlSessionHeader | null = null;
-  const turns: SessionTurn[] = [];
-  let turnIndex = 0;
-
-  for (const line of lines) {
-    let obj: Record<string, unknown>;
-    try {
-      obj = JSON.parse(line) as Record<string, unknown>;
-    } catch {
-      continue;
-    }
-
-    if (obj.type === "session" && !header) {
-      header = obj as unknown as JsonlSessionHeader;
-      continue;
-    }
-
-    if (obj.type === "message") {
-      const msg = (obj as unknown as JsonlMessageLine).message;
-      if (!msg?.role || !msg.content) continue;
-      if (msg.role !== "user" && msg.role !== "assistant") continue;
-      const text = messageText(msg.content);
-      if (!text.trim()) continue;
-      turns.push({ role: msg.role, content: text, turnIndex: turnIndex++ });
-    }
-  }
-
-  if (turns.length === 0) return null;
-
-  return {
-    id: header?.id ?? path.basename(filePath, path.extname(filePath)),
-    title: header?.title ?? "",
-    createdAt: header?.timestamp ?? "",
-    filePath,
-    modifiedAt: new Date(),
-    turns,
-  };
-}
-
 function parseJsonSession(raw: string, filePath: string): LoadedSession | null {
   let session: PiSessionFile;
   try {
@@ -167,6 +94,16 @@ function parseJsonSession(raw: string, filePath: string): LoadedSession | null {
   };
 }
 
+function parseJsonlSession(raw: string, filePath: string): LoadedSession | null {
+  const parsed = parseActiveBranchJsonl(raw, filePath);
+  if (!parsed) return null;
+  return {
+    ...parsed,
+    filePath,
+    modifiedAt: new Date(),
+  };
+}
+
 /**
  * Scan session files (JSON + JSONL, recursive subdirectories), parse Pi session
  * format, optionally filter by modified-after timestamp for incremental training.
@@ -190,17 +127,7 @@ export async function loadSessions(
     if (!st.isFile()) continue;
     if (modifiedAfter && st.mtime <= modifiedAfter) continue;
 
-    let raw: string;
-    try {
-      raw = await fs.readFile(filePath, "utf8");
-    } catch {
-      continue;
-    }
-
-    const isJsonl = filePath.endsWith(".jsonl");
-    const parsed = isJsonl
-      ? parseJsonlSession(raw, filePath)
-      : parseJsonSession(raw, filePath);
+    const parsed = await loadSessionFile(filePath);
 
     if (!parsed) continue;
     parsed.modifiedAt = st.mtime;
@@ -211,6 +138,30 @@ export async function loadSessions(
   return deduplicateSessions(sessions);
 }
 
+export async function loadSessionFile(filePath: string): Promise<LoadedSession | null> {
+  let st: Awaited<ReturnType<typeof fs.stat>>;
+  try {
+    st = await fs.stat(filePath);
+  } catch {
+    return null;
+  }
+  if (!st.isFile()) return null;
+
+  let raw: string;
+  try {
+    raw = await fs.readFile(filePath, "utf8");
+  } catch {
+    return null;
+  }
+
+  const parsed = filePath.endsWith(".jsonl")
+    ? parseJsonlSession(raw, filePath)
+    : parseJsonSession(raw, filePath);
+  if (!parsed) return null;
+  parsed.modifiedAt = st.mtime;
+  return parsed;
+}
+
 function deduplicateSessions(sessions: LoadedSession[]): LoadedSession[] {
   const seen = new Set<string>();
   return sessions.filter((s) => {