From 3d36fa4905316a74d6107eb8cde6cc77e73ffd1f Mon Sep 17 00:00:00 2001
From: Jeremy lewi <jeremy@lewi.us>
Date: Wed, 13 May 2026 14:42:45 -0700
Subject: [PATCH 1/6] Restore notebook cell focus

Signed-off-by: Jeremy lewi <jeremy@lewi.us>
---
 app/src/components/Actions/Actions.tsx        | 166 +++++++++++++++++-
 app/src/components/Actions/Editor.tsx         |  12 +-
 .../components/Actions/MarkdownCell.test.tsx  | 144 +++++++++++++++
 app/src/components/Actions/MarkdownCell.tsx   |  67 ++++++-
 4 files changed, 378 insertions(+), 11 deletions(-)
 create mode 100644 app/src/components/Actions/MarkdownCell.test.tsx

diff --git a/app/src/components/Actions/Actions.tsx b/app/src/components/Actions/Actions.tsx
index 565206b3..643f8168 100644
--- a/app/src/components/Actions/Actions.tsx
+++ b/app/src/components/Actions/Actions.tsx
@@ -287,6 +287,13 @@ type SupportedLanguage =
   | "markdown"
   | "python";
 
+type CellFocusRole = "editor" | "rendered";
+
+type NotebookActiveCell = {
+  refId: string;
+  focusRole: CellFocusRole;
+};
+
 const outputTextDecoder = new TextDecoder();
 const ALWAYS_SKIP_MIMES = new Set<string>([MimeType.StatefulRunmeTerminal]);
 
@@ -481,12 +488,18 @@ export function ActionOutputItems({ outputs }: { outputs: parser_pb.CellOutput[]
 
 export function Action({
   cellData,
-  docUri,
+  docUri = "",
   isFirst,
+  restoreFocusRequest = 0,
+  restoreFocusRole = "editor",
+  onFocusStateChange,
 }: {
   cellData: CellData;
-  docUri: string;
+  docUri?: string;
   isFirst: boolean;
+  restoreFocusRequest?: number;
+  restoreFocusRole?: CellFocusRole;
+  onFocusStateChange?: (state: NotebookActiveCell) => void;
 }) {
   const { store } = useNotebookStore();
   const { listRunners, defaultRunnerName } = useRunners();
@@ -621,6 +634,28 @@ export function Action({
     [],
   );
 
+  const handleFocusCapture = useCallback(
+    (event: React.FocusEvent<HTMLDivElement>) => {
+      if (!cell?.refId || !onFocusStateChange) {
+        return;
+      }
+      const target = event.target;
+      if (!(target instanceof HTMLElement)) {
+        return;
+      }
+      const focusRole =
+        target.closest<HTMLElement>("[data-cell-focus-role]")?.dataset
+          .cellFocusRole === "rendered"
+          ? "rendered"
+          : "editor";
+      onFocusStateChange({
+        refId: cell.refId,
+        focusRole,
+      });
+    },
+    [cell?.refId, onFocusStateChange],
+  );
+
   const handleRemoveCell = useCallback(() => {
     cellData.remove();
     setContextMenu(null);
@@ -963,7 +998,9 @@ export function Action({
         id={`markdown-action-${cell.refId}`}
         className="group/cell relative flex min-w-0"
         onContextMenu={handleContextMenu}
+        onFocusCapture={handleFocusCapture}
         data-testid="markdown-action"
+        data-cell-ref-id={cell.refId}
       >
         {/* Left gutter: top + bottom add-cell buttons */}
         <div id={`markdown-gutter-${cell.refId}`} className="flex w-7 shrink-0 flex-col items-center justify-between py-1">
@@ -994,6 +1031,8 @@ export function Action({
               languageOptions={LANGUAGE_OPTIONS}
               onLanguageChange={handleLanguageChange}
               forceEditRequest={markdownEditRequest}
+              restoreFocusRequest={restoreFocusRequest}
+              restoreFocusRole={restoreFocusRole}
             />
             {/* Trash icon on the right, visible on hover */}
             <button
@@ -1051,7 +1090,9 @@ export function Action({
       id={`code-action-${cell.refId}`}
       className="group/cell relative flex"
       onContextMenu={handleContextMenu}
+      onFocusCapture={handleFocusCapture}
       data-testid="code-action"
+      data-cell-ref-id={cell.refId}
     >
       {/* Left gutter: top + bottom add-cell buttons */}
       <div id={`code-gutter-${cell.refId}`} className="flex w-7 shrink-0 flex-col items-center justify-between py-1">
@@ -1080,7 +1121,10 @@ export function Action({
           className="cell-card"
         >
           {/* Code editor section — overflow-hidden keeps border-radius clipping on the editor */}
-          <div className="overflow-hidden rounded-t-nb-md">
+          <div
+            className="overflow-hidden rounded-t-nb-md"
+            data-cell-focus-role="editor"
+          >
             <Editor
               key={`editor-${cell.refId}-${selectedLanguage}`}
               id={cell.refId}
@@ -1088,6 +1132,7 @@ export function Action({
               language={editorLanguage}
               fontSize={fontSettings.fontSize}
               fontFamily={fontSettings.fontFamily}
+              focusRequest={restoreFocusRequest}
               onChange={(v) => {
                 const updated = create(parser_pb.CellSchema, cell);
                 updated.value = v;
@@ -1270,7 +1315,17 @@ export function Action({
   );
 }
 
-function NotebookTabContent({ docUri }: { docUri: string }) {
+function NotebookTabContent({
+  docUri,
+  activeCell,
+  restoreFocusRequest,
+  onCellFocus,
+}: {
+  docUri: string;
+  activeCell: NotebookActiveCell | null;
+  restoreFocusRequest: number;
+  onCellFocus: (docUri: string, state: NotebookActiveCell) => void;
+}) {
   const { getNotebookData, useNotebookSnapshot } = useNotebookContext();
   const notebookSnapshot = useNotebookSnapshot(docUri);
   const cellDatas = useMemo(() => {
@@ -1329,6 +1384,11 @@ function NotebookTabContent({ docUri }: { docUri: string }) {
                   cellData={cellData}
                   docUri={docUri}
                   isFirst={index === 0}
+                  restoreFocusRequest={
+                    activeCell?.refId === refId ? restoreFocusRequest : 0
+                  }
+                  restoreFocusRole={activeCell?.focusRole ?? "editor"}
+                  onFocusStateChange={(state) => onCellFocus(docUri, state)}
                 />
               );
             })}
@@ -1363,6 +1423,12 @@ export default function Actions() {
     Boolean(driveLinkSnapshot.lastErrorMessage);
   const [mountedTabs, setMountedTabs] = useState<Set<string>>(() => new Set());
   const [selectedTabUri, setSelectedTabUri] = useState<string | null>(null);
+  const [activeCellsByDoc, setActiveCellsByDoc] = useState<
+    Record<string, NotebookActiveCell>
+  >({});
+  const [restoreFocusRequestsByDoc, setRestoreFocusRequestsByDoc] = useState<
+    Record<string, number>
+  >({});
   // Empty-state hint visibility is stored locally so the hint panel can be
   // revealed on demand without cluttering the default view.
   const [showConsoleHints, setShowConsoleHints] = useState(false);
@@ -1405,11 +1471,45 @@ export default function Actions() {
   // }, [cellsInitialized, currentDocUri, ensureNotebook, run, runName, setCurrentDoc]);
 
   const { registerRenderer, unregisterRenderer } = useOutput();
+  const currentDocUriRef = useRef(currentDocUri);
   const resolvedSelectedTabUri =
     selectedTabUri ??
     currentDocUri ??
     (statusTabVisible ? DRIVE_LINK_STATUS_TAB_URI : openNotebooks[0]?.uri ?? "");
 
+  useEffect(() => {
+    currentDocUriRef.current = currentDocUri;
+  }, [currentDocUri]);
+
+  const handleCellFocus = useCallback(
+    (docUri: string, state: NotebookActiveCell) => {
+      setActiveCellsByDoc((prev) => {
+        const current = prev[docUri];
+        if (
+          current?.refId === state.refId &&
+          current.focusRole === state.focusRole
+        ) {
+          return prev;
+        }
+        return {
+          ...prev,
+          [docUri]: state,
+        };
+      });
+    },
+    [],
+  );
+
+  const requestFocusRestore = useCallback((docUri: string | null | undefined) => {
+    if (!docUri) {
+      return;
+    }
+    setRestoreFocusRequestsByDoc((prev) => ({
+      ...prev,
+      [docUri]: (prev[docUri] ?? 0) + 1,
+    }));
+  }, []);
+
   // Ensure the active tab is tracked as mounted on first render/whenever it changes.
   useEffect(() => {
     if (!currentDocUri) {
@@ -1439,6 +1539,57 @@ export default function Actions() {
     }
   }, [currentDocUri, openNotebooks, selectedTabUri, statusTabVisible]);
 
+  useEffect(() => {
+    const openUris = new Set(openNotebooks.map((doc) => doc.uri));
+    setActiveCellsByDoc((prev) => {
+      let changed = false;
+      const next: Record<string, NotebookActiveCell> = {};
+      for (const [uri, state] of Object.entries(prev)) {
+        if (!openUris.has(uri)) {
+          changed = true;
+          continue;
+        }
+        next[uri] = state;
+      }
+      return changed ? next : prev;
+    });
+    setRestoreFocusRequestsByDoc((prev) => {
+      let changed = false;
+      const next: Record<string, number> = {};
+      for (const [uri, request] of Object.entries(prev)) {
+        if (!openUris.has(uri)) {
+          changed = true;
+          continue;
+        }
+        next[uri] = request;
+      }
+      return changed ? next : prev;
+    });
+  }, [openNotebooks]);
+
+  useEffect(() => {
+    if (typeof window === "undefined" || typeof document === "undefined") {
+      return;
+    }
+
+    const restoreCurrentNotebookFocus = () => {
+      if (document.visibilityState !== "visible") {
+        return;
+      }
+      requestFocusRestore(currentDocUriRef.current);
+    };
+
+    window.addEventListener("focus", restoreCurrentNotebookFocus);
+    document.addEventListener("visibilitychange", restoreCurrentNotebookFocus);
+    return () => {
+      window.removeEventListener("focus", restoreCurrentNotebookFocus);
+      document.removeEventListener(
+        "visibilitychange",
+        restoreCurrentNotebookFocus,
+      );
+    };
+  }, [requestFocusRestore]);
+
   // Keep the active tab discoverable when the tab rail overflows horizontally.
   // We track each rendered tab node and ask the browser to reveal the selected
   // one so keyboard/mouse tab changes do not leave the active notebook clipped.
@@ -1884,7 +2035,12 @@ export default function Actions() {
               asChild
             >
               <TabPanel className="flex-1 min-h-0" data-document-id={doc.uri}>
-                <NotebookTabContent docUri={doc.uri} />
+                <NotebookTabContent
+                  docUri={doc.uri}
+                  activeCell={activeCellsByDoc[doc.uri] ?? null}
+                  restoreFocusRequest={restoreFocusRequestsByDoc[doc.uri] ?? 0}
+                  onCellFocus={handleCellFocus}
+                />
               </TabPanel>
             </Tabs.Content>
           ))}
diff --git a/app/src/components/Actions/Editor.tsx b/app/src/components/Actions/Editor.tsx
index bed4da5c..82777192 100644
--- a/app/src/components/Actions/Editor.tsx
+++ b/app/src/components/Actions/Editor.tsx
@@ -17,6 +17,7 @@ const Editor = memo(
     readOnly = false,
     ariaLabel,
     autoFocusWhenEmpty = true,
+    focusRequest = 0,
     onChange,
     onEnter,
     onMount,
@@ -29,6 +30,7 @@ const Editor = memo(
     readOnly?: boolean;
     ariaLabel?: string;
     autoFocusWhenEmpty?: boolean;
+    focusRequest?: number;
     onChange: (value: string) => void;
     onEnter: () => void;
     onMount?: (editor: any, monaco: any) => void;
@@ -160,6 +162,13 @@ const Editor = memo(
       editorRef.current.layout?.({ width, height });
     }, [width, height]);
 
+    useEffect(() => {
+      if (!focusRequest || readOnly || !editorRef.current) {
+        return;
+      }
+      editorRef.current.focus?.();
+    }, [focusRequest, readOnly]);
+
     useEffect(() => {
       return () => {
         contentSizeListener.current?.dispose?.();
@@ -220,7 +229,8 @@ const Editor = memo(
       prevProps.language === nextProps.language &&
       prevProps.readOnly === nextProps.readOnly &&
       prevProps.ariaLabel === nextProps.ariaLabel &&
-      prevProps.autoFocusWhenEmpty === nextProps.autoFocusWhenEmpty
+      prevProps.autoFocusWhenEmpty === nextProps.autoFocusWhenEmpty &&
+      prevProps.focusRequest === nextProps.focusRequest
     );
   },
 );
diff --git a/app/src/components/Actions/MarkdownCell.test.tsx b/app/src/components/Actions/MarkdownCell.test.tsx
new file mode 100644
index 00000000..a3f0562c
--- /dev/null
+++ b/app/src/components/Actions/MarkdownCell.test.tsx
@@ -0,0 +1,144 @@
+// @vitest-environment jsdom
+import { describe, expect, it, vi } from "vitest";
+import { act, render, screen } from "@testing-library/react";
+import { clone, create } from "@bufbuild/protobuf";
+import React from "react";
+
+import { parser_pb } from "../../runme/client";
+import type { CellData } from "../../lib/notebookData";
+import MarkdownCell from "./MarkdownCell";
+
+vi.mock("./Editor", () => ({
+  default: ({
+    id,
+    value,
+    onChange,
+    focusRequest = 0,
+  }: {
+    id: string;
+    value: string;
+    onChange: (value: string) => void;
+    focusRequest?: number;
+  }) => {
+    const ref = React.useRef<HTMLTextAreaElement | null>(null);
+
+    React.useEffect(() => {
+      if (!focusRequest) {
+        return;
+      }
+      ref.current?.focus();
+    }, [focusRequest]);
+
+    return (
+      <textarea
+        ref={ref}
+        data-testid="mock-markdown-editor-input"
+        id={id}
+        value={value}
+        onChange={(event) => onChange(event.target.value)}
+      />
+    );
+  },
+}));
+
+class StubCellData {
+  snapshot: parser_pb.Cell;
+  private listeners = new Set<() => void>();
+
+  constructor(cell: parser_pb.Cell) {
+    this.snapshot = cell;
+  }
+
+  subscribeToContentChange(listener: () => void) {
+    this.listeners.add(listener);
+    return () => this.listeners.delete(listener);
+  }
+
+  update(nextCell: parser_pb.Cell) {
+    this.snapshot = clone(parser_pb.CellSchema, nextCell);
+    this.listeners.forEach((listener) => listener());
+  }
+}
+
+describe("MarkdownCell", () => {
+  it("restores focus into editor mode for markdown cells", async () => {
+    const cell = create(parser_pb.CellSchema, {
+      refId: "md-focus-editor",
+      kind: parser_pb.CellKind.MARKUP,
+      languageId: "markdown",
+      outputs: [],
+      metadata: {},
+      value: "hello",
+    });
+    const stub = new StubCellData(cell);
+
+    const { rerender } = render(
+      <MarkdownCell
+        cellData={stub as unknown as CellData}
+        selectedLanguage="markdown"
+        languageSelectId="lang-md-focus-editor"
+        languageOptions={[{ label: "Markdown", value: "markdown" }]}
+        onLanguageChange={() => {}}
+      />,
+    );
+
+    expect(screen.getByTestId("markdown-rendered")).toBeTruthy();
+
+    await act(async () => {
+      rerender(
+        <MarkdownCell
+          cellData={stub as unknown as CellData}
+          selectedLanguage="markdown"
+          languageSelectId="lang-md-focus-editor"
+          languageOptions={[{ label: "Markdown", value: "markdown" }]}
+          onLanguageChange={() => {}}
+          restoreFocusRequest={1}
+          restoreFocusRole="editor"
+        />,
+      );
+      await Promise.resolve();
+    });
+
+    expect(screen.getByTestId("markdown-editor")).toBeTruthy();
+    expect(screen.getByTestId("mock-markdown-editor-input")).toHaveFocus();
+  });
+
+  it("restores focus onto rendered markdown when requested", async () => {
+    const cell = create(parser_pb.CellSchema, {
+      refId: "md-focus-rendered",
+      kind: parser_pb.CellKind.MARKUP,
+      languageId: "markdown",
+      outputs: [],
+      metadata: {},
+      value: "hello",
+    });
+    const stub = new StubCellData(cell);
+
+    const { rerender } = render(
+      <MarkdownCell
+        cellData={stub as unknown as CellData}
+        selectedLanguage="markdown"
+        languageSelectId="lang-md-focus-rendered"
+        languageOptions={[{ label: "Markdown", value: "markdown" }]}
+        onLanguageChange={() => {}}
+      />,
+    );
+
+    await act(async () => {
+      rerender(
+        <MarkdownCell
+          cellData={stub as unknown as CellData}
+          selectedLanguage="markdown"
+          languageSelectId="lang-md-focus-rendered"
+          languageOptions={[{ label: "Markdown", value: "markdown" }]}
+          onLanguageChange={() => {}}
+          restoreFocusRequest={1}
+          restoreFocusRole="rendered"
+        />,
+      );
+      await Promise.resolve();
+    });
+
+    expect(screen.getByTestId("markdown-rendered")).toHaveFocus();
+  });
+});
diff --git a/app/src/components/Actions/MarkdownCell.tsx b/app/src/components/Actions/MarkdownCell.tsx
index 03081965..8eec56b0 100644
--- a/app/src/components/Actions/MarkdownCell.tsx
+++ b/app/src/components/Actions/MarkdownCell.tsx
@@ -23,6 +23,7 @@ import {
   useCallback,
   useEffect,
   useMemo,
+  useRef,
   useState,
   useSyncExternalStore,
   type FocusEvent,
@@ -167,6 +168,10 @@ interface MarkdownCellProps {
   onLanguageChange: (event: ChangeEvent<HTMLSelectElement>) => void;
   /** Incrementing signal to force editor mode after conversion to markdown */
   forceEditRequest?: number;
+  /** Incrementing signal to restore focus into this cell */
+  restoreFocusRequest?: number;
+  /** Which markdown surface should regain focus */
+  restoreFocusRole?: "editor" | "rendered";
 }
 
 /**
@@ -186,6 +191,8 @@ const MarkdownCell = memo(
     languageOptions,
     onLanguageChange,
     forceEditRequest = 0,
+    restoreFocusRequest = 0,
+    restoreFocusRole = "editor",
   }: MarkdownCellProps) => {
     // Subscribe to cell data changes using useSyncExternalStore for tearing-safe reads
     const cell = useSyncExternalStore(
@@ -203,9 +210,17 @@ const MarkdownCell = memo(
       const value = cell?.value ?? "";
       return value.trim().length > 0;
     });
+    const renderedRef = useRef<HTMLDivElement | null>(null);
+    const [editorFocusRequest, setEditorFocusRequest] = useState(0);
+    const [pendingRestoreTarget, setPendingRestoreTarget] = useState<
+      "editor" | "rendered" | null
+    >(null);
 
     // Get the current cell value
     const value = cell?.value ?? "";
+    const requestEditorFocus = useCallback(() => {
+      setEditorFocusRequest((request) => request + 1);
+    }, []);
 
     // Enforce invariant: empty cells must be in edit mode.
     // This handles external changes (undo/redo, sync) that clear the value.
@@ -218,15 +233,46 @@ const MarkdownCell = memo(
     useEffect(() => {
       if (forceEditRequest > 0) {
         setRendered(false);
+        requestEditorFocus();
+      }
+    }, [forceEditRequest, requestEditorFocus]);
+
+    useEffect(() => {
+      if (restoreFocusRequest <= 0) {
+        return;
+      }
+      if (restoreFocusRole === "editor" || !value.trim()) {
+        setRendered(false);
+        setPendingRestoreTarget("editor");
+        return;
       }
-    }, [forceEditRequest]);
+      setRendered(true);
+      setPendingRestoreTarget("rendered");
+    }, [restoreFocusRequest, restoreFocusRole, value]);
+
+    useEffect(() => {
+      if (pendingRestoreTarget !== "editor" || rendered) {
+        return;
+      }
+      requestEditorFocus();
+      setPendingRestoreTarget(null);
+    }, [pendingRestoreTarget, rendered, requestEditorFocus]);
+
+    useEffect(() => {
+      if (pendingRestoreTarget !== "rendered" || !rendered) {
+        return;
+      }
+      renderedRef.current?.focus();
+      setPendingRestoreTarget(null);
+    }, [pendingRestoreTarget, rendered]);
 
     /**
      * Handle switching to edit mode when user double-clicks rendered content.
      */
     const handleDoubleClick = useCallback(() => {
       setRendered(false);
-    }, []);
+      requestEditorFocus();
+    }, [requestEditorFocus]);
 
     /**
      * Handle keyboard activation on the rendered container for accessibility.
@@ -243,9 +289,10 @@ const MarkdownCell = memo(
         if (event.key === "Enter" || event.key === " ") {
           event.preventDefault();
           setRendered(false);
+          requestEditorFocus();
         }
       },
-      []
+      [requestEditorFocus]
     );
 
     /**
@@ -348,10 +395,12 @@ const MarkdownCell = memo(
             className="cursor-text rounded-nb-md border border-transparent p-4 transition-[border-color,background-color,box-shadow] duration-200 hover:border-nb-border hover:bg-nb-surface-2/60 hover:shadow-nb-xs"
             onDoubleClick={handleDoubleClick}
             onKeyDown={handleRenderedKeyDown}
+            ref={renderedRef}
             tabIndex={0}
             role="button"
             aria-label="Double-click or press Enter to edit markdown"
             data-testid="markdown-rendered"
+            data-cell-focus-role="rendered"
           >
             {renderedMarkdown}
           </div>
@@ -363,6 +412,7 @@ const MarkdownCell = memo(
             onBlur={handleBlur}
             onKeyDown={handleEditorKeyDown}
             data-testid="markdown-editor"
+            data-cell-focus-role="editor"
           >
             <Editor
               id={`md-editor-${cell.refId}`}
@@ -370,6 +420,7 @@ const MarkdownCell = memo(
               language="markdown"
               fontSize={fontSettings.fontSize}
               fontFamily={fontSettings.fontFamily}
+              focusRequest={editorFocusRequest}
               onChange={handleEditorChange}
               onEnter={handleRun}
             />
@@ -399,8 +450,14 @@ const MarkdownCell = memo(
     );
   },
   (prevProps, nextProps) => {
-    // Skip re-render if the cellData reference hasn't changed
-    return prevProps.cellData === nextProps.cellData;
+    return (
+      prevProps.cellData === nextProps.cellData &&
+      prevProps.selectedLanguage === nextProps.selectedLanguage &&
+      prevProps.languageSelectId === nextProps.languageSelectId &&
+      prevProps.forceEditRequest === nextProps.forceEditRequest &&
+      prevProps.restoreFocusRequest === nextProps.restoreFocusRequest &&
+      prevProps.restoreFocusRole === nextProps.restoreFocusRole
+    );
   }
 );
 

From 6cf46c382d00868e44f4aa09e0a58ba7df1fdea8 Mon Sep 17 00:00:00 2001
From: Jeremy lewi <jeremy@lewi.us>
Date: Wed, 13 May 2026 23:12:13 -0700
Subject: [PATCH 2/6] Skip focus restore on status tab

Signed-off-by: Jeremy lewi <jeremy@lewi.us>
---
 app/src/components/Actions/Actions.tsx | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/app/src/components/Actions/Actions.tsx b/app/src/components/Actions/Actions.tsx
index 643f8168..24792750 100644
--- a/app/src/components/Actions/Actions.tsx
+++ b/app/src/components/Actions/Actions.tsx
@@ -1471,15 +1471,15 @@ export default function Actions() {
   // }, [cellsInitialized, currentDocUri, ensureNotebook, run, runName, setCurrentDoc]);
 
   const { registerRenderer, unregisterRenderer } = useOutput();
-  const currentDocUriRef = useRef(currentDocUri);
   const resolvedSelectedTabUri =
     selectedTabUri ??
     currentDocUri ??
     (statusTabVisible ? DRIVE_LINK_STATUS_TAB_URI : openNotebooks[0]?.uri ?? "");
+  const resolvedSelectedTabUriRef = useRef(resolvedSelectedTabUri);
 
   useEffect(() => {
-    currentDocUriRef.current = currentDocUri;
-  }, [currentDocUri]);
+    resolvedSelectedTabUriRef.current = resolvedSelectedTabUri;
+  }, [resolvedSelectedTabUri]);
 
   const handleCellFocus = useCallback(
     (docUri: string, state: NotebookActiveCell) => {
@@ -1576,7 +1576,11 @@ export default function Actions() {
       if (document.visibilityState !== "visible") {
         return;
       }
-      requestFocusRestore(currentDocUriRef.current);
+      const visibleTabUri = resolvedSelectedTabUriRef.current;
+      if (!visibleTabUri || visibleTabUri === DRIVE_LINK_STATUS_TAB_URI) {
+        return;
+      }
+      requestFocusRestore(visibleTabUri);
     };
 
     window.addEventListener("focus", restoreCurrentNotebookFocus);

From c5108e9ebd16a345267c3a00caaf91722419c0fd Mon Sep 17 00:00:00 2001
From: Jeremy lewi <jeremy@lewi.us>
Date: Thu, 14 May 2026 16:24:10 -0700
Subject: [PATCH 3/6] Persist notebook active cell state

Signed-off-by: Jeremy lewi <jeremy@lewi.us>
---
 app/src/components/Actions/Actions.tsx        | 156 ++++------
 app/src/components/Actions/Editor.tsx         |  13 +-
 .../components/Actions/MarkdownCell.test.tsx  | 170 +++++++++--
 app/src/components/Actions/MarkdownCell.tsx   |  90 +++---
 app/src/lib/notebookActiveCellState.test.ts   |  77 +++++
 app/src/lib/notebookActiveCellState.ts        | 106 +++++++
 app/test/browser/run-cuj-scenarios.ts         |   1 +
 ...est-scenario-notebook-focus-persistence.ts | 278 ++++++++++++++++++
 .../20260514_notebook_focus_persistence.md    | 149 ++++++++++
 9 files changed, 871 insertions(+), 169 deletions(-)
 create mode 100644 app/src/lib/notebookActiveCellState.test.ts
 create mode 100644 app/src/lib/notebookActiveCellState.ts
 create mode 100644 app/test/browser/test-scenario-notebook-focus-persistence.ts
 create mode 100644 docs-dev/design/20260514_notebook_focus_persistence.md

diff --git a/app/src/components/Actions/Actions.tsx b/app/src/components/Actions/Actions.tsx
index 24792750..4bb66bab 100644
--- a/app/src/components/Actions/Actions.tsx
+++ b/app/src/components/Actions/Actions.tsx
@@ -5,7 +5,6 @@ import {
   useEffect,
   useLayoutEffect,
   useMemo,
-  useRef,
   useState,
   useSyncExternalStore,
 } from "react";
@@ -29,6 +28,14 @@ import Editor from "./Editor";
 import MarkdownCell from "./MarkdownCell";
 import { IOPUB_INCOMPLETE_METADATA_KEY } from "../../lib/ipykernel";
 import { appLogger } from "../../lib/logging/runtime";
+import {
+  createNotebookActiveCellState,
+  loadNotebookActiveCellMap,
+  persistNotebookActiveCellMap,
+  type CellFocusRole,
+  type NotebookActiveCellMap,
+  type NotebookActiveCellState,
+} from "../../lib/notebookActiveCellState";
 import { copyNotebookShareUrl } from "../../lib/shareLinks";
 import {
   PlayIcon,
@@ -287,13 +294,6 @@ type SupportedLanguage =
   | "markdown"
   | "python";
 
-type CellFocusRole = "editor" | "rendered";
-
-type NotebookActiveCell = {
-  refId: string;
-  focusRole: CellFocusRole;
-};
-
 const outputTextDecoder = new TextDecoder();
 const ALWAYS_SKIP_MIMES = new Set<string>([MimeType.StatefulRunmeTerminal]);
 
@@ -490,16 +490,18 @@ export function Action({
   cellData,
   docUri = "",
   isFirst,
-  restoreFocusRequest = 0,
-  restoreFocusRole = "editor",
+  isActiveCell = false,
+  activeFocusRole = "editor",
+  isWindowFocused = false,
   onFocusStateChange,
 }: {
   cellData: CellData;
   docUri?: string;
   isFirst: boolean;
-  restoreFocusRequest?: number;
-  restoreFocusRole?: CellFocusRole;
-  onFocusStateChange?: (state: NotebookActiveCell) => void;
+  isActiveCell?: boolean;
+  activeFocusRole?: CellFocusRole;
+  isWindowFocused?: boolean;
+  onFocusStateChange?: (state: NotebookActiveCellState) => void;
 }) {
   const { store } = useNotebookStore();
   const { listRunners, defaultRunnerName } = useRunners();
@@ -648,10 +650,11 @@ export function Action({
           .cellFocusRole === "rendered"
           ? "rendered"
           : "editor";
-      onFocusStateChange({
-        refId: cell.refId,
-        focusRole,
-      });
+      const nextState = createNotebookActiveCellState(cell.refId, focusRole);
+      if (!nextState) {
+        return;
+      }
+      onFocusStateChange(nextState);
     },
     [cell?.refId, onFocusStateChange],
   );
@@ -1031,8 +1034,9 @@ export function Action({
               languageOptions={LANGUAGE_OPTIONS}
               onLanguageChange={handleLanguageChange}
               forceEditRequest={markdownEditRequest}
-              restoreFocusRequest={restoreFocusRequest}
-              restoreFocusRole={restoreFocusRole}
+              isActiveCell={isActiveCell}
+              activeFocusRole={activeFocusRole}
+              isWindowFocused={isWindowFocused}
             />
             {/* Trash icon on the right, visible on hover */}
             <button
@@ -1132,7 +1136,7 @@ export function Action({
               language={editorLanguage}
               fontSize={fontSettings.fontSize}
               fontFamily={fontSettings.fontFamily}
-              focusRequest={restoreFocusRequest}
+              shouldFocus={isActiveCell && isWindowFocused}
               onChange={(v) => {
                 const updated = create(parser_pb.CellSchema, cell);
                 updated.value = v;
@@ -1318,13 +1322,13 @@ export function Action({
 function NotebookTabContent({
   docUri,
   activeCell,
-  restoreFocusRequest,
+  isWindowFocused,
   onCellFocus,
 }: {
   docUri: string;
-  activeCell: NotebookActiveCell | null;
-  restoreFocusRequest: number;
-  onCellFocus: (docUri: string, state: NotebookActiveCell) => void;
+  activeCell: NotebookActiveCellState | null;
+  isWindowFocused: boolean;
+  onCellFocus: (docUri: string, state: NotebookActiveCellState) => void;
 }) {
   const { getNotebookData, useNotebookSnapshot } = useNotebookContext();
   const notebookSnapshot = useNotebookSnapshot(docUri);
@@ -1384,10 +1388,9 @@ function NotebookTabContent({
                   cellData={cellData}
                   docUri={docUri}
                   isFirst={index === 0}
-                  restoreFocusRequest={
-                    activeCell?.refId === refId ? restoreFocusRequest : 0
-                  }
-                  restoreFocusRole={activeCell?.focusRole ?? "editor"}
+                  isActiveCell={activeCell?.refId === refId}
+                  activeFocusRole={activeCell?.focusRole ?? "editor"}
+                  isWindowFocused={isWindowFocused}
                   onFocusStateChange={(state) => onCellFocus(docUri, state)}
                 />
               );
@@ -1423,12 +1426,15 @@ export default function Actions() {
     Boolean(driveLinkSnapshot.lastErrorMessage);
   const [mountedTabs, setMountedTabs] = useState<Set<string>>(() => new Set());
   const [selectedTabUri, setSelectedTabUri] = useState<string | null>(null);
-  const [activeCellsByDoc, setActiveCellsByDoc] = useState<
-    Record<string, NotebookActiveCell>
-  >({});
-  const [restoreFocusRequestsByDoc, setRestoreFocusRequestsByDoc] = useState<
-    Record<string, number>
-  >({});
+  const [activeCellsByDoc, setActiveCellsByDoc] = useState<NotebookActiveCellMap>(
+    () => loadNotebookActiveCellMap(),
+  );
+  const [isWindowFocused, setIsWindowFocused] = useState(() => {
+    if (typeof document === "undefined") {
+      return false;
+    }
+    return document.visibilityState === "visible" && document.hasFocus();
+  });
   // Empty-state hint visibility is stored locally so the hint panel can be
   // revealed on demand without cluttering the default view.
   const [showConsoleHints, setShowConsoleHints] = useState(false);
@@ -1475,14 +1481,9 @@ export default function Actions() {
     selectedTabUri ??
     currentDocUri ??
     (statusTabVisible ? DRIVE_LINK_STATUS_TAB_URI : openNotebooks[0]?.uri ?? "");
-  const resolvedSelectedTabUriRef = useRef(resolvedSelectedTabUri);
-
-  useEffect(() => {
-    resolvedSelectedTabUriRef.current = resolvedSelectedTabUri;
-  }, [resolvedSelectedTabUri]);
 
   const handleCellFocus = useCallback(
-    (docUri: string, state: NotebookActiveCell) => {
+    (docUri: string, state: NotebookActiveCellState) => {
       setActiveCellsByDoc((prev) => {
         const current = prev[docUri];
         if (
@@ -1491,25 +1492,17 @@ export default function Actions() {
         ) {
           return prev;
         }
-        return {
+        const next = {
           ...prev,
           [docUri]: state,
         };
+        persistNotebookActiveCellMap(next);
+        return next;
       });
     },
     [],
   );
 
-  const requestFocusRestore = useCallback((docUri: string | null | undefined) => {
-    if (!docUri) {
-      return;
-    }
-    setRestoreFocusRequestsByDoc((prev) => ({
-      ...prev,
-      [docUri]: (prev[docUri] ?? 0) + 1,
-    }));
-  }, []);
-
   // Ensure the active tab is tracked as mounted on first render/whenever it changes.
   useEffect(() => {
     if (!currentDocUri) {
@@ -1539,60 +1532,27 @@ export default function Actions() {
     }
   }, [currentDocUri, openNotebooks, selectedTabUri, statusTabVisible]);
 
-  useEffect(() => {
-    const openUris = new Set(openNotebooks.map((doc) => doc.uri));
-    setActiveCellsByDoc((prev) => {
-      let changed = false;
-      const next: Record<string, NotebookActiveCell> = {};
-      for (const [uri, state] of Object.entries(prev)) {
-        if (!openUris.has(uri)) {
-          changed = true;
-          continue;
-        }
-        next[uri] = state;
-      }
-      return changed ? next : prev;
-    });
-    setRestoreFocusRequestsByDoc((prev) => {
-      let changed = false;
-      const next: Record<string, number> = {};
-      for (const [uri, request] of Object.entries(prev)) {
-        if (!openUris.has(uri)) {
-          changed = true;
-          continue;
-        }
-        next[uri] = request;
-      }
-      return changed ? next : prev;
-    });
-  }, [openNotebooks]);
-
   useEffect(() => {
     if (typeof window === "undefined" || typeof document === "undefined") {
       return;
     }
 
-    const restoreCurrentNotebookFocus = () => {
-      if (document.visibilityState !== "visible") {
-        return;
-      }
-      const visibleTabUri = resolvedSelectedTabUriRef.current;
-      if (!visibleTabUri || visibleTabUri === DRIVE_LINK_STATUS_TAB_URI) {
-        return;
-      }
-      requestFocusRestore(visibleTabUri);
+    const syncWindowFocus = () => {
+      setIsWindowFocused(
+        document.visibilityState === "visible" && document.hasFocus(),
+      );
     };
 
-    window.addEventListener("focus", restoreCurrentNotebookFocus);
-    document.addEventListener("visibilitychange", restoreCurrentNotebookFocus);
+    syncWindowFocus();
+    window.addEventListener("focus", syncWindowFocus);
+    window.addEventListener("blur", syncWindowFocus);
+    document.addEventListener("visibilitychange", syncWindowFocus);
     return () => {
-      window.removeEventListener("focus", restoreCurrentNotebookFocus);
-      document.removeEventListener(
-        "visibilitychange",
-        restoreCurrentNotebookFocus,
-      );
+      window.removeEventListener("focus", syncWindowFocus);
+      window.removeEventListener("blur", syncWindowFocus);
+      document.removeEventListener("visibilitychange", syncWindowFocus);
     };
-  }, [requestFocusRestore]);
+  }, []);
 
   // Keep the active tab discoverable when the tab rail overflows horizontally.
   // We track each rendered tab node and ask the browser to reveal the selected
@@ -2042,7 +2002,9 @@ export default function Actions() {
                 <NotebookTabContent
                   docUri={doc.uri}
                   activeCell={activeCellsByDoc[doc.uri] ?? null}
-                  restoreFocusRequest={restoreFocusRequestsByDoc[doc.uri] ?? 0}
+                  isWindowFocused={
+                    isWindowFocused && resolvedSelectedTabUri === doc.uri
+                  }
                   onCellFocus={handleCellFocus}
                 />
               </TabPanel>
diff --git a/app/src/components/Actions/Editor.tsx b/app/src/components/Actions/Editor.tsx
index 82777192..3db01257 100644
--- a/app/src/components/Actions/Editor.tsx
+++ b/app/src/components/Actions/Editor.tsx
@@ -17,7 +17,7 @@ const Editor = memo(
     readOnly = false,
     ariaLabel,
     autoFocusWhenEmpty = true,
-    focusRequest = 0,
+    shouldFocus = false,
     onChange,
     onEnter,
     onMount,
@@ -30,7 +30,7 @@ const Editor = memo(
     readOnly?: boolean;
     ariaLabel?: string;
     autoFocusWhenEmpty?: boolean;
-    focusRequest?: number;
+    shouldFocus?: boolean;
     onChange: (value: string) => void;
     onEnter: () => void;
     onMount?: (editor: any, monaco: any) => void;
@@ -48,6 +48,7 @@ const Editor = memo(
     // don't render an empty scroll track.
     const [isClamped, setIsClamped] = useState(false);
     const contentSizeListener = useRef<{ dispose: () => void } | null>(null);
+    const previousShouldFocusRef = useRef(false);
 
     // Keep the ref updated with the latest onEnter
     useEffect(() => {
@@ -163,11 +164,13 @@ const Editor = memo(
     }, [width, height]);
 
     useEffect(() => {
-      if (!focusRequest || readOnly || !editorRef.current) {
+      const wasFocused = previousShouldFocusRef.current;
+      previousShouldFocusRef.current = shouldFocus;
+      if (!shouldFocus || wasFocused || readOnly || !editorRef.current) {
         return;
       }
       editorRef.current.focus?.();
-    }, [focusRequest, readOnly]);
+    }, [readOnly, shouldFocus]);
 
     useEffect(() => {
       return () => {
@@ -230,7 +233,7 @@ const Editor = memo(
       prevProps.readOnly === nextProps.readOnly &&
       prevProps.ariaLabel === nextProps.ariaLabel &&
       prevProps.autoFocusWhenEmpty === nextProps.autoFocusWhenEmpty &&
-      prevProps.focusRequest === nextProps.focusRequest
+      prevProps.shouldFocus === nextProps.shouldFocus
     );
   },
 );
diff --git a/app/src/components/Actions/MarkdownCell.test.tsx b/app/src/components/Actions/MarkdownCell.test.tsx
index a3f0562c..22476661 100644
--- a/app/src/components/Actions/MarkdownCell.test.tsx
+++ b/app/src/components/Actions/MarkdownCell.test.tsx
@@ -1,6 +1,6 @@
 // @vitest-environment jsdom
 import { describe, expect, it, vi } from "vitest";
-import { act, render, screen } from "@testing-library/react";
+import { act, fireEvent, render, screen } from "@testing-library/react";
 import { clone, create } from "@bufbuild/protobuf";
 import React from "react";
 
@@ -13,21 +13,24 @@ vi.mock("./Editor", () => ({
     id,
     value,
     onChange,
-    focusRequest = 0,
+    shouldFocus = false,
   }: {
     id: string;
     value: string;
     onChange: (value: string) => void;
-    focusRequest?: number;
+    shouldFocus?: boolean;
   }) => {
     const ref = React.useRef<HTMLTextAreaElement | null>(null);
+    const previousShouldFocusRef = React.useRef(false);
 
     React.useEffect(() => {
-      if (!focusRequest) {
+      const wasFocused = previousShouldFocusRef.current;
+      previousShouldFocusRef.current = shouldFocus;
+      if (!shouldFocus || wasFocused) {
         return;
       }
       ref.current?.focus();
-    }, [focusRequest]);
+    }, [shouldFocus]);
 
     return (
       <textarea
@@ -60,10 +63,26 @@ class StubCellData {
   }
 }
 
+function renderMarkdownCell(
+  cellData: CellData,
+  props?: Partial<React.ComponentProps<typeof MarkdownCell>>,
+) {
+  return render(
+    <MarkdownCell
+      cellData={cellData}
+      selectedLanguage="markdown"
+      languageSelectId="lang-md-focus"
+      languageOptions={[{ label: "Markdown", value: "markdown" }]}
+      onLanguageChange={() => {}}
+      {...props}
+    />,
+  );
+}
+
 describe("MarkdownCell", () => {
-  it("restores focus into editor mode for markdown cells", async () => {
+  it("starts markdown cells in editor mode when persisted focus target is editor", () => {
     const cell = create(parser_pb.CellSchema, {
-      refId: "md-focus-editor",
+      refId: "md-active-editor",
       kind: parser_pb.CellKind.MARKUP,
       languageId: "markdown",
       outputs: [],
@@ -72,17 +91,31 @@ describe("MarkdownCell", () => {
     });
     const stub = new StubCellData(cell);
 
-    const { rerender } = render(
-      <MarkdownCell
-        cellData={stub as unknown as CellData}
-        selectedLanguage="markdown"
-        languageSelectId="lang-md-focus-editor"
-        languageOptions={[{ label: "Markdown", value: "markdown" }]}
-        onLanguageChange={() => {}}
-      />,
-    );
+    renderMarkdownCell(stub as unknown as CellData, {
+      isActiveCell: true,
+      activeFocusRole: "editor",
+    });
+
+    expect(screen.getByTestId("markdown-editor")).toBeTruthy();
+    expect(screen.queryByTestId("markdown-rendered")).toBeNull();
+  });
+
+  it("focuses the editor when the window regains focus", async () => {
+    const cell = create(parser_pb.CellSchema, {
+      refId: "md-focus-editor",
+      kind: parser_pb.CellKind.MARKUP,
+      languageId: "markdown",
+      outputs: [],
+      metadata: {},
+      value: "hello",
+    });
+    const stub = new StubCellData(cell);
 
-    expect(screen.getByTestId("markdown-rendered")).toBeTruthy();
+    const { rerender } = renderMarkdownCell(stub as unknown as CellData, {
+      isActiveCell: true,
+      activeFocusRole: "editor",
+      isWindowFocused: false,
+    });
 
     await act(async () => {
       rerender(
@@ -92,18 +125,96 @@ describe("MarkdownCell", () => {
           languageSelectId="lang-md-focus-editor"
           languageOptions={[{ label: "Markdown", value: "markdown" }]}
           onLanguageChange={() => {}}
-          restoreFocusRequest={1}
-          restoreFocusRole="editor"
+          isActiveCell
+          activeFocusRole="editor"
+          isWindowFocused
         />,
       );
       await Promise.resolve();
     });
 
+    expect(screen.getByTestId("mock-markdown-editor-input")).toHaveFocus();
+  });
+
+  it("does not replay restore behavior on each editor change", async () => {
+    const cell = create(parser_pb.CellSchema, {
+      refId: "md-typing",
+      kind: parser_pb.CellKind.MARKUP,
+      languageId: "markdown",
+      outputs: [],
+      metadata: {},
+      value: "hello",
+    });
+    const stub = new StubCellData(cell);
+
+    renderMarkdownCell(stub as unknown as CellData, {
+      isActiveCell: true,
+      activeFocusRole: "editor",
+      isWindowFocused: true,
+    });
+
+    const input = screen.getByTestId("mock-markdown-editor-input");
+    fireEvent.change(input, { target: { value: "hello world" } });
+
+    await act(async () => {
+      await Promise.resolve();
+    });
+
     expect(screen.getByTestId("markdown-editor")).toBeTruthy();
+    expect(screen.queryByTestId("markdown-rendered")).toBeNull();
+  });
+
+  it("focuses the editor when rendered markdown is opened for editing", async () => {
+    const cell = create(parser_pb.CellSchema, {
+      refId: "md-manual-edit",
+      kind: parser_pb.CellKind.MARKUP,
+      languageId: "markdown",
+      outputs: [],
+      metadata: {},
+      value: "hello",
+    });
+    const stub = new StubCellData(cell);
+
+    renderMarkdownCell(stub as unknown as CellData);
+
+    fireEvent.doubleClick(screen.getByTestId("markdown-rendered"));
+
+    await act(async () => {
+      await Promise.resolve();
+    });
+
     expect(screen.getByTestId("mock-markdown-editor-input")).toHaveFocus();
   });
 
-  it("restores focus onto rendered markdown when requested", async () => {
+  it("focuses rendered markdown when leaving editor mode with escape", async () => {
+    const cell = create(parser_pb.CellSchema, {
+      refId: "md-escape-rendered",
+      kind: parser_pb.CellKind.MARKUP,
+      languageId: "markdown",
+      outputs: [],
+      metadata: {},
+      value: "hello",
+    });
+    const stub = new StubCellData(cell);
+
+    renderMarkdownCell(stub as unknown as CellData, {
+      isActiveCell: true,
+      activeFocusRole: "editor",
+      isWindowFocused: true,
+    });
+
+    fireEvent.keyDown(screen.getByTestId("markdown-editor"), {
+      key: "Escape",
+    });
+
+    await act(async () => {
+      await Promise.resolve();
+    });
+
+    expect(screen.getByTestId("markdown-rendered")).toHaveFocus();
+  });
+
+  it("focuses rendered markdown when the window regains focus", async () => {
     const cell = create(parser_pb.CellSchema, {
       refId: "md-focus-rendered",
       kind: parser_pb.CellKind.MARKUP,
@@ -114,15 +225,11 @@ describe("MarkdownCell", () => {
     });
     const stub = new StubCellData(cell);
 
-    const { rerender } = render(
-      <MarkdownCell
-        cellData={stub as unknown as CellData}
-        selectedLanguage="markdown"
-        languageSelectId="lang-md-focus-rendered"
-        languageOptions={[{ label: "Markdown", value: "markdown" }]}
-        onLanguageChange={() => {}}
-      />,
-    );
+    const { rerender } = renderMarkdownCell(stub as unknown as CellData, {
+      isActiveCell: true,
+      activeFocusRole: "rendered",
+      isWindowFocused: false,
+    });
 
     await act(async () => {
       rerender(
@@ -132,8 +239,9 @@ describe("MarkdownCell", () => {
           languageSelectId="lang-md-focus-rendered"
           languageOptions={[{ label: "Markdown", value: "markdown" }]}
           onLanguageChange={() => {}}
-          restoreFocusRequest={1}
-          restoreFocusRole="rendered"
+          isActiveCell
+          activeFocusRole="rendered"
+          isWindowFocused
         />,
       );
       await Promise.resolve();
diff --git a/app/src/components/Actions/MarkdownCell.tsx b/app/src/components/Actions/MarkdownCell.tsx
index 8eec56b0..07c10e24 100644
--- a/app/src/components/Actions/MarkdownCell.tsx
+++ b/app/src/components/Actions/MarkdownCell.tsx
@@ -36,6 +36,7 @@ import remarkGfm from "remark-gfm";
 import { create } from "@bufbuild/protobuf";
 import { parser_pb } from "../../runme/client";
 import type { CellData } from "../../lib/notebookData";
+import type { CellFocusRole } from "../../lib/notebookActiveCellState";
 import Editor from "./Editor";
 import { fontSettings } from "./CellConsole";
 
@@ -168,10 +169,12 @@ interface MarkdownCellProps {
   onLanguageChange: (event: ChangeEvent<HTMLSelectElement>) => void;
   /** Incrementing signal to force editor mode after conversion to markdown */
   forceEditRequest?: number;
-  /** Incrementing signal to restore focus into this cell */
-  restoreFocusRequest?: number;
-  /** Which markdown surface should regain focus */
-  restoreFocusRole?: "editor" | "rendered";
+  /** Whether this cell is the persisted active cell for the visible notebook */
+  isActiveCell?: boolean;
+  /** Which markdown surface was last active for this cell */
+  activeFocusRole?: CellFocusRole;
+  /** Whether the visible browser tab is currently focused */
+  isWindowFocused?: boolean;
 }
 
 /**
@@ -191,8 +194,9 @@ const MarkdownCell = memo(
     languageOptions,
     onLanguageChange,
     forceEditRequest = 0,
-    restoreFocusRequest = 0,
-    restoreFocusRole = "editor",
+    isActiveCell = false,
+    activeFocusRole = "editor",
+    isWindowFocused = false,
   }: MarkdownCellProps) => {
     // Subscribe to cell data changes using useSyncExternalStore for tearing-safe reads
     const cell = useSyncExternalStore(
@@ -208,19 +212,19 @@ const MarkdownCell = memo(
     // Start in rendered mode unless the cell is empty (new cell)
     const [rendered, setRendered] = useState(() => {
       const value = cell?.value ?? "";
+      if (value.trim() && isActiveCell && activeFocusRole === "editor") {
+        return false;
+      }
       return value.trim().length > 0;
     });
     const renderedRef = useRef<HTMLDivElement | null>(null);
-    const [editorFocusRequest, setEditorFocusRequest] = useState(0);
-    const [pendingRestoreTarget, setPendingRestoreTarget] = useState<
-      "editor" | "rendered" | null
-    >(null);
+    const previousShouldOwnFocusRef = useRef(false);
+    const [editorFocusIntent, setEditorFocusIntent] = useState(false);
+    const [renderedFocusIntent, setRenderedFocusIntent] = useState(false);
 
     // Get the current cell value
     const value = cell?.value ?? "";
-    const requestEditorFocus = useCallback(() => {
-      setEditorFocusRequest((request) => request + 1);
-    }, []);
+    const shouldOwnFocus = isActiveCell && isWindowFocused;
 
     // Enforce invariant: empty cells must be in edit mode.
     // This handles external changes (undo/redo, sync) that clear the value.
@@ -233,46 +237,53 @@ const MarkdownCell = memo(
     useEffect(() => {
       if (forceEditRequest > 0) {
         setRendered(false);
-        requestEditorFocus();
+        setEditorFocusIntent(true);
       }
-    }, [forceEditRequest, requestEditorFocus]);
+    }, [forceEditRequest]);
 
     useEffect(() => {
-      if (restoreFocusRequest <= 0) {
+      if (!editorFocusIntent || rendered) {
         return;
       }
-      if (restoreFocusRole === "editor" || !value.trim()) {
-        setRendered(false);
-        setPendingRestoreTarget("editor");
+      setEditorFocusIntent(false);
+    }, [editorFocusIntent, rendered]);
+
+    useEffect(() => {
+      if (!renderedFocusIntent || !rendered) {
         return;
       }
-      setRendered(true);
-      setPendingRestoreTarget("rendered");
-    }, [restoreFocusRequest, restoreFocusRole, value]);
+      renderedRef.current?.focus();
+      setRenderedFocusIntent(false);
+    }, [rendered, renderedFocusIntent]);
 
     useEffect(() => {
-      if (pendingRestoreTarget !== "editor" || rendered) {
+      const previouslyOwnedFocus = previousShouldOwnFocusRef.current;
+      previousShouldOwnFocusRef.current = shouldOwnFocus;
+      if (!shouldOwnFocus || previouslyOwnedFocus) {
         return;
       }
-      requestEditorFocus();
-      setPendingRestoreTarget(null);
-    }, [pendingRestoreTarget, rendered, requestEditorFocus]);
+      if (activeFocusRole === "editor" || !value.trim()) {
+        setRendered(false);
+        return;
+      }
+      setRendered(true);
+      renderedRef.current?.focus();
+    }, [activeFocusRole, shouldOwnFocus, value]);
 
     useEffect(() => {
-      if (pendingRestoreTarget !== "rendered" || !rendered) {
+      if (!shouldOwnFocus || activeFocusRole !== "rendered" || !rendered) {
         return;
       }
       renderedRef.current?.focus();
-      setPendingRestoreTarget(null);
-    }, [pendingRestoreTarget, rendered]);
+    }, [activeFocusRole, rendered, shouldOwnFocus]);
 
     /**
      * Handle switching to edit mode when user double-clicks rendered content.
      */
     const handleDoubleClick = useCallback(() => {
       setRendered(false);
-      requestEditorFocus();
-    }, [requestEditorFocus]);
+      setEditorFocusIntent(true);
+    }, []);
 
     /**
      * Handle keyboard activation on the rendered container for accessibility.
@@ -289,10 +300,10 @@ const MarkdownCell = memo(
         if (event.key === "Enter" || event.key === " ") {
           event.preventDefault();
           setRendered(false);
-          requestEditorFocus();
+          setEditorFocusIntent(true);
         }
       },
-      [requestEditorFocus]
+      []
     );
 
     /**
@@ -325,6 +336,7 @@ const MarkdownCell = memo(
           // Only render if there's content; empty cells stay in edit mode
           if (value.trim()) {
             setRendered(true);
+            setRenderedFocusIntent(true);
           }
         }
       },
@@ -353,6 +365,7 @@ const MarkdownCell = memo(
     const handleRun = useCallback(() => {
       if (value.trim()) {
         setRendered(true);
+        setRenderedFocusIntent(true);
       }
     }, [value]);
 
@@ -420,7 +433,11 @@ const MarkdownCell = memo(
               language="markdown"
               fontSize={fontSettings.fontSize}
               fontFamily={fontSettings.fontFamily}
-              focusRequest={editorFocusRequest}
+              shouldFocus={
+                !rendered &&
+                ((shouldOwnFocus && activeFocusRole === "editor") ||
+                  editorFocusIntent)
+              }
               onChange={handleEditorChange}
               onEnter={handleRun}
             />
@@ -455,8 +472,9 @@ const MarkdownCell = memo(
       prevProps.selectedLanguage === nextProps.selectedLanguage &&
       prevProps.languageSelectId === nextProps.languageSelectId &&
       prevProps.forceEditRequest === nextProps.forceEditRequest &&
-      prevProps.restoreFocusRequest === nextProps.restoreFocusRequest &&
-      prevProps.restoreFocusRole === nextProps.restoreFocusRole
+      prevProps.isActiveCell === nextProps.isActiveCell &&
+      prevProps.activeFocusRole === nextProps.activeFocusRole &&
+      prevProps.isWindowFocused === nextProps.isWindowFocused
     );
   }
 );
diff --git a/app/src/lib/notebookActiveCellState.test.ts b/app/src/lib/notebookActiveCellState.test.ts
new file mode 100644
index 00000000..01482e2e
--- /dev/null
+++ b/app/src/lib/notebookActiveCellState.test.ts
@@ -0,0 +1,77 @@
+// @vitest-environment jsdom
+import { afterEach, describe, expect, it } from "vitest";
+
+import {
+  createNotebookActiveCellState,
+  loadNotebookActiveCellMap,
+  persistNotebookActiveCellMap,
+} from "./notebookActiveCellState";
+
+const STORAGE_KEY = "runme/notebook-active-cells";
+
+describe("notebookActiveCellState", () => {
+  afterEach(() => {
+    localStorage.removeItem(STORAGE_KEY);
+  });
+
+  it("normalizes persisted notebook active-cell state", () => {
+    localStorage.setItem(
+      STORAGE_KEY,
+      JSON.stringify({
+        " local://file/demo.json ": {
+          refId: " cell-a ",
+          focusRole: "rendered",
+          updatedAt: "2026-05-14T00:00:00.000Z",
+        },
+        "": {
+          refId: "ignored",
+          focusRole: "editor",
+          updatedAt: "",
+        },
+        "local://file/bad.json": {
+          refId: "",
+          focusRole: "rendered",
+          updatedAt: "",
+        },
+      }),
+    );
+
+    expect(loadNotebookActiveCellMap()).toEqual({
+      "local://file/demo.json": {
+        refId: "cell-a",
+        focusRole: "rendered",
+        updatedAt: "2026-05-14T00:00:00.000Z",
+      },
+    });
+  });
+
+  it("persists only valid entries", () => {
+    persistNotebookActiveCellMap({
+      " local://file/demo.json ": {
+        refId: " cell-a ",
+        focusRole: "editor",
+        updatedAt: "2026-05-14T00:00:00.000Z",
+      },
+      "": {
+        refId: "",
+        focusRole: "rendered",
+        updatedAt: "",
+      },
+    });
+
+    expect(JSON.parse(localStorage.getItem(STORAGE_KEY) ?? "{}")).toEqual({
+      "local://file/demo.json": {
+        refId: "cell-a",
+        focusRole: "editor",
+        updatedAt: "2026-05-14T00:00:00.000Z",
+      },
+    });
+  });
+
+  it("creates timestamped active-cell entries", () => {
+    const snapshot = createNotebookActiveCellState("cell-a", "rendered");
+    expect(snapshot?.refId).toBe("cell-a");
+    expect(snapshot?.focusRole).toBe("rendered");
+    expect(snapshot?.updatedAt).toMatch(/^20\d\d-/);
+  });
+});
diff --git a/app/src/lib/notebookActiveCellState.ts b/app/src/lib/notebookActiveCellState.ts
new file mode 100644
index 00000000..b95f28da
--- /dev/null
+++ b/app/src/lib/notebookActiveCellState.ts
@@ -0,0 +1,106 @@
+export type CellFocusRole = "editor" | "rendered";
+
+export type NotebookActiveCellState = {
+  refId: string;
+  focusRole: CellFocusRole;
+  updatedAt: string;
+};
+
+export type NotebookActiveCellMap = Record<string, NotebookActiveCellState>;
+
+const STORAGE_KEY = "runme/notebook-active-cells";
+
+function normalizeString(value: unknown): string {
+  return typeof value === "string" ? value.trim() : "";
+}
+
+function normalizeFocusRole(value: unknown): CellFocusRole {
+  return value === "rendered" ? "rendered" : "editor";
+}
+
+function normalizeEntry(value: unknown): NotebookActiveCellState | null {
+  if (!value || typeof value !== "object") {
+    return null;
+  }
+  const candidate = value as Partial<NotebookActiveCellState>;
+  const refId = normalizeString(candidate.refId);
+  if (!refId) {
+    return null;
+  }
+  return {
+    refId,
+    focusRole: normalizeFocusRole(candidate.focusRole),
+    updatedAt: normalizeString(candidate.updatedAt),
+  };
+}
+
+export function loadNotebookActiveCellMap(): NotebookActiveCellMap {
+  if (typeof window === "undefined" || !window.localStorage) {
+    return {};
+  }
+  try {
+    const raw = window.localStorage.getItem(STORAGE_KEY);
+    if (!raw) {
+      return {};
+    }
+    const parsed = JSON.parse(raw);
+    if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) {
+      return {};
+    }
+    const next: NotebookActiveCellMap = {};
+    for (const [docUri, entry] of Object.entries(parsed)) {
+      const normalizedDocUri = normalizeString(docUri);
+      if (!normalizedDocUri) {
+        continue;
+      }
+      const normalizedEntry = normalizeEntry(entry);
+      if (!normalizedEntry) {
+        continue;
+      }
+      next[normalizedDocUri] = normalizedEntry;
+    }
+    return next;
+  } catch {
+    return {};
+  }
+}
+
+export function persistNotebookActiveCellMap(
+  snapshot: NotebookActiveCellMap,
+): void {
+  if (typeof window === "undefined" || !window.localStorage) {
+    return;
+  }
+  try {
+    const next: NotebookActiveCellMap = {};
+    for (const [docUri, entry] of Object.entries(snapshot)) {
+      const normalizedDocUri = normalizeString(docUri);
+      if (!normalizedDocUri) {
+        continue;
+      }
+      const normalizedEntry = normalizeEntry(entry);
+      if (!normalizedEntry) {
+        continue;
+      }
+      next[normalizedDocUri] = normalizedEntry;
+    }
+    window.localStorage.setItem(STORAGE_KEY, JSON.stringify(next));
+  } catch {
+    // Ignore localStorage write failures; focus restore should degrade safely.
+  }
+}
+
+export function createNotebookActiveCellState(
+  refId: string,
+  focusRole: CellFocusRole,
+): NotebookActiveCellState | null {
+  const normalizedRefId = normalizeString(refId);
+  if (!normalizedRefId) {
+    return null;
+  }
+  return {
+    refId: normalizedRefId,
+    focusRole: normalizeFocusRole(focusRole),
+    updatedAt: new Date().toISOString(),
+  };
+}
diff --git a/app/test/browser/run-cuj-scenarios.ts b/app/test/browser/run-cuj-scenarios.ts
index 2952d68e..0ee3a0ad 100644
--- a/app/test/browser/run-cuj-scenarios.ts
+++ b/app/test/browser/run-cuj-scenarios.ts
@@ -87,6 +87,7 @@ const SCENARIO_DRIVERS = [
   join(SCRIPT_DIR, "test-scenario-ai.ts"),
   join(SCRIPT_DIR, "test-scenario-ai-codex.ts"),
   join(SCRIPT_DIR, "test-scenario-chatkit-thread-persistence.ts"),
+  join(SCRIPT_DIR, "test-scenario-notebook-focus-persistence.ts"),
 ];
 
 type BackendCommandConfig = {
diff --git a/app/test/browser/test-scenario-notebook-focus-persistence.ts b/app/test/browser/test-scenario-notebook-focus-persistence.ts
new file mode 100644
index 00000000..f572b420
--- /dev/null
+++ b/app/test/browser/test-scenario-notebook-focus-persistence.ts
@@ -0,0 +1,278 @@
+import { spawnSync } from "node:child_process";
+import { mkdirSync, rmSync, writeFileSync } from "node:fs";
+import { dirname, join } from "node:path";
+import { fileURLToPath } from "node:url";
+
+const FRONTEND_URL = process.env.CUJ_FRONTEND_URL ?? "http://localhost:5173";
+const NOTEBOOK_NAME = "scenario-notebook-focus.json";
+const NOTEBOOK_URI = `local://file/${NOTEBOOK_NAME}`;
+const MARKDOWN_CELL_ID = "md_focus_cell";
+const STORAGE_KEY = "runme/notebook-active-cells";
+
+const CURRENT_FILE_DIR = dirname(fileURLToPath(import.meta.url));
+const SCRIPT_DIR =
+  CURRENT_FILE_DIR.endsWith("/.generated") || CURRENT_FILE_DIR.endsWith("\\.generated")
+    ? dirname(CURRENT_FILE_DIR)
+    : CURRENT_FILE_DIR;
+const OUTPUT_DIR = join(SCRIPT_DIR, "test-output");
+const MOVIE_PATH = join(OUTPUT_DIR, "scenario-notebook-focus-persistence-walkthrough.webm");
+const AGENT_BROWSER_SESSION = process.env.AGENT_BROWSER_SESSION?.trim() ?? "";
+const AGENT_BROWSER_PROFILE = process.env.AGENT_BROWSER_PROFILE?.trim() ?? "";
+const AGENT_BROWSER_HEADED = (process.env.AGENT_BROWSER_HEADED ?? "false")
+  .trim()
+  .toLowerCase() === "true";
+const AGENT_BROWSER_KEEP_OPEN = (process.env.AGENT_BROWSER_KEEP_OPEN ?? "false")
+  .trim()
+  .toLowerCase() === "true";
+
+let passCount = 0;
+let failCount = 0;
+let totalCount = 0;
+
+function run(
+  command: string,
+  options?: { timeoutMs?: number; throwOnAgentBrowserTimeout?: boolean },
+): { status: number; stdout: string; stderr: string } {
+  const effectiveCommand = withAgentBrowserOptions(command);
+  const timeoutMs = options?.timeoutMs ?? Number(process.env.CUJ_SCENARIO_CMD_TIMEOUT_MS ?? "15000");
+  const result = spawnSync(effectiveCommand, {
+    shell: true,
+    encoding: "utf-8",
+    timeout: timeoutMs,
+    killSignal: "SIGKILL",
+  });
+  const errorCode =
+    typeof result.error === "object" && result.error !== null && "code" in result.error
+      ? String((result.error as { code?: string }).code ?? "")
+      : "";
+  const timedOut = errorCode === "ETIMEDOUT";
+  const timeoutHint = timedOut
+    ? `\n[scenario-timeout] command timed out after ${timeoutMs}ms: ${effectiveCommand}\n`
+    : "";
+  const shouldThrowOnTimeout = options?.throwOnAgentBrowserTimeout ?? true;
+  if (
+    timedOut &&
+    shouldThrowOnTimeout &&
+    effectiveCommand.trim().startsWith("agent-browser ")
+  ) {
+    throw new Error(timeoutHint.trim());
+  }
+  return {
+    status: result.status ?? (timedOut ? 124 : 1),
+    stdout: result.stdout ?? "",
+    stderr: `${result.stderr ?? ""}${timeoutHint}`,
+  };
+}
+
+function shellQuote(value: string): string {
+  return `'${value.replace(/'/g, `'\"'\"'`)}'`;
+}
+
+function withAgentBrowserOptions(command: string): string {
+  const trimmed = command.trimStart();
+  if (!trimmed.startsWith("agent-browser ")) {
+    return command;
+  }
+  const leadingWhitespace = command.slice(0, command.length - trimmed.length);
+  const subcommand = trimmed.slice("agent-browser ".length);
+  const args: string[] = [];
+  if (AGENT_BROWSER_SESSION) {
+    args.push("--session", shellQuote(AGENT_BROWSER_SESSION));
+  }
+  if (AGENT_BROWSER_PROFILE) {
+    args.push("--profile", shellQuote(AGENT_BROWSER_PROFILE));
+  }
+  if (AGENT_BROWSER_HEADED) {
+    args.push("--headed");
+  }
+  const prefix = ["agent-browser", ...args].join(" ");
+  return `${leadingWhitespace}${prefix} ${subcommand}`;
+}
+
+function runOrThrow(command: string): string {
+  const result = run(command);
+  if (result.status !== 0) {
+    throw new Error(`Command failed: ${command}\n${result.stderr}`);
+  }
+  return result.stdout;
+}
+
+function pass(message: string): void {
+  totalCount += 1;
+  passCount += 1;
+  console.log(`[PASS] ${message}`);
+}
+
+function fail(message: string): void {
+  totalCount += 1;
+  failCount += 1;
+  console.log(`[FAIL] ${message}`);
+}
+
+function writeArtifact(name: string, content: string): void {
+  writeFileSync(join(OUTPUT_DIR, name), content, "utf-8");
+}
+
+function escapeDoubleQuotes(value: string): string {
+  return value.replace(/\\/g, "\\\\").replace(/"/g, '\\"');
+}
+
+mkdirSync(OUTPUT_DIR, { recursive: true });
+rmSync(MOVIE_PATH, { force: true });
+for (const file of [
+  "scenario-notebook-focus-persistence-01-seed.json",
+  "scenario-notebook-focus-persistence-02-storage-after-edit.json",
+  "scenario-notebook-focus-persistence-03-after-reload-snapshot.txt",
+  "scenario-notebook-focus-persistence-04-after-reload-state.json",
+  "scenario-notebook-focus-persistence-05-after-reload.png",
+]) {
+  rmSync(join(OUTPUT_DIR, file), { force: true });
+}
+
+if (run("command -v agent-browser").status !== 0) {
+  console.error("ERROR: agent-browser is required on PATH");
+  process.exit(2);
+}
+
+if (run(`curl -sf ${FRONTEND_URL}`).status !== 0) {
+  console.error(`ERROR: frontend is not running at ${FRONTEND_URL}`);
+  process.exit(1);
+}
+
+let startedRecording = false;
+
+try {
+  runOrThrow(`agent-browser open ${FRONTEND_URL}`);
+  runOrThrow(`agent-browser record restart ${MOVIE_PATH}`);
+  startedRecording = true;
+  run("agent-browser wait 2200");
+
+  const seedResult = run(
+    `agent-browser eval "(async () => {
+      const notebook = {
+        metadata: {},
+        cells: [
+          {
+            refId: '${MARKDOWN_CELL_ID}',
+            kind: 1,
+            languageId: 'markdown',
+            value: '# Focus persistence\\n\\nEdit me after reload.',
+            metadata: {},
+            outputs: []
+          }
+        ]
+      };
+      const ln = window.app?.localNotebooks;
+      if (!ln) return JSON.stringify({ status: 'missing-local-notebooks' });
+      await ln.files.put({
+        id: '${NOTEBOOK_URI}',
+        uri: '${NOTEBOOK_URI}',
+        name: '${NOTEBOOK_NAME}',
+        doc: JSON.stringify(notebook),
+        updatedAt: new Date().toISOString(),
+        parent: 'local://folder/local',
+        lastSynced: '',
+        remoteId: '',
+        lastRemoteChecksum: ''
+      });
+      localStorage.setItem('runme/openNotebooks', JSON.stringify([
+        { uri: '${NOTEBOOK_URI}', name: '${NOTEBOOK_NAME}', type: 'file', children: [] }
+      ]));
+      localStorage.setItem('runme/currentDoc', '${NOTEBOOK_URI}');
+      localStorage.removeItem('${STORAGE_KEY}');
+      return JSON.stringify({ status: 'ok' });
+    })()"`,
+  ).stdout.trim();
+  writeArtifact("scenario-notebook-focus-persistence-01-seed.json", seedResult);
+  if (seedResult.includes('"status":"ok"')) {
+    pass("Seeded local markdown notebook");
+  } else {
+    fail("Failed to seed local markdown notebook");
+  }
+
+  run("agent-browser reload");
+  run("agent-browser wait 2400");
+
+  const openEditorResult = run(
+    `agent-browser eval "(async () => {
+      const rendered = document.querySelector('#markdown-rendered-${MARKDOWN_CELL_ID}');
+      if (!(rendered instanceof HTMLElement)) return 'missing-rendered-markdown';
+      rendered.dispatchEvent(new MouseEvent('dblclick', { bubbles: true }));
+      rendered.focus();
+      return 'ok';
+    })()"`,
+  ).stdout.trim();
+  if (openEditorResult.includes("ok")) {
+    pass("Opened markdown cell in editor mode");
+  } else {
+    fail("Failed to open markdown cell in editor mode");
+  }
+
+  run("agent-browser wait 600");
+
+  const storageAfterEdit = run(
+    `agent-browser eval "${escapeDoubleQuotes(`(() => {
+      const raw = localStorage.getItem('${STORAGE_KEY}');
+      return raw || '';
+    })()`)}"`,
+  ).stdout.trim();
+  writeArtifact(
+    "scenario-notebook-focus-persistence-02-storage-after-edit.json",
+    storageAfterEdit,
+  );
+  if (storageAfterEdit.includes(MARKDOWN_CELL_ID) && storageAfterEdit.includes('"focusRole":"editor"')) {
+    pass("Persisted active markdown editor state");
+  } else {
+    fail("Did not persist active markdown editor state");
+  }
+
+  run("agent-browser reload");
+  run("agent-browser wait 2400");
+
+  const afterReloadSnapshot = run("agent-browser snapshot -i").stdout;
+  writeArtifact(
+    "scenario-notebook-focus-persistence-03-after-reload-snapshot.txt",
+    afterReloadSnapshot,
+  );
+
+  const afterReloadState = run(
+    `agent-browser eval "${escapeDoubleQuotes(`(() => {
+      const editor = document.querySelector('#markdown-editor-${MARKDOWN_CELL_ID}');
+      const rendered = document.querySelector('#markdown-rendered-${MARKDOWN_CELL_ID}');
+      const active = document.activeElement;
+      return JSON.stringify({
+        hasEditor: Boolean(editor),
+        hasRendered: Boolean(rendered),
+        activeElementTag: active?.tagName ?? '',
+        activeElementId: active?.id ?? '',
+      });
+    })()`)}"`,
+  ).stdout.trim();
+  writeArtifact(
+    "scenario-notebook-focus-persistence-04-after-reload-state.json",
+    afterReloadState,
+  );
+  run(`agent-browser screenshot ${join(OUTPUT_DIR, "scenario-notebook-focus-persistence-05-after-reload.png")}`);
+
+  if (afterReloadState.includes('"hasEditor":true') && !afterReloadState.includes('"hasRendered":true')) {
+    pass("Reload restored markdown cell in editor mode");
+  } else {
+    fail("Reload did not restore markdown cell in editor mode");
+  }
+} catch (error) {
+  fail(`Scenario execution error: ${String(error)}`);
+} finally {
+  if (startedRecording) {
+    run("agent-browser record stop", {
+      timeoutMs: 30000,
+      throwOnAgentBrowserTimeout: false,
+    });
+  }
+  if (!AGENT_BROWSER_KEEP_OPEN) {
+    run("agent-browser close");
+  }
+}
+
+console.log(`Movie: ${MOVIE_PATH}`);
+console.log(`Assertions: ${totalCount}, Passed: ${passCount}, Failed: ${failCount}`);
+process.exit(failCount > 0 ? 1 : 0);
diff --git a/docs-dev/design/20260514_notebook_focus_persistence.md b/docs-dev/design/20260514_notebook_focus_persistence.md
new file mode 100644
index 00000000..8e9d581f
--- /dev/null
+++ b/docs-dev/design/20260514_notebook_focus_persistence.md
@@ -0,0 +1,149 @@
+# Notebook Focus Persistence
+
+## Scope
+
+This document proposes a simpler model for restoring notebook cell focus across:
+
+- switching away from the browser tab and back
+- switching between notebook tabs in the app
+- refreshing the page
+
+The target behavior is:
+
+1. Each open notebook remembers its last active cell.
+2. Markdown cells remember whether the last active surface was the editor or the
+   rendered view.
+3. Returning to the browser tab restores focus to the visible notebook's last
+   active cell.
+4. Refreshing the page restores the same cell and reopens markdown in edit mode
+   when the editor was last active.
+
+## Current Problem
+
+The current implementation models focus restore as a React event counter.
+
+That creates two problems:
+
+- it mixes durable UI state with transient focus events
+- it lets unrelated re-renders replay restore logic
+
+The reviewer concern is valid. If restore is encoded as `restoreFocusRequest`,
+the code must answer "has this request already been consumed?" in every effect
+that reads it. That makes markdown edit/render behavior fragile.
+
+## Decision
+
+We will persist per-notebook active-cell UI state in `localStorage`.
+
+We will not persist focus-restore event counters.
+
+We will track window/tab focus separately as transient React state and use it
+only to decide when to apply the already-persisted active-cell state.
+
+## Why `localStorage`
+
+`localStorage` is the right default here because:
+
+- the state is tiny
+- it is JSON-shaped
+- it must be readable during initial render
+- it should survive refresh
+- it does not need IndexedDB queries or transactions
+
+IndexedDB remains the right place for notebook content and mirrored storage.
+This feature is UI state, not notebook data.
+
+## Stored Shape
+
+Suggested payload:
+
+```ts
+type CellFocusRole = "editor" | "rendered";
+
+interface NotebookActiveCellState {
+  refId: string;
+  focusRole: CellFocusRole;
+  updatedAt: string;
+}
+
+type NotebookActiveCellMap = Record<string, NotebookActiveCellState>;
+```
+
+Storage key:
+
+```ts
+const STORAGE_KEY = "runme/notebook-active-cells";
+```
+
+The map is keyed by notebook URI.
+
+## Behavior
+
+### On cell focus
+
+When focus moves within a notebook cell:
+
+- record the notebook URI
+- record the cell ref id
+- record whether focus is in the editor or rendered markdown surface
+- persist the map to `localStorage`
+
+### On browser focus return
+
+When the browser tab becomes focused and visible:
+
+- look at the visible notebook tab, not `currentDoc`
+- read that notebook's stored active-cell state
+- if the stored cell still exists, restore focus to it
+- if the stored focus role is `editor` for a markdown cell, reopen that cell in
+  edit mode before focusing it
+- if the visible tab is a non-notebook synthetic tab, do nothing
+
+### On page refresh
+
+During initial render of a notebook tab:
+
+- read stored active-cell state for that notebook
+- if the active cell is a markdown cell and the stored role is `editor`, mount
+  that cell in edit mode
+- if the page is already focused, place focus back into that cell
+
+### On stale state
+
+If the stored cell no longer exists:
+
+- ignore the stored entry
+- fall back to normal rendering
+
+We do not need aggressive cleanup when notebooks close. Keeping the last active
+cell for a notebook is useful if the user reopens it later.
+
+## Implementation Plan
+
+1. Add a small helper module for reading and writing the persisted map.
+2. Load persisted active-cell state in `Actions`.
+3. Replace restore counters with:
+   - durable active-cell state from storage
+   - transient `window focused` state
+4. Change `MarkdownCell` to derive initial edit mode from persisted active-cell
+   state instead of replaying a request counter.
+5. Add focused unit tests for:
+   - storage normalization
+   - markdown initial edit-mode restore
+   - focus restore on browser refocus without replaying on every keystroke
+6. Verify behavior in a browser scenario and record a video artifact.
+
+## Tradeoffs
+
+- `localStorage` is synchronous. That is acceptable because the payload is tiny
+  and the feature is startup-sensitive.
+- The stored state is browser-local. That is correct for per-browser UI focus.
+- We still keep a small amount of React state for "is the browser focused?".
+  That state is transient and does not encode restore events, which keeps the
+  model simpler than a counter.
+
+## Recommendation
+
+Persist `{ docUri -> activeCellId, focusRole }` in `localStorage` and treat
+window focus as a transient signal that reapplies that durable state to the
+visible notebook only.

From 38d8f010a48a34af661fa4cc934785cc41df7555 Mon Sep 17 00:00:00 2001
From: Jeremy lewi <jeremy@lewi.us>
Date: Thu, 14 May 2026 16:42:00 -0700
Subject: [PATCH 4/6] Fix notebook focus validation harness

Signed-off-by: Jeremy lewi <jeremy@lewi.us>
---
 .../components/Actions/MarkdownCell.test.tsx  | 16 +++-
 ...est-scenario-notebook-focus-persistence.ts | 92 +++++++++++++++----
 2 files changed, 87 insertions(+), 21 deletions(-)

diff --git a/app/src/components/Actions/MarkdownCell.test.tsx b/app/src/components/Actions/MarkdownCell.test.tsx
index 22476661..0e1df9ae 100644
--- a/app/src/components/Actions/MarkdownCell.test.tsx
+++ b/app/src/components/Actions/MarkdownCell.test.tsx
@@ -133,7 +133,9 @@ describe("MarkdownCell", () => {
       await Promise.resolve();
     });
 
-    expect(screen.getByTestId("mock-markdown-editor-input")).toHaveFocus();
+    expect(document.activeElement).toBe(
+      screen.getByTestId("mock-markdown-editor-input"),
+    );
   });
 
   it("does not replay restore behavior on each editor change", async () => {
@@ -183,7 +185,9 @@ describe("MarkdownCell", () => {
       await Promise.resolve();
     });
 
-    expect(screen.getByTestId("mock-markdown-editor-input")).toHaveFocus();
+    expect(document.activeElement).toBe(
+      screen.getByTestId("mock-markdown-editor-input"),
+    );
   });
 
   it("focuses rendered markdown when leaving editor mode with escape", async () => {
@@ -211,7 +215,9 @@ describe("MarkdownCell", () => {
       await Promise.resolve();
     });
 
-    expect(screen.getByTestId("markdown-rendered")).toHaveFocus();
+    expect(document.activeElement).toBe(
+      screen.getByTestId("markdown-rendered"),
+    );
   });
 
   it("focuses rendered markdown when the window regains focus", async () => {
@@ -247,6 +253,8 @@ describe("MarkdownCell", () => {
       await Promise.resolve();
     });
 
-    expect(screen.getByTestId("markdown-rendered")).toHaveFocus();
+    expect(document.activeElement).toBe(
+      screen.getByTestId("markdown-rendered"),
+    );
   });
 });
diff --git a/app/test/browser/test-scenario-notebook-focus-persistence.ts b/app/test/browser/test-scenario-notebook-focus-persistence.ts
index f572b420..ac6e8c86 100644
--- a/app/test/browser/test-scenario-notebook-focus-persistence.ts
+++ b/app/test/browser/test-scenario-notebook-focus-persistence.ts
@@ -8,6 +8,7 @@ const NOTEBOOK_NAME = "scenario-notebook-focus.json";
 const NOTEBOOK_URI = `local://file/${NOTEBOOK_NAME}`;
 const MARKDOWN_CELL_ID = "md_focus_cell";
 const STORAGE_KEY = "runme/notebook-active-cells";
+const LOCAL_DB_NAME = "runme-local-notebooks";
 
 const CURRENT_FILE_DIR = dirname(fileURLToPath(import.meta.url));
 const SCRIPT_DIR =
@@ -117,6 +118,19 @@ function escapeDoubleQuotes(value: string): string {
   return value.replace(/\\/g, "\\\\").replace(/"/g, '\\"');
 }
 
+function decodeAgentBrowserEvalOutput(output: string): string {
+  const trimmed = output.trim();
+  if (!trimmed) {
+    return "";
+  }
+  try {
+    const parsed = JSON.parse(trimmed);
+    return typeof parsed === "string" ? parsed : JSON.stringify(parsed);
+  } catch {
+    return trimmed;
+  }
+}
+
 mkdirSync(OUTPUT_DIR, { recursive: true });
 rmSync(MOVIE_PATH, { force: true });
 for (const file of [
@@ -147,8 +161,31 @@ try {
   startedRecording = true;
   run("agent-browser wait 2200");
 
-  const seedResult = run(
+  const seedResult = decodeAgentBrowserEvalOutput(
+    run(
     `agent-browser eval "(async () => {
+      const openLocalDb = () => new Promise((resolve, reject) => {
+        const request = indexedDB.open('${LOCAL_DB_NAME}');
+        request.onupgradeneeded = () => {
+          const db = request.result;
+          if (!db.objectStoreNames.contains('files')) {
+            db.createObjectStore('files', { keyPath: 'id' });
+          }
+          if (!db.objectStoreNames.contains('folders')) {
+            db.createObjectStore('folders', { keyPath: 'id' });
+          }
+        };
+        request.onsuccess = () => resolve(request.result);
+        request.onerror = () => reject(request.error);
+      });
+      const runTransaction = (db, storeName, mode, callback) => new Promise((resolve, reject) => {
+        const tx = db.transaction(storeName, mode);
+        const store = tx.objectStore(storeName);
+        const request = callback(store);
+        tx.oncomplete = () => resolve(request?.result);
+        tx.onerror = () => reject(tx.error ?? request?.error);
+        tx.onabort = () => reject(tx.error ?? request?.error);
+      });
       const notebook = {
         metadata: {},
         cells: [
@@ -162,27 +199,34 @@ try {
           }
         ]
       };
-      const ln = window.app?.localNotebooks;
-      if (!ln) return JSON.stringify({ status: 'missing-local-notebooks' });
-      await ln.files.put({
+      const notebookDoc = JSON.stringify(notebook);
+      const db = await openLocalDb();
+      await runTransaction(db, 'folders', 'readwrite', (store) => store.put({
+        id: 'local://folder/local',
+        name: 'Local Notebooks',
+        remoteId: '',
+        children: ['${NOTEBOOK_URI}'],
+        lastSynced: ''
+      }));
+      await runTransaction(db, 'files', 'readwrite', (store) => store.put({
         id: '${NOTEBOOK_URI}',
-        uri: '${NOTEBOOK_URI}',
         name: '${NOTEBOOK_NAME}',
-        doc: JSON.stringify(notebook),
-        updatedAt: new Date().toISOString(),
-        parent: 'local://folder/local',
+        remoteId: '${NOTEBOOK_URI}',
         lastSynced: '',
-        remoteId: '',
-        lastRemoteChecksum: ''
-      });
+        lastRemoteChecksum: '',
+        doc: notebookDoc,
+        md5Checksum: ''
+      }));
+      db.close();
       localStorage.setItem('runme/openNotebooks', JSON.stringify([
-        { uri: '${NOTEBOOK_URI}', name: '${NOTEBOOK_NAME}', type: 'file', children: [] }
+        { uri: '${NOTEBOOK_URI}', name: '${NOTEBOOK_NAME}', type: 'file', children: [], parents: ['local://folder/local'] }
       ]));
       localStorage.setItem('runme/currentDoc', '${NOTEBOOK_URI}');
       localStorage.removeItem('${STORAGE_KEY}');
       return JSON.stringify({ status: 'ok' });
     })()"`,
-  ).stdout.trim();
+    ).stdout,
+  );
   writeArtifact("scenario-notebook-focus-persistence-01-seed.json", seedResult);
   if (seedResult.includes('"status":"ok"')) {
     pass("Seeded local markdown notebook");
@@ -210,12 +254,24 @@ try {
 
   run("agent-browser wait 600");
 
-  const storageAfterEdit = run(
+  run(
+    `agent-browser eval "(async () => {
+      const textarea = document.querySelector('#markdown-editor-${MARKDOWN_CELL_ID} textarea');
+      if (!(textarea instanceof HTMLElement)) return 'missing-editor-textarea';
+      textarea.focus();
+      return document.activeElement === textarea ? 'ok' : 'focus-missed';
+    })()"`,
+  );
+  run("agent-browser wait 400");
+
+  const storageAfterEdit = decodeAgentBrowserEvalOutput(
+    run(
     `agent-browser eval "${escapeDoubleQuotes(`(() => {
       const raw = localStorage.getItem('${STORAGE_KEY}');
       return raw || '';
     })()`)}"`,
-  ).stdout.trim();
+    ).stdout,
+  );
   writeArtifact(
     "scenario-notebook-focus-persistence-02-storage-after-edit.json",
     storageAfterEdit,
@@ -235,7 +291,8 @@ try {
     afterReloadSnapshot,
   );
 
-  const afterReloadState = run(
+  const afterReloadState = decodeAgentBrowserEvalOutput(
+    run(
     `agent-browser eval "${escapeDoubleQuotes(`(() => {
       const editor = document.querySelector('#markdown-editor-${MARKDOWN_CELL_ID}');
       const rendered = document.querySelector('#markdown-rendered-${MARKDOWN_CELL_ID}');
@@ -247,7 +304,8 @@ try {
         activeElementId: active?.id ?? '',
       });
     })()`)}"`,
-  ).stdout.trim();
+    ).stdout,
+  );
   writeArtifact(
     "scenario-notebook-focus-persistence-04-after-reload-state.json",
     afterReloadState,

From 080481d993ad5c33435d469e0cfe8707d6a9950a Mon Sep 17 00:00:00 2001
From: Jeremy lewi <jeremy@lewi.us>
Date: Thu, 14 May 2026 16:49:33 -0700
Subject: [PATCH 5/6] Persist markdown editor focus state

Signed-off-by: Jeremy lewi <jeremy@lewi.us>
---
 app/src/components/Actions/Actions.tsx        | 16 +++++++++++++++
 app/src/components/Actions/Editor.tsx         |  8 ++++++++
 app/src/components/Actions/MarkdownCell.tsx   | 20 ++++++++++++++-----
 ...est-scenario-notebook-focus-persistence.ts |  6 ++++--
 4 files changed, 43 insertions(+), 7 deletions(-)

diff --git a/app/src/components/Actions/Actions.tsx b/app/src/components/Actions/Actions.tsx
index 4bb66bab..e3f03de1 100644
--- a/app/src/components/Actions/Actions.tsx
+++ b/app/src/components/Actions/Actions.tsx
@@ -5,6 +5,7 @@ import {
   useEffect,
   useLayoutEffect,
   useMemo,
+  useRef,
   useState,
   useSyncExternalStore,
 } from "react";
@@ -659,6 +660,20 @@ export function Action({
     [cell?.refId, onFocusStateChange],
   );
 
+  const handleMarkdownFocusRoleChange = useCallback(
+    (focusRole: CellFocusRole) => {
+      if (!cell?.refId || !onFocusStateChange) {
+        return;
+      }
+      const nextState = createNotebookActiveCellState(cell.refId, focusRole);
+      if (!nextState) {
+        return;
+      }
+      onFocusStateChange(nextState);
+    },
+    [cell?.refId, onFocusStateChange],
+  );
+
   const handleRemoveCell = useCallback(() => {
     cellData.remove();
     setContextMenu(null);
@@ -1037,6 +1052,7 @@ export function Action({
               isActiveCell={isActiveCell}
               activeFocusRole={activeFocusRole}
               isWindowFocused={isWindowFocused}
+              onFocusRoleChange={handleMarkdownFocusRoleChange}
             />
             {/* Trash icon on the right, visible on hover */}
             <button
diff --git a/app/src/components/Actions/Editor.tsx b/app/src/components/Actions/Editor.tsx
index 3db01257..a312be60 100644
--- a/app/src/components/Actions/Editor.tsx
+++ b/app/src/components/Actions/Editor.tsx
@@ -18,6 +18,7 @@ const Editor = memo(
     ariaLabel,
     autoFocusWhenEmpty = true,
     shouldFocus = false,
+    onFocus,
     onChange,
     onEnter,
     onMount,
@@ -31,6 +32,7 @@ const Editor = memo(
     ariaLabel?: string;
     autoFocusWhenEmpty?: boolean;
     shouldFocus?: boolean;
+    onFocus?: () => void;
     onChange: (value: string) => void;
     onEnter: () => void;
     onMount?: (editor: any, monaco: any) => void;
@@ -48,6 +50,7 @@ const Editor = memo(
     // don't render an empty scroll track.
     const [isClamped, setIsClamped] = useState(false);
     const contentSizeListener = useRef<{ dispose: () => void } | null>(null);
+    const focusListener = useRef<{ dispose: () => void } | null>(null);
     const previousShouldFocusRef = useRef(false);
 
     // Keep the ref updated with the latest onEnter
@@ -149,6 +152,9 @@ const Editor = memo(
       contentSizeListener.current = editor.onDidContentSizeChange(() => {
         adjustHeight();
       });
+      focusListener.current = editor.onDidFocusEditorText(() => {
+        onFocus?.();
+      });
       onMount?.(editor, monaco);
     };
 
@@ -176,6 +182,8 @@ const Editor = memo(
       return () => {
         contentSizeListener.current?.dispose?.();
         contentSizeListener.current = null;
+        focusListener.current?.dispose?.();
+        focusListener.current = null;
       };
     }, []);
 
diff --git a/app/src/components/Actions/MarkdownCell.tsx b/app/src/components/Actions/MarkdownCell.tsx
index 07c10e24..95b4eb02 100644
--- a/app/src/components/Actions/MarkdownCell.tsx
+++ b/app/src/components/Actions/MarkdownCell.tsx
@@ -175,6 +175,8 @@ interface MarkdownCellProps {
   activeFocusRole?: CellFocusRole;
   /** Whether the visible browser tab is currently focused */
   isWindowFocused?: boolean;
+  /** Notify parent when markdown focus target changes explicitly */
+  onFocusRoleChange?: (focusRole: CellFocusRole) => void;
 }
 
 /**
@@ -197,6 +199,7 @@ const MarkdownCell = memo(
     isActiveCell = false,
     activeFocusRole = "editor",
     isWindowFocused = false,
+    onFocusRoleChange,
   }: MarkdownCellProps) => {
     // Subscribe to cell data changes using useSyncExternalStore for tearing-safe reads
     const cell = useSyncExternalStore(
@@ -283,7 +286,8 @@ const MarkdownCell = memo(
     const handleDoubleClick = useCallback(() => {
       setRendered(false);
       setEditorFocusIntent(true);
-    }, []);
+      onFocusRoleChange?.("editor");
+    }, [onFocusRoleChange]);
 
     /**
      * Handle keyboard activation on the rendered container for accessibility.
@@ -301,9 +305,10 @@ const MarkdownCell = memo(
           event.preventDefault();
           setRendered(false);
           setEditorFocusIntent(true);
+          onFocusRoleChange?.("editor");
         }
       },
-      []
+      [onFocusRoleChange]
     );
 
     /**
@@ -322,8 +327,9 @@ const MarkdownCell = memo(
           return;
         }
         setRendered(true);
+        onFocusRoleChange?.("rendered");
       },
-      [value]
+      [onFocusRoleChange, value]
     );
 
     /**
@@ -337,10 +343,11 @@ const MarkdownCell = memo(
           if (value.trim()) {
             setRendered(true);
             setRenderedFocusIntent(true);
+            onFocusRoleChange?.("rendered");
           }
         }
       },
-      [value]
+      [onFocusRoleChange, value]
     );
 
     /**
@@ -414,6 +421,7 @@ const MarkdownCell = memo(
             aria-label="Double-click or press Enter to edit markdown"
             data-testid="markdown-rendered"
             data-cell-focus-role="rendered"
+            onFocus={() => onFocusRoleChange?.("rendered")}
           >
             {renderedMarkdown}
           </div>
@@ -438,6 +446,7 @@ const MarkdownCell = memo(
                 ((shouldOwnFocus && activeFocusRole === "editor") ||
                   editorFocusIntent)
               }
+              onFocus={() => onFocusRoleChange?.("editor")}
               onChange={handleEditorChange}
               onEnter={handleRun}
             />
@@ -474,7 +483,8 @@ const MarkdownCell = memo(
       prevProps.forceEditRequest === nextProps.forceEditRequest &&
       prevProps.isActiveCell === nextProps.isActiveCell &&
       prevProps.activeFocusRole === nextProps.activeFocusRole &&
-      prevProps.isWindowFocused === nextProps.isWindowFocused
+      prevProps.isWindowFocused === nextProps.isWindowFocused &&
+      prevProps.onFocusRoleChange === nextProps.onFocusRoleChange
     );
   }
 );
diff --git a/app/test/browser/test-scenario-notebook-focus-persistence.ts b/app/test/browser/test-scenario-notebook-focus-persistence.ts
index ac6e8c86..06e2afc9 100644
--- a/app/test/browser/test-scenario-notebook-focus-persistence.ts
+++ b/app/test/browser/test-scenario-notebook-focus-persistence.ts
@@ -242,8 +242,10 @@ try {
       const rendered = document.querySelector('#markdown-rendered-${MARKDOWN_CELL_ID}');
       if (!(rendered instanceof HTMLElement)) return 'missing-rendered-markdown';
       rendered.dispatchEvent(new MouseEvent('dblclick', { bubbles: true }));
-      rendered.focus();
-      return 'ok';
+      await new Promise((resolve) => setTimeout(resolve, 250));
+      return document.querySelector('#markdown-editor-${MARKDOWN_CELL_ID}')
+        ? 'ok'
+        : 'missing-markdown-editor';
     })()"`,
   ).stdout.trim();
   if (openEditorResult.includes("ok")) {

From e6c5b23e9c76284217258aea16d37e198ec9db78 Mon Sep 17 00:00:00 2001
From: Jeremy lewi <jeremy@lewi.us>
Date: Thu, 14 May 2026 17:39:43 -0700
Subject: [PATCH 6/6] Ignore non-surface focus for active-cell restore

Signed-off-by: Jeremy lewi <jeremy@lewi.us>
---
 app/src/components/Actions/Actions.test.tsx | 25 +++++++++++++++++++++
 app/src/components/Actions/Actions.tsx      |  9 ++++++--
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/app/src/components/Actions/Actions.test.tsx b/app/src/components/Actions/Actions.test.tsx
index daf33fe6..6722f702 100644
--- a/app/src/components/Actions/Actions.test.tsx
+++ b/app/src/components/Actions/Actions.test.tsx
@@ -372,4 +372,29 @@ describe("Action component", () => {
     expect(runnerSelect).toBeNull();
     expect(kernelSelect).toBeTruthy();
   });
+
+  it("ignores focus on rendered markdown controls that are outside a focus-role surface", () => {
+    const cell = create(parser_pb.CellSchema, {
+      refId: "cell-md-rendered-controls",
+      kind: parser_pb.CellKind.MARKUP,
+      languageId: "markdown",
+      outputs: [],
+      metadata: {},
+      value: "hello",
+    });
+    const stub = new StubCellData(cell);
+    const onFocusStateChange = vi.fn();
+
+    render(
+      <Action
+        cellData={stub as unknown as CellData}
+        isFirst={false}
+        onFocusStateChange={onFocusStateChange}
+      />,
+    );
+
+    fireEvent.focus(screen.getByRole("button", { name: "Delete cell" }));
+
+    expect(onFocusStateChange).not.toHaveBeenCalled();
+  });
 });
diff --git a/app/src/components/Actions/Actions.tsx b/app/src/components/Actions/Actions.tsx
index e3f03de1..8ad0500f 100644
--- a/app/src/components/Actions/Actions.tsx
+++ b/app/src/components/Actions/Actions.tsx
@@ -646,9 +646,14 @@ export function Action({
       if (!(target instanceof HTMLElement)) {
         return;
       }
+      const focusRoleElement = target.closest<HTMLElement>(
+        "[data-cell-focus-role]",
+      );
+      if (!focusRoleElement) {
+        return;
+      }
       const focusRole =
-        target.closest<HTMLElement>("[data-cell-focus-role]")?.dataset
-          .cellFocusRole === "rendered"
+        focusRoleElement.dataset.cellFocusRole === "rendered"
           ? "rendered"
           : "editor";
       const nextState = createNotebookActiveCellState(cell.refId, focusRole);