Implement withEnv to propagate modified workers env

src/cloudflare/internal/env.d.ts

@@ -1,2 +1,3 @@
 // Get the current environment, if any
 export function getCurrent(): object | undefined;
+export function withEnv(newEnv: any, fn: () => any): any;

src/cloudflare/workers.ts

@@ -19,6 +19,11 @@ export const WorkflowEntrypoint = entrypoints.WorkflowEntrypoint;
 // is intentional.
 /* eslint-disable @typescript-eslint/no-unsafe-assignment */
 /* eslint-disable @typescript-eslint/no-explicit-any */
+
+export function withEnv(newEnv: any, fn: () => any): any {
+  return innerEnv.withEnv(newEnv, fn);
+}
+
 // A proxy for the workers env/bindings. The proxy is io-context
 // aware in that it will only return values when there is an active
 // IoContext. Mutations to the env via this proxy propagate to the

src/workerd/api/modules.h

@@ -38,8 +38,15 @@ class EnvModule final: public jsg::Object {
     return kj::none;
   }
 
+  jsg::JsValue withEnv(jsg::Lock& js, jsg::Value newEnv, jsg::Function<jsg::JsValue()> fn) {
+    auto& key = jsg::IsolateBase::from(js.v8Isolate).getEnvAsyncContextKey();
+    jsg::AsyncContextFrame::StorageScope storage(js, key, kj::mv(newEnv));
+    return fn(js);
+  }
+
   JSG_RESOURCE_TYPE(EnvModule) {
     JSG_METHOD(getCurrent);
+    JSG_METHOD(withEnv);
   }
 };
 

src/workerd/api/tests/importable-env-test.js

@@ -1,11 +1,36 @@
-import { strictEqual, deepStrictEqual, notStrictEqual } from 'node:assert';
-import { env } from 'cloudflare:workers';
+import {
+  strictEqual,
+  deepStrictEqual,
+  notStrictEqual,
+  ok,
+  rejects,
+} from 'node:assert';
+import { env, withEnv } from 'cloudflare:workers';
 
 // The env is populated at the top level scope.
 strictEqual(env.FOO, 'BAR');
 
+// Cache exists and is accessible.
+ok(env.CACHE);
+// But fails when used because we're not in an io-context
+await rejects(
+  env.CACHE.read('hello', async () => {}),
+  {
+    message: /Disallowed operation called within global scope./,
+  }
+);
+
 export const importableEnv = {
   async test(_, argEnv) {
+    // Accessing the cache initially at the global scope didn't break anything
+    const cached = await argEnv.CACHE.read('hello', async () => {
+      return {
+        value: 123,
+        expiration: Date.now() + 10000,
+      };
+    });
+    strictEqual(cached, 123);
+
     // They aren't the same objects...
     notStrictEqual(env, argEnv);
     // But have all the same stuff...
@@ -23,7 +48,17 @@ export const importableEnv = {
     const { env: otherEnv } = await import('child');
     strictEqual(otherEnv.FOO, 'BAR');
     strictEqual(otherEnv.BAR, 123);
-    strictEqual(otherEnv, env);
     deepStrictEqual(argEnv, otherEnv);
+
+    // Using withEnv to replace the env...
+    const { env: otherEnv2 } = await withEnv({ BAZ: 1 }, async () => {
+      await scheduler.wait(0);
+      return import('child2');
+    });
+    strictEqual(otherEnv2.FOO, undefined);
+    strictEqual(otherEnv2.BAZ, 1);
+
+    // Original env is unmodified
+    strictEqual(env.BAZ, undefined);
   },
 };

src/workerd/api/tests/importable-env-test.wd-test

@@ -6,7 +6,8 @@ const unitTests :Workerd.Config = (
       worker = (
         modules = [
           (name = "worker", esModule = embed "importable-env-test.js"),
-          (name = "child", esModule = "export {env} from 'cloudflare:workers'"),
+          (name = "child", esModule = "import {env as live} from 'cloudflare:workers'; export const env = {...live};"),
+          (name = "child2", esModule = "import {env as live} from 'cloudflare:workers'; export const env = {...live};"),
         ],
         compatibilityDate = "2025-02-01",
         compatibilityFlags = [
@@ -15,6 +16,14 @@ const unitTests :Workerd.Config = (
         ],
         bindings = [
           (name = "FOO", text = "BAR"),
+          (name = "CACHE", memoryCache = (
+            id = "abc123",
+            limits = (
+              maxKeys = 10,
+              maxValueSize = 1024,
+              maxTotalValueSize = 1024,
+            ),
+          ))
         ],
       )
     ),