workflow: support to using a shared sandbox for all workflow invocations

changes/vercel-workflow/shared-sandboxes.feature.md

@@ -0,0 +1 @@
+Add a `share_sandboxes` parameter to `SandboxPolicy` to enable reusing already created sandboxes instead of creating a new one on each invocation. This speeds up workflows but means that modifications to global state may persist between invocations.

src/vercel-workflow/tests/unit/test_py_sandbox.py

@@ -14,10 +14,11 @@
 
 import platform
 import sys
+from contextlib import contextmanager
 
 import pytest
 
-from vercel.workflow._internal.py_sandbox import SandboxRestrictionError, workflow_sandbox
+from vercel.workflow._internal.py_sandbox import Sandbox, SandboxRestrictionError
 from vercel.workflow.sandbox import (
     ALL_CLEANUPS,
     SandboxCleanupContext,
@@ -27,6 +28,12 @@
 CLEANUP_ALL = SandboxPolicy(cleanups=ALL_CLEANUPS)
 
 
+@contextmanager
+def workflow_sandbox(*, policy: SandboxPolicy | None = None):
+    with Sandbox(policy=policy).enter():
+        yield
+
+
 # ── helpers ────────────────────────────────────────────────────
 
 
@@ -1165,3 +1172,33 @@ def test_workflows_threads_policy_to_registry(self):
         registry = core.Workflows(as_vercel_job=False, sandbox_policy=policy)
         assert registry._sandbox_policy is policy
         assert core.Workflows(as_vercel_job=False)._sandbox_policy == SandboxPolicy()
+
+    @staticmethod
+    def _modules_seen_by_two_runs(policy: SandboxPolicy | None) -> list:
+        from vercel.workflow._internal import core
+
+        registry = core.Workflows(as_vercel_job=False, sandbox_policy=policy)
+        modules = []
+        for _ in range(2):
+            with registry._get_sandbox() as sandbox, sandbox.enter():
+                import uuid  # noqa: PLC0415
+
+                modules.append(uuid)
+        return modules
+
+    def test_shared_sandboxes_reuse_modules_across_runs(self):
+        """With share_sandboxes=True every run gets the one cached sandbox,
+        so a re-imported module is the same object run to run."""
+        first, second = self._modules_seen_by_two_runs(SandboxPolicy(share_sandboxes=True))
+        assert first is second
+
+    def test_unshared_sandboxes_import_modules_freshly_each_run(self):
+        """With share_sandboxes=False every run gets its own sandbox, and
+        with it its own fresh imports."""
+        first, second = self._modules_seen_by_two_runs(SandboxPolicy(share_sandboxes=False))
+        assert first is not second
+
+    def test_sandboxes_are_not_shared_by_default(self):
+        assert SandboxPolicy().share_sandboxes is False
+        first, second = self._modules_seen_by_two_runs(None)
+        assert first is not second

src/vercel-workflow/tests/unit/test_workflow_encryption.py

@@ -308,7 +308,7 @@ def test_decryption_works_inside_the_workflow_sandbox() -> None:
     key = encryption.derive_run_key(DEPLOYMENT_KEY, project_id=PROJECT_ID, run_id=RUN_ID)
     payload = seal(key, ser.dehydrate([{"amount": 21}]))
 
-    with py_sandbox.workflow_sandbox():
+    with py_sandbox.Sandbox().enter():
         assert ser.hydrate(payload, what="the input of run wrun_1", key=key) == [{"amount": 21}]
 
 
@@ -340,7 +340,7 @@ def _first_call_of_a_fresh_interpreter(body: str, payload: bytes) -> str:
 from vercel.workflow._internal import encryption, py_sandbox, serialization as ser
 
 key = encryption.derive_run_key(bytes(range(32)), project_id="prj_test", run_id="wrun_test")
-with py_sandbox.workflow_sandbox():
+with py_sandbox.Sandbox().enter():
     print(ser.hydrate(bytes.fromhex(sys.argv[1]), what="the input of run wrun_1", key=key))
 """
 
@@ -349,7 +349,7 @@ def _first_call_of_a_fresh_interpreter(body: str, payload: bytes) -> str:
 from vercel.workflow._internal import encryption, py_sandbox
 
 key = encryption.derive_run_key(bytes(range(32)), project_id="prj_test", run_id="wrun_test")
-with py_sandbox.workflow_sandbox():
+with py_sandbox.Sandbox().enter():
     try:
         encryption.open_envelope(key, bytes.fromhex(sys.argv[1]))
     except encryption.DecryptionError as exc:
@@ -496,7 +496,7 @@ def test_x25519_is_bound_once_at_import(monkeypatch) -> None:
 def test_opening_a_sealed_payload_works_inside_the_workflow_sandbox() -> None:
     payload = seal_to(RUN_KEY, ser.dehydrate({"type": "approve"}))
 
-    with py_sandbox.workflow_sandbox():
+    with py_sandbox.Sandbox().enter():
         assert ser.hydrate(payload, what="the payload of hook hook_1", key=RUN_KEY) == {
             "type": "approve"
         }

src/vercel-workflow/tests/unit/test_workflow_serde.py

@@ -30,6 +30,13 @@ def _round_trip(value):
     return ser.hydrate(ser.dehydrate(value), what="a payload")
 
 
+def _sandbox_registry() -> serde.Registry:
+    """A registry as `Sandbox` builds one: fresh, seeded with the built-ins."""
+    registry = serde.Registry()
+    serde._register_builtins(registry)
+    return registry
+
+
 # ═══════════════════════════════════════════════════════════════════════════
 # the wire form
 # ═══════════════════════════════════════════════════════════════════════════
@@ -251,7 +258,7 @@ def test_replays_do_not_accumulate_in_the_host_registry() -> None:
     classes, class_ids = len(serde._HOST.by_class), len(serde._HOST.by_class_id)
 
     for _ in range(50):
-        with serde.sandboxed_registrations():
+        with serde.sandboxed_registrations(_sandbox_registry()):
             _define_and_register()
     gc.collect()
 
@@ -270,7 +277,7 @@ def test_a_sandbox_registration_does_not_outlive_the_sandbox() -> None:
     """
     host = _define_and_register()
 
-    with serde.sandboxed_registrations():
+    with serde.sandboxed_registrations(_sandbox_registry()):
         sandboxed = _define_and_register()
         assert type(_round_trip(sandboxed())) is sandboxed, "its own class applies inside"
         # The classId is the same on both sides, which is what lets a payload
@@ -284,7 +291,7 @@ def test_a_sandbox_registration_does_not_outlive_the_sandbox() -> None:
 def test_a_sandbox_registration_does_not_reach_the_host() -> None:
     # Nothing carries a sandbox-registered class out: a run's payloads are all
     # serialized inside the sandbox, its return value included.
-    with serde.sandboxed_registrations():
+    with serde.sandboxed_registrations(_sandbox_registry()):
         serde.register_serializable(Point, class_id="class//sandbox-only//Point")
 
     with pytest.raises(ser.SerializationError, match="Register Point with @serializable"):
@@ -302,7 +309,7 @@ def test_a_host_registration_is_not_visible_inside_a_sandbox() -> None:
     serde.register_serializable(Point)
     payload = ser.dehydrate(Point(1, 2))
 
-    with serde.sandboxed_registrations():
+    with serde.sandboxed_registrations(_sandbox_registry()):
         with pytest.raises(ser.SerializationError, match="unknown class"):
             ser.hydrate(payload, what="a payload")
 
@@ -313,7 +320,7 @@ def test_the_sandbox_registers_the_stdlib_classes_it_imported_itself() -> None:
     Its `uuid.UUID` is a different class object, so a registry holding only the
     host's would not cover a `UUID` built inside a workflow.
     """
-    with py_sandbox.workflow_sandbox():
+    with py_sandbox.Sandbox().enter():
         import uuid as sandboxed  # noqa: PLC0415
 
         value = sandboxed.UUID("6ba7b810-9dad-11d1-80b4-00c04fd430c8")
@@ -330,7 +337,7 @@ def test_a_restricted_stdlib_class_is_not_swapped_for_the_hosts() -> None:
     """
     payload = ser.dehydrate(datetime.date(2026, 8, 4))
 
-    with py_sandbox.workflow_sandbox():
+    with py_sandbox.Sandbox().enter():
         import datetime as sandboxed  # noqa: PLC0415
 
         revived = ser.hydrate(payload, what="a payload")
@@ -342,7 +349,7 @@ def test_a_restricted_stdlib_class_is_not_swapped_for_the_hosts() -> None:
 def test_the_wire_id_does_not_follow_the_sandbox_class_name() -> None:
     # Derived from the class it would read `class//...//_RestrictedDate`, which
     # the other side has never heard of.
-    with py_sandbox.workflow_sandbox():
+    with py_sandbox.Sandbox().enter():
         import datetime as sandboxed  # noqa: PLC0415
 
         wire = _wire(sandboxed.date(2026, 8, 4))
@@ -353,18 +360,17 @@ def test_the_wire_id_does_not_follow_the_sandbox_class_name() -> None:
 def test_datetime_stays_native_inside_a_sandbox() -> None:
     # `registry.native` has to hold the sandbox's `datetime`, not the host's,
     # or the `date` registration would capture it through the MRO.
-    with py_sandbox.workflow_sandbox():
+    with py_sandbox.Sandbox().enter():
         import datetime as sandboxed  # noqa: PLC0415
 
         value = sandboxed.datetime(2026, 7, 30, tzinfo=datetime.timezone.utc)
         assert "Instance" not in _wire(value)
 
 
 def test_built_ins_are_visible_inside_a_sandbox() -> None:
-    # They are registered when this module is imported, and `serde` is reached
-    # through to the host rather than re-imported, so a registry that started
-    # empty would not have them.
-    with serde.sandboxed_registrations():
+    # Each sandbox registry is seeded with them at construction; a registry
+    # that started empty would not have them.
+    with serde.sandboxed_registrations(_sandbox_registry()):
         assert _round_trip(decimal.Decimal("1.50")) == decimal.Decimal("1.50")
         assert _round_trip(pathlib.Path("/tmp/x")) == pathlib.Path("/tmp/x")
 

src/vercel-workflow/vercel/workflow/_internal/core.py

@@ -1,11 +1,12 @@
 from __future__ import annotations
 
+import contextlib
 import dataclasses
 import datetime
 import functools
 import inspect
 import random as _random
-from collections.abc import AsyncIterator, Callable, Coroutine, Generator
+from collections.abc import AsyncIterator, Callable, Coroutine, Generator, Iterator
 from typing import Any, Generic, ParamSpec, TypeVar, overload
 
 import pydantic
@@ -285,12 +286,24 @@ def __init__(
         if sandbox_policy is None:
             sandbox_policy = py_sandbox.SandboxPolicy()
         self._sandbox_policy = sandbox_policy
+
+        self._cached_sandbox = None
+        if sandbox_policy.share_sandboxes and not py_sandbox.in_sandbox():
+            self._cached_sandbox = py_sandbox.Sandbox(policy=sandbox_policy, run_cleanups=False)
+
         self._http_handler: w.HTTPHandler | None = None
         if as_vercel_job and not py_sandbox.in_sandbox():
             from . import runtime
 
             self._http_handler = runtime.workflow_entrypoint(self)
 
+    @contextlib.contextmanager
+    def _get_sandbox(self) -> Iterator[py_sandbox.Sandbox]:
+        if self._cached_sandbox:
+            yield self._cached_sandbox
+        else:
+            yield py_sandbox.Sandbox(policy=self._sandbox_policy)
+
     @property
     def namespace(self) -> str | None:
         """The immutable queue namespace for this registry."""