fix(server): run after_handler_run when a handler call is cancelled

src/a2a/server/request_handlers/_plugin_hooks.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import functools
 import inspect
 
@@ -82,7 +83,7 @@ async def stream_wrapper(
                                 )
                             )
                             yield event
-                except Exception as e:
+                except (Exception, asyncio.CancelledError) as e:
                     error = e
                     raise
                 finally:
@@ -110,7 +111,8 @@ async def unary_wrapper(
             )
             try:
                 result = await func(self, params, context, *args, **kwargs)
-            except Exception as e:
+            except (Exception, asyncio.CancelledError) as e:
+                # A cancelled call still gets its after call, as streams do.
                 await plugins.after_handler_run(
                     ServerResponse(method, None, context, error=e)
                 )

tests/server/request_handlers/test_default_request_handler_v2_plugins.py

@@ -50,6 +50,7 @@
     Role,
     SendMessageConfiguration,
     SendMessageRequest,
+    SubscribeToTaskRequest,
     Task,
     TaskPushNotificationConfig,
     TaskState,
@@ -353,6 +354,64 @@ async def test_stream_error_reaches_final_call():
     assert isinstance(final.error, TaskNotFoundError)
 
 
+class BlockingTaskStore(InMemoryTaskStore):
+    """Blocks every `get` until cancelled."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.waiting = asyncio.Event()
+
+    async def get(self, task_id, context):
+        self.waiting.set()
+        await asyncio.Event().wait()
+
+
+@pytest.mark.asyncio
+async def test_cancelled_call_reaches_after_handler_run():
+    recorder = Recorder()
+    store = BlockingTaskStore()
+    handler = make_handler(recorder, task_store=store)
+
+    call = asyncio.create_task(
+        handler.on_get_task(GetTaskRequest(id='t-1'), ServerCallContext())
+    )
+    await asyncio.wait_for(store.waiting.wait(), timeout=5)
+    call.cancel()
+
+    with pytest.raises(asyncio.CancelledError):
+        await call
+
+    (response,) = recorder.hooks('after_handler_run')
+    assert response.method == 'get_task'
+    assert isinstance(response.error, asyncio.CancelledError)
+    assert response.is_final
+
+
+@pytest.mark.asyncio
+async def test_cancelled_stream_reports_the_cancellation():
+    recorder = Recorder()
+    store = BlockingTaskStore()
+    handler = make_handler(recorder, task_store=store)
+
+    async def consume() -> None:
+        async for _ in handler.on_subscribe_to_task(
+            SubscribeToTaskRequest(id='t-1'), ServerCallContext()
+        ):
+            pass
+
+    reader = asyncio.create_task(consume())
+    await asyncio.wait_for(store.waiting.wait(), timeout=5)
+    reader.cancel()
+
+    with pytest.raises(asyncio.CancelledError):
+        await reader
+
+    (response,) = recorder.hooks('after_handler_run')
+    assert response.method == 'subscribe'
+    assert isinstance(response.error, asyncio.CancelledError)
+    assert response.is_final
+
+
 @pytest.mark.asyncio
 async def test_before_handler_run_can_reject_the_request():
     class Deny(plugins.ServerPlugin):