@pi-unipi/compactor 0.1.7 → 0.2.1

package/src/session/db.ts

@@ -69,7 +69,7 @@ const DEDUP_WINDOW = 5;
 
 export class SessionDB {
   private db: any;
-  private stmts: Map<string, PreparedStatement> = new Map();
+  private stmts: Map<string, PreparedStatement> | null = null;
   private dbPath: string;
 
   constructor(opts?: { dbPath?: string }) {
@@ -101,11 +101,8 @@ export class SessionDB {
         priority INTEGER NOT NULL DEFAULT 2,
         data TEXT NOT NULL,
         project_dir TEXT NOT NULL DEFAULT '',
-        attribution_source TEXT NOT NULL DEFAULT 'unknown',
-        attribution_confidence REAL NOT NULL DEFAULT 0,
         source_hook TEXT NOT NULL,
-        created_at TEXT NOT NULL DEFAULT (datetime('now')),
-        data_hash TEXT NOT NULL DEFAULT ''
+        created_at TEXT NOT NULL DEFAULT (datetime('now'))
       );
       CREATE INDEX IF NOT EXISTS idx_session_events_session ON session_events(session_id);
       CREATE INDEX IF NOT EXISTS idx_session_events_type ON session_events(session_id, type);
@@ -129,11 +126,42 @@ export class SessionDB {
         consumed INTEGER NOT NULL DEFAULT 0
       );
     `);
+
+    // Run version-gated schema migrations
+    this.runMigrations();
+  }
+
+  /** Run version-gated schema migrations using PRAGMA user_version. */
+  private runMigrations(): void {
+    const currentVersion = this.db.pragma("user_version", { simple: true }) as number;
+
+    if (currentVersion < 1) {
+      // V1: Add columns introduced by compactor gap analysis (2026-04-30)
+      // Each ALTER TABLE is wrapped individually — SQLite auto-commits DDL,
+      // so a partial failure from a prior run would leave some columns added
+      // and others not. We catch "duplicate column" to handle this safely.
+      const safeAddColumn = (table: string, col: string, def: string) => {
+        try {
+          this.db.exec(`ALTER TABLE ${table} ADD COLUMN ${col} ${def}`);
+        } catch (e: any) {
+          if (e?.message?.includes("duplicate column")) return;
+          throw e;
+        }
+      };
+      safeAddColumn("session_meta", "total_chars_before", "INTEGER NOT NULL DEFAULT 0");
+      safeAddColumn("session_meta", "total_chars_kept", "INTEGER NOT NULL DEFAULT 0");
+      safeAddColumn("session_meta", "total_messages_summarized", "INTEGER NOT NULL DEFAULT 0");
+      safeAddColumn("session_events", "attribution_source", "TEXT NOT NULL DEFAULT 'unknown'");
+      safeAddColumn("session_events", "attribution_confidence", "REAL NOT NULL DEFAULT 0");
+      safeAddColumn("session_events", "data_hash", "TEXT NOT NULL DEFAULT ''");
+      this.db.pragma("user_version = 1");
+    }
   }
 
   private prepareStatements(): void {
+    this.stmts = new Map();
     const p = (key: string, sql: string) => {
-      this.stmts.set(key, this.db.prepare(sql) as PreparedStatement);
+      this.stmts!.set(key, this.db.prepare(sql) as PreparedStatement);
     };
 
     p("insertEvent", `INSERT INTO session_events (session_id, type, category, priority, data, project_dir, attribution_source, attribution_confidence, source_hook, data_hash) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`);
@@ -143,8 +171,10 @@ export class SessionDB {
     p("evictLowestPriority", `DELETE FROM session_events WHERE id = (SELECT id FROM session_events WHERE session_id = ? ORDER BY priority ASC, id ASC LIMIT 1)`);
     p("updateMetaLastEvent", `UPDATE session_meta SET last_event_at = datetime('now'), event_count = event_count + 1 WHERE session_id = ?`);
     p("ensureSession", `INSERT OR IGNORE INTO session_meta (session_id, project_dir) VALUES (?, ?)`);
-    p("getSessionStats", `SELECT session_id, project_dir, started_at, last_event_at, event_count, compact_count FROM session_meta WHERE session_id = ?`);
+    p("getSessionStats", `SELECT session_id, project_dir, started_at, last_event_at, event_count, compact_count, total_chars_before, total_chars_kept, total_messages_summarized FROM session_meta WHERE session_id = ?`);
     p("incrementCompactCount", `UPDATE session_meta SET compact_count = compact_count + 1 WHERE session_id = ?`);
+    p("addCompactionStats", `UPDATE session_meta SET total_chars_before = total_chars_before + ?, total_chars_kept = total_chars_kept + ?, total_messages_summarized = total_messages_summarized + ? WHERE session_id = ?`);
+    p("getAllTimeStats", `SELECT COALESCE(SUM(total_chars_before), 0) AS all_chars_before, COALESCE(SUM(total_chars_kept), 0) AS all_chars_kept, COALESCE(SUM(total_messages_summarized), 0) AS all_messages_summarized, COALESCE(SUM(compact_count), 0) AS all_compactions FROM session_meta`);
     p("upsertResume", `INSERT INTO session_resume (session_id, snapshot, event_count) VALUES (?, ?, ?) ON CONFLICT(session_id) DO UPDATE SET snapshot = excluded.snapshot, event_count = excluded.event_count, created_at = datetime('now'), consumed = 0`);
     p("getResume", `SELECT snapshot, event_count, consumed FROM session_resume WHERE session_id = ?`);
     p("markResumeConsumed", `UPDATE session_resume SET consumed = 1 WHERE session_id = ?`);
@@ -155,10 +185,11 @@ export class SessionDB {
   }
 
   private stmt(key: string): PreparedStatement {
-    return this.stmts.get(key)!;
+    return this.stmts!.get(key)!;
   }
 
   insertEvent(sessionId: string, event: SessionEvent, sourceHook: string = "PostToolUse"): void {
+    if (!this.stmts) return;
     const dataHash = createHash("sha256").update(event.data).digest("hex").slice(0, 16).toUpperCase();
     const projectDir = String(event.project_dir ?? "").trim();
     const attributionSource = String(event.attribution_source ?? "unknown");
@@ -185,42 +216,67 @@ export class SessionDB {
   }
 
   getEvents(sessionId: string, opts?: { type?: string; minPriority?: number; limit?: number }): StoredEvent[] {
+    if (!this.stmts) return [];
     const limit = opts?.limit ?? 1000;
     return this.stmt("getEvents").all(sessionId, limit) as StoredEvent[];
   }
 
   getEventCount(sessionId: string): number {
+    if (!this.stmts) return 0;
     const row = this.stmt("getEventCount").get(sessionId) as { cnt: number };
     return row.cnt;
   }
 
   ensureSession(sessionId: string, projectDir: string): void {
+    if (!this.stmts) return;
     this.stmt("ensureSession").run(sessionId, projectDir);
   }
 
   getSessionStats(sessionId: string): SessionMeta | null {
+    if (!this.stmts) return null;
     const row = this.stmt("getSessionStats").get(sessionId) as SessionMeta | undefined;
     return row ?? null;
   }
 
   incrementCompactCount(sessionId: string): void {
+    if (!this.stmts) return;
     this.stmt("incrementCompactCount").run(sessionId);
   }
 
+  addCompactionStats(sessionId: string, charsBefore: number, charsKept: number, messagesSummarized: number): void {
+    if (!this.stmts) return;
+    this.stmt("addCompactionStats").run(charsBefore, charsKept, messagesSummarized, sessionId);
+  }
+
+  getAllTimeStats(): { allCharsBefore: number; allCharsKept: number; allMessagesSummarized: number; allCompactions: number } {
+    if (!this.stmts) return { allCharsBefore: 0, allCharsKept: 0, allMessagesSummarized: 0, allCompactions: 0 };
+    const row = this.stmt("getAllTimeStats").get() as { all_chars_before: number; all_chars_kept: number; all_messages_summarized: number; all_compactions: number };
+    return {
+      allCharsBefore: row?.all_chars_before ?? 0,
+      allCharsKept: row?.all_chars_kept ?? 0,
+      allMessagesSummarized: row?.all_messages_summarized ?? 0,
+      allCompactions: row?.all_compactions ?? 0,
+    };
+  }
+
   upsertResume(sessionId: string, snapshot: string, eventCount?: number): void {
+    if (!this.stmts) return;
     this.stmt("upsertResume").run(sessionId, snapshot, eventCount ?? 0);
   }
 
   getResume(sessionId: string): ResumeRow | null {
+    if (!this.stmts) return null;
     const row = this.stmt("getResume").get(sessionId) as ResumeRow | undefined;
     return row ?? null;
   }
 
   markResumeConsumed(sessionId: string): void {
+    if (!this.stmts) return;
     this.stmt("markResumeConsumed").run(sessionId);
   }
 
   deleteSession(sessionId: string): void {
+    if (!this.stmts) return;
     this.db.transaction(() => {
       this.stmt("deleteEvents").run(sessionId);
       this.stmt("deleteResume").run(sessionId);
@@ -229,6 +285,7 @@ export class SessionDB {
   }
 
   cleanupOldSessions(maxAgeDays: number = 7): number {
+    if (!this.stmts) return 0;
     const oldSessions = this.stmt("getOldSessions").all(`-${maxAgeDays}`) as Array<{ session_id: string }>;
     for (const { session_id } of oldSessions) {
       this.deleteSession(session_id);

package/src/session/resume-inject.ts

@@ -4,6 +4,8 @@
 
 import type { SessionDB } from "./db.js";
 import { buildResumeSnapshot } from "./snapshot.js";
+import { buildAutoInjection } from "./auto-inject.js";
+import { loadConfig } from "../config/manager.js";
 
 export async function injectResumeSnapshot(
   db: SessionDB,
@@ -20,6 +22,16 @@ export async function injectResumeSnapshot(
     searchTool: opts?.searchTool ?? "ctx_search",
   });
 
+  // Auto-injection: add behavioral state after compaction (if enabled)
+  const config = loadConfig();
+  let fullSnapshot = snapshot;
+  if (config.pipeline.autoInjection) {
+    const autoInjection = buildAutoInjection(events);
+    if (autoInjection) {
+      fullSnapshot = `${snapshot}\n\n${autoInjection}`;
+    }
+  }
+
   db.markResumeConsumed(sessionId);
-  return snapshot;
+  return fullSnapshot;
 }

package/src/store/db-base.ts

@@ -42,6 +42,17 @@ export async function loadSQLite() {
 export function applyWALPragmas(db: any): void {
   db.exec("PRAGMA journal_mode = WAL;");
   db.exec("PRAGMA synchronous = NORMAL;");
+  // Memory-map the DB file for read-heavy FTS5 search workloads (if enabled)
+  try {
+    const { loadConfig } = require("../config/manager.js");
+    const config = loadConfig();
+    if (config.pipeline?.mmapPragma !== false) {
+      db.exec("PRAGMA mmap_size = 268435456;"); // 256MB
+    }
+  } catch {
+    // Fallback: always apply mmap if config can't be loaded
+    try { db.exec("PRAGMA mmap_size = 268435456;"); } catch { /* unsupported runtime */ }
+  }
 }
 
 export function withRetry<T>(fn: () => T, maxRetries = 3): T {

package/src/store/index.ts

@@ -8,6 +8,7 @@ import { loadSQLite, applyWALPragmas, withRetry, isSQLiteCorruptionError, defaul
 import type { PreparedStatement } from "./db-base.js";
 import { autoChunk } from "./chunking.js";
 import type { IndexResult, SearchResult, StoreStats } from "../types.js";
+import { loadConfig } from "../config/manager.js";
 
 // --- Fuzzy correction ---
 
@@ -131,6 +132,118 @@ function rrfMerge(
     .map((s) => ({ ...s.result, rank: s.score }));
 }
 
+// ── Proximity Reranking (from context-mode) ──────────────────
+
+/** Find all character positions of a term in text */
+function findAllPositions(text: string, term: string): number[] {
+  const positions: number[] = [];
+  let idx = text.indexOf(term);
+  while (idx !== -1) {
+    positions.push(idx);
+    idx = text.indexOf(term, idx + 1);
+  }
+  return positions;
+}
+
+/** Sweep-line algorithm to find minimum span covering all terms */
+function findMinSpan(positionLists: number[][]): number {
+  if (positionLists.length === 0) return Infinity;
+  if (positionLists.length === 1) return 0;
+
+  const sorted = positionLists.map((p) => [...p].sort((a, b) => a - b));
+  const ptrs = new Array(sorted.length).fill(0);
+  let minSpan = Infinity;
+
+  while (true) {
+    let curMin = Infinity;
+    let curMax = -Infinity;
+    let minIdx = 0;
+
+    for (let i = 0; i < sorted.length; i++) {
+      const val = sorted[i][ptrs[i]];
+      if (val < curMin) { curMin = val; minIdx = i; }
+      if (val > curMax) { curMax = val; }
+    }
+
+    const span = curMax - curMin;
+    if (span < minSpan) minSpan = span;
+
+    ptrs[minIdx]++;
+    if (ptrs[minIdx] >= sorted[minIdx].length) break;
+  }
+
+  return minSpan;
+}
+
+/** Count adjacent term pairs within a character gap */
+function countAdjacentPairs(
+  positionLists: number[][],
+  terms: string[],
+  gap: number = 30,
+): number {
+  if (positionLists.length < 2 || terms.length < 2) return 0;
+  let total = 0;
+  const pairs = Math.min(positionLists.length, terms.length) - 1;
+  for (let i = 0; i < pairs; i++) {
+    const left = positionLists[i];
+    const right = positionLists[i + 1];
+    const leftLen = terms[i].length;
+    let j = 0;
+    for (const p of left) {
+      const minStart = p + leftLen;
+      const maxStart = minStart + gap;
+      while (j < right.length && right[j] < minStart) j++;
+      if (j < right.length && right[j] <= maxStart) {
+        total++;
+        j++;
+      }
+    }
+  }
+  return total;
+}
+
+/** Apply proximity reranking to RRF results */
+function applyProximityReranking(
+  results: SearchResult[],
+  query: string,
+): SearchResult[] {
+  const allTerms = query
+    .toLowerCase()
+    .split(/\s+/)
+    .filter((w) => w.length >= 2);
+  const filtered = allTerms.filter((w) => !STOPWORDS.has(w));
+  const terms = filtered.length > 0 ? filtered : allTerms;
+
+  if (terms.length < 2) return results; // Single-term queries skip proximity
+
+  const scored = results.map((r) => {
+    const titleLower = r.title.toLowerCase();
+    const titleHits = terms.filter((t) => titleLower.includes(t)).length;
+    const titleWeight = r.contentType === "code" ? 0.6 : 0.3;
+    const titleBoost = titleHits > 0 ? titleWeight * (titleHits / terms.length) : 0;
+
+    let proximityBoost = 0;
+    let phraseBoost = 0;
+
+    const content = r.content.toLowerCase();
+    const positions = terms.map((t) => findAllPositions(content, t));
+
+    if (!positions.some((p) => p.length === 0)) {
+      const minSpan = findMinSpan(positions);
+      proximityBoost = 1 / (1 + minSpan / Math.max(content.length, 1));
+
+      const adjacentPairs = countAdjacentPairs(positions, terms);
+      phraseBoost = 0.5 * Math.min(1, adjacentPairs / 4);
+    }
+
+    return { result: r, boost: titleBoost + proximityBoost + phraseBoost };
+  });
+
+  return scored
+    .sort((a, b) => b.boost - a.boost || a.result.rank - b.result.rank)
+    .map((s) => s.result);
+}
+
 const STOPWORDS = new Set([
   "the", "and", "for", "are", "but", "not", "you", "all", "can", "had",
   "her", "was", "one", "our", "out", "has", "his", "how", "its", "may",
@@ -176,6 +289,7 @@ export class ContentStore {
   private stmts: Map<string, PreparedStatement> = new Map();
   private dbPath: string;
   private ready = false;
+  private writeCount = 0;
 
   constructor(opts?: { dbPath?: string }) {
     this.dbPath = opts?.dbPath ?? defaultDBPath("content");
@@ -227,6 +341,7 @@ export class ContentStore {
     p("deleteByLabel", `DELETE FROM content_fts WHERE label = ?`);
     p("insertSource", `INSERT INTO content_sources (label, source, content_type, mtime, sha256, chunk_count) VALUES (?, ?, ?, ?, ?, ?) ON CONFLICT(label) DO UPDATE SET source=excluded.source, content_type=excluded.content_type, mtime=excluded.mtime, sha256=excluded.sha256, chunk_count=excluded.chunk_count, indexed_at=datetime('now')`);
     p("getSource", `SELECT label, source, content_type, mtime, sha256, chunk_count, indexed_at FROM content_sources WHERE label = ?`);
+    p("getSourceMeta", `SELECT label, chunk_count, indexed_at FROM content_sources WHERE label = ?`);
     p("deleteSource", `DELETE FROM content_sources WHERE label = ?`);
     p("countSources", `SELECT COUNT(*) AS cnt FROM content_sources`);
     p("countFTS", `SELECT COUNT(*) AS cnt FROM content_fts`);
@@ -271,6 +386,7 @@ export class ContentStore {
     });
 
     withRetry(() => transaction());
+    this.afterWrite();
 
     return { sourceId: 1, label, totalChunks: chunks.length, codeChunks };
   }
@@ -312,7 +428,13 @@ export class ContentStore {
     // RRF fusion
     const rrfResults = rrfMerge(porterResults, trigramResults);
 
-    if (mode === "rrf") return rrfResults.slice(0, limit);
+    // Apply proximity reranking to all RRF results (if enabled)
+    const config = loadConfig();
+    const rerankedResults = config.pipeline.proximityReranking
+      ? applyProximityReranking(rrfResults, query)
+      : rrfResults;
+
+    if (mode === "rrf") return rerankedResults.slice(0, limit);
 
     // Fuzzy mode: apply fuzzy correction to query terms
     const vocab = buildVocabulary(allRows);
@@ -336,11 +458,11 @@ export class ContentStore {
         matchLayer: "fuzzy" as const,
         rank: r.rank * 0.9, // slightly lower confidence
       }));
-      const merged = rrfMerge(rrfResults, correctedResults);
-      return merged.slice(0, limit);
+      const merged = rrfMerge(rerankedResults, correctedResults);
+      return applyProximityReranking(merged, query).slice(0, limit);
     }
 
-    return rrfResults.slice(0, limit);
+    return rerankedResults.slice(0, limit);
   }
 
   async getStats(): Promise<StoreStats> {
@@ -354,13 +476,37 @@ export class ContentStore {
     };
   }
 
+  /** Get source metadata for TTL cache check */
+  getSourceMeta(label: string): { label: string; chunkCount: number; indexedAt: string } | null {
+    const row = this.stmt("getSourceMeta").get(label) as { label: string; chunk_count: number; indexed_at: string } | undefined;
+    if (!row) return null;
+    return { label: row.label, chunkCount: row.chunk_count, indexedAt: row.indexed_at };
+  }
+
   async purge(): Promise<number> {
     if (!this.ready) await this.init();
     this.db.exec(`DELETE FROM content_fts; DELETE FROM content_sources;`);
+    this.afterWrite();
     const row = this.stmt("countSources").get() as { cnt: number };
     return row.cnt;
   }
 
+  /** Run WAL checkpoint to prevent unbounded WAL file growth. */
+  checkpointWAL(mode: "PASSIVE" | "TRUNCATE" = "PASSIVE"): void {
+    if (!this.db) return;
+    try {
+      this.db.exec(`PRAGMA wal_checkpoint(${mode});`);
+    } catch { /* ignore */ }
+  }
+
+  /** Increment write counter and trigger PASSIVE checkpoint every 10th write. */
+  private afterWrite(): void {
+    this.writeCount++;
+    if (this.writeCount % 10 === 0) {
+      this.checkpointWAL("PASSIVE");
+    }
+  }
+
   close(): void {
     try { this.db.close(); } catch { /* ignore */ }
   }

package/src/store/unified.ts

@@ -0,0 +1,109 @@
+/**
+ * Unified search across ContentStore + SessionDB events
+ * Supports timeline (chronological) and relevance sorting
+ * (from context-mode src/search/unified.ts)
+ */
+
+import type { ContentStore } from "./index.js";
+import type { SessionDB } from "../session/db.js";
+import type { SearchResult } from "../types.js";
+
+export interface UnifiedSearchResult {
+  title: string;
+  content: string;
+  source: string;
+  origin: "current-session" | "prior-session";
+  timestamp: string;
+  rank: number;
+  matchLayer: string;
+  contentType: "prose" | "code";
+}
+
+export interface UnifiedSearchOptions {
+  query: string;
+  limit?: number;
+  sort?: "relevance" | "timeline";
+  source?: string;
+  contentType?: string;
+  projectDir?: string;
+}
+
+/**
+ * Search across multiple sources and optionally sort chronologically.
+ * - relevance: ContentStore only, ranked by RRF
+ * - timeline: ContentStore + SessionDB events, sorted by timestamp
+ */
+export async function searchAllSources(
+  store: ContentStore,
+  sessionDB: SessionDB | null,
+  opts: UnifiedSearchOptions,
+): Promise<UnifiedSearchResult[]> {
+  const limit = opts.limit ?? 10;
+  const sort = opts.sort ?? "relevance";
+  const sessionStartTime = new Date().toISOString();
+
+  const results: UnifiedSearchResult[] = [];
+
+  // Source 1: ContentStore (always, both modes)
+  try {
+    const storeResults = await store.search(opts.query, { limit });
+    results.push(
+      ...storeResults.map((r) => ({
+        title: r.title,
+        content: r.content,
+        source: r.source,
+        origin: "current-session" as const,
+        timestamp: sessionStartTime, // ContentStore doesn't track per-result timestamps yet
+        rank: r.rank,
+        matchLayer: r.matchLayer ?? "porter",
+        contentType: r.contentType,
+      })),
+    );
+  } catch {
+    // ContentStore search failed — continue with other sources
+  }
+
+  // Source 2: SessionDB events (timeline mode only)
+  if (sort === "timeline" && sessionDB) {
+    try {
+      const sessionId = opts.projectDir ?? "";
+      const events = sessionDB.getEvents(sessionId, { limit: 100 });
+      const queryLower = opts.query.toLowerCase();
+      const matchingEvents = events.filter((e) => {
+        const data = String(e.data ?? "").toLowerCase();
+        const type = String(e.type ?? "").toLowerCase();
+        const category = String(e.category ?? "").toLowerCase();
+        return data.includes(queryLower) || type.includes(queryLower) || category.includes(queryLower);
+      });
+
+      results.push(
+        ...matchingEvents.slice(0, limit).map((e) => ({
+          title: `[${e.category}] ${e.type}`,
+          content: String(e.data ?? "").slice(0, 500),
+          source: "prior-session",
+          origin: "prior-session" as const,
+          timestamp: e.created_at ?? sessionStartTime,
+          rank: 0,
+          matchLayer: "event",
+          contentType: "prose" as const,
+        })),
+      );
+    } catch {
+      // SessionDB search failed — continue
+    }
+  }
+
+  // Normalize SQLite datetime format to ISO 8601
+  for (const r of results) {
+    if (r.timestamp && !r.timestamp.includes("T")) {
+      r.timestamp = r.timestamp.replace(" ", "T") + "Z";
+    }
+  }
+
+  // Sort: timeline = chronological, relevance = by rank
+  if (sort === "timeline") {
+    results.sort((a, b) => (a.timestamp || "").localeCompare(b.timestamp || ""));
+  }
+
+  return results.slice(0, limit);
+}

package/src/tools/context-budget.ts

@@ -0,0 +1,50 @@
+/**
+ * context_budget tool — estimate remaining context window
+ */
+
+export interface ContextBudgetResult {
+  percentFull: number;
+  remainingTokens: number;
+  totalTokens: number;
+  message: string;
+  advice: string;
+}
+
+export function estimateContextBudget(
+  tokensBefore?: number,
+  contextWindowSize?: number,
+): ContextBudgetResult | null {
+  const windowSize = contextWindowSize ?? 200000; // Default 200K context
+  const used = tokensBefore ?? 0;
+
+  if (used <= 0 && tokensBefore === undefined) return null;
+
+  const remaining = Math.max(0, windowSize - used);
+  const percentFull = windowSize > 0 ? Math.round((used / windowSize) * 100) : 0;
+
+  let advice: string;
+  if (percentFull >= 90) {
+    advice = "CRITICAL: Compact immediately. Very little room for complex tasks.";
+  } else if (percentFull >= 75) {
+    advice = "Context is filling up. Compact before starting complex work.";
+  } else if (percentFull >= 50) {
+    advice = "Moderate context usage. Compact before large multi-step tasks.";
+  } else {
+    advice = "Context has plenty of room. No compaction needed yet.";
+  }
+
+  const message = `Context: ~${percentFull}% full (estimated ${remaining.toLocaleString()} tokens remaining)`;
+
+  return { percentFull, remainingTokens: remaining, totalTokens: windowSize, message, advice };
+}
+
+/**
+ * The context_budget tool handler.
+ * Called from the tool registration — receives tokensBefore from Pi context.
+ */
+export function contextBudgetTool(tokensBefore?: number): string {
+  const budget = estimateContextBudget(tokensBefore);
+  if (!budget) return "Context budget: Unknown (no token data available from session).";
+
+  return `${budget.message}\nAdvice: ${budget.advice}`;
+}

package/src/tools/ctx-batch-execute.ts

@@ -3,7 +3,7 @@
  */
 
 import { PolyglotExecutor } from "../executor/executor.js";
-import { ContentStore } from "../store/index.js";
+import type { ContentStore } from "../store/index.js";
 import type { Language, ExecResult, SearchResult } from "../types.js";
 
 export interface BatchCommand {
@@ -28,11 +28,9 @@ export interface BatchResult {
   >;
 }
 
-export async function ctxBatchExecute(items: BatchItem[]): Promise<BatchResult> {
+export async function ctxBatchExecute(store: ContentStore, items: BatchItem[]): Promise<BatchResult> {
   const results: BatchResult["results"] = [];
   const executor = new PolyglotExecutor();
-  const store = new ContentStore();
-  await store.init();
 
   for (const item of items) {
     if (item.type === "execute") {
@@ -48,6 +46,5 @@ export async function ctxBatchExecute(items: BatchItem[]): Promise<BatchResult>
     }
   }
 
-  store.close();
   return { results };
 }

package/src/tools/ctx-fetch-and-index.ts

@@ -2,7 +2,7 @@
  * ctx_fetch_and_index tool — fetch URL → markdown → index
  */
 
-import { ContentStore } from "../store/index.js";
+import type { ContentStore } from "../store/index.js";
 import type { IndexResult } from "../types.js";
 
 export interface CtxFetchAndIndexInput {
@@ -11,7 +11,7 @@ export interface CtxFetchAndIndexInput {
   chunkSize?: number;
 }
 
-export async function ctxFetchAndIndex(input: CtxFetchAndIndexInput): Promise<IndexResult> {
+export async function ctxFetchAndIndex(store: ContentStore, input: CtxFetchAndIndexInput): Promise<IndexResult> {
   const label = input.label ?? input.url;
 
   const response = await fetch(input.url, {
@@ -23,15 +23,10 @@ export async function ctxFetchAndIndex(input: CtxFetchAndIndexInput): Promise<In
   }
 
   const text = await response.text();
-  const store = new ContentStore();
-  await store.init();
 
-  const result = await store.index(label, text, {
+  return store.index(label, text, {
     contentType: "plain",
     source: input.url,
     chunkSize: input.chunkSize,
   });
-
-  store.close();
-  return result;
 }

package/src/tools/ctx-index.ts

@@ -2,7 +2,7 @@
  * ctx_index tool — chunk content → index into FTS5
  */
 
-import { ContentStore } from "../store/index.js";
+import type { ContentStore } from "../store/index.js";
 import type { IndexResult } from "../types.js";
 import { readFileSync } from "node:fs";
 
@@ -14,10 +14,7 @@ export interface CtxIndexInput {
   chunkSize?: number;
 }
 
-export async function ctxIndex(input: CtxIndexInput): Promise<IndexResult> {
-  const store = new ContentStore();
-  await store.init();
-
+export async function ctxIndex(store: ContentStore, input: CtxIndexInput): Promise<IndexResult> {
   let text: string;
   let source: string;
 
@@ -31,12 +28,9 @@ export async function ctxIndex(input: CtxIndexInput): Promise<IndexResult> {
     throw new Error("Either content or filePath must be provided");
   }
 
-  const result = await store.index(input.label, text, {
+  return store.index(input.label, text, {
     contentType: input.contentType ?? "plain",
     source,
     chunkSize: input.chunkSize,
   });
-
-  store.close();
-  return result;
 }