feat(tui): maintain session family index in data context

packages/tui/src/component/dialog-session-list.tsx

@@ -83,7 +83,7 @@ export function DialogSessionList() {
         category,
         footer,
         gutter:
-          data.session.status(session.id) === "running"
+          data.session.family(session.id).some((id) => data.session.status(id) === "running")
             ? () => <Spinner />
             : slot === undefined
               ? undefined

packages/tui/src/component/prompt/index.tsx

@@ -160,13 +160,12 @@ export function Prompt(props: PromptProps) {
   const dialog = useDialog()
   const toast = useToast()
   const status = createMemo(() => data.session.status(props.sessionID ?? ""))
-  const activeSubagents = createMemo(
-    () =>
-      data.session
-        .list()
-        .filter((session) => session.parentID === props.sessionID && data.session.status(session.id) === "running")
-        .length,
-  )
+  const activeSubagents = createMemo(() => {
+    if (!props.sessionID) return 0
+    return data.session.family(props.sessionID).filter(
+      (id) => id !== props.sessionID && data.session.status(id) === "running",
+    ).length
+  })
   const runningShells = createMemo(
     () => data.shell.list(currentLocation()).filter((shell) => shell.metadata.sessionID === props.sessionID).length,
   )

packages/tui/src/context/data.tsx

@@ -44,6 +44,10 @@ type LocationData = {
 type Data = {
   session: {
     info: Record<string, SessionV2Info>
+    // Family index keyed by a family's root (or furthest-known-ancestor when the
+    // true root is not yet loaded). The value is a flat deduplicated list of every
+    // session ID in that family, including the key itself once its info arrives.
+    family: Record<string, string[]>
     status: Record<string, DataSessionStatus>
     message: Record<string, SessionMessage[]>
     permission: Record<string, PermissionV2Request[]>
@@ -77,6 +81,7 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
     const [store, setStore] = createStore<Data>({
       session: {
         info: {},
+        family: {},
         status: {},
         message: {},
         permission: {},
@@ -149,6 +154,46 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
       return created
     }
 
+    // Walk parentID upward through loaded session info to the family root. When a
+    // parent's info is missing, that missing ID is the furthest-known ancestor and
+    // is returned so orphan subtrees group under it until the parent arrives. A
+    // seen set guards against parent cycles, stopping at the last non-repeating
+    // ancestor.
+    function resolveRoot(sessionID: string) {
+      let current = sessionID
+      let parentID = store.session.info[sessionID]?.parentID
+      const seen = new Set([sessionID])
+      while (parentID) {
+        if (seen.has(parentID)) break
+        seen.add(parentID)
+        current = parentID
+        parentID = store.session.info[parentID]?.parentID
+      }
+      return current
+    }
+
+    // Register one session into the family index. Idempotent: refreshing an
+    // existing session never duplicates its ID. When a tentative family keyed by
+    // sessionID exists (descendants arrived while sessionID's own info was
+    // absent) but sessionID turns out to have a parent, fold the orphan subtree
+    // into the resolved root's family and drop the tentative entry.
+    function registerSession(sessionID: string) {
+      const info = store.session.info[sessionID]
+      if (!info) return
+      const rootID = resolveRoot(sessionID)
+      setStore("session", "family", produce((draft) => {
+        if (sessionID !== rootID && draft[sessionID]) {
+          const members = draft[rootID] ??= []
+          for (const id of draft[sessionID]) {
+            if (!members.includes(id)) members.push(id)
+          }
+          delete draft[sessionID]
+        }
+        const family = draft[rootID] ??= []
+        if (!family.includes(sessionID)) family.push(sessionID)
+      }))
+    }
+
     function handleEvent(event: V2Event) {
       switch (event.type) {
         case "session.created":
@@ -599,11 +644,18 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
         get(sessionID: string) {
           return store.session.info[sessionID]
         },
+        root(sessionID: string) {
+          return resolveRoot(sessionID)
+        },
+        family(sessionID: string) {
+          return store.session.family[resolveRoot(sessionID)] ?? []
+        },
         status(sessionID: string) {
           return store.session.status[sessionID] ?? "idle"
         },
         async refresh(sessionID: string) {
           setStore("session", "info", sessionID, mutable(await sdk.api.session.get({ sessionID })))
+          registerSession(sessionID)
         },
         message: {
           ids(sessionID: string) {
@@ -795,15 +847,16 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
             directory: defaultLocation().directory,
             workspace: defaultLocation().workspaceID,
           })
-          .then((response) =>
+          .then((response) => {
             setStore(
               "session",
               "info",
               produce((draft) => {
                 for (const session of response.data) draft[session.id] = mutable(session)
               }),
-            ),
-          ),
+            )
+            for (const session of response.data) registerSession(session.id)
+          }),
         sdk.api.session
           .active()
           .then((active) =>

packages/tui/src/routes/session/index.tsx

@@ -172,16 +172,7 @@ export function Session() {
   const messages = sessionMessages
   const descendantSessionIDs = createMemo(() => {
     if (session()?.parentID) return []
-    const sessions = data.session.list()
-    const childrenByParent = sessions.reduce((acc, item) => {
-      if (!item.parentID) return acc
-      acc.set(item.parentID, [...(acc.get(item.parentID) ?? []), item.id])
-      return acc
-    }, new Map<string, string[]>())
-    function collect(sessionID: string): string[] {
-      return (childrenByParent.get(sessionID) ?? []).flatMap((id) => [id, ...collect(id)])
-    }
-    return collect(route.sessionID)
+    return data.session.family(route.sessionID).filter((id) => id !== route.sessionID)
   })
   const permissions = createMemo(() => {
     if (session()?.parentID) return []