feat(workers-runtime-sdk): support embedding javascript files in sdk

.github/workflows/tests.yml

@@ -35,11 +35,16 @@ jobs:
         uv pip install --system -e packages/runtime-sdk -e packages/cli
         uv pip install --system --group dev --project packages/cli
 
-    - name: Run tests
+    - name: Run cli tests
       working-directory: packages/cli
       run: |
         pytest -v tests
 
+    - name: Run runtime-sdk tests
+      working-directory: packages/runtime-sdk
+      run: |
+        pytest tests
+
     - name: Verify that pywrangler can be run globally
       run: |
         pywrangler --help

packages/cli/tests/test_in_workerd.py

@@ -41,8 +41,11 @@ def embed(dir: Path, root: Path, level: int = 0):
             modules.append(
                 f'(name = "{module_path}", pythonModule = embed "{embed_path}")'
             )
+        elif path.suffix == ".mjs":
+            modules.append(f'(name = "{module_path}", esModule = embed "{embed_path}")')
         else:
             modules.append(f'(name = "{module_path}", data = embed "{embed_path}")')
+
     return modules
 
 

packages/runtime-sdk/scripts/compile_js_sdk.py

@@ -0,0 +1,87 @@
+"""
+Compile ts/sdk.ts -> src/workers/sdk.mjs using esbuild.
+
+Usage:
+    python scripts/compile_js_sdk.py          # compile and write
+    python scripts/compile_js_sdk.py --check  # verify sdk.mjs is up to date
+"""
+
+import shutil
+import subprocess
+import sys
+from pathlib import Path
+
+RUNTIME_SDK_DIR = Path(__file__).resolve().parent.parent
+TS_SOURCE = RUNTIME_SDK_DIR / "ts" / "sdk.ts"
+MJS_OUTPUT = RUNTIME_SDK_DIR / "src" / "workers" / "sdk.mjs"
+
+HEADER = """\
+// AUTO-GENERATED FILE - DO NOT EDIT DIRECTLY
+// Source: ts/sdk.ts
+// Regenerate: python scripts/compile_js_sdk.py
+"""
+
+
+def compile_ts() -> str:
+    """Compile ts/sdk.ts to JavaScript and return the output string (with header)."""
+    npx = shutil.which("npx")
+    if npx is None:
+        print(
+            "error: npx not found. Install Node.js to compile TypeScript.",
+            file=sys.stderr,
+        )
+        sys.exit(1)
+
+    result = subprocess.run(
+        [
+            npx,
+            "--yes",
+            "esbuild@0.28.0",
+            str(TS_SOURCE),
+            "--format=esm",
+            "--log-level=error",
+        ],
+        capture_output=True,
+        text=True,
+        cwd=str(RUNTIME_SDK_DIR),
+        check=False,
+    )
+
+    if result.returncode != 0:
+        print(f"esbuild failed:\n{result.stderr}", file=sys.stderr)
+        sys.exit(result.returncode)
+
+    return HEADER + result.stdout
+
+
+def main() -> None:
+    compiled = compile_ts()
+
+    if "--check" in sys.argv:
+        if not MJS_OUTPUT.exists():
+            print(
+                f"error: {MJS_OUTPUT.relative_to(RUNTIME_SDK_DIR)} does not exist.",
+                file=sys.stderr,
+            )
+            print("Run: python scripts/compile_js_sdk.py", file=sys.stderr)
+            sys.exit(1)
+
+        current = MJS_OUTPUT.read_text()
+        if current != compiled:
+            print(
+                f"error: {MJS_OUTPUT.relative_to(RUNTIME_SDK_DIR)} is out of date.",
+                file=sys.stderr,
+            )
+            print("Run: python scripts/compile_js_sdk.py", file=sys.stderr)
+            sys.exit(1)
+
+        print(f"{MJS_OUTPUT.relative_to(RUNTIME_SDK_DIR)} is up to date.")
+    else:
+        MJS_OUTPUT.write_text(compiled)
+        print(
+            f"Compiled {TS_SOURCE.relative_to(RUNTIME_SDK_DIR)} -> {MJS_OUTPUT.relative_to(RUNTIME_SDK_DIR)}"
+        )
+
+
+if __name__ == "__main__":
+    main()

packages/runtime-sdk/src/workers/__init__.py

@@ -19,11 +19,13 @@
     WorkflowEntrypoint,
     fetch,
     handler,
-    import_from_javascript,
     patch_env,
     python_from_rpc,
     python_to_rpc,
 )
+from .js_sdk import (
+    import_from_javascript,
+)
 
 __all__ = [
     "Blob",

packages/runtime-sdk/src/workers/_workers.py

@@ -39,6 +39,9 @@
 )
 from pyodide.http import pyfetch
 
+from workers.js_sdk import (
+    patch_wait_until,
+)
 from workers.workflows import NonRetryableError
 
 if TYPE_CHECKING:
@@ -61,54 +64,6 @@ def _jsnull_to_none(x):
     return x
 
 
-def import_from_javascript(module_name: str) -> Any:
-    """
-    Import a JavaScript ES module from Python.
-
-    Args:
-        module_name: The name of the module to import. This can be a module name or a path.
-
-    Returns:
-        The imported module object.
-
-    Example:
-        cloudflare_workers = import_from_javascript("cloudflare:workers")
-        env = cloudflare_workers.env
-
-    Note:
-        Behind the scenes import_from_javascript uses JSPI to do imports but that means we need an
-        async context. To enable importing cloudflare:workers and cloudflare:sockets in the global
-        scope we specifically imported them in the global scope and exposed them here.
-    """
-    # Special case for global scope available modules
-    # JSPI won't work in the global scope in 0.26.0a2 so we need modules importable in the global
-    # scope to be imported beforehand.
-    if module_name == "cloudflare:workers":
-        return _pyodide_entrypoint_helper.cloudflareWorkersModule
-    elif module_name == "cloudflare:sockets":
-        return _pyodide_entrypoint_helper.cloudflareSocketsModule
-
-    try:
-        from pyodide.ffi import run_sync
-
-        # Call the JavaScript import function
-        return run_sync(_pyodide_entrypoint_helper.doAnImport(module_name))
-    except JsException as e:
-        raise ImportError(f"Failed to import '{module_name}': {e}") from e
-    except RuntimeError as e:
-        if e.args[0] == "No suspender":
-            raise ImportError(
-                f"Failed to import '{module_name}': Only 'cloudflare:workers' and 'cloudflare:sockets' are available in the global scope."
-            ) from e
-        raise
-    except ImportError as e:
-        if e.args[0].startswith("cannot import name 'run_sync' from 'pyodide.ffi'"):
-            raise ImportError(
-                f"Failed to import '{module_name}': Only 'cloudflare:workers' and 'cloudflare:sockets' are available until the next python runtime version."
-            ) from e
-        raise
-
-
 @contextmanager
 def patch_env(
     d: dict[str, Any] | Sequence[tuple[str, Any]] | None = None, **kwds: dict[str, Any]
@@ -1367,7 +1322,7 @@ def _wrap_subclass(cls):
     def wrapped_init(self, *args, **kwargs):
         args = list(args)
         if len(args) > 0:
-            _pyodide_entrypoint_helper.patchWaitUntil(args[0])
+            patch_wait_until(args[0])
             if issubclass(cls, DurableObject):
                 args[0] = DurableObjectContext(args[0])
         if len(args) > 1:

packages/runtime-sdk/src/workers/js_sdk.py

@@ -0,0 +1,75 @@
+import functools
+from typing import Any
+
+import _pyodide_entrypoint_helper
+from pyodide import __version__ as pyodide_version
+from pyodide.ffi import JsException
+
+
+def import_from_javascript(module_name: str) -> Any:
+    """
+    Import a JavaScript ES module from Python.
+
+    Args:
+        module_name: The name of the module to import. This can be a module name or a path.
+
+    Returns:
+        The imported module object.
+
+    Example:
+        cloudflare_workers = import_from_javascript("cloudflare:workers")
+        env = cloudflare_workers.env
+
+    Note:
+        Behind the scenes import_from_javascript uses JSPI to do imports but that means we need an
+        async context. To enable importing cloudflare:workers and cloudflare:sockets in the global
+        scope we specifically imported them in the global scope and exposed them here.
+    """
+    # Special case for global scope available modules
+    # JSPI won't work in the global scope in 0.26.0a2 so we need modules importable in the global
+    # scope to be imported beforehand.
+    if module_name == "cloudflare:workers":
+        return _pyodide_entrypoint_helper.cloudflareWorkersModule
+    elif module_name == "cloudflare:sockets":
+        return _pyodide_entrypoint_helper.cloudflareSocketsModule
+
+    try:
+        from pyodide.ffi import run_sync
+
+        # Call the JavaScript import function
+        return run_sync(_pyodide_entrypoint_helper.doAnImport(module_name))
+    except JsException as e:
+        raise ImportError(f"Failed to import '{module_name}': {e}") from e
+    except RuntimeError as e:
+        if e.args[0] == "No suspender":
+            raise ImportError(
+                f"Failed to import '{module_name}': Only 'cloudflare:workers' and 'cloudflare:sockets' are available in the global scope."
+            ) from e
+        raise
+    except ImportError as e:
+        if e.args[0].startswith("cannot import name 'run_sync' from 'pyodide.ffi'"):
+            raise ImportError(
+                f"Failed to import '{module_name}': Only 'cloudflare:workers' and 'cloudflare:sockets' are available until the next python runtime version."
+            ) from e
+        raise
+
+
+@functools.cache
+def get_js_sdk():
+    # IMPORTANT:
+    # The module name here must match how wrangler registers the JS modules
+    # while vendoring the python_modules directory.
+    # See: https://github.com/cloudflare/workers-sdk/pull/13311
+    return import_from_javascript("python_modules/workers/sdk.mjs")
+
+
+def patch_wait_until(ctx):
+    """
+    Patch the waitUntil method of the given context to ensure that async operations are properly handled.
+    """
+    if pyodide_version == "0.26.0a2":
+        _pyodide_entrypoint_helper.patchWaitUntil(ctx)
+        return
+
+    js_sdk = get_js_sdk()
+    js_sdk.patchWaitUntil(ctx)

packages/runtime-sdk/src/workers/sdk.mjs

@@ -0,0 +1,36 @@
+// AUTO-GENERATED FILE - DO NOT EDIT DIRECTLY
+// Source: ts/sdk.ts
+// Regenerate: python scripts/compile_js_sdk.py
+const waitUntilPatched = /* @__PURE__ */ new WeakSet();
+function patchWaitUntil(ctx) {
+  let tag;
+  try {
+    tag = Object.prototype.toString.call(ctx);
+  } catch (_e) {
+  }
+  if (tag !== "[object ExecutionContext]") {
+    return;
+  }
+  if (waitUntilPatched.has(ctx)) {
+    return;
+  }
+  const origWaitUntil = ctx.waitUntil.bind(ctx);
+  function waitUntil(p) {
+    origWaitUntil(
+      (async function() {
+        if ("copy" in p) {
+          p = p.copy();
+        }
+        await p;
+        if ("destroy" in p) {
+          p.destroy();
+        }
+      })()
+    );
+  }
+  ctx.waitUntil = waitUntil;
+  waitUntilPatched.add(ctx);
+}
+export {
+  patchWaitUntil
+};

packages/runtime-sdk/tests/test_js_sdk_compile.py

@@ -0,0 +1,26 @@
+"""Verify that sdk.mjs is up to date with the TypeScript source."""
+
+import subprocess
+import sys
+from pathlib import Path
+
+RUNTIME_SDK_DIR = Path(__file__).resolve().parent.parent
+COMPILE_SCRIPT = RUNTIME_SDK_DIR / "scripts" / "compile_js_sdk.py"
+
+
+def test_sdk_mjs_up_to_date() -> None:
+    """sdk.mjs must match the output of compiling ts/sdk.ts.
+
+    If this test fails, run:
+        python scripts/compile_js_sdk.py
+    """
+    result = subprocess.run(
+        [sys.executable, str(COMPILE_SCRIPT), "--check"],
+        capture_output=True,
+        text=True,
+        cwd=str(RUNTIME_SDK_DIR),
+        check=False,
+    )
+    assert result.returncode == 0, (
+        f"sdk.mjs is out of date. Run: python scripts/compile_js_sdk.py\n{result.stderr}"
+    )

packages/runtime-sdk/ts/sdk.ts

@@ -0,0 +1,41 @@
+// Javascript helper functions for workers-runtime-sdk
+// This file is compiled to src/workers/sdk.mjs via scripts/compile_js_sdk.py
+
+// Pyodide proxy future — supports copy/destroy for proxy lifecycle management
+type PyFuture<T> = Promise<T> & {
+  copy(): PyFuture<T>;
+  destroy(): void;
+};
+
+const waitUntilPatched = new WeakSet();
+
+export function patchWaitUntil(ctx: {
+  waitUntil: (p: Promise<void> | PyFuture<void>) => void;
+}): void {
+  let tag;
+  try {
+    tag = Object.prototype.toString.call(ctx);
+  } catch (_e) {}
+  if (tag !== '[object ExecutionContext]') {
+    return;
+  }
+  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;
+  waitUntilPatched.add(ctx);
+}