@bastani/atomic 0.8.26 → 0.8.27-alpha.1

package/dist/core/compaction/context-compaction.js

@@ -1,7 +1,6 @@
 import { Agent } from "@earendil-works/pi-agent-core";
-import { createAssistantMessageEventStream, getSupportedThinkingLevels, isContextOverflow, streamSimple, StringEnum, } from "@earendil-works/pi-ai";
+import { createAssistantMessageEventStream, isContextOverflow, streamSimple, StringEnum, } from "@earendil-works/pi-ai";
 import { mkdtempSync, rmSync, writeFileSync } from "node:fs";
-import { createRequire } from "node:module";
 import { tmpdir } from "node:os";
 import { join } from "node:path";
 import { Type } from "typebox";
@@ -9,7 +8,6 @@ import { createBranchSummaryMessage, createCompactionSummaryMessage, createCusto
 import { buildContextDeletionFilteredPath, buildContextDeletionFilters, } from "../session-manager.js";
 import { estimateTokens } from "./compaction.js";
 export const CONTEXT_COMPACTION_PROMPT_VERSION = 1;
-const CONTEXT_COMPACTION_THINKING_LEVEL_ORDER = ["off", "minimal", "low", "medium", "high", "xhigh"];
 const CONTEXT_DELETE_TOOL_NAME = "context_delete";
 const CONTEXT_GREP_DELETE_TOOL_NAME = "context_grep_delete";
 const CONTEXT_SEARCH_TRANSCRIPT_TOOL_NAME = "context_search_transcript";
@@ -854,204 +852,156 @@ function addGrepCandidate(candidates, matches, seenTargets, candidate, match) {
     candidates.push(candidate);
     matches.push(match);
 }
-const moduleRequire = createRequire(import.meta.url);
-function createTransientSqliteDatabase() {
-    // better-sqlite3 is the portable/package dependency for Node-based Atomic installs.
-    // Bun cannot dlopen better-sqlite3 yet, so Bun runtime/tests use the API-compatible builtin.
-    if (process.versions.bun) {
-        const sqlite = moduleRequire("bun:sqlite");
-        return new sqlite.Database(":memory:");
-    }
-    const BetterSqliteDatabase = moduleRequire("better-sqlite3");
-    return new BetterSqliteDatabase(":memory:");
-}
-class SqliteAdapter {
-    constructor(db) {
-        this.db = db;
-    }
-    static createTransient() {
-        return new SqliteAdapter(createTransientSqliteDatabase());
-    }
-    exec(sql) {
-        this.db.exec(sql);
-    }
-    run(sql, ...params) {
-        this.db.prepare(sql).run(...params);
-    }
-    all(sql, ...params) {
-        return this.db.prepare(sql).all(...params);
-    }
-    get(sql, ...params) {
-        return this.db.prepare(sql).get(...params);
+function copyDeletionTarget(target) {
+    return target.kind === "entry"
+        ? { kind: "entry", entryId: target.entryId }
+        : { kind: "content_block", entryId: target.entryId, blockIndex: target.blockIndex };
+}
+class ContextDeletionMemoryStore {
+    constructor(transcript) {
+        this.deletionTargets = [];
+        this.callCount = 0;
+        const entryIds = new Set();
+        const blockKeys = new Set();
+        this.entries = transcript.entries.map((entry) => {
+            if (entryIds.has(entry.entryId)) {
+                throw new Error(`Duplicate transcript entry id: ${entry.entryId}`);
+            }
+            entryIds.add(entry.entryId);
+            return {
+                entryId: entry.entryId,
+                role: entry.role,
+                protected: entry.protected,
+                tokenEstimate: entry.tokenEstimate,
+                text: entry.text,
+            };
+        });
+        this.entriesById = new Map(this.entries.map((entry) => [entry.entryId, entry]));
+        this.contentBlocks = transcript.entries.flatMap((entry, entryPosition) => entry.contentBlocks.map((block) => {
+            if (block.entryId !== entry.entryId) {
+                throw new Error(`Transcript content block ${block.entryId}:${block.blockIndex} does not belong to entry ${entry.entryId}`);
+            }
+            const blockKey = `${block.entryId}:${block.blockIndex}`;
+            if (blockKeys.has(blockKey)) {
+                throw new Error(`Duplicate transcript content block: ${blockKey}`);
+            }
+            blockKeys.add(blockKey);
+            return {
+                entryPosition,
+                entryId: block.entryId,
+                blockIndex: block.blockIndex,
+                type: block.type,
+                protected: block.protected,
+                tokenEstimate: block.tokenEstimate,
+                text: block.text,
+            };
+        }));
+        this.contentBlockCountByEntryId = new Map();
+        for (const block of this.contentBlocks) {
+            this.contentBlockCountByEntryId.set(block.entryId, (this.contentBlockCountByEntryId.get(block.entryId) ?? 0) + 1);
+        }
     }
     transaction(operation) {
-        this.exec("BEGIN IMMEDIATE");
+        const snapshot = this.snapshot();
         try {
-            const result = operation();
-            this.exec("COMMIT");
-            return result;
+            return operation();
         }
         catch (error) {
-            try {
-                this.exec("ROLLBACK");
-            }
-            catch { }
+            this.restore(snapshot);
             throw error;
         }
     }
-    close() {
-        this.db.close();
-    }
-}
-class ContextDeletionSqliteStore {
-    constructor(sqlite) {
-        this.sqlite = sqlite;
-    }
-    initialize(transcript) {
-        this.sqlite.transaction(() => {
-            this.sqlite.exec(`
-				PRAGMA foreign_keys = ON;
-				CREATE TABLE transcript_entries (
-					position INTEGER PRIMARY KEY,
-					entry_id TEXT NOT NULL UNIQUE,
-					role TEXT NOT NULL,
-					is_protected INTEGER NOT NULL,
-					token_estimate INTEGER NOT NULL,
-					text TEXT NOT NULL,
-					tool_result_for TEXT
-				);
-				CREATE TABLE transcript_content_blocks (
-					entry_id TEXT NOT NULL,
-					block_index INTEGER NOT NULL,
-					type TEXT NOT NULL,
-					is_protected INTEGER NOT NULL,
-					token_estimate INTEGER NOT NULL,
-					text TEXT NOT NULL,
-					tool_call_id TEXT,
-					PRIMARY KEY (entry_id, block_index),
-					FOREIGN KEY (entry_id) REFERENCES transcript_entries(entry_id) ON DELETE CASCADE
-				);
-				CREATE TABLE deletion_targets (
-					position INTEGER PRIMARY KEY AUTOINCREMENT,
-					target_key TEXT NOT NULL UNIQUE,
-					kind TEXT NOT NULL CHECK (kind IN ('entry', 'content_block')),
-					entry_id TEXT NOT NULL,
-					block_index INTEGER
-				);
-				CREATE TABLE context_compaction_state (
-					key TEXT PRIMARY KEY,
-					value TEXT NOT NULL
-				);
-				INSERT INTO context_compaction_state (key, value) VALUES ('call_count', '0');
-			`);
-            for (const [position, entry] of transcript.entries.entries()) {
-                this.sqlite.run("INSERT INTO transcript_entries (position, entry_id, role, is_protected, token_estimate, text, tool_result_for) VALUES (?, ?, ?, ?, ?, ?, ?)", position, entry.entryId, entry.role, entry.protected ? 1 : 0, entry.tokenEstimate, entry.text, entry.toolResultFor ?? null);
-                for (const block of entry.contentBlocks) {
-                    this.sqlite.run("INSERT INTO transcript_content_blocks (entry_id, block_index, type, is_protected, token_estimate, text, tool_call_id) VALUES (?, ?, ?, ?, ?, ?, ?)", block.entryId, block.blockIndex, block.type, block.protected ? 1 : 0, block.tokenEstimate, block.text, block.toolCallId ?? null);
-                }
-            }
-        });
-    }
-    transaction(operation) {
-        return this.sqlite.transaction(operation);
-    }
     readTargets() {
-        return this.sqlite
-            .all("SELECT kind, entry_id, block_index FROM deletion_targets ORDER BY position")
-            .map((row) => row.kind === "entry"
-            ? { kind: "entry", entryId: row.entry_id }
-            : { kind: "content_block", entryId: row.entry_id, blockIndex: row.block_index });
+        return this.deletionTargets.map(copyDeletionTarget);
     }
     replaceTargets(targets) {
-        this.sqlite.run("DELETE FROM deletion_targets");
-        for (const target of targets) {
-            this.sqlite.run("INSERT INTO deletion_targets (target_key, kind, entry_id, block_index) VALUES (?, ?, ?, ?)", targetKey(target), target.kind, target.entryId, target.kind === "content_block" ? target.blockIndex : null);
-        }
+        this.deletionTargets = targets.map(copyDeletionTarget);
     }
     listEntriesForGrep() {
-        return this.sqlite.all("SELECT entry_id, text, is_protected FROM transcript_entries ORDER BY position");
+        return this.entries.map((entry) => ({
+            entry_id: entry.entryId,
+            text: entry.text,
+            is_protected: entry.protected ? 1 : 0,
+        }));
     }
     listContentBlocksForGrep() {
-        return this.sqlite.all(`
-			SELECT
-				blocks.entry_id,
-				blocks.block_index,
-				blocks.text,
-				entries.is_protected AS entry_protected,
-				blocks.is_protected AS block_protected,
-				(
-					SELECT COUNT(*)
-					FROM transcript_content_blocks sibling
-					WHERE sibling.entry_id = blocks.entry_id
-				) AS block_count
-			FROM transcript_content_blocks blocks
-			JOIN transcript_entries entries ON entries.entry_id = blocks.entry_id
-			ORDER BY entries.position, blocks.block_index
-		`);
+        return [...this.contentBlocks]
+            .sort((a, b) => a.entryPosition - b.entryPosition || a.blockIndex - b.blockIndex)
+            .map((block) => ({
+            entry_id: block.entryId,
+            block_index: block.blockIndex,
+            text: block.text,
+            entry_protected: this.entriesById.get(block.entryId)?.protected ? 1 : 0,
+            block_protected: block.protected ? 1 : 0,
+            block_count: this.contentBlockCountByEntryId.get(block.entryId) ?? 0,
+        }));
     }
     getEntryForRead(entryId) {
-        return this.sqlite.get("SELECT entry_id, role, is_protected, token_estimate, text FROM transcript_entries WHERE entry_id = ?", entryId);
+        const entry = this.entriesById.get(entryId);
+        if (!entry)
+            return undefined;
+        return {
+            entry_id: entry.entryId,
+            role: entry.role,
+            is_protected: entry.protected ? 1 : 0,
+            token_estimate: entry.tokenEstimate,
+            text: entry.text,
+        };
     }
     getContentBlockForRead(entryId, blockIndex) {
-        return this.sqlite.get(`
-				SELECT
-					blocks.entry_id,
-					blocks.block_index,
-					blocks.type,
-					blocks.token_estimate,
-					blocks.text,
-					entries.is_protected AS entry_protected,
-					blocks.is_protected AS block_protected,
-					(
-						SELECT COUNT(*)
-						FROM transcript_content_blocks sibling
-						WHERE sibling.entry_id = blocks.entry_id
-					) AS block_count
-				FROM transcript_content_blocks blocks
-				JOIN transcript_entries entries ON entries.entry_id = blocks.entry_id
-				WHERE blocks.entry_id = ? AND blocks.block_index = ?
-			`, entryId, blockIndex);
+        const block = this.contentBlocks.find((candidate) => candidate.entryId === entryId && candidate.blockIndex === blockIndex);
+        if (!block)
+            return undefined;
+        return {
+            entry_id: block.entryId,
+            block_index: block.blockIndex,
+            type: block.type,
+            token_estimate: block.tokenEstimate,
+            text: block.text,
+            entry_protected: this.entriesById.get(block.entryId)?.protected ? 1 : 0,
+            block_protected: block.protected ? 1 : 0,
+            block_count: this.contentBlockCountByEntryId.get(block.entryId) ?? 0,
+        };
     }
     getGrepScanTextLength(target) {
-        const table = target === "entry" ? "transcript_entries" : "transcript_content_blocks";
-        const row = this.sqlite.get(`SELECT SUM(LENGTH(text)) AS scan_chars FROM ${table}`);
-        return row?.scan_chars ?? 0;
+        const texts = target === "entry" ? this.entries : this.contentBlocks;
+        return texts.reduce((sum, item) => sum + item.text.length, 0);
     }
     incrementCallCount() {
-        const next = this.getCallCount() + 1;
-        this.setState("call_count", String(next));
-        return next;
+        this.callCount += 1;
+        return this.callCount;
     }
     getCallCount() {
-        return Number(this.getState("call_count") ?? "0");
+        return this.callCount;
     }
     setLastError(message) {
-        this.setState("last_error", message);
+        this.lastError = message;
     }
     clearLastError() {
-        this.sqlite.run("DELETE FROM context_compaction_state WHERE key = ?", "last_error");
+        this.lastError = undefined;
     }
     getLastError() {
-        return this.getState("last_error");
+        return this.lastError;
     }
-    getState(key) {
-        return this.sqlite.get("SELECT value FROM context_compaction_state WHERE key = ?", key)?.value;
-    }
-    setState(key, value) {
-        this.sqlite.run("INSERT INTO context_compaction_state (key, value) VALUES (?, ?) ON CONFLICT(key) DO UPDATE SET value = excluded.value", key, value);
+    snapshot() {
+        return {
+            deletionTargets: this.readTargets(),
+            callCount: this.callCount,
+            ...(this.lastError === undefined ? {} : { lastError: this.lastError }),
+        };
     }
-    close() {
-        this.sqlite.close();
+    restore(snapshot) {
+        this.deletionTargets = snapshot.deletionTargets.map(copyDeletionTarget);
+        this.callCount = snapshot.callCount;
+        this.lastError = snapshot.lastError;
     }
 }
-function createContextDeletionSqliteStore(transcript) {
-    const store = new ContextDeletionSqliteStore(SqliteAdapter.createTransient());
-    store.initialize(transcript);
-    return store;
+function createContextDeletionStore(transcript) {
+    return new ContextDeletionMemoryStore(transcript);
 }
 export function createContextDeletionTool(transcript, options = {}) {
     const mode = options.mode ?? "standard";
-    const store = createContextDeletionSqliteStore(transcript);
+    const store = createContextDeletionStore(transcript);
     let validatedResult;
     function readTargets() {
         return store.readTargets();
@@ -1554,16 +1504,8 @@ function createContextCompactionStopStream(model, text) {
 function isContextCompactionOverflowError(model, errorMessage) {
     return isContextOverflow(createContextCompactionAssistantMessage(model, [], "error", errorMessage), model.contextWindow);
 }
-export function getLowestContextCompactionThinkingLevel(model) {
-    const supportedLevels = getSupportedThinkingLevels(model);
-    for (const level of CONTEXT_COMPACTION_THINKING_LEVEL_ORDER) {
-        if (supportedLevels.includes(level))
-            return level;
-    }
-    return "off";
-}
-async function runContextDeletionAssistant(transcript, model, apiKey, headers, signal, mode = "standard") {
-    const maxTokens = Math.min(4096, model.maxTokens > 0 ? model.maxTokens : Number.POSITIVE_INFINITY);
+async function runContextDeletionAssistant(transcript, model, apiKey, headers, signal, thinkingLevel = "off", mode = "standard") {
+    const maxTokens = model.maxTokens > 0 ? model.maxTokens : Number.POSITIVE_INFINITY;
     if (signal?.aborted) {
         throw new Error("Context compaction failed: Request was aborted");
     }
@@ -1574,13 +1516,12 @@ async function runContextDeletionAssistant(transcript, model, apiKey, headers, s
         timestamp: Date.now(),
     };
     const deletionTool = createContextDeletionTool(transcript, { mode });
-    const effectiveThinkingLevel = getLowestContextCompactionThinkingLevel(model);
     let compactionTurnCount = 0;
     const agent = new Agent({
         initialState: {
             systemPrompt: CONTEXT_COMPACTION_SYSTEM_PROMPT,
             model,
-            thinkingLevel: effectiveThinkingLevel,
+            thinkingLevel,
             tools: deletionTool.tools,
         },
         toolExecution: "parallel",
@@ -1626,8 +1567,8 @@ async function runContextDeletionAssistant(transcript, model, apiKey, headers, s
         lastToolError: deletionTool.getLastError(),
     };
 }
-export async function contextCompact(preparation, model, apiKey, headers, signal, _thinkingLevel, mode = preparation.mode ?? "standard") {
-    const { validatedResult, lastToolError } = await runContextDeletionAssistant(preparation.transcript, model, apiKey, headers, signal, mode);
+export async function contextCompact(preparation, model, apiKey, headers, signal, thinkingLevel = "off", mode = preparation.mode ?? "standard") {
+    const { validatedResult, lastToolError } = await runContextDeletionAssistant(preparation.transcript, model, apiKey, headers, signal, thinkingLevel, mode);
     if (!validatedResult || validatedResult.deletedTargets.length === 0) {
         throw new Error(lastToolError ? `No safe context deletions proposed; last deletion tool error: ${lastToolError}` : "No safe context deletions proposed");
     }