Python: Patch workers.wait_until as well as ctx.waitUntil

src/pyodide/python-entrypoint-helper.ts

@@ -35,11 +35,34 @@ import {
 } from 'pyodide-internal:util';
 import { PyodideVersion } from 'pyodide-internal:const';
 import { default as introspectionSource } from 'pyodide-internal:introspection.py';
-export { createImportProxy } from 'pyodide-internal:serializeJsModule';
+import { createImportProxy } from 'pyodide-internal:serializeJsModule';
+export { createImportProxy };
 
 type PyFuture<T> = Promise<T> & { copy(): PyFuture<T>; destroy(): void };
 
 const waitUntilPatched = new WeakSet();
+function getPatchedWaitUntil(ctx: {
+  waitUntil: (p: Promise<void> | PyFuture<void>) => void;
+}): (p: Promise<void> | PyFuture<void>) => void {
+  const origWaitUntil: (p: Promise<void>) => void = ctx.waitUntil.bind(ctx);
+  function waitUntil(p: Promise<void> | PyFuture<void>): void {
+    origWaitUntil(
+      (async function (): Promise<void> {
+        if ('copy' in p) {
+          p = p.copy();
+        }
+        try {
+          await p;
+        } finally {
+          if ('destroy' in p) {
+            p.destroy();
+          }
+        }
+      })()
+    );
+  }
+  return waitUntil;
+}
 
 function patchWaitUntil(ctx: {
   waitUntil: (p: Promise<void> | PyFuture<void>) => void;
@@ -54,21 +77,7 @@ function patchWaitUntil(ctx: {
   if (waitUntilPatched.has(ctx)) {
     return;
   }
-  const origWaitUntil: (p: Promise<void>) => void = ctx.waitUntil.bind(ctx);
-  function waitUntil(p: Promise<void> | PyFuture<void>): void {
-    origWaitUntil(
-      (async function (): Promise<void> {
-        if ('copy' in p) {
-          p = p.copy();
-        }
-        await p;
-        if ('destroy' in p) {
-          p.destroy();
-        }
-      })()
-    );
-  }
-  ctx.waitUntil = waitUntil;
+  ctx.waitUntil = getPatchedWaitUntil(ctx);
   waitUntilPatched.add(ctx);
 }
 
@@ -93,13 +102,35 @@ function get_pyodide_entrypoint_helper(): PyodideEntrypointHelper {
   return _pyodide_entrypoint_helper;
 }
 
+async function getCloudflareWorkersModule(
+  doAnImport: (mod: string) => Promise<any>
+): Promise<{ env: any }> {
+  const cloudflareWorkersModule = await doAnImport('cloudflare:workers');
+  const waitUntil = createImportProxy(
+    'cloudflare:workers',
+    getPatchedWaitUntil(cloudflareWorkersModule),
+    ['waitUntil']
+  );
+  return new Proxy(cloudflareWorkersModule, {
+    get(_target: any, prop: string | symbol): any {
+      if (prop === 'waitUntil') {
+        return waitUntil;
+      }
+      // @ts-expect-error untyped Reflect.get
+      // eslint-disable-next-line prefer-rest-params
+      return Reflect.get(...arguments);
+    },
+  });
+}
+
 export async function setDoAnImport(
   doAnImport: (mod: string) => Promise<any>,
   workerEntrypoint: any
 ): Promise<void> {
+  const cloudflareWorkersModule = await getCloudflareWorkersModule(doAnImport);
   _pyodide_entrypoint_helper = {
     doAnImport,
-    cloudflareWorkersModule: await doAnImport('cloudflare:workers'),
+    cloudflareWorkersModule,
     cloudflareSocketsModule: await doAnImport('cloudflare:sockets'),
     workerEntrypoint,
     patchWaitUntil,
@@ -112,7 +143,12 @@ export async function setDoAnImport(
       throw new PythonWorkersInternalError(message);
     },
   };
-  await fillSnapshotJsModules(doAnImport);
+  await fillSnapshotJsModules(async (mod: string): Promise<any> => {
+    if (mod === 'cloudflare:workers') {
+      return cloudflareWorkersModule;
+    }
+    return await doAnImport(mod);
+  });
 }
 
 function handleSrcImport(pyodide: Pyodide, e: any): never {

src/workerd/server/tests/python/python-rpc/worker.py

@@ -10,14 +10,16 @@
 from unittest import TestCase
 
 import js
-from workers import Blob, Request, Response, WorkerEntrypoint, handler
+from workers import Blob, Request, Response, WorkerEntrypoint, handler, waitUntil
 
-from pyodide.ffi import JsException, JsProxy, to_js
+from pyodide.ffi import JsException, JsProxy, create_proxy, to_js
 
 assertRaises = TestCase().assertRaises
 assertRaisesRegex = TestCase().assertRaisesRegex
 
 testFuture = Future()
+globalWaitUntilFuture = Future()
+explicitProxyWaitUntilFuture = Future()
 
 
 class PythonRpcTester(WorkerEntrypoint):
@@ -66,6 +68,20 @@ async def sleep_then_set_result(self):
     async def test_wait_until_coroutine_lifetime(self):
         self.ctx.waitUntil(self.sleep_then_set_result())
 
+    async def explicit_proxy_wait_until(self):
+        async def set_result():
+            await sleep(0.1)
+            explicitProxyWaitUntilFuture.set_result(100)
+
+        self.ctx.waitUntil(create_proxy(set_result()))
+
+    async def sleep_then_set_global_wait_until_result(self):
+        await sleep(0.1)
+        globalWaitUntilFuture.set_result(100)
+
+    async def test_global_wait_until_coroutine_lifetime(self):
+        waitUntil(self.sleep_then_set_global_wait_until_result())
+
 
 class CustomType:
     def __init__(self, x):
@@ -280,3 +296,15 @@ def my_func():
     assert not testFuture.done()
     await sleep(0.2)
     assert testFuture.result() == 100
+
+    # Check that waitUntil() accepts a PyProxy explicitly created by the user.
+    await env.PythonRpc.explicit_proxy_wait_until()
+    assert not explicitProxyWaitUntilFuture.done()
+    await sleep(0.2)
+    assert explicitProxyWaitUntilFuture.result() == 100
+
+    # Check that the module-level waitUntil() also keeps its coroutine alive.
+    await env.PythonRpc.test_global_wait_until_coroutine_lifetime()
+    assert not globalWaitUntilFuture.done()
+    await sleep(0.2)
+    assert globalWaitUntilFuture.result() == 100