From 9fa7dd59a6673dd014741531b9a226047b46740d Mon Sep 17 00:00:00 2001 From: Hood Chatham Date: Wed, 9 Sep 2026 10:04:46 -0700 Subject: [PATCH 1/2] Python: Patch `workers.wait_until` as well as `ctx.waitUntil` We've been patching `ctx.waitUntil()` so that when it is passed a Python awaitable it will keep it alive until it resolves. We forgot to patch `workers.wait_until()` in the same way. This fixes the problem. Thanks to audreyfeldroy for reporting this in cloudflare/workers-py # 249. --- src/pyodide/python-entrypoint-helper.ts | 67 ++++++++++++++----- .../server/tests/python/python-rpc/worker.py | 32 ++++++++- 2 files changed, 80 insertions(+), 19 deletions(-) diff --git a/src/pyodide/python-entrypoint-helper.ts b/src/pyodide/python-entrypoint-helper.ts index e06b202576f..3f7f584bec8 100644 --- a/src/pyodide/python-entrypoint-helper.ts +++ b/src/pyodide/python-entrypoint-helper.ts @@ -35,25 +35,15 @@ 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 = Promise & { copy(): PyFuture; destroy(): void }; const waitUntilPatched = new WeakSet(); - -function patchWaitUntil(ctx: { +function getPatchedWaitUntil(ctx: { waitUntil: (p: Promise | PyFuture) => void; -}): void { - let tag; - try { - tag = Object.prototype.toString.call(ctx); - } catch (_e) {} - if (tag !== '[object ExecutionContext]') { - return; - } - if (waitUntilPatched.has(ctx)) { - return; - } +}): (p: Promise | PyFuture) => void { const origWaitUntil: (p: Promise) => void = ctx.waitUntil.bind(ctx); function waitUntil(p: Promise | PyFuture): void { origWaitUntil( @@ -68,7 +58,23 @@ function patchWaitUntil(ctx: { })() ); } - ctx.waitUntil = waitUntil; + return waitUntil; +} + +function patchWaitUntil(ctx: { + waitUntil: (p: Promise | PyFuture) => void; +}): void { + let tag; + try { + tag = Object.prototype.toString.call(ctx); + } catch (_e) {} + if (tag !== '[object ExecutionContext]') { + return; + } + if (waitUntilPatched.has(ctx)) { + return; + } + ctx.waitUntil = getPatchedWaitUntil(ctx); waitUntilPatched.add(ctx); } @@ -93,13 +99,35 @@ function get_pyodide_entrypoint_helper(): PyodideEntrypointHelper { return _pyodide_entrypoint_helper; } +async function getCloudflareWorkersModule( + doAnImport: (mod: string) => Promise +): 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, workerEntrypoint: any ): Promise { + const cloudflareWorkersModule = await getCloudflareWorkersModule(doAnImport); _pyodide_entrypoint_helper = { doAnImport, - cloudflareWorkersModule: await doAnImport('cloudflare:workers'), + cloudflareWorkersModule, cloudflareSocketsModule: await doAnImport('cloudflare:sockets'), workerEntrypoint, patchWaitUntil, @@ -112,7 +140,12 @@ export async function setDoAnImport( throw new PythonWorkersInternalError(message); }, }; - await fillSnapshotJsModules(doAnImport); + await fillSnapshotJsModules(async (mod: string): Promise => { + if (mod === 'cloudflare:workers') { + return cloudflareWorkersModule; + } + return await doAnImport(mod); + }); } function handleSrcImport(pyodide: Pyodide, e: any): never { diff --git a/src/workerd/server/tests/python/python-rpc/worker.py b/src/workerd/server/tests/python/python-rpc/worker.py index f7c3b7d3d61..37f45ae3505 100644 --- a/src/workerd/server/tests/python/python-rpc/worker.py +++ b/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 From e23e55dab16e8b9fb20bdedd992b99255df0ccdf Mon Sep 17 00:00:00 2001 From: Hood Chatham Date: Wed, 9 Sep 2026 10:46:21 -0700 Subject: [PATCH 2/2] Python: Make sure that `waitUntil()` releases a rejected future --- src/pyodide/python-entrypoint-helper.ts | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/src/pyodide/python-entrypoint-helper.ts b/src/pyodide/python-entrypoint-helper.ts index 3f7f584bec8..f29807198b2 100644 --- a/src/pyodide/python-entrypoint-helper.ts +++ b/src/pyodide/python-entrypoint-helper.ts @@ -51,9 +51,12 @@ function getPatchedWaitUntil(ctx: { if ('copy' in p) { p = p.copy(); } - await p; - if ('destroy' in p) { - p.destroy(); + try { + await p; + } finally { + if ('destroy' in p) { + p.destroy(); + } } })() );