fix(sandbox): suppress policy warnings during initial workflow load

CHANGELOG.md

@@ -71,6 +71,9 @@ to include examples, links to docs, or any other relevant information.
 
 - Marked system Nexus envelope payloads so nested payloads can be detected and
   visited after the envelope is already stored as a payload.
+- Suppressed intentional passthrough notifications for the initial workflow
+  module import while retaining configured notifications for that module's
+  import-time dependencies.
 
 ### Security
 

temporalio/worker/workflow_sandbox/_importer.py

@@ -66,6 +66,7 @@ def __init__(
             "__main__": types.ModuleType("__main__"),
         }
         self.modules_checked_for_restrictions: set[str] = set()
+        self._initial_workflow_module_import: str | None = None
         self.import_func = self._import if not LOG_TRACE else self._traced_import
         # Pre-collect restricted builtins
         self.restricted_builtins: list[tuple[str, _ThreadLocalCallable, Callable]] = []
@@ -159,6 +160,16 @@ def applied(self) -> Iterator[None]:
         finally:
             Importer._thread_local_current.importer = orig_importer
 
+    @contextmanager
+    def initial_workflow_module_import(self, module_name: str) -> Iterator[None]:
+        """Suppress notifications only for the initial workflow module import."""
+        previous_module = self._initial_workflow_module_import
+        self._initial_workflow_module_import = module_name
+        try:
+            yield None
+        finally:
+            self._initial_workflow_module_import = previous_module
+
     @contextmanager
     def _unapplied(self) -> Iterator[None]:
         orig_importer = Importer.current_importer()
@@ -307,19 +318,30 @@ def _is_import_notification_policy_applied(
         return policy in self.restrictions.import_notification_policy
 
     def _maybe_passthrough_module(self, name: str) -> types.ModuleType | None:
+        # The workflow module itself must be loaded into every sandbox instance.
+        # That import is intentional, but its import-time dependencies still need
+        # to follow the configured notification policy.
+        is_initial_workflow_module = name == self._initial_workflow_module_import
+
         # If imports not passed through and all modules are not passed through
         # and name not in passthrough modules, check parents
         if (
             not temporalio.workflow.unsafe.is_imports_passed_through()
             and not self.module_configured_passthrough(name)
         ):
-            if self._is_import_notification_policy_applied(
-                temporalio.workflow.SandboxImportNotificationPolicy.RAISE_ON_UNINTENTIONAL_PASSTHROUGH
+            if (
+                not is_initial_workflow_module
+                and self._is_import_notification_policy_applied(
+                    temporalio.workflow.SandboxImportNotificationPolicy.RAISE_ON_UNINTENTIONAL_PASSTHROUGH
+                )
             ):
                 raise UnintentionalPassthroughError(name)
 
-            if self._is_import_notification_policy_applied(
-                temporalio.workflow.SandboxImportNotificationPolicy.WARN_ON_UNINTENTIONAL_PASSTHROUGH
+            if (
+                not is_initial_workflow_module
+                and self._is_import_notification_policy_applied(
+                    temporalio.workflow.SandboxImportNotificationPolicy.WARN_ON_UNINTENTIONAL_PASSTHROUGH
+                )
             ):
                 warnings.warn(
                     f"Module {name} was not intentionally passed through to the sandbox."

temporalio/worker/workflow_sandbox/_runner.py

@@ -135,14 +135,17 @@ def _create_instance(self) -> None:
         if module_name == "__main__":
             module_name = "__temporal_main__"
         try:
-            # Import user code
-            self._run_code(
-                "with __temporal_importer.applied():\n"
-                # Import the workflow code
-                f"  from {module_name} import {self.instance_details.defn.cls.__name__} as __temporal_workflow_class\n"
-                f"  from {self.runner_class.__module__} import {self.runner_class.__name__} as __temporal_runner_class\n",
-                __temporal_importer=self.importer,
-            )
+            # Import user code. The workflow module is intentionally loaded
+            # into every sandbox instance, while its import-time dependencies
+            # must retain the configured warning/error policy.
+            with self.importer.initial_workflow_module_import(module_name):
+                self._run_code(
+                    "with __temporal_importer.applied():\n"
+                    # Import the workflow code
+                    f"  from {module_name} import {self.instance_details.defn.cls.__name__} as __temporal_workflow_class\n"
+                    f"  from {self.runner_class.__module__} import {self.runner_class.__name__} as __temporal_runner_class\n",
+                    __temporal_importer=self.importer,
+                )
 
             # Set context as in runtime
             self.importer.restriction_context.is_runtime = True

tests/worker/workflow_sandbox/test_runner.py

@@ -525,6 +525,59 @@ async def run(self) -> None:
             ) from err
 
 
+async def test_workflow_sandbox_initial_workflow_import_does_not_warn() -> None:
+    restrictions = dataclasses.replace(
+        SandboxRestrictions.default,
+        import_notification_policy=SandboxImportNotificationPolicy.WARN_ON_UNINTENTIONAL_PASSTHROUGH,
+    )
+
+    with warnings.catch_warnings(record=True) as recorder:
+        warnings.simplefilter("always")
+        definition = workflow._Definition.from_class(LazyImportWorkflow)
+        assert definition is not None
+        SandboxedWorkflowRunner(restrictions).prepare_workflow(definition)
+
+    assert (
+        "Module tests.worker.workflow_sandbox.test_runner was not intentionally "
+        "passed through to the sandbox."
+        not in {str(warning.message) for warning in recorder}
+    )
+
+
+async def test_workflow_sandbox_initial_workflow_import_warns_for_dependencies() -> (
+    None
+):
+    from tests.worker.workflow_sandbox.testmodules.initial_import_warning_dependency import (
+        InitialImportWarningDependencyWorkflow,
+    )
+
+    restrictions = dataclasses.replace(
+        SandboxRestrictions.default,
+        import_notification_policy=SandboxImportNotificationPolicy.WARN_ON_UNINTENTIONAL_PASSTHROUGH,
+    )
+
+    with warnings.catch_warnings(record=True) as recorder:
+        warnings.simplefilter("always")
+        definition = workflow._Definition.from_class(
+            InitialImportWarningDependencyWorkflow
+        )
+        assert definition is not None
+        SandboxedWorkflowRunner(restrictions).prepare_workflow(definition)
+
+    assert (
+        "Module tests.worker.workflow_sandbox.testmodules.lazy_module was not "
+        "intentionally passed through to the sandbox."
+        in {str(warning.message) for warning in recorder}
+    )
+
+    raising_restrictions = dataclasses.replace(
+        restrictions,
+        import_notification_policy=SandboxImportNotificationPolicy.RAISE_ON_UNINTENTIONAL_PASSTHROUGH,
+    )
+    with pytest.raises(UnintentionalPassthroughError, match="lazy_module"):
+        SandboxedWorkflowRunner(raising_restrictions).prepare_workflow(definition)
+
+
 async def test_workflow_sandbox_import_default_warnings(client: Client):
     restrictions = dataclasses.replace(
         SandboxRestrictions.default,

tests/worker/workflow_sandbox/testmodules/initial_import_warning_dependency.py

@@ -0,0 +1,9 @@
+import tests.worker.workflow_sandbox.testmodules.lazy_module  # noqa: F401
+from temporalio import workflow
+
+
+@workflow.defn
+class InitialImportWarningDependencyWorkflow:
+    @workflow.run
+    async def run(self) -> None:
+        pass