Merge branch 'main' into remove-camel-case

Author: holtskinner

.ruff.toml

@@ -136,7 +136,7 @@ inline-quotes = "single"
     "PLR2004",
     "SLF001",
 ]
-"types.py" = ["D", "E501", "N815"]  # Ignore docstring and annotation issues in types.py
+"types.py" = ["D", "E501"]  # Ignore docstring and annotation issues in types.py
 "proto_utils.py" = ["D102", "PLR0911"]
 "helpers.py" = ["ANN001", "ANN201", "ANN202"]
 

scripts/format.sh

@@ -2,37 +2,84 @@
 set -e
 set -o pipefail
 
+# --- Argument Parsing ---
+# Initialize flags
+FORMAT_ALL=false
+RUFF_UNSAFE_FIXES_FLAG=""
+
+# Process command-line arguments
+# We use a while loop with shift to process each argument
+while [[ "$#" -gt 0 ]]; do
+    case "$1" in
+        --all)
+            FORMAT_ALL=true
+            echo "Detected --all flag: Formatting all Python files."
+            shift # Consume the argument
+            ;;
+        --unsafe-fixes)
+            RUFF_UNSAFE_FIXES_FLAG="--unsafe-fixes"
+            echo "Detected --unsafe-fixes flag: Ruff will run with unsafe fixes."
+            shift # Consume the argument
+            ;;
+        *)
+            # Handle unknown arguments or just ignore them if we only care about specific ones
+            echo "Warning: Unknown argument '$1'. Ignoring."
+            shift # Consume the argument
+            ;;
+    esac
+done
+
 # Sort Spelling Allowlist
-# The user did not provide this file, so we check for its existence.
 SPELLING_ALLOW_FILE=".github/actions/spelling/allow.txt"
 if [ -f "$SPELLING_ALLOW_FILE" ]; then
+    echo "Sorting and de-duplicating $SPELLING_ALLOW_FILE"
     sort -u "$SPELLING_ALLOW_FILE" -o "$SPELLING_ALLOW_FILE"
 fi
 
-TARGET_BRANCH="origin/${GITHUB_BASE_REF:-main}"
-git fetch origin "${GITHUB_BASE_REF:-main}" --depth=1
+CHANGED_FILES=""
+
+if $FORMAT_ALL; then
+    echo "Formatting all Python files in the repository."
+    # Find all Python files, excluding grpc generated files as per original logic.
+    # `sort -u` ensures unique files and consistent ordering for display/xargs.
+    CHANGED_FILES=$(find . -name '*.py' -not -path './src/a2a/grpc/*' | sort -u)
 
-# Find merge base between HEAD and target branch
-MERGE_BASE=$(git merge-base HEAD "$TARGET_BRANCH")
+    if [ -z "$CHANGED_FILES" ]; then
+        echo "No Python files found to format."
+        exit 0
+    fi
+else
+    echo "No '--all' flag found. Formatting changed Python files based on git diff."
+    TARGET_BRANCH="origin/${GITHUB_BASE_REF:-main}"
+    git fetch origin "${GITHUB_BASE_REF:-main}" --depth=1
 
-# Get python files changed in this PR, excluding grpc generated files
-CHANGED_FILES=$(git diff --name-only --diff-filter=ACMRTUXB "$MERGE_BASE" HEAD -- '*.py' ':!src/a2a/grpc/*')
+    MERGE_BASE=$(git merge-base HEAD "$TARGET_BRANCH")
 
-if [ -z "$CHANGED_FILES" ]; then
-    echo "No changed Python files to format."
-    exit 0
+    # Get python files changed in this PR, excluding grpc generated files
+    CHANGED_FILES=$(git diff --name-only --diff-filter=ACMRTUXB "$MERGE_BASE" HEAD -- '*.py' ':!src/a2a/grpc/*')
+
+    if [ -z "$CHANGED_FILES" ]; then
+        echo "No changed Python files to format."
+        exit 0
+    fi
 fi
 
-echo "Formatting changed files:"
+echo "Files to be formatted:"
 echo "$CHANGED_FILES"
 
-# Formatters are already installed in the activated venv from the GHA step.
-# Use xargs to pass the file list to the formatters.
+# Helper function to run formatters with the list of files.
+# The list of files is passed to xargs via stdin.
 run_formatter() {
     echo "$CHANGED_FILES" | xargs -r "$@"
 }
 
+echo "Running pyupgrade..."
 run_formatter pyupgrade --exit-zero-even-if-changed --py310-plus
+echo "Running autoflake..."
 run_formatter autoflake -i -r --remove-all-unused-imports
-run_formatter ruff check --fix-only
+echo "Running ruff check (fix-only)..."
+run_formatter ruff check --fix-only $RUFF_UNSAFE_FIXES_FLAG
+echo "Running ruff format..."
 run_formatter ruff format
+
+echo "Formatting complete."

src/a2a/extensions/common.py

@@ -0,0 +1,27 @@
+from a2a.types import AgentCard, AgentExtension
+
+
+HTTP_EXTENSION_HEADER = 'X-A2A-Extensions'
+
+
+def get_requested_extensions(values: list[str]) -> set[str]:
+    """Get the set of requested extensions from an input list.
+
+    This handles the list containing potentially comma-separated values, as
+    occurs when using a list in an HTTP header.
+    """
+    return {
+        stripped
+        for v in values
+        for ext in v.split(',')
+        if (stripped := ext.strip())
+    }
+
+
+def find_extension_by_uri(card: AgentCard, uri: str) -> AgentExtension | None:
+    """Find an AgentExtension in an AgentCard given a uri."""
+    for ext in card.capabilities.extensions or []:
+        if ext.uri == uri:
+            return ext
+
+    return None

src/a2a/grpc/a2a_pb2.py

@@ -143,6 +143,24 @@ def metadata(self) -> dict[str, Any]:
             return {}
         return self._params.metadata or {}
 
+    def add_activated_extension(self, uri: str) -> None:
+        """Add an extension to the set of activated extensions for this request.
+
+        This causes the extension to be indicated back to the client in the
+        response.
+        """
+        if self._call_context:
+            self._call_context.activated_extensions.add(uri)
+
+    @property
+    def requested_extensions(self) -> set[str]:
+        """Extensions that the client requested to activate."""
+        return (
+            self._call_context.requested_extensions
+            if self._call_context
+            else set()
+        )
+
     def _check_or_generate_task_id(self) -> None:
         """Ensures a task ID is present, generating one if necessary."""
         if not self._params:

src/a2a/server/agent_execution/context.py

@@ -19,6 +19,10 @@
 
 from a2a.auth.user import UnauthenticatedUser
 from a2a.auth.user import User as A2AUser
+from a2a.extensions.common import (
+    HTTP_EXTENSION_HEADER,
+    get_requested_extensions,
+)
 from a2a.server.context import ServerCallContext
 from a2a.server.request_handlers.jsonrpc_handler import JSONRPCHandler
 from a2a.server.request_handlers.request_handler import RequestHandler
@@ -99,7 +103,13 @@ def build(self, request: Request) -> ServerCallContext:
             user = StarletteUserProxy(request.user)
             state['auth'] = request.auth
         state['headers'] = dict(request.headers)
-        return ServerCallContext(user=user, state=state)
+        return ServerCallContext(
+            user=user,
+            state=state,
+            requested_extensions=get_requested_extensions(
+                request.headers.getlist(HTTP_EXTENSION_HEADER)
+            ),
+        )
 
 
 class JSONRPCApplication(ABC):
@@ -281,7 +291,7 @@ async def _process_streaming_request(
                 request_obj, context
             )
 
-        return self._create_response(handler_result)
+        return self._create_response(context, handler_result)
 
     async def _process_non_streaming_request(
         self,
@@ -353,10 +363,11 @@ async def _process_non_streaming_request(
                     id=request_id, error=error
                 )
 
-        return self._create_response(handler_result)
+        return self._create_response(context, handler_result)
 
     def _create_response(
         self,
+        context: ServerCallContext,
         handler_result: (
             AsyncGenerator[SendStreamingMessageResponse]
             | JSONRPCErrorResponse
@@ -372,12 +383,16 @@ def _create_response(
         payloads.
 
         Args:
+            context: The ServerCallContext provided to the request handler.
             handler_result: The result from a request handler method. Can be an
                 async generator for streaming or a Pydantic model for non-streaming.
 
         Returns:
             A Starlette JSONResponse or EventSourceResponse.
         """
+        headers = {}
+        if exts := context.activated_extensions:
+            headers[HTTP_EXTENSION_HEADER] = ', '.join(sorted(exts))
         if isinstance(handler_result, AsyncGenerator):
             # Result is a stream of SendStreamingMessageResponse objects
             async def event_generator(
@@ -386,17 +401,21 @@ async def event_generator(
                 async for item in stream:
                     yield {'data': item.root.model_dump_json(exclude_none=True)}
 
-            return EventSourceResponse(event_generator(handler_result))
+            return EventSourceResponse(
+                event_generator(handler_result), headers=headers
+            )
         if isinstance(handler_result, JSONRPCErrorResponse):
             return JSONResponse(
                 handler_result.model_dump(
                     mode='json',
                     exclude_none=True,
-                )
+                ),
+                headers=headers,
             )
 
         return JSONResponse(
-            handler_result.root.model_dump(mode='json', exclude_none=True)
+            handler_result.root.model_dump(mode='json', exclude_none=True),
+            headers=headers,
         )
 
     async def _handle_get_agent_card(self, request: Request) -> JSONResponse:

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

@@ -21,3 +21,5 @@ class ServerCallContext(BaseModel):
 
     state: State = Field(default={})
     user: User = Field(default=UnauthenticatedUser())
+    requested_extensions: set[str] = Field(default_factory=set)
+    activated_extensions: set[str] = Field(default_factory=set)

src/a2a/server/context.py

@@ -3,12 +3,14 @@
 import logging
 
 from abc import ABC, abstractmethod
-from collections.abc import AsyncIterable
+from collections.abc import AsyncIterable, Sequence
 
 
 try:
     import grpc
     import grpc.aio
+
+    from grpc.aio import Metadata
 except ImportError as e:
     raise ImportError(
         'GrpcHandler requires grpcio and grpcio-tools to be installed. '
@@ -20,6 +22,10 @@
 
 from a2a import types
 from a2a.auth.user import UnauthenticatedUser
+from a2a.extensions.common import (
+    HTTP_EXTENSION_HEADER,
+    get_requested_extensions,
+)
 from a2a.grpc import a2a_pb2
 from a2a.server.context import ServerCallContext
 from a2a.server.request_handlers.request_handler import RequestHandler
@@ -42,6 +48,19 @@ def build(self, context: grpc.aio.ServicerContext) -> ServerCallContext:
         """Builds a ServerCallContext from a gRPC Request."""
 
 
+def _get_metadata_value(
+    context: grpc.aio.ServicerContext, key: str
+) -> list[str]:
+    md = context.invocation_metadata
+    raw_values: list[str | bytes] = []
+    if isinstance(md, Metadata):
+        raw_values = md.get_all(key)
+    elif isinstance(md, Sequence):
+        lower_key = key.lower()
+        raw_values = [e for (k, e) in md if k.lower() == lower_key]
+    return [e if isinstance(e, str) else e.decode('utf-8') for e in raw_values]
+
+
 class DefaultCallContextBuilder(CallContextBuilder):
     """A default implementation of CallContextBuilder."""
 
@@ -51,7 +70,13 @@ def build(self, context: grpc.aio.ServicerContext) -> ServerCallContext:
         state = {}
         with contextlib.suppress(Exception):
             state['grpc_context'] = context
-        return ServerCallContext(user=user, state=state)
+        return ServerCallContext(
+            user=user,
+            state=state,
+            requested_extensions=get_requested_extensions(
+                _get_metadata_value(context, HTTP_EXTENSION_HEADER)
+            ),
+        )
 
 
 class GrpcHandler(a2a_grpc.A2AServiceServicer):
@@ -102,6 +127,7 @@ async def SendMessage(
             task_or_message = await self.request_handler.on_message_send(
                 a2a_request, server_context
             )
+            self._set_extension_metadata(context, server_context)
             return proto_utils.ToProto.task_or_message(task_or_message)
         except ServerError as e:
             await self.abort_context(e, context)
@@ -140,6 +166,7 @@ async def SendStreamingMessage(
                 a2a_request, server_context
             ):
                 yield proto_utils.ToProto.stream_response(event)
+            self._set_extension_metadata(context, server_context)
         except ServerError as e:
             await self.abort_context(e, context)
         return
@@ -371,3 +398,16 @@ async def abort_context(
                     grpc.StatusCode.UNKNOWN,
                     f'Unknown error type: {error.error}',
                 )
+
+    def _set_extension_metadata(
+        self,
+        context: grpc.aio.ServicerContext,
+        server_context: ServerCallContext,
+    ) -> None:
+        if server_context.activated_extensions:
+            context.set_trailing_metadata(
+                [
+                    (HTTP_EXTENSION_HEADER, e)
+                    for e in sorted(server_context.activated_extensions)
+                ]
+            )

src/a2a/server/request_handlers/grpc_handler.py

@@ -18,7 +18,7 @@ def new_task(request: Message) -> Task:
 
     Raises:
         TypeError: If the message role is None.
-        ValueError: If the message parts are empty or if any part has empty content.
+        ValueError: If the message parts are empty, if any part has empty content, or if the provided context_id is invalid.
     """
     if not request.role:
         raise TypeError('Message role cannot be None')
@@ -28,12 +28,22 @@ def new_task(request: Message) -> Task:
         if isinstance(part.root, TextPart) and not part.root.text:
             raise ValueError('TextPart content cannot be empty')
 
+    context_id_str = request.context_id
+    if context_id_str is not None:
+        try:
+            uuid.UUID(context_id_str)
+            context_id = context_id_str
+        except (ValueError, AttributeError, TypeError) as e:
+            raise ValueError(
+                f"Invalid context_id: '{context_id_str}' is not a valid UUID."
+            ) from e
+    else:
+        context_id = str(uuid.uuid4())
+
     return Task(
         status=TaskStatus(state=TaskState.submitted),
         id=(request.task_id if request.task_id else str(uuid.uuid4())),
-        context_id=(
-            request.context_id if request.context_id else str(uuid.uuid4())
-        ),
+        context_id=context_id,
         history=[request],
     )