Fix CSS HMR for SSR (#85)

Since we used to build the HTML using our own `<Document>` component, we
were previously adding a data-turbopack-chunk-id attribute to our
`<link>` tags to reconcile chunk paths with their chunk ids when
initializing HMR. However, Next.js is now responsible for building the
HTML, and it has no such mechanism.

**NOTE:** HMR is currently broken for non-Next-SSR rendering
(HtmlAsset). This PR does not fix that.

crates/next-core/js/src/dev/client.ts

@@ -2,7 +2,9 @@ import { connect } from "./hmr-client";
 import { connectHMR } from "./websocket";
 
 export function initializeHMR(options: { assetPrefix: string }) {
-  connect();
+  connect({
+    assetPrefix: options.assetPrefix,
+  });
   connectHMR({
     path: "/turbopack-hmr",
     assetPrefix: options.assetPrefix,

crates/next-core/js/src/dev/hmr-client.ts

@@ -11,7 +11,11 @@ import { addEventListener, sendMessage } from "./websocket";
 
 declare var globalThis: TurbopackGlobals;
 
-export function connect() {
+export type ClientOptions = {
+  assetPrefix: string;
+};
+
+export function connect({ assetPrefix }: ClientOptions) {
   addEventListener((event) => {
     switch (event.type) {
       case "connected":
@@ -39,7 +43,7 @@ export function connect() {
     }
   }
 
-  subscribeToInitialCssChunksUpdates();
+  subscribeToInitialCssChunksUpdates(assetPrefix);
 }
 
 const chunkUpdateCallbacks: Map<string, ChunkUpdateCallback[]> = new Map();
@@ -102,15 +106,21 @@ function triggerChunkUpdate(update: ServerMessage) {
 
 // Unlike ES chunks, CSS chunks cannot contain the logic to accept updates.
 // They must be reloaded here instead.
-function subscribeToInitialCssChunksUpdates() {
+function subscribeToInitialCssChunksUpdates(assetPrefix: string) {
   const initialCssChunkLinks: NodeListOf<HTMLLinkElement> =
     document.head.querySelectorAll("link");
+  const cssChunkPrefix = `${assetPrefix}/`;
   initialCssChunkLinks.forEach((link) => {
-    if (!link.href) return;
-    const url = new URL(link.href);
-    if (url.origin !== location.origin) return;
-    const chunkPath = url.pathname.slice(1);
+    const href = link.href;
+    if (href == null) {
+      return;
+    }
+    const { pathname, origin } = new URL(href);
+    if (origin !== location.origin || !pathname.startsWith(cssChunkPrefix)) {
+      return;
+    }
 
+    const chunkPath = pathname.slice(cssChunkPrefix.length);
     onChunkUpdate(chunkPath, (update) => {
       switch (update.type) {
         case "restart": {

crates/next-core/js/src/entry/next-hydrate.tsx

@@ -2,30 +2,31 @@ import "@vercel/turbopack-next/internal/shims";
 
 import { initialize, hydrate } from "next/dist/client";
 import { initializeHMR } from "@vercel/turbopack-next/dev/client";
+import { displayContent } from "next/dist/client/dev/fouc";
 
 import * as _app from "@vercel/turbopack-next/pages/_app";
 import * as page from ".";
 
 (async () => {
   console.debug("Initializing Next.js");
 
-  initializeHMR({
-    assetPrefix: "",
-  });
-
-  await initialize({
+  const { assetPrefix } = await initialize({
     webpackHMR: {
       // Expected when `process.env.NODE_ENV === 'development'`
       onUnrecoverableError() {},
     },
   });
 
+  initializeHMR({
+    assetPrefix,
+  });
+
   window.__NEXT_P.push(["/_app", () => _app]);
   window.__NEXT_P.push([window.__NEXT_DATA__.page, () => page]);
 
   console.debug("Hydrating the page");
 
-  await hydrate();
+  await hydrate({ beforeRender: displayContent });
 
   console.debug("The page has been hydrated");
 })().catch((err) => console.error(err));

crates/next-core/js/src/entry/server-renderer.tsx

@@ -10,6 +10,8 @@ import Document from "@vercel/turbopack-next/pages/_document";
 import Component, * as otherExports from ".";
 ("TURBOPACK { transition: next-client }");
 import chunkGroup from ".";
+import { BuildManifest } from "next/dist/server/get-page-files";
+import { ChunkGroup } from "types/next";
 
 const END_OF_OPERATION = process.argv[2];
 const NEW_LINE = "\n".charCodeAt(0);
@@ -168,14 +170,14 @@ async function operation(renderData: RenderData) {
   // TODO(alexkirsz) This is missing *a lot* of data, but it's enough to get a
   // basic render working.
 
-  /* BuildManifest */
-  const buildManifest = {
+  const group = chunkGroup as ChunkGroup;
+  const buildManifest: BuildManifest = {
     pages: {
       // TODO(alexkirsz) We should separate _app and page chunks. Right now, we
       // computing the chunk items of `next-hydrate.js`, so they contain both
       // _app and page chunks.
       "/_app": [],
-      [renderData.path]: chunkGroup.map((c: { path: string }) => c.path),
+      [renderData.path]: group.map((chunk) => chunk.path),
     },
 
     devFiles: [],
@@ -202,6 +204,8 @@ async function operation(renderData: RenderData) {
     buildId: "",
 
     /* RenderOptsPartial */
+    dev: true,
+    runtimeConfig: {},
     assetPrefix: "",
     canonicalBase: "",
     previewProps: {

crates/next-core/js/types/next.d.ts

@@ -0,0 +1 @@
+export type ChunkGroup = Array<{ path: string; chunkId: string }>;

crates/turbopack-dev-server/src/update/stream.rs

@@ -29,7 +29,6 @@ async fn compute_update_stream(
 struct VersionState {
     #[turbo_tasks(debug_ignore)]
     inner: Mutex<(VersionVc, Option<Invalidator>)>,
-    id: VersionStateId,
 }
 
 #[turbo_tasks::value_impl]
@@ -43,23 +42,17 @@ impl VersionStateVc {
     }
 }
 
-#[turbo_tasks::value(transparent, serialization = "auto_for_input")]
-#[derive(Debug, PartialOrd, Ord, Hash, Clone)]
-struct VersionStateId(String);
-
 impl VersionStateVc {
-    async fn new(inner: VersionVc, chunk_path: &str) -> Result<Self> {
-        let id = VersionStateId(chunk_path.to_string());
-        let inner = inner.keyed_cell_local(id.clone()).await?;
+    async fn new(inner: VersionVc) -> Result<Self> {
+        let inner = inner.cell_local().await?;
         Ok(Self::cell(VersionState {
             inner: Mutex::new((inner, None)),
-            id,
         }))
     }
 
     async fn set(&self, new_inner: VersionVc) -> Result<()> {
         let this = self.await?;
-        let new_inner = new_inner.keyed_cell_local(this.id.clone()).await?;
+        let new_inner = new_inner.cell_local().await?;
         let mut lock = this.inner.lock().unwrap();
         if let (_, Some(invalidator)) = std::mem::replace(&mut *lock, (new_inner, None)) {
             invalidator.invalidate();
@@ -77,7 +70,7 @@ impl UpdateStream {
     pub async fn new(chunk_path: String, content: VersionedContentVc) -> Result<UpdateStream> {
         let (sx, rx) = tokio::sync::mpsc::channel(32);
 
-        let version_state = VersionStateVc::new(content.version(), &chunk_path).await?;
+        let version_state = VersionStateVc::new(content.version()).await?;
 
         compute_update_stream(
             version_state,

crates/turbopack-ecmascript/js/types/index.d.ts

@@ -4,7 +4,7 @@ import { Hot } from "./hot";
 
 export type RefreshHelpers = RefreshRuntimeGlobals["$RefreshHelpers$"];
 
-type ChunkId = string;
+type ChunkPath = string;
 type ModuleId = string;
 
 interface Chunk {}
@@ -18,7 +18,7 @@ interface Exports {
 export type ChunkModule = () => void;
 export type Runnable = (...args: any[]) => boolean;
 export declare type ChunkRegistration = [
-  chunkPath: string,
+  chunkPath: ChunkPath,
   chunkModules: ChunkModule[],
   ...run: Runnable[]
 ];
@@ -49,7 +49,7 @@ type EsmImport = (
 type EsmExport = (exportGetters: Record<string, () => any>) => void;
 type ExportValue = (value: any) => void;
 
-type LoadFile = (id: ChunkId, path: string) => Promise<any> | undefined;
+type LoadFile = (path: ChunkPath) => Promise<any> | undefined;
 
 interface TurbopackContext {
   e: Module["exports"];
@@ -72,7 +72,7 @@ type ModuleFactory = (
 type ModuleFactoryString = string;
 
 interface Runtime {
-  loadedChunks: Set<ChunkId>;
+  loadedChunks: Set<ChunkPath>;
   modules: Record<ModuleId, ModuleFactory>;
   cache: Record<string, Module>;
 
@@ -81,14 +81,14 @@ interface Runtime {
 
 export type ChunkUpdateCallback = (update: ServerMessage) => void;
 export type ChunkUpdateProvider = {
-  push: (registration: [ChunkId, ChunkUpdateCallback]) => void;
+  push: (registration: [ChunkPath, ChunkUpdateCallback]) => void;
 };
 
 export interface TurbopackGlobals {
   TURBOPACK?: ChunkRegistrations | ChunkRegistration[];
   TURBOPACK_CHUNK_UPDATE_LISTENERS?:
     | ChunkUpdateProvider
-    | [ChunkId, ChunkUpdateCallback][];
+    | [ChunkPath, ChunkUpdateCallback][];
 }
 
 declare global {

crates/turbopack-ecmascript/js/types/protocol.d.ts

@@ -1,7 +1,7 @@
-import { ChunkId } from "./index";
+import { ChunkPath } from "./index";
 
 export type ServerMessage = {
-  chunkId: ChunkId;
+  chunkPath: ChunkPath;
 } & (
   | {
       type: "restart";
@@ -17,5 +17,5 @@ export type ServerMessage = {
 
 export type ClientMessage = {
   type: "subscribe";
-  chunkId: ChunkId;
+  chunkPath: ChunkPath;
 };