microsoft · larohra · Dec 31, 2025 · Dec 31, 2025 · Dec 31, 2025
diff --git a/durabletask/task.py b/durabletask/task.py
@@ -537,7 +537,7 @@ def task_id(self) -> int:


 # Orchestrators are generators that yield tasks and receive/return any type
-Orchestrator = Callable[[OrchestrationContext, TInput], Union[Generator[Task, Any, Any], TOutput]]
+Orchestrator = Callable[[OrchestrationContext, TInput], Union[Generator[Task[Any], Any, TOutput], TOutput]]

 # Activities are simple functions that can be scheduled by orchestrators
 Activity = Callable[[ActivityContext, TInput], TOutput]

diff --git a/durabletask/worker.py b/durabletask/worker.py
@@ -150,7 +150,7 @@ def __init__(self):
        self.entities = {}
        self.entity_instances = {}

-    def add_orchestrator(self, fn: task.Orchestrator) -> str:
+    def add_orchestrator(self, fn: task.Orchestrator[TInput, TOutput]) -> str:
        if fn is None:
            raise ValueError("An orchestrator function argument is required.")

@@ -166,7 +166,7 @@ def add_named_orchestrator(self, name: str, fn: task.Orchestrator) -> None:

        self.orchestrators[name] = fn

-    def get_orchestrator(self, name: str) -> Optional[task.Orchestrator]:
+    def get_orchestrator(self, name: str) -> Optional[task.Orchestrator[Any, Any]]:
        return self.orchestrators.get(name)

    def add_activity(self, fn: task.Activity) -> str: