From 17624b06aad01f4c48ca93305753ce601ce9a2e0 Mon Sep 17 00:00:00 2001 From: Wondr Date: Sun, 25 Jan 2026 17:04:48 +0100 Subject: [PATCH 1/4] feat: add auth header to push notifications --- .../tasks/base_push_notification_sender.py | 31 ++++++++++++-- .../tasks/test_push_notification_sender.py | 40 ++++++++++++++++++- 2 files changed, 67 insertions(+), 4 deletions(-) diff --git a/src/a2a/server/tasks/base_push_notification_sender.py b/src/a2a/server/tasks/base_push_notification_sender.py index 087d2973d..004546e5b 100644 --- a/src/a2a/server/tasks/base_push_notification_sender.py +++ b/src/a2a/server/tasks/base_push_notification_sender.py @@ -52,9 +52,7 @@ async def _dispatch_notification( ) -> bool: url = push_info.url try: - headers = None - if push_info.token: - headers = {'X-A2A-Notification-Token': push_info.token} + headers = self._build_headers(push_info) response = await self._client.post( url, json=task.model_dump(mode='json', exclude_none=True), @@ -72,3 +70,30 @@ async def _dispatch_notification( ) return False return True + + @staticmethod + def _authorization_header( + push_info: PushNotificationConfig, + ) -> str | None: + auth = push_info.authentication + if not auth or not auth.credentials: + return None + schemes = [scheme for scheme in auth.schemes if scheme] + if not schemes: + return None + scheme = next( + (scheme for scheme in schemes if scheme.lower() == 'bearer'), + schemes[0], + ) + return f'{scheme} {auth.credentials}' + + def _build_headers( + self, push_info: PushNotificationConfig + ) -> dict[str, str] | None: + headers: dict[str, str] = {} + if push_info.token: + headers['X-A2A-Notification-Token'] = push_info.token + authorization = self._authorization_header(push_info) + if authorization: + headers['Authorization'] = authorization + return headers or None diff --git a/tests/server/tasks/test_push_notification_sender.py b/tests/server/tasks/test_push_notification_sender.py index a3272c2c1..1068efd75 100644 --- a/tests/server/tasks/test_push_notification_sender.py +++ b/tests/server/tasks/test_push_notification_sender.py @@ -8,6 +8,7 @@ BasePushNotificationSender, ) from a2a.types import ( + PushNotificationAuthenticationInfo, PushNotificationConfig, Task, TaskState, @@ -29,8 +30,14 @@ def create_sample_push_config( url: str = 'http://example.com/callback', config_id: str = 'cfg1', token: str | None = None, + authentication: PushNotificationAuthenticationInfo | None = None, ) -> PushNotificationConfig: - return PushNotificationConfig(id=config_id, url=url, token=token) + return PushNotificationConfig( + id=config_id, + url=url, + token=token, + authentication=authentication, + ) class TestBasePushNotificationSender(unittest.IsolatedAsyncioTestCase): @@ -92,6 +99,37 @@ async def test_send_notification_with_token_success(self) -> None: ) mock_response.raise_for_status.assert_called_once() + async def test_send_notification_with_auth_header(self) -> None: + task_id = 'task_send_auth' + task_data = create_sample_task(task_id=task_id) + auth = PushNotificationAuthenticationInfo( + schemes=['Basic', 'Bearer'], credentials='token_or_jwt' + ) + config = create_sample_push_config( + url='http://notify.me/here', + token='unique_token', + authentication=auth, + ) + self.mock_config_store.get_info.return_value = [config] + + mock_response = AsyncMock(spec=httpx.Response) + mock_response.status_code = 200 + self.mock_httpx_client.post.return_value = mock_response + + await self.sender.send_notification(task_data) + + self.mock_config_store.get_info.assert_awaited_once_with + + self.mock_httpx_client.post.assert_awaited_once_with( + config.url, + json=task_data.model_dump(mode='json', exclude_none=True), + headers={ + 'X-A2A-Notification-Token': 'unique_token', + 'Authorization': 'Bearer token_or_jwt', + }, + ) + mock_response.raise_for_status.assert_called_once() + async def test_send_notification_no_config(self) -> None: task_id = 'task_send_no_config' task_data = create_sample_task(task_id=task_id) From 1267f82a3a0a7a3d7e784fe047bca1bbf3944863 Mon Sep 17 00:00:00 2001 From: Wondr Date: Sun, 25 Jan 2026 17:12:42 +0100 Subject: [PATCH 2/4] fix: adjust telemetry no-op context typing --- src/a2a/utils/telemetry.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/a2a/utils/telemetry.py b/src/a2a/utils/telemetry.py index c73d2ac92..77d2b3e9f 100644 --- a/src/a2a/utils/telemetry.py +++ b/src/a2a/utils/telemetry.py @@ -62,6 +62,12 @@ def internal_method(self): from typing import TYPE_CHECKING, Any +try: + from typing import Self +except ImportError: # pragma: no cover - for Python < 3.11 + from typing_extensions import Self + + if TYPE_CHECKING: from opentelemetry.trace import SpanKind as SpanKindType else: @@ -86,7 +92,7 @@ class _NoOp: def __call__(self, *args: Any, **kwargs: Any) -> Any: return self - def __enter__(self) -> '_NoOp': + def __enter__(self) -> Self: return self def __exit__(self, *args: object, **kwargs: Any) -> None: From 7496dd5eee9bb60969e06510e58773b2f82a11be Mon Sep 17 00:00:00 2001 From: Wondr Date: Sun, 25 Jan 2026 17:22:27 +0100 Subject: [PATCH 3/4] fix: use typing_extensions Self for py310 --- pyproject.toml | 1 + src/a2a/utils/telemetry.py | 9 ++++++--- uv.lock | 4 +++- 3 files changed, 10 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 561a5a45c..8dc564353 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,6 +13,7 @@ dependencies = [ "pydantic>=2.11.3", "protobuf>=5.29.5", "google-api-core>=1.26.0", + "typing-extensions>=4.0.0", ] classifiers = [ diff --git a/src/a2a/utils/telemetry.py b/src/a2a/utils/telemetry.py index 77d2b3e9f..99fa6b1e8 100644 --- a/src/a2a/utils/telemetry.py +++ b/src/a2a/utils/telemetry.py @@ -62,10 +62,13 @@ def internal_method(self): from typing import TYPE_CHECKING, Any -try: - from typing import Self -except ImportError: # pragma: no cover - for Python < 3.11 +if TYPE_CHECKING: from typing_extensions import Self +else: + try: + from typing import Self + except ImportError: # pragma: no cover - for Python < 3.11 + from typing_extensions import Self if TYPE_CHECKING: diff --git a/uv.lock b/uv.lock index 710f12989..86bd45ff1 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 3 +revision = 2 requires-python = ">=3.10" resolution-markers = [ "python_full_version >= '3.13'", @@ -15,6 +15,7 @@ dependencies = [ { name = "httpx-sse" }, { name = "protobuf" }, { name = "pydantic" }, + { name = "typing-extensions" }, ] [package.optional-dependencies] @@ -125,6 +126,7 @@ requires-dist = [ { name = "sse-starlette", marker = "extra == 'http-server'" }, { name = "starlette", marker = "extra == 'all'" }, { name = "starlette", marker = "extra == 'http-server'" }, + { name = "typing-extensions", specifier = ">=4.0.0" }, ] provides-extras = ["all", "encryption", "grpc", "http-server", "mysql", "postgresql", "signing", "sql", "sqlite", "telemetry"] From 21af7340e367e0693cb505c95763017c3057849c Mon Sep 17 00:00:00 2001 From: Wondr Date: Sun, 25 Jan 2026 17:26:26 +0100 Subject: [PATCH 4/4] test: cover auth header edge cases --- .../tasks/test_push_notification_sender.py | 36 +++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/tests/server/tasks/test_push_notification_sender.py b/tests/server/tasks/test_push_notification_sender.py index 1068efd75..8be1e1014 100644 --- a/tests/server/tasks/test_push_notification_sender.py +++ b/tests/server/tasks/test_push_notification_sender.py @@ -3,6 +3,8 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx +import pytest +from pydantic import ValidationError from a2a.server.tasks.base_push_notification_sender import ( BasePushNotificationSender, @@ -130,6 +132,40 @@ async def test_send_notification_with_auth_header(self) -> None: ) mock_response.raise_for_status.assert_called_once() + def test_authorization_header_no_credentials(self) -> None: + auth = PushNotificationAuthenticationInfo( + schemes=['Bearer'], credentials=None + ) + config = create_sample_push_config(authentication=auth) + assert self.sender._authorization_header(config) is None + + def test_authorization_header_empty_schemes(self) -> None: + auth = PushNotificationAuthenticationInfo( + schemes=[], credentials='token' + ) + config = create_sample_push_config(authentication=auth) + assert self.sender._authorization_header(config) is None + + def test_authorization_header_non_bearer_scheme(self) -> None: + auth = PushNotificationAuthenticationInfo( + schemes=['Basic'], credentials='token' + ) + config = create_sample_push_config(authentication=auth) + assert self.sender._authorization_header(config) == 'Basic token' + + def test_authorization_header_filters_empty_schemes(self) -> None: + auth = PushNotificationAuthenticationInfo( + schemes=['', 'Bearer'], credentials='token' + ) + config = create_sample_push_config(authentication=auth) + assert self.sender._authorization_header(config) == 'Bearer token' + + def test_authorization_header_none_scheme_rejected(self) -> None: + with pytest.raises(ValidationError): + PushNotificationAuthenticationInfo( + schemes=['Bearer', None], credentials='token' + ) + async def test_send_notification_no_config(self) -> None: task_id = 'task_send_no_config' task_data = create_sample_task(task_id=task_id)