feat: add useSubagentCostRollup hook

2026-08-24 14:50:34 +01:00
parent 9ae30bd412
commit 7fde0a01b5
2 changed files with 100 additions and 0 deletions
@@ -0,0 +1,54 @@
+import React from 'react';
+import type { Session } from '@opencode-ai/sdk/v2';
+import { useAllLiveSessions } from '@/sync/sync-context';
+import { buildChildrenIndex, computeSubtreeCost } from './subagentCost';
+
+export type SubagentCostRollup = {
+  totalCost: number | null;
+  subagentCount: number;
+  perChildCost: Map<string, number>;
+};
+
+const EMPTY_ROLLUP: SubagentCostRollup = { totalCost: null, subagentCount: 0, perChildCost: new Map() };
+
+function countDescendants(id: string, childrenByParent: Map<string, Session[]>, visited: Set<string>): number {
+  if (visited.has(id)) return 0;
+  visited.add(id);
+  const kids = childrenByParent.get(id) ?? [];
+  let count = kids.length;
+  for (const kid of kids) count += countDescendants(kid.id, childrenByParent, visited);
+  return count;
+}
+
+/**
+ * Pure core of useSubagentCostRollup, kept separate so it can be unit-tested
+ * directly against a plain session array instead of rendering the hook.
+ */
+export function computeRollup(liveSessions: Session[], sessionId: string | null): SubagentCostRollup {
+  if (!sessionId) return EMPTY_ROLLUP;
+
+  const sessionsById = new Map(liveSessions.map((session) => [session.id, session]));
+  if (!sessionsById.has(sessionId)) return EMPTY_ROLLUP;
+
+  const childrenByParent = buildChildrenIndex(liveSessions);
+  const totalCost = computeSubtreeCost(sessionId, sessionsById, childrenByParent);
+
+  const perChildCost = new Map<string, number>();
+  for (const child of childrenByParent.get(sessionId) ?? []) {
+    perChildCost.set(child.id, computeSubtreeCost(child.id, sessionsById, childrenByParent));
+  }
+
+  const subagentCount = countDescendants(sessionId, childrenByParent, new Set());
+
+  return { totalCost, subagentCount, perChildCost };
+}
+
+/**
+ * Own cost plus every descendant subagent's cost, recursively summed, for a
+ * given root session. Reads the same `useAllLiveSessions()` subscription
+ * WorkStatusSubagentsSection already holds — no new store subscription.
+ */
+export function useSubagentCostRollup(sessionId: string | null): SubagentCostRollup {
+  const liveSessions = useAllLiveSessions();
+  return React.useMemo(() => computeRollup(liveSessions, sessionId), [liveSessions, sessionId]);
+}