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

.github/workflows/tests.yml

@@ -30,11 +30,16 @@ jobs:
       run: |
         uv pip install --system -e packages/runtime-sdk -e packages/cli pytest
 
-    - name: Run tests
+    - name: Run cli tests
       working-directory: packages/cli
       run: |
         pytest 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

@@ -38,8 +38,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,89 @@
+"""
+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
+"""
+
+from __future__ import annotations
+
+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/_workers.py

@@ -108,6 +108,13 @@ def import_from_javascript(module_name: str) -> Any:
         raise
 
 
+@functools.cache
+def get_js_sdk():
+    # IMPORTANT: The module name here must match how wrangler registers the module
+    # while vendoring the python_modules.
+    return import_from_javascript("python_modules/workers/sdk.mjs")
+
+
 @contextmanager
 def patch_env(
     d: dict[str, Any] | Sequence[tuple[str, Any]] | None = None, **kwds: dict[str, Any]
@@ -1332,7 +1339,8 @@ def _wrap_subclass(cls):
 
     def wrapped_init(self, *args, **kwargs):
         if len(args) > 0:
-            _pyodide_entrypoint_helper.patchWaitUntil(args[0])
+            js_sdk = get_js_sdk()
+            js_sdk.patchWaitUntil(args[0])
         if len(args) > 1:
             args = list(args)
             args[1] = _EnvWrapper(args[1])

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,28 @@
+"""Verify that sdk.mjs is up to date with the TypeScript source."""
+
+from __future__ import annotations
+
+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);
+}