@build-astron-co/nimbus 0.2.0 → 0.4.0

package/src/agent/context-manager.ts

@@ -190,6 +190,34 @@ export function estimateMessageTokens(message: LLMMessage): number {
   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',
+] as const;
+
+/**
+ * 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: string): boolean {
+  return TERRAFORM_PLAN_INDICATORS.some(indicator => text.includes(indicator));
+}
+
 // ---------------------------------------------------------------------------
 // ContextManager Class
 // ---------------------------------------------------------------------------
@@ -206,6 +234,8 @@ export class ContextManager {
   private autoCompactThreshold: number;
   private preserveRecentMessages: number;
   private alwaysInContext: string[];
+  /** Per-message token count cache: key = `role:contentLen:toolCallCount:toolCallId`. */
+  private _tokenCache: Map<string, number> = new Map();
 
   constructor(options?: ContextManagerOptions) {
     // Try loading from config DB, fall back to options/defaults
@@ -216,7 +246,31 @@ export class ContextManager {
     this.maxContextTokens = options?.maxContextTokens ?? modelWindow ?? 200_000;
     this.autoCompactThreshold = configThreshold ?? options?.autoCompactThreshold ?? 0.85;
     this.preserveRecentMessages = options?.preserveRecentMessages ?? 5;
-    this.alwaysInContext = options?.alwaysInContext ?? [];
+    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.
+   */
+  private _getCachedTokens(msg: LLMMessage): number {
+    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(): void {
+    this._tokenCache.clear();
   }
 
   /**
@@ -270,7 +324,7 @@ export class ContextManager {
       baseSystemTokens = systemPromptTokens - nimbusInstructionsTokens;
     }
 
-    const messagesTokens = messages.reduce((sum, msg) => sum + estimateMessageTokens(msg), 0);
+    const messagesTokens = messages.reduce((sum, msg) => sum + this._getCachedTokens(msg), 0);
 
     const total = systemPromptTokens + messagesTokens + toolDefinitionsTokens;
     const usagePercent =
@@ -295,6 +349,8 @@ export class ContextManager {
    * - 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.
@@ -321,7 +377,15 @@ export class ContextManager {
       // Always preserve summary blocks (from previous compactions)
       const isSummary = getTextContent(msg.content).startsWith('[Context Summary]');
 
-      if (isFirst || isRecent || hasActiveTools || isSummary) {
+      // 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);

package/src/agent/deploy-preview.ts

@@ -180,11 +180,28 @@ async function generateTerraformPreview(workdir: string): Promise<DeployPreview>
     const changes = parseTerraformPlanOutput(output);
     const summary = summarizeChanges(changes);
 
+    // Try infracost for real cost delta (best-effort, non-blocking)
+    let costImpact: string | undefined;
+    try {
+      const { stdout: icOut } = await execAsync(
+        `infracost breakdown --path ${workdir} --format json`,
+        { timeout: 60_000, maxBuffer: 5 * 1024 * 1024 }
+      );
+      const ic = JSON.parse(icOut);
+      const totalMonthly = parseFloat(ic.totalMonthlyCost ?? '0').toFixed(2);
+      const diff = parseFloat(ic.diffTotalMonthlyCost ?? '0');
+      const sign = diff > 0 ? '+' : '';
+      costImpact = `$${totalMonthly}/month${diff !== 0 ? ` (${sign}$${diff.toFixed(2)} change)` : ' (no change)'}`;
+    } catch {
+      // infracost not installed or failed — omit cost impact
+    }
+
     return {
       ...base,
       changes,
       summary,
       rawOutput: output,
+      costImpact,
       success: true,
     };
   } catch (error: unknown) {
@@ -280,6 +297,50 @@ async function generateHelmPreview(action: string, workdir: string): Promise<Dep
   }
 }
 
+// ---------------------------------------------------------------
+// G9: Multi-file diff batch builder for Terraform plan output
+// ---------------------------------------------------------------
+
+/**
+ * G9: Build a FileDiffBatch from a Terraform plan preview.
+ *
+ * Parses the resource changes in the plan and represents each modified
+ * resource as a synthetic diff entry so they can be reviewed via the
+ * existing FileDiffModal batch flow.
+ *
+ * @param preview - DeployPreview from generateTerraformPreview
+ * @returns Array of file-diff entries suitable for FileDiffBatch.files
+ */
+export function buildFileDiffBatchFromPlan(
+  preview: DeployPreview
+): Array<{ filePath: string; diff: string; toolName: string }> {
+  return preview.changes.map(change => {
+    const actionLine = getChangeSymbolStr(change.action);
+    const diff = [
+      `--- a/${change.resource}`,
+      `+++ b/${change.resource}`,
+      `@@ -0,0 +1,1 @@`,
+      `${actionLine} ${change.resource}${change.details ? `  # ${change.details}` : ''}`,
+    ].join('\n');
+
+    return {
+      filePath: change.resource,
+      diff,
+      toolName: 'terraform',
+    };
+  });
+}
+
+function getChangeSymbolStr(action: ResourceChange['action']): string {
+  switch (action) {
+    case 'create':  return '+';
+    case 'update':  return '~';
+    case 'destroy': return '-';
+    case 'replace': return '+/-';
+    default:        return ' ';
+  }
+}
+
 // ---------------------------------------------------------------
 // Output parsers
 // ---------------------------------------------------------------
@@ -297,7 +358,7 @@ async function generateHelmPreview(action: string, workdir: string): Promise<Dep
  * Falls back to the summary line (`Plan: X to add, Y to change, Z to destroy.`)
  * when no individual resource lines are found.
  */
-function parseTerraformPlanOutput(output: string): ResourceChange[] {
+export function parseTerraformPlanOutput(output: string): ResourceChange[] {
   const changes: ResourceChange[] = [];
   const lines = output.split('\n');
 

package/src/agent/expand-files.ts

@@ -0,0 +1,108 @@
+/**
+ * Shared @file reference expansion utility.
+ *
+ * Replaces @path/to/file references in a prompt with the file contents
+ * wrapped in <file> tags. Files larger than 10KB are truncated.
+ *
+ * Security: blocks path traversal (../../etc/passwd) and sensitive file patterns.
+ *
+ * Special tokens (G22):
+ *   @workspace — concatenates all .tf files in cwd (up to 50KB)
+ *   @cluster   — instructs the agent to use the kubectl tool (async call not possible here)
+ */
+
+import * as fs from 'node:fs';
+import * as path from 'node:path';
+
+/** Sensitive file patterns that should never be expanded. */
+const SENSITIVE_FILE_RE = /(\.env|\.pem|\.key|\.p12|\.pfx|id_rsa|id_ed25519|credentials|\.netrc)$/i;
+
+/**
+ * Expand @file references in a prompt string.
+ * Replaces @path/to/file with the file contents wrapped in <file> tags.
+ * Files larger than 10KB are truncated.
+ *
+ * Security: blocks path traversal outside the project root and sensitive
+ * file patterns (.env, .pem, .key, credentials, etc.).
+ */
+export function expandFileReferences(text: string, cwd: string = process.cwd()): string {
+  const fileRefs = text.match(/@([\w./_-]+)/g);
+  if (!fileRefs) return text;
+
+  let expanded = text;
+
+  // G22: Handle @workspace special token — concatenate all .tf files
+  if (expanded.includes('@workspace')) {
+    try {
+      const tfFiles = fs.readdirSync(cwd).filter(f => f.endsWith('.tf'));
+      if (tfFiles.length > 0) {
+        const parts: string[] = [];
+        let totalSize = 0;
+        const MAX_SIZE = 50_000;
+        for (const file of tfFiles) {
+          if (totalSize >= MAX_SIZE) { parts.push('... [truncated — 50KB limit reached]'); break; }
+          try {
+            const content = fs.readFileSync(path.join(cwd, file), 'utf-8');
+            parts.push(`\n<file path="${file}">\n${content.slice(0, Math.min(content.length, MAX_SIZE - totalSize))}\n</file>`);
+            totalSize += content.length;
+          } catch { /* skip unreadable files */ }
+        }
+        expanded = expanded.replace('@workspace', parts.join('\n'));
+      } else {
+        expanded = expanded.replace('@workspace', '[No .tf files found in current directory]');
+      }
+    } catch {
+      expanded = expanded.replace('@workspace', '[Error reading workspace files]');
+    }
+  }
+
+  // G22: Handle @cluster special token — kubectl get all
+  // @cluster requires an async kubectl call — guide agent to use the kubectl tool instead
+  if (expanded.includes('@cluster')) {
+    expanded = expanded.replace(
+      '@cluster',
+      '[Use the kubectl tool with action=get and args="all -A" to retrieve cluster resources]'
+    );
+  }
+
+  for (const ref of fileRefs) {
+    // Skip special tokens already handled above
+    if (ref === '@workspace' || ref === '@cluster') {
+      continue;
+    }
+
+    const filePath = ref.slice(1);
+    try {
+      const resolved = path.resolve(cwd, filePath);
+
+      // Security: block path traversal outside project root
+      const rel = path.relative(cwd, resolved);
+      if (rel.startsWith('..') || path.isAbsolute(rel)) {
+        expanded = expanded.replace(
+          ref,
+          `[File blocked: path "${filePath}" is outside the project directory]`
+        );
+        continue;
+      }
+
+      // Security: block sensitive file patterns
+      if (SENSITIVE_FILE_RE.test(resolved)) {
+        expanded = expanded.replace(
+          ref,
+          `[File blocked: "${filePath}" matches sensitive file pattern]`
+        );
+        continue;
+      }
+
+      const content = fs.readFileSync(resolved, 'utf-8');
+      const truncated =
+        content.length > 10_000
+          ? `${content.slice(0, 10_000)}\n... (truncated — showing 10,000 of ${content.length.toLocaleString()} chars)`
+          : content;
+      expanded = expanded.replace(ref, `\n<file path="${filePath}">\n${truncated}\n</file>`);
+    } catch {
+      // File not found — leave the @reference as-is
+    }
+  }
+  return expanded;
+}