@assistant-ui/metro 0.0.1

package/src/transformer.test.ts

@@ -0,0 +1,101 @@
+import { describe, it, expect, beforeAll } from "vitest";
+import { fileURLToPath } from "node:url";
+import { dirname, join } from "node:path";
+import { withAui, UPSTREAM_TRANSFORMER_ENV } from "./index";
+
+const here = dirname(fileURLToPath(import.meta.url));
+const mockUpstream = join(here, "__fixtures__/mock-upstream.cjs");
+
+const GENERATIVE = `"use generative";
+import { defineToolkit } from "@assistant-ui/react";
+import { z } from "zod";
+import { db } from "@/db";
+export default defineToolkit({
+  weather: {
+    description: "Show the weather.",
+    parameters: z.object({ city: z.string() }),
+    execute: async ({ city }) => db.get(city),
+    render: () => null,
+  },
+});
+`;
+
+describe("withAui", () => {
+  it("points babelTransformerPath at our transformer, preserving the rest", () => {
+    const config = withAui({
+      projectRoot: "/app",
+      transformer: {
+        babelTransformerPath: "/up/stream.js",
+        assetPlugins: ["x"],
+      },
+      resolver: { sourceExts: ["ts"] },
+    });
+
+    expect(config.transformer?.babelTransformerPath).toContain("transformer");
+    expect(config.transformer?.babelTransformerPath).not.toBe("/up/stream.js");
+    expect(config.transformer?.assetPlugins).toEqual(["x"]);
+    expect(config.resolver).toEqual({ sourceExts: ["ts"] });
+    expect(config.projectRoot).toBe("/app");
+    // stashes the upstream so the worker can delegate to it
+    expect(process.env[UPSTREAM_TRANSFORMER_ENV]).toBe("/up/stream.js");
+  });
+
+  it("is idempotent: a double-wrap keeps the real upstream, not the wrapper", () => {
+    const once = withAui({
+      transformer: { babelTransformerPath: "/up/stream.js" },
+    });
+    const self = once.transformer?.babelTransformerPath;
+
+    const twice = withAui(once);
+    // still points at our transformer, and the env still names the real upstream
+    expect(twice.transformer?.babelTransformerPath).toBe(self);
+    expect(process.env[UPSTREAM_TRANSFORMER_ENV]).toBe("/up/stream.js");
+  });
+});
+
+describe("transformer", () => {
+  beforeAll(() => {
+    process.env[UPSTREAM_TRANSFORMER_ENV] = mockUpstream;
+  });
+
+  it("compiles a use-generative module to the client build by default", async () => {
+    const { transform } = await import("./transformer");
+    const out = transform({
+      filename: "/app/toolkit.tsx",
+      src: GENERATIVE,
+      options: {},
+    }) as { src: string };
+
+    // client build keeps render + schema, drops the backend execute + its import
+    expect(out.src).toContain("render");
+    expect(out.src).toContain("z.object");
+    expect(out.src).toContain('type: "backend"');
+    expect(out.src).not.toContain("db.get");
+    expect(out.src).not.toContain("@/db");
+  });
+
+  it("compiles to the server build for a node environment", async () => {
+    const { transform } = await import("./transformer");
+    const out = transform({
+      filename: "/app/toolkit.tsx",
+      src: GENERATIVE,
+      options: { customTransformOptions: { environment: "node" } },
+    }) as { src: string };
+
+    // server build keeps the backend execute, drops render
+    expect(out.src).toContain("db.get");
+    expect(out.src).not.toMatch(/render\s*:/);
+  });
+
+  it("passes non-generative modules through untouched", async () => {
+    const { transform } = await import("./transformer");
+    const src = `export const x = 1;\n`;
+    const out = transform({
+      filename: "/app/x.ts",
+      src,
+      options: {},
+    }) as { src: string };
+
+    expect(out.src).toBe(src);
+  });
+});

package/src/transformer.ts

@@ -0,0 +1,116 @@
+import { createRequire } from "node:module";
+import { statSync } from "node:fs";
+import { fileURLToPath } from "node:url";
+import {
+  compileGenerative,
+  isGenerativeModule,
+} from "@assistant-ui/x-generative-compiler";
+import { UPSTREAM_TRANSFORMER_ENV } from "./index";
+
+const require = createRequire(import.meta.url);
+
+/** Source modules a `"use generative"` directive can appear in. */
+const SOURCE_RE = /\.[cm]?[jt]sx?$/;
+
+/** Minimal shape of a Metro babel transformer. */
+type BabelTransformer = {
+  transform: (props: {
+    filename: string;
+    src: string;
+    options?:
+      | { customTransformOptions?: { environment?: string } | undefined }
+      | undefined;
+    [key: string]: unknown;
+  }) => unknown;
+  getCacheKey?: (() => string) | undefined;
+  [key: string]: unknown;
+};
+
+/**
+ * The project's upstream babel transformer: what Metro would have used without
+ * this wrapper. Taken from {@link UPSTREAM_TRANSFORMER_ENV} (set by `withAui`),
+ * else auto-detected: Expo's transformer, then React Native's.
+ */
+function resolveUpstream(): BabelTransformer {
+  const fromEnv = process.env[UPSTREAM_TRANSFORMER_ENV];
+  if (fromEnv) return require(fromEnv) as BabelTransformer;
+
+  for (const id of [
+    "@expo/metro-config/babel-transformer",
+    "@react-native/metro-babel-transformer",
+  ]) {
+    try {
+      return require(id) as BabelTransformer;
+    } catch {
+      // try the next candidate
+    }
+  }
+
+  throw new Error(
+    "[@assistant-ui/metro] Could not resolve an upstream Metro babel " +
+      "transformer. Make sure `transformer.babelTransformerPath` is set before " +
+      "`withAui(...)` (Expo's `getDefaultConfig` sets it), or install " +
+      "`@expo/metro-config` or `@react-native/metro-babel-transformer`.",
+  );
+}
+
+// Resolved lazily on first use, by which point Metro has loaded `metro.config` and
+// `withAui` has populated the env var, and forked workers have inherited it.
+let cachedUpstream: BabelTransformer | undefined;
+function upstreamTransformer(): BabelTransformer {
+  return (cachedUpstream ??= resolveUpstream());
+}
+
+/** Expo bundles `+api` routes / RSC for a server environment; the app is client. */
+function isServerEnvironment(environment: string | undefined): boolean {
+  return environment === "node" || environment === "react-server";
+}
+
+export function transform(
+  props: Parameters<BabelTransformer["transform"]>[0],
+): unknown {
+  const upstream = upstreamTransformer();
+  const { filename, src, options } = props;
+
+  if (SOURCE_RE.test(filename) && isGenerativeModule(src)) {
+    const target = isServerEnvironment(
+      options?.customTransformOptions?.environment,
+    )
+      ? "server"
+      : "client";
+
+    const { code } = compileGenerative(src, {
+      target,
+      filename,
+      // Metro has no `react-server` layer, so `import "server-only"` would not
+      // resolve; the environment split already keeps a server `execute` out of
+      // the client bundle.
+      injectServerOnly: false,
+    });
+
+    return upstream.transform({ ...props, src: code });
+  }
+
+  return upstream.transform(props);
+}
+
+// The bundled compiler lives in this file, so its own mtime moves whenever the
+// transform output could change (a newly published version, or an in-place
+// monorepo rebuild). Deriving the token from it busts Metro's cache
+// automatically.
+let cachedSelfToken: string | undefined;
+function selfCacheToken(): string {
+  if (cachedSelfToken !== undefined) return cachedSelfToken;
+  try {
+    cachedSelfToken = String(statSync(fileURLToPath(import.meta.url)).mtimeMs);
+  } catch {
+    cachedSelfToken = "0";
+  }
+  return cachedSelfToken;
+}
+
+export function getCacheKey(): string {
+  const upstream = upstreamTransformer();
+  const base = upstream.getCacheKey ? upstream.getCacheKey() : "";
+  return `${base}$aui-metro:${selfCacheToken()}`;
+}