From 7fde0a01b5e8975050663d4216215bdf7ea63e89 Mon Sep 17 00:00:00 2001
From: Igor Velho <velhoi@ryanair.com>
Date: Mon, 24 Aug 2026 14:50:34 +0100
Subject: [PATCH] feat: add useSubagentCostRollup hook

---
 .../work-status/useSubagentCostRollup.test.ts | 46 ++++++++++++++++
 .../chat/work-status/useSubagentCostRollup.ts | 54 +++++++++++++++++++
 2 files changed, 100 insertions(+)
 create mode 100644 packages/ui/src/components/chat/work-status/useSubagentCostRollup.test.ts
 create mode 100644 packages/ui/src/components/chat/work-status/useSubagentCostRollup.ts

diff --git a/packages/ui/src/components/chat/work-status/useSubagentCostRollup.test.ts b/packages/ui/src/components/chat/work-status/useSubagentCostRollup.test.ts
new file mode 100644
index 00000000..9368fe4b
--- /dev/null
+++ b/packages/ui/src/components/chat/work-status/useSubagentCostRollup.test.ts
@@ -0,0 +1,46 @@
+import { describe, expect, test } from 'bun:test';
+import type { Session } from '@opencode-ai/sdk/v2';
+import { computeRollup } from './useSubagentCostRollup';
+
+function makeSession(id: string, cost: number, parentID?: string): Session {
+  return { id, cost, parentID } as unknown as Session;
+}
+
+const sessions: Session[] = [
+  makeSession('root', 1),
+  makeSession('a', 2, 'root'),
+  makeSession('b', 3, 'root'),
+  makeSession('a1', 5, 'a'),
+];
+
+describe('computeRollup', () => {
+  test('sums own cost plus every descendant', () => {
+    const result = computeRollup(sessions, 'root');
+    expect(result.totalCost).toBe(11);
+    expect(result.subagentCount).toBe(3);
+  });
+
+  test('maps each direct child to its own subtree cost', () => {
+    const result = computeRollup(sessions, 'root');
+    expect(result.perChildCost.get('a')).toBe(7);
+    expect(result.perChildCost.get('b')).toBe(3);
+  });
+
+  test('returns null total for a null sessionId', () => {
+    const result = computeRollup(sessions, null);
+    expect(result.totalCost).toBeNull();
+    expect(result.subagentCount).toBe(0);
+  });
+
+  test('returns null total for an unknown sessionId', () => {
+    const result = computeRollup(sessions, 'missing');
+    expect(result.totalCost).toBeNull();
+  });
+
+  test('sum of perChildCost plus root cost equals totalCost', () => {
+    const result = computeRollup(sessions, 'root');
+    const childSum = Array.from(result.perChildCost.values()).reduce((sum, v) => sum + v, 0);
+    const rootOwnCost = 1;
+    expect(childSum + rootOwnCost).toBe(result.totalCost);
+  });
+});
diff --git a/packages/ui/src/components/chat/work-status/useSubagentCostRollup.ts b/packages/ui/src/components/chat/work-status/useSubagentCostRollup.ts
new file mode 100644
index 00000000..189da940
--- /dev/null
+++ b/packages/ui/src/components/chat/work-status/useSubagentCostRollup.ts
@@ -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]);
+}