Source map server actions to their server location

This enables right-clicking on an action and clicking "Go to definition" in the React DevTools. (The serving of the source map files is currently only implemented in the AWS example app's dev server.)
unstubbable · Sep 26, 2024 · 4e683b7 · 4e683b7
1 parent 59cee12
commit 4e683b7
Show file tree

Hide file tree

Showing 8 changed files with 237 additions and 89 deletions.
diff --git a/apps/aws-app/dev-server/run.ts b/apps/aws-app/dev-server/run.ts
@@ -1,3 +1,6 @@
+import fs from 'fs/promises';
+import path from 'path';
+import url from 'url';
 import {serve} from '@hono/node-server';
 import {serveStatic} from '@hono/node-server/serve-static';
 import {Hono} from 'hono';
@@ -11,6 +14,29 @@ const app = new Hono();
 
 app.use(authMiddleware);
 app.use(`/client/*`, serveStatic({root: `dist/static`}));
+
+app.get(`/source-maps`, async (context) => {
+  const filenameQueryParam = context.req.query(`filename`);
+
+  if (!filenameQueryParam) {
+    return context.newResponse(`Missing query parameter "filename"`, 400);
+  }
+
+  const filename = filenameQueryParam.startsWith(`file://`)
+    ? url.fileURLToPath(filenameQueryParam)
+    : path.join(import.meta.dirname, `../dist/static`, filenameQueryParam);
+
+  try {
+    const sourceMapFilename = `${filename}.map`;
+    const sourceMapContents = await fs.readFile(sourceMapFilename);
+
+    return context.newResponse(sourceMapContents);
+  } catch (error) {
+    console.error(error);
+    return context.newResponse(null, 404);
+  }
+});
+
 app.route(`/`, handlerApp);
 
 const server = serve({fetch: app.fetch, port: 3002}, ({address, port}) => {

diff --git a/apps/shared-app/src/client/button.tsx b/apps/shared-app/src/client/button.tsx
@@ -1,13 +1,17 @@
 'use client';
 
 import * as React from 'react';
-import {trackClick} from '../server/track-click.js';
 
 export type ButtonProps = React.PropsWithChildren<{
   readonly disabled?: boolean;
+  readonly trackClick: () => Promise<number>;
 }>;
 
-export function Button({children, disabled}: ButtonProps): React.ReactNode {
+export function Button({
+  children,
+  disabled,
+  trackClick,
+}: ButtonProps): React.ReactNode {
   return (
     <button
       onClick={() => void trackClick()}

diff --git a/apps/shared-app/src/client/product.tsx b/apps/shared-app/src/client/product.tsx
@@ -3,6 +3,7 @@
 import {clsx} from 'clsx';
 import * as React from 'react';
 import type {BuyResult} from '../server/buy.js';
+import {trackClick} from '../server/track-click.js';
 import {Notification} from '../shared/notification.js';
 import {Button} from './button.js';
 
@@ -43,7 +44,9 @@ export function Product({buy}: ProductProps): React.ReactNode {
         )}
       />
       {` `}
-      <Button disabled={isPending}>Buy now</Button>
+      <Button disabled={isPending} trackClick={trackClick}>
+        Buy now
+      </Button>
       {result && (
         <Notification status={result.status}>
           {result.status === `success` ? (

diff --git a/packages/core/src/client/hydrate-app.tsx b/packages/core/src/client/hydrate-app.tsx
@@ -12,11 +12,19 @@ export interface RscAppResult {
   readonly formState?: ReactFormState;
 }
 
+const originRegExp = new RegExp(`^${document.location.origin}`);
+
+export function findSourceMapUrl(filename: string): string | null {
+  return `${document.location.origin}/source-maps?filename=${encodeURIComponent(
+    filename.replace(originRegExp, ``),
+  )}`;
+}
+
 export async function hydrateApp(): Promise<void> {
   const {root: initialRoot, formState} =
     await ReactServerDOMClient.createFromReadableStream<RscAppResult>(
       self.initialRscResponseStream,
-      {callServer},
+      {callServer, findSourceMapURL: findSourceMapUrl},
     );
 
   const initialUrlPath = createUrlPath(document.location);
@@ -25,7 +33,7 @@ export async function hydrateApp(): Promise<void> {
     async function fetchRoot(urlPath: string): Promise<React.ReactElement> {
       const {root} = await ReactServerDOMClient.createFromFetch<RscAppResult>(
         fetch(urlPath, {headers: {accept: `text/x-component`}}),
-        {callServer},
+        {callServer, findSourceMapURL: findSourceMapUrl},
       );
 
       return root;

diff --git a/packages/webpack-rsc/src/webpack-rsc-client-loader.cts b/packages/webpack-rsc/src/webpack-rsc-client-loader.cts
@@ -14,6 +14,11 @@ namespace webpackRscClientLoader {
 
 type SourceMap = Parameters<LoaderDefinitionFunction>[1];
 
+interface FunctionInfo {
+  readonly exportName: string;
+  readonly loc: t.SourceLocation | null | undefined;
+}
+
 function webpackRscClientLoader(
   this: LoaderContext<webpackRscClientLoader.WebpackRscClientLoaderOptions>,
   source: string,
@@ -35,80 +40,101 @@ function webpackRscClientLoader(
     plugins: [`importAssertions`],
   });
 
+  let moduleId: string | number | undefined;
   let hasUseServerDirective = false;
+  let addedRegisterServerReferenceCall = false;
+  const importNodes = new Set<t.Node>();
 
   traverse.default(ast, {
-    Program(path) {
+    enter(path) {
       const {node} = path;
 
-      if (!node.directives.some(isUseServerDirective)) {
-        return;
-      }
+      if (t.isProgram(node)) {
+        if (node.directives.some(isUseServerDirective)) {
+          hasUseServerDirective = true;
 
-      hasUseServerDirective = true;
+          const moduleInfo = serverReferencesMap.get(resourcePath);
 
-      const moduleInfo = serverReferencesMap.get(resourcePath);
+          if (!moduleInfo) {
+            loaderContext.emitError(
+              new Error(
+                `Could not find server references module info in \`serverReferencesMap\` for ${resourcePath}.`,
+              ),
+            );
 
-      if (!moduleInfo) {
-        loaderContext.emitError(
-          new Error(
-            `Could not find server references module info in \`serverReferencesMap\` for ${resourcePath}.`,
-          ),
-        );
+            path.replaceWith(t.program([]));
+          } else if (!moduleInfo.moduleId) {
+            loaderContext.emitError(
+              new Error(
+                `Could not find server references module ID in \`serverReferencesMap\` for ${resourcePath}.`,
+              ),
+            );
 
-        path.replaceWith(t.program([]));
+            path.replaceWith(t.program([]));
+          } else {
+            moduleId = moduleInfo.moduleId;
+          }
+        } else {
+          path.skip();
+        }
 
         return;
       }
 
-      const {moduleId, exportNames} = moduleInfo;
+      if (importNodes.has(node)) {
+        return path.skip();
+      }
 
-      if (!moduleId) {
-        loaderContext.emitError(
-          new Error(
-            `Could not find server references module ID in \`serverReferencesMap\` for ${resourcePath}.`,
-          ),
-        );
+      const functionInfo = getFunctionInfo(node);
 
-        path.replaceWith(t.program([]));
+      if (moduleId && functionInfo) {
+        path.replaceWith(
+          createNamedExportedServerReference(functionInfo, moduleId),
+        );
+        path.skip();
+        addedRegisterServerReferenceCall = true;
+      } else {
+        path.remove();
+      }
+    },
+    exit(path) {
+      if (!t.isProgram(path.node) || !addedRegisterServerReferenceCall) {
+        path.skip();
 
         return;
       }
 
-      path.replaceWith(
-        t.program([
-          t.importDeclaration(
-            [
-              t.importSpecifier(
-                t.identifier(`createServerReference`),
-                t.identifier(`createServerReference`),
-              ),
-            ],
-            t.stringLiteral(`react-server-dom-webpack/client`),
-          ),
-          t.importDeclaration(
-            [
-              t.importSpecifier(
-                t.identifier(`callServer`),
-                t.identifier(`callServer`),
-              ),
-            ],
-            t.stringLiteral(callServerImportSource),
-          ),
-          ...exportNames.map((exportName) =>
-            t.exportNamedDeclaration(
-              t.variableDeclaration(`const`, [
-                t.variableDeclarator(
-                  t.identifier(exportName),
-                  t.callExpression(t.identifier(`createServerReference`), [
-                    t.stringLiteral(`${moduleId}#${exportName}`),
-                    t.identifier(`callServer`),
-                  ]),
-                ),
-              ]),
+      importNodes.add(
+        t.importDeclaration(
+          [
+            t.importSpecifier(
+              t.identifier(`createServerReference`),
+              t.identifier(`createServerReference`),
             ),
-          ),
-        ]),
+          ],
+          t.stringLiteral(`react-server-dom-webpack/client`),
+        ),
+      );
+
+      importNodes.add(
+        t.importDeclaration(
+          [
+            t.importSpecifier(
+              t.identifier(`callServer`),
+              t.identifier(`callServer`),
+            ),
+            t.importSpecifier(
+              t.identifier(`findSourceMapUrl`),
+              t.identifier(`findSourceMapUrl`),
+            ),
+          ],
+          t.stringLiteral(callServerImportSource),
+        ),
+      );
+
+      (path as traverse.NodePath<t.Program>).unshiftContainer(
+        `body`,
+        Array.from(importNodes),
       );
     },
   });
@@ -117,11 +143,43 @@ function webpackRscClientLoader(
     return this.callback(null, source, sourceMap);
   }
 
-  // TODO: Handle source maps.
+  const {code, map} = generate.default(
+    ast,
+    {
+      sourceFileName: this.resourcePath,
+      sourceMaps: this.sourceMap,
+      // @ts-expect-error
+      inputSourceMap: sourceMap,
+    },
+    source,
+  );
 
-  const {code} = generate.default(ast, {sourceFileName: this.resourcePath});
+  this.callback(null, code, map ?? sourceMap);
+}
 
-  this.callback(null, code);
+function createNamedExportedServerReference(
+  functionInfo: FunctionInfo,
+  moduleId: string | number,
+) {
+  const {exportName, loc} = functionInfo;
+  const exportIdentifier = t.identifier(exportName);
+
+  exportIdentifier.loc = loc;
+
+  return t.exportNamedDeclaration(
+    t.variableDeclaration(`const`, [
+      t.variableDeclarator(
+        exportIdentifier,
+        t.callExpression(t.identifier(`createServerReference`), [
+          t.stringLiteral(`${moduleId}#${exportName}`),
+          t.identifier(`callServer`),
+          t.identifier(`undefined`), // encodeFormAction
+          t.identifier(`findSourceMapUrl`),
+          t.stringLiteral(exportName),
+        ]),
+      ),
+    ]),
+  );
 }
 
 function isUseServerDirective(directive: t.Directive): boolean {
@@ -131,4 +189,33 @@ function isUseServerDirective(directive: t.Directive): boolean {
   );
 }
 
+function getFunctionInfo(node: t.Node): FunctionInfo | undefined {
+  let exportName: string | undefined;
+  let loc: t.SourceLocation | null | undefined;
+
+  if (t.isExportNamedDeclaration(node)) {
+    if (t.isFunctionDeclaration(node.declaration)) {
+      exportName = node.declaration.id?.name;
+      loc = node.declaration.id?.loc;
+    } else if (t.isVariableDeclaration(node.declaration)) {
+      const declarator = node.declaration.declarations[0];
+
+      if (!declarator) {
+        return undefined;
+      }
+
+      if (
+        (t.isFunctionExpression(declarator.init) ||
+          t.isArrowFunctionExpression(declarator.init)) &&
+        t.isIdentifier(declarator.id)
+      ) {
+        exportName = declarator.id.name;
+        loc = declarator.id.loc;
+      }
+    }
+  }
+
+  return exportName ? {exportName, loc} : undefined;
+}
+
 export = webpackRscClientLoader;
diff --git a/packages/webpack-rsc/src/webpack-rsc-client-loader.test.ts b/packages/webpack-rsc/src/webpack-rsc-client-loader.test.ts
@@ -54,17 +54,18 @@ describe(`webpackRscClientLoader`, () => {
     );
 
     const serverReferencesMap: ServerReferencesMap = new Map([
-      [resourcePath, {moduleId: `test`, exportNames: [`foo`, `bar`]}],
+      [resourcePath, {moduleId: `test`, exportNames: []}],
     ]);
 
     const output = await callLoader(resourcePath, {serverReferencesMap});
 
     expect(output).toEqual(
       `
 import { createServerReference } from "react-server-dom-webpack/client";
-import { callServer } from "@mfng/core/client/browser";
-export const foo = createServerReference("test#foo", callServer);
-export const bar = createServerReference("test#bar", callServer);
+import { callServer, findSourceMapUrl } from "@mfng/core/client/browser";
+export const foo = createServerReference("test#foo", callServer, undefined, findSourceMapUrl, "foo");
+export const bar = createServerReference("test#bar", callServer, undefined, findSourceMapUrl, "bar");
+export const baz = createServerReference("test#baz", callServer, undefined, findSourceMapUrl, "baz");
 `.trim(),
     );
   });
@@ -86,12 +87,8 @@ export const bar = createServerReference("test#bar", callServer);
       callServerImportSource,
     });
 
-    expect(output).toEqual(
-      `
-import { createServerReference } from "react-server-dom-webpack/client";
-import { callServer } from "some-router/call-server";
-export const foo = createServerReference("test#foo", callServer);
-`.trim(),
+    expect(output).toMatch(
+      `import { callServer, findSourceMapUrl } from "some-router/call-server";`,
     );
   });