Add context manager support to TaskHubGrpcClient (#145)

* Add context manager support to TaskHubGrpcClient (#134)

Add `__enter__`/`__exit__` to the sync `TaskHubGrpcClient` so callers can
use it with a `with` statement, mirroring the existing
`AsyncTaskHubGrpcClient` async-context-manager support and the
`TaskHubGrpcWorker` pattern. `DurableTaskSchedulerClient` inherits this
behavior automatically.

`__exit__` delegates to `close()`, so the resiliency-aware teardown
introduced in #135 (in-flight recreate thread join, retired-channel
timer cancellation, SDK-owned channel cleanup) runs unchanged through
the new `with` path. Caller-owned channels remain untouched.

Migrate every test and example callsite that previously instantiated
`TaskHubGrpcClient(...)` and never closed it to the `with` form so the
gRPC channel is deterministically released. Unit tests in
`test_client.py` that intentionally test construction (with mocked
stubs) are left unchanged.

Add focused unit tests for the new context-manager behavior, including
a regression test that exits a `with` block while a fire-and-forget
channel recreate is pending and asserts the #135 resiliency invariants
(retired-channel timers cancelled, recreate thread joined) still hold.

Fixes #134

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Address auto-generated code-quality review comments on PR #145

1. tests/durabletask/test_orchestration_e2e.py (test_suspend_and_resume):
   Drop the unused `state =` assignment around the expected-timeout
   `wait_for_orchestration_completion` call. The value is never read
   because the next line asserts False and the only non-failing path
   raises TimeoutError; `state` is reassigned a few lines down. Silences
   the "variable defined multiple times" warning that CodeQL flagged
   because this previously-untouched line was pulled into the diff by
   the indent change.

2. tests/durabletask/test_client.py
   (test_sync_client_context_manager_propagates_exception_and_calls_close):
   Replace the nested `with pytest.raises(...): with client: raise ...`
   pattern with an explicit try/except so CodeQL no longer reports the
   post-block assertions as unreachable. Test intent (exception
   propagation + cleanup verification) is preserved.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Align __exit__ typing with repo convention (BlobPayloadStore)

Use `*args: object` for the `__exit__` parameters instead of leaving
them untyped. This matches the most recent context-manager class in
the repo (`durabletask/extensions/azure_blob_payloads/blob_payload_store.py`),
is more type-safe under Pylance/mypy, and avoids the parameter
shadowing of the builtin `type` that exists in
`TaskHubGrpcWorker.__exit__`. Behavior is unchanged.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Align pre-existing context manager methods with repo idiom

Apply the same `__enter__`/`__exit__` typing pattern that PR #145 uses
for `TaskHubGrpcClient` (and that `BlobPayloadStore` already follows)
to three pre-existing context managers that were either untyped or
shadowed a builtin:

* `AsyncTaskHubGrpcClient.__aenter__` now returns the concrete type and
  `__aexit__` takes `*args: object` with `-> None`.
* `TaskHubGrpcWorker.__enter__/__exit__` get the same treatment, also
  removing the `type` parameter that shadowed the builtin.
* `EntityLock.__enter__/__exit__` get the same treatment; the file
  already has `from __future__ import annotations` so the return
  annotation is the bare class name.

Behavior is unchanged: each `__exit__` still delegates to its existing
teardown method (`close`/`stop`/`_exit_critical_section`), so the
gRPC resiliency teardown added in PR #135 continues to flow through
`TaskHubGrpcClient.close()` unchanged.

No changelog entry: per the repo's contributor guidance, internal-only
type-annotation refactors with no externally observable behavior change
are excluded from the changelog.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

---------

Co-authored-by: Bernd Verst <beverst@microsoft.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: 3 people

CHANGELOG.md

@@ -7,6 +7,17 @@ adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## Unreleased
 
+ADDED
+
+- Added context-manager support (`__enter__` / `__exit__`) to
+  `TaskHubGrpcClient` so it can be used with `with` statements, mirroring the
+  existing `AsyncTaskHubGrpcClient` async-context-manager support and the
+  `TaskHubGrpcWorker` pattern. `DurableTaskSchedulerClient` inherits this
+  behavior automatically. `__exit__` delegates to `close()`, so the
+  resiliency-aware teardown introduced in v1.5.0 (in-flight recreate
+  thread join, retired-channel timer cancellation, and SDK-owned channel
+  cleanup) runs unchanged through the new `with` path.
+
 ## v1.5.0
 
 ADDED

durabletask/client.py

@@ -366,6 +366,12 @@ def close(self) -> None:
                 retired_channel.close()
             current_channel.close()
 
+    def __enter__(self) -> "TaskHubGrpcClient":
+        return self
+
+    def __exit__(self, *args: object) -> None:
+        self.close()
+
     def schedule_new_orchestration(self, orchestrator: task.Orchestrator[TInput, TOutput] | str, *,
                                    input: TInput | None = None,
                                    instance_id: str | None = None,
@@ -783,10 +789,10 @@ async def close(self) -> None:
                 await retired_channel.close()
             await self._channel.close()
 
-    async def __aenter__(self):
+    async def __aenter__(self) -> "AsyncTaskHubGrpcClient":
         return self
 
-    async def __aexit__(self, exc_type, exc_val, exc_tb):
+    async def __aexit__(self, *args: object) -> None:
         await self.close()
 
     def _schedule_recreate(self) -> None:

durabletask/entities/entity_lock.py

@@ -12,8 +12,8 @@ class EntityLock:
     def __init__(self, context: OrchestrationContext):
         self._context = context
 
-    def __enter__(self):
+    def __enter__(self) -> EntityLock:
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
+    def __exit__(self, *args: object) -> None:
         self._context._exit_critical_section()

durabletask/worker.py

@@ -581,10 +581,10 @@ def maximum_timer_interval(self) -> timedelta | None:
         """Get the configured maximum timer interval for long timer chunking."""
         return self._maximum_timer_interval
 
-    def __enter__(self):
+    def __enter__(self) -> "TaskHubGrpcWorker":
         return self
 
-    def __exit__(self, type, value, traceback):
+    def __exit__(self, *args: object) -> None:
         self.stop()
 
     def _classify_stream_outcome(

examples/human_interaction.py

@@ -79,31 +79,31 @@ def purchase_order_workflow(ctx: task.OrchestrationContext, order: Order):
             w.add_activity(place_order)
             w.start()
 
-            c = client.TaskHubGrpcClient()
-
-            # Start a purchase order workflow using the user input
-            order = Order(args.cost, "MyProduct", 1)
-            instance_id = c.schedule_new_orchestration(purchase_order_workflow, input=order)
-
-            def prompt_for_approval():
-                input("Press [ENTER] to approve the order...\n")
-                approval_event = namedtuple("Approval", ["approver"])(args.approver)
-                c.raise_orchestration_event(instance_id, "approval_received", data=approval_event)
-
-            # Prompt the user for approval on a background thread
-            threading.Thread(target=prompt_for_approval, daemon=True).start()
-
-            # Wait for the orchestration to complete
-            try:
-                state = c.wait_for_orchestration_completion(instance_id, timeout=args.timeout + 2)
-                if not state:
-                    print("Workflow not found!")  # not expected
-                elif state.runtime_status == client.OrchestrationStatus.COMPLETED:
-                    print(f'Orchestration completed! Result: {state.serialized_output}')
-                else:
-                    state.raise_if_failed()  # raises an exception
-            except TimeoutError:
-                print("*** Orchestration timed out!")
+            with client.TaskHubGrpcClient() as c:
+
+                # Start a purchase order workflow using the user input
+                order = Order(args.cost, "MyProduct", 1)
+                instance_id = c.schedule_new_orchestration(purchase_order_workflow, input=order)
+
+                def prompt_for_approval():
+                    input("Press [ENTER] to approve the order...\n")
+                    approval_event = namedtuple("Approval", ["approver"])(args.approver)
+                    c.raise_orchestration_event(instance_id, "approval_received", data=approval_event)
+
+                # Prompt the user for approval on a background thread
+                threading.Thread(target=prompt_for_approval, daemon=True).start()
+
+                # Wait for the orchestration to complete
+                try:
+                    state = c.wait_for_orchestration_completion(instance_id, timeout=args.timeout + 2)
+                    if not state:
+                        print("Workflow not found!")  # not expected
+                    elif state.runtime_status == client.OrchestrationStatus.COMPLETED:
+                        print(f'Orchestration completed! Result: {state.serialized_output}')
+                    else:
+                        state.raise_if_failed()  # raises an exception
+                except TimeoutError:
+                    print("*** Orchestration timed out!")
     else:
         # Use DurableTaskScheduler
         # Use environment variables if provided, otherwise use default emulator values

examples/in_memory_backend_example/test/test_workflows.py

@@ -72,9 +72,9 @@ def test_single_item_order(self):
 
         with _create_worker() as w:
             w.start()
-            c = client.TaskHubGrpcClient(host_address=HOST)
-            instance_id = c.schedule_new_orchestration(process_order, input=order)
-            state = c.wait_for_orchestration_completion(instance_id, timeout=30)
+            with client.TaskHubGrpcClient(host_address=HOST) as c:
+                instance_id = c.schedule_new_orchestration(process_order, input=order)
+                state = c.wait_for_orchestration_completion(instance_id, timeout=30)
 
         assert state is not None
         assert state.runtime_status == client.OrchestrationStatus.COMPLETED
@@ -98,9 +98,9 @@ def test_multi_item_order(self):
 
         with _create_worker() as w:
             w.start()
-            c = client.TaskHubGrpcClient(host_address=HOST)
-            instance_id = c.schedule_new_orchestration(process_order, input=order)
-            state = c.wait_for_orchestration_completion(instance_id, timeout=30)
+            with client.TaskHubGrpcClient(host_address=HOST) as c:
+                instance_id = c.schedule_new_orchestration(process_order, input=order)
+                state = c.wait_for_orchestration_completion(instance_id, timeout=30)
 
         assert state is not None
         assert state.runtime_status == client.OrchestrationStatus.COMPLETED
@@ -118,9 +118,9 @@ def test_empty_order_fails(self):
 
         with _create_worker() as w:
             w.start()
-            c = client.TaskHubGrpcClient(host_address=HOST)
-            instance_id = c.schedule_new_orchestration(process_order, input=order)
-            state = c.wait_for_orchestration_completion(instance_id, timeout=30)
+            with client.TaskHubGrpcClient(host_address=HOST) as c:
+                instance_id = c.schedule_new_orchestration(process_order, input=order)
+                state = c.wait_for_orchestration_completion(instance_id, timeout=30)
 
         assert state is not None
         assert state.runtime_status == client.OrchestrationStatus.FAILED
@@ -136,9 +136,9 @@ def test_invalid_quantity_fails(self):
 
         with _create_worker() as w:
             w.start()
-            c = client.TaskHubGrpcClient(host_address=HOST)
-            instance_id = c.schedule_new_orchestration(process_order, input=order)
-            state = c.wait_for_orchestration_completion(instance_id, timeout=30)
+            with client.TaskHubGrpcClient(host_address=HOST) as c:
+                instance_id = c.schedule_new_orchestration(process_order, input=order)
+                state = c.wait_for_orchestration_completion(instance_id, timeout=30)
 
         assert state is not None
         assert state.runtime_status == client.OrchestrationStatus.FAILED
@@ -163,9 +163,9 @@ def test_low_value_auto_approved(self):
 
         with _create_worker() as w:
             w.start()
-            c = client.TaskHubGrpcClient(host_address=HOST)
-            instance_id = c.schedule_new_orchestration(order_with_approval, input=order)
-            state = c.wait_for_orchestration_completion(instance_id, timeout=30)
+            with client.TaskHubGrpcClient(host_address=HOST) as c:
+                instance_id = c.schedule_new_orchestration(order_with_approval, input=order)
+                state = c.wait_for_orchestration_completion(instance_id, timeout=30)
 
         assert state is not None
         assert state.runtime_status == client.OrchestrationStatus.COMPLETED
@@ -184,12 +184,12 @@ def test_high_value_approved(self):
 
         with _create_worker() as w:
             w.start()
-            c = client.TaskHubGrpcClient(host_address=HOST)
-            instance_id = c.schedule_new_orchestration(order_with_approval, input=order)
+            with client.TaskHubGrpcClient(host_address=HOST) as c:
+                instance_id = c.schedule_new_orchestration(order_with_approval, input=order)
 
-            # Raise the approval event
-            c.raise_orchestration_event(instance_id, "approval", data=True)
-            state = c.wait_for_orchestration_completion(instance_id, timeout=30)
+                # Raise the approval event
+                c.raise_orchestration_event(instance_id, "approval", data=True)
+                state = c.wait_for_orchestration_completion(instance_id, timeout=30)
 
         assert state is not None
         assert state.runtime_status == client.OrchestrationStatus.COMPLETED
@@ -208,12 +208,12 @@ def test_high_value_rejected(self):
 
         with _create_worker() as w:
             w.start()
-            c = client.TaskHubGrpcClient(host_address=HOST)
-            instance_id = c.schedule_new_orchestration(order_with_approval, input=order)
+            with client.TaskHubGrpcClient(host_address=HOST) as c:
+                instance_id = c.schedule_new_orchestration(order_with_approval, input=order)
 
-            # Reject the order
-            c.raise_orchestration_event(instance_id, "approval", data=False)
-            state = c.wait_for_orchestration_completion(instance_id, timeout=30)
+                # Reject the order
+                c.raise_orchestration_event(instance_id, "approval", data=False)
+                state = c.wait_for_orchestration_completion(instance_id, timeout=30)
 
         assert state is not None
         assert state.runtime_status == client.OrchestrationStatus.COMPLETED
@@ -232,11 +232,11 @@ def test_high_value_timeout(self):
 
         with _create_worker() as w:
             w.start()
-            c = client.TaskHubGrpcClient(host_address=HOST)
-            instance_id = c.schedule_new_orchestration(order_with_approval, input=order)
+            with client.TaskHubGrpcClient(host_address=HOST) as c:
+                instance_id = c.schedule_new_orchestration(order_with_approval, input=order)
 
-            # Don't raise any event — let the timer fire
-            state = c.wait_for_orchestration_completion(instance_id, timeout=60)
+                # Don't raise any event — let the timer fire
+                state = c.wait_for_orchestration_completion(instance_id, timeout=60)
 
         assert state is not None
         assert state.runtime_status == client.OrchestrationStatus.COMPLETED

tests/durabletask/entities/test_class_based_entities_e2e.py

@@ -37,10 +37,10 @@ def do_nothing(self, _):
         w.add_entity(EmptyEntity, name="EntityNameCustom")
         w.start()
 
-        c = client.TaskHubGrpcClient(host_address=HOST)
-        entity_id = entities.EntityInstanceId("EntityNameCustom", "testEntity")
-        c.signal_entity(entity_id, "do_nothing")
-        time.sleep(2)  # wait for the signal to be processed
+        with client.TaskHubGrpcClient(host_address=HOST) as c:
+            entity_id = entities.EntityInstanceId("EntityNameCustom", "testEntity")
+            c.signal_entity(entity_id, "do_nothing")
+            time.sleep(2)  # wait for the signal to be processed
 
     assert invoked
 
@@ -59,14 +59,14 @@ def do_nothing(self, _):
         w.add_entity(EmptyEntity)
         w.start()
 
-        c = client.TaskHubGrpcClient(host_address=HOST)
-        entity_id = entities.EntityInstanceId("EmptyEntity", "testEntity")
-        c.signal_entity(entity_id, "do_nothing")
-        time.sleep(2)  # wait for the signal to be processed
-        state = c.get_entity(entity_id, include_state=True)
-        assert state is not None
-        assert state.id == entity_id
-        assert state.get_state(int) == 1
+        with client.TaskHubGrpcClient(host_address=HOST) as c:
+            entity_id = entities.EntityInstanceId("EmptyEntity", "testEntity")
+            c.signal_entity(entity_id, "do_nothing")
+            time.sleep(2)  # wait for the signal to be processed
+            state = c.get_entity(entity_id, include_state=True)
+            assert state is not None
+            assert state.id == entity_id
+            assert state.get_state(int) == 1
 
     assert invoked
 
@@ -89,10 +89,10 @@ def empty_orchestrator(ctx: task.OrchestrationContext, _):
         w.add_entity(EmptyEntity, name="EntityNameCustom")
         w.start()
 
-        c = client.TaskHubGrpcClient(host_address=HOST)
-        id = c.schedule_new_orchestration(empty_orchestrator)
-        state = c.wait_for_orchestration_completion(id, timeout=30)
-        time.sleep(2)  # wait for the signal to be processed
+        with client.TaskHubGrpcClient(host_address=HOST) as c:
+            id = c.schedule_new_orchestration(empty_orchestrator)
+            state = c.wait_for_orchestration_completion(id, timeout=30)
+            time.sleep(2)  # wait for the signal to be processed
 
     assert invoked
     assert state is not None
@@ -123,9 +123,9 @@ def empty_orchestrator(ctx: task.OrchestrationContext, _):
         w.add_entity(EmptyEntity)
         w.start()
 
-        c = client.TaskHubGrpcClient(host_address=HOST)
-        id = c.schedule_new_orchestration(empty_orchestrator)
-        state = c.wait_for_orchestration_completion(id, timeout=30)
+        with client.TaskHubGrpcClient(host_address=HOST) as c:
+            id = c.schedule_new_orchestration(empty_orchestrator)
+            state = c.wait_for_orchestration_completion(id, timeout=30)
 
     assert invoked
     assert state is not None

tests/durabletask/entities/test_entity_failure_handling.py

@@ -39,9 +39,9 @@ def test_orchestrator(ctx: task.OrchestrationContext, _):
         w.add_entity(FailingEntity)
         w.start()
 
-        c = client.TaskHubGrpcClient(host_address=HOST)
-        id = c.schedule_new_orchestration(test_orchestrator)
-        state = c.wait_for_orchestration_completion(id, timeout=30)
+        with client.TaskHubGrpcClient(host_address=HOST) as c:
+            id = c.schedule_new_orchestration(test_orchestrator)
+            state = c.wait_for_orchestration_completion(id, timeout=30)
 
     assert state is not None
     assert state.name == task.get_name(test_orchestrator)
@@ -66,9 +66,9 @@ def test_orchestrator(ctx: task.OrchestrationContext, _):
         w.add_entity(failing_entity)
         w.start()
 
-        c = client.TaskHubGrpcClient(host_address=HOST)
-        id = c.schedule_new_orchestration(test_orchestrator)
-        state = c.wait_for_orchestration_completion(id, timeout=30)
+        with client.TaskHubGrpcClient(host_address=HOST) as c:
+            id = c.schedule_new_orchestration(test_orchestrator)
+            state = c.wait_for_orchestration_completion(id, timeout=30)
 
     assert state is not None
     assert state.name == task.get_name(test_orchestrator)
@@ -97,9 +97,9 @@ def test_orchestrator(ctx: task.OrchestrationContext, _):
         w.add_entity(FailingEntity)
         w.start()
 
-        c = client.TaskHubGrpcClient(host_address=HOST)
-        id = c.schedule_new_orchestration(test_orchestrator)
-        state = c.wait_for_orchestration_completion(id, timeout=30)
+        with client.TaskHubGrpcClient(host_address=HOST) as c:
+            id = c.schedule_new_orchestration(test_orchestrator)
+            state = c.wait_for_orchestration_completion(id, timeout=30)
 
     assert state is not None
     assert state.name == task.get_name(test_orchestrator)
@@ -129,9 +129,9 @@ def test_orchestrator(ctx: task.OrchestrationContext, _):
         w.add_entity(failing_entity)
         w.start()
 
-        c = client.TaskHubGrpcClient(host_address=HOST)
-        id = c.schedule_new_orchestration(test_orchestrator)
-        state = c.wait_for_orchestration_completion(id, timeout=30)
+        with client.TaskHubGrpcClient(host_address=HOST) as c:
+            id = c.schedule_new_orchestration(test_orchestrator)
+            state = c.wait_for_orchestration_completion(id, timeout=30)
 
     assert state is not None
     assert state.name == task.get_name(test_orchestrator)
@@ -168,9 +168,9 @@ def test_orchestrator(ctx: task.OrchestrationContext, _):
         w.add_entity(failing_entity)
         w.start()
 
-        c = client.TaskHubGrpcClient(host_address=HOST)
-        id = c.schedule_new_orchestration(test_orchestrator)
-        state = c.wait_for_orchestration_completion(id, timeout=30)
+        with client.TaskHubGrpcClient(host_address=HOST) as c:
+            id = c.schedule_new_orchestration(test_orchestrator)
+            state = c.wait_for_orchestration_completion(id, timeout=30)
 
     assert state is not None
     assert state.name == task.get_name(test_orchestrator)