From c4c5c1dfc8b32c2b29cbdfc282138019e036ea73 Mon Sep 17 00:00:00 2001
From: Holt Skinner <holtskinner@google.com>
Date: Tue, 12 Aug 2025 21:30:07 +0100
Subject: [PATCH 1/2] add fix for `TRY301` lint error

---
 .../default_request_handler.py                | 18 ++++----
 .../request_handlers/jsonrpc_handler.py       | 44 +++++++++++--------
 2 files changed, 34 insertions(+), 28 deletions(-)

diff --git a/src/a2a/server/request_handlers/default_request_handler.py b/src/a2a/server/request_handlers/default_request_handler.py
index 25fe25258..1e1dc2f77 100644
--- a/src/a2a/server/request_handlers/default_request_handler.py
+++ b/src/a2a/server/request_handlers/default_request_handler.py
@@ -283,7 +283,6 @@ async def on_message_send(
         blocking = True  # Default to blocking behavior
         if params.configuration and params.configuration.blocking is False:
             blocking = False
-
         interrupted_or_non_blocking = False
         try:
             (
@@ -292,15 +291,6 @@ async def on_message_send(
             ) = await result_aggregator.consume_and_break_on_interrupt(
                 consumer, blocking=blocking
             )
-            if not result:
-                raise ServerError(error=InternalError())  # noqa: TRY301
-
-            if isinstance(result, Task):
-                self._validate_task_id_match(task_id, result.id)
-
-            await self._send_push_notification_if_needed(
-                task_id, result_aggregator
-            )
 
         except Exception:
             logger.exception('Agent execution failed')
@@ -314,6 +304,14 @@ async def on_message_send(
             else:
                 await self._cleanup_producer(producer_task, task_id)
 
+        if not result:
+            raise ServerError(error=InternalError())
+
+        if isinstance(result, Task):
+            self._validate_task_id_match(task_id, result.id)
+
+        await self._send_push_notification_if_needed(task_id, result_aggregator)
+
         return result
 
     async def on_message_send_stream(
diff --git a/src/a2a/server/request_handlers/jsonrpc_handler.py b/src/a2a/server/request_handlers/jsonrpc_handler.py
index 3beb4e4f6..0aef47aa7 100644
--- a/src/a2a/server/request_handlers/jsonrpc_handler.py
+++ b/src/a2a/server/request_handlers/jsonrpc_handler.py
@@ -179,15 +179,6 @@ async def on_cancel_task(
             task = await self.request_handler.on_cancel_task(
                 request.params, context
             )
-            if task:
-                return prepare_response_object(
-                    request.id,
-                    task,
-                    (Task,),
-                    CancelTaskSuccessResponse,
-                    CancelTaskResponse,
-                )
-            raise ServerError(error=TaskNotFoundError())  # noqa: TRY301
         except ServerError as e:
             return CancelTaskResponse(
                 root=JSONRPCErrorResponse(
@@ -195,6 +186,19 @@ async def on_cancel_task(
                 )
             )
 
+        if task:
+            return prepare_response_object(
+                request.id,
+                task,
+                (Task,),
+                CancelTaskSuccessResponse,
+                CancelTaskResponse,
+            )
+
+        return CancelTaskResponse(
+            root=JSONRPCErrorResponse(id=request.id, error=TaskNotFoundError())
+        )
+
     async def on_resubscribe_to_task(
         self,
         request: TaskResubscriptionRequest,
@@ -331,15 +335,6 @@ async def on_get_task(
             task = await self.request_handler.on_get_task(
                 request.params, context
             )
-            if task:
-                return prepare_response_object(
-                    request.id,
-                    task,
-                    (Task,),
-                    GetTaskSuccessResponse,
-                    GetTaskResponse,
-                )
-            raise ServerError(error=TaskNotFoundError())  # noqa: TRY301
         except ServerError as e:
             return GetTaskResponse(
                 root=JSONRPCErrorResponse(
@@ -347,6 +342,19 @@ async def on_get_task(
                 )
             )
 
+        if task:
+            return prepare_response_object(
+                request.id,
+                task,
+                (Task,),
+                GetTaskSuccessResponse,
+                GetTaskResponse,
+            )
+
+        return GetTaskResponse(
+            root=JSONRPCErrorResponse(id=request.id, error=TaskNotFoundError())
+        )
+
     async def list_push_notification_config(
         self,
         request: ListTaskPushNotificationConfigRequest,

From 2e1417b6bd854bb25b76e29e600573f5498fb2c5 Mon Sep 17 00:00:00 2001
From: Holt Skinner <13262395+holtskinner@users.noreply.github.com>
Date: Wed, 3 Sep 2025 09:51:47 -0500
Subject: [PATCH 2/2] Update
 src/a2a/server/request_handlers/default_request_handler.py

---
 src/a2a/server/request_handlers/default_request_handler.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/a2a/server/request_handlers/default_request_handler.py b/src/a2a/server/request_handlers/default_request_handler.py
index 1e1dc2f77..47069c5f1 100644
--- a/src/a2a/server/request_handlers/default_request_handler.py
+++ b/src/a2a/server/request_handlers/default_request_handler.py
@@ -283,6 +283,7 @@ async def on_message_send(
         blocking = True  # Default to blocking behavior
         if params.configuration and params.configuration.blocking is False:
             blocking = False
+
         interrupted_or_non_blocking = False
         try:
             (