[codex] Restore notebook cell focus

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();
+  });
 });

app/src/components/Actions/Actions.tsx

@@ -29,6 +29,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,
@@ -481,12 +489,20 @@ export function ActionOutputItems({ outputs }: { outputs: parser_pb.CellOutput[]
 
 export function Action({
   cellData,
-  docUri,
+  docUri = "",
   isFirst,
+  isActiveCell = false,
+  activeFocusRole = "editor",
+  isWindowFocused = false,
+  onFocusStateChange,
 }: {
   cellData: CellData;
-  docUri: string;
+  docUri?: string;
   isFirst: boolean;
+  isActiveCell?: boolean;
+  activeFocusRole?: CellFocusRole;
+  isWindowFocused?: boolean;
+  onFocusStateChange?: (state: NotebookActiveCellState) => void;
 }) {
   const { store } = useNotebookStore();
   const { listRunners, defaultRunnerName } = useRunners();
@@ -621,6 +637,48 @@ 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 focusRoleElement = target.closest<HTMLElement>(
+        "[data-cell-focus-role]",
+      );
+      if (!focusRoleElement) {
+        return;
+      }
+      const focusRole =
+        focusRoleElement.dataset.cellFocusRole === "rendered"
+          ? "rendered"
+          : "editor";
+      const nextState = createNotebookActiveCellState(cell.refId, focusRole);
+      if (!nextState) {
+        return;
+      }
+      onFocusStateChange(nextState);
+    },
+    [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);
@@ -963,7 +1021,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 +1054,10 @@ export function Action({
               languageOptions={LANGUAGE_OPTIONS}
               onLanguageChange={handleLanguageChange}
               forceEditRequest={markdownEditRequest}
+              isActiveCell={isActiveCell}
+              activeFocusRole={activeFocusRole}
+              isWindowFocused={isWindowFocused}
+              onFocusRoleChange={handleMarkdownFocusRoleChange}
             />
             {/* Trash icon on the right, visible on hover */}
             <button
@@ -1051,7 +1115,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,14 +1146,18 @@ 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}
               value={cell.value}
               language={editorLanguage}
               fontSize={fontSettings.fontSize}
               fontFamily={fontSettings.fontFamily}
+              shouldFocus={isActiveCell && isWindowFocused}
               onChange={(v) => {
                 const updated = create(parser_pb.CellSchema, cell);
                 updated.value = v;
@@ -1270,7 +1340,17 @@ export function Action({
   );
 }
 
-function NotebookTabContent({ docUri }: { docUri: string }) {
+function NotebookTabContent({
+  docUri,
+  activeCell,
+  isWindowFocused,
+  onCellFocus,
+}: {
+  docUri: string;
+  activeCell: NotebookActiveCellState | null;
+  isWindowFocused: boolean;
+  onCellFocus: (docUri: string, state: NotebookActiveCellState) => void;
+}) {
   const { getNotebookData, useNotebookSnapshot } = useNotebookContext();
   const notebookSnapshot = useNotebookSnapshot(docUri);
   const cellDatas = useMemo(() => {
@@ -1329,6 +1409,10 @@ function NotebookTabContent({ docUri }: { docUri: string }) {
                   cellData={cellData}
                   docUri={docUri}
                   isFirst={index === 0}
+                  isActiveCell={activeCell?.refId === refId}
+                  activeFocusRole={activeCell?.focusRole ?? "editor"}
+                  isWindowFocused={isWindowFocused}
+                  onFocusStateChange={(state) => onCellFocus(docUri, state)}
                 />
               );
             })}
@@ -1363,6 +1447,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<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);
@@ -1410,6 +1503,27 @@ export default function Actions() {
     currentDocUri ??
     (statusTabVisible ? DRIVE_LINK_STATUS_TAB_URI : openNotebooks[0]?.uri ?? "");
 
+  const handleCellFocus = useCallback(
+    (docUri: string, state: NotebookActiveCellState) => {
+      setActiveCellsByDoc((prev) => {
+        const current = prev[docUri];
+        if (
+          current?.refId === state.refId &&
+          current.focusRole === state.focusRole
+        ) {
+          return prev;
+        }
+        const next = {
+          ...prev,
+          [docUri]: state,
+        };
+        persistNotebookActiveCellMap(next);
+        return next;
+      });
+    },
+    [],
+  );
+
   // Ensure the active tab is tracked as mounted on first render/whenever it changes.
   useEffect(() => {
     if (!currentDocUri) {
@@ -1439,6 +1553,28 @@ export default function Actions() {
     }
   }, [currentDocUri, openNotebooks, selectedTabUri, statusTabVisible]);
 
+  useEffect(() => {
+    if (typeof window === "undefined" || typeof document === "undefined") {
+      return;
+    }
+
+    const syncWindowFocus = () => {
+      setIsWindowFocused(
+        document.visibilityState === "visible" && document.hasFocus(),
+      );
+    };
+
+    syncWindowFocus();
+    window.addEventListener("focus", syncWindowFocus);
+    window.addEventListener("blur", syncWindowFocus);
+    document.addEventListener("visibilitychange", syncWindowFocus);
+    return () => {
+      window.removeEventListener("focus", syncWindowFocus);
+      window.removeEventListener("blur", syncWindowFocus);
+      document.removeEventListener("visibilitychange", syncWindowFocus);
+    };
+  }, []);
+
   // 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 +2020,14 @@ 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}
+                  isWindowFocused={
+                    isWindowFocused && resolvedSelectedTabUri === doc.uri
+                  }
+                  onCellFocus={handleCellFocus}
+                />
               </TabPanel>
             </Tabs.Content>
           ))}

app/src/components/Actions/Editor.tsx

@@ -17,6 +17,8 @@ const Editor = memo(
     readOnly = false,
     ariaLabel,
     autoFocusWhenEmpty = true,
+    shouldFocus = false,
+    onFocus,
     onChange,
     onEnter,
     onMount,
@@ -29,6 +31,8 @@ const Editor = memo(
     readOnly?: boolean;
     ariaLabel?: string;
     autoFocusWhenEmpty?: boolean;
+    shouldFocus?: boolean;
+    onFocus?: () => void;
     onChange: (value: string) => void;
     onEnter: () => void;
     onMount?: (editor: any, monaco: any) => void;
@@ -46,6 +50,8 @@ 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
     useEffect(() => {
@@ -146,6 +152,9 @@ const Editor = memo(
       contentSizeListener.current = editor.onDidContentSizeChange(() => {
         adjustHeight();
       });
+      focusListener.current = editor.onDidFocusEditorText(() => {
+        onFocus?.();
+      });
       onMount?.(editor, monaco);
     };
 
@@ -160,10 +169,21 @@ const Editor = memo(
       editorRef.current.layout?.({ width, height });
     }, [width, height]);
 
+    useEffect(() => {
+      const wasFocused = previousShouldFocusRef.current;
+      previousShouldFocusRef.current = shouldFocus;
+      if (!shouldFocus || wasFocused || readOnly || !editorRef.current) {
+        return;
+      }
+      editorRef.current.focus?.();
+    }, [readOnly, shouldFocus]);
+
     useEffect(() => {
       return () => {
         contentSizeListener.current?.dispose?.();
         contentSizeListener.current = null;
+        focusListener.current?.dispose?.();
+        focusListener.current = null;
       };
     }, []);
 
@@ -220,7 +240,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.shouldFocus === nextProps.shouldFocus
     );
   },
 );