Make JSPI work correctly

Author: hoodmane

src/pyodide/internal/pool/builtin_wrappers.ts

@@ -74,6 +74,9 @@ export function finishSetup() {
 }
 
 export function newWasmModule(buffer: Uint8Array): WebAssembly.Module {
+  if (!UnsafeEval) {
+    return new WebAssembly.Module(buffer);
+  }
   if (finishedSetup) {
     checkCallee();
   }

src/pyodide/python-entrypoint-helper.ts

@@ -135,6 +135,13 @@ async function preparePython(): Promise<PyModule> {
   return mainModule;
 }
 
+function callHandler(handler: PyCallable, args: any[]): any {
+  if (handler.callWithOptions) {
+    return handler.callWithOptions({ relaxed: true, promising: true }, ...args);
+  }
+  return handler.callRelaxed(...args);
+}
+
 function makeHandler(pyHandlerName: string): Handler {
   return async function (...args: any[]) {
     try {
@@ -149,7 +156,7 @@ function makeHandler(pyHandlerName: string): Handler {
         );
       }
       const result = await enterJaegerSpan('python_code', () => {
-        return handler.callRelaxed(...args);
+        return callHandler(handler, args);
       });
 
       // Support returning a pyodide.ffi.FetchResponse.

src/pyodide/types/Pyodide.d.ts

@@ -5,6 +5,10 @@ declare type Handler = (...args: any[]) => any;
 declare type PyCallable = ((...args: any[]) => any) & {
   call: (...args: any[]) => any;
   callRelaxed: (...args: any[]) => any;
+  callWithOptions?: (
+    options: { relaxed?: boolean; promising?: boolean; kwargs?: boolean },
+    ...args: any[]
+  ) => any;
 };
 
 interface PyModule {

src/workerd/server/tests/python/BUILD.bazel

@@ -72,3 +72,8 @@ py_wd_test("filter-non-py-files")
 py_wd_test("durable-object")
 
 py_wd_test("worker-entrypoint")
+
+py_wd_test(
+    "jspi",
+    skip_python_flags = ["0.26.0a2"],
+)

src/workerd/server/tests/python/jspi/jspi.wd-test

@@ -0,0 +1,15 @@
+using Workerd = import "/workerd/workerd.capnp";
+
+const unitTests :Workerd.Config = (
+  services = [
+    ( name = "jspi",
+      worker = (
+        modules = [
+          (name = "worker.py", pythonModule = embed "worker.py")
+        ],
+        compatibilityDate = "2024-01-15",
+        compatibilityFlags = [%PYTHON_FEATURE_FLAGS],
+      )
+    ),
+  ],
+);

src/workerd/server/tests/python/jspi/worker.py

@@ -0,0 +1,14 @@
+from asyncio import sleep
+
+from pyodide import __version__
+
+
+def test():
+    if __version__ == "0.26.0a2":
+        print("JSPI not supported on 0.26.0a2, skipping")
+        return
+
+    from pyodide.ffi import run_sync
+
+    run_sync(sleep(0.01))
+    print("Okay")