From 10bfae876aec43e8ca129cfebeb5b8e5b0fd3a08 Mon Sep 17 00:00:00 2001
From: Nayan <xNayanKumar@gmail.com>
Date: Sat, 15 Aug 2026 13:41:29 +0530
Subject: [PATCH] fix(overlay): give a nested same-origin document its own
 surface
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Picking inside a nested same-origin iframe — a Storybook preview is the
one everyone hits — selected the <iframe> itself: the hit test stopped at
the first document, and a node from a nested document belonged to no
surface, so selection bailed before it began. Reported in #17 by
nikbrunner, whose diagnosis (recursive elementFromPoint; each iframe's
rect already in its parent's coordinates, so scale composes only at the
outer boundary) this follows.

Rather than teaching FrameSurface to carry foreign nodes — too much of
the editor reads surface.doc expecting the node's own document — a nested
document becomes its own Surface. NestedSurface keeps `doc` honest, so
the marquee, drag ghost, DOM tree and every other within-one-document
consumer keep working untouched; only the coordinate methods compose the
iframe chain (re-derived per call, since an HMR reload replaces the
elements with no event), bounds() clips to the nested viewport, and
extract/scanTokens route to the nested document's own agent.

The resolvers do the walking: at() steps into the same-origin iframe
under the point up to MAX_NEST_DEPTH deep, of() rides frameOf's chain
walk, and nested surfaces cache weakly per document so the picker's
identity comparisons hold. Inline mode gets the same descent with the
page as root — no agent there, deliberately, since the proxy passes
third-party iframes through under inline. A cross-origin boundary stays
an opaque leaf, exactly today's behaviour. The one shared hitTestIn keeps
the drag-proxy filter in a single place; its six regression cases are
untouched and green.
---
 packages/overlay/src/surface.test.ts | 154 ++++++++++-
 packages/overlay/src/surface.ts      | 367 +++++++++++++++++++++++++--
 2 files changed, 499 insertions(+), 22 deletions(-)

diff --git a/packages/overlay/src/surface.test.ts b/packages/overlay/src/surface.test.ts
index 29b0580..1a3c8b6 100644
--- a/packages/overlay/src/surface.test.ts
+++ b/packages/overlay/src/surface.test.ts
@@ -1,6 +1,13 @@
 import { afterEach, describe, expect, it } from "vitest";
+import { type Frame, FrameManager } from "./canvas/frames";
+import type { CanvasViewport } from "./canvas/viewport";
 import { cls, PREFIX } from "./dom";
-import { InlineSurface } from "./surface";
+import {
+  CanvasResolver,
+  InlineResolver,
+  InlineSurface,
+  NestedSurface,
+} from "./surface";
 
 /*
  * Hit-testing through the editor's own overlay.
@@ -109,3 +116,148 @@ describe("InlineSurface.elementAtScreen", () => {
     expect(at()).toBeNull();
   });
 });
+
+// -- nested documents ---------------------------------------------------------
+// happy-dom does no layout, so each document's single-element hit test is
+// stubbed; what is under test is the descent, the ascent and the composition.
+
+/** Pin one document's hit test to a fixed answer. */
+function stubHit(doc: Document, node: Element | null): void {
+  Object.defineProperty(doc, "elementFromPoint", {
+    configurable: true,
+    value: () => node,
+    writable: true,
+  });
+}
+
+/** Pin an element's screen rect — see space.test.ts for the same idiom. */
+function stubRect(el: Element, left: number, top: number): void {
+  (el as HTMLElement).getBoundingClientRect = () =>
+    ({ height: 0, left, top, width: 0 }) as DOMRect;
+}
+
+interface NestedFixture {
+  frame: Frame;
+  frames: FrameManager;
+  nested: HTMLIFrameElement;
+  nestedDoc: Document;
+  resolver: CanvasResolver;
+  target: Element;
+  world: HTMLElement;
+}
+
+function nestedFixture(scale = 1): NestedFixture {
+  const world = document.createElement("div");
+  document.body.append(world);
+  const frames = new FrameManager({
+    pathname: "/",
+    storageKey: "__airship-test:surface",
+    world,
+  });
+  const frame = frames.add({ name: "one" });
+  if (!frame?.doc) {
+    throw new Error("frame did not build");
+  }
+  const nested = frame.doc.createElement("iframe");
+  frame.doc.body.append(nested);
+  const nestedDoc = nested.contentDocument as Document;
+  const target = nestedDoc.createElement("button");
+  nestedDoc.body.append(target);
+
+  stubHit(frame.doc, nested);
+  stubHit(nestedDoc, target);
+
+  const resolver = new CanvasResolver(frames, {
+    scale,
+  } as unknown as CanvasViewport);
+  return { frame, frames, nested, nestedDoc, resolver, target, world };
+}
+
+describe("CanvasResolver, over a nested same-origin iframe", () => {
+  afterEach(() => {
+    // The fixture's manager owns the global ready-hook; drop it.
+    for (const child of Array.from(document.body.children)) {
+      child.remove();
+    }
+  });
+
+  it("descends to the deepest surface instead of returning the iframe", () => {
+    const { resolver, nestedDoc, target } = nestedFixture();
+    // happy-dom's zero-size rects put {0,0} inside the frame.
+    const surface = resolver.at({ x: 0, y: 0 });
+
+    expect(surface).toBeInstanceOf(NestedSurface);
+    expect(surface?.doc).toBe(nestedDoc);
+    expect(surface?.elementAtScreen({ x: 0, y: 0 })).toBe(target);
+  });
+
+  it("resolves a nested node back to the same nested surface", () => {
+    const { resolver, nestedDoc, target } = nestedFixture();
+    const surface = resolver.of(target);
+
+    expect(surface).toBeInstanceOf(NestedSurface);
+    expect(surface?.doc).toBe(nestedDoc);
+    // Identity is stable across calls — the picker compares surfaces by it.
+    expect(resolver.of(target)).toBe(surface);
+    expect(resolver.at({ x: 0, y: 0 })).toBe(surface);
+  });
+
+  it("keeps the frame surface for the frame's own nodes", () => {
+    const { resolver, frame } = nestedFixture();
+    const own = frame.doc?.body as Element;
+
+    expect(resolver.of(own)).not.toBeInstanceOf(NestedSurface);
+    expect(resolver.of(own)?.doc).toBe(frame.doc);
+  });
+
+  it("composes the nested offset under the canvas scale", () => {
+    const { resolver, target, frame, nested } = nestedFixture(2);
+    stubRect(frame.el, 100, 50);
+    stubRect(nested, 30, 40);
+
+    const surface = resolver.of(target);
+    const screen = surface?.toScreen({ height: 4, left: 10, top: 5, width: 6 });
+
+    // left = frame.left + (rect.left + offset.x) * scale = 100 + 40*2.
+    expect(screen).toEqual({ height: 8, left: 180, top: 140, width: 12 });
+
+    const back = surface?.toLocal({ x: 180, y: 140 });
+    expect(back?.x).toBeCloseTo(10, 6);
+    expect(back?.y).toBeCloseTo(5, 6);
+  });
+
+  it("stops at a boundary whose document is unreachable", () => {
+    const { resolver, frame, nested } = nestedFixture();
+    Object.defineProperty(nested, "contentDocument", {
+      configurable: true,
+      get() {
+        return null;
+      },
+    });
+
+    const surface = resolver.at({ x: 0, y: 0 });
+    expect(surface).not.toBeInstanceOf(NestedSurface);
+    expect(surface?.doc).toBe(frame.doc);
+  });
+});
+
+describe("InlineResolver, over a nested same-origin iframe", () => {
+  it("descends and ascends with the page as the root", () => {
+    const iframe = document.createElement("iframe");
+    document.body.append(iframe);
+    const innerDoc = iframe.contentDocument as Document;
+    const target = innerDoc.createElement("p");
+    innerDoc.body.append(target);
+
+    stack([iframe, document.body]);
+    stubHit(innerDoc, target);
+
+    const resolver = new InlineResolver();
+    const surface = resolver.at({ x: 10, y: 10 });
+    expect(surface.doc).toBe(innerDoc);
+    expect(surface.elementAtScreen({ x: 10, y: 10 })).toBe(target);
+    expect(resolver.of(target)).toBe(surface);
+    // The top document keeps the one inline surface.
+    expect(resolver.of(document.body).doc).toBe(document);
+  });
+});
diff --git a/packages/overlay/src/surface.ts b/packages/overlay/src/surface.ts
index 48e0156..74c34e3 100644
--- a/packages/overlay/src/surface.ts
+++ b/packages/overlay/src/surface.ts
@@ -4,8 +4,12 @@ import { extractElementInfo } from "@airship/source/browser";
 import type { Frame, FrameManager } from "./canvas/frames";
 import {
   clipTo,
+  frameChain,
   frameScreenRect,
   frameToScreen,
+  intersectRects,
+  MAX_NEST_DEPTH,
+  nestOffset,
   type Point,
   type Rect,
   screenToFrame,
@@ -65,6 +69,64 @@ export interface SurfaceResolver {
   of: (node: Node | null) => Surface | null;
 }
 
+/**
+ * The app node under a point in one document, looking *through* the editor's
+ * own chrome. `elementsFromPoint` rather than `elementFromPoint`, because
+ * chrome that is deliberately opaque to the pointer (the reorder drag proxy,
+ * see `InlineSurface.elementAtScreen`'s history) must not win the hit; the
+ * single-element form is the fallback for engines without the stack API —
+ * which includes happy-dom, where this package's DOM tests run.
+ */
+function hitTestIn(doc: Document, point: Point): Element | null {
+  const probe = doc.elementsFromPoint?.bind(doc);
+  if (!probe) {
+    return doc.elementFromPoint(point.x, point.y);
+  }
+  for (const node of probe(point.x, point.y)) {
+    if (!isEditorNode(node)) {
+      return node;
+    }
+  }
+  return null;
+}
+
+/**
+ * Step from a surface into the same-origin iframe under the point, repeatedly,
+ * and return the deepest surface. This is what makes a document nested inside
+ * a frame (Storybook's preview) pickable rather than one opaque `<iframe>`
+ * box. A cross-origin iframe stays an opaque leaf: `contentDocument` is null
+ * or throws, and the descent stops at the surface above it.
+ */
+function descend(
+  surface: Surface,
+  point: Point,
+  nestedFor: (win: Window, doc: Document) => Surface | null
+): Surface {
+  let current = surface;
+  for (let depth = 0; depth < MAX_NEST_DEPTH; depth += 1) {
+    const hit = current.elementAtScreen(point);
+    if (hit?.tagName !== "IFRAME") {
+      return current;
+    }
+    let doc: Document | null;
+    try {
+      doc = (hit as HTMLIFrameElement).contentDocument;
+    } catch {
+      return current;
+    }
+    const win = doc?.defaultView;
+    if (!(doc && win)) {
+      return current;
+    }
+    const next = nestedFor(win, doc);
+    if (!next) {
+      return current;
+    }
+    current = next;
+  }
+  return current;
+}
+
 /** Convenience: a rect from a node, measured in its own surface's space. */
 export function localRect(node: Element): Rect {
   const r = node.getBoundingClientRect();
@@ -134,19 +196,7 @@ export class InlineSurface implements Surface {
    * the proxy the first one armed.
    */
   elementAtScreen(point: Point): Element | null {
-    // Falls back to the single-element hit test where the stack API is missing,
-    // which is every pre-2017 engine and — the reason it is written down —
-    // happy-dom, where the rest of the package's DOM tests run.
-    const probe = document.elementsFromPoint?.bind(document);
-    if (!probe) {
-      return document.elementFromPoint(point.x, point.y);
-    }
-    for (const node of probe(point.x, point.y)) {
-      if (!isEditorNode(node)) {
-        return node;
-      }
-    }
-    return null;
+    return hitTestIn(document, point);
   }
 
   extract(
@@ -161,19 +211,141 @@ export class InlineSurface implements Surface {
   }
 }
 
+let nestedSeq = 0;
+
+/** Debugger-friendly ids; nothing in production reads `Surface.id`. */
+function nextNestedId(prefix: string): string {
+  nestedSeq += 1;
+  return `${prefix}/nested:${nestedSeq}`;
+}
+
+/**
+ * A same-origin document nested inside the inline surface — the app's own
+ * iframe holding the content. Scale is 1 and the composition is a plain
+ * translation. No agent is ever injected here: under inline mode the proxy
+ * deliberately passes third-party iframes through, so `extract` runs the
+ * shell's copy and source resolution may be missing where the fiber carries
+ * no owner stack.
+ */
+class InlineNestedSurface implements Surface {
+  readonly id = nextNestedId("inline");
+  readonly scale = 1;
+  private readonly nestedWin: Window;
+
+  constructor(win: Window) {
+    this.nestedWin = win;
+  }
+
+  get doc(): Document {
+    return this.nestedWin.document;
+  }
+
+  get win(): Window {
+    return this.nestedWin;
+  }
+
+  get isLive(): boolean {
+    try {
+      return Boolean(this.chain() && this.nestedWin.document?.body);
+    } catch {
+      return false;
+    }
+  }
+
+  bounds(): Rect | null {
+    const chain = this.chain();
+    if (!chain) {
+      return null;
+    }
+    const [inner] = chain;
+    const offset = nestOffset(chain);
+    return {
+      height: inner.clientHeight,
+      left: offset.x,
+      top: offset.y,
+      width: inner.clientWidth,
+    };
+  }
+
+  toScreen(rect: Rect): Rect {
+    const offset = this.offset();
+    return {
+      height: rect.height,
+      left: rect.left + offset.x,
+      top: rect.top + offset.y,
+      width: rect.width,
+    };
+  }
+
+  toLocal(point: Point): Point {
+    const offset = this.offset();
+    return { x: point.x - offset.x, y: point.y - offset.y };
+  }
+
+  elementAtScreen(point: Point): Element | null {
+    if (!this.isLive) {
+      return null;
+    }
+    return hitTestIn(this.doc, this.toLocal(point));
+  }
+
+  extract(
+    node: Element
+  ): Promise<{ context: ElementContext; source: SourceLocation | null }> {
+    return extractElementInfo(node);
+  }
+
+  scanTokens(): TokenScanResult {
+    // Same origin and same realm access, so a direct scan works inline.
+    return scanRuntimeTokens(this.doc, this.win);
+  }
+
+  /** Re-derived per call: an HMR reload or a navigation replaces the iframe
+   * elements with no event to invalidate on. */
+  private chain(): HTMLIFrameElement[] | null {
+    const chain = frameChain(this.nestedWin, window);
+    return chain && chain.length > 0 ? chain : null;
+  }
+
+  private offset(): Point {
+    const chain = this.chain();
+    return chain ? nestOffset(chain) : { x: 0, y: 0 };
+  }
+}
+
 export class InlineResolver implements SurfaceResolver {
   private readonly surface = new InlineSurface();
+  /** Keyed weakly by document, so surface identity is stable across pointer
+   * moves (the picker compares surfaces by identity) and dies with the doc. */
+  private readonly nested = new WeakMap<Document, InlineNestedSurface>();
 
   all(): Surface[] {
     return [this.surface];
   }
 
-  at(): Surface {
-    return this.surface;
+  at(point: Point): Surface {
+    return descend(this.surface, point, (win, doc) => this.nestedFor(win, doc));
   }
 
-  of(): Surface {
-    return this.surface;
+  of(node: Node | null): Surface {
+    const doc = node ? (node.ownerDocument ?? (node as Document)) : null;
+    if (!doc || doc === document) {
+      return this.surface;
+    }
+    const win = doc.defaultView;
+    if (!(win && frameChain(win, window)?.length)) {
+      return this.surface;
+    }
+    return this.nestedFor(win, doc);
+  }
+
+  private nestedFor(win: Window, doc: Document): InlineNestedSurface {
+    let surface = this.nested.get(doc);
+    if (!surface || surface.win !== win) {
+      surface = new InlineNestedSurface(win);
+      this.nested.set(doc, surface);
+    }
+    return surface;
   }
 }
 
@@ -228,8 +400,7 @@ export class FrameSurface implements Surface {
     if (!doc) {
       return null;
     }
-    const local = this.toLocal(point);
-    return doc.elementFromPoint(local.x, local.y);
+    return hitTestIn(doc, this.toLocal(point));
   }
 
   extract(
@@ -259,8 +430,140 @@ export class FrameSurface implements Surface {
   }
 }
 
+/**
+ * A same-origin document nested inside a frame — Storybook's preview iframe
+ * is the canonical case. Its own document stays `doc`, so every
+ * within-one-document consumer (the marquee, the drag ghost, the DOM tree)
+ * keeps working untouched; only the coordinate methods compose the chain of
+ * ancestor iframes, and `extract`/`scanTokens` route to the *nested*
+ * document's own agent.
+ */
+export class NestedSurface implements Surface {
+  readonly frame: Frame;
+  readonly id: string;
+  private readonly nestedWin: Window;
+  private readonly getScale: () => number;
+
+  constructor(frame: Frame, nestedWin: Window, getScale: () => number) {
+    this.frame = frame;
+    this.id = nextNestedId(frame.id);
+    this.nestedWin = nestedWin;
+    this.getScale = getScale;
+  }
+
+  get scale(): number {
+    return this.getScale();
+  }
+
+  get doc(): Document {
+    return this.nestedWin.document;
+  }
+
+  get win(): Window {
+    return this.nestedWin;
+  }
+
+  get isLive(): boolean {
+    try {
+      return Boolean(
+        this.frame.doc?.body && this.chain() && this.nestedWin.document?.body
+      );
+    } catch {
+      return false;
+    }
+  }
+
+  /** The frame's rect clipped to the nested viewport, so chrome over an
+   * element scrolled out of the preview pane clips correctly. */
+  bounds(): Rect | null {
+    const chain = this.chain();
+    if (!chain) {
+      return null;
+    }
+    const [inner] = chain;
+    const viewport: Rect = {
+      height: inner.clientHeight,
+      left: 0,
+      top: 0,
+      width: inner.clientWidth,
+    };
+    const nested = frameToScreen(
+      this.frame.el,
+      viewport,
+      this.scale,
+      nestOffset(chain)
+    );
+    return intersectRects(frameScreenRect(this.frame.el), nested);
+  }
+
+  toScreen(rect: Rect): Rect {
+    const chain = this.chain();
+    return frameToScreen(
+      this.frame.el,
+      rect,
+      this.scale,
+      chain ? nestOffset(chain) : undefined
+    );
+  }
+
+  toLocal(point: Point): Point {
+    const chain = this.chain();
+    return screenToFrame(
+      this.frame.el,
+      point,
+      this.scale,
+      chain ? nestOffset(chain) : undefined
+    );
+  }
+
+  elementAtScreen(point: Point): Element | null {
+    if (!this.isLive) {
+      return null;
+    }
+    return hitTestIn(this.doc, this.toLocal(point));
+  }
+
+  extract(
+    node: Element
+  ): Promise<{ context: ElementContext; source: SourceLocation | null }> {
+    const agent = this.frame.agents.get(this.nestedWin);
+    if (agent) {
+      return agent.extract(node);
+    }
+    // The same pre-registration sliver FrameSurface documents: the shell's
+    // copy resolves DOM context and may miss the source location.
+    return extractElementInfo(node);
+  }
+
+  scanTokens(): TokenScanResult {
+    const agent = this.frame.agents.get(this.nestedWin);
+    if (agent) {
+      return agent.scanTokens();
+    }
+    return { framework: "unknown", tokens: [] };
+  }
+
+  /** Re-derived per call — an HMR reload or a Storybook navigation replaces
+   * the iframe elements with no event to invalidate on. The walk is a couple
+   * of `getBoundingClientRect` calls, the order of cost `frameScreenRect`
+   * already pays per conversion. */
+  private chain(): HTMLIFrameElement[] | null {
+    const root = this.frame.win;
+    if (!root) {
+      return null;
+    }
+    const chain = frameChain(this.nestedWin, root);
+    // Length 0 would mean the frame's own document — FrameSurface's job.
+    return chain && chain.length > 0 ? chain : null;
+  }
+}
+
 export class CanvasResolver implements SurfaceResolver {
   private readonly cache = new Map<string, FrameSurface>();
+  /** Keyed weakly by document, so nested surface identity is stable across
+   * pointer moves (the picker compares surfaces by identity) and dies with
+   * the document. */
+  private readonly nested = new WeakMap<Document, NestedSurface>();
 
   private readonly frames: FrameManager;
   private readonly viewport: CanvasViewport;
@@ -276,12 +579,25 @@ export class CanvasResolver implements SurfaceResolver {
 
   at(point: Point): Surface | null {
     const frame = this.frames.frameAt(point);
-    return frame ? this.surfaceFor(frame) : null;
+    if (!frame) {
+      return null;
+    }
+    return descend(this.surfaceFor(frame), point, (win, doc) =>
+      this.nestedFor(frame, win, doc)
+    );
   }
 
   of(node: Node | null): Surface | null {
     const frame = this.frames.frameOf(node);
-    return frame ? this.surfaceFor(frame) : null;
+    if (!frame) {
+      return null;
+    }
+    const doc = node ? (node.ownerDocument ?? (node as Document)) : null;
+    if (!doc || doc === frame.doc) {
+      return this.surfaceFor(frame);
+    }
+    const win = doc.defaultView;
+    return win ? this.nestedFor(frame, win, doc) : this.surfaceFor(frame);
   }
 
   /** Surfaces are cached per frame so identity comparisons stay meaningful. */
@@ -293,4 +609,13 @@ export class CanvasResolver implements SurfaceResolver {
     }
     return surface;
   }
+
+  private nestedFor(frame: Frame, win: Window, doc: Document): NestedSurface {
+    let surface = this.nested.get(doc);
+    if (!surface || surface.frame !== frame || surface.win !== win) {
+      surface = new NestedSurface(frame, win, () => this.viewport.scale);
+      this.nested.set(doc, surface);
+    }
+    return surface;
+  }
 }