@dreb/coding-agent 2.67.0 → 2.68.1

package/dist/modes/rpc/rpc-mode.js

@@ -15,13 +15,14 @@ import { basename } from "node:path";
 import { isValidThinkingLevel, VALID_THINKING_LEVELS } from "../../cli/args.js";
 import { VERSION } from "../../config.js";
 import { canonicalizeTrustedRoots, matchContextTrust, validateTrustedContextFolder, validateTrustedContextFolders, } from "../../core/context-trust.js";
-import { DailyCostTracker } from "../../core/daily-cost-tracker.js";
+import { DailyCostTracker, dailyCostSourcesForSession, } from "../../core/daily-cost-tracker.js";
 import { acquireDreamLock, buildDreamPrompt, cleanupDreamTmpDirs, parseDreamCommand, performDreamBackup, pruneOldBackups, resolveDreamContext, validateArchivePath, validateMemoryLinks, } from "../../core/dream.js";
 import { getGitBranch } from "../../core/git-branch.js";
 import { parseModelPattern, resolveModelScopePatterns } from "../../core/model-resolver.js";
 import { takeOverStdout, writeRawStdout } from "../../core/output-guard.js";
 import { SessionManager } from "../../core/session-manager.js";
 import { BUILTIN_SLASH_COMMANDS, parseBuiltinSlashCommand } from "../../core/slash-commands.js";
+import { SubagentSessionCostTracker } from "../../core/subagent-session-cost-tracker.js";
 import { TabTitleGenerator } from "../../core/tab-title.js";
 import { validateThinkingLevelForModel } from "../../core/thinking.js";
 import { resolveToCwd } from "../../core/tools/path-utils.js";
@@ -1446,6 +1447,9 @@ export async function runRpcMode(session, modelFallbackMessage) {
     let shutdownRequested = false;
     let dailyCostTracker;
     let dailyCostTrackerPrimed = false;
+    let subagentSessionCostTracker;
+    let subagentCostGeneration = 0;
+    let refreshedSubagentCostGeneration = -1;
     // Extension UI context uses the RPC protocol; built by a module-scope
     // factory so the dialog round trip is unit-testable (see createRpcExtensionUIContext).
     const createExtensionUIContext = () => createRpcExtensionUIContext(output, pendingExtensionRequests, () => void session.abort());
@@ -1487,13 +1491,13 @@ export async function runRpcMode(session, modelFallbackMessage) {
             output({ type: "extension_error", extensionPath: err.extensionPath, event: err.event, error: err.error });
         },
     });
-    const getDailyCost = async () => {
-        dailyCostTracker ??= new DailyCostTracker();
+    const getDailyCostBreakdown = async () => {
+        dailyCostTracker ??= new DailyCostTracker(dailyCostSourcesForSession(session));
         if (!dailyCostTrackerPrimed) {
             await dailyCostTracker.refresh();
             dailyCostTrackerPrimed = true;
         }
-        return dailyCostTracker.getDailyCost();
+        return dailyCostTracker.getDailyCostBreakdown();
     };
     const cwd = session.cwd;
     const tabTitleSettings = session.settingsManager.getTabTitleSettings();
@@ -1533,6 +1537,10 @@ export async function runRpcMode(session, modelFallbackMessage) {
     const dashboardEventProjector = session.uiType === "dashboard" ? createDashboardRpcEventProjector() : undefined;
     // Output all agent events as JSON
     session.subscribe((event) => {
+        if (event.type === "background_agent_end") {
+            // Force the next stats request to await a fresh per-session cost scan.
+            subagentCostGeneration++;
+        }
         if (tabTitleGenerator && !session.sessionName) {
             if (event.type === "tool_execution_end") {
                 tabTitleGenerator.onToolEnd({
@@ -1650,7 +1658,12 @@ export async function runRpcMode(session, modelFallbackMessage) {
                 return success(id, "get_git_branch", { branch: getGitBranch(session.cwd) });
             }
             case "get_daily_cost": {
-                return success(id, "get_daily_cost", { cost: await getDailyCost() });
+                const breakdown = await getDailyCostBreakdown();
+                return success(id, "get_daily_cost", {
+                    cost: breakdown.total,
+                    main: breakdown.main,
+                    subagent: breakdown.subagent,
+                });
             }
             // =================================================================
             // Model
@@ -1791,8 +1804,17 @@ export async function runRpcMode(session, modelFallbackMessage) {
             // Session
             // =================================================================
             case "get_session_stats": {
+                subagentSessionCostTracker ??= new SubagentSessionCostTracker();
+                if (refreshedSubagentCostGeneration !== subagentCostGeneration) {
+                    const refreshGeneration = subagentCostGeneration;
+                    await subagentSessionCostTracker.refresh();
+                    refreshedSubagentCostGeneration = refreshGeneration;
+                }
                 const stats = session.getSessionStats();
-                return success(id, "get_session_stats", stats);
+                return success(id, "get_session_stats", {
+                    ...stats,
+                    subagentCost: subagentSessionCostTracker.getCost(),
+                });
             }
             case "get_performance_stats": {
                 return success(id, "get_performance_stats", getPerformanceStatsData(session));
@@ -1968,6 +1990,7 @@ export async function runRpcMode(session, modelFallbackMessage) {
             await currentRunner.emit({ type: "session_shutdown" });
         }
         dailyCostTracker?.dispose();
+        subagentSessionCostTracker?.dispose();
         detachInput();
         process.stdin.pause();
         process.exit(0);