@chendpoc/pi-memory 0.1.12 → 0.1.13

package/src/session/activeBranch.ts

@@ -0,0 +1,294 @@
+import path from "node:path";
+
+interface UnknownRecord {
+  [key: string]: unknown;
+}
+
+export interface ActiveBranchTurn {
+  role: string;
+  content: string;
+  turnIndex: number;
+}
+
+export interface ParsedSession {
+  id: string;
+  title: string;
+  createdAt: string;
+  turns: ActiveBranchTurn[];
+  parentSessionId?: string;
+  parentSessionFile?: string;
+}
+
+interface JsonlNodeInfo {
+  parentId?: string;
+  parentSessionId?: string;
+  parentSessionFile?: string;
+}
+
+interface ParsedJsonlLine {
+  id?: string;
+  parentId?: string;
+  parentSessionId?: string;
+  parentSessionFile?: string;
+  message?: {
+    role?: string;
+    content?: unknown;
+  };
+  role?: string;
+  content?: unknown;
+  timestamp?: string;
+  createdAt?: string;
+}
+
+function asString(value: unknown): string | undefined {
+  if (typeof value === "string") {
+    const text = value.trim();
+    return text.length > 0 ? text : undefined;
+  }
+  return undefined;
+}
+
+function getStringField(obj: UnknownRecord, names: string[]): string | undefined {
+  for (const name of names) {
+    const found = asString(obj[name]);
+    if (found) return found;
+  }
+  return undefined;
+}
+
+function getBooleanField(obj: UnknownRecord, names: string[]): boolean | undefined {
+  for (const name of names) {
+    if (typeof obj[name] === "boolean") return obj[name];
+  }
+  return undefined;
+}
+
+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 UnknownRecord;
+      if (typeof b.text === "string") parts.push(b.text);
+      else if (typeof b.content === "string") parts.push(b.content);
+    }
+  }
+  return parts.join("\n");
+}
+
+function extractMessage(line: ParsedJsonlLine): { role: string; content: unknown } | null {
+  if (line.message?.role && line.message.content !== undefined) {
+    return { role: line.message.role, content: line.message.content };
+  }
+  if (line.role && line.content !== undefined) {
+    return { role: line.role, content: line.content };
+  }
+  return null;
+}
+
+function pickActiveLeafHint(obj: UnknownRecord): string | undefined {
+  const explicitByName = getStringField(obj, [
+    "activeSessionId",
+    "active_session_id",
+    "activeSession",
+    "active_session",
+    "activeLeafId",
+    "active_leaf_id",
+    "activeLeafSessionId",
+    "active_leaf_session_id",
+    "currentSessionId",
+    "current_session_id",
+    "leafSessionId",
+    "leaf_session_id",
+  ]);
+  if (explicitByName) return explicitByName;
+
+  if (getBooleanField(obj, ["active", "is_active", "isActive", "activeLeaf"])
+    && typeof obj.id === "string") {
+    return obj.id;
+  }
+  return undefined;
+}
+
+function parseSessionIdFromPath(filePath: string | undefined): string | undefined {
+  const normalized = asString(filePath);
+  if (!normalized) return undefined;
+  const base = path.basename(normalized, path.extname(normalized));
+  const underscore = base.lastIndexOf("_");
+  if (underscore >= 0 && underscore < base.length - 1) {
+    return base.slice(underscore + 1);
+  }
+  return base || undefined;
+}
+
+export function parseJsonlSession(raw: string, filePath: string): ParsedSession | null {
+  const lines = raw.split("\n").filter((line) => line.trim());
+  if (lines.length === 0) return null;
+
+  const parsedLines: ParsedJsonlLine[] = [];
+  const nodes = new Map<string, JsonlNodeInfo>();
+
+  let headerSessionId: string | undefined;
+  let headerTitle = "";
+  let headerCreatedAt = "";
+  let lastNodeId: string | undefined;
+  let explicitLeafHint: string | undefined;
+  let fallbackParentSessionId: string | undefined;
+  let fallbackParentSessionFile: string | undefined;
+
+  for (const rawLine of lines) {
+    let obj: UnknownRecord;
+    try {
+      obj = JSON.parse(rawLine) as UnknownRecord;
+    } catch {
+      continue;
+    }
+
+    const id = getStringField(obj, ["id", "entryId", "entry_id"]);
+    const parentId = getStringField(obj, ["parentId", "parent_id"]);
+    const parentSessionId = getStringField(obj, [
+      "parentSessionId",
+      "parent_session_id",
+    ]);
+    const parentSessionFile = getStringField(obj, [
+      "parentSessionFile",
+      "parent_session_file",
+      "parentSessionPath",
+      "parent_session_path",
+      "parentSession",
+      "parent_session",
+    ]);
+    const type = asString(obj.type);
+    const message: ParsedJsonlLine["message"] = obj.message as ParsedJsonlLine["message"] | undefined;
+
+    const line: ParsedJsonlLine = {
+      id,
+      parentId,
+      parentSessionId,
+      parentSessionFile,
+      message,
+      role: asString(obj.role),
+      content: obj.content,
+      timestamp: asString(obj.timestamp),
+      createdAt: asString(obj.createdAt),
+    };
+
+    const explicitLeafId = pickActiveLeafHint(obj);
+    if (explicitLeafId) explicitLeafHint = explicitLeafId;
+
+    if (type === "session") {
+      headerSessionId = asString(obj.id) ?? headerSessionId;
+      headerTitle = asString(obj.title) ?? headerTitle;
+      headerCreatedAt = asString(obj.timestamp) ?? asString(obj.created_at) ?? asString(obj.createdAt) ?? headerCreatedAt;
+      if (!fallbackParentSessionId) {
+        fallbackParentSessionId = getStringField(obj, ["parentSessionId", "parent_session_id"]);
+      }
+      if (!fallbackParentSessionFile) {
+        fallbackParentSessionFile = getStringField(obj, [
+          "parentSessionFile",
+          "parent_session_file",
+          "parentSessionPath",
+          "parent_session_path",
+          "parentSession",
+          "parent_session",
+        ]);
+      }
+      if (!fallbackParentSessionId && fallbackParentSessionFile) {
+        fallbackParentSessionId = parseSessionIdFromPath(fallbackParentSessionFile);
+      }
+    }
+
+    if (id && type !== "session") {
+      nodes.set(id, {
+        parentId,
+        parentSessionId,
+        parentSessionFile,
+      });
+      lastNodeId = id;
+    }
+
+    parsedLines.push(line);
+  }
+
+  if (parsedLines.length === 0) return null;
+
+  const hasBranchNodes = nodes.size > 0;
+  let activeLeafId = explicitLeafHint ?? lastNodeId;
+  if (activeLeafId && explicitLeafHint && explicitLeafHint !== lastNodeId && !nodes.has(activeLeafId)) {
+    activeLeafId = lastNodeId;
+  }
+  const activeSessionId = activeLeafId ?? headerSessionId ?? path.basename(filePath, path.extname(filePath));
+  let parentSessionId: string | undefined = fallbackParentSessionId;
+  let parentSessionFile: string | undefined = fallbackParentSessionFile;
+
+  const activeNodeIds = new Set<string>();
+  if (hasBranchNodes && activeLeafId) {
+    let cursor: string | undefined = activeLeafId;
+    const leafNode = nodes.get(activeLeafId);
+    if (leafNode) {
+      if (leafNode.parentSessionId) {
+        parentSessionId = leafNode.parentSessionId;
+      }
+      if (leafNode.parentSessionFile) {
+        parentSessionFile = leafNode.parentSessionFile;
+      }
+      if (!parentSessionId && parentSessionFile) {
+        parentSessionId = parseSessionIdFromPath(parentSessionFile);
+      }
+    }
+    for (let i = 0; i < nodes.size + 5; i++) {
+      if (!cursor || activeNodeIds.has(cursor)) break;
+      activeNodeIds.add(cursor);
+      const node = nodes.get(cursor);
+      if (!node) break;
+      if (!node.parentId) break;
+      cursor = node.parentId;
+      if (!nodes.has(cursor)) break;
+    }
+  }
+
+  const includeLine = (line: ParsedJsonlLine): boolean => {
+    if (!hasBranchNodes || !activeLeafId) return true;
+    if (!line.id && !line.parentId) return false;
+    if (line.id) return activeNodeIds.has(line.id);
+    if (line.parentId) return activeNodeIds.has(line.parentId);
+    return false;
+  };
+
+  const turns: ActiveBranchTurn[] = [];
+  let turnIndex = 0;
+  for (const line of parsedLines) {
+    if (!includeLine(line)) continue;
+
+    const messageObj = extractMessage(line);
+    if (!messageObj?.role || !messageObj.content) continue;
+    if (messageObj.role !== "user" && messageObj.role !== "assistant") continue;
+    const text = messageText(messageObj.content);
+    if (!text.trim()) continue;
+
+    turns.push({
+      role: messageObj.role,
+      content: text,
+      turnIndex: turnIndex++,
+    });
+  }
+
+  if (turns.length === 0) return null;
+
+  return {
+    id: activeSessionId,
+    title: headerTitle,
+    createdAt: headerCreatedAt,
+    turns,
+    parentSessionId,
+    parentSessionFile,
+  };
+}
+
+export { messageText };

package/src/settings.ts

@@ -7,6 +7,7 @@ import {
   defaultMemoryConfig,
   normalizeMemoryConfig,
   type ExtractorType,
+  type ConsolidationConfig,
   type MemoryConfig,
   type MemoryProvider,
   type TrainerConfig,
@@ -31,6 +32,7 @@ export interface MemorySettingsFile {
   /** OpenAI-compatible endpoint config (vLLM, SGLang, LM Studio, etc.) */
   vllm?: Partial<OpenAICompatConfig>;
   trainer?: Partial<TrainerConfig>;
+  consolidation?: Partial<ConsolidationConfig>;
 }
 
 export interface LoadedMemorySettings {
@@ -70,11 +72,19 @@ export function loadMemorySettings(
   configPath = defaultMemoryConfigPath(),
 ): LoadedMemorySettings {
   const fileSettings = readMemorySettingsFile(configPath);
-  const { helperModel, ollama, vllm, trainer, ...configFields } = fileSettings;
+  const {
+    helperModel,
+    ollama,
+    vllm,
+    trainer,
+    consolidation,
+    ...configFields
+  } = fileSettings;
 
   const config = normalizeMemoryConfig({
     ...configFields,
     ...(trainer ? { trainer } : {}),
+    ...(consolidation ? { consolidation } : {}),
     ...overrides,
   } as Partial<MemoryConfig> & Record<string, unknown>);
 
@@ -122,5 +132,30 @@ export function resolveHelperModelSpec(
   return settingsHelperModel;
 }
 
+/** Write memory settings back to disk (preserves unknown fields). */
+export function saveMemorySettings(
+  updates: Partial<MemorySettingsFile>,
+  configPath = defaultMemoryConfigPath(),
+): void {
+  let existing: Record<string, unknown> = {};
+  try {
+    const raw = fs.readFileSync(configPath, "utf8");
+    const parsed = JSON.parse(raw);
+    if (parsed && typeof parsed === "object" && !Array.isArray(parsed)) {
+      existing = parsed as Record<string, unknown>;
+    }
+  } catch { /* start fresh */ }
+
+  const merged = { ...existing, ...updates };
+  const dir = path.dirname(configPath);
+  fs.mkdirSync(dir, { recursive: true });
+  fs.writeFileSync(configPath, JSON.stringify(merged, null, 2) + "\n", "utf8");
+}
+
 export { defaultMemoryConfig };
-export type { ExtractorType, MemoryProvider, TrainerConfig };
+export type {
+  ConsolidationConfig,
+  ExtractorType,
+  MemoryProvider,
+  TrainerConfig,
+};

package/src/tools/memoryAppend.ts

@@ -2,18 +2,20 @@ import fs from "node:fs/promises";
 import { open } from "node:fs/promises";
 import path from "node:path";
 
+import { openConsolidationStore } from "../consolidation/stage1/store.js";
+import type { MemoryScope } from "../consolidation/scope.js";
 import type { ToolResult } from "../types.js";
 
 export const MEMORY_APPEND_NAME = "memory_append";
 
 export const MEMORY_APPEND_DESCRIPTION =
-  "Append new entries to MEMORY.md. Use this instead of file_write or file_edit for memory updates. " +
-  "Writes are append-only with a simple lock to avoid concurrent clobber.";
+  "Persist durable memory only when the user explicitly asks to remember something. " +
+  "Writes go through the pi-memory consolidation stage instead of direct file edits when configured.";
 
-export const MEMORY_APPEND_PROMPT_SNIPPET = "Append durable notes to MEMORY.md";
+export const MEMORY_APPEND_PROMPT_SNIPPET = "Persist explicit user memory";
 
 export const MEMORY_APPEND_PROMPT_GUIDELINES = [
-  "Use memory_append to persist user preferences or facts the user explicitly asked to remember — not for transient task state.",
+  "Use memory_append only when the user explicitly asks to remember a preference or durable fact — not for transient task state.",
 ] as const;
 
 export const MEMORY_APPEND_PARAMETERS = {
@@ -71,8 +73,52 @@ export async function appendToMemoryMd(
   });
 }
 
+export interface Stage1AppendOptions {
+  dbPath: string;
+  sessionId?: string;
+  sessionFile?: string;
+  scope?: MemoryScope;
+  now?: string;
+}
+
+export async function appendToStage1(
+  content: string,
+  opts: Stage1AppendOptions,
+): Promise<string> {
+  const trimmed = content.trim();
+  if (!trimmed) throw new Error("content must not be empty");
+  const store = openConsolidationStore(opts.dbPath);
+  if (!store) throw new Error("could not open consolidation store");
+  const now = opts.now ?? new Date().toISOString();
+  const sessionId = opts.sessionId
+    ? `manual_${opts.sessionId}_${Date.now()}`
+    : `manual_${Date.now()}`;
+  try {
+    store.upsertStage1Output({
+      session_id: sessionId,
+      session_file: opts.sessionFile ?? "manual-memory-append",
+      source_mtime_ms: Date.now(),
+      generated_at: now,
+      raw_memory: trimmed,
+      rollout_summary: trimmed,
+      scope: opts.scope ?? "global",
+      status: "done",
+      selected_for_phase2: false,
+      usage_count: 0,
+      last_usage: now,
+      error_message: null,
+    });
+    return sessionId;
+  } finally {
+    store.close();
+  }
+}
+
 export class MemoryAppendTool {
-  constructor(private readonly memoryMdPath: string) {}
+  constructor(
+    private readonly memoryMdPath: string,
+    private readonly stage1?: Stage1AppendOptions,
+  ) {}
 
   info() {
     return {
@@ -97,6 +143,10 @@ export class MemoryAppendTool {
       return { content: "content must not be empty", isError: true };
     }
     try {
+      if (this.stage1) {
+        const id = await appendToStage1(content, this.stage1);
+        return { content: `queued memory for consolidation: ${id}` };
+      }
       await appendToMemoryMd(this.memoryMdPath, content);
       return { content: `appended to ${this.memoryMdPath}` };
     } catch (e) {
@@ -108,6 +158,9 @@ export class MemoryAppendTool {
   }
 }
 
-export function createMemoryAppendTool(memoryMdPath: string): MemoryAppendTool {
-  return new MemoryAppendTool(memoryMdPath);
+export function createMemoryAppendTool(
+  memoryMdPath: string,
+  stage1?: Stage1AppendOptions,
+): MemoryAppendTool {
+  return new MemoryAppendTool(memoryMdPath, stage1);
 }

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/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) => {