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 []

packages/tui/test/cli/tui/data.test.tsx

@@ -915,3 +915,122 @@ test("projects live context updates with their message ID", async () => {
     app.renderer.destroy()
   }
 })
+
+function sessionInfo(id: string, parentID: string | undefined) {
+  return {
+    id,
+    parentID,
+    projectID: "proj_test",
+    cost: 0,
+    tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } },
+    time: { created: 0, updated: 0 },
+    title: id,
+    location: { directory },
+  }
+}
+
+// Mounts a DataProvider whose `/api/session/:id` responses are driven by the
+// given parent map (sessionID -> parentID). Roots omit the entry. Reused across
+// the family-index tests below.
+async function mountData(parents: Record<string, string>) {
+  const calls = createFetch((url) => {
+    const match = url.pathname.match(/^\/api\/session\/([^/]+)$/)
+    if (match && match[1] !== "active") return json({ data: sessionInfo(match[1], parents[match[1]]) })
+  })
+  let data!: ReturnType<typeof useData>
+  let ready!: () => void
+  const mounted = new Promise<void>((resolve) => {
+    ready = resolve
+  })
+  function Probe() {
+    data = useData()
+    onMount(ready)
+    return <box />
+  }
+  const app = await testRender(() => (
+    <TestTuiContexts>
+      <SDKProvider client={createClient(calls.fetch)} api={createApi(calls.fetch)}>
+        <ProjectProvider>
+          <DataProvider>
+            <Probe />
+          </DataProvider>
+        </ProjectProvider>
+      </SDKProvider>
+    </TestTuiContexts>
+  ))
+  await mounted
+  return { data, app }
+}
+
+test("groups an orphan child under its missing parent until the root arrives", async () => {
+  const { data, app } = await mountData({ child: "root" })
+  try {
+    await data.session.refresh("child")
+    // Parent info is absent, so the missing parent is the furthest-known ancestor.
+    expect(data.session.root("child")).toBe("root")
+    expect(data.session.family("child")).toEqual(["child"])
+    expect(data.session.family("root")).toEqual(["child"])
+
+    await data.session.refresh("root")
+    expect(data.session.root("root")).toBe("root")
+    // The tentative root entry folds into the now-known root's family.
+    expect(data.session.family("child")).toEqual(["child", "root"])
+    expect(data.session.family("root")).toEqual(["child", "root"])
+  } finally {
+    app.renderer.destroy()
+  }
+})
+
+test("indexes arbitrarily deep nesting under a single root", async () => {
+  const { data, app } = await mountData({ grandchild: "child", child: "root" })
+  try {
+    await data.session.refresh("grandchild")
+    expect(data.session.root("grandchild")).toBe("child")
+    expect(data.session.family("grandchild")).toEqual(["grandchild"])
+
+    await data.session.refresh("child")
+    // grandchild's tentative family (keyed by the missing "child") merges up
+    // toward the still-missing "root".
+    expect(data.session.root("child")).toBe("root")
+    expect(data.session.family("grandchild")).toEqual(["grandchild", "child"])
+
+    await data.session.refresh("root")
+    expect(data.session.root("grandchild")).toBe("root")
+    expect(data.session.root("child")).toBe("root")
+    expect(data.session.family("root")).toEqual(["grandchild", "child", "root"])
+  } finally {
+    app.renderer.destroy()
+  }
+})
+
+test("re-registering an existing session is idempotent", async () => {
+  const { data, app } = await mountData({ grandchild: "child", child: "root" })
+  try {
+    await data.session.refresh("grandchild")
+    await data.session.refresh("child")
+    await data.session.refresh("root")
+    const before = data.session.family("root")
+    expect(before).toEqual(["grandchild", "child", "root"])
+
+    await data.session.refresh("child")
+    await data.session.refresh("root")
+    await data.session.refresh("grandchild")
+    expect(data.session.family("root")).toEqual(before)
+    expect(data.session.family("root")).toHaveLength(3)
+  } finally {
+    app.renderer.destroy()
+  }
+})
+
+test("stops at the last non-repeating ancestor on a parent cycle", async () => {
+  const { data, app } = await mountData({ x: "y", y: "x" })
+  try {
+    await data.session.refresh("x")
+    await data.session.refresh("y")
+    // Does not hang; walking up from "y" stops before re-entering "x".
+    expect(data.session.root("y")).toBe("x")
+    expect(data.session.family("y")).toEqual(["x", "y"])
+  } finally {
+    app.renderer.destroy()
+  }
+})