@build-astron-co/nimbus 0.2.0 → 0.4.0

package/dist/src/agent/context-manager.js

@@ -0,0 +1,385 @@
+/**
+ * Context Manager — Token Tracking & Auto-Compact
+ *
+ * Tracks cumulative token usage across the agent loop and triggers
+ * automatic context compaction when usage exceeds a configurable
+ * threshold (default 85% of the model's context window).
+ *
+ * The manager provides:
+ * - Token estimation for messages, system prompts, and tool definitions.
+ * - A breakdown of how the context budget is being consumed.
+ * - Message selection logic for deciding what to preserve vs. summarize.
+ * - A builder for reassembling messages after compaction.
+ *
+ * Configuration can be supplied via constructor options or read from the
+ * Nimbus config database (keys: `context.auto_compact_threshold`,
+ * `context.max_file_injection`).
+ *
+ * @module agent/context-manager
+ */
+import { getTextContent } from '../llm/types';
+import { getConfig } from '../state/config';
+// ---------------------------------------------------------------------------
+// Per-Model Context Window Sizes
+// ---------------------------------------------------------------------------
+/**
+ * Known context window sizes (in tokens) for popular models.
+ *
+ * When a model is not listed here, the manager falls back to the
+ * `maxContextTokens` option (default: 200 000).
+ */
+const MODEL_CONTEXT_WINDOWS = {
+    // Anthropic
+    'claude-opus-4-20250514': 200_000,
+    'claude-sonnet-4-20250514': 200_000,
+    'claude-haiku-4-20250514': 200_000,
+    'claude-3-5-sonnet-20241022': 200_000,
+    'claude-3-5-haiku-20241022': 200_000,
+    'claude-3-opus-20240229': 200_000,
+    'claude-3-sonnet-20240229': 200_000,
+    'claude-3-haiku-20240307': 200_000,
+    // OpenAI
+    'gpt-4o': 128_000,
+    'gpt-4o-mini': 128_000,
+    'gpt-4-turbo': 128_000,
+    'gpt-4': 8_192,
+    'gpt-3.5-turbo': 16_385,
+    o1: 200_000,
+    'o1-mini': 128_000,
+    'o1-preview': 128_000,
+    'o3-mini': 200_000,
+    // Google
+    'gemini-2.0-flash-exp': 1_048_576,
+    'gemini-1.5-pro': 2_097_152,
+    'gemini-1.5-flash': 1_048_576,
+    // Groq (Llama)
+    'llama-3.1-70b-versatile': 131_072,
+    'llama-3.1-8b-instant': 131_072,
+    'llama-3.3-70b-versatile': 131_072,
+    // DeepSeek
+    'deepseek-chat': 64_000,
+    'deepseek-coder': 64_000,
+    'deepseek-reasoner': 64_000,
+    // Local (Ollama defaults — dynamic lookup can override)
+    'llama3.2': 128_000,
+    mistral: 32_768,
+    codellama: 16_384,
+};
+/**
+ * Look up the context window size for a model identifier.
+ *
+ * Tries exact match first, then prefix match (for versioned model IDs
+ * like `claude-sonnet-4-20250514`), then returns `null` if unknown.
+ */
+export function getModelContextWindow(model) {
+    // Exact match
+    if (MODEL_CONTEXT_WINDOWS[model] !== undefined) {
+        return MODEL_CONTEXT_WINDOWS[model];
+    }
+    // Prefix match: e.g., "gpt-4o-2024-08-06" should match "gpt-4o"
+    for (const [key, value] of Object.entries(MODEL_CONTEXT_WINDOWS)) {
+        if (model.startsWith(key)) {
+            return value;
+        }
+    }
+    return null;
+}
+// ---------------------------------------------------------------------------
+// Token Estimation Utilities
+// ---------------------------------------------------------------------------
+/**
+ * Rough token estimate based on character count.
+ *
+ * Uses the common heuristic of ~4 characters per token, which is a
+ * reasonable average across English text and source code.
+ *
+ * @param text - The text to estimate.
+ * @returns Approximate token count (rounded up).
+ */
+export function estimateTokens(text) {
+    return Math.ceil(text.length / 4);
+}
+/**
+ * Estimate token count for a single LLM message.
+ *
+ * Accounts for the message content, structural overhead (role, framing),
+ * and any tool calls embedded in the message.
+ *
+ * @param message - The LLM message to estimate.
+ * @returns Approximate token count.
+ */
+export function estimateMessageTokens(message) {
+    let tokens = 0;
+    tokens += estimateTokens(getTextContent(message.content));
+    // Add overhead for role and message structure
+    tokens += 4;
+    // Tool calls add extra tokens for name, arguments, and JSON structure
+    if (message.toolCalls) {
+        for (const tc of message.toolCalls) {
+            tokens += estimateTokens(tc.function.name);
+            tokens += estimateTokens(tc.function.arguments);
+            tokens += 10; // structural overhead per tool call
+        }
+    }
+    return tokens;
+}
+// ---------------------------------------------------------------------------
+// H2: Terraform plan output detection
+// ---------------------------------------------------------------------------
+/**
+ * Patterns that indicate a message contains a terraform plan output.
+ * These messages are critical context that should survive compaction.
+ */
+const TERRAFORM_PLAN_INDICATORS = [
+    'Plan:',
+    'will be created',
+    'will be destroyed',
+    'to add,',
+    'to change,',
+    'to destroy',
+];
+/**
+ * Return true if the given text content contains terraform plan output.
+ * Used by selectPreservedMessages to protect plan results from compaction.
+ *
+ * @param text - The text content to inspect.
+ * @returns `true` if the text looks like terraform plan output.
+ */
+function containsTerraformPlanOutput(text) {
+    return TERRAFORM_PLAN_INDICATORS.some(indicator => text.includes(indicator));
+}
+// ---------------------------------------------------------------------------
+// ContextManager Class
+// ---------------------------------------------------------------------------
+/**
+ * Manages context window budget and auto-compaction decisions.
+ *
+ * Create one instance per agent session. The manager does not hold
+ * conversation state itself -- it operates on message arrays passed in
+ * by the caller.
+ */
+export class ContextManager {
+    maxContextTokens;
+    autoCompactThreshold;
+    preserveRecentMessages;
+    alwaysInContext;
+    /** Per-message token count cache: key = `role:contentLen:toolCallCount:toolCallId`. */
+    _tokenCache = new Map();
+    constructor(options) {
+        // Try loading from config DB, fall back to options/defaults
+        const configThreshold = getConfigSafe('context.auto_compact_threshold');
+        // Auto-detect context window from model if provided, then options, then default
+        const modelWindow = options?.model ? getModelContextWindow(options.model) : null;
+        this.maxContextTokens = options?.maxContextTokens ?? modelWindow ?? 200_000;
+        this.autoCompactThreshold = configThreshold ?? options?.autoCompactThreshold ?? 0.85;
+        this.preserveRecentMessages = options?.preserveRecentMessages ?? 5;
+        this.alwaysInContext = options?.alwaysInContext ?? ['Infrastructure Context', 'Tool Timeouts'];
+    }
+    /**
+     * Get cached token count for a message, computing + caching on miss.
+     * Key encodes role + content length + tool call count + tool call IDs so
+     * any meaningful change busts the cache entry.
+     */
+    _getCachedTokens(msg) {
+        const contentLen = typeof msg.content === 'string'
+            ? msg.content.length
+            : JSON.stringify(msg.content).length;
+        const tcCount = msg.toolCalls?.length ?? 0;
+        const tcIds = msg.toolCalls?.map(tc => tc.id).join(',') ?? '';
+        const key = `${msg.role}:${contentLen}:${tcCount}:${tcIds}`;
+        const cached = this._tokenCache.get(key);
+        if (cached !== undefined)
+            return cached;
+        const val = estimateMessageTokens(msg);
+        this._tokenCache.set(key, val);
+        return val;
+    }
+    /** Clear the token cache. Call after compaction when old messages are removed. */
+    clearTokenCache() {
+        this._tokenCache.clear();
+    }
+    /**
+     * Check whether auto-compaction should be triggered.
+     *
+     * Returns `true` if the estimated token usage is at or above the
+     * configured threshold percentage of the context window.
+     *
+     * @param systemPrompt - The full system prompt string.
+     * @param messages - Current conversation messages.
+     * @param toolDefinitionsTokens - Pre-computed token count for tool schemas.
+     * @returns `true` if compaction should run.
+     */
+    shouldCompact(systemPrompt, messages, toolDefinitionsTokens) {
+        const usage = this.calculateUsage(systemPrompt, messages, toolDefinitionsTokens);
+        return usage.usagePercent >= this.autoCompactThreshold * 100;
+    }
+    /**
+     * Calculate a detailed context usage breakdown.
+     *
+     * Separates the system prompt into base instructions and NIMBUS.md
+     * content (if present), and sums up messages and tool definitions
+     * to produce a full picture of context window consumption.
+     *
+     * @param systemPrompt - The full system prompt string.
+     * @param messages - Current conversation messages.
+     * @param toolDefinitionsTokens - Pre-computed token count for tool schemas.
+     * @returns A {@link ContextBreakdown} with per-category token counts.
+     */
+    calculateUsage(systemPrompt, messages, toolDefinitionsTokens) {
+        const systemPromptTokens = estimateTokens(systemPrompt);
+        // Separate NIMBUS.md instructions if they appear in system prompt
+        const nimbusMarker = '# NIMBUS.md';
+        const nimbusIdx = systemPrompt.indexOf(nimbusMarker);
+        let nimbusInstructionsTokens = 0;
+        let baseSystemTokens = systemPromptTokens;
+        if (nimbusIdx >= 0) {
+            const nimbusContent = systemPrompt.slice(nimbusIdx);
+            nimbusInstructionsTokens = estimateTokens(nimbusContent);
+            baseSystemTokens = systemPromptTokens - nimbusInstructionsTokens;
+        }
+        const messagesTokens = messages.reduce((sum, msg) => sum + this._getCachedTokens(msg), 0);
+        const total = systemPromptTokens + messagesTokens + toolDefinitionsTokens;
+        const usagePercent = this.maxContextTokens > 0 ? Math.round((total / this.maxContextTokens) * 100) : 0;
+        return {
+            systemPrompt: baseSystemTokens,
+            nimbusInstructions: nimbusInstructionsTokens,
+            messages: messagesTokens,
+            toolDefinitions: toolDefinitionsTokens,
+            total,
+            budget: this.maxContextTokens,
+            usagePercent,
+        };
+    }
+    /**
+     * Select which messages to preserve during compaction.
+     *
+     * Preservation rules:
+     * - The first message is always kept (initial user context).
+     * - The last N messages are always kept (recent conversation).
+     * - Tool messages near the recent window are kept (active tool state).
+     * - Previous compaction summary blocks are always kept.
+     * - H2: Tool result messages containing terraform plan output are always
+     *   kept so the agent retains critical plan context across compaction.
+     * - Everything else is marked for summarization.
+     *
+     * @param messages - The full conversation message array.
+     * @returns An object with `preserved` and `toSummarize` arrays.
+     */
+    selectPreservedMessages(messages) {
+        if (messages.length <= this.preserveRecentMessages + 1) {
+            return { preserved: [...messages], toSummarize: [] };
+        }
+        const preserved = [];
+        const toSummarize = [];
+        for (let i = 0; i < messages.length; i++) {
+            const msg = messages[i];
+            const isFirst = i === 0;
+            const isRecent = i >= messages.length - this.preserveRecentMessages;
+            const hasActiveTools = msg.role === 'tool' && i >= messages.length - this.preserveRecentMessages - 2;
+            // Always preserve summary blocks (from previous compactions)
+            const isSummary = getTextContent(msg.content).startsWith('[Context Summary]');
+            // H2: Always preserve messages containing terraform plan output.
+            // These are typically `tool` role messages or `assistant` messages with
+            // tool_use content blocks whose text contains plan output markers.
+            const textContent = getTextContent(msg.content);
+            const isTerraformPlan = (msg.role === 'tool' || msg.role === 'assistant') &&
+                containsTerraformPlanOutput(textContent);
+            if (isFirst || isRecent || hasActiveTools || isSummary || isTerraformPlan) {
+                preserved.push(msg);
+            }
+            else {
+                toSummarize.push(msg);
+            }
+        }
+        return { preserved, toSummarize };
+    }
+    /**
+     * Build the compacted message array by inserting a summary.
+     *
+     * Places the summary as a user message immediately after the first
+     * preserved message, then appends all remaining preserved messages.
+     * The summary is wrapped with `[Context Summary]` markers so future
+     * compaction passes can identify and preserve it.
+     *
+     * @param preserved - Messages to keep verbatim.
+     * @param summary - The LLM-generated (or fallback) summary text.
+     * @returns A new message array ready to replace the original.
+     */
+    buildCompactedMessages(preserved, summary) {
+        const result = [];
+        // Keep the first preserved message (typically the first user message)
+        if (preserved.length > 0) {
+            result.push(preserved[0]);
+        }
+        // Insert the summary as a user message with a clear marker
+        result.push({
+            role: 'user',
+            content: `[Context Summary] The following is a summary of the earlier conversation:\n\n${summary}\n\n---\nThe conversation continues below.`,
+        });
+        // Append remaining preserved messages
+        for (let i = 1; i < preserved.length; i++) {
+            result.push(preserved[i]);
+        }
+        return result;
+    }
+    /**
+     * Get the current configuration values.
+     *
+     * Useful for displaying context status in the TUI.
+     */
+    getConfig() {
+        return {
+            maxContextTokens: this.maxContextTokens,
+            autoCompactThreshold: this.autoCompactThreshold,
+            preserveRecentMessages: this.preserveRecentMessages,
+        };
+    }
+    /**
+     * Update the max context tokens.
+     *
+     * Call this when the model changes mid-session so the compaction
+     * threshold adjusts to the new model's context window.
+     *
+     * @param tokens - The new maximum context window size.
+     */
+    setMaxContextTokens(tokens) {
+        this.maxContextTokens = tokens;
+    }
+    /**
+     * Update the context window based on a model identifier.
+     *
+     * Looks up the model's known context window size. If the model is
+     * not in the built-in map, the current budget is left unchanged.
+     *
+     * @param model - The model identifier (e.g., "gpt-4o", "claude-sonnet-4-20250514").
+     * @returns `true` if the budget was updated, `false` if model is unknown.
+     */
+    setModel(model) {
+        // Strip provider prefix (e.g., "openai/gpt-4o" → "gpt-4o")
+        const stripped = model.includes('/') ? model.split('/').slice(1).join('/') : model;
+        const window = getModelContextWindow(stripped);
+        if (window !== null) {
+            this.maxContextTokens = window;
+            return true;
+        }
+        return false;
+    }
+}
+// ---------------------------------------------------------------------------
+// Internal Helpers
+// ---------------------------------------------------------------------------
+/**
+ * Safely read a config value without crashing if the DB is not ready.
+ *
+ * During early initialization the SQLite database may not yet be open.
+ * This wrapper catches any error and returns `null` so the constructor
+ * can fall back to provided options or built-in defaults.
+ */
+function getConfigSafe(key) {
+    try {
+        return getConfig(key);
+    }
+    catch {
+        return null;
+    }
+}