[Markdoc] Fix global asset bleed (#6758)

* wip: propagatedAssets flag per-component

* Propagate in TreeNode

* fix: remove unused inject comment

* feat: make asset propagation an integration opt-in

* fix: remove crawlGraph stopper

* wip: logs to understand what's happening

* SSR mdoc files in dev

* feat: add astroPropagatedAssets flag with vite

* chore: remove console logs

* chore: cleanup hasContentFlag

* fix: set handlePropagation default for legacy integrations

* chore: changeset

* temp: silence acorn type error

* chore: revert pnpm-lock changes

* fix: check correct flag

* We need to handle propagation on markdown because of layouts

* Remove use of renderStyleElement

* Fix heading tests

* Fix merge conflict

* typeof function

* Switch the check

* Add comment on injection detection regexp

---------

Co-authored-by: Matthew Phillips <matthew@skypack.dev>

5 files changed, 153 insertions, 25 deletions

diff --git a/packages/integrations/markdoc/components/Renderer.astro b/packages/integrations/markdoc/components/Renderer.astro
index 5e2b6833a..6571e8c71 100644
--- a/packages/integrations/markdoc/components/Renderer.astro
+++ b/packages/integrations/markdoc/components/Renderer.astro
@@ -1,4 +1,5 @@
 ---
+//! astro-head-inject
 import type { Config } from '@markdoc/markdoc';
 import Markdoc from '@markdoc/markdoc';
 import { ComponentNode, createTreeNode } from './TreeNode.js';
@@ -14,4 +15,4 @@ const ast = Markdoc.Ast.fromJSON(stringifiedAst);
 const content = Markdoc.transform(ast, config);
 ---
 
-<ComponentNode treeNode={createTreeNode(content)} />
+<ComponentNode treeNode={await createTreeNode(content)} />
diff --git a/packages/integrations/markdoc/components/TreeNode.ts b/packages/integrations/markdoc/components/TreeNode.ts
index a60597a0d..3f9740af1 100644
--- a/packages/integrations/markdoc/components/TreeNode.ts
+++ b/packages/integrations/markdoc/components/TreeNode.ts
@@ -2,7 +2,16 @@ import type { AstroInstance } from 'astro';
 import { Fragment } from 'astro/jsx-runtime';
 import type { RenderableTreeNode } from '@markdoc/markdoc';
 import Markdoc from '@markdoc/markdoc';
-import { createComponent, renderComponent, render } from 'astro/runtime/server/index.js';
+import {
+	createComponent,
+	renderComponent,
+	render,
+	renderScriptElement,
+	renderUniqueStylesheet,
+	createHeadAndContent,
+	unescapeHTML,
+	renderTemplate,
+} from 'astro/runtime/server/index.js';
 
 export type TreeNode =
 	| {
@@ -12,6 +21,9 @@ export type TreeNode =
 	| {
 			type: 'component';
 			component: AstroInstance['default'];
+			collectedLinks?: string[];
+			collectedStyles?: string[];
+			collectedScripts?: string[];
 			props: Record<string, any>;
 			children: TreeNode[];
 	  }
@@ -32,20 +44,63 @@ export const ComponentNode = createComponent({
 				)}`,
 		};
 		if (treeNode.type === 'component') {
-			return renderComponent(
-				result,
-				treeNode.component.name,
-				treeNode.component,
-				treeNode.props,
-				slots
+			let styles = '',
+				links = '',
+				scripts = '';
+			if (Array.isArray(treeNode.collectedStyles)) {
+				styles = treeNode.collectedStyles.map((style: any) => renderUniqueStylesheet({
+					type: 'inline',
+					content: style,
+				})).join('');
+			}
+			if (Array.isArray(treeNode.collectedLinks)) {
+				links = treeNode.collectedLinks
+					.map((link: any) => {
+						return renderUniqueStylesheet(result, {
+							href: link[0] === '/' ? link : '/' + link,
+						});
+					})
+					.join('');
+			}
+			if (Array.isArray(treeNode.collectedScripts)) {
+				scripts = treeNode.collectedScripts
+					.map((script: any) => renderScriptElement(script))
+					.join('');
+			}
+
+			const head = unescapeHTML(styles + links + scripts);
+
+			let headAndContent = createHeadAndContent(
+				head,
+				renderTemplate`${renderComponent(
+					result,
+					treeNode.component.name,
+					treeNode.component,
+					treeNode.props,
+					slots
+				)}`
 			);
+
+			// Let the runtime know that this component is being used.
+			result.propagators.set(
+				{},
+				{
+					init() {
+						return headAndContent;
+					},
+				}
+			);
+
+			return headAndContent;
 		}
 		return renderComponent(result, treeNode.tag, treeNode.tag, treeNode.attributes, slots);
 	},
-	propagation: 'none',
+	propagation: 'self',
 });
 
-export function createTreeNode(node: RenderableTreeNode | RenderableTreeNode[]): TreeNode {
+export async function createTreeNode(
+	node: RenderableTreeNode | RenderableTreeNode[]
+): Promise<TreeNode> {
 	if (typeof node === 'string' || typeof node === 'number') {
 		return { type: 'text', content: String(node) };
 	} else if (Array.isArray(node)) {
@@ -53,16 +108,17 @@ export function createTreeNode(node: RenderableTreeNode | RenderableTreeNode[]):
 			type: 'component',
 			component: Fragment,
 			props: {},
-			children: node.map((child) => createTreeNode(child)),
+			children: await Promise.all(node.map((child) => createTreeNode(child))),
 		};
 	} else if (node === null || typeof node !== 'object' || !Markdoc.Tag.isTag(node)) {
 		return { type: 'text', content: '' };
 	}
 
+	const children = await Promise.all(node.children.map((child) => createTreeNode(child)));
+
 	if (typeof node.name === 'function') {
 		const component = node.name;
 		const props = node.attributes;
-		const children = node.children.map((child) => createTreeNode(child));
 
 		return {
 			type: 'component',
@@ -70,12 +126,38 @@ export function createTreeNode(node: RenderableTreeNode | RenderableTreeNode[]):
 			props,
 			children,
 		};
+	} else if (isPropagatedAssetsModule(node.name)) {
+		const { collectedStyles, collectedLinks, collectedScripts } = node.name;
+		const component = (await node.name.getMod())?.default ?? Fragment;
+		const props = node.attributes;
+
+		return {
+			type: 'component',
+			component,
+			collectedStyles,
+			collectedLinks,
+			collectedScripts,
+			props,
+			children,
+		};
 	} else {
 		return {
 			type: 'element',
 			tag: node.name,
 			attributes: node.attributes,
-			children: node.children.map((child) => createTreeNode(child)),
+			children,
 		};
 	}
 }
+
+type PropagatedAssetsModule = {
+	__astroPropagation: true;
+	getMod: () => Promise<AstroInstance['default']>;
+	collectedStyles: string[];
+	collectedLinks: string[];
+	collectedScripts: string[];
+};
+
+function isPropagatedAssetsModule(module: any): module is PropagatedAssetsModule {
+	return typeof module === 'object' && module != null && '__astroPropagation' in module;
+}
diff --git a/packages/integrations/markdoc/src/index.ts b/packages/integrations/markdoc/src/index.ts
index 627f08c77..ba8a0af84 100644
--- a/packages/integrations/markdoc/src/index.ts
+++ b/packages/integrations/markdoc/src/index.ts
@@ -32,7 +32,11 @@ export default function markdocIntegration(legacyConfig?: any): AstroIntegration
 		name: '@astrojs/markdoc',
 		hooks: {
 			'astro:config:setup': async (params) => {
-				const { config: astroConfig, addContentEntryType } = params as SetupHookParams;
+				const {
+					config: astroConfig,
+					updateConfig,
+					addContentEntryType,
+				} = params as SetupHookParams;
 
 				markdocConfigResult = await loadMarkdocConfig(astroConfig);
 				const userMarkdocConfig = markdocConfigResult?.config ?? {};
@@ -49,6 +53,9 @@ export default function markdocIntegration(legacyConfig?: any): AstroIntegration
 				addContentEntryType({
 					extensions: ['.mdoc'],
 					getEntryInfo,
+					// Markdoc handles script / style propagation
+					// for Astro components internally
+					handlePropagation: false,
 					async getRenderModule({ entry, viteId }) {
 						const ast = Markdoc.parse(entry.body);
 						const pluginContext = this;
@@ -88,7 +95,10 @@ export default function markdocIntegration(legacyConfig?: any): AstroIntegration
 							});
 						}
 
-						const res = `import { jsx as h } from 'astro/jsx-runtime';
+						const res = `import {
+							createComponent,
+							renderComponent,
+						} from 'astro/runtime/server/index.js';
 						import { Renderer } from '@astrojs/markdoc/components';
 						import { collectHeadings, setupConfig, Markdoc } from '@astrojs/markdoc/runtime';
 import * as entry from ${JSON.stringify(viteId + '?astroContentCollectionEntry')};
@@ -119,14 +129,24 @@ export function getHeadings() {
 	const content = Markdoc.transform(ast, config);
 	return collectHeadings(Array.isArray(content) ? content : content.children);
 }
-export async function Content (props) {
-	const config = setupConfig({
-		...userConfig,
-		variables: { ...userConfig.variables, ...props },
-	}, entry);
 
-	return h(Renderer, { config, stringifiedAst });
-}`;
+export const Content = createComponent({
+	factory(result, props) {
+		const config = setupConfig({
+			...userConfig,
+			variables: { ...userConfig.variables, ...props },
+		}, entry);
+		
+		return renderComponent(
+			result,
+			Renderer.name,
+			Renderer,
+			{ stringifiedAst, config },
+			{}
+		);
+	},
+	propagation: 'self',
+});`;
 						return { code: res };
 					},
 					contentModuleTypes: await fs.promises.readFile(
@@ -134,6 +154,27 @@ export async function Content (props) {
 						'utf-8'
 					),
 				});
+
+				updateConfig({
+					vite: {
+						plugins: [
+							{
+								name: '@astrojs/markdoc:astro-propagated-assets',
+								enforce: 'pre',
+								// Astro component styles and scripts should only be injected
+								// When a given Markdoc file actually uses that component.
+								// Add the `astroPropagatedAssets` flag to inject only when rendered.
+								resolveId(this: rollup.TransformPluginContext, id: string, importer: string) {
+									if (importer === markdocConfigResult?.fileUrl.pathname && id.endsWith('.astro')) {
+										return this.resolve(id + '?astroPropagatedAssets', importer, {
+											skipSelf: true,
+										});
+									}
+								},
+							},
+						],
+					},
+				});
 			},
 			'astro:server:setup': async ({ server }) => {
 				server.watcher.on('all', (event, entry) => {
diff --git a/packages/integrations/markdoc/src/nodes/heading.ts b/packages/integrations/markdoc/src/nodes/heading.ts
index 0210e9b90..cb50dd231 100644
--- a/packages/integrations/markdoc/src/nodes/heading.ts
+++ b/packages/integrations/markdoc/src/nodes/heading.ts
@@ -37,13 +37,14 @@ export const heading: Schema = {
 		const slug = getSlug(attributes, children, config.ctx.headingSlugger);
 
 		const render = config.nodes?.heading?.render ?? `h${level}`;
+
 		const tagProps =
 			// For components, pass down `level` as a prop,
 			// alongside `__collectHeading` for our `headings` collector.
 			// Avoid accidentally rendering `level` as an HTML attribute otherwise!
-			typeof render === 'function'
-				? { ...attributes, id: slug, __collectHeading: true, level }
-				: { ...attributes, id: slug };
+			typeof render === 'string'
+				? { ...attributes, id: slug }
+				: { ...attributes, id: slug, __collectHeading: true, level };
 
 		return new Markdoc.Tag(render, tagProps, children);
 	},
diff --git a/packages/integrations/mdx/src/index.ts b/packages/integrations/mdx/src/index.ts
index 2ccf66266..1ef23e1af 100644
--- a/packages/integrations/mdx/src/index.ts
+++ b/packages/integrations/mdx/src/index.ts
@@ -55,6 +55,9 @@ export default function mdx(partialMdxOptions: Partial<MdxOptions> = {}): AstroI
 						new URL('../template/content-module-types.d.ts', import.meta.url),
 						'utf-8'
 					),
+					// MDX can import scripts and styles,
+					// so wrap all MDX files with script / style propagation checks
+					handlePropagation: true,
 				});
 
 				const extendMarkdownConfig =