Validate programmatic workflow handlers at every admission surface

Author: durable-workflow-ops

docs/sdk-reference.md

@@ -421,7 +421,8 @@ Within a workflow class, signal, query, update and validator names are checked
 independently. One name may be shared across these kinds. Python's normal
 method resolution applies to inheritance: overriding the same Python attribute
 is allowed, while two resolved attributes with the same handler kind and wire
-name are rejected at decoration time. The test environment rejects a different
+name are rejected at decoration time and at Worker, Replayer and test-environment
+admission, including programmatically defined classes. The test environment rejects a different
 workflow class replacing an existing name, while repeated registration of the
 same class remains idempotent for automatic execution. Explicit activity mock
 and canned-result setters continue to support replacement.

src/durable_workflow/testing.py

@@ -52,6 +52,7 @@ def test_my_workflow():
     StartTimer,
     UpsertSearchAttributes,
     _validate_workflow_entrypoint,
+    _workflow_handler_registries,
     registry,
     replay,
 )
@@ -136,6 +137,7 @@ def register_workflow(self, workflow_cls: type) -> None:
                 "decorate it with @workflow.defn(name=...) first."
             )
         _validate_workflow_entrypoint(workflow_cls)
+        _workflow_handler_registries(workflow_cls)
         if name in self._registered_workflows and self._registered_workflows[name] is not workflow_cls:
             raise DuplicateRegistrationError("workflow", name, self._registered_workflows[name], workflow_cls)
         self._registered_workflows[name] = workflow_cls

src/durable_workflow/worker.py

@@ -107,6 +107,7 @@
     StartChildWorkflow,
     UpsertMemo,
     _validate_workflow_entrypoint,
+    _workflow_handler_registries,
     apply_update,
     commands_to_server_commands,
     query_state,
@@ -1029,6 +1030,7 @@ def __init__(
         self.workflows: dict[str, type] = {}
         for definition in workflows:
             _validate_workflow_entrypoint(definition)
+            _workflow_handler_registries(definition)
             name = _workflow_name(definition)
             if name in self.workflows:
                 raise DuplicateRegistrationError("workflow", name, self.workflows[name], definition)

src/durable_workflow/workflow.py

@@ -125,6 +125,29 @@ def _validate_workflow_result(value: Any, cls: type) -> None:
         )
 
 
+def _workflow_handler_registries(
+    cls: type,
+) -> tuple[dict[str, str], dict[str, str], dict[str, str], dict[str, str]]:
+    handlers: dict[str, dict[str, str]] = {kind: {} for kind in ("signal", "query", "update", "update_validator")}
+    for attr in dir(cls):
+        if attr.startswith("_"):
+            continue
+        member = getattr(cls, attr, None)
+        for kind, registry in handlers.items():
+            wire_name = getattr(member, f"__{kind}_name__", None)
+            if not isinstance(wire_name, str) or not wire_name:
+                continue
+            if kind == "signal" and wire_name == MESSAGE_STREAM_SIGNAL:
+                raise ValueError(f"signal name {wire_name!r} is reserved by the workflow runtime")
+            if wire_name in registry:
+                owner = f"{cls.__module__}.{cls.__qualname__}"
+                raise DuplicateRegistrationError(
+                    kind, wire_name, f"{owner}.{registry[wire_name]}", f"{owner}.{attr}",
+                )
+            registry[wire_name] = attr
+    return handlers["signal"], handlers["query"], handlers["update"], handlers["update_validator"]
+
+
 def defn(*, name: str) -> Callable[[type[_WorkflowT]], type[_WorkflowT]]:
     """Register a class as a workflow type under a language-neutral name.
 
@@ -137,37 +160,7 @@ def defn(*, name: str) -> Callable[[type[_WorkflowT]], type[_WorkflowT]]:
     def wrap(cls: type[_WorkflowT]) -> type[_WorkflowT]:
         cls.__workflow_name__ = name  # type: ignore[attr-defined]
         _validate_workflow_entrypoint(cls)
-        signals: dict[str, str] = {}
-        queries: dict[str, str] = {}
-        updates: dict[str, str] = {}
-        update_validators: dict[str, str] = {}
-
-        def register_handler(kind: str, handlers: dict[str, str], wire_name: str, attr: str) -> None:
-            if wire_name in handlers:
-                owner = f"{cls.__module__}.{cls.__qualname__}"
-                raise DuplicateRegistrationError(
-                    kind, wire_name, f"{owner}.{handlers[wire_name]}", f"{owner}.{attr}",
-                )
-            handlers[wire_name] = attr
-
-        for attr in dir(cls):
-            if attr.startswith("_"):
-                continue
-            member = getattr(cls, attr, None)
-            signal_name = getattr(member, "__signal_name__", None)
-            if isinstance(signal_name, str) and signal_name:
-                if signal_name == MESSAGE_STREAM_SIGNAL:
-                    raise ValueError(f"signal name {signal_name!r} is reserved by the workflow runtime")
-                register_handler("signal", signals, signal_name, attr)
-            query_name = getattr(member, "__query_name__", None)
-            if isinstance(query_name, str) and query_name:
-                register_handler("query", queries, query_name, attr)
-            update_name = getattr(member, "__update_name__", None)
-            if isinstance(update_name, str) and update_name:
-                register_handler("update", updates, update_name, attr)
-            update_validator_name = getattr(member, "__update_validator_name__", None)
-            if isinstance(update_validator_name, str) and update_validator_name:
-                register_handler("update_validator", update_validators, update_validator_name, attr)
+        signals, queries, updates, update_validators = _workflow_handler_registries(cls)
         cls.__workflow_signals__ = signals  # type: ignore[attr-defined]
         cls.__workflow_queries__ = queries  # type: ignore[attr-defined]
         cls.__workflow_updates__ = updates  # type: ignore[attr-defined]
@@ -2588,6 +2581,7 @@ class Replayer:
     def __init__(self, *, workflows: Iterable[type]) -> None:
         self._workflows: dict[str, type] = {}
         for workflow_cls in workflows:
+            _workflow_handler_registries(workflow_cls)
             workflow_type = str(getattr(workflow_cls, "__workflow_name__", workflow_cls.__name__))
             if workflow_type in self._workflows:
                 raise DuplicateRegistrationError(

tests/test_registration_ambiguity.py

@@ -114,6 +114,34 @@ def replacement(self: Any) -> str:
     assert child().handler() == "replacement"
 
 
+@pytest.mark.parametrize("kind", ["signal", "query", "update", "update_validator"])
+@pytest.mark.parametrize("consumer", ["worker", "replayer", "environment"])
+def test_programmatic_workflow_admission_cannot_bypass_handler_validation(kind: str, consumer: str) -> None:
+    decorate: Callable[..., Any] = getattr(workflow, kind)
+    def first_handler(self: Any) -> None:
+        pass
+
+    def second_handler(self: Any) -> None:
+        pass
+
+    cls = type("ProgrammaticWorkflow", (), {
+        "__workflow_name__": "programmatic-workflow", "run": run,
+        "first_handler": decorate("approve")(first_handler),
+        "second_handler": decorate("approve")(second_handler),
+    })
+    client = AsyncMock(spec=Client)
+    with pytest.raises(DuplicateRegistrationError) as captured:
+        if consumer == "worker":
+            Worker(client, task_queue="registration", workflows=[cls])
+        elif consumer == "replayer":
+            workflow.Replayer(workflows=[cls])
+        else:
+            WorkflowEnvironment().register_workflow(cls)
+    assert captured.value.registration_kind == kind
+    assert captured.value.wire_name == "approve"
+    assert client.mock_calls == []
+
+
 def test_handler_names_are_scoped_by_kind() -> None:
     cls = definition("scoped-handlers", "ScopedHandlers", {
         "signal": workflow.signal("approve")(lambda self: None),
@@ -157,7 +185,8 @@ def test_mock_replacement_is_an_explicit_test_setting() -> None:
 
 def test_duplicate_error_is_a_public_structured_sdk_error() -> None:
     error = DuplicateRegistrationError("workflow", "duplicate", run, first)
-    assert isinstance(error, (DurableWorkflowError, ValueError))
+    assert isinstance(error, DurableWorkflowError)
+    assert isinstance(error, ValueError)
     assert error.registration_kind == "workflow"
     assert error.wire_name == "duplicate"
     assert error.first_definition.endswith(".run")