fix(ui): White screen when hedgehog mode loses its WebGL context

apps/code/src/renderer/platform-adapters/hedgehog-mode-host.ts

@@ -28,8 +28,23 @@ export class RendererHedgehogModeHost implements HedgehogModeHost {
 
     await game.render(container);
 
+    const canvas = game.app.canvas;
+    const notifyContextLost = () => options.onContextLost?.();
+    canvas.addEventListener("webglcontextlost", notifyContextLost, {
+      once: true,
+    });
+
     return {
-      destroy: () => game.destroy(),
+      destroy: () => {
+        canvas.removeEventListener("webglcontextlost", notifyContextLost);
+        game.destroy();
+      },
+      isContextLost: () => {
+        const renderer = game.app.renderer as unknown as {
+          context?: { isLost?: boolean };
+        };
+        return renderer.context?.isLost === true;
+      },
     };
   }
 }

packages/ui/src/shell/HedgehogMode.test.tsx

@@ -0,0 +1,199 @@
+import { act, cleanup, render } from "@testing-library/react";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+
+const mocks = vi.hoisted(() => ({
+  mount: vi.fn(),
+  destroy: vi.fn(),
+  captureException: vi.fn(),
+  contextLost: false,
+  onContextLost: undefined as (() => void) | undefined,
+}));
+
+const settingsState = vi.hoisted(() => ({
+  hedgehogMode: true,
+  setHedgehogMode: () => {},
+}));
+
+vi.mock("@posthog/di/react", () => ({
+  useService: () => ({ mount: mocks.mount }),
+}));
+
+vi.mock("../features/auth/useMeQuery", () => ({
+  useMeQuery: () => ({ data: undefined }),
+}));
+
+vi.mock("../features/settings/settingsStore", () => ({
+  useSettingsStore: (selector: (state: unknown) => unknown) =>
+    selector(settingsState),
+}));
+
+vi.mock("./analytics", () => ({
+  captureException: mocks.captureException,
+}));
+
+vi.mock("./logger", () => ({
+  logger: {
+    scope: () => ({
+      error: vi.fn(),
+      warn: vi.fn(),
+      info: vi.fn(),
+      debug: vi.fn(),
+    }),
+  },
+}));
+
+import { HedgehogMode } from "./HedgehogMode";
+import { useRendererWindowFocusStore } from "./rendererWindowFocusStore";
+
+function mountGameInto(
+  container: HTMLDivElement,
+  options: { onContextLost?: () => void },
+) {
+  const canvas = document.createElement("canvas");
+  container.appendChild(canvas);
+  mocks.onContextLost = options.onContextLost;
+  mocks.destroy.mockImplementation(() => canvas.remove());
+  return Promise.resolve({
+    destroy: mocks.destroy,
+    isContextLost: () => mocks.contextLost,
+  });
+}
+
+async function loseContext() {
+  mocks.contextLost = true;
+  await act(async () => {
+    mocks.onContextLost?.();
+  });
+}
+
+async function renderHedgehogMode() {
+  const view = render(<HedgehogMode />);
+  await act(async () => {});
+  return { view, overlay: view.container.firstElementChild as HTMLDivElement };
+}
+
+async function remountAfterDelay() {
+  mocks.contextLost = false;
+  await act(async () => {
+    vi.advanceTimersByTime(2000);
+  });
+}
+
+beforeEach(() => {
+  vi.useFakeTimers();
+  mocks.mount.mockImplementation(mountGameInto);
+  mocks.contextLost = false;
+  settingsState.hedgehogMode = true;
+});
+
+afterEach(() => {
+  cleanup();
+  vi.runOnlyPendingTimers();
+  vi.useRealTimers();
+  vi.clearAllMocks();
+});
+
+describe("HedgehogMode", () => {
+  it("mounts the game into the overlay container", async () => {
+    const { overlay } = await renderHedgehogMode();
+
+    expect(mocks.mount).toHaveBeenCalledTimes(1);
+    expect(overlay.querySelector("canvas")).not.toBeNull();
+    expect(overlay.style.visibility).toBe("visible");
+  });
+
+  it("destroys the game and reports when the context loss callback fires", async () => {
+    const { overlay } = await renderHedgehogMode();
+
+    await loseContext();
+
+    expect(mocks.destroy).toHaveBeenCalledTimes(1);
+    expect(overlay.querySelector("canvas")).toBeNull();
+    expect(mocks.captureException).toHaveBeenCalledWith(
+      expect.any(Error),
+      expect.objectContaining({ source: "hedgehog-mode", losses: 1 }),
+    );
+  });
+
+  it("tears down when polling detects a lost context without a callback", async () => {
+    const { overlay } = await renderHedgehogMode();
+
+    mocks.contextLost = true;
+    await act(async () => {
+      vi.advanceTimersByTime(10_000);
+    });
+
+    expect(mocks.destroy).toHaveBeenCalledTimes(1);
+    expect(overlay.querySelector("canvas")).toBeNull();
+    expect(mocks.captureException).toHaveBeenCalledTimes(1);
+  });
+
+  it("tears down when a lost context is detected on window focus", async () => {
+    const { overlay } = await renderHedgehogMode();
+
+    mocks.contextLost = true;
+    await act(async () => {
+      useRendererWindowFocusStore.setState({ focused: true });
+    });
+
+    expect(mocks.destroy).toHaveBeenCalledTimes(1);
+    expect(overlay.querySelector("canvas")).toBeNull();
+  });
+
+  it("remounts the game after the context loss delay", async () => {
+    const { overlay } = await renderHedgehogMode();
+
+    await loseContext();
+    await remountAfterDelay();
+
+    expect(mocks.mount).toHaveBeenCalledTimes(2);
+    expect(overlay.querySelector("canvas")).not.toBeNull();
+    expect(overlay.style.visibility).toBe("visible");
+  });
+
+  it("hides the overlay after repeated context losses", async () => {
+    const { overlay } = await renderHedgehogMode();
+
+    for (let loss = 0; loss < 4; loss += 1) {
+      await loseContext();
+      await remountAfterDelay();
+    }
+
+    expect(mocks.mount).toHaveBeenCalledTimes(4);
+    expect(overlay.querySelector("canvas")).toBeNull();
+    expect(overlay.style.visibility).toBe("hidden");
+
+    await act(async () => {
+      vi.advanceTimersByTime(10_000);
+    });
+    expect(mocks.mount).toHaveBeenCalledTimes(4);
+  });
+
+  it("destroys the game on toggle off and remounts armed on re-enable", async () => {
+    const { view, overlay } = await renderHedgehogMode();
+    expect(mocks.mount).toHaveBeenCalledTimes(1);
+
+    settingsState.hedgehogMode = false;
+    view.rerender(<HedgehogMode />);
+    expect(mocks.destroy).toHaveBeenCalledTimes(1);
+    expect(overlay.querySelector("canvas")).toBeNull();
+    expect(overlay.style.visibility).toBe("hidden");
+
+    settingsState.hedgehogMode = true;
+    view.rerender(<HedgehogMode />);
+    await act(async () => {});
+    expect(mocks.mount).toHaveBeenCalledTimes(2);
+
+    await loseContext();
+    expect(mocks.captureException).toHaveBeenCalledTimes(1);
+    expect(overlay.querySelector("canvas")).toBeNull();
+  });
+
+  it("destroys the game on unmount", async () => {
+    const { view } = await renderHedgehogMode();
+
+    view.unmount();
+
+    expect(mocks.destroy).toHaveBeenCalledTimes(1);
+  });
+});

packages/ui/src/shell/HedgehogMode.tsx

@@ -1,15 +1,20 @@
 import { useService } from "@posthog/di/react";
-import { useEffect, useRef } from "react";
+import { useEffect, useRef, useState } from "react";
 import { useMeQuery } from "../features/auth/useMeQuery";
 import { useSettingsStore } from "../features/settings/settingsStore";
+import { captureException } from "./analytics";
 import {
   HEDGEHOG_MODE_HOST,
   type HedgehogModeHandle,
   type HedgehogModeHost,
 } from "./hedgehogModeHost";
 import { logger } from "./logger";
+import { useRendererWindowFocusStore } from "./rendererWindowFocusStore";
 
 const log = logger.scope("hedgehog-mode");
+const MAX_CONTEXT_LOSS_REMOUNTS = 3;
+const REMOUNT_DELAY_MS = 2000;
+const CONTEXT_CHECK_INTERVAL_MS = 10_000;
 
 export function HedgehogMode() {
   const hedgehogMode = useSettingsStore((s) => s.hedgehogMode);
@@ -18,12 +23,19 @@ export function HedgehogMode() {
   const host = useService<HedgehogModeHost>(HEDGEHOG_MODE_HOST);
   const containerRef = useRef<HTMLDivElement>(null);
   const handleRef = useRef<HedgehogModeHandle | null>(null);
+  const [gameDead, setGameDead] = useState(false);
 
   useEffect(() => {
-    if (!hedgehogMode || !containerRef.current || handleRef.current) return;
-    if (!host) return;
+    if (hedgehogMode) return;
+    setGameDead(false);
+  }, [hedgehogMode]);
+
+  useEffect(() => {
+    if (!hedgehogMode || gameDead || !containerRef.current || !host) return;
 
     let cancelled = false;
+    let losses = 0;
+    let remountTimer: ReturnType<typeof setTimeout> | null = null;
     const container = containerRef.current;
 
     const hedgehogConfig = user?.hedgehog_config as Record<
@@ -32,42 +44,95 @@ export function HedgehogMode() {
     > | null;
     const actorOptions = hedgehogConfig?.actor_options;
 
-    host
-      .mount(container, {
-        actorOptions,
-        onQuit: () => setHedgehogMode(false),
-      })
-      .then((handle) => {
-        if (cancelled) {
-          handle.destroy();
-          return;
-        }
-        handleRef.current = handle;
-      })
-      .catch((err) => {
-        log.error("Failed to mount hedgehog mode", err);
+    const destroyGame = () => {
+      try {
+        handleRef.current?.destroy();
+      } catch (err) {
+        log.error("Failed to destroy hedgehog mode game", err);
+      }
+      handleRef.current = null;
+      container.replaceChildren();
+    };
+
+    // A game whose rendering context died composites its full-window canvas
+    // as an opaque sheet over the whole app, so it must leave the DOM
+    // immediately.
+    const handleContextLost = () => {
+      if (!handleRef.current) return;
+      losses += 1;
+      log.error("Hedgehog mode WebGL context lost", { losses });
+      captureException(new Error("Hedgehog mode WebGL context lost"), {
+        source: "hedgehog-mode",
+        losses,
       });
+      destroyGame();
+      if (losses > MAX_CONTEXT_LOSS_REMOUNTS) {
+        setGameDead(true);
+        return;
+      }
+      remountTimer = setTimeout(() => {
+        log.warn("Remounting hedgehog mode after WebGL context loss", {
+          attempt: losses,
+        });
+        mountGame();
+      }, REMOUNT_DELAY_MS);
+    };
 
-    return () => {
-      cancelled = true;
+    // Backup for a missed context-loss callback (e.g. swallowed across
+    // sleep/wake), so a dead canvas can never linger on screen undetected.
+    const checkContext = () => {
+      if (document.hidden) return;
+      if (handleRef.current?.isContextLost()) handleContextLost();
     };
-  }, [hedgehogMode, user?.hedgehog_config, setHedgehogMode, host]);
 
-  useEffect(() => {
+    const mountGame = () => {
+      if (cancelled || handleRef.current) return;
+      host
+        .mount(container, {
+          actorOptions,
+          onQuit: () => setHedgehogMode(false),
+          onContextLost: handleContextLost,
+        })
+        .then((handle) => {
+          if (cancelled) {
+            handle.destroy();
+            return;
+          }
+          handleRef.current = handle;
+        })
+        .catch((err) => {
+          log.error("Failed to mount hedgehog mode", err);
+        });
+    };
+
+    mountGame();
+    const contextCheckInterval = setInterval(
+      checkContext,
+      CONTEXT_CHECK_INTERVAL_MS,
+    );
+    const unsubscribeFocusCheck = useRendererWindowFocusStore.subscribe(
+      (state) => {
+        if (state.focused) checkContext();
+      },
+    );
+
     return () => {
-      if (handleRef.current) {
-        handleRef.current.destroy();
-        handleRef.current = null;
+      cancelled = true;
+      clearInterval(contextCheckInterval);
+      unsubscribeFocusCheck();
+      if (remountTimer) {
+        clearTimeout(remountTimer);
       }
+      destroyGame();
     };
-  }, []);
+  }, [hedgehogMode, gameDead, user?.hedgehog_config, setHedgehogMode, host]);
 
   return (
     <div
       ref={containerRef}
       style={{
         zIndex: 999998,
-        visibility: hedgehogMode ? "visible" : "hidden",
+        visibility: hedgehogMode && !gameDead ? "visible" : "hidden",
       }}
       className="pointer-events-none fixed inset-0"
     />