Return focus to the play surface when a clicked menu button unmounts

changes/reclaim-focus-removed.md

@@ -0,0 +1,3 @@
+### Fixed
+
+- After a click on a control that then unmounts (a start screen's "Play", a menu's last button), `@jgengine/shell` returns focus to the play surface instead of leaving it on `<body>`, so gameplay keys work without an extra click on the canvas.

packages/shell/src/GamePlayerShell.tsx

@@ -67,6 +67,7 @@ import { observableShellTracker, attachShellInputPublication } from "./shellInpu
 import { attachShellControlSuspension } from "./shellControlSuspension";
 import { LazyShell3dPresentation } from "./lazyShell3dPresentation";
 import { createPresentationDiagnosticOwnership, createPlaySurfaceFocusOwnership, PresentationRecovery, type PresentationRecoveryEvent } from "./presentationRecovery";
+import { installPlayFocusReclaim } from "./playSurfaceFocus";
 
 const DEV_USER_ID = "dev-player";
 
@@ -135,6 +136,8 @@ export function GamePlayerShell({
   const presentationDiagnostics = useRef(createPresentationDiagnosticOwnership());
   const [remotePlayers, setRemotePlayers] = useState<PresencePoseRow[]>([]);
   const wrapperRef = useRef<HTMLDivElement | null>(null);
+  // A start/menu button that unmounts on click drops focus to <body>, where gameplay keys never arrive.
+  useEffect(() => (poster ? undefined : installPlayFocusReclaim(document, () => wrapperRef.current)), [poster]);
   const recoveredFocus = useRef(createPlaySurfaceFocusOwnership<GameContext>());
   const yawRef = useRef(0);
   const pitchRef = useRef(0);

packages/shell/src/playSurfaceFocus.test.ts

@@ -0,0 +1,24 @@
+import { describe, expect, test } from "bun:test";
+import { shouldReclaimPlayFocus } from "./playSurfaceFocus";
+
+const el = (isConnected: boolean) => ({ isConnected }) as Element;
+const body = el(true);
+
+describe("shouldReclaimPlayFocus", () => {
+  test("reclaims when the focused button unmounted and focus fell to the body", () => {
+    expect(shouldReclaimPlayFocus(el(false), body, body)).toBe(true);
+    expect(shouldReclaimPlayFocus(el(false), null, body)).toBe(true);
+  });
+
+  test("leaves focus alone when the old element still exists (a deliberate blur)", () => {
+    expect(shouldReclaimPlayFocus(el(true), body, body)).toBe(false);
+  });
+
+  test("leaves focus alone when something else already took it", () => {
+    expect(shouldReclaimPlayFocus(el(false), el(true), body)).toBe(false);
+  });
+
+  test("ignores focusout events without an element target", () => {
+    expect(shouldReclaimPlayFocus(null, body, body)).toBe(false);
+  });
+});

packages/shell/src/playSurfaceFocus.ts

@@ -0,0 +1,39 @@
+/**
+ * Focus recovery for the play surface. Gameplay keys are read on the shell's focusable wrapper, so
+ * when a focused menu button unmounts (a start screen's "Play", a closed panel's last control) the
+ * browser drops focus to `<body>` and WASD silently does nothing until the player clicks the canvas.
+ * The one decision: reclaim only when the clicked element is gone from the document and
+ * nothing else took focus — a deliberate click elsewhere on the page keeps its focus.
+ * @internal
+ */
+export function shouldReclaimPlayFocus(lost: Element | null, active: Element | null, body: Element | null): boolean {
+  if (lost === null || lost.isConnected) return false;
+  return active === null || active === body;
+}
+
+/**
+ * Watches clicks (pointer, Enter or Space on a button) and, right after the click's update commits, returns focus to `surface()` if the clicked control unmounted and focus fell to the body.
+ * Chromium moves focus to `<body>` on removal without a `focusout`, so the click is the hook.
+ * Returns the cleanup.
+ * @internal
+ */
+export function installPlayFocusReclaim(doc: Document, surface: () => HTMLElement | null): () => void {
+  let timer: ReturnType<typeof setTimeout> | undefined;
+  const onClick = (event: MouseEvent) => {
+    const clicked = event.target instanceof Element ? event.target : null;
+    if (clicked === null) return;
+    clearTimeout(timer);
+    // React commits a click's update synchronously at the end of the event; a macrotask is enough and,
+    // unlike a frame callback, does not stall when the page renders slowly.
+    timer = setTimeout(() => {
+      const target = surface();
+      if (target === null || !target.isConnected) return;
+      if (shouldReclaimPlayFocus(clicked, doc.activeElement, doc.body)) target.focus({ preventScroll: true });
+    });
+  };
+  doc.addEventListener("click", onClick, true);
+  return () => {
+    clearTimeout(timer);
+    doc.removeEventListener("click", onClick, true);
+  };
+}

scripts/export-manifest.json

@@ -862,6 +862,7 @@
     "./map/terrainMap",
     "./materialOverride",
     "./multiplayer",
+    "./playSurfaceFocus",
     "./pointer/PointerOverlays",
     "./pointer/PointerProbe",
     "./pointer/pointerService",