fix(runtime): inject comment overlay into html previews

Sun-sunshine06 · Sun-sunshine06 · commit bed4458bf3ed · 2026-04-21T00:53:43.000+08:00
diff --git a/packages/runtime/src/index.test.ts b/packages/runtime/src/index.test.ts
@@ -9,15 +9,29 @@ describe('buildSrcdoc', () => {
     expect(out).not.toContain('Content-Security-Policy');
   });
 
-  it('passes legacy full-HTML documents through verbatim (pre-JSX-only snapshots)', () => {
+  it('keeps legacy full-HTML documents as HTML but injects the preview overlay', () => {
     // Snapshots written before the JSX-only switchover contain raw HTML
     // documents. Wrapping those as JSX makes Babel bark on the DOCTYPE /
-    // <html> tokens, so buildSrcdoc short-circuits when it detects a
-    // full HTML document and returns it unchanged.
+    // <html> tokens, so buildSrcdoc injects the preview overlay without
+    // routing them through the React+Babel wrapper.
     const html = '<html><body><p>x</p></body></html>';
-    expect(buildSrcdoc(html)).toBe(html);
+    const out = buildSrcdoc(html);
+    expect(out).toContain('<p>x</p>');
+    expect(out).toContain('CODESIGN_OVERLAY_SCRIPT');
+    expect(out).toContain('ELEMENT_SELECTED');
+    expect(out).not.toContain('AGENT_BODY_BEGIN');
+
     const doctyped = '<!DOCTYPE html><html><body><p>y</p></body></html>';
-    expect(buildSrcdoc(doctyped)).toBe(doctyped);
+    const doctypedOut = buildSrcdoc(doctyped);
+    expect(doctypedOut).toContain('<p>y</p>');
+    expect(doctypedOut).toContain('CODESIGN_OVERLAY_SCRIPT');
+    expect(doctypedOut).not.toContain('AGENT_BODY_BEGIN');
+  });
+
+  it('does not duplicate the overlay when a full-HTML document is rebuilt', () => {
+    const once = buildSrcdoc('<html><body><p>x</p></body></html>');
+    const twice = buildSrcdoc(once);
+    expect(twice).toBe(once);
   });
 
   it('wraps a fragment via the JSX path (no legacy HTML branch)', () => {
diff --git a/packages/runtime/src/index.ts b/packages/runtime/src/index.ts
@@ -33,6 +33,7 @@ export type { IframeErrorMessage } from './iframe-errors';
 
 const JSX_TEMPLATE_BEGIN = '<!-- AGENT_BODY_BEGIN -->';
 const JSX_TEMPLATE_END = '<!-- AGENT_BODY_END -->';
+const OVERLAY_MARKER = '<!-- CODESIGN_OVERLAY_SCRIPT -->';
 
 function escapeForScriptLiteral(jsx: string): string {
   // JSON.stringify handles quotes/newlines; the </script> escape prevents the
@@ -80,6 +81,24 @@ ${JSX_TEMPLATE_END}
 </html>`;
 }
 
+function overlayScriptTag(): string {
+  return `${OVERLAY_MARKER}<script>${OVERLAY_SCRIPT}</script>`;
+}
+
+function injectOverlayIntoHtmlDocument(html: string): string {
+  if (html.includes(OVERLAY_MARKER) || html.includes("type: 'ELEMENT_SELECTED'")) {
+    return html;
+  }
+  const script = overlayScriptTag();
+  if (/<\/body\s*>/i.test(html)) {
+    return html.replace(/<\/body\s*>/i, `${script}</body>`);
+  }
+  if (/<\/html\s*>/i.test(html)) {
+    return html.replace(/<\/html\s*>/i, `${script}</html>`);
+  }
+  return `${html}${script}`;
+}
+
 /**
  * Wrap an agent artifact in the vendored React + Babel skeleton, ready for
  * use as an iframe `srcdoc`. Already-wrapped payloads pass through unchanged.
@@ -108,7 +127,7 @@ export function buildSrcdoc(userSource: string): string {
   // Legacy HTML document (pre-JSX-only-switchover snapshots) — render as-is.
   const head = stripped.trimStart().slice(0, 2048).toLowerCase();
   if (head.startsWith('<!doctype') || head.startsWith('<html')) {
-    return stripped;
+    return injectOverlayIntoHtmlDocument(stripped);
   }
   return wrapJsxAsSrcdoc(stripped);
 }
