fix: PromptBus wrapper-recapture crash + TUI auto-cancel (#136)

packages/extension/src/__tests__/ctx-ui-session-start-repatch-regression.test.ts

@@ -32,20 +32,40 @@ describe("ctx.ui session_start re-patch & error handling regression (#115)", ()
     expect(patchedSelect).not.toHaveBeenCalled();
   });
 
-  it("warns when capturing an un-cached ui object with an already-wrapped method", () => {
+  it("never stores an already-wrapped method (no pristine stash) as an original; TUI arm disabled for it (#136)", () => {
     const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {});
     const wrappedSelect = vi.fn();
     (wrappedSelect as any).__isPromptBusWrapper = true;
     const uiWithWrapped = { select: wrappedSelect };
 
-    getOrCreatePristineOriginals(uiWithWrapped);
+    const orig = getOrCreatePristineOriginals(uiWithWrapped);
 
+    // Wrapper must never become the "original" — that stored wrapper is the
+    // infinite-recursion bug (#136). With no pristine stash it degrades.
+    expect(orig.select).toBeUndefined();
     expect(warnSpy).toHaveBeenCalledWith(
-      "[bridge] getOrCreatePristineOriginals: captured an already-wrapped ui method"
+      "[bridge] getOrCreatePristineOriginals: ui.select already-wrapped with no pristine original — TUI arm disabled for it",
     );
     warnSpy.mockRestore();
   });
 
+  it("recovers the true native from an already-wrapped method via __pristineOriginal stash (#136)", async () => {
+    const nativeSelect = vi.fn().mockResolvedValue("option1");
+    // Simulates a wrapper installed by a previous bridge incarnation: tagged,
+    // carrying its bound pristine native, and NOT in the process WeakMap cache
+    // (the isolated-vm-context cache-miss path that #136 crashes on).
+    const wrappedSelect = vi.fn();
+    (wrappedSelect as any).__isPromptBusWrapper = true;
+    (wrappedSelect as any).__pristineOriginal = nativeSelect;
+    const uiWithWrapped = { select: wrappedSelect };
+
+    const orig = getOrCreatePristineOriginals(uiWithWrapped);
+
+    expect(orig.select).toBeDefined();
+    await orig.select!("q", ["a"]);
+    expect(nativeSelect).toHaveBeenCalledWith("q", ["a"]);
+    expect(wrappedSelect).not.toHaveBeenCalled();
+  });
   it("survives module reload without recapturing wrappers (jiti)", async () => {
     const nativeSelect = vi.fn().mockResolvedValue("option1");
     const mockUi = {

packages/extension/src/bridge.ts

@@ -2327,15 +2327,23 @@ function initBridge(pi: ExtensionAPI) {
                 return;
               }
 
-              if (!ac.signal.aborted) {
-                const answerStr = typeof answer === "boolean" ? (answer ? "true" : "false") : answer;
-                bus.respond({
-                  id: prompt.id,
-                  answer: answerStr ?? undefined,
-                  cancelled: answerStr == null,
-                  source: "tui",
-                });
+            if (!ac.signal.aborted) {
+              const answerStr = typeof answer === "boolean" ? (answer ? "true" : "false") : answer;
+              // When the native ui method returned undefined without user
+              // interaction (RPC-mode no-op stub) AND a dashboard adapter
+              // claimed this prompt with a component, don't auto-cancel.
+              // Let the dashboard handle the prompt instead of winning the
+              // race with a spurious cancellation (#136 fresh-session cancel).
+              if (answerStr == null && bus.hasComponentClaim(prompt.id)) {
+                return;
               }
+              bus.respond({
+                id: prompt.id,
+                answer: answerStr ?? undefined,
+                cancelled: answerStr == null,
+                source: "tui",
+              });
+            }
             } catch (err) {
               if (!ac.signal.aborted) {
                 bus.respond({
@@ -2406,6 +2414,7 @@ function initBridge(pi: ExtensionAPI) {
           )
           .then((r) => (r.cancelled ? undefined : r.answer));
       (selectWrapper as any).__isPromptBusWrapper = true;
+      (selectWrapper as any).__pristineOriginal = originals.select;
       (ctx.ui as any).select = selectWrapper;
 
       const inputWrapper = (title: string, placeholder?: string, opts?: any) =>
@@ -2422,6 +2431,7 @@ function initBridge(pi: ExtensionAPI) {
           )
           .then((r) => (r.cancelled ? undefined : r.answer));
       (inputWrapper as any).__isPromptBusWrapper = true;
+      (inputWrapper as any).__pristineOriginal = originals.input;
       (ctx.ui as any).input = inputWrapper;
 
       // Persist pasted images for an ask_user input answer to disk + emit one
@@ -2492,6 +2502,7 @@ function initBridge(pi: ExtensionAPI) {
           )
           .then((r) => !r.cancelled && r.answer === "true");
       (confirmWrapper as any).__isPromptBusWrapper = true;
+      (confirmWrapper as any).__pristineOriginal = originals.confirm;
       (ctx.ui as any).confirm = confirmWrapper;
 
       const editorWrapper = (title: string, prefill?: string, opts?: any) =>
@@ -2508,6 +2519,7 @@ function initBridge(pi: ExtensionAPI) {
           )
           .then((r) => (r.cancelled ? undefined : r.answer));
       (editorWrapper as any).__isPromptBusWrapper = true;
+      (editorWrapper as any).__pristineOriginal = originals.editor;
       (ctx.ui as any).editor = editorWrapper;
 
       // ── Multiselect ──────────────────────────────────────────────
@@ -2597,6 +2609,7 @@ function initBridge(pi: ExtensionAPI) {
         });
       };
       (notifyWrapper as any).__isPromptBusWrapper = true;
+      (notifyWrapper as any).__pristineOriginal = originals.notify;
       (ctx.ui as any).notify = notifyWrapper;
     }
 

packages/extension/src/ctx-ui-originals.ts

@@ -3,6 +3,15 @@
  * Keyed by ctx.ui object identity. Stored on `process` so jiti module-cache
  * invalidation cannot recapture PromptBus wrappers as natives (same pattern
  * as bridge.ts BRIDGE_KEY).
+ *
+ * Recovery is double-keyed: the WeakMap is the fast path (same module
+ * instance), but every wrapper also carries its own pristine original
+ * (`__pristineOriginal`) so capture can unwrap an already-patched method even
+ * when the WeakMap misses — isolated extension vm contexts, an extension
+ * reload that loses module-global state, or a bridge re-initialization on a
+ * shared ctx.ui. A wrapper is never stored as an "original"; if a genuine
+ * native cannot be unwrapped, that method is left undefined (the TUI arm
+ * degrades to the dashboard adapter) instead of recursing infinitely.
  */
 
 export interface PristineUiOriginals {
@@ -15,40 +24,81 @@ export interface PristineUiOriginals {
 
 const PRISTINE_UI_KEY = "__pi_dashboard_pristine_ui__";
 
-interface PromptBusTagged {
+interface PromptBusTagged extends Function {
   __isPromptBusWrapper?: boolean;
+  __pristineOriginal?: (...args: unknown[]) => unknown;
+}
+
+/**
+ * Walk a wrapper chain down to the underlying native. Bounded by a visited set
+ * so a stale cycle cannot loop forever. Returns `undefined` when the value is
+ * a wrapper with no recoverable pristine native.
+ */
+function unwrapPristine(fn: unknown): ((...args: unknown[]) => unknown) | undefined {
+  if (typeof fn !== "function") {
+    return undefined;
+  }
+  const seen = new Set<PromptBusTagged>();
+  let cur: ((...args: unknown[]) => unknown) | undefined = fn as (
+    ...args: unknown[]
+  ) => unknown;
+  while (cur) {
+    const tagged = cur as unknown as PromptBusTagged;
+    if (!tagged.__isPromptBusWrapper) {
+      return cur;
+    }
+    if (seen.has(tagged)) {
+      return undefined; // wrapper cycle — never usable as an original
+    }
+    seen.add(tagged);
+    const inner = tagged.__pristineOriginal;
+    if (typeof inner !== "function") {
+      return undefined;
+    }
+    cur = inner;
+  }
+  return undefined;
 }
 
 export function getOrCreatePristineOriginals(ui: unknown): PristineUiOriginals {
   if (!ui || typeof ui !== "object") {
     return {};
   }
 
-  const proc = process as unknown as Record<PropertyKey, WeakMap<object, PristineUiOriginals> | undefined>;
+  const proc = process as unknown as Record<string, WeakMap<object, PristineUiOriginals> | undefined>;
   const map = proc[PRISTINE_UI_KEY] ?? (proc[PRISTINE_UI_KEY] = new WeakMap<object, PristineUiOriginals>());
   const existing = map.get(ui);
   if (existing) {
     return existing;
   }
 
   const record = ui as Record<string, unknown>;
-  const keys: (keyof PristineUiOriginals)[] = ["notify", "select", "input", "confirm", "editor"];
+  const keys = ["notify", "select", "input", "confirm", "editor"] as const;
+
+  const captured: Record<string, unknown> = {};
+  let warnedUnrecoverable = false;
   for (const key of keys) {
     const fn = record[key];
-    if (typeof fn === "function" && (fn as PromptBusTagged).__isPromptBusWrapper) {
-      console.warn("[bridge] getOrCreatePristineOriginals: captured an already-wrapped ui method");
-      break;
+    if (typeof fn !== "function") {
+      continue;
+    }
+    const original = unwrapPristine(fn);
+    if (!original) {
+      if ((fn as PromptBusTagged).__isPromptBusWrapper && !warnedUnrecoverable) {
+        console.warn(
+          `[bridge] getOrCreatePristineOriginals: ui.${key} already-wrapped with no pristine original — TUI arm disabled for it`,
+        );
+        warnedUnrecoverable = true;
+      }
+      continue;
     }
+    if ((fn as PromptBusTagged).__isPromptBusWrapper && !(fn as PromptBusTagged).__pristineOriginal) {
+      (fn as PromptBusTagged).__pristineOriginal = original;
+    }
+    captured[key] = original.bind(ui);
   }
 
-  const captured: PristineUiOriginals = {
-    notify: typeof record.notify === "function" ? (record.notify as PristineUiOriginals["notify"])!.bind(ui) : undefined,
-    select: typeof record.select === "function" ? (record.select as PristineUiOriginals["select"])!.bind(ui) : undefined,
-    input: typeof record.input === "function" ? (record.input as PristineUiOriginals["input"])!.bind(ui) : undefined,
-    confirm: typeof record.confirm === "function" ? (record.confirm as PristineUiOriginals["confirm"])!.bind(ui) : undefined,
-    editor: typeof record.editor === "function" ? (record.editor as PristineUiOriginals["editor"])!.bind(ui) : undefined,
-  };
-
-  map.set(ui, captured);
-  return captured;
-}
+  const result = captured as unknown as PristineUiOriginals;
+  map.set(ui, result);
+  return result;
+}

packages/extension/src/prompt-bus.ts

@@ -243,6 +243,18 @@ export class PromptBus {
     entry.resolve(response);
   }
 
+
+  /** True when the pending prompt has a dashboard component claim — meaning a
+   *  non-TUI adapter returned a component for it. Used by the TUI adapter to
+   *  avoid auto-cancelling a prompt when the dashboard is rendering it and the
+   *  native ui method returned undefined as a no-op rather than a genuine user
+   *  cancel.
+   *
+   *  Returns false if the prompt is not pending (already resolved or unknown). */
+  hasComponentClaim(id: string): boolean {
+    const entry = this.pending.get(id);
+    return entry ? entry.resolvedComponent !== undefined : false;
+  }
   /**
    * Cancel a pending prompt (e.g. on timeout or abort).
    */