From 56241625a49de627174ae7b2b998943805f61f87 Mon Sep 17 00:00:00 2001
From: ezl-keygraph <ezhil@keygraph.io>
Date: Mon, 15 Jun 2026 15:59:55 +0530
Subject: [PATCH] fix(worker): count sub-agent cost and surface compaction
 failures

---
 apps/worker/src/ai/claude-executor.ts | 14 +++++++++++++-
 apps/worker/src/ai/tools.ts           | 15 +++++++++++++++
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/apps/worker/src/ai/claude-executor.ts b/apps/worker/src/ai/claude-executor.ts
index da5e447..0914f18 100644
--- a/apps/worker/src/ai/claude-executor.ts
+++ b/apps/worker/src/ai/claude-executor.ts
@@ -230,12 +230,16 @@ export async function runClaudePrompt(
   // Load the code_path deny extension only when a deny config was written; the same
   // loader is reused by child task sessions so they inherit the policy.
   const resourceLoader = await buildPermissionResourceLoader(sourceDir, logger);
+  // Accumulates cost from in-process `task` child sessions so the parent's reported
+  // cost includes sub-agent spend (their getSessionStats is separate from ours).
+  const childUsage = { cost: 0 };
   const customTools: ToolDefinition[] = [
     createTaskTool({
       model: selection.model,
       thinkingLevel: selection.thinkingLevel,
       authStorage: selection.authStorage,
       cwd: sourceDir,
+      childUsage,
       ...(resourceLoader && { resourceLoader }),
     }),
     createTodoWriteTool(auditLogger),
@@ -292,6 +296,14 @@ export async function runClaudePrompt(
         case 'tool_execution_end':
           void auditLogger.logToolEnd(event.result);
           break;
+        case 'compaction_end':
+          if (!event.aborted && !event.willRetry && event.errorMessage) {
+            pendingError =
+              pendingError ??
+              classifyErrorText(event.errorMessage) ??
+              new PentestError(`Context compaction failed: ${event.errorMessage.slice(0, 200)}`, 'unknown', true);
+          }
+          break;
         default:
           break;
       }
@@ -306,7 +318,7 @@ export async function runClaudePrompt(
 
     // 8. Read usage/cost and final text.
     const stats = session.getSessionStats();
-    const totalCost = stats.cost;
+    const totalCost = stats.cost + childUsage.cost;
     const result = session.getLastAssistantText() ?? null;
 
     // 9. Defense-in-depth: detect a spending cap that produced an empty/cheap run.
diff --git a/apps/worker/src/ai/tools.ts b/apps/worker/src/ai/tools.ts
index 492aab0..cbc2997 100644
--- a/apps/worker/src/ai/tools.ts
+++ b/apps/worker/src/ai/tools.ts
@@ -37,6 +37,12 @@ export interface TaskToolContext {
   cwd: string;
   /** When set, child sessions inherit the code_path deny policy. */
   resourceLoader?: ResourceLoader;
+  /**
+   * Mutable accumulator: each child (sub-agent) session's cost is added here so the
+   * parent executor can include sub-agent spend in its reported cost. Child sessions
+   * keep their own `getSessionStats`, separate from the parent's.
+   */
+  childUsage?: { cost: number };
 }
 
 /**
@@ -80,6 +86,15 @@ export function createTaskTool(ctx: TaskToolContext): ToolDefinition {
         const text = child.getLastAssistantText() ?? '(sub-agent produced no output)';
         return { content: [{ type: 'text' as const, text }], details: {} };
       } finally {
+        // Roll the child's cost up to the parent before disposing (best-effort, and
+        // captured in `finally` so a failed child's partial spend still counts).
+        if (ctx.childUsage) {
+          try {
+            ctx.childUsage.cost += child.getSessionStats().cost;
+          } catch {
+            // ignore — cost capture is best-effort
+          }
+        }
         child.dispose();
       }
     },