Return to parent tabs on close

packages/app/src/stores/workspace-layout-actions.ts

@@ -26,6 +26,7 @@ export type SplitNode = { kind: "pane"; pane: SplitPane } | { kind: "group"; gro
 export interface WorkspaceLayout {
   root: SplitNode;
   focusedPaneId: string | null;
+  parentTabIdByTabId?: Record<string, string>;
 }
 
 interface SplitPaneInternal extends SplitPane {
@@ -903,10 +904,16 @@ export function normalizeLayout(layout: unknown): WorkspaceLayout {
         collectAllPanes(root)[0]?.id ??
         DEFAULT_PANE_ID);
 
-  return {
+  const normalizedLayout = {
     root,
     focusedPaneId: resolvedFocusedPaneId,
   };
+  const parentTabIdByTabId = normalizeParentTabMap({
+    raw: rawLayout.parentTabIdByTabId,
+    openTabIds: new Set(collectAllTabs(root).map((tab) => tab.tabId)),
+  });
+
+  return parentTabIdByTabId ? { ...normalizedLayout, parentTabIdByTabId } : normalizedLayout;
 }
 
 export function findPaneById(root: SplitNode, paneId: string | null | undefined): SplitPane | null {
@@ -1040,14 +1047,15 @@ function insertNewTabIntoFocusedPane(input: {
 
   return {
     tabId,
-    layout: {
+    layout: withNormalizedParentTabMap({
       root: insertTabIntoPane(layout.root, {
         paneId: focusedPane.id,
         tab: nextTab,
         focusTabId: input.focus ? tabId : preservedFocusTabId,
       }),
       focusedPaneId: input.focus ? focusedPane.id : layout.focusedPaneId,
-    },
+      parentTabIdByTabId: input.layout.parentTabIdByTabId,
+    }),
   };
 }
 
@@ -1061,28 +1069,28 @@ function findExistingTabForTarget(root: SplitNodeInternal, target: WorkspaceTabT
 }
 
 function updateExistingTabTarget(
-  layout: { root: SplitNodeInternal; focusedPaneId: string | null },
+  layout: WorkspaceLayout,
   tab: WorkspaceTab,
   target: WorkspaceTabTarget,
-): { root: SplitNodeInternal; focusedPaneId: string | null } {
+): WorkspaceLayout {
   if (workspaceTabTargetsEqual(tab.target, target)) {
     return layout;
   }
-  return {
+  return withNormalizedParentTabMap({
     ...layout,
-    root: replaceTabInTree(layout.root, {
+    root: replaceTabInTree(asInternalNode(layout.root), {
       tabId: tab.tabId,
       nextTabId: tab.tabId,
       target,
     }),
-  };
+  });
 }
 
 export function openTabInLayoutFocused(input: OpenTabInLayoutInput): OpenTabInLayoutResult {
   const layout = asInternalLayout(input.layout);
   const existingTab = findExistingTabForTarget(layout.root, input.target);
   if (existingTab) {
-    const nextLayout = updateExistingTabTarget(layout, existingTab, input.target);
+    const nextLayout = updateExistingTabTarget(input.layout, existingTab, input.target);
     return {
       tabId: existingTab.tabId,
       layout:
@@ -1102,7 +1110,7 @@ export function openTabInLayoutBackground(input: OpenTabInLayoutInput): OpenTabI
   if (existingTab) {
     return {
       tabId: existingTab.tabId,
-      layout: updateExistingTabTarget(layout, existingTab, input.target),
+      layout: updateExistingTabTarget(input.layout, existingTab, input.target),
     };
   }
 
@@ -1116,20 +1124,42 @@ export function closeTabInLayout(input: CloseTabInLayoutInput): WorkspaceLayout
     return null;
   }
 
+  const closeSuccessorTabId = getCloseSuccessorTabId({
+    pane,
+    tabId: input.tabId,
+    openTabIds: new Set(collectAllTabs(internalLayout.root).map((tab) => tab.tabId)),
+    parentTabIdByTabId: input.layout.parentTabIdByTabId,
+  });
   const fallbackPaneId = findNearestSiblingPaneId(internalLayout.root, pane.id);
   const nextRoot = removeTabFromTree(internalLayout.root, input.tabId) as SplitNodeInternal;
-  const nextFocusedPaneId =
-    internalLayout.focusedPaneId === null
-      ? null
-      : (findPaneById(nextRoot, internalLayout.focusedPaneId)?.id ??
-        (fallbackPaneId && findPaneById(nextRoot, fallbackPaneId)?.id) ??
-        collectAllPanes(nextRoot)[0]?.id ??
-        DEFAULT_PANE_ID);
+  const parentTabIdByTabId = normalizeParentTabMap({
+    raw: input.layout.parentTabIdByTabId,
+    openTabIds: new Set(collectAllTabs(nextRoot).map((tab) => tab.tabId)),
+  });
+  const nextFocusedPaneId = getFocusedPaneIdAfterTabClose({
+    root: nextRoot,
+    focusedPaneId: internalLayout.focusedPaneId,
+    fallbackPaneId,
+  });
 
-  return {
+  const nextLayout = {
     root: nextRoot,
     focusedPaneId: nextFocusedPaneId,
   };
+  const nextLayoutWithParentMap = parentTabIdByTabId
+    ? { ...nextLayout, parentTabIdByTabId }
+    : nextLayout;
+
+  if (closeSuccessorTabId && findPaneContainingTab(nextRoot, closeSuccessorTabId)) {
+    const focusedLayout =
+      focusTabInLayout({
+        layout: nextLayoutWithParentMap,
+        tabId: closeSuccessorTabId,
+      }) ?? nextLayoutWithParentMap;
+    return parentTabIdByTabId ? { ...focusedLayout, parentTabIdByTabId } : focusedLayout;
+  }
+
+  return nextLayoutWithParentMap;
 }
 
 export function focusTabInLayout(input: FocusTabInLayoutInput): WorkspaceLayout | null {
@@ -1143,10 +1173,11 @@ export function focusTabInLayout(input: FocusTabInLayoutInput): WorkspaceLayout
     return null;
   }
 
-  return {
+  return withNormalizedParentTabMap({
     root: focusTabInPane(layout.root, pane.id, input.tabId),
     focusedPaneId: pane.id,
-  };
+    parentTabIdByTabId: input.layout.parentTabIdByTabId,
+  });
 }
 
 export function retargetTabInLayout(
@@ -1197,14 +1228,15 @@ export function retargetTabInLayout(
     // target identity immediately so local tab state cannot masquerade as the
     // previous agent/terminal/file.
     tabId: nextTabId,
-    layout: {
+    layout: withNormalizedParentTabMap({
       root: replaceTabInTree(layout.root, {
         tabId: input.tabId,
         nextTabId,
         target: input.target,
       }),
       focusedPaneId: layout.focusedPaneId,
-    },
+      parentTabIdByTabId: input.layout.parentTabIdByTabId,
+    }),
   };
 }
 
@@ -1243,14 +1275,15 @@ export function convertDraftToAgentInLayout(
 
   return {
     tabId: canonicalTabId,
-    layout: {
+    layout: withNormalizedParentTabMap({
       root: replaceTabInTree(layout.root, {
         tabId: input.tabId,
         nextTabId: canonicalTabId,
         target,
       }),
       focusedPaneId: layout.focusedPaneId,
-    },
+      parentTabIdByTabId: input.layout.parentTabIdByTabId,
+    }),
   };
 }
 
@@ -1262,13 +1295,14 @@ export function reorderFocusedPaneTabsInLayout(
     return null;
   }
 
-  return {
+  return withNormalizedParentTabMap({
     root: updatePaneInTree(layout.root, {
       paneId: layout.focusedPaneId,
       updater: (pane) => reorderTabsForPane({ pane, tabIds: input.tabIds }),
     }),
     focusedPaneId: layout.focusedPaneId,
-  };
+    parentTabIdByTabId: input.layout.parentTabIdByTabId,
+  });
 }
 
 export function splitPaneInLayout(input: SplitPaneInLayoutInput): SplitPaneInLayoutResult | null {
@@ -1293,10 +1327,11 @@ export function splitPaneInLayout(input: SplitPaneInLayoutInput): SplitPaneInLay
 
   return {
     paneId: result.newPaneId,
-    layout: {
+    layout: withNormalizedParentTabMap({
       root: result.root,
       focusedPaneId: result.newPaneId,
-    },
+      parentTabIdByTabId: input.layout.parentTabIdByTabId,
+    }),
   };
 }
 
@@ -1350,10 +1385,11 @@ export function splitPaneEmptyInLayout(
 
   return {
     paneId: newPaneId,
-    layout: {
+    layout: withNormalizedParentTabMap({
       root: nextRoot,
       focusedPaneId: newPaneId,
-    },
+      parentTabIdByTabId: input.layout.parentTabIdByTabId,
+    }),
   };
 }
 
@@ -1372,14 +1408,15 @@ export function moveTabToPaneInLayout(input: MoveTabToPaneInLayoutInput): Worksp
     return null;
   }
 
-  return {
+  return withNormalizedParentTabMap({
     root: insertTabIntoPane(detached.root, {
       paneId: input.toPaneId,
       tab: detached.tab,
       focusTabId: input.tabId,
     }),
     focusedPaneId: input.toPaneId,
-  };
+    parentTabIdByTabId: input.layout.parentTabIdByTabId,
+  });
 }
 
 export function focusPaneInLayout(input: FocusPaneInLayoutInput): WorkspaceLayout | null {
@@ -1389,21 +1426,23 @@ export function focusPaneInLayout(input: FocusPaneInLayoutInput): WorkspaceLayou
   if (input.layout.focusedPaneId === input.paneId) {
     return null;
   }
-  return {
+  return withNormalizedParentTabMap({
     root: input.layout.root,
     focusedPaneId: input.paneId,
-  };
+    parentTabIdByTabId: input.layout.parentTabIdByTabId,
+  });
 }
 
 export function resizeSplitInLayout(input: ResizeSplitInLayoutInput): WorkspaceLayout {
   const layout = asInternalLayout(input.layout);
-  return {
+  return withNormalizedParentTabMap({
     root: updateGroupSizesInTree(layout.root, {
       groupId: input.groupId,
       sizes: input.sizes,
     }),
     focusedPaneId: layout.focusedPaneId,
-  };
+    parentTabIdByTabId: input.layout.parentTabIdByTabId,
+  });
 }
 
 export function reorderPaneTabsInLayout(
@@ -1414,13 +1453,14 @@ export function reorderPaneTabsInLayout(
     return null;
   }
 
-  return {
+  return withNormalizedParentTabMap({
     root: updatePaneInTree(layout.root, {
       paneId: input.paneId,
       updater: (pane) => reorderTabsForPane({ pane, tabIds: input.tabIds }),
     }),
     focusedPaneId: layout.focusedPaneId,
-  };
+    parentTabIdByTabId: input.layout.parentTabIdByTabId,
+  });
 }
 
 function normalizeStringSet(values: Iterable<string>): Set<string> {
@@ -1434,6 +1474,86 @@ function normalizeStringSet(values: Iterable<string>): Set<string> {
   return next;
 }
 
+function normalizeParentTabMap(input: {
+  raw: unknown;
+  openTabIds: ReadonlySet<string>;
+}): Record<string, string> | undefined {
+  if (!input.raw || typeof input.raw !== "object" || Array.isArray(input.raw)) {
+    return undefined;
+  }
+
+  const next: Record<string, string> = {};
+  for (const [rawChildId, rawParentId] of Object.entries(input.raw)) {
+    const childId = trimNonEmpty(rawChildId);
+    const parentId = trimNonEmpty(typeof rawParentId === "string" ? rawParentId : null);
+    if (
+      !childId ||
+      !parentId ||
+      childId === parentId ||
+      !input.openTabIds.has(childId) ||
+      !input.openTabIds.has(parentId)
+    ) {
+      continue;
+    }
+    next[childId] = parentId;
+  }
+
+  return Object.keys(next).length > 0 ? next : undefined;
+}
+
+function normalizeLayoutParentTabMap(layout: WorkspaceLayout): Record<string, string> | undefined {
+  return normalizeParentTabMap({
+    raw: layout.parentTabIdByTabId,
+    openTabIds: new Set(collectAllTabs(layout.root).map((tab) => tab.tabId)),
+  });
+}
+
+function withNormalizedParentTabMap(layout: WorkspaceLayout): WorkspaceLayout {
+  const parentTabIdByTabId = normalizeLayoutParentTabMap(layout);
+  return parentTabIdByTabId
+    ? { ...layout, parentTabIdByTabId }
+    : { root: layout.root, focusedPaneId: layout.focusedPaneId };
+}
+
+function getCloseSuccessorTabId(input: {
+  pane: SplitPane;
+  tabId: string;
+  openTabIds: ReadonlySet<string>;
+  parentTabIdByTabId?: Record<string, string>;
+}): string | null {
+  if (input.pane.focusedTabId !== input.tabId) {
+    return null;
+  }
+
+  const tabIndex = input.pane.tabIds.indexOf(input.tabId);
+  const parentTabId = input.parentTabIdByTabId?.[input.tabId] ?? null;
+  if (parentTabId && input.openTabIds.has(parentTabId)) {
+    return parentTabId;
+  }
+
+  return (
+    input.pane.tabIds[tabIndex + 1] ??
+    (tabIndex > 0 ? input.pane.tabIds[tabIndex - 1] : null) ??
+    null
+  );
+}
+
+function getFocusedPaneIdAfterTabClose(input: {
+  root: SplitNode;
+  focusedPaneId: string | null;
+  fallbackPaneId: string | null;
+}): string | null {
+  if (input.focusedPaneId === null) {
+    return null;
+  }
+  return (
+    findPaneById(input.root, input.focusedPaneId)?.id ??
+    (input.fallbackPaneId && findPaneById(input.root, input.fallbackPaneId)?.id) ??
+    collectAllPanes(input.root)[0]?.id ??
+    DEFAULT_PANE_ID
+  );
+}
+
 function isEntityTarget(
   target: WorkspaceTabTarget,
 ): target is Extract<WorkspaceTabTarget, { kind: "agent" | "terminal" }> {
@@ -1464,7 +1584,7 @@ function openEntityTabWithoutFocusing(
   invariant(focusedPane, "Workspace layout must always have a pane");
 
   const tabId = buildDeterministicWorkspaceTabId(target);
-  return {
+  return withNormalizedParentTabMap({
     root: insertTabIntoPane(internalLayout.root, {
       paneId: focusedPane.id,
       tab: {
@@ -1475,7 +1595,8 @@ function openEntityTabWithoutFocusing(
       focusTabId: focusedPane.focusedTabId ?? tabId,
     }),
     focusedPaneId: internalLayout.focusedPaneId,
-  };
+    parentTabIdByTabId: layout.parentTabIdByTabId,
+  });
 }
 
 interface EntityTabGroup {
@@ -1654,14 +1775,15 @@ export function reconcileWorkspaceTabs(
       reconciledFocusedTabId = canonicalTabId;
     }
     if (keeper.tabId !== canonicalTabId || !workspaceTabTargetsEqual(keeper.target, group.target)) {
-      nextLayout = {
+      nextLayout = withNormalizedParentTabMap({
         root: replaceTabInTree(asInternalLayout(nextLayout).root, {
           tabId: keeper.tabId,
           nextTabId: canonicalTabId,
           target: group.target,
         }),
         focusedPaneId: nextLayout.focusedPaneId,
-      };
+        parentTabIdByTabId: nextLayout.parentTabIdByTabId,
+      });
     }
     for (const tab of group.tabs) {
       if (tab.tabId === keeper.tabId) {