88import logging
99import os
1010import time
11- from collections .abc import Callable , Generator , Sequence
11+ from collections .abc import Callable , Coroutine , Generator , Sequence
1212from concurrent .futures import ThreadPoolExecutor
1313from dataclasses import dataclass , field
1414from datetime import datetime , timedelta , timezone
1515from threading import Event , Lock , Thread
16- from types import GeneratorType
16+ from types import CoroutineType , GeneratorType
1717from enum import Enum
1818from typing import Any , TypeVar , cast , overload
1919import uuid
@@ -1503,7 +1503,7 @@ def _cancel_entity_batch(
15031503
15041504
15051505class _RuntimeOrchestrationContext (task .OrchestrationContext ):
1506- _generator : Generator [task .Task [Any ], Any , Any ] | None
1506+ _generator : Generator [task .Task [Any ], Any , Any ] | Coroutine [ Any , Any , Any ] | None
15071507 _previous_task : task .Task [Any ] | None
15081508
15091509 def __init__ (self ,
@@ -1549,21 +1549,26 @@ def __init__(self,
15491549 self ._exception_properties_provider = exception_properties_provider
15501550 self ._logger = logger
15511551
1552- def run (self , generator : Generator [task .Task [Any ], Any , Any ]) -> None :
1552+ def run (self , generator : Generator [task .Task [Any ], Any , Any ] | Coroutine [ Any , Any , Any ] ) -> None :
15531553 self ._generator = generator
1554- # TODO: Do something with this task
1555- task = next (generator ) # this starts the generator
1556- # TODO: Check if the task is null?
1557- self ._previous_task = task
1554+ self ._set_next_task (generator .send (None ))
1555+
1556+ def _set_next_task (self , next_task : Any ) -> None :
1557+ if not isinstance (next_task , task .Task ):
1558+ raise TypeError (
1559+ "Orchestrators must only yield or await durabletask.Task objects. "
1560+ "Use durable timers and task.when_all/when_any instead of asyncio APIs."
1561+ )
1562+ self ._previous_task = next_task
15581563
15591564 def resume (self ) -> None :
15601565 if self ._generator is None :
15611566 # This is never expected unless maybe there's an issue with the history
15621567 raise TypeError (
1563- "The orchestrator generator is not initialized! Was the orchestration history corrupted?"
1568+ "The orchestrator is not initialized! Was the orchestration history corrupted?"
15641569 )
15651570
1566- # We can resume the generator only if the previously yielded task
1571+ # We can resume the generator or coroutine only if the previous task
15671572 # has reached a completed state. The only time this won't be the
15681573 # case is if the user yielded on a WhenAll task and there are still
15691574 # outstanding child tasks that need to be completed.
@@ -1578,9 +1583,7 @@ def resume(self) -> None:
15781583 # This will either return a Task or raise StopIteration if it's done.
15791584 next_task = self ._generator .send (self ._previous_task .get_result ())
15801585
1581- if not isinstance (next_task , task .Task ):
1582- raise TypeError ("The orchestrator generator yielded a non-Task object" )
1583- self ._previous_task = next_task
1586+ self ._set_next_task (next_task )
15841587
15851588 def set_complete (
15861589 self ,
@@ -2521,10 +2524,15 @@ def process_event(
25212524
25222525 result = fn (
25232526 ctx , input
2524- ) # this does not execute the generator, only creates it
2525- if isinstance (result , GeneratorType ):
2526- # Start the orchestrator's generator function
2527- ctx .run (cast (Generator [task .Task [Any ], Any , Any ], result ))
2527+ ) # this creates the generator or coroutine without executing it
2528+ if isinstance (result , (GeneratorType , CoroutineType )):
2529+ # Both protocols suspend on durable tasks and finish with StopIteration.
2530+ ctx .run (cast (Generator [task .Task [Any ], Any , Any ] | Coroutine [Any , Any , Any ], result ))
2531+ elif inspect .isasyncgen (result ):
2532+ raise TypeError (
2533+ "Async generator orchestrators are not supported. "
2534+ "Use async def with await, not yield."
2535+ )
25282536 else :
25292537 # This is an orchestrator that doesn't schedule any tasks
25302538 ctx .set_complete (result , pb .ORCHESTRATION_STATUS_COMPLETED )
0 commit comments