Merge branch 'main' into kthota/typechanges

Author: kthota-g

.github/actions/spelling/excludes.txt

@@ -81,8 +81,7 @@
 \.xz$
 \.zip$
 ^\.github/actions/spelling/
-^\Q.github/workflows/spelling.yaml\E$
-^\Q.github/workflows/linter.yaml\E$
+^\.github/workflows/
 \.gitignore\E$
 \.vscode/
 noxfile.py

.github/linters/.mypy.ini

@@ -0,0 +1,6 @@
+[mypy]
+exclude = examples/
+disable_error_code = import-not-found
+
+[mypy-examples.*]
+follow_imports = skip

.github/workflows/linter.yaml

@@ -13,10 +13,9 @@ name: Lint Code Base
 #############################
 # Start the job on all push #
 #############################
-# on:
-#   pull_request:
-#     branches: [main]
-on: workflow_dispatch
+on:
+  pull_request:
+    branches: [main]
 
 ###############
 # Set the Job #
@@ -64,3 +63,5 @@ jobs:
           VALIDATE_TYPESCRIPT_STANDARD: false
           VALIDATE_GIT_COMMITLINT: false
           MARKDOWN_CONFIG_FILE: .markdownlint.json
+          PYTHON_MYPY_CONFIG_FILE: .mypy.ini
+          FILTER_REGEX_EXCLUDE: "^examples/.*"

.github/workflows/python-publish.yml

@@ -24,7 +24,7 @@ jobs:
 
       - name: Build
         run: uv build
-      
+
       - name: Upload distributions
         uses: actions/upload-artifact@v4
         with:
@@ -49,6 +49,3 @@ jobs:
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
           packages-dir: dist/
-  
-        
-      

examples/google_adk/birthday_planner/adk_agent_executor.py

@@ -1,3 +1,4 @@
+# mypy: ignore-errors
 import asyncio
 import logging
 

examples/google_adk/calendar_agent/adk_agent_executor.py

@@ -1,3 +1,4 @@
+# mypy: ignore-errors
 import asyncio
 import logging
 
@@ -53,7 +54,7 @@ def __init__(self, runner: Runner, card: AgentCard):
 
     def _run_agent(
         self, session_id, new_message: types.Content
-    ) -> AsyncGenerator[Event, None]:
+    ) -> AsyncGenerator[Event]:
         return self.runner.run_async(
             session_id=session_id, user_id='self', new_message=new_message
         )

examples/langgraph/__main__.py

@@ -2,14 +2,15 @@
 import sys
 
 import click
+import httpx
 
 from agent import CurrencyAgent
 from agent_executor import CurrencyAgentExecutor
 from dotenv import load_dotenv
 
 from a2a.server.apps import A2AStarletteApplication
 from a2a.server.request_handlers import DefaultRequestHandler
-from a2a.server.tasks import InMemoryTaskStore
+from a2a.server.tasks import InMemoryPushNotifier, InMemoryTaskStore
 from a2a.types import (
     AgentAuthentication,
     AgentCapabilities,
@@ -29,9 +30,11 @@ def main(host: str, port: int):
         print('GOOGLE_API_KEY environment variable not set.')
         sys.exit(1)
 
+    client = httpx.AsyncClient()
     request_handler = DefaultRequestHandler(
         agent_executor=CurrencyAgentExecutor(),
         task_store=InMemoryTaskStore(),
+        push_notifier=InMemoryPushNotifier(client),
     )
 
     server = A2AStarletteApplication(

noxfile.py

@@ -114,9 +114,15 @@ def format(session):
         'pyupgrade',
         'autoflake',
         'ruff',
+        'no_implicit_optional',
     )
 
     if lint_paths_py:
+        session.run(
+            'no_implicit_optional',
+            '--use-union-or',
+            *lint_paths_py,
+        )
         if not format_all:
             session.run(
                 'pyupgrade',

src/a2a/server/agent_execution/context.py

@@ -3,8 +3,8 @@
 from a2a.types import (
     InvalidParamsError,
     Message,
-    MessageSendParams,
     MessageSendConfiguration,
+    MessageSendParams,
     Task,
 )
 from a2a.utils import get_message_text
@@ -82,6 +82,8 @@ def context_id(self) -> str | None:
 
     @property
     def configuration(self) -> MessageSendConfiguration | None:
+        if not self._params:
+            return None
         return self._params.configuration
 
     def _check_or_generate_task_id(self) -> None:

src/a2a/server/request_handlers/default_request_handler.py

@@ -14,11 +14,18 @@
     TaskQueueExists,
 )
 from a2a.server.request_handlers.request_handler import RequestHandler
-from a2a.server.tasks import ResultAggregator, TaskManager, TaskStore
+from a2a.server.tasks import (
+    PushNotifier,
+    ResultAggregator,
+    TaskManager,
+    TaskStore,
+)
 from a2a.types import (
     InternalError,
     Message,
+    MessageSendConfiguration,
     MessageSendParams,
+    PushNotificationConfig,
     Task,
     TaskIdParams,
     TaskNotFoundError,
@@ -44,10 +51,12 @@ def __init__(
         agent_executor: AgentExecutor,
         task_store: TaskStore,
         queue_manager: QueueManager | None = None,
+        push_notifier: PushNotifier | None = None,
     ) -> None:
         self.agent_executor = agent_executor
         self.task_store = task_store
         self._queue_manager = queue_manager or InMemoryQueueManager()
+        self._push_notifier = push_notifier
         # TODO: Likely want an interface for managing this, like AgentExecutionManager.
         self._running_agents = {}
         self._running_agents_lock = asyncio.Lock()
@@ -118,6 +127,18 @@ async def on_message_send(
         task: Task | None = await task_manager.get_task()
         if task:
             task = task_manager.update_with_message(params.message, task)
+            if self.should_add_push_info(params):
+                assert isinstance(self._push_notifier, PushNotifier)
+                assert isinstance(
+                    params.configuration, MessageSendConfiguration
+                )
+                assert isinstance(
+                    params.configuration.pushNotificationConfig,
+                    PushNotificationConfig,
+                )
+                await self._push_notifier.set_info(
+                    task.id, params.configuration.pushNotificationConfig
+                )
         request_context = RequestContext(
             params,
             task.id if task else None,
@@ -176,6 +197,20 @@ async def on_message_send_stream(
         if task:
             task = task_manager.update_with_message(params.message, task)
 
+            if self.should_add_push_info(params):
+                assert isinstance(self._push_notifier, PushNotifier)
+                assert isinstance(
+                    params.configuration, MessageSendConfiguration
+                )
+                assert isinstance(
+                    params.configuration.pushNotificationConfig,
+                    PushNotificationConfig,
+                )
+                await self._push_notifier.set_info(
+                    task.id, params.configuration.pushNotificationConfig
+                )
+        else:
+            queue = EventQueue()
         result_aggregator = ResultAggregator(task_manager)
         request_context = RequestContext(
             params,
@@ -202,12 +237,26 @@ async def on_message_send_stream(
                         f'Agent generated task_id={event.id} does not match the RequestContext task_id={task_id}.'
                     )
                     try:
-                        await self._queue_manager.add(event.id, queue)
-                        task_id = event.id
+                        created_task: Task = event
+                        await self._queue_manager.add(created_task.id, queue)
+                        task_id = created_task.id
                     except TaskQueueExists:
                         logging.info(
                             'Multiple Task objects created in event stream.'
                         )
+                    if (
+                        self._push_notifier
+                        and params.configuration
+                        and params.configuration.pushNotificationConfig
+                    ):
+                        await self._push_notifier.set_info(
+                            created_task.id,
+                            params.configuration.pushNotificationConfig,
+                        )
+                if self._push_notifier and task_id:
+                    latest_task = await result_aggregator.current_result
+                    if isinstance(latest_task, Task):
+                        await self._push_notifier.send_notification(latest_task)
                 yield event
         finally:
             await self._cleanup_producer(producer_task, task_id)
@@ -226,13 +275,38 @@ async def on_set_task_push_notification_config(
         self, params: TaskPushNotificationConfig
     ) -> TaskPushNotificationConfig:
         """Default handler for 'tasks/pushNotificationConfig/set'."""
-        raise ServerError(error=UnsupportedOperationError())
+        if not self._push_notifier:
+            raise ServerError(error=UnsupportedOperationError())
+
+        task: Task | None = await self.task_store.get(params.taskId)
+        if not task:
+            raise ServerError(error=TaskNotFoundError())
+
+        await self._push_notifier.set_info(
+            params.taskId,
+            params.pushNotificationConfig,
+        )
+
+        return params
 
     async def on_get_task_push_notification_config(
         self, params: TaskIdParams
     ) -> TaskPushNotificationConfig:
         """Default handler for 'tasks/pushNotificationConfig/get'."""
-        raise ServerError(error=UnsupportedOperationError())
+        if not self._push_notifier:
+            raise ServerError(error=UnsupportedOperationError())
+
+        task: Task | None = await self.task_store.get(params.id)
+        if not task:
+            raise ServerError(error=TaskNotFoundError())
+
+        push_notification_config = await self._push_notifier.get_info(params.id)
+        if not push_notification_config:
+            raise ServerError(error=InternalError())
+
+        return TaskPushNotificationConfig(
+            taskId=params.id, pushNotificationConfig=push_notification_config
+        )
 
     async def on_resubscribe_to_task(
         self, params: TaskIdParams
@@ -258,3 +332,10 @@ async def on_resubscribe_to_task(
         consumer = EventConsumer(queue)
         async for event in result_aggregator.consume_and_emit(consumer):
             yield event
+
+    def should_add_push_info(self, params: MessageSendParams) -> bool:
+        return bool(
+            self._push_notifier
+            and params.configuration
+            and params.configuration.pushNotificationConfig
+        )