pi-mega-compact 0.4.0

package/src/store.ts

@@ -0,0 +1,192 @@
+/**
+ * store.ts — persistence primitives for checkpoints + session state.
+ *
+ * Mirrors memory-mcp session_context.py: sessions normalize to `sess_xxx`,
+ * checkpoints are sequential `chkpt_001` per session. State lives under
+ * ~/.pi/agent/extensions/mega-compact/ as gzipped JSON.
+ */
+
+import { randomBytes } from "node:crypto";
+import { readFileSync, writeFileSync, mkdirSync, existsSync } from "node:fs";
+import { join } from "node:path";
+import { homedir } from "node:os";
+import { compressSmart, decompressSmart } from "./store/compression.js";
+
+// Re-export the compression primitives from their extracted home so existing
+// imports (`import { compressSmart } from "./store.js"`) keep working.
+export {
+  compressSmart,
+  decompressSmart,
+  compressZstd,
+  compressZstdMax,
+  decompressZstd,
+  isVersioned,
+  isZstd,
+  detectFormat,
+  decompressSyncAuto,
+} from "./store/compression.js";
+export type { CompressedFormat } from "./store/compression.js";
+
+/**
+ * State directory. Read lazily (per call) so tests can redirect it via
+ * MEGACOMPACT_STATE_DIR without re-importing the module. Defaults to
+ * ~/.pi/agent/extensions/mega-compact/.
+ */
+export function getStateDir(): string {
+  return process.env.MEGACOMPACT_STATE_DIR ?? join(homedir(), ".pi", "agent", "extensions", "mega-compact");
+}
+
+/** Normalize an arbitrary session id to the `sess_xxx` form (port of memory-mcp). */
+export function normalizeSessionId(sessionId: string | undefined | null): string {
+  if (!sessionId) return `sess_${randomBytes(8).toString("hex")}`;
+  if (sessionId.startsWith("sess_")) return sessionId;
+  if (sessionId.length >= 32 && sessionId.includes("-")) {
+    return `sess_${sessionId.replace(/-/g, "").slice(0, 16)}`;
+  }
+  return `sess_${sessionId}`;
+}
+
+export interface StoredCheckpoint {
+  checkpointId: string;
+  sessionId: string;
+  summary: string;
+  /** Compressed topic summary (extractive, ~2K tokens vs ~70K raw). */
+  topicSummary?: string;
+  /** SHA-256 of topicSummary — summary-content dedup key. */
+  summaryHash?: string;
+  keyDecisions: string[];
+  nextSteps: string[];
+  filesModified: string[];
+  tokenEstimate: number;
+  regionHash: string;
+  /** Primary content-addressable hash (full 64-hex SHA-256 of normalized text). */
+  contentHash?: string;
+  /** Secondary independent hash (reversed-text view) — guards single-hash collisions. */
+  contentHash2?: string;
+  contentHashVersion?: number;
+  /** Whitespace/ANSI-normalized text the content hashes were computed over. */
+  normalizedText?: string;
+  /** Reconstructible raw region (sync-compressed) for audit/replay/re-summarize. */
+  compressedOriginal?: Buffer;
+  embedding: number[];
+  timestamp: number;
+  /** Dedup lifecycle: 'active' | 'removed' (SemDeDup) | 'dup-resolved' (backfill). */
+  dedupStatus?: string;
+}
+
+// ---------------------------------------------------------------------------
+// JSON persistence (DR snapshots + migration source).
+// Compression lives in ./store/compression.ts; compressSmart/decompressSmart
+// are imported above and re-exported for backward-compatible call sites.
+// ---------------------------------------------------------------------------
+
+/**
+ * Read smart-compressed JSON, returning `fallback` on missing/corrupt file.
+ *
+ * Backward-compatible: detects legacy gzip files (magic byte 0x1f) and
+ * decompresses them correctly alongside new tagged files.
+ */
+export function readGzJson<T>(path: string, fallback: T): T {
+  try {
+    if (!existsSync(path)) return fallback;
+    const buf = readFileSync(path);
+    const out = decompressSmart(buf);
+    return JSON.parse(out.toString("utf-8")) as T;
+  } catch {
+    return fallback;
+  }
+}
+
+/** Write JSON with dynamic compression; creates parent dirs. */
+export function writeGzJson(path: string, data: unknown): void {
+  mkdirSync(join(path, ".."), { recursive: true });
+  const jsonBuf = Buffer.from(JSON.stringify(data), "utf-8");
+  const compressed = compressSmart(jsonBuf);
+  writeFileSync(path, compressed);
+}
+
+/** Append a checkpoint to the per-session checkpoint file (gzipped). */
+export function appendCheckpoint(cp: StoredCheckpoint, stateDir: string = getStateDir()): void {
+  const file = join(stateDir, `${cp.sessionId}.checkpoints.json.gz`);
+  const existing = readGzJson<StoredCheckpoint[]>(file, []);
+  existing.push(cp);
+  writeGzJson(file, existing);
+}
+
+/** All checkpoints for a session (across branches). */
+export function listCheckpoints(sessionId: string, stateDir: string = getStateDir()): StoredCheckpoint[] {
+  const file = join(stateDir, `${normalizeSessionId(sessionId)}.checkpoints.json.gz`);
+  return readGzJson<StoredCheckpoint[]>(file, []);
+}
+
+/**
+ * Rewrite ALL checkpoints for a session (in-place update).
+ *
+ * Used by VectorStore when summaryHash or contentSimilarity dedup updates an
+ * existing checkpoint's timestamp/metadata instead of creating a new one.
+ */
+export function rewriteCheckpoints(
+  sessionId: string,
+  checkpoints: StoredCheckpoint[],
+  stateDir: string = getStateDir(),
+): void {
+  const file = join(stateDir, `${normalizeSessionId(sessionId)}.checkpoints.json.gz`);
+  writeGzJson(file, checkpoints);
+}
+
+/**
+ * Generate the next sequential checkpoint id for a session (chkpt_001 ...).
+ */
+export function nextCheckpointId(sessionId: string, stateDir: string = getStateDir()): string {
+  const list = listCheckpoints(sessionId, stateDir);
+  const max = list.reduce((m, c) => {
+    const n = parseInt(c.checkpointId.replace("chkpt_", ""), 10);
+    return Number.isFinite(n) && n > m ? n : m;
+  }, 0);
+  return `chkpt_${String(max + 1).padStart(3, "0")}`;
+}
+
+export interface SessionState {
+  /** checkpointIds already injected into the current window this session. */
+  injectedCheckpointIds: string[];
+  /** regionHashes already represented (for sentinel dedup). */
+  storedRegionHashes: string[];
+}
+
+/** Load mutable session state (created on demand). */
+export function loadSessionState(sessionId: string, stateDir: string = getStateDir()): SessionState {
+  const file = join(stateDir, `${normalizeSessionId(sessionId)}.state.json.gz`);
+  return readGzJson<SessionState>(file, {
+    injectedCheckpointIds: [],
+    storedRegionHashes: [],
+  });
+}
+
+export function saveSessionState(sessionId: string, state: SessionState, stateDir: string = getStateDir()): void {
+  const file = join(stateDir, `${normalizeSessionId(sessionId)}.state.json.gz`);
+  writeGzJson(file, state);
+}
+
+/**
+ * Cumulative store-wide dedup accounting. Lives outside the per-session
+ * SessionState (which resets on every session instance) so the dedup rate
+ * stays stable across session restarts. Persisted as plain JSON in the
+ * state dir.
+ */
+export interface DedupStats {
+  /** Total add() calls (new checkpoints + deduped collapses). */
+  attempts: number;
+  /** add() calls that collapsed onto an existing checkpoint. */
+  deduped: number;
+}
+
+const DEDUP_STATS_FILE = "dedup-stats.json";
+
+export function loadDedupStats(stateDir: string = getStateDir()): DedupStats {
+  const file = join(stateDir, DEDUP_STATS_FILE);
+  return readGzJson<DedupStats>(file, { attempts: 0, deduped: 0 });
+}
+
+export function saveDedupStats(stats: DedupStats, stateDir: string = getStateDir()): void {
+  writeGzJson(join(stateDir, DEDUP_STATS_FILE), stats);
+}

package/src/supersede.test.ts

@@ -0,0 +1,42 @@
+import { test } from "node:test";
+import assert from "node:assert/strict";
+import type { EngineMessage } from "./types.js";
+import { findSuperseded, supersede } from "./supersede.js";
+
+function msg(role: EngineMessage["role"], text: string): EngineMessage { return { role, text }; }
+
+test("read superseded by later write to same path is pruned", () => {
+  const messages = [
+    msg("assistant", "read src/server.ts"),
+    msg("user", "now change it"),
+    msg("assistant", "write src/server.ts with the fix"),
+  ];
+  assert.deepEqual(findSuperseded(messages), [0]);
+});
+
+test("older read superseded by newer read of same path (keep latest)", () => {
+  const messages = [
+    msg("assistant", "read src/a.ts"),
+    msg("assistant", "read src/a.ts again"),
+  ];
+  assert.deepEqual(findSuperseded(messages), [0]);
+});
+
+test("unrelated reads are kept", () => {
+  const messages = [
+    msg("assistant", "read src/a.ts"),
+    msg("assistant", "read src/b.ts"),
+  ];
+  assert.deepEqual(findSuperseded(messages), []);
+});
+
+test("supersede() drops the obsolete read and preserves order", () => {
+  const messages = [
+    msg("assistant", "read src/server.ts"),
+    msg("user", "change it"),
+    msg("assistant", "write src/server.ts done"),
+  ];
+  const out = supersede(messages);
+  assert.equal(out.length, 2);
+  assert.equal(out[0].text, "change it");
+});

package/src/supersede.ts

@@ -0,0 +1,67 @@
+/**
+ * supersede.ts — Layer 1 (SUPERSEDE): zero-cost factual pruning.
+ *
+ * If you read server.py, the previous read of server.py is factually obsolete
+ * once you write it. We detect file-read turns that are superseded by a later
+ * write (or a later read) to the same path and mark them for pruning — no
+ * summarization, no token cost. Mirrors memory-mcp MemoryCompactor Stage 1.
+ */
+
+import type { EngineMessage } from "./types.js";
+import { extractFileCandidates } from "./compact.js";
+
+/** Classify a message's relationship to a file path. */
+function fileOps(msg: EngineMessage): { path: string; op: "read" | "write" }[] {
+  const paths = extractFileCandidates(msg.text);
+  if (paths.length === 0) return [];
+  const low = msg.text.toLowerCase();
+  const isWrite = /\b(write|edit|create|save|append|overwrite|update|patch|modify)\b/.test(low);
+  return paths.map((p) => ({ path: p, op: isWrite ? "write" : "read" }));
+}
+
+/**
+ * Return the indexes (into `messages`) of file-read turns superseded by a later
+ * operation on the same path. A read is obsolete once a write touches the same
+ * path, or once a newer read of the same path exists (keep only the latest).
+ */
+export function findSuperseded(messages: EngineMessage[]): number[] {
+  // Build, per path, the latest operation index and whether a write occurred.
+  const lastWriteAt = new Map<string, number>();
+  const lastReadAt = new Map<string, number>();
+  messages.forEach((m, i) => {
+    for (const { path, op } of fileOps(m)) {
+      if (op === "write") lastWriteAt.set(path, i);
+      else lastReadAt.set(path, i);
+    }
+  });
+
+  const superseded = new Set<number>();
+  // Reads before a write to the same path are obsolete.
+  messages.forEach((m, i) => {
+    for (const { path, op } of fileOps(m)) {
+      if (op !== "read") continue;
+      const writeAt = lastWriteAt.get(path);
+      if (writeAt !== undefined && i < writeAt) superseded.add(i);
+    }
+  });
+  // For paths with multiple reads and no write, keep only the latest read.
+  const readsByPath = new Map<string, number[]>();
+  messages.forEach((m, i) => {
+    for (const { path, op } of fileOps(m)) {
+      if (op !== "read") continue;
+      if (!readsByPath.has(path)) readsByPath.set(path, []);
+      readsByPath.get(path)!.push(i);
+    }
+  });
+  for (const [, idxs] of readsByPath) {
+    if (idxs.length > 1) idxs.slice(0, -1).forEach((i) => superseded.add(i));
+  }
+
+  return [...superseded].sort((a, b) => a - b);
+}
+
+/** Convenience: drop superseded messages, preserving order. */
+export function supersede(messages: EngineMessage[]): EngineMessage[] {
+  const drop = new Set(findSuperseded(messages));
+  return messages.filter((_m, i) => !drop.has(i));
+}

package/src/tokens.ts

@@ -0,0 +1,35 @@
+/**
+ * Token estimation — rough, deterministic, no LLM.
+ *
+ * Ported from claw-code rusty-claude-cli compact.rs estimate_message_tokens
+ * (len/4 + 1 per content block). Good enough to gate compaction and to report
+ * tokens-saved; NOT a substitute for a real tokenizer.
+ */
+
+/** Estimate tokens for a single text/tool block already as a string. */
+export function estimateBlockTokens(text: string): number {
+  return Math.floor(text.length / 4) + 1;
+}
+
+/**
+ * Estimate tokens for an EngineMessage. Tool-use/result blocks carry name +
+ * input/output strings, mirrored from the claw-code block accounting.
+ */
+export function estimateMessageTokens(msg: {
+  text?: string;
+  toolName?: string;
+  input?: string;
+  output?: string;
+}): number {
+  let t = 0;
+  if (msg.text) t += estimateBlockTokens(msg.text);
+  if (msg.toolName) t += estimateBlockTokens(msg.toolName);
+  if (msg.input) t += estimateBlockTokens(msg.input);
+  if (msg.output) t += estimateBlockTokens(msg.output);
+  return t;
+}
+
+/** Sum tokens over a list of messages. */
+export function estimateSessionTokens(messages: Array<{ text?: string; toolName?: string; input?: string; output?: string }>): number {
+  return messages.reduce((acc, m) => acc + estimateMessageTokens(m), 0);
+}

package/src/types.test.ts

@@ -0,0 +1,10 @@
+import { test } from "node:test";
+import assert from "node:assert/strict";
+
+test("EngineMessage roles are constrained to the pi message contract (no system role)", () => {
+  const roles = ["user", "assistant", "tool", "custom"] as const;
+  // pi Message = user|assistant|tool (+ custom for markers). There is no system role.
+  const probe: string = "system";
+  assert.equal((roles as readonly string[]).includes(probe), false);
+  assert.equal(roles.length, 4);
+});

package/src/types.ts

@@ -0,0 +1,49 @@
+/**
+ * Shared internal types for the pi-mega-compact engine.
+ *
+ * Kept independent of pi's runtime types (src/ is pi-agnostic; the extension
+ * entry in extensions/ adapts between the two). See RESEARCH.md for the pi
+ * AgentMessage contract this must eventually satisfy.
+ */
+
+/** A lightweight message shape the engine reasons about (role only matters). */
+export interface EngineMessage {
+  role: "user" | "assistant" | "tool" | "custom";
+  text: string;
+  toolName?: string;
+  /** Tool input/output payload (for tool-use / tool-result roles). */
+  input?: string;
+  output?: string;
+}
+
+/** A persisted compaction checkpoint (Layer 4 → vector store). */
+export interface Checkpoint {
+  checkpointId: string; // chkpt_001
+  sessionId: string; // sess_xxx (normalized)
+  summary: string;
+  keyDecisions: string[];
+  nextSteps: string[];
+  filesModified: string[];
+  tokenEstimate: number;
+  regionHash: string; // dedup sentinel key
+  timestamp: number;
+}
+
+/** Result of a Trident run over a set of messages. */
+export interface TridentResult {
+  superseded: string[];
+  collapsed: string;
+  checkpoints: Checkpoint[];
+}
+
+/** Compaction config (mirrors claw-code CompactionConfig + memory-mcp knobs). */
+export interface CompactionConfig {
+  /** Preserve the most recent N messages verbatim. */
+  preserveRecent: number;
+  /** Token budget that triggers compaction. */
+  maxEstimatedTokens: number;
+  /** Number of recent user messages to never drop (anchor floor). */
+  anchorUserMessages: number;
+  /** Min chatty messages before COLLAPSE summarizes them. */
+  collapseThreshold: number;
+}