fix(aiassistant): wait for the Macros panel when saving a macro

src/App.tsx

@@ -165,6 +165,7 @@ import {
 } from "./utils/saveFile";
 import { formatBytes } from "./utils/memory";
 import { imageCenter, imageExtent } from "./utils/pixelCenters";
+import { PendingHandle } from "./utils/pendingHandle";
 import { listRecent } from "./storage/recentStore";
 import type {
   AnalysisResult,
@@ -188,6 +189,9 @@ import type {
 const MacroPanel = lazy(() =>
   import("./components/MacroPanel").then((m) => ({ default: m.MacroPanel })),
 );
+// Maximum wait for the Macros panel when a macro is sent to it: once its
+// chunk is loaded, the panel restores its macros before accepting new ones.
+const MACRO_PANEL_TIMEOUT_MS = 60_000;
 const NotebookPanel = lazy(() =>
   import("./components/notebook/NotebookPanel").then((m) => ({
     default: m.NotebookPanel,
@@ -484,7 +488,9 @@ export default function App() {
   treeKindRef.current = treeKind;
   const [centralView, setCentralView] = useState<CentralView>("plot");
   const notebookPanelRef = useRef<NotebookPanelHandle | null>(null);
-  const macroPanelRef = useRef<MacroPanelHandle | null>(null);
+  const [macroPanel] = useState(
+    () => new PendingHandle<MacroPanelHandle>("Macros panel"),
+  );
   const [leftPanelWidth, setLeftPanelWidth, persistLeftPanelWidth] =
     usePersistedSize("datalab-web.leftPanelWidth", 280);
   const [sidePanelWidth, setSidePanelWidth, persistSidePanelWidth] =
@@ -567,15 +573,16 @@ export default function App() {
   }, [macroFloating, setMacroFloating]);
   // "Save to Macros" in the AI Assistant: the Macros panel saves the
   // macro (whatever the runtime mode) and shows it. A floating Macros
-  // panel is already visible, so only the docked one gets focus.
+  // panel is already visible, so only the docked one gets focus. A
+  // conversation restored after a reload shows its button before the
+  // panel is ready: an early click waits for it.
   const saveMacroFromAssistant = useCallback(
     async (title: string, code: string) => {
-      const panel = macroPanelRef.current;
-      if (!panel) throw new Error("Macros panel is not available yet");
+      const panel = await macroPanel.wait(MACRO_PANEL_TIMEOUT_MS);
       await panel.importMacro(title, code);
       if (!macroFloating) setCentralView("macro");
     },
-    [macroFloating],
+    [macroFloating, macroPanel],
   );
   // Stable DOM containers for the Notebook / Macro panels.  They are
   // moved between the central area and the FloatingDockStack via
@@ -4804,7 +4811,9 @@ export default function App() {
                   // it was floating, then focus the central tab).
                   if (macroFloating) setMacroFloating(false);
                   setCentralView("macro");
-                  void macroPanelRef.current?.importMacro(title, code);
+                  void macroPanel
+                    .wait(MACRO_PANEL_TIMEOUT_MS)
+                    .then((panel) => panel.importMacro(title, code));
                 }}
               />
             </Suspense>,
@@ -4816,7 +4825,7 @@ export default function App() {
             <Suspense fallback={null}>
               <MacroPanel
                 key={`macro-${workspaceVersion}`}
-                ref={macroPanelRef}
+                ref={macroPanel.ref}
                 runtime={runtime}
                 onCountChanged={setMacroCount}
                 placement={macroFloating ? "floating" : "tab"}

src/components/MacroPanel.tsx

@@ -148,6 +148,10 @@ interface MacroState extends MacroMeta {
  * (notebook → macro conversion, the AI Assistant's "Save to Macros"
  * buttons) can inject a ready-made macro into the panel without having
  * to be visible.
+ *
+ * The handle is exposed only once the panel has restored the macros of the
+ * workspace, which would otherwise drop a macro imported earlier: callers
+ * wait for it (see :class:`PendingHandle`).
  */
 export interface MacroPanelHandle {
   /** Persist a fresh macro and open it as the active tab. */
@@ -792,7 +796,11 @@ export const MacroPanel = forwardRef<MacroPanelHandle, Props>(
       [runtime, openMacro, persistMirror],
     );
 
-    useImperativeHandle(ref, () => ({ importMacro }), [importMacro]);
+    // Exposed once the initial load has restored the macros (see
+    // ``MacroPanelHandle``): ``ref`` is attached when ``loaded`` turns true.
+    useImperativeHandle(loaded ? ref : null, () => ({ importMacro }), [
+      importMacro,
+    ]);
 
     // Close Recent… menu on outside click.
     useEffect(() => {

src/utils/pendingHandle.ts

@@ -0,0 +1,53 @@
+/**
+ * Imperative handle of a component that becomes available later.
+ *
+ * A lazy-loaded panel exposes its handle only once its chunk is loaded and
+ * it is mounted (or, in the Macros panel, once its macros are restored),
+ * whereas other panels can call it as soon as they are shown. Pass `ref`
+ * as the component's `ref` and call `wait()` instead of reading the handle:
+ * calls made too early then wait for the handle instead of failing.
+ */
+export class PendingHandle<T> {
+  private handle: T | null = null;
+  private readonly waiters = new Set<(handle: T) => void>();
+
+  /**
+   * @param name Name of the component, used in the timeout error.
+   */
+  constructor(private readonly name: string) {}
+
+  /** Callback ref receiving the handle, `null` while it is unavailable. */
+  readonly ref = (handle: T | null): void => {
+    this.handle = handle;
+    if (handle === null) return;
+    const waiters = [...this.waiters];
+    this.waiters.clear();
+    for (const resolve of waiters) resolve(handle);
+  };
+
+  /**
+   * Wait for the handle.
+   *
+   * @param timeoutMs Maximum waiting time, in milliseconds.
+   * @returns The handle, at once when it is available.
+   * @throws Error when the handle is still unavailable after `timeoutMs`.
+   */
+  wait(timeoutMs: number): Promise<T> {
+    if (this.handle !== null) return Promise.resolve(this.handle);
+    return new Promise<T>((resolve, reject) => {
+      const waiter = (handle: T): void => {
+        clearTimeout(timer);
+        resolve(handle);
+      };
+      const timer = setTimeout(() => {
+        this.waiters.delete(waiter);
+        reject(
+          new Error(
+            `${this.name} is not available after ${timeoutMs / 1000} s`,
+          ),
+        );
+      }, timeoutMs);
+      this.waiters.add(waiter);
+    });
+  }
+}

tests/e2e/ai_assistant_macro.spec.ts

@@ -18,6 +18,7 @@ import { waitForRuntimeReady } from "./fixtures";
  *     the Macros panel,
  *   * the same works after a page reload, from the conversation restored
  *     from the browser history (the demo scenario of the original report),
+ *     even when clicked before the Macros panel is loaded,
  *   * a Python code block written in a reply can be saved too.
  */
 
@@ -114,14 +115,33 @@ test.describe.serial("AI Assistant — Save to Macros (worker mode)", () => {
   });
 
   test("saves the macro of a conversation restored after a reload", async () => {
-    await page.reload();
-    await waitForRuntimeReady(page);
-    // Macros live in the (unsaved) workspace: the reload dropped the one
-    // saved above, while the conversation is restored from IndexedDB.
-    await expect(page.getByText("Here is the macro.")).toBeVisible();
-    await page.getByRole("button", { name: /Save to Macros/ }).click();
+    // The Macros panel is loaded on demand, whereas the restored
+    // conversation shows its "Save to Macros" button at once: hold the
+    // panel's module until the click, so that the save has to wait for
+    // the panel to load, then to restore its macros.
+    let releasePanel!: () => void;
+    const panelHeld = new Promise<void>((resolve) => {
+      releasePanel = resolve;
+    });
+    const panelModule = /\/src\/components\/MacroPanel\.tsx(?:\?|$)/;
+    await page.route(panelModule, async (route) => {
+      await panelHeld;
+      await route.continue();
+    });
+    try {
+      await page.reload();
+      await waitForRuntimeReady(page);
+      // Macros live in the (unsaved) workspace: the reload dropped the one
+      // saved above, while the conversation is restored from IndexedDB.
+      await expect(page.getByText("Here is the macro.")).toBeVisible();
+      await expect(page.locator(".macro-tab")).toHaveCount(0);
+      await page.getByRole("button", { name: /Save to Macros/ }).click();
+    } finally {
+      releasePanel();
+    }
     await expect(activeMacroTitle(page)).toBeVisible();
     await expect(activeMacroTitle(page)).toHaveText(TOOL_MACRO_TITLE);
+    await page.unroute(panelModule);
   });
 
   test("saves a Python code block written in a reply", async () => {

tests/ts/components/MacroPanel.test.tsx

@@ -0,0 +1,108 @@
+import { describe, expect, it, vi } from "vitest";
+import { render, screen, waitFor } from "@testing-library/react";
+import "@testing-library/jest-dom/vitest";
+import "fake-indexeddb/auto";
+
+// The macro worker plays no part in restoring or importing macros.
+vi.mock("../../../src/runtime/MacroRuntime", () => ({
+  MacroRuntime: class {
+    externalCallbacks = {};
+    preload(): Promise<void> {
+      return Promise.resolve();
+    }
+    dispose(): void {}
+  },
+}));
+
+import {
+  MacroPanel,
+  type MacroPanelHandle,
+} from "../../../src/components/MacroPanel";
+import { DialogProvider } from "../../../src/components/ConfirmDialog";
+import { ToastProvider } from "../../../src/components/Toast";
+import type { MacroRecord, RuntimeApi } from "../../../src/runtime/runtime";
+import { PendingHandle } from "../../../src/utils/pendingHandle";
+
+/**
+ * Runtime holding the macros of the workspace, whose ``listMacros`` answers
+ * only once released. Like the runtime worker, which handles calls in
+ * order, it lists the macros as they are when it is called.
+ */
+function fakeRuntime(): { runtime: RuntimeApi; releaseList: () => void } {
+  const macros: MacroRecord[] = [
+    { id: "m1", title: "Workspace macro", code: "print(1)" },
+  ];
+  let releaseList!: () => void;
+  const listHeld = new Promise<void>((resolve) => {
+    releaseList = resolve;
+  });
+  const runtime = {
+    listMacros: vi.fn(async () => {
+      const metas = macros.map(({ id, title }) => ({ id, title }));
+      await listHeld;
+      return metas;
+    }),
+    getMacro: vi.fn(async (id: string) => macros.find((m) => m.id === id)),
+    createMacro: vi.fn(async (title?: string, code?: string) => {
+      const rec = {
+        id: `m${macros.length + 1}`,
+        title: title ?? "Untitled",
+        code: code ?? "",
+      };
+      macros.push(rec);
+      return rec;
+    }),
+    setMacroCode: vi.fn(async () => undefined),
+  };
+  return { runtime: runtime as unknown as RuntimeApi, releaseList };
+}
+
+describe("MacroPanel — importing a macro", () => {
+  it("keeps a macro imported while the workspace macros are restored", async () => {
+    window.localStorage.clear();
+    const { runtime, releaseList } = fakeRuntime();
+    const panel = new PendingHandle<MacroPanelHandle>("Macros panel");
+    render(
+      <ToastProvider>
+        <DialogProvider>
+          <MacroPanel
+            ref={panel.ref}
+            runtime={runtime}
+            onSetCurrentPanel={() => undefined}
+            getSelection={() => []}
+            getCurrentPanel={() => "signal"}
+            selectObjects={() => undefined}
+            onModelChanged={() => undefined}
+            theme="light"
+          />
+        </DialogProvider>
+      </ToastProvider>,
+    );
+    await waitFor(() => expect(runtime.listMacros).toHaveBeenCalled());
+    // E.g. "Save to Macros" in the AI Assistant right after a reload.
+    let importDone = false;
+    const imported = panel
+      .wait(5000)
+      .then((handle) => handle.importMacro("Saved macro", "print(2)"))
+      .then(() => {
+        importDone = true;
+      });
+    releaseList();
+    // Both done: the import, and the restoration of the workspace macros.
+    await waitFor(() => {
+      expect(importDone).toBe(true);
+      expect(
+        screen.getByRole("tab", { name: /Workspace macro/ }),
+      ).toBeInTheDocument();
+    });
+    await imported;
+    const tabs = screen.getAllByRole("tab");
+    expect(tabs.map((tab) => tab.textContent)).toEqual([
+      expect.stringContaining("Workspace macro"),
+      expect.stringContaining("Saved macro"),
+    ]);
+    expect(screen.getByRole("tab", { selected: true })).toHaveTextContent(
+      "Saved macro",
+    );
+  });
+});

tests/ts/utils/pendingHandle.test.ts

@@ -0,0 +1,57 @@
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+
+import { PendingHandle } from "../../../src/utils/pendingHandle";
+
+interface Handle {
+  name: string;
+}
+
+describe("PendingHandle", () => {
+  beforeEach(() => {
+    vi.useFakeTimers();
+  });
+
+  afterEach(() => {
+    vi.useRealTimers();
+  });
+
+  it("resolves at once with an available handle", async () => {
+    const pending = new PendingHandle<Handle>("Panel");
+    const handle = { name: "first" };
+    pending.ref(handle);
+    await expect(pending.wait(1000)).resolves.toBe(handle);
+  });
+
+  it("waits for the handle when it is not available yet", async () => {
+    const pending = new PendingHandle<Handle>("Panel");
+    const first = pending.wait(1000);
+    const second = pending.wait(1000);
+    const handle = { name: "first" };
+    pending.ref(handle);
+    await expect(first).resolves.toBe(handle);
+    await expect(second).resolves.toBe(handle);
+  });
+
+  it("waits for the new handle of a remounted component", async () => {
+    const pending = new PendingHandle<Handle>("Panel");
+    pending.ref({ name: "first" });
+    pending.ref(null);
+    const waiting = pending.wait(1000);
+    const handle = { name: "second" };
+    pending.ref(handle);
+    await expect(waiting).resolves.toBe(handle);
+  });
+
+  it("rejects when the handle is still unavailable after the timeout", async () => {
+    const pending = new PendingHandle<Handle>("Macros panel");
+    const waiting = pending.wait(2000);
+    const outcome = expect(waiting).rejects.toThrow(
+      "Macros panel is not available after 2 s",
+    );
+    await vi.advanceTimersByTimeAsync(2000);
+    await outcome;
+    // A late handle neither resolves the expired wait nor fails.
+    pending.ref({ name: "late" });
+    await expect(pending.wait(1000)).resolves.toEqual({ name: "late" });
+  });
+});