Initial version of A2A version header validation.

Author: bartek-w

src/a2a/compat/v0_3/jsonrpc_adapter.py

@@ -38,8 +38,9 @@
 from a2a.server.jsonrpc_models import (
     JSONRPCError as CoreJSONRPCError,
 )
+from a2a.utils import constants
 from a2a.utils.errors import ExtendedAgentCardNotConfiguredError
-from a2a.utils.helpers import maybe_await
+from a2a.utils.helpers import maybe_await, validate_version
 
 
 logger = logging.getLogger(__name__)
@@ -152,6 +153,7 @@ async def handle_request(
                 request_id, CoreInternalError(message=str(e))
             )
 
+    @validate_version(constants.PROTOCOL_VERSION_0_3)
     async def _process_non_streaming_request(
         self,
         request_id: 'str | int | None',
@@ -266,6 +268,7 @@ async def get_authenticated_extended_card(
 
         return conversions.to_compat_agent_card(card_to_serve)
 
+    @validate_version(constants.PROTOCOL_VERSION_0_3)
     async def _process_streaming_request(
         self,
         request_id: 'str | int | None',

src/a2a/compat/v0_3/rest_handler.py

@@ -28,7 +28,12 @@
 from a2a.compat.v0_3 import types as types_v03
 from a2a.compat.v0_3.request_handler import RequestHandler03
 from a2a.server.context import ServerCallContext
-from a2a.utils.helpers import validate, validate_async_generator
+from a2a.utils import constants
+from a2a.utils.helpers import (
+    validate,
+    validate_async_generator,
+    validate_version,
+)
 from a2a.utils.telemetry import SpanKind, trace_class
 
 
@@ -53,6 +58,7 @@ def __init__(
         self.agent_card = agent_card
         self.handler03 = RequestHandler03(request_handler=request_handler)
 
+    @validate_version(constants.PROTOCOL_VERSION_0_3)
     async def on_message_send(
         self,
         request: Request,
@@ -78,6 +84,7 @@ async def on_message_send(
         pb2_v03_resp = proto_utils.ToProto.task_or_message(v03_resp)
         return MessageToDict(pb2_v03_resp)
 
+    @validate_version(constants.PROTOCOL_VERSION_0_3)
     @validate_async_generator(
         lambda self: self.agent_card.capabilities.streaming,
         'Streaming is not supported by the agent',

src/a2a/server/apps/jsonrpc/jsonrpc_app.py

@@ -254,7 +254,7 @@ def __init__(  # noqa: PLR0913
                 agent_card=agent_card,
                 http_handler=http_handler,
                 extended_agent_card=extended_agent_card,
-                context_builder=context_builder,
+                context_builder=self._context_builder,
                 card_modifier=card_modifier,
                 extended_card_modifier=extended_card_modifier,
             )

src/a2a/server/request_handlers/jsonrpc_handler.py

@@ -31,7 +31,7 @@
     Task,
     TaskPushNotificationConfig,
 )
-from a2a.utils import proto_utils
+from a2a.utils import constants, proto_utils
 from a2a.utils.errors import (
     JSON_RPC_ERROR_CODE_MAP,
     A2AError,
@@ -49,7 +49,12 @@
     UnsupportedOperationError,
     VersionNotSupportedError,
 )
-from a2a.utils.helpers import maybe_await, validate, validate_async_generator
+from a2a.utils.helpers import (
+    maybe_await,
+    validate,
+    validate_async_generator,
+    validate_version,
+)
 from a2a.utils.telemetry import SpanKind, trace_class
 
 
@@ -142,6 +147,7 @@ def _get_request_id(
             return None
         return context.state.get('request_id')
 
+    @validate_version(constants.PROTOCOL_VERSION_1_0)
     async def on_message_send(
         self,
         request: SendMessageRequest,
@@ -171,6 +177,11 @@ async def on_message_send(
         except A2AError as e:
             return _build_error_response(request_id, e)
 
+    @validate_async_generator(
+        lambda self: self.agent_card.capabilities.streaming,
+        'Streaming is not supported by the agent',
+    )
+    @validate_version(constants.PROTOCOL_VERSION_1_0)
     @validate_async_generator(
         lambda self: self.agent_card.capabilities.streaming,
         'Streaming is not supported by the agent',

src/a2a/server/request_handlers/rest_handler.py

@@ -27,9 +27,13 @@
     GetTaskPushNotificationConfigRequest,
     SubscribeToTaskRequest,
 )
-from a2a.utils import proto_utils
+from a2a.utils import constants, proto_utils
 from a2a.utils.errors import TaskNotFoundError
-from a2a.utils.helpers import validate, validate_async_generator
+from a2a.utils.helpers import (
+    validate,
+    validate_async_generator,
+    validate_version,
+)
 from a2a.utils.telemetry import SpanKind, trace_class
 
 
@@ -61,6 +65,7 @@ def __init__(
         self.agent_card = agent_card
         self.request_handler = request_handler
 
+    @validate_version(constants.PROTOCOL_VERSION_1_0)
     async def on_message_send(
         self,
         request: Request,
@@ -87,6 +92,7 @@ async def on_message_send(
             response = a2a_pb2.SendMessageResponse(message=task_or_message)
         return MessageToDict(response)
 
+    @validate_version(constants.PROTOCOL_VERSION_1_0)
     @validate_async_generator(
         lambda self: self.agent_card.capabilities.streaming,
         'Streaming is not supported by the agent',

src/a2a/utils/error_handlers.py

@@ -1,4 +1,5 @@
 import functools
+import inspect
 import logging
 
 from collections.abc import Awaitable, Callable, Coroutine
@@ -130,11 +131,8 @@ def rest_stream_error_handler(
 ) -> Callable[..., Coroutine[Any, Any, Any]]:
     """Decorator to catch A2AError for a streaming method, log it and then rethrow it to be handled by framework."""
 
-    @functools.wraps(func)
-    async def wrapper(*args: Any, **kwargs: Any) -> Any:
-        try:
-            return await func(*args, **kwargs)
-        except A2AError as error:
+    def _log_error(error: Exception) -> None:
+        if isinstance(error, A2AError):
             log_level = (
                 logging.ERROR
                 if isinstance(error, InternalError)
@@ -147,14 +145,33 @@ async def wrapper(*args: Any, **kwargs: Any) -> Any:
                 getattr(error, 'message', str(error)),
                 f', Data={error.data}' if error.data else '',
             )
-            # Since the stream has started, we can't return a JSONResponse.
-            # Instead, we run the error handling logic (provides logging)
-            # and reraise the error and let server framework manage
-            raise error
+        else:
+            logger.exception('Unknown streaming error occurred')
+
+    @functools.wraps(func)
+    async def wrapper(*args: Any, **kwargs: Any) -> Any:
+        try:
+            response = await func(*args, **kwargs)
+
+            # If the response has an async generator body (like EventSourceResponse),
+            # we must wrap it to catch errors that occur during stream execution.
+            if hasattr(response, 'body_iterator') and inspect.isasyncgen(response.body_iterator):
+                original_iterator = response.body_iterator
+
+                async def error_catching_iterator():
+                    try:
+                        async for item in original_iterator:
+                            yield item
+                    except Exception as stream_error:
+                        _log_error(stream_error)
+                        raise stream_error
+
+                response.body_iterator = error_catching_iterator()
+
+            return response
+
         except Exception as e:
-            # Since the stream has started, we can't return a JSONResponse.
-            # Instead, we run the error handling logic (provides logging)
-            # and reraise the error and let server framework manage
+            _log_error(e)
             raise e
 
     return wrapper

src/a2a/utils/helpers.py

@@ -5,12 +5,14 @@
 import json
 import logging
 
-from collections.abc import Awaitable, Callable
-from typing import Any, TypeVar
+from collections.abc import AsyncIterator, Awaitable, Callable
+from typing import Any, TypeVar, cast
 from uuid import uuid4
 
 from google.protobuf.json_format import MessageToDict
+from packaging.version import InvalidVersion, Version
 
+from a2a.server.context import ServerCallContext
 from a2a.types.a2a_pb2 import (
     AgentCard,
     Artifact,
@@ -21,11 +23,13 @@
     TaskState,
     TaskStatus,
 )
-from a2a.utils.errors import UnsupportedOperationError
+from a2a.utils import constants
+from a2a.utils.errors import UnsupportedOperationError, VersionNotSupportedError
 from a2a.utils.telemetry import trace_function
 
 
 T = TypeVar('T')
+F = TypeVar('F', bound=Callable[..., Any])
 
 
 logger = logging.getLogger(__name__)
@@ -297,7 +301,6 @@ def validate_async_generator(
         This decorator is specifically for async generator methods (async def with yield).
         The validation happens before the generator starts yielding values.
     """
-
     def decorator(function):
         @functools.wraps(function)
         async def wrapper(self, *args, **kwargs):
@@ -378,3 +381,117 @@ async def maybe_await(value: T | Awaitable[T]) -> T:
     if inspect.isawaitable(value):
         return await value
     return value
+
+
+def validate_version(expected_version: str) -> Callable[[F], F]:
+    """Decorator that validates the A2A-Version header in the request context.
+
+    The header name is defined by `constants.VERSION_HEADER` ('A2A-Version').
+    If the header is missing or empty, it is interpreted as `constants.PROTOCOL_VERSION_0_3` ('0.3').
+    If the version in the header does not match the `expected_version` (major and minor parts),
+    a `VersionNotSupportedError` is raised. Patch version is ignored.
+
+    This decorator supports both async methods and async generator methods. It
+    expects a `ServerCallContext` to be present either in the arguments or
+    keyword arguments of the decorated method.
+
+    Args:
+        expected_version: The A2A protocol version string expected by the method.
+
+    Returns:
+        The decorated function.
+
+    Raises:
+        VersionNotSupportedError: If the version in the request does not match `expected_version`.
+    """
+    try:
+        expected_v = Version(expected_version)
+    except InvalidVersion:
+        # If the expected version is not a valid semver, we can't do major/minor comparison.
+        # This shouldn't happen with our constants.
+        expected_v = None
+
+    def decorator(func: F) -> F:
+        def _get_actual_version(
+            args: tuple[Any, ...], kwargs: dict[str, Any]
+        ) -> str:
+            context = kwargs.get('context')
+            if context is None:
+                for arg in args:
+                    if isinstance(arg, ServerCallContext):
+                        context = arg
+                        break
+
+            if context is None:
+                # If no context is found, we can't validate the version.
+                # In a real scenario, this shouldn't happen for properly routed requests.
+                # We default to the expected version to allow test call to proceed.
+                return expected_version
+
+            headers = context.state.get('headers', {})
+            # Header names are usually case-insensitive in most frameworks, but dict lookup is case-sensitive.
+            # We check both standard and lowercase versions.
+            actual_version = headers.get(
+                constants.VERSION_HEADER
+            ) or headers.get(constants.VERSION_HEADER.lower())
+
+            if not actual_version:
+                return constants.PROTOCOL_VERSION_0_3
+
+            return str(actual_version)
+
+        def _is_version_compatible(actual: str) -> bool:
+            if actual == expected_version:
+                return True
+            if not expected_v:
+                return False
+            try:
+                actual_v = Version(actual)
+            except InvalidVersion:
+                return False
+            else:
+                return (
+                    actual_v.major == expected_v.major
+                    and actual_v.minor == expected_v.minor
+                )
+
+        if inspect.isasyncgenfunction(inspect.unwrap(func)):
+
+            @functools.wraps(func)
+            async def async_gen_wrapper(
+                self: Any, *args: Any, **kwargs: Any
+            ) -> AsyncIterator[Any]:
+                actual_version = _get_actual_version(args, kwargs)
+                if not _is_version_compatible(actual_version):
+                    logger.warning(
+                        "Version mismatch: actual='%s', expected='%s'",
+                        actual_version,
+                        expected_version,
+                    )
+                    raise VersionNotSupportedError(
+                        message=f"A2A version '{actual_version}' is not supported by this handler. "
+                        f"Expected version '{expected_version}'."
+                    )
+                async for item in func(self, *args, **kwargs):
+                    yield item
+
+            return cast('F', async_gen_wrapper)
+
+        @functools.wraps(func)
+        async def async_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
+            actual_version = _get_actual_version(args, kwargs)
+            if not _is_version_compatible(actual_version):
+                logger.warning(
+                    "Version mismatch: actual='%s', expected='%s'",
+                    actual_version,
+                    expected_version,
+                )
+                raise VersionNotSupportedError(
+                    message=f"A2A version '{actual_version}' is not supported by this handler. "
+                    f"Expected version '{expected_version}'."
+                )
+            return await func(self, *args, **kwargs)
+
+        return cast('F', async_wrapper)
+
+    return decorator

tests/compat/v0_3/test_rest_handler.py

@@ -37,7 +37,9 @@ def rest_handler(agent_card, mock_core_handler):
 
 @pytest.fixture
 def mock_context():
-    return MagicMock(spec=ServerCallContext)
+    m = MagicMock(spec=ServerCallContext)
+    m.state = {'headers': {'A2A-Version': '0.3'}}
+    return m
 
 
 @pytest.fixture