Miscellaneous fixes

Change-Id: Ic545b230c7dd153e105c855b68c43fab4a010fa3

Author: kthota-g

examples/helloworld/agent_executor.py

@@ -1,28 +1,18 @@
-import asyncio
-
-from collections.abc import AsyncGenerator
-from typing import Any
-from uuid import uuid4
-
 from typing_extensions import override
+
 from a2a.server.agent_execution import AgentExecutor, RequestContext
-from a2a.server.events import Event, EventQueue
-from a2a.utils import new_agent_text_message
+from a2a.server.events import EventQueue
 from a2a.types import (
-    Message,
-    Part,
-    Role,
-    SendMessageRequest,
-    SendStreamingMessageRequest,
     Task,
-    TextPart,
 )
+from a2a.utils import new_agent_text_message
+
 
 class HelloWorldAgent:
     """Hello World Agent."""
 
     async def invoke(self) -> str:
-      return 'Hello World'
+        return 'Hello World'
 
 
 class HelloWorldAgentExecutor(AgentExecutor):
@@ -41,5 +31,7 @@ async def execute(
         event_queue.enqueue_event(new_agent_text_message(result))
 
     @override
-    async def cancel(self, request: RequestContext, event_queue: EventQueue) -> Task | None:
-        raise Exception("cancel not supported")
+    async def cancel(
+        self, request: RequestContext, event_queue: EventQueue
+    ) -> Task | None:
+        raise Exception('cancel not supported')

src/a2a/server/agent_execution/agent_executor.py

@@ -1,23 +1,16 @@
 from abc import ABC, abstractmethod
+
 from a2a.server.agent_execution.context import RequestContext
-from a2a.server.events.event_queue import Event, EventQueue
-from a2a.types import (
-    Task,
-    TaskStatusUpdateEvent,
-    TaskArtifactUpdateEvent,
-    Message
-)
+from a2a.server.events.event_queue import EventQueue
 
-class AgentExecutor(ABC):
 
-  @abstractmethod
-  async def execute(
-      self, context: RequestContext, event_queue: EventQueue
-  ):
-      pass
+class AgentExecutor(ABC):
+    """Agent Executor interface."""
 
-  @abstractmethod
-  async def cancel(
-      self, context: RequestContext, event_queue: EventQueue):
-      pass
+    @abstractmethod
+    async def execute(self, context: RequestContext, event_queue: EventQueue):
+        pass
 
+    @abstractmethod
+    async def cancel(self, context: RequestContext, event_queue: EventQueue):
+        pass

src/a2a/server/agent_execution/context.py

@@ -1,36 +1,38 @@
 import uuid
-from a2a.utils.errors import ServerError
+
 from a2a.types import (
-    Task,
+    InvalidParamsError,
     Message,
     MessageSendParams,
-    TextPart
+    Task,
+    TextPart,
 )
+from a2a.utils.errors import ServerError
+
 
 class RequestContext:
+    """Request Context."""
 
     def __init__(
         self,
         request: MessageSendParams | None = None,
         task_id: str | None = None,
         context_id: str | None = None,
         task: Task | None = None,
-        related_tasks: list[Task] = []
+        related_tasks: list[Task] = [],
     ):
         self._params = request
         self._task_id = task_id
         self._context_id = context_id
         self._current_task = task
-        self._related_tasks: related_tasks
+        self._related_tasks = related_tasks
         # If the task id and context id were provided, make sure they
         # match the request. Otherwise, create them
         if self._params:
             if task_id:
                 self._params.message.taskId = task_id
                 if task and task.id != task_id:
-                    raise ServerError(
-                        InvalidParamsError(message='bad task id')
-                    )
+                    raise ServerError(InvalidParamsError(message='bad task id'))
             else:
                 self._check_or_generate_task_id()
             if context_id:
@@ -43,6 +45,9 @@ def __init__(
                 self._check_or_generate_context_id()
 
     def get_user_input(self, delimiter='\n') -> str:
+        if not self._params:
+            return ''
+
         parts = []
         for part in self._params.message.parts:
             if isinstance(part.root, TextPart):
@@ -61,29 +66,35 @@ def related_tasks(self) -> list[Task]:
         return self._related_tasks
 
     @property
-    def current_task(self) -> Task:
+    def current_task(self) -> Task | None:
         return self._current_task
 
     @current_task.setter
-    def current_task(self, task: Task):
+    def current_task(self, task: Task) -> None:
         self._current_task = task
 
     @property
-    def task_id(self) -> str:
+    def task_id(self) -> str | None:
         return self._task_id
 
     @property
-    def context_id(self) -> str:
+    def context_id(self) -> str | None:
         return self._context_id
 
-    def _check_or_generate_task_id(self):
-      if not self._task_id and not self._params.message.taskId:
-          self._params.message.taskId = str(uuid.uuid4())
-      if self._params.message.taskId:
-          self._task_id = self._params.message.taskId
-
-    def _check_or_generate_context_id(self):
-      if not self._context_id and not self._params.message.contextId:
-          self._params.message.taskId = str(uuid.uuid4())
-      if self._params.message.contextId:
-          self._context_id = self._params.message.contextId
+    def _check_or_generate_task_id(self) -> None:
+        if not self._params:
+            return
+
+        if not self._task_id and not self._params.message.taskId:
+            self._params.message.taskId = str(uuid.uuid4())
+        if self._params.message.taskId:
+            self._task_id = self._params.message.taskId
+
+    def _check_or_generate_context_id(self) -> None:
+        if not self._params:
+            return
+
+        if not self._context_id and not self._params.message.contextId:
+            self._params.message.contextId = str(uuid.uuid4())
+        if self._params.message.contextId:
+            self._context_id = self._params.message.contextId

src/a2a/server/apps/http_app.py

@@ -1,8 +1,12 @@
 from abc import ABC, abstractmethod
+from typing import Any
+
 from starlette.applications import Starlette
 
+
 class HttpApp(ABC):
+    """A2A Server application interface."""
 
-  @abstractmethod
-  def build(self, **kwargs) -> Starlette:
-    pass
+    @abstractmethod
+    def build(self, **kwargs: Any) -> Starlette:
+        pass

src/a2a/server/apps/starlette_app.py

@@ -12,9 +12,10 @@
 from starlette.responses import JSONResponse, Response
 from starlette.routing import Route
 
-from a2a.utils.errors import MethodNotImplementedError
-from a2a.server.request_handlers.jsonrpc_handler import RequestHandler, JSONRPCHandler
-from a2a.server.events.event_queue import Event
+from a2a.server.request_handlers.jsonrpc_handler import (
+    JSONRPCHandler,
+    RequestHandler,
+)
 from a2a.types import (
     A2AError,
     A2ARequest,
@@ -35,6 +36,7 @@
     TaskResubscriptionRequest,
     UnsupportedOperationError,
 )
+from a2a.utils.errors import MethodNotImplementedError
 
 
 logger = logging.getLogger(__name__)
@@ -47,17 +49,17 @@ class A2AStarletteApplication:
     handler methods, and manages response generation including Server-Sent Events (SSE).
     """
 
-    def __init__(
-        self, agent_card: AgentCard, http_handler: RequestHandler
-    ):
+    def __init__(self, agent_card: AgentCard, http_handler: RequestHandler):
         """Initializes the A2AApplication.
 
         Args:
             agent_card: The AgentCard describing the agent's capabilities.
             http_handler: The handler instance responsible for processing A2A requests via http.
         """
         self.agent_card = agent_card
-        self.handler = JSONRPCHandler(agent_card=agent_card, request_handler=http_handler)
+        self.handler = JSONRPCHandler(
+            agent_card=agent_card, request_handler=http_handler
+        )
 
     def _generate_error_response(
         self, request_id: str | int | None, error: JSONRPCError | A2AError
@@ -154,7 +156,7 @@ async def _process_streaming_request(
         ):
             handler_result = self.handler.on_message_send_stream(request_obj)
         elif isinstance(request_obj, TaskResubscriptionRequest):
-            handler_result = self.handler.on_resubscribe(request_obj)
+            handler_result = self.handler.on_resubscribe_to_task(request_obj)
 
         return self._create_response(handler_result)
 
@@ -173,15 +175,17 @@ async def _process_non_streaming_request(
             case SendMessageRequest():
                 handler_result = await self.handler.on_message_send(request_obj)
             case CancelTaskRequest():
-                handler_result = await self.handler.on_cancel(request_obj)
+                handler_result = await self.handler.on_cancel_task(request_obj)
             case GetTaskRequest():
                 handler_result = await self.handler.on_get_task(request_obj)
             case SetTaskPushNotificationConfigRequest():
-                handler_result = await self.handler.on_set_push_notification(
-                    request_obj)
+                handler_result = await self.handler.set_push_notification(
+                    request_obj
+                )
             case GetTaskPushNotificationConfigRequest():
-                handler_result = await self.handler.on_get_push_notification(
-                        request_obj)
+                handler_result = await self.handler.get_push_notification(
+                    request_obj
+                )
             case _:
                 logger.error(
                     f'Unhandled validated request type: {type(request_obj)}'
@@ -218,7 +222,7 @@ def _create_response(
         if isinstance(handler_result, AsyncGenerator):
             # Result is a stream of SendStreamingMessageResponse objects
             async def event_generator(
-                stream: AsyncGenerator[Event, None],
+                stream: AsyncGenerator[SendStreamingMessageResponse, None],
             ) -> AsyncGenerator[dict[str, str], None]:
                 async for item in stream:
                     yield {'data': item.root.model_dump_json(exclude_none=True)}

src/a2a/server/events/event_consumer.py

@@ -2,18 +2,16 @@
 import logging
 
 from collections.abc import AsyncGenerator
-from a2a.utils.errors import ServerError
+
 from a2a.server.events.event_queue import Event, EventQueue
 from a2a.types import (
-    A2AError,
     InternalError,
-    JSONRPCError,
     Message,
     Task,
-    TaskArtifactUpdateEvent,
-    TaskStatusUpdateEvent,
     TaskState,
+    TaskStatusUpdateEvent,
 )
+from a2a.utils.errors import ServerError
 
 
 logger = logging.getLogger(__name__)
@@ -35,13 +33,14 @@ async def consume_one(self) -> Event:
             logger.warning('Event queue was empty in consume_one.')
             raise ServerError(
                 InternalError(message='Agent did not return any response')
-            )
+            ) from None
+
+        logger.debug(f'Dequeued event of type: {type(event)} in consume_one.')
 
         self.queue.task_done()
 
         return event
 
-
     async def consume_all(self) -> AsyncGenerator[Event]:
         """Consume all the generated streaming events from the agent."""
         logger.debug('Starting to consume all events from the queue.')
@@ -53,17 +52,22 @@ async def consume_all(self) -> AsyncGenerator[Event]:
                 )
                 yield event
                 self.queue.task_done()
-                logger.debug('Marked task as done in event queue in consume_all')
-
+                logger.debug(
+                    'Marked task as done in event queue in consume_all'
+                )
 
                 is_final_event = (
-                    (isinstance(event, TaskStatusUpdateEvent) and event.final) or
-                    isinstance(event, Message) or
-                    (isinstance(event, Task) and
-                     event.status.state in (
-                         TaskState.completed, TaskState.canceled, TaskState.failed
-                     )
-                     )
+                    (isinstance(event, TaskStatusUpdateEvent) and event.final)
+                    or isinstance(event, Message)
+                    or (
+                        isinstance(event, Task)
+                        and event.status.state
+                        in (
+                            TaskState.completed,
+                            TaskState.canceled,
+                            TaskState.failed,
+                        )
+                    )
                 )
 
                 if is_final_event:

src/a2a/server/events/event_queue.py

@@ -1,7 +1,7 @@
 import asyncio
 import logging
+
 from typing import Any
-from pydantic import RootModel
 
 from a2a.types import (
     A2AError,