feat: implement next/document with class-based Document and getInitialProps support

packages/vinext/src/server/dev-server.ts

@@ -83,11 +83,50 @@ const STREAM_BODY_MARKER = "<!--VINEXT_STREAM_BODY-->";
  * deferring them reduces TTFB and lets the browser start parsing the
  * shell sooner).
  */
+/**
+ * Build a minimal DocumentContext for calling _document.getInitialProps.
+ *
+ * Next.js DocumentContext includes renderPage, defaultGetInitialProps,
+ * req/res, pathname, query, asPath, and locale info. We provide the
+ * subset that custom documents typically use.
+ */
+function buildDocumentContext(
+  req: IncomingMessage,
+  url: string,
+): import("../shims/document.js").DocumentContext {
+  const defaultGetInitialProps = async (): Promise<
+    import("../shims/document.js").DocumentInitialProps
+  > => ({
+    html: "",
+  });
+  const [pathname, search] = url.split("?");
+  const query: Record<string, string | string[] | undefined> = {};
+  if (search) {
+    for (const [k, v] of new URLSearchParams(search)) {
+      const existing = query[k];
+      if (existing !== undefined) {
+        query[k] = Array.isArray(existing) ? [...existing, v] : [existing, v];
+      } else {
+        query[k] = v;
+      }
+    }
+  }
+  return {
+    pathname: pathname || "/",
+    query,
+    asPath: url,
+    req,
+    renderPage: async () => ({ html: "" }),
+    defaultGetInitialProps,
+  };
+}
+
 async function streamPageToResponse(
   res: ServerResponse,
   element: React.ReactElement,
   options: {
     url: string;
+    req: IncomingMessage;
     server: ViteDevServer;
     fontHeadHTML: string;
     scripts: string;
@@ -100,6 +139,7 @@ async function streamPageToResponse(
 ): Promise<void> {
   const {
     url,
+    req,
     server,
     fontHeadHTML,
     scripts,
@@ -121,7 +161,26 @@ async function streamPageToResponse(
   let shellTemplate: string;
 
   if (DocumentComponent) {
-    const docElement = React.createElement(DocumentComponent);
+    // Call getInitialProps if the custom Document class defines it.
+    // This allows documents to augment props (e.g. inject a theme prop).
+    let docProps: Record<string, unknown> = {};
+    const DocClass = DocumentComponent as unknown as {
+      getInitialProps?: (
+        ctx: import("../shims/document.js").DocumentContext,
+      ) => Promise<Record<string, unknown>>;
+    };
+    if (typeof DocClass.getInitialProps === "function") {
+      try {
+        const ctx = buildDocumentContext(req, url);
+        docProps = await DocClass.getInitialProps(ctx);
+      } catch {
+        // If getInitialProps fails, fall back to rendering with no props
+      }
+    }
+    const docElement = React.createElement(
+      DocumentComponent,
+      docProps as React.ComponentProps<typeof DocumentComponent>,
+    );
     let docHtml = await renderToStringAsync(docElement);
     // Replace __NEXT_MAIN__ with our stream marker
     docHtml = docHtml.replace("__NEXT_MAIN__", STREAM_BODY_MARKER);
@@ -942,6 +1001,7 @@ hydrate();
         // Suspense content streams in as it resolves.
         await streamPageToResponse(res, element, {
           url,
+          req,
           server,
           fontHeadHTML,
           scripts: allScripts,

packages/vinext/src/shims/document.tsx

@@ -4,8 +4,39 @@
  * Provides Html, Head, Main, NextScript components for custom _document.tsx.
  * During SSR these render placeholder markers that the dev server replaces
  * with actual content.
+ *
+ * Also exports DocumentContext, DocumentInitialProps, and the base Document
+ * class for typed custom document classes that use getInitialProps.
  */
 import React from "react";
+import type { IncomingMessage, ServerResponse } from "node:http";
+
+// ─── Types ────────────────────────────────────────────────────────────────────
+
+export type DocumentInitialProps = {
+  html: string;
+  head?: Array<React.ReactElement | null>;
+  styles?: React.ReactElement[] | Iterable<React.ReactNode> | React.ReactElement;
+};
+
+export type DocumentContext = {
+  pathname: string;
+  query: Record<string, string | string[] | undefined>;
+  asPath?: string;
+  req?: IncomingMessage;
+  res?: ServerResponse;
+  err?: (Error & { statusCode?: number }) | null;
+  locale?: string;
+  locales?: readonly string[];
+  defaultLocale?: string;
+  renderPage: () => DocumentInitialProps | Promise<DocumentInitialProps>;
+  defaultGetInitialProps(
+    ctx: DocumentContext,
+    options?: { nonce?: string },
+  ): Promise<DocumentInitialProps>;
+};
+
+// ─── Components ───────────────────────────────────────────────────────────────
 
 export function Html({
   children,
@@ -49,17 +80,36 @@ export function NextScript() {
   return <span dangerouslySetInnerHTML={{ __html: "<!-- __NEXT_SCRIPTS__ -->" }} />;
 }
 
+// ─── Base Document class ──────────────────────────────────────────────────────
+
 /**
- * Default Document component - used when no custom _document.tsx exists.
+ * Base Document class that custom _document.tsx classes extend.
+ *
+ * Provides a default getInitialProps implementation that mirrors Next.js:
+ * it calls ctx.defaultGetInitialProps(ctx), which in our shim just returns
+ * an empty html string (the dev server injects actual content separately).
+ *
+ * Custom documents can override getInitialProps to augment props:
+ *
+ *   static async getInitialProps(ctx: DocumentContext): Promise<DocumentProps> {
+ *     const initialProps = await Document.getInitialProps(ctx);
+ *     return { ...initialProps, theme: "light" };
+ *   }
  */
-export default function Document() {
-  return (
-    <Html>
-      <Head />
-      <body>
-        <Main />
-        <NextScript />
-      </body>
-    </Html>
-  );
+export default class Document<P = {}> extends React.Component<DocumentInitialProps & P> {
+  static async getInitialProps(ctx: DocumentContext): Promise<DocumentInitialProps> {
+    return ctx.defaultGetInitialProps(ctx);
+  }
+
+  render() {
+    return (
+      <Html>
+        <Head />
+        <body>
+          <Main />
+          <NextScript />
+        </body>
+      </Html>
+    );
+  }
 }

packages/vinext/src/shims/next-shims.d.ts

@@ -5,6 +5,48 @@
  * satisfies TypeScript when one shim imports another (e.g. link -> router).
  */
 
+declare module "next/document" {
+  import type { IncomingMessage, ServerResponse } from "node:http";
+  import { Component, type HTMLAttributes, type ReactNode, type ReactElement } from "react";
+
+  export type DocumentInitialProps = {
+    html: string;
+    head?: Array<ReactElement | null>;
+    styles?: ReactElement[] | Iterable<ReactNode> | ReactElement;
+  };
+
+  export type DocumentContext = {
+    pathname: string;
+    query: Record<string, string | string[] | undefined>;
+    asPath?: string;
+    req?: IncomingMessage;
+    res?: ServerResponse;
+    err?: (Error & { statusCode?: number }) | null;
+    locale?: string;
+    locales?: readonly string[];
+    defaultLocale?: string;
+    renderPage: () => DocumentInitialProps | Promise<DocumentInitialProps>;
+    defaultGetInitialProps(
+      ctx: DocumentContext,
+      options?: { nonce?: string },
+    ): Promise<DocumentInitialProps>;
+  };
+
+  export type DocumentProps = DocumentInitialProps & { [key: string]: unknown };
+
+  export function Html(
+    props: HTMLAttributes<HTMLHtmlElement> & { children?: ReactNode },
+  ): ReactElement;
+  export function Head(props: { children?: ReactNode }): ReactElement;
+  export function Main(): ReactElement;
+  export function NextScript(): ReactElement;
+
+  export default class Document<P = {}> extends Component<DocumentInitialProps & P> {
+    static getInitialProps(ctx: DocumentContext): Promise<DocumentInitialProps>;
+    render(): ReactElement;
+  }
+}
+
 declare module "next/router" {
   export function useRouter(): any;
   export function setSSRContext(ctx: any): void;

tests/e2e/pages-router/document.spec.ts

@@ -0,0 +1,11 @@
+import { test, expect } from "@playwright/test";
+
+const BASE = "http://localhost:4173";
+
+test.describe("Document", () => {
+  test("page includes theme attribute on the body", async ({ page }) => {
+    await page.goto(`${BASE}/`);
+
+    await expect(page.getAttribute("body", "data-theme-prop")).resolves.toBe("light");
+  });
+});

tests/fixtures/pages-basic/pages/_document.tsx

@@ -1,15 +1,26 @@
-import { Html, Head, Main, NextScript } from "next/document";
-
-export default function Document() {
-  return (
-    <Html lang="en">
-      <Head>
-        <meta name="description" content="A vinext test app" />
-      </Head>
-      <body className="custom-body">
-        <Main />
-        <NextScript />
-      </body>
-    </Html>
-  );
+import DocumentImpl, { Html, Head, Main, NextScript } from "next/document";
+import type { DocumentContext, DocumentInitialProps } from "next/document";
+
+type DocumentProps = DocumentInitialProps & { theme: string };
+
+export default class Document extends DocumentImpl<DocumentProps> {
+  static async getInitialProps(ctx: DocumentContext): Promise<DocumentProps> {
+    const initialProps = await DocumentImpl.getInitialProps(ctx);
+
+    return Promise.resolve({ ...initialProps, theme: "light" });
+  }
+
+  render() {
+    return (
+      <Html lang="en">
+        <Head>
+          <meta name="description" content="A vinext test app" />
+        </Head>
+        <body className="custom-body" data-theme-prop={this.props.theme}>
+          <Main />
+          <NextScript />
+        </body>
+      </Html>
+    );
+  }
 }