openclaw-memory-hierarchical 0.1.0

package/worker.ts

@@ -0,0 +1,457 @@
+/**
+ * Background worker for hierarchical memory summarization.
+ *
+ * Runs on a timer, finds eligible chunks, summarizes them,
+ * and merges summaries when thresholds are reached.
+ */
+
+import fs from "node:fs";
+import path from "node:path";
+import { SessionManager } from "@mariozechner/pi-coding-agent";
+import type { PluginConfig } from "./config.js";
+import { resolveHierarchicalMemoryConfig } from "./config.js";
+import { acquireSummaryLock } from "./lock.js";
+import {
+  generateNextSummaryId,
+  loadSummaryContents,
+  loadSummaryIndex,
+  saveSummaryIndex,
+  writeSummary,
+} from "./storage.js";
+import {
+  type ChunkToSummarize,
+  estimateMessagesTokens,
+  getNextLevel,
+  getSourceLevel,
+  mergeSummaries,
+  summarizeChunk,
+  type SummarizationParams,
+} from "./summarize.js";
+import {
+  getAllSummariesForContext,
+  getUnmergedSummaries,
+  type HierarchicalMemoryConfig,
+  type SummaryEntry,
+  type SummaryIndex,
+} from "./types.js";
+
+const DEFAULT_AGENT_ID = "main";
+
+export type WorkerResult = {
+  success: boolean;
+  skipped?: "disabled" | "lock_held" | "no_session";
+  chunksProcessed?: number;
+  mergesPerformed?: number;
+  error?: string;
+  durationMs?: number;
+};
+
+/**
+ * Run the hierarchical memory worker for an agent.
+ */
+export async function runHierarchicalMemoryWorker(params: {
+  agentId: string;
+  pluginConfig: PluginConfig;
+  stateDir: string;
+  signal?: AbortSignal;
+}): Promise<WorkerResult> {
+  const startTime = Date.now();
+  const memoryConfig = resolveHierarchicalMemoryConfig(params.pluginConfig);
+
+  // Acquire lock to prevent concurrent runs
+  const lock = await acquireSummaryLock(params.agentId);
+  if (!lock) {
+    return { success: true, skipped: "lock_held" };
+  }
+
+  try {
+    const result = await runWorkerWithLock({
+      agentId: params.agentId,
+      pluginConfig: params.pluginConfig,
+      memoryConfig,
+      stateDir: params.stateDir,
+      signal: params.signal,
+    });
+
+    return {
+      ...result,
+      durationMs: Date.now() - startTime,
+    };
+  } finally {
+    await lock.release();
+  }
+}
+
+async function runWorkerWithLock(params: {
+  agentId: string;
+  pluginConfig: PluginConfig;
+  memoryConfig: HierarchicalMemoryConfig;
+  stateDir: string;
+  signal?: AbortSignal;
+}): Promise<Omit<WorkerResult, "durationMs">> {
+  const { agentId, pluginConfig, memoryConfig, stateDir, signal } = params;
+
+  try {
+    // Load current state
+    const index = await loadSummaryIndex(agentId);
+
+    // Find the current session
+    const storePath = resolveSessionStorePath(stateDir, agentId);
+    const sessionStore = loadSessionStoreSimple(storePath);
+    const mainSessionKey = `agent:${agentId}:main`;
+    const sessionEntry = sessionStore[mainSessionKey];
+
+    if (!sessionEntry?.sessionId) {
+      return { success: true, skipped: "no_session" };
+    }
+
+    const sessionFile = resolveSessionTranscriptPath(stateDir, agentId, sessionEntry.sessionId);
+
+    // Get summarization params
+    const summarization = resolveSummarizationParams({
+      pluginConfig,
+      memoryConfig,
+      sessionEntry,
+    });
+
+    if (!summarization) {
+      return {
+        success: false,
+        error: "Failed to resolve summarization parameters (no API key?)",
+      };
+    }
+
+    // Phase 1: Find and summarize eligible chunks
+    const chunks = await findEligibleChunks({
+      sessionFile,
+      lastSummarizedEntryId: index.lastSummarizedEntryId,
+      memoryConfig,
+      sessionId: sessionEntry.sessionId,
+    });
+
+    let chunksProcessed = 0;
+    for (const chunk of chunks) {
+      if (signal?.aborted) {
+        break;
+      }
+
+      // Load prior summaries for context
+      const summaryContext = getAllSummariesForContext(index);
+      const priorSummaries = [
+        ...(await loadSummaryContents(summaryContext.L3, agentId)),
+        ...(await loadSummaryContents(summaryContext.L2, agentId)),
+        ...(await loadSummaryContents(summaryContext.L1, agentId)),
+      ];
+
+      // Summarize the chunk
+      const summaryContent = await summarizeChunk({
+        chunk,
+        priorSummaries,
+        config: memoryConfig,
+        summarization,
+      });
+
+      // Create and save the summary entry
+      const summaryId = generateNextSummaryId(index, "L1");
+      const entry: SummaryEntry = {
+        id: summaryId,
+        level: "L1",
+        createdAt: Date.now(),
+        tokenEstimate: Math.ceil(summaryContent.length / 4), // Rough estimate
+        sourceLevel: "L0",
+        sourceIds: chunk.entryIds,
+        sourceSessionId: chunk.sessionId,
+        mergedInto: null,
+      };
+
+      await writeSummary(entry, summaryContent, agentId);
+      index.levels.L1.push(entry);
+      index.lastSummarizedEntryId = chunk.entryIds[chunk.entryIds.length - 1];
+      index.lastSummarizedSessionId = chunk.sessionId;
+
+      chunksProcessed++;
+    }
+
+    // Phase 2: Check for merges at each level
+    let mergesPerformed = 0;
+
+    for (const level of ["L1", "L2"] as const) {
+      if (signal?.aborted) {
+        break;
+      }
+
+      const merged = await maybeMergeLevel({
+        index,
+        level,
+        memoryConfig,
+        summarization,
+        agentId,
+      });
+
+      if (merged) {
+        mergesPerformed++;
+      }
+    }
+
+    // Save updated index
+    index.worker.lastRunAt = Date.now();
+    index.worker.lastError = null;
+    await saveSummaryIndex(index, agentId);
+
+    return {
+      success: true,
+      chunksProcessed,
+      mergesPerformed,
+    };
+  } catch (err) {
+    const errorMessage = err instanceof Error ? err.message : String(err);
+
+    // Try to save error state
+    try {
+      const index = await loadSummaryIndex(agentId);
+      index.worker.lastRunAt = Date.now();
+      index.worker.lastError = errorMessage;
+      await saveSummaryIndex(index, agentId);
+    } catch {
+      // Ignore save errors
+    }
+
+    return {
+      success: false,
+      error: errorMessage,
+    };
+  }
+}
+
+/**
+ * Find chunks of messages eligible for summarization.
+ */
+async function findEligibleChunks(params: {
+  sessionFile: string;
+  lastSummarizedEntryId: string | null;
+  memoryConfig: HierarchicalMemoryConfig;
+  sessionId: string;
+}): Promise<ChunkToSummarize[]> {
+  const { sessionFile, lastSummarizedEntryId, memoryConfig, sessionId } = params;
+
+  let sessionManager: SessionManager;
+  try {
+    sessionManager = SessionManager.open(sessionFile);
+  } catch {
+    return []; // Session file doesn't exist or is invalid
+  }
+
+  try {
+    const entries = sessionManager.getEntries();
+
+    if (entries.length === 0) {
+      return [];
+    }
+
+    // Walk entries from the end to find the cutoff index: everything before
+    // this index is "old enough" to summarize (at least pruningBoundaryTokens
+    // behind the conversation head). This uses the same entry stream we iterate
+    // so the boundary is correctly aligned.
+    let tailTokens = 0;
+    let cutoffIndex = -1;
+    for (let i = entries.length - 1; i >= 0; i--) {
+      const entry = entries[i];
+      if (entry.type !== "message") {
+        continue;
+      }
+      const msg = entry.message as { role: string; content?: unknown };
+      tailTokens += estimateMessagesTokens([msg]);
+      if (tailTokens >= memoryConfig.pruningBoundaryTokens) {
+        cutoffIndex = i;
+        break;
+      }
+    }
+
+    if (cutoffIndex <= 0) {
+      return []; // Not enough history yet
+    }
+
+    // Find start position (after lastSummarizedEntryId)
+    let startIndex = 0;
+    if (lastSummarizedEntryId) {
+      const lastIdx = entries.findIndex((e) => e.id === lastSummarizedEntryId);
+      if (lastIdx >= 0) {
+        startIndex = lastIdx + 1;
+      }
+    }
+
+    // Build chunks from eligible entries (startIndex..cutoffIndex)
+    const chunks: ChunkToSummarize[] = [];
+    let currentChunk: ChunkToSummarize = {
+      messages: [],
+      entryIds: [],
+      sessionId,
+      tokenEstimate: 0,
+    };
+
+    for (let i = startIndex; i < cutoffIndex; i++) {
+      const entry = entries[i];
+
+      // Skip non-message entries
+      if (entry.type !== "message") {
+        continue;
+      }
+
+      const msg = entry.message as { role: string; content?: unknown };
+      const msgTokens = estimateMessagesTokens([msg]);
+
+      currentChunk.messages.push(msg);
+      currentChunk.entryIds.push(entry.id);
+      currentChunk.tokenEstimate += msgTokens;
+
+      // Check if chunk is big enough
+      if (currentChunk.tokenEstimate >= memoryConfig.chunkTokens) {
+        // Ensure we end on a complete turn (assistant message)
+        if (msg.role === "assistant") {
+          chunks.push(currentChunk);
+          currentChunk = {
+            messages: [],
+            entryIds: [],
+            sessionId,
+            tokenEstimate: 0,
+          };
+        }
+      }
+    }
+
+    // Don't add partial chunks - they'll be picked up next time
+
+    return chunks;
+  } finally {
+    // SessionManager doesn't have dispose in all versions, but we're done with it
+  }
+}
+
+/**
+ * Merge summaries at a level if threshold is reached.
+ */
+async function maybeMergeLevel(params: {
+  index: SummaryIndex;
+  level: "L1" | "L2";
+  memoryConfig: HierarchicalMemoryConfig;
+  summarization: SummarizationParams;
+  agentId: string;
+}): Promise<boolean> {
+  const { index, level, memoryConfig, summarization, agentId } = params;
+
+  const unmerged = getUnmergedSummaries(index, level);
+
+  if (unmerged.length < memoryConfig.mergeThreshold) {
+    return false; // Not enough to merge yet
+  }
+
+  const nextLevel = getNextLevel(level);
+  if (!nextLevel) {
+    return false; // Already at max level
+  }
+
+  // Take exactly mergeThreshold entries to maintain fixed merge cadence
+  const toMerge = unmerged.slice(0, memoryConfig.mergeThreshold);
+
+  // Load summary contents
+  const summaryContents = await loadSummaryContents(toMerge, agentId);
+
+  // Load older context (unmerged higher-level summaries only)
+  const olderContext: string[] = [];
+  if (nextLevel === "L2") {
+    olderContext.push(...(await loadSummaryContents(getUnmergedSummaries(index, "L3"), agentId)));
+  }
+  if (nextLevel === "L3") {
+    // L3 has no older context
+  }
+
+  // Merge summaries
+  const mergedContent = await mergeSummaries({
+    summaries: summaryContents,
+    olderContext,
+    config: memoryConfig,
+    summarization,
+  });
+
+  // Create merged entry
+  const mergedId = generateNextSummaryId(index, nextLevel);
+  const mergedEntry: SummaryEntry = {
+    id: mergedId,
+    level: nextLevel,
+    createdAt: Date.now(),
+    tokenEstimate: Math.ceil(mergedContent.length / 4),
+    sourceLevel: getSourceLevel(nextLevel),
+    sourceIds: toMerge.map((s) => s.id),
+    mergedInto: null,
+  };
+
+  // Save merged summary
+  await writeSummary(mergedEntry, mergedContent, agentId);
+  index.levels[nextLevel].push(mergedEntry);
+
+  // Mark source summaries as merged
+  for (const summary of toMerge) {
+    summary.mergedInto = mergedId;
+  }
+
+  return true;
+}
+
+/**
+ * Resolve parameters needed for summarization.
+ * Uses the API key from plugin config instead of the complex auth system.
+ */
+function resolveSummarizationParams(params: {
+  pluginConfig: PluginConfig;
+  memoryConfig: HierarchicalMemoryConfig;
+  sessionEntry: { model?: string };
+}): SummarizationParams | null {
+  const { pluginConfig, memoryConfig, sessionEntry } = params;
+
+  // Determine model to use
+  const modelSpec =
+    memoryConfig.model ?? sessionEntry.model ?? "anthropic/claude-sonnet-4-5-20250929";
+  const [provider, model] = modelSpec.includes("/")
+    ? modelSpec.split("/", 2)
+    : ["anthropic", modelSpec];
+
+  const apiKey = pluginConfig.apiKey;
+  if (!apiKey) {
+    return null;
+  }
+
+  return {
+    model,
+    provider,
+    apiKey,
+  };
+}
+
+// --- Inline session helpers (replacing core imports) ---
+
+/** Resolve the session store path for an agent */
+function resolveSessionStorePath(stateDir: string, agentId?: string): string {
+  const id = agentId ?? DEFAULT_AGENT_ID;
+  return path.join(stateDir, "agents", id, "sessions", "sessions.json");
+}
+
+/** Resolve the session transcript JSONL path */
+function resolveSessionTranscriptPath(
+  stateDir: string,
+  agentId: string,
+  sessionId: string,
+): string {
+  return path.join(stateDir, "agents", agentId, "sessions", `${sessionId}.jsonl`);
+}
+
+/** Load the session store (simplified — just reads JSON) */
+function loadSessionStoreSimple(
+  storePath: string,
+): Record<string, { sessionId?: string; model?: string }> {
+  try {
+    const raw = fs.readFileSync(storePath, "utf-8");
+    return JSON.parse(raw) as Record<string, { sessionId?: string; model?: string }>;
+  } catch {
+    return {};
+  }
+}