[ssr] Add option to defer hydration on top level custom element (#2940)

.changeset/rich-monkeys-flow.md

@@ -0,0 +1,5 @@
+---
+'@lit-labs/ssr': minor
+---
+
+Add option to defer hydration of top level custom elements.

packages/labs/ssr/src/lib/render-lit-html.ts

@@ -520,12 +520,18 @@ export type RenderInfo = {
    * template, even in the case of conditional templates.
    */
   customElementRendered?: (tagName: string) => void;
+
+  /**
+   * Flag to defer hydration of top level custom element. Defaults to false.
+   */
+  deferHydration: boolean;
 };
 
 const defaultRenderInfo = {
   elementRenderers: [LitElementRenderer],
   customElementInstanceStack: [],
   customElementHostStack: [],
+  deferHydration: false,
 };
 
 declare global {
@@ -710,10 +716,13 @@ function* renderTemplateResult(
         // Render out any attributes on the instance (both static and those
         // that may have been dynamically set by the renderer)
         yield* instance.renderAttributes();
-        // If this element is nested in another, add the `defer-hydration`
-        // attribute, so that it does not enable before the host element
-        // hydrates
-        if (renderInfo.customElementHostStack.length > 0) {
+        // If deferHydration flag is true or if this element is nested in
+        // another, add the `defer-hydration` attribute, so that it does not
+        // enable before the host element hydrates
+        if (
+          renderInfo.deferHydration ||
+          renderInfo.customElementHostStack.length > 0
+        ) {
           yield ' defer-hydration';
         }
         break;

packages/labs/ssr/src/test/integration/server/server.ts

@@ -49,7 +49,10 @@ export const ssrMiddleware = () => {
     if (test.registerElements) {
       await test.registerElements();
     }
-    const result = render(test.render(...test.expectations[0].args));
+    const result = render(
+      test.render(...test.expectations[0].args),
+      test.serverRenderOptions
+    );
     context.type = 'text/html';
     context.body = Readable.from(result);
   });

packages/labs/ssr/src/test/integration/tests/basic.ts

@@ -4866,4 +4866,51 @@ export const tests: {[name: string]: SSRTest} = {
       stableSelectors: ['le-order1'],
     };
   },
+
+  'LitElement: defer hydration': () => {
+    return {
+      registerElements() {
+        class LEDefer extends LitElement {
+          @property({type: Number})
+          clicked = 0;
+          handleClick() {
+            this.clicked += 1;
+          }
+          override render() {
+            return html`<button @click=${this.handleClick}>X</button>`;
+          }
+        }
+        customElements.define('le-defer', LEDefer);
+      },
+      render() {
+        return html`<le-defer></le-defer>`;
+      },
+      serverRenderOptions: {
+        deferHydration: true,
+      },
+      expectations: [
+        {
+          args: [],
+          async check(assert: Chai.Assert, dom: HTMLElement) {
+            const el = dom.querySelector('le-defer') as LitElement & {
+              clicked: number;
+            };
+            const button = el.shadowRoot!.querySelector('button')!;
+            button.click();
+            assert.equal(el.clicked, 0);
+            el.removeAttribute('defer-hydration');
+            await el.updateComplete;
+            button.click();
+            await el.updateComplete;
+            assert.equal(el.clicked, 1);
+          },
+          html: {
+            root: `<le-defer></le-defer>`,
+            'le-defer': `<button>X</button>`,
+          },
+        },
+      ],
+      stableSelectors: ['le-defer'],
+    };
+  },
 };

packages/labs/ssr/src/test/integration/tests/ssr-test.ts

@@ -4,7 +4,8 @@
  * SPDX-License-Identifier: BSD-3-Clause
  */
 
-import {TemplateResult} from 'lit';
+import type {TemplateResult} from 'lit';
+import type {RenderInfo} from '../../../lib/render-lit-html.js';
 
 export type SSRExpectedHTML =
   | string
@@ -39,6 +40,7 @@ export interface SSRTestDescription {
   skip?: boolean;
   only?: boolean;
   registerElements?(): void | Promise<unknown>;
+  serverRenderOptions?: Partial<RenderInfo>;
 }
 
 export type SSRTestFactory = () => SSRTestDescription;

packages/labs/ssr/src/test/lib/render-lit_test.ts

@@ -209,6 +209,18 @@ test('simple custom element', async () => {
   assert.is(customElementsRendered[0], 'test-simple');
 });
 
+test('simple custom element with deferHydration', async () => {
+  const {render, simpleTemplateWithElement} = await setup();
+
+  const result = await render(simpleTemplateWithElement, {
+    deferHydration: true,
+  });
+  assert.is(
+    result,
+    `<!--lit-part tjmYe1kHIVM=--><test-simple defer-hydration><template shadowroot="open"><!--lit-part UNbWrd8S5FY=--><main></main><!--/lit-part--></template></test-simple><!--/lit-part-->`
+  );
+});
+
 test('element with property', async () => {
   const {render, elementWithProperty} = await setup();
   const result = await render(elementWithProperty);