refactor waku internals (#2115)

dai-shi · web-flow · commit d8d823a5330c · 2026-06-01T23:00:16.000+09:00
diff --git a/packages/waku/src/adapters/aws-lambda.ts b/packages/waku/src/adapters/aws-lambda.ts
@@ -8,11 +8,19 @@ import { unstable_createServerEntryAdapter as createServerEntryAdapter } from 'w
 import {
   unstable_constants as constants,
   unstable_honoMiddleware as honoMiddleware,
+  unstable_runWithContext as runWithContext,
 } from 'waku/internals';
 import type { BuildOptions } from './aws-lambda-build-enhancer.js';
 
 const { DIST_PUBLIC } = constants;
-const { contextMiddleware, rscMiddleware, middlewareRunner } = honoMiddleware;
+const { rscMiddleware, middlewareRunner } = honoMiddleware;
+
+function contextMiddleware(): MiddlewareHandler {
+  return (c, next) => {
+    const req = c.req.raw;
+    return runWithContext(req, next);
+  };
+}
 
 const DEFAULT_BODY_LIMIT_MAX_SIZE = 100 * 1024 * 1024;
 
diff --git a/packages/waku/src/adapters/bun.ts b/packages/waku/src/adapters/bun.ts
@@ -7,11 +7,19 @@ import { unstable_createServerEntryAdapter as createServerEntryAdapter } from 'w
 import {
   unstable_constants as constants,
   unstable_honoMiddleware as honoMiddleware,
+  unstable_runWithContext as runWithContext,
 } from 'waku/internals';
 import type { BuildOptions } from './bun-build-enhancer.js';
 
 const { DIST_PUBLIC } = constants;
-const { contextMiddleware, rscMiddleware, middlewareRunner } = honoMiddleware;
+const { rscMiddleware, middlewareRunner } = honoMiddleware;
+
+function contextMiddleware(): MiddlewareHandler {
+  return (c, next) => {
+    const req = c.req.raw;
+    return runWithContext(req, next);
+  };
+}
 
 const DEFAULT_BODY_LIMIT_MAX_SIZE = 100 * 1024 * 1024;
 
diff --git a/packages/waku/src/adapters/cloudflare.ts b/packages/waku/src/adapters/cloudflare.ts
@@ -10,11 +10,19 @@ import {
   unstable_consumeMultiplexedStream as consumeMultiplexedStream,
   unstable_honoMiddleware as honoMiddleware,
   unstable_produceMultiplexedStream as produceMultiplexedStream,
+  unstable_runWithContext as runWithContext,
 } from 'waku/internals';
 import type { BuildOptions } from './cloudflare-build-enhancer.js';
 
 const { DIST_PUBLIC } = constants;
-const { contextMiddleware, rscMiddleware, middlewareRunner } = honoMiddleware;
+const { rscMiddleware, middlewareRunner } = honoMiddleware;
+
+function contextMiddleware(): MiddlewareHandler {
+  return (c, next) => {
+    const req = c.req.raw;
+    return runWithContext(req, next);
+  };
+}
 
 const DEFAULT_BODY_LIMIT_MAX_SIZE = 100 * 1024 * 1024;
 
diff --git a/packages/waku/src/adapters/deno.ts b/packages/waku/src/adapters/deno.ts
@@ -7,11 +7,19 @@ import { unstable_createServerEntryAdapter as createServerEntryAdapter } from 'w
 import {
   unstable_constants as constants,
   unstable_honoMiddleware as honoMiddleware,
+  unstable_runWithContext as runWithContext,
 } from 'waku/internals';
 import type { BuildOptions } from './deno-build-enhancer.js';
 
 const { DIST_PUBLIC } = constants;
-const { contextMiddleware, rscMiddleware, middlewareRunner } = honoMiddleware;
+const { rscMiddleware, middlewareRunner } = honoMiddleware;
+
+function contextMiddleware(): MiddlewareHandler {
+  return (c, next) => {
+    const req = c.req.raw;
+    return runWithContext(req, next);
+  };
+}
 
 const DEFAULT_BODY_LIMIT_MAX_SIZE = 100 * 1024 * 1024;
 
diff --git a/packages/waku/src/adapters/edge.ts b/packages/waku/src/adapters/edge.ts
@@ -3,15 +3,25 @@ import { bodyLimit } from 'hono/body-limit';
 import { Hono } from 'hono/tiny';
 import type { ImportGlobFunction } from 'vite/types/importGlob.d.ts';
 import { unstable_createServerEntryAdapter as createServerEntryAdapter } from 'waku/adapter-builders';
-import { unstable_honoMiddleware as honoMiddleware } from 'waku/internals';
+import {
+  unstable_honoMiddleware as honoMiddleware,
+  unstable_runWithContext as runWithContext,
+} from 'waku/internals';
 
 declare global {
   interface ImportMeta {
     glob: ImportGlobFunction;
   }
 }
 
-const { contextMiddleware, rscMiddleware, middlewareRunner } = honoMiddleware;
+const { rscMiddleware, middlewareRunner } = honoMiddleware;
+
+function contextMiddleware(): MiddlewareHandler {
+  return (c, next) => {
+    const req = c.req.raw;
+    return runWithContext(req, next);
+  };
+}
 
 const DEFAULT_BODY_LIMIT_MAX_SIZE = 100 * 1024 * 1024;
 
diff --git a/packages/waku/src/adapters/netlify.ts b/packages/waku/src/adapters/netlify.ts
@@ -5,11 +5,19 @@ import { unstable_createServerEntryAdapter as createServerEntryAdapter } from 'w
 import {
   unstable_constants as constants,
   unstable_honoMiddleware as honoMiddleware,
+  unstable_runWithContext as runWithContext,
 } from 'waku/internals';
 import type { BuildOptions } from './netlify-build-enhancer.js';
 
 const { DIST_PUBLIC } = constants;
-const { contextMiddleware, rscMiddleware, middlewareRunner } = honoMiddleware;
+const { rscMiddleware, middlewareRunner } = honoMiddleware;
+
+function contextMiddleware(): MiddlewareHandler {
+  return (c, next) => {
+    const req = c.req.raw;
+    return runWithContext(req, next);
+  };
+}
 
 const DEFAULT_BODY_LIMIT_MAX_SIZE = 100 * 1024 * 1024;
 
diff --git a/packages/waku/src/adapters/node.ts b/packages/waku/src/adapters/node.ts
@@ -8,11 +8,19 @@ import { unstable_createServerEntryAdapter as createServerEntryAdapter } from 'w
 import {
   unstable_constants as constants,
   unstable_honoMiddleware as honoMiddleware,
+  unstable_runWithContext as runWithContext,
 } from 'waku/internals';
 import type { BuildOptions } from './node-build-enhancer.js';
 
 const { DIST_PUBLIC } = constants;
-const { contextMiddleware, rscMiddleware, middlewareRunner } = honoMiddleware;
+const { rscMiddleware, middlewareRunner } = honoMiddleware;
+
+function contextMiddleware(): MiddlewareHandler {
+  return (c, next) => {
+    const req = c.req.raw;
+    return runWithContext(req, next);
+  };
+}
 
 const DEFAULT_BODY_LIMIT_MAX_SIZE = 100 * 1024 * 1024;
 
diff --git a/packages/waku/src/adapters/vercel.ts b/packages/waku/src/adapters/vercel.ts
@@ -6,11 +6,19 @@ import { unstable_createServerEntryAdapter as createServerEntryAdapter } from 'w
 import {
   unstable_constants as constants,
   unstable_honoMiddleware as honoMiddleware,
+  unstable_runWithContext as runWithContext,
 } from 'waku/internals';
 import type { BuildOptions } from './vercel-build-enhancer.js';
 
 const { DIST_PUBLIC } = constants;
-const { contextMiddleware, rscMiddleware, middlewareRunner } = honoMiddleware;
+const { rscMiddleware, middlewareRunner } = honoMiddleware;
+
+function contextMiddleware(): MiddlewareHandler {
+  return (c, next) => {
+    const req = c.req.raw;
+    return runWithContext(req, next);
+  };
+}
 
 const DEFAULT_BODY_LIMIT_MAX_SIZE = 100 * 1024 * 1024;
 (global as any).__WAKU_HONO_NODE_SERVER_GET_REQUEST_LISTENER__ =
diff --git a/packages/waku/src/internals.ts b/packages/waku/src/internals.ts
@@ -5,3 +5,4 @@ export {
   produceMultiplexedStream as unstable_produceMultiplexedStream,
   consumeMultiplexedStream as unstable_consumeMultiplexedStream,
 } from './lib/utils/stream.js';
+export { runWithContext as unstable_runWithContext } from './lib/context.js';
diff --git a/packages/waku/src/lib/context.ts b/packages/waku/src/lib/context.ts
@@ -8,7 +8,10 @@ type Context = {
 
 const contextStorage = new AsyncLocalStorage<Context>();
 
-export function INTERNAL_runWithContext<T>(req: Request, next: () => T): T {
+/**
+ * This is an internal function and not for public use.
+ */
+export function runWithContext<T>(req: Request, next: () => T): T {
   const context: Context = {
     req,
     nonce: undefined,
@@ -20,9 +23,7 @@ export function INTERNAL_runWithContext<T>(req: Request, next: () => T): T {
 export function getContext() {
   const context = contextStorage.getStore();
   if (!context) {
-    throw new Error(
-      'Context is not available. Make sure to use the context middleware.',
-    );
+    throw new Error('Context is not available.');
   }
   return context;
 }
diff --git a/packages/waku/src/lib/env.ts b/packages/waku/src/lib/env.ts
@@ -0,0 +1,13 @@
+// The use of `globalThis` in this file is more or less a hack.
+// It should be revisited with a better solution.
+
+/**
+ * This is an internal function and not for public use.
+ */
+export function setAllEnv(newEnv: Readonly<Record<string, string>>) {
+  (globalThis as any).__WAKU_SERVER_ENV__ = newEnv;
+}
+
+export function getEnv(key: string): string | undefined {
+  return (globalThis as any).__WAKU_SERVER_ENV__?.[key];
+}
diff --git a/packages/waku/src/lib/hono/middleware.ts b/packages/waku/src/lib/hono/middleware.ts
@@ -1,15 +1,7 @@
 import type { MiddlewareHandler } from 'hono';
 import type { Hono } from 'hono/tiny';
-import { INTERNAL_runWithContext } from '../context.js';
 import type { Unstable_ProcessRequest as ProcessRequest } from '../types.js';
 
-export function contextMiddleware(): MiddlewareHandler {
-  return (c, next) => {
-    const req = c.req.raw;
-    return INTERNAL_runWithContext(req, next);
-  };
-}
-
 export function rscMiddleware({
   processRequest,
 }: {
diff --git a/packages/waku/src/lib/vite-entries/entry.build.ts b/packages/waku/src/lib/vite-entries/entry.build.ts
@@ -1,5 +1,5 @@
 import serverEntry from 'virtual:vite-rsc-waku/server-entry';
-import { INTERNAL_setAllEnv } from '../../server.js';
+import { setAllEnv } from '../env.js';
 import type { Unstable_EmitFile } from '../types.js';
 import { joinPath } from '../utils/path.js';
 
@@ -31,7 +31,7 @@ export async function INTERNAL_runBuild({
   rootDir: string;
   emitFile: Unstable_EmitFile;
 }) {
-  INTERNAL_setAllEnv(process.env as any);
+  setAllEnv(process.env as any);
   const prunableFiles = new Set<string>();
   let build = serverEntry.build;
   for (const enhancer of serverEntry.buildEnhancers || []) {
diff --git a/packages/waku/src/lib/vite-entries/entry.server.tsx b/packages/waku/src/lib/vite-entries/entry.server.tsx
@@ -1,5 +1,5 @@
 import serverEntry from 'virtual:vite-rsc-waku/server-entry';
-import { INTERNAL_setAllEnv } from '../../server.js';
+import { setAllEnv } from '../env.js';
 
 export { serverEntry as unstable_serverEntry };
 
@@ -8,7 +8,7 @@ export async function INTERNAL_runFetch(
   req: Request,
   ...args: any[]
 ) {
-  INTERNAL_setAllEnv(env);
+  setAllEnv(env);
   return serverEntry.fetch(req, ...args);
 }
 
diff --git a/packages/waku/src/lib/vite-rsc/handler.ts b/packages/waku/src/lib/vite-rsc/handler.ts
@@ -10,9 +10,9 @@ import {
 import { buildMetadata } from 'virtual:vite-rsc-waku/build-metadata';
 import { config, isBuild } from 'virtual:vite-rsc-waku/config';
 import notFoundHtml from 'virtual:vite-rsc-waku/not-found';
-import { INTERNAL_setAllEnv } from '../../server.js';
 import { BUILD_METADATA_FILE, DIST_PUBLIC, DIST_SERVER } from '../constants.js';
-import { INTERNAL_runWithContext } from '../context.js';
+import { runWithContext } from '../context.js';
+import { setAllEnv } from '../env.js';
 import type {
   Unstable_CreateServerEntryAdapter as CreateServerEntryAdapter,
   Unstable_HandleBuild as HandleBuild,
@@ -144,7 +144,7 @@ const toProcessBuild =
       saveBuildMetadata: async (key, value) => {
         buildMetadata.set(key, value);
       },
-      withRequest: (req, fn) => INTERNAL_runWithContext(req, fn),
+      withRequest: (req, fn) => runWithContext(req, fn),
       generateFile: async (fileName, body) => {
         await emitFile(
           joinPath(DIST_PUBLIC, fileName),
@@ -178,7 +178,7 @@ export const createServerEntryAdapter: CreateServerEntryAdapter =
         handlers,
         processRequest,
         processBuild,
-        setAllEnv: INTERNAL_setAllEnv,
+        setAllEnv,
         config,
         isBuild,
         notFoundHtml,
diff --git a/packages/waku/src/server.ts b/packages/waku/src/server.ts
@@ -8,19 +8,7 @@ export {
   getContextData as unstable_getContextData,
 } from './lib/context.js';
 
-// The use of `globalThis` in this file is more or less a hack.
-// It should be revisited with a better solution.
-
-/**
- * This is an internal function and not for public use.
- */
-export function INTERNAL_setAllEnv(newEnv: Readonly<Record<string, string>>) {
-  (globalThis as any).__WAKU_SERVER_ENV__ = newEnv;
-}
-
-export function getEnv(key: string): string | undefined {
-  return (globalThis as any).__WAKU_SERVER_ENV__?.[key];
-}
+export { getEnv } from './lib/env.js';
 
 export function unstable_getHeaders(): Readonly<Record<string, string>> {
   return Object.fromEntries(getContext().req.headers.entries());
