Merge pull request #82 from unstubbable/server-action-source-maps

Source map server actions to their server location
unstubbable · Sep 26, 2024 · c708f51 · c708f51
2 parents 59cee12 + 4e683b7
commit c708f51
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";`,
     );
   });