Harden async client error handling

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

Copilot-Session: 76f2163a-0773-4889-b627-1bd4e5c7bea0

Author: andystaples

azure-functions-durable/azure/durable_functions/decorators/durable_app.py

@@ -408,8 +408,10 @@ def attach_client_function(user_fn: Callable[..., Any]) -> FunctionBuilder:
                         if isinstance(user_fn, FunctionBuilder) else user_fn)
             signature = inspect.signature(function)
 
-            @wraps(function)
-            async def client_bound(*args: Any, **kwargs: Any) -> Any:
+            def bind_client(
+                    args: tuple[Any, ...],
+                    kwargs: dict[str, Any],
+            ) -> tuple[inspect.BoundArguments, DurableFunctionsClient | SyncDurableFunctionsClient]:
                 bound = signature.bind(*args, **kwargs)
                 if client_name not in bound.arguments:
                     raise TypeError(
@@ -422,15 +424,35 @@ async def client_bound(*args: Any, **kwargs: Any) -> Any:
                 client = (SyncDurableFunctionsClient.get_cached(raw_client) if sync
                           else DurableFunctionsClient(raw_client))
                 bound.arguments[client_name] = client
-                try:
-                    result = function(*bound.args, **bound.kwargs)
-                    return await result if inspect.isawaitable(result) else result
-                finally:
-                    if isinstance(client, DurableFunctionsClient):
-                        client.schedule_close()
-
-            client_bound.__annotations__[client_name] = str
-            setattr(client_bound, "client_function", function)
+                return bound, client
+
+            def set_client_metadata(client_bound: Callable[..., Any]) -> None:
+                annotations = dict(function.__annotations__)
+                annotations.setdefault(client_name, str)
+                client_bound.__annotations__ = annotations
+                setattr(client_bound, "client_function", function)
+
+            if inspect.iscoroutinefunction(function):
+                @wraps(function)
+                async def client_bound(*args: Any, **kwargs: Any) -> Any:
+                    bound, client = bind_client(args, kwargs)
+                    try:
+                        result = function(*bound.args, **bound.kwargs)
+                        return await result
+                    finally:
+                        if isinstance(client, DurableFunctionsClient):
+                            client.schedule_close()
+            else:
+                @wraps(function)
+                def client_bound(*args: Any, **kwargs: Any) -> Any:
+                    bound, client = bind_client(args, kwargs)
+                    try:
+                        return function(*bound.args, **bound.kwargs)
+                    finally:
+                        if isinstance(client, DurableFunctionsClient):
+                            client.schedule_close()
+
+            set_client_metadata(client_bound)
             if isinstance(user_fn, FunctionBuilder):
                 user_fn._function._func = client_bound  # pyright: ignore[reportPrivateUsage]
                 return wrap(user_fn)

tests/azure-functions-durable/test_decorator_compat.py

@@ -2,6 +2,7 @@
 # Licensed under the MIT License.
 
 import azure.durable_functions as df
+import inspect
 import pytest
 from azure.durable_functions.constants import (
     ACTIVITY_TRIGGER,
@@ -134,7 +135,7 @@ async def starter(client):
     assert await fb._function._func(client="{}") == "DurableFunctionsClient"
 
 
-async def test_durable_client_input_sync_injects_sync_client():
+def test_durable_client_input_sync_injects_sync_client():
     app = df.DFApp()
     clients = []
 
@@ -143,8 +144,9 @@ def starter(client):
         return type(client).__name__
 
     fb = app.durable_client_input_sync(client_name="client")(starter)
-    assert await fb._function._func(client="{}") == "SyncDurableFunctionsClient"
-    assert await fb._function._func(client="{}") == "SyncDurableFunctionsClient"
+    assert not inspect.iscoroutinefunction(fb._function._func)
+    assert fb._function._func(client="{}") == "SyncDurableFunctionsClient"
+    assert fb._function._func(client="{}") == "SyncDurableFunctionsClient"
     assert clients[0] is clients[1]
 
 
@@ -159,6 +161,17 @@ async def starter(other):
         await fb._function._func(other="{}")
 
 
+async def test_durable_client_input_preserves_client_annotation():
+    app = df.DFApp()
+
+    async def starter(client: df.DurableFunctionsClient):
+        return type(client).__name__
+
+    fb = app.durable_client_input(client_name="client")(starter)
+    assert starter.__annotations__["client"] is df.DurableFunctionsClient
+    assert await fb._function._func(client="{}") == "DurableFunctionsClient"
+
+
 # ---------------------------------------------------------------------------
 # All decorators register a function builder
 # ---------------------------------------------------------------------------

tests/durabletask/extensions/history_export/test_client.py

@@ -17,6 +17,7 @@
 from datetime import datetime, timedelta, timezone
 
 import pytest
+import grpc
 
 from durabletask import client, task, worker
 from durabletask.extensions.history_export import (
@@ -192,6 +193,47 @@ async def test_async_client_get_job_returns_none_for_empty_state(writer):
     assert await export_client.get_job("empty-state") is None
 
 
+async def test_async_wait_for_job_raises_lookup_when_job_never_exists(writer):
+    export_client = AsyncExportHistoryClient(MagicMock(), writer)
+    export_client.get_job = AsyncMock(return_value=None)
+
+    with pytest.raises(ExportJobNotFoundError):
+        await export_client.wait_for_job(
+            "never-created", timeout=0.01, poll_interval=0.001)
+
+
+async def test_async_wait_for_job_times_out_when_status_stays_active(writer):
+    export_client = AsyncExportHistoryClient(MagicMock(), writer)
+    export_client.get_job = AsyncMock(
+        return_value=MagicMock(status=ExportJobStatus.ACTIVE))
+
+    with pytest.raises(TimeoutError, match="did not reach a terminal status"):
+        await export_client.wait_for_job(
+            "stuck-job", timeout=0.01, poll_interval=0.001)
+
+
+class _NotFoundRpcError(grpc.RpcError):
+    def code(self):
+        return grpc.StatusCode.NOT_FOUND
+
+
+async def test_async_delete_job_ignores_not_found_errors(writer):
+    dt_client = MagicMock()
+    dt_client.signal_entity = AsyncMock()
+    dt_client.terminate_orchestration = AsyncMock(side_effect=_NotFoundRpcError())
+    dt_client.wait_for_orchestration_completion = AsyncMock(
+        side_effect=_NotFoundRpcError())
+    dt_client.purge_orchestration = AsyncMock(side_effect=_NotFoundRpcError())
+    export_client = AsyncExportHistoryClient(dt_client, writer)
+
+    await export_client.delete_job("missing-job")
+
+    dt_client.signal_entity.assert_awaited_once()
+    dt_client.terminate_orchestration.assert_awaited_once()
+    dt_client.wait_for_orchestration_completion.assert_awaited_once()
+    dt_client.purge_orchestration.assert_awaited_once()
+
+
 def test_get_job_returns_none_for_unknown_id(export_client):
     assert export_client.get_job("does-not-exist") is None
 

tests/durabletask/scheduled/test_client_filter_and_capability.py

@@ -4,9 +4,13 @@
 """Unit tests for ScheduleClient filtering and scheduled-tasks worker capability."""
 
 from datetime import datetime, timedelta, timezone
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, MagicMock
 
+import pytest
 import durabletask.internal.orchestrator_service_pb2 as pb
-from durabletask.scheduled.client import ScheduledTaskClient
+from durabletask.client import OrchestrationStatus
+from durabletask.scheduled.client import AsyncScheduleClient, ScheduledTaskClient
 from durabletask.scheduled.models import (ScheduleConfiguration,
                                           ScheduleCreationOptions, ScheduleQuery,
                                           ScheduleState)
@@ -76,8 +80,35 @@ def test_capability_absent_by_default(self):
         assert pb.WORKER_CAPABILITY_SCHEDULED_TASKS not in worker._capabilities  # pyright: ignore[reportPrivateUsage]
 
     def test_add_capability_rejected_while_running(self):
-        import pytest
         worker = TaskHubGrpcWorker()
         worker._is_running = True  # pyright: ignore[reportPrivateUsage]
         with pytest.raises(RuntimeError):
             worker.add_capability(pb.WORKER_CAPABILITY_SCHEDULED_TASKS)
+
+
+async def test_async_schedule_operation_raises_failure_details():
+    client = MagicMock()
+    client.schedule_new_orchestration = AsyncMock(return_value="operation-id")
+    client.wait_for_orchestration_completion = AsyncMock(return_value=SimpleNamespace(
+        runtime_status=OrchestrationStatus.FAILED,
+        failure_details=SimpleNamespace(message="operation failed"),
+    ))
+    schedule = AsyncScheduleClient(client, "schedule-id")
+
+    with pytest.raises(RuntimeError, match="operation failed"):
+        await schedule.create(ScheduleCreationOptions(
+            schedule_id="schedule-id",
+            orchestration_name="orchestration",
+            interval=timedelta(minutes=1),
+        ))
+
+
+async def test_async_schedule_operation_propagates_timeout():
+    client = MagicMock()
+    client.schedule_new_orchestration = AsyncMock(return_value="operation-id")
+    client.wait_for_orchestration_completion = AsyncMock(
+        side_effect=TimeoutError("operation timed out"))
+    schedule = AsyncScheduleClient(client, "schedule-id")
+
+    with pytest.raises(TimeoutError, match="operation timed out"):
+        await schedule.pause()