[Lit] Fix hydration not having the same reactive values as server (#6080)

* Fix lit hydration not having the same reactive values

* add changeset

* add clientEntrypoint to package exports

* update tests

* add changeset

* only add defer-hydration when strictly necessary

* remove second changest

* fix test typos

13 files changed, 181 insertions, 35 deletions

diff --git a/.changeset/dry-sloths-flash.md b/.changeset/dry-sloths-flash.md
new file mode 100644
index 000000000..9e8564cf9
--- /dev/null
+++ b/.changeset/dry-sloths-flash.md
@@ -0,0 +1,5 @@
+---
+'@astrojs/lit': patch
+---
+
+Fixes Lit hydration not having the same reactive values as server (losing state upon hydration)
diff --git a/packages/astro/e2e/fixtures/lit-component/src/components/NonDeferredCounter.js b/packages/astro/e2e/fixtures/lit-component/src/components/NonDeferredCounter.js
new file mode 100644
index 000000000..283d4063c
--- /dev/null
+++ b/packages/astro/e2e/fixtures/lit-component/src/components/NonDeferredCounter.js
@@ -0,0 +1,35 @@
+import { LitElement, html } from 'lit';
+
+export default class NonDeferredCounter extends LitElement {
+	static get properties() {
+		return {
+			count: {
+				type: Number,
+				// All set properties are reflected to attributes so its hydration is
+				// not deferred.
+				reflect: true,
+			},
+		};
+	}
+
+	constructor() {
+		super();
+		this.count = 0;
+	}
+
+	increment() {
+		this.count++;
+	}
+
+	render() {
+		return html`
+			<div>
+				<p>Count: ${this.count}</p>
+
+				<button type="button" @click=${this.increment}>Increment</button>
+			</div>
+		`;
+	}
+}
+
+customElements.define('non-deferred-counter', NonDeferredCounter);
diff --git a/packages/astro/e2e/fixtures/lit-component/src/pages/index.astro b/packages/astro/e2e/fixtures/lit-component/src/pages/index.astro
index ef86839d6..f15e832dc 100644
--- a/packages/astro/e2e/fixtures/lit-component/src/pages/index.astro
+++ b/packages/astro/e2e/fixtures/lit-component/src/pages/index.astro
@@ -1,8 +1,9 @@
 ---
 import MyCounter from '../components/Counter.js';
+import NonDeferredCounter from '../components/NonDeferredCounter.js';
 
 const someProps = {
-	count: 0,
+	count: 10,
 };
 ---
 
@@ -15,6 +16,9 @@ const someProps = {
 			<h1>Hello, client:idle!</h1>
 		</MyCounter>
 
+		<NonDeferredCounter id="non-deferred" client:idle {...someProps}>
+		</NonDeferredCounter>
+
 		<MyCounter id="client-load" {...someProps} client:load>
 			<h1>Hello, client:load!</h1>
 		</MyCounter>
diff --git a/packages/astro/e2e/fixtures/lit-component/src/pages/media.astro b/packages/astro/e2e/fixtures/lit-component/src/pages/media.astro
index a05d52863..69e1d7005 100644
--- a/packages/astro/e2e/fixtures/lit-component/src/pages/media.astro
+++ b/packages/astro/e2e/fixtures/lit-component/src/pages/media.astro
@@ -2,7 +2,7 @@
 import MyCounter from '../components/Counter.js';
 
 const someProps = {
-	count: 0,
+	count: 10,
 };
 ---
 
diff --git a/packages/astro/e2e/fixtures/lit-component/src/pages/solo.astro b/packages/astro/e2e/fixtures/lit-component/src/pages/solo.astro
index 1d2745e47..c9dfe0ba0 100644
--- a/packages/astro/e2e/fixtures/lit-component/src/pages/solo.astro
+++ b/packages/astro/e2e/fixtures/lit-component/src/pages/solo.astro
@@ -2,7 +2,7 @@
 import MyCounter from '../components/Counter.js';
 
 const someProps = {
-	count: 0,
+	count: 10,
 };
 ---
 
diff --git a/packages/astro/e2e/lit-component.test.js b/packages/astro/e2e/lit-component.test.js
index 6603d5b80..85af631b9 100644
--- a/packages/astro/e2e/lit-component.test.js
+++ b/packages/astro/e2e/lit-component.test.js
@@ -32,12 +32,25 @@ test.describe('Lit components', () => {
 			await expect(counter).toHaveCount(1);
 
 			const count = counter.locator('p');
-			await expect(count, 'initial count is 0').toHaveText('Count: 0');
+			await expect(count, 'initial count is 10').toHaveText('Count: 10');
 
 			const inc = counter.locator('button');
 			await inc.click();
 
-			await expect(count, 'count incremented by 1').toHaveText('Count: 1');
+			await expect(count, 'count incremented by 1').toHaveText('Count: 11');
+		});
+
+		t('non-deferred attribute serialization', async ({ page, astro }) => {
+			await page.goto(astro.resolveUrl('/'));
+
+			const counter = page.locator('#non-deferred');
+			const count = counter.locator('p');
+			await expect(count, 'initial count is 10').toHaveText('Count: 10');
+
+			const inc = counter.locator('button');
+			await inc.click();
+
+			await expect(count, 'count incremented by 1').toHaveText('Count: 11');
 		});
 
 		t('client:load', async ({ page, astro }) => {
@@ -47,12 +60,12 @@ test.describe('Lit components', () => {
 			await expect(counter, 'component is visible').toBeVisible();
 
 			const count = counter.locator('p');
-			await expect(count, 'initial count is 0').toHaveText('Count: 0');
+			await expect(count, 'initial count is 10').toHaveText('Count: 10');
 
 			const inc = counter.locator('button');
 			await inc.click();
 
-			await expect(count, 'count incremented by 1').toHaveText('Count: 1');
+			await expect(count, 'count incremented by 1').toHaveText('Count: 11');
 		});
 
 		t('client:visible', async ({ page, astro }) => {
@@ -64,12 +77,12 @@ test.describe('Lit components', () => {
 			await expect(counter, 'component is visible').toBeVisible();
 
 			const count = counter.locator('p');
-			await expect(count, 'initial count is 0').toHaveText('Count: 0');
+			await expect(count, 'initial count is 10').toHaveText('Count: 10');
 
 			const inc = counter.locator('button');
 			await inc.click();
 
-			await expect(count, 'count incremented by 1').toHaveText('Count: 1');
+			await expect(count, 'count incremented by 1').toHaveText('Count: 11');
 		});
 
 		t('client:media', async ({ page, astro }) => {
@@ -79,18 +92,18 @@ test.describe('Lit components', () => {
 			await expect(counter, 'component is visible').toBeVisible();
 
 			const count = counter.locator('p');
-			await expect(count, 'initial count is 0').toHaveText('Count: 0');
+			await expect(count, 'initial count is 10').toHaveText('Count: 10');
 
 			const inc = counter.locator('button');
 			await inc.click();
 
-			await expect(count, 'component not hydrated yet').toHaveText('Count: 0');
+			await expect(count, 'component not hydrated yet').toHaveText('Count: 10');
 
 			// Reset the viewport to hydrate the component (max-width: 50rem)
 			await page.setViewportSize({ width: 414, height: 1124 });
 
 			await inc.click();
-			await expect(count, 'count incremented by 1').toHaveText('Count: 1');
+			await expect(count, 'count incremented by 1').toHaveText('Count: 11');
 		});
 
 		t.skip('HMR', async ({ page, astro }) => {
diff --git a/packages/astro/test/fixtures/lit-element/src/components/non-deferred-element.ts b/packages/astro/test/fixtures/lit-element/src/components/non-deferred-element.ts
new file mode 100644
index 000000000..0cd36576b
--- /dev/null
+++ b/packages/astro/test/fixtures/lit-element/src/components/non-deferred-element.ts
@@ -0,0 +1,23 @@
+import { LitElement, html } from 'lit';
+import { property, customElement } from 'lit/decorators.js';
+
+@customElement('non-deferred-counter')
+export class NonDeferredCounter extends LitElement {
+	// All set properties are reflected to attributes so its hydration is not
+	// hydration-deferred should always be set.
+	@property({ type: Number, reflect: true }) count = 0;
+
+	increment() {
+		this.count++;
+	}
+
+	render() {
+		return html`
+			<div>
+				<p>Count: ${this.count}</p>
+
+				<button type="button" @click=${this.increment}>Increment</button>
+			</div>
+		`;
+	}
+}
diff --git a/packages/astro/test/fixtures/lit-element/src/pages/index.astro b/packages/astro/test/fixtures/lit-element/src/pages/index.astro
index 408360157..394885fca 100644
--- a/packages/astro/test/fixtures/lit-element/src/pages/index.astro
+++ b/packages/astro/test/fixtures/lit-element/src/pages/index.astro
@@ -1,18 +1,25 @@
 ---
 import {MyElement} from '../components/my-element.js';
+import {NonDeferredCounter} from '../components/non-deferred-element.js';
 ---
 
 <html>
 <head>
-  <title>LitElements</title>
+	<title>LitElements</title>
 </head>
 <body>
-  <MyElement
-      foo="bar"
-      str-attr={'initialized'}
-      bool={false}
-      obj={{data: 1}}
-      reflectedStrProp={'initialized reflected'}>
-  </MyElement>
+	<MyElement
+			id="default"
+			foo="bar"
+			str-attr={'initialized'}
+			bool={false}
+			obj={{data: 1}}
+			reflectedStrProp={'initialized reflected'}>
+	</MyElement>
+	<NonDeferredCounter
+			id="non-deferred"
+			count={10}
+			foo="bar">
+	</NonDeferredCounter>
 </body>
 </html>
diff --git a/packages/astro/test/lit-element.test.js b/packages/astro/test/lit-element.test.js
index a74fb0fc5..0b42c5699 100644
--- a/packages/astro/test/lit-element.test.js
+++ b/packages/astro/test/lit-element.test.js
@@ -30,36 +30,59 @@ describe('LitElement test', function () {
 		const $ = cheerio.load(html);
 
 		// test 1: attributes rendered – non reactive properties
-		expect($('my-element').attr('foo')).to.equal('bar');
+		expect($('#default').attr('foo')).to.equal('bar');
 
 		// test 2: shadow rendered
-		expect($('my-element').html()).to.include(`<div>Testing...</div>`);
+		expect($('#default').html()).to.include(`<div>Testing...</div>`);
 
 		// test 3: string reactive property set
-		expect(stripExpressionMarkers($('my-element').html())).to.include(
+		expect(stripExpressionMarkers($('#default').html())).to.include(
 			`<div id="str">initialized</div>`
 		);
 
 		// test 4: boolean reactive property correctly set
 		// <my-element bool="false"> Lit will equate to true because it uses
 		// this.hasAttribute to determine its value
-		expect(stripExpressionMarkers($('my-element').html())).to.include(`<div id="bool">B</div>`);
+		expect(stripExpressionMarkers($('#default').html())).to.include(`<div id="bool">B</div>`);
 
 		// test 5: object reactive property set
 		// by default objects will be stringified to [object Object]
-		expect(stripExpressionMarkers($('my-element').html())).to.include(
+		expect(stripExpressionMarkers($('#default').html())).to.include(
 			`<div id="data">data: 1</div>`
 		);
 
 		// test 6: reactive properties are not rendered as attributes
-		expect($('my-element').attr('obj')).to.equal(undefined);
-		expect($('my-element').attr('bool')).to.equal(undefined);
-		expect($('my-element').attr('str')).to.equal(undefined);
+		expect($('#default').attr('obj')).to.equal(undefined);
+		expect($('#default').attr('bool')).to.equal(undefined);
+		expect($('#default').attr('str')).to.equal(undefined);
 
 		// test 7: reflected reactive props are rendered as attributes
-		expect($('my-element').attr('reflectedbool')).to.equal('');
-		expect($('my-element').attr('reflected-str')).to.equal('default reflected string');
-		expect($('my-element').attr('reflected-str-prop')).to.equal('initialized reflected');
+		expect($('#default').attr('reflectedbool')).to.equal('');
+		expect($('#default').attr('reflected-str')).to.equal('default reflected string');
+		expect($('#default').attr('reflected-str-prop')).to.equal('initialized reflected');
+	});
+
+	it('Sets defer-hydration on element only when necessary', async () => {
+		// @lit-labs/ssr/ requires Node 13.9 or higher
+		if (NODE_VERSION < 13.9) {
+			return;
+		}
+		const html = await fixture.readFile('/index.html');
+		const $ = cheerio.load(html);
+
+		// test 1: reflected reactive props are rendered as attributes
+		expect($('#non-deferred').attr('count')).to.equal('10');
+
+		// test 2: non-reactive props are set as attributes
+		expect($('#non-deferred').attr('foo')).to.equal('bar');
+
+		// test 3: components with only reflected reactive props set are not
+		// deferred because their state can be completely serialized via attributes
+		expect($('#non-deferred').attr('defer-hydration')).to.equal(undefined);
+
+		// test 4: components with non-reflected reactive props set are deferred because
+		// their state needs to be synced with the server on the client.
+		expect($('#default').attr('defer-hydration')).to.equal('');
 	});
 
 	it('Correctly passes child slots', async () => {
@@ -74,7 +97,7 @@ describe('LitElement test', function () {
 		const $slottedMyElement = $('#slotted');
 		const $slottedSlottedMyElement = $('#slotted-slotted');
 
-		expect($('my-element').length).to.equal(3);
+		expect($('#default').length).to.equal(3);
 
 		// Root my-element
 		expect($rootMyElement.children('.default').length).to.equal(2);
diff --git a/packages/integrations/lit/package.json b/packages/integrations/lit/package.json
index ac0a8608c..1136570ee 100644
--- a/packages/integrations/lit/package.json
+++ b/packages/integrations/lit/package.json
@@ -23,6 +23,7 @@
     ".": "./dist/index.js",
     "./server.js": "./server.js",
     "./client-shim.js": "./client-shim.js",
+    "./dist/client.js": "./dist/client.js",
     "./hydration-support.js": "./hydration-support.js",
     "./package.json": "./package.json"
   },
diff --git a/packages/integrations/lit/server.js b/packages/integrations/lit/server.js
index 83737c183..48a3c646f 100644
--- a/packages/integrations/lit/server.js
+++ b/packages/integrations/lit/server.js
@@ -36,10 +36,18 @@ function* render(Component, attrs, slots) {
 
 	// LitElementRenderer creates a new element instance, so copy over.
 	const Ctr = getCustomElementConstructor(tagName);
+	let shouldDeferHydration = false;
+
 	if (attrs) {
 		for (let [name, value] of Object.entries(attrs)) {
-			// check if this is a reactive property
-			if (name in Ctr.prototype) {
+			const isReactiveProperty = name in Ctr.prototype;
+			const isReflectedReactiveProperty = Ctr.elementProperties.get(name)?.reflect;
+
+			// Only defer hydration if we are setting a reactive property that cannot
+			// be reflected / serialized as a property.
+			shouldDeferHydration ||= isReactiveProperty && !isReflectedReactiveProperty;
+
+			if (isReactiveProperty) {
 				instance.setProperty(name, value);
 			} else {
 				instance.setAttribute(name, value);
@@ -49,7 +57,7 @@ function* render(Component, attrs, slots) {
 
 	instance.connectedCallback();
 
-	yield `<${tagName}`;
+	yield `<${tagName}${shouldDeferHydration ? ' defer-hydration' : ''}`;
 	yield* instance.renderAttributes();
 	yield `>`;
 	const shadowContents = instance.renderShadow({});
diff --git a/packages/integrations/lit/src/client.ts b/packages/integrations/lit/src/client.ts
new file mode 100644
index 000000000..00f126e34
--- /dev/null
+++ b/packages/integrations/lit/src/client.ts
@@ -0,0 +1,25 @@
+export default (element: HTMLElement) =>
+	async (
+		Component: any,
+		props: Record<string, any>,
+	) => {
+		// Get the LitElement element instance (may or may not be upgraded).
+		const component = element.children[0] as HTMLElement;
+
+		// If there is no deferral of hydration, then all reactive properties are
+		// already serialzied as reflected attributes, or no reactive props were set
+		if (!component || !component.hasAttribute('defer-hydration')) {
+			return;
+		}
+
+		// Set properties on the LitElement instance for resuming hydration.
+		for (let [name, value] of Object.entries(props)) {
+			// Check if reactive property or class property.
+			if (name in Component.prototype) {
+				(component as any)[name] = value;
+			}
+		}
+
+		// Tell LitElement to resume hydration.
+		component.removeAttribute('defer-hydration');
+	};
diff --git a/packages/integrations/lit/src/index.ts b/packages/integrations/lit/src/index.ts
index 5b428ef8d..de6d5b0f9 100644
--- a/packages/integrations/lit/src/index.ts
+++ b/packages/integrations/lit/src/index.ts
@@ -5,6 +5,7 @@ function getViteConfiguration() {
 	return {
 		optimizeDeps: {
 			include: [
+				'@astrojs/lit/dist/client.js',
 				'@astrojs/lit/client-shim.js',
 				'@astrojs/lit/hydration-support.js',
 				'@webcomponents/template-shadowroot/template-shadowroot.js',
@@ -34,6 +35,7 @@ export default function (): AstroIntegration {
 				addRenderer({
 					name: '@astrojs/lit',
 					serverEntrypoint: '@astrojs/lit/server.js',
+					clientEntrypoint: '@astrojs/lit/dist/client.js',
 				});
 				// Update the vite configuration.
 				updateConfig({