Fix mobile session pane cleanup

Author: pallyoung

packages/web/src/features/agent-panes/actions/use-pane-actions.ts

@@ -8,6 +8,7 @@ import {
   assignSessionToPane,
   closeDraftPaneById,
   closePaneBySessionId,
+  removePaneBySessionId,
   splitPaneByPaneId,
   splitPaneBySessionId,
 } from "../pane-layout-tree";
@@ -56,6 +57,13 @@ export function usePaneActions(workspaceId: string) {
     [applyLayout]
   );
 
+  const removeSessionPane = useCallback(
+    (sessionId: string) => {
+      applyLayout((current) => removePaneBySessionId(current, sessionId));
+    },
+    [applyLayout]
+  );
+
   const assignSession = useCallback(
     (paneId: string, sessionId: string) => {
       applyLayout((current) => assignSessionToPane(current, paneId, sessionId));
@@ -92,6 +100,7 @@ export function usePaneActions(workspaceId: string) {
     assignSession,
     closeDraftPane,
     closeSessionPane,
+    removeSessionPane,
     replaceWithSession,
     splitDraftPane,
     splitSessionPane,

packages/web/src/features/agent-panes/actions/use-session-actions.test.tsx

@@ -121,7 +121,7 @@ describe("useSessionActions", () => {
       });
 
       await vi.advanceTimersByTimeAsync(100);
-      await expect(closePromise).resolves.toBeUndefined();
+      await expect(closePromise).resolves.toBe(true);
     } finally {
       if (windowDescriptor) {
         Object.defineProperty(globalThis, "window", windowDescriptor);

packages/web/src/features/agent-panes/actions/use-session-actions.ts

@@ -32,40 +32,43 @@ export function useSessionActions() {
     async (sessionId: string) => {
       const session = store.get(sessionByIdAtomFamily(sessionId));
       if (!session) {
-        return;
+        return false;
       }
 
       if (session.state === "ended") {
         const removeResult = await dispatch<void>("session.remove", { sessionId });
         if (!removeResult.ok) {
           console.error("Failed to remove ended session:", removeResult.error?.message);
+          return false;
         }
-        return;
+        return true;
       }
 
       const stopResult = await dispatch<void>("session.stop", { sessionId });
       if (!stopResult.ok && stopResult.error?.code !== "invalid_state") {
         console.error("Failed to stop session before removal:", stopResult.error?.message);
-        return;
+        return false;
       }
 
       const deadline = Date.now() + SESSION_REMOVAL_TIMEOUT_MS;
       while (Date.now() < deadline) {
         const current = store.get(sessionByIdAtomFamily(sessionId));
         if (!current) {
-          return;
+          return true;
         }
         if (current.state === "ended") {
           const removeResult = await dispatch<void>("session.remove", { sessionId });
           if (!removeResult.ok) {
             console.error("Failed to remove ended session:", removeResult.error?.message);
+            return false;
           }
-          return;
+          return true;
         }
         await delay(SESSION_REMOVAL_POLL_INTERVAL_MS);
       }
 
       console.error("Timed out waiting for session to end before removal:", sessionId);
+      return false;
     },
     [dispatch, store]
   );

packages/web/src/features/agent-panes/actions/use-workspace-sessions.ts

@@ -94,7 +94,9 @@ export function useWorkspaceSessions(
 
         const currentLayout = store.get(paneLayoutAtomFamily(workspaceId));
         const workspacePaneLayout = normalizePaneLayout(workspace?.uiState.paneLayout);
-        const legacyPaneLayout = workspacePaneLayout ? null : readLegacyPaneLayout(workspace.id);
+        const legacyPaneLayout = workspacePaneLayout
+          ? null
+          : normalizePaneLayout(readLegacyPaneLayout(workspace.id));
         const baseLayout =
           workspacePaneLayout ?? legacyPaneLayout ?? currentLayout ?? defaultPaneLayout;
         const displayableSessionIds = new Set(
@@ -172,13 +174,19 @@ export function useWorkspaceSessions(
   };
 }
 
-function normalizePaneLayout(layout: Workspace["uiState"]["paneLayout"]): PaneNode | null {
+function normalizePaneLayout(
+  layout: Workspace["uiState"]["paneLayout"] | PaneNode | null | undefined
+): PaneNode | null {
   if (!layout) {
     return null;
   }
 
   return {
-    ...layout,
+    id: layout.id,
+    type: layout.type,
+    sessionId: layout.sessionId,
+    direction: layout.direction,
+    ratio: "ratio" in layout ? layout.ratio : undefined,
     children: layout.children?.map((child) => normalizePaneLayout(child) ?? defaultPaneLayout),
   };
 }

packages/web/src/features/agent-panes/pane-layout-tree.test.ts

@@ -6,6 +6,7 @@ import {
   closeDraftPaneById,
   closePaneBySessionId,
   createFallbackPaneLayout,
+  removePaneBySessionId,
   splitPaneByPaneId,
   splitPaneBySessionId,
 } from "./pane-layout-tree";
@@ -59,6 +60,25 @@ describe("pane-layout-tree", () => {
     });
   });
 
+  it("removes a session pane and collapses the split when explicitly requested", () => {
+    const layout: PaneNode = {
+      id: "root",
+      type: "split",
+      direction: "vertical",
+      ratio: 0.5,
+      children: [
+        { id: "left", type: "leaf", sessionId: "sess_1" },
+        { id: "right", type: "leaf", sessionId: "sess_2" },
+      ],
+    };
+
+    expect(removePaneBySessionId(layout, "sess_2")).toEqual({
+      id: "left",
+      type: "leaf",
+      sessionId: "sess_1",
+    });
+  });
+
   it("assigns a session to the matching draft pane without touching siblings", () => {
     const layout: PaneNode = {
       id: "root",

packages/web/src/features/agent-panes/pane-layout-tree.ts

@@ -186,7 +186,7 @@ export function closeDraftPaneById(node: PaneNode, paneId: string): PaneNode {
 
 /**
  * Close a session pane by turning it into a draft leaf while preserving the
- * existing split structure. This matches the session-card close behavior:
+ * existing split structure. This matches the desktop session-card close behavior:
  * the session ends, but the workspace layout remains stable so the user can
  * immediately launch a replacement session in the same pane.
  */
@@ -221,6 +221,49 @@ function replaceSessionWithDraft(node: PaneNode, sessionId: string): PaneNode {
   };
 }
 
+export function removePaneBySessionId(node: PaneNode, sessionId: string): PaneNode {
+  return removeSessionPane(node, sessionId) ?? { id: node.id, type: "leaf" };
+}
+
+function removeSessionPane(node: PaneNode, sessionId: string): PaneNode | null {
+  if (node.type === "leaf") {
+    if (node.sessionId === sessionId) {
+      return null;
+    }
+    return node;
+  }
+
+  const children = node.children ?? [];
+  let changed = false;
+  const nextChildren: PaneNode[] = [];
+  for (const child of children) {
+    const nextChild = removeSessionPane(child, sessionId);
+    if (nextChild !== child) {
+      changed = true;
+    }
+    if (nextChild !== null) {
+      nextChildren.push(nextChild);
+    }
+  }
+
+  if (!changed) {
+    return node;
+  }
+
+  if (nextChildren.length === 1) {
+    return nextChildren[0]!;
+  }
+
+  if (nextChildren.length === 0) {
+    return null;
+  }
+
+  return {
+    ...node,
+    children: nextChildren,
+  };
+}
+
 export function paneLayoutHasSession(node: PaneNode, sessionIds: Set<string>): boolean {
   if (node.type === "leaf") {
     return node.sessionId ? sessionIds.has(node.sessionId) : false;
@@ -317,14 +360,13 @@ export function createFallbackPaneLayout(sessionIds: string[]): PaneNode {
  */
 export function sanitizePaneLayout(node: PaneNode, liveSessionIds: Set<string>): PaneNode {
   if (node.type === "leaf") {
-    // If this leaf references a session that is ended or removed, turn it into a draft
+    // If this leaf references a session that is ended or removed, turn it into a draft.
     if (node.sessionId && !liveSessionIds.has(node.sessionId)) {
       return { id: node.id, type: "leaf" };
     }
     return node;
   }
 
-  // For splits, recursively sanitize all children and keep the structure intact
   const children = node.children ?? [];
   let changed = false;
   const nextChildren = children.map((child) => {
@@ -335,16 +377,10 @@ export function sanitizePaneLayout(node: PaneNode, liveSessionIds: Set<string>):
     return nextChild;
   });
 
-  // If no children changed, return the same node to preserve reference equality
   if (!changed) {
     return node;
   }
 
-  // If all children collapsed to a single leaf, simplify
-  if (nextChildren.length === 1) {
-    return nextChildren[0]!;
-  }
-
   return {
     ...node,
     children: nextChildren,

packages/web/src/features/workspace/actions/use-workspace-screen-model.ts

@@ -1,6 +1,6 @@
 import type { GitStatus, Session } from "@coder-studio/core";
 import { useAtomValue, useSetAtom, useStore } from "jotai";
-import { useCallback, useEffect, useMemo, useState } from "react";
+import { useCallback, useEffect, useMemo, useRef, useState } from "react";
 import { dispatchCommandAtom } from "../../../atoms/connection";
 import {
   activeWorkspaceAtom,
@@ -63,6 +63,7 @@ export function useWorkspaceScreenModel() {
   const [mobileSheet, setMobileSheet] = useState<MobileWorkspaceSheetKind>(null);
   const [mobileFilesRoute, setMobileFilesRoute] = useState<MobileFilesRoute>({ kind: "root" });
   const [mobileActiveSessionId, setMobileActiveSessionId] = useState<string | null>(null);
+  const mobileSelectionVersionRef = useRef(0);
 
   useEffect(() => {
     if (!workspace) {
@@ -170,6 +171,7 @@ export function useWorkspaceScreenModel() {
         paneActions.appendSession(sessionId, mobileActiveSessionId, "vertical");
       }
 
+      mobileSelectionVersionRef.current += 1;
       setMobileActiveSessionId(sessionId);
     },
     [mobileActiveSessionId, orderedSessions, paneActions, sessions]
@@ -178,23 +180,38 @@ export function useWorkspaceScreenModel() {
   const handleMobileSessionCreated = useCallback(
     (sessionId: string) => {
       paneActions.appendSession(sessionId, mobileActiveSessionId, "vertical");
+      mobileSelectionVersionRef.current += 1;
       setMobileActiveSessionId(sessionId);
     },
     [mobileActiveSessionId, paneActions]
   );
 
   const closeMobileSession = useCallback(
     async (sessionId: string) => {
+      const wasActive = mobileActiveSessionId === sessionId;
+      const selectionVersionAtCloseStart = mobileSelectionVersionRef.current;
       const remainingSessions = mobileAgentSessions.filter((session) => session.id !== sessionId);
       const nextActiveSessionId = remainingSessions[0]?.id ?? null;
 
-      paneActions.closeSessionPane(sessionId);
-      setMobileActiveSessionId((current) =>
-        current === sessionId ? nextActiveSessionId : current
-      );
-      await sessionActions.closeSession(sessionId);
+      if (wasActive) {
+        setMobileActiveSessionId(nextActiveSessionId);
+      }
+
+      const closed = await sessionActions.closeSession(sessionId);
+      if (!closed) {
+        if (!wasActive || mobileSelectionVersionRef.current !== selectionVersionAtCloseStart) {
+          return;
+        }
+
+        setMobileActiveSessionId((current) =>
+          current === nextActiveSessionId ? sessionId : current
+        );
+        return;
+      }
+
+      paneActions.removeSessionPane(sessionId);
     },
-    [mobileAgentSessions, paneActions, sessionActions]
+    [mobileActiveSessionId, mobileAgentSessions, paneActions, sessionActions]
   );
 
   const restoreMobileSession = useCallback(
@@ -207,6 +224,7 @@ export function useWorkspaceScreenModel() {
         paneActions.appendSession(sessionId, mobileActiveSessionId, "vertical");
       }
 
+      mobileSelectionVersionRef.current += 1;
       setMobileActiveSessionId(sessionId);
     },
     [mobileActiveSessionId, orderedSessions, paneActions, sessions]