@dreb/coding-agent 2.66.3 → 2.68.0

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

@@ -22,6 +22,7 @@ import { parseModelPattern, resolveModelScopePatterns } from "../../core/model-r
 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";
@@ -247,9 +248,9 @@ export async function deleteSessionForRpc(sessionManager, sessionPath) {
     }
     return { ok: true, method: result.method };
 }
-/** Handle the `list_all_sessions` RPC command: list every project's sessions as RPC DTOs. */
-export async function listAllSessionsForRpc() {
-    const sessions = await SessionManager.listAll();
+/** Handle `list_all_sessions` using a custom flat store when one is active. */
+export async function listAllSessionsForRpc(customDir) {
+    const sessions = customDir !== undefined ? await SessionManager.listAllFromDir(customDir) : await SessionManager.listAll();
     return sessions.map(toRpcSessionInfo);
 }
 /**
@@ -1433,7 +1434,7 @@ export async function runRpcMode(session, modelFallbackMessage) {
         return { id, type: "response", command, success: false, error: message };
     };
     if (session.sessionFile && session.messages.length > 0) {
-        const rehydratedCount = rehydrateBackgroundAgentsFromDisk(session.sessionFile);
+        const rehydratedCount = rehydrateBackgroundAgentsFromDisk(session.sessionFile, session.subagentSessionDiscoveryRoots);
         if (rehydratedCount > 0) {
             console.error(`[rpc] Rehydrated ${rehydratedCount} background subagent${rehydratedCount === 1 ? "" : "s"} from disk`);
         }
@@ -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 () => {
+    const getDailyCostBreakdown = async () => {
         dailyCostTracker ??= new DailyCostTracker();
         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));
@@ -1874,7 +1896,9 @@ export async function runRpcMode(session, modelFallbackMessage) {
                 return success(id, "list_sessions", { sessions: sessions.map(toRpcSessionInfo) });
             }
             case "list_all_sessions": {
-                return success(id, "list_all_sessions", { sessions: await listAllSessionsForRpc() });
+                return success(id, "list_all_sessions", {
+                    sessions: await listAllSessionsForRpc(session.sessionManager.getCustomSessionInventoryRoot()),
+                });
             }
             // =================================================================
             // Background agents
@@ -1966,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);