fix: enable hydration for Lit elements used as raw HTML tags

dist/index.js

@@ -1,14 +1,26 @@
 // src/index.ts
 import { readFileSync } from "node:fs";
-function getViteConfiguration() {
+function litHydrationPlugin() {
   return {
+    name: "astro-lit-hydration",
+    transform(code, id, options) {
+      if (options?.ssr) return;
+      if (!id.includes("type=script")) return;
+      if (code.includes("hydration-support")) return;
+      return `import '@semantic-ui/astro-lit/hydration-support.js';
+` + code;
+    }
+  };
+}
+function getViteConfiguration(usePlugin) {
+  return {
+    plugins: usePlugin ? [litHydrationPlugin()] : [],
     optimizeDeps: {
       include: [
         "@semantic-ui/astro-lit/dist/client.js",
         "@semantic-ui/astro-lit/client-shim.js",
         "@semantic-ui/astro-lit/hydration-support.js",
-        "@webcomponents/template-shadowroot/template-shadowroot.js",
-        "@lit-labs/ssr-client/lit-element-hydrate-support.js"
+        "@webcomponents/template-shadowroot/template-shadowroot.js"
       ],
       exclude: ["@semantic-ui/astro-lit/server.js"]
     },
@@ -17,6 +29,15 @@ function getViteConfiguration() {
     }
   };
 }
+function litHandlesDeferredHydration() {
+  try {
+    const resolved = import.meta.resolve("@lit-labs/ssr-client/lit-element-hydrate-support.js");
+    const src = readFileSync(new URL(resolved), "utf-8");
+    return src.includes("skip-hydration") || src.includes("deferredBySSR");
+  } catch {
+    return false;
+  }
+}
 function getContainerRenderer() {
   return {
     name: "@semantic-ui/astro-lit",
@@ -32,14 +53,22 @@ function index_default() {
           "head-inline",
           readFileSync(new URL("../client-shim.min.js", import.meta.url), { encoding: "utf-8" })
         );
-        injectScript("before-hydration", `import '@semantic-ui/astro-lit/hydration-support.js';`);
+        const litNative = litHandlesDeferredHydration();
+        if (litNative) {
+          injectScript("before-hydration", `import '@lit-labs/ssr-client/lit-element-hydrate-support.js';`);
+        } else {
+          injectScript(
+            "head-inline",
+            readFileSync(new URL("../hydration-support-global.js", import.meta.url), { encoding: "utf-8" })
+          );
+        }
         addRenderer({
           name: "@semantic-ui/astro-lit",
           serverEntrypoint: "@semantic-ui/astro-lit/server.js",
           clientEntrypoint: "@semantic-ui/astro-lit/dist/client.js"
         });
         updateConfig({
-          vite: getViteConfiguration()
+          vite: getViteConfiguration(!litNative)
         });
       },
       "astro:build:setup": ({ vite, target }) => {

hydration-support-global.js

@@ -0,0 +1,43 @@
+// Inline script that sets up globalThis.litElementHydrateSupport synchronously.
+// This MUST run before any module script that imports 'lit', because LitElement
+// checks for this global during class definition and it's a one-shot opportunity.
+//
+// This sets up the STRUCTURAL patches (shadow root reuse, defer-hydration).
+// The RENDER patches (hydrate vs render on update) must be loaded from the same
+// <script> as the component library to share the same Lit instance — see
+// hydration-support.js.
+"use strict";
+globalThis.litElementHydrateSupport = function(ref) {
+  var LitElement = ref.LitElement;
+
+  // Make LitElement observe the defer-hydration attribute
+  var origObserved = Object.getOwnPropertyDescriptor(
+    Object.getPrototypeOf(LitElement), 'observedAttributes'
+  ).get;
+  Object.defineProperty(LitElement, 'observedAttributes', {
+    get: function() { return [].concat(origObserved.call(this), ['defer-hydration']); }
+  });
+
+  // Defer connectedCallback when defer-hydration is present
+  var origConnected = LitElement.prototype.connectedCallback;
+  LitElement.prototype.connectedCallback = function() {
+    if (!this.hasAttribute('defer-hydration')) origConnected.call(this);
+  };
+
+  // Resume hydration when defer-hydration attribute is removed
+  var origAttrChanged = LitElement.prototype.attributeChangedCallback;
+  LitElement.prototype.attributeChangedCallback = function(name, oldVal, newVal) {
+    if (name === 'defer-hydration' && newVal === null) origConnected.call(this);
+    origAttrChanged.call(this, name, oldVal, newVal);
+  };
+
+  // Reuse existing DSD shadow root instead of calling attachShadow
+  var origCreateRenderRoot = LitElement.prototype.createRenderRoot;
+  LitElement.prototype.createRenderRoot = function() {
+    if (this.shadowRoot) {
+      this._$AG = true;
+      return this.shadowRoot;
+    }
+    return origCreateRenderRoot.call(this);
+  };
+};

hydration-support.js

@@ -1,2 +1,105 @@
 // @ts-check
-import '@lit-labs/ssr-client/lit-element-hydrate-support.js';
+//
+// This module is only loaded when Lit doesn't handle deferred hydration
+// natively (checked at build time in src/index.ts). The structural patches
+// (shadow root reuse, defer-hydration attribute observation) are set up by
+// the inline hydration-support-global.js. This module adds the render-time
+// patches that need lit-html imports.
+
+import { LitElement } from 'lit';
+import { render } from 'lit-html';
+import { hydrate } from '@lit-labs/ssr-client';
+
+// Patch update() to use hydrate() on first render when a DSD shadow root
+// exists, and to do a clean replace when the SSR output can't match the
+// client render. Workaround for lit/lit#4822.
+
+const origUpdate = Object.getPrototypeOf(LitElement.prototype).update;
+
+function replaceSSRContent(element, templateResult) {
+	const root = element.shadowRoot;
+	root.replaceChildren();
+
+	const ctor = /** @type {typeof LitElement} */ (element.constructor);
+	if (ctor.elementStyles) {
+		const sheets = [];
+		for (const s of ctor.elementStyles) {
+			if (s instanceof CSSStyleSheet) {
+				sheets.push(s);
+			} else if (s.styleSheet) {
+				sheets.push(s.styleSheet);
+			}
+		}
+		if (sheets.length) {
+			root.adoptedStyleSheets = sheets;
+		}
+	}
+
+	// Reset renderBefore since we cleared the shadow root
+	element.renderOptions.renderBefore = root.firstChild;
+	element.__childPart = render(templateResult, root, element.renderOptions);
+}
+
+// The SSR renderer only has access to reflected attributes. If any non-
+// reflected property holds a value that would change the render output
+// compared to the default, hydrate() will hit a mismatch. Similarly, if
+// the shadow root contains child elements with defer-hydration, the SSR'd
+// DOM structure won't match the client's template expectations.
+function canHydrate(element) {
+	// Shadow root contains deferred children whose DOM won't match
+	if (element.shadowRoot?.querySelector('[defer-hydration]')) return false;
+
+	// Check non-reflected properties for non-default values. These couldn't
+	// be serialized to HTML, so the SSR output used defaults only.
+	const ctor = /** @type {typeof LitElement} */ (element.constructor);
+	if (ctor.elementProperties) {
+		for (const [name, options] of ctor.elementProperties) {
+			if (options.reflect) continue;
+			const value = element[name];
+			if (value === undefined || value === null || value === '' || value === false) continue;
+			if (Array.isArray(value) && value.length === 0) continue;
+			if (typeof value === 'function') return false;
+			if (Array.isArray(value) && value.length > 0) return false;
+			if (typeof value === 'object' && Object.keys(value).length > 0) return false;
+		}
+	}
+
+	return true;
+}
+
+LitElement.prototype.update = function update(changedProperties) {
+	const templateResult = this.render();
+	origUpdate.call(this, changedProperties);
+
+	if (this._$AG) {
+		this._$AG = false;
+
+		if (canHydrate(this)) {
+			for (let i = this.attributes.length - 1; i >= 0; i--) {
+				const attr = this.attributes[i];
+				if (attr.name.startsWith('hydrate-internals-')) {
+					this.removeAttribute(attr.name.slice(18));
+					this.removeAttribute(attr.name);
+				}
+			}
+			this.__childPart = hydrate(templateResult, this.renderRoot, this.renderOptions);
+		} else {
+			replaceSSRContent(this, templateResult);
+		}
+	} else {
+		this.__childPart = render(templateResult, this.renderRoot, this.renderOptions);
+	}
+};
+
+// Remove defer-hydration from SSR'd elements so they can initialize.
+// queueMicrotask so inline scripts that set properties run first.
+// Guard against server-side execution where document doesn't exist.
+if (typeof document !== 'undefined') {
+	queueMicrotask(() => {
+		const deferred = document.querySelectorAll('[defer-hydration]');
+		if (deferred.length === 0) return;
+		deferred.forEach((el) => {
+			el.removeAttribute('defer-hydration');
+		});
+	});
+}

package.json

@@ -22,13 +22,15 @@
     "./client-shim.js": "./client-shim.js",
     "./dist/client.js": "./dist/client.js",
     "./hydration-support.js": "./hydration-support.js",
+    "./hydration-support-global.js": "./hydration-support-global.js",
     "./package.json": "./package.json"
   },
   "files": [
     "dist",
     "client-shim.js",
     "client-shim.min.js",
     "hydration-support.js",
+    "hydration-support-global.js",
     "server.js",
     "server.d.ts",
     "server-shim.js"

src/index.ts

@@ -1,15 +1,45 @@
 import { readFileSync } from 'node:fs';
 import type { AstroIntegration, ContainerRenderer } from 'astro';
+import type { Plugin } from 'vite';
 
-function getViteConfiguration() {
+/**
+ * Vite plugin that prepends the hydration support import to any module
+ * that imports from 'lit' or 'lit-element'. This ensures the hydration
+ * patches and the component definitions end up in the same Vite chunk,
+ * sharing one LitElement prototype.
+ *
+ * Without this, injectScript('page') creates a separate entry point
+ * that Vite may split into a different chunk — the patches would target
+ * a different LitElement class than the one components use.
+ */
+function litHydrationPlugin(): Plugin {
 	return {
+		name: 'astro-lit-hydration',
+		transform(code, id, options) {
+			// Skip server-side transforms (both dev SSR and production build)
+			if (options?.ssr) return;
+			// Target only Astro <script> tags. Astro compiles them into
+			// virtual modules with IDs like:
+			//   /path/Layout.astro?astro&type=script&index=0&lang.ts
+			// Skip frontmatter modules (?id=N) and style modules (?type=style).
+			if (!id.includes('type=script')) return;
+			if (code.includes('hydration-support')) return;
+			// Prepend hydration support so it's in the same Vite chunk as the
+			// user's component import — sharing one LitElement prototype.
+			return `import '@semantic-ui/astro-lit/hydration-support.js';\n` + code;
+		},
+	};
+}
+
+function getViteConfiguration(usePlugin: boolean) {
+	return {
+		plugins: usePlugin ? [litHydrationPlugin()] : [],
 		optimizeDeps: {
 			include: [
 				'@semantic-ui/astro-lit/dist/client.js',
 				'@semantic-ui/astro-lit/client-shim.js',
 				'@semantic-ui/astro-lit/hydration-support.js',
 				'@webcomponents/template-shadowroot/template-shadowroot.js',
-				'@lit-labs/ssr-client/lit-element-hydrate-support.js',
 			],
 			exclude: ['@semantic-ui/astro-lit/server.js'],
 		},
@@ -19,6 +49,20 @@ function getViteConfiguration() {
 	};
 }
 
+/**
+ * Check whether the installed @lit-labs/ssr-client already handles deferred
+ * hydration natively. If so, we skip our workaround entirely.
+ */
+function litHandlesDeferredHydration(): boolean {
+	try {
+		const resolved = import.meta.resolve('@lit-labs/ssr-client/lit-element-hydrate-support.js');
+		const src = readFileSync(new URL(resolved), 'utf-8');
+		return src.includes('skip-hydration') || src.includes('deferredBySSR');
+	} catch {
+		return false;
+	}
+}
+
 export function getContainerRenderer(): ContainerRenderer {
 	return {
 		name: '@semantic-ui/astro-lit',
@@ -31,22 +75,38 @@ export default function (): AstroIntegration {
 		name: '@semantic-ui/astro-lit',
 		hooks: {
 			'astro:config:setup': ({ updateConfig, addRenderer, injectScript }) => {
-				// Inject the necessary polyfills on every page (inlined for speed).
+				// DSD polyfill for browsers that don't support declarative shadow DOM.
 				injectScript(
 					'head-inline',
 					readFileSync(new URL('../client-shim.min.js', import.meta.url), { encoding: 'utf-8' })
 				);
-				// Inject the hydration code, before a component is hydrated.
-				injectScript('before-hydration', `import '@semantic-ui/astro-lit/hydration-support.js';`);
-				// Add the lit renderer so that Astro can understand lit components.
+
+				const litNative = litHandlesDeferredHydration();
+
+				if (litNative) {
+					// Lit handles deferred hydration natively.
+					injectScript('before-hydration', `import '@lit-labs/ssr-client/lit-element-hydrate-support.js';`);
+				} else {
+					// Inline <script> in <head> sets up the one-shot
+					// globalThis.litElementHydrateSupport callback before any
+					// module imports lit.
+					injectScript(
+						'head-inline',
+						readFileSync(new URL('../hydration-support-global.js', import.meta.url), { encoding: 'utf-8' })
+					);
+					// The Vite plugin (litHydrationPlugin) prepends the
+					// hydration-support.js import to any module that imports
+					// from lit, ensuring they share one chunk.
+				}
+
 				addRenderer({
 					name: '@semantic-ui/astro-lit',
 					serverEntrypoint: '@semantic-ui/astro-lit/server.js',
 					clientEntrypoint: '@semantic-ui/astro-lit/dist/client.js',
 				});
-				// Update the vite configuration.
+
 				updateConfig({
-					vite: getViteConfiguration(),
+					vite: getViteConfiguration(!litNative),
 				});
 			},
 			'astro:build:setup': ({ vite, target }) => {