Skip to content

Commit f7035f6

Browse files
hqhq1025hqhq1025
committed
fix(runtime): handle all 4 body tag shapes correctly
Codex flagged a malformed reconstruction when input had </body> but no <body>: the prior cascade fell through to the bare-fragment branch, wrapping content that already ended with </body> inside another <body>...</body>, producing nested/crossed tags. Refactor into a single normalizeBodyTags step that explicitly handles all four shapes (both, opener-only, closer-only, neither), then runs baseline + overlay injection on the normalized document. Adds unit tests for each shape.
1 parent 52f346a commit f7035f6

2 files changed

Lines changed: 123 additions & 23 deletions

File tree

‎packages/runtime/src/index.test.ts‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
import { describe, expect, it } from 'vitest';
22
import { buildSrcdoc } from './index';
33

4+
const BODY_OPEN_RE = /<body[^>]*>/i;
5+
46
describe('buildSrcdoc', () => {
57
it('wraps a fragment in a full document', () => {
68
const out = buildSrcdoc('<div>hi</div>');
@@ -69,6 +71,59 @@ describe('buildSrcdoc', () => {
6971
expect(out).toContain('ELEMENT_SELECTED');
7072
});
7173

74+
describe('body tag shape normalization', () => {
75+
it('shape 1: leaves matched <body>...</body> intact', () => {
76+
const html = '<html><head></head><body><p>both</p></body></html>';
77+
const out = buildSrcdoc(html);
78+
expect(out).toContain('<p>both</p>');
79+
expect(out.match(/<body[^>]*>/gi)?.length).toBe(1);
80+
expect(out.match(/<\/body\s*>/gi)?.length).toBe(1);
81+
expect(out.indexOf('<body')).toBeLessThan(out.indexOf('<p>both</p>'));
82+
expect(out.indexOf('<p>both</p>')).toBeLessThan(out.indexOf('</body>'));
83+
});
84+
85+
it('shape 2: appends </body> when only opener is present', () => {
86+
const html = '<html><head></head><body><p>open only</p>';
87+
const out = buildSrcdoc(html);
88+
expect(out.match(/<body[^>]*>/gi)?.length).toBe(1);
89+
expect(out.match(/<\/body\s*>/gi)?.length).toBe(1);
90+
expect(out.indexOf('<p>open only</p>')).toBeLessThan(out.indexOf('</body>'));
91+
expect(out.indexOf('ELEMENT_SELECTED')).toBeLessThan(out.indexOf('</body>'));
92+
});
93+
94+
it('shape 3: injects <body> opener before content when only </body> is present', () => {
95+
const html = '<p>close only</p></body>';
96+
const out = buildSrcdoc(html);
97+
const bodyOpen = out.search(BODY_OPEN_RE);
98+
const contentIdx = out.indexOf('<p>close only</p>');
99+
const bodyClose = out.indexOf('</body>');
100+
expect(bodyOpen).toBeGreaterThanOrEqual(0);
101+
expect(bodyOpen).toBeLessThan(contentIdx);
102+
expect(contentIdx).toBeLessThan(bodyClose);
103+
expect(out.match(/<body[^>]*>/gi)?.length).toBe(1);
104+
expect(out.match(/<\/body\s*>/gi)?.length).toBe(1);
105+
});
106+
107+
it('shape 3: injects <body> after <head> when </body> present without opener', () => {
108+
const html = '<html><head><title>t</title></head><p>x</p></body></html>';
109+
const out = buildSrcdoc(html);
110+
const headClose = out.indexOf('</head>');
111+
const bodyOpen = out.search(BODY_OPEN_RE);
112+
const contentIdx = out.indexOf('<p>x</p>');
113+
expect(headClose).toBeGreaterThanOrEqual(0);
114+
expect(bodyOpen).toBeGreaterThan(headClose);
115+
expect(bodyOpen).toBeLessThan(contentIdx);
116+
});
117+
118+
it('shape 4: wraps content in <body>...</body> when neither tag is present', () => {
119+
const out = buildSrcdoc('<div>none</div>');
120+
expect(out).toContain('<body>');
121+
expect(out).toContain('</body>');
122+
expect(out.indexOf('<body>')).toBeLessThan(out.indexOf('<div>none</div>'));
123+
expect(out.indexOf('<div>none</div>')).toBeLessThan(out.indexOf('</body>'));
124+
});
125+
});
126+
72127
it('injects baseline into a full document with <head>', () => {
73128
const html = '<!doctype html><html><head><title>t</title></head><body>y</body></html>';
74129
const out = buildSrcdoc(html);

‎packages/runtime/src/index.ts‎

Lines changed: 68 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,47 @@ export type { IframeErrorMessage } from './iframe-errors';
1616
const BASELINE_STYLE =
1717
'<style>html,body{margin:0;padding:0;background:var(--color-artifact-bg, #ffffff);min-height:100%;}</style>';
1818

19+
const HTML_RE = /<html[^>]*>/i;
20+
const HEAD_OPEN_RE = /<head[^>]*>/i;
21+
const HEAD_CLOSE_RE = /<\/head\s*>/i;
22+
const BODY_OPEN_RE = /<body[^>]*>/i;
23+
const BODY_CLOSE_RE = /<\/body\s*>/i;
24+
25+
/**
26+
* Ensure the document has matching <body>...</body> tags so downstream
27+
* baseline/overlay injection can rely on them. Handles all four input
28+
* shapes: both tags, opener only, closer only, neither.
29+
*/
30+
function normalizeBodyTags(html: string): string {
31+
const hasOpen = BODY_OPEN_RE.test(html);
32+
const hasClose = BODY_CLOSE_RE.test(html);
33+
34+
if (hasOpen && hasClose) return html;
35+
36+
if (hasOpen && !hasClose) {
37+
return `${html}</body>`;
38+
}
39+
40+
if (!hasOpen && hasClose) {
41+
if (HEAD_CLOSE_RE.test(html)) {
42+
return html.replace(HEAD_CLOSE_RE, (m) => `${m}<body>`);
43+
}
44+
if (HTML_RE.test(html)) {
45+
return html.replace(HTML_RE, (m) => `${m}<body>`);
46+
}
47+
return `<body>${html}`;
48+
}
49+
50+
// Neither opener nor closer.
51+
if (HEAD_CLOSE_RE.test(html)) {
52+
return `${html.replace(HEAD_CLOSE_RE, (m) => `${m}<body>`)}</body>`;
53+
}
54+
if (HTML_RE.test(html)) {
55+
return `${html.replace(HTML_RE, (m) => `${m}<body>`)}</body>`;
56+
}
57+
return `<body>${html}</body>`;
58+
}
59+
1960
/**
2061
* Build a complete srcdoc HTML string for the preview iframe.
2162
* Strips CSP <meta> tags from user content to allow overlay injection.
@@ -29,30 +70,14 @@ export function buildSrcdoc(userHtml: string): string {
2970
'',
3071
);
3172

32-
if (/<\/body\s*>/i.test(stripped)) {
33-
let withBaseline: string;
34-
if (/<head[^>]*>/i.test(stripped)) {
35-
withBaseline = stripped.replace(/<head[^>]*>/i, (match) => `${match}${BASELINE_STYLE}`);
36-
} else if (/<html[^>]*>/i.test(stripped)) {
37-
withBaseline = stripped.replace(
38-
/<html[^>]*>/i,
39-
(match) => `${match}<head>${BASELINE_STYLE}</head>`,
40-
);
41-
} else if (/<body[^>]*>/i.test(stripped)) {
42-
withBaseline = `${stripped.replace(
43-
/<body[^>]*>/i,
44-
(match) => `<html><head>${BASELINE_STYLE}</head>${match}`,
45-
)}</html>`;
46-
} else {
47-
withBaseline = `<!doctype html><html><head>${BASELINE_STYLE}</head><body>${stripped}</body></html>`;
48-
}
49-
return withBaseline.replace(
50-
/<\/body\s*>(?![\s\S]*<\/body\s*>)/i,
51-
`<script>${OVERLAY_SCRIPT}</script></body>`,
52-
);
53-
}
73+
const hasAnyStructure =
74+
HTML_RE.test(stripped) ||
75+
HEAD_OPEN_RE.test(stripped) ||
76+
BODY_OPEN_RE.test(stripped) ||
77+
BODY_CLOSE_RE.test(stripped);
5478

55-
return `<!doctype html>
79+
if (!hasAnyStructure) {
80+
return `<!doctype html>
5681
<html lang="en">
5782
<head>
5883
<meta charset="utf-8" />
@@ -65,6 +90,26 @@ ${stripped}
6590
<script>${OVERLAY_SCRIPT}</script>
6691
</body>
6792
</html>`;
93+
}
94+
95+
const normalized = normalizeBodyTags(stripped);
96+
97+
let withBaseline: string;
98+
if (HEAD_OPEN_RE.test(normalized)) {
99+
withBaseline = normalized.replace(HEAD_OPEN_RE, (match) => `${match}${BASELINE_STYLE}`);
100+
} else if (HTML_RE.test(normalized)) {
101+
withBaseline = normalized.replace(
102+
HTML_RE,
103+
(match) => `${match}<head>${BASELINE_STYLE}</head>`,
104+
);
105+
} else {
106+
withBaseline = `<html><head>${BASELINE_STYLE}</head>${normalized}</html>`;
107+
}
108+
109+
return withBaseline.replace(
110+
/<\/body\s*>(?![\s\S]*<\/body\s*>)/i,
111+
`<script>${OVERLAY_SCRIPT}</script></body>`,
112+
);
68113
}
69114

70115
/**

0 commit comments

Comments
 (0)