diff --git a/packages/tui/src/component/dialog-session-list.tsx b/packages/tui/src/component/dialog-session-list.tsx index 3713e9fc47..b6117c5ca6 100644 --- a/packages/tui/src/component/dialog-session-list.tsx +++ b/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") ? () => : slot === undefined ? undefined diff --git a/packages/tui/src/component/prompt/index.tsx b/packages/tui/src/component/prompt/index.tsx index 4d810309f8..e182d537ff 100644 --- a/packages/tui/src/component/prompt/index.tsx +++ b/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, ) diff --git a/packages/tui/src/context/data.tsx b/packages/tui/src/context/data.tsx index f8e8993142..88e3649e63 100644 --- a/packages/tui/src/context/data.tsx +++ b/packages/tui/src/context/data.tsx @@ -44,6 +44,10 @@ type LocationData = { type Data = { session: { info: Record + // 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 status: Record message: Record permission: Record @@ -77,6 +81,7 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({ const [store, setStore] = createStore({ 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) => diff --git a/packages/tui/src/routes/session/index.tsx b/packages/tui/src/routes/session/index.tsx index 17c4d76cb4..d9db6d26a4 100644 --- a/packages/tui/src/routes/session/index.tsx +++ b/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()) - 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 [] diff --git a/packages/tui/test/cli/tui/data.test.tsx b/packages/tui/test/cli/tui/data.test.tsx index efeab233f3..ed4f407993 100644 --- a/packages/tui/test/cli/tui/data.test.tsx +++ b/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) { + 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 + let ready!: () => void + const mounted = new Promise((resolve) => { + ready = resolve + }) + function Probe() { + data = useData() + onMount(ready) + return + } + const app = await testRender(() => ( + + + + + + + + + + )) + 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() + } +})