refactor: refactor base client

[guglielmo-san]

Description

This PR refactors the base_client API definitions, introducing a new object to store the extensions, avoiding passing it as external parameter.
Interceptors have been removed as their logic will be changed in next PR

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly refactors the client-side API by centralizing the management of service-level parameters, such as extensions, within a dedicated ClientCallContext object. This change simplifies method signatures across the BaseClient and its various transport layers, promoting cleaner code and better encapsulation of request-specific configurations. The update enhances the flexibility and maintainability of the client by providing a unified mechanism for handling call parameters.

Highlights

• API Refactoring: The BaseClient API definitions were refactored to streamline method signatures by removing the extensions parameter from all client-facing methods and their underlying transport implementations.
• Centralized Extension Management: A new ServiceParameters object and ServiceParametersFactory were introduced to centralize the management of service-level parameters, including extensions, within the ClientCallContext.
• ClientCallContext Enhancement: The ClientCallContext was updated to include a service_parameters field, allowing extensions and other call-specific metadata to be passed implicitly through the context rather than explicitly as method arguments.
• Transport Layer Updates: All transport implementations (gRPC, JSON-RPC, REST) were modified to retrieve extension information and other metadata from the ClientCallContext.service_parameters, ensuring consistent handling across different communication protocols.
• Simplified SendMessage API: The send_message method in BaseClient and its transports now accepts a single SendMessageRequest object, which encapsulates both the message and its configuration, including blocking behavior and accepted output modes.

Changelog

• src/a2a/client/base_client.py
  • Removed Any import.
  • Removed ClientCallContext import from a2a.client.client and imported it from a2a.client.middleware.
  • Removed Message and SendMessageConfiguration imports.
  • Updated the send_message method signature to accept SendMessageRequest and removed configuration, request_metadata, and extensions parameters.
  • Removed logic for constructing SendMessageConfiguration and SendMessageRequest internally.
  • Added a new private method _apply_client_config to apply client configuration to a SendMessageRequest.
  • Removed extensions parameter from get_task, cancel_task, create_task_push_notification_config, get_task_push_notification_config, list_task_push_notification_configs, delete_task_push_notification_config, subscribe, and get_extended_agent_card methods.
• src/a2a/client/client.py
  • Removed Message and SendMessageConfiguration imports and added SendMessageRequest import.
  • Removed the extensions field from the ClientConfig class.
  • Updated the send_message abstract method signature to accept SendMessageRequest and removed configuration, request_metadata, and extensions parameters.
  • Removed extensions parameter from get_task, cancel_task, create_task_push_notification_config, get_task_push_notification_config, list_task_push_notification_configs, delete_task_push_notification_config, subscribe, and get_extended_agent_card abstract methods.
  • Removed logic related to all_extensions in the create method.
• src/a2a/client/client_factory.py
  • Removed config.extensions from the _register_defaults method calls for JsonRpcTransport and RestTransport.
  • Removed the extensions parameter from the connect method.
  • Removed the extensions parameter from the create method.
  • Removed logic related to all_extensions in the create method.
• src/a2a/client/middleware.py
  • Imported ServiceParameters from a2a.client.service_parameters.
  • Added service_parameters: ServiceParameters | None = None field to the ClientCallContext class.
• src/a2a/client/service_parameters.py
  • Added new file service_parameters.py which defines ServiceParameters (TypeAlias for dict[str, str]), ServiceParametersUpdate (TypeAlias for Callable[[ServiceParameters], None]), ServiceParametersFactory class with create and create_from static methods, and with_a2a_extensions function.
• src/a2a/client/transports/base.py
  • Removed extensions parameter from the abstract methods send_message, send_message_streaming, get_task, list_tasks, cancel_task, create_task_push_notification_config, get_task_push_notification_config, list_task_push_notification_configs, delete_task_push_notification_config, subscribe, and get_extended_agent_card.
• src/a2a/client/transports/grpc.py
  • Removed A2AClientError and A2AClientTimeoutError imports from a2a.client.errors and imported ClientCallContext from a2a.client.middleware.
  • Removed ClientCallContext import from a2a.client.middleware and imported A2AClientError and A2AClientTimeoutError from a2a.client.errors.
  • Removed HTTP_EXTENSION_HEADER import from a2a.extensions.common.
  • Removed the extensions parameter from the __init__ method.
  • Removed self.extensions attribute.
  • Removed config.extensions from the create class method.
  • Removed extensions parameter from send_message, send_message_streaming, subscribe, get_task, list_tasks, cancel_task, create_task_push_notification_config, get_task_push_notification_config, list_task_push_notification_configs, delete_task_push_notification_config, and get_extended_agent_card methods.
  • Modified _get_grpc_metadata to retrieve metadata from context.service_parameters instead of a direct extensions parameter or self.extensions.
  • Modified _call_grpc and _call_grpc_stream to pass context directly for metadata generation.
• src/a2a/client/transports/jsonrpc.py
  • Removed cast import from typing.
  • Removed update_extension_header import from a2a.extensions.common.
  • Removed the extensions parameter from the __init__ method.
  • Removed self.extensions attribute.
  • Simplified send_message by removing update_extension_header and _apply_interceptors calls, and passing rpc_request.data and context directly to _send_request.
  • Simplified send_message_streaming by removing update_extension_header and _apply_interceptors calls, and passing rpc_request.data and context directly to _send_stream_request.
  • Simplified get_task, list_tasks, cancel_task, create_task_push_notification_config, get_task_push_notification_config, list_task_push_notification_configs, delete_task_push_notification_config, subscribe, and get_extended_agent_card methods by removing extensions parameter, update_extension_header and _apply_interceptors calls, and passing rpc_request.data and context directly to _send_request or _send_stream_request.
  • Removed the _apply_interceptors method.
  • Modified _get_http_args to include context.service_parameters in headers if present.
  • Modified _send_request and _send_stream_request to accept context and use _get_http_args to build http_kwargs.
• src/a2a/client/transports/rest.py
  • Removed update_extension_header import from a2a.extensions.common.
  • Removed the extensions parameter from the __init__ method.
  • Removed self.extensions attribute.
  • Simplified send_message by replacing _prepare_send_message and _send_post_request calls with a unified _execute_request call.
  • Simplified send_message_streaming by removing _prepare_send_message and directly using _get_http_args and _send_stream_request.
  • Simplified get_task, list_tasks, cancel_task, create_task_push_notification_config, get_task_push_notification_config, list_task_push_notification_configs, delete_task_push_notification_config, subscribe, and get_extended_agent_card methods by replacing specific _send_get_request, _send_post_request, _send_delete_request calls with a unified _execute_request call and removing extensions parameter and _apply_interceptors calls.
  • Removed the _apply_interceptors method.
  • Removed the _prepare_send_message method.
  • Modified _get_http_args to include context.service_parameters in headers if present.
  • Added headers update logic to _send_stream_request.
  • Replaced _send_post_request, _send_get_request, and _send_delete_request with a single _execute_request method that handles different HTTP methods and payload types.
• src/a2a/client/transports/tenant_decorator.py
  • Removed extensions parameter from send_message, send_message_streaming, get_task, list_tasks, cancel_task, create_task_push_notification_config, get_task_push_notification_config, list_task_push_notification_configs, delete_task_push_notification_config, subscribe, and get_extended_agent_card methods.
• src/a2a/extensions/common.py
  • Removed Any import from typing.
  • Removed the update_extension_header function.
• tests/client/test_auth_middleware.py
  • Added SendMessageRequest import.
  • Updated send_message calls to pass SendMessageRequest object.
  • Added @pytest.mark.skip decorators to interceptor tests, indicating they are not explicitly being tested.
• tests/client/test_base_client.py
  • Updated send_message calls to use SendMessageRequest object, encapsulating the message and configuration.
• tests/client/test_client_factory.py
  • Removed extensions parameter from ClientConfig initialization in tests.
  • Removed assertions related to client._transport.extensions.
• tests/client/transports/test_grpc_client.py
  • Imported ClientCallContext.
  • Removed extensions parameter from GrpcTransport initialization.
  • Updated send_message and cancel_task calls to pass extensions via ClientCallContext's service_parameters.
  • Adjusted metadata assertions to reflect the removal of default extensions and the use of service_parameters.
• tests/client/transports/test_jsonrpc_client.py
  • Removed test_init_with_extensions.
  • Removed test_interceptor_called.
  • Removed extensions parameter from JsonRpcTransport initialization.
  • Updated test_extensions_added_to_request and test_get_card_with_extended_card_support_with_extensions to pass extensions via ClientCallContext's service_parameters and adjusted assertions.
• tests/client/transports/test_rest_client.py
  • Removed extensions parameter from RestTransport initialization.
  • Updated test_send_message_with_default_extensions, test_send_message_streaming_with_new_extensions, and test_get_card_with_extended_card_support_with_extensions to pass extensions via ClientCallContext's service_parameters and adjusted assertions.
  • Replaced _send_get_request with _execute_request in test_get_card_with_extended_card_support_with_extensions.
• tests/e2e/push_notifications/test_default_push_notification_support.py
  • Added SendMessageConfiguration and SendMessageRequest imports.
  • Updated send_message calls to use SendMessageRequest object, encapsulating the message and configuration.
• tests/extensions/test_common.py
  • Removed update_extension_header import.
  • Removed all tests related to update_extension_header.
• tests/integration/test_client_server_integration.py
  • Added imports for ClientCallContext, ServiceParametersFactory, with_a2a_extensions, and A2ACardResolver.
  • Added a duplicate SendMessageRequest import.
  • Updated test_json_transport_base_client_send_message_with_extensions to construct ClientCallContext with service_parameters for extensions and adjusted assertions.
• tests/integration/test_end_to_end.py
  • Added SendMessageRequest import.
  • Updated send_message calls to use SendMessageRequest object, encapsulating the message and configuration.

Activity

• No specific activity was recorded for this pull request in the provided context.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request refactors how client-side extensions are handled by introducing ServiceParameters within the ClientCallContext. This is a good architectural improvement that centralizes service-level parameters and cleans up method signatures across the client and transport layers.

However, there are a few significant points to consider:

• Breaking API Change: The send_message method signature has been changed to accept a SendMessageRequest object instead of a Message and other parameters. This is a major breaking change for users of the SDK and should be clearly communicated.
• Interceptor Removal: The refactoring appears to have removed support for ClientCallInterceptor from the JsonRpcTransport. This is another significant breaking change that is not mentioned in the PR description. The skipping of interceptor-related tests confirms this. This needs to be addressed, either by restoring the functionality or by explicitly documenting its removal.
• Potential Bug: I've identified a potential logic reversal in base_client.py for handling the blocking parameter in polling mode, which could lead to incorrect behavior.

I've left specific comments on the code with more details. Overall, the direction of the refactoring is positive, but these critical issues should be resolved before merging.

_{Note: Security Review did not run due to the size of the PR.}

I am having trouble creating individual review comments. Click here to see my feedback.

src/a2a/client/transports/jsonrpc.py (436-457)

The _apply_interceptors method has been removed from this transport. This is a significant breaking change as it appears to remove support for client-side call interceptors for the JsonRpcTransport. This was not mentioned in the pull request description. If this removal was intentional, it should be documented. If it was accidental, it should be restored.

src/a2a/client/base_client.py (98-99)

There appears to be a logic reversal in how the blocking configuration is applied. The previous implementation set blocking=not self._config.polling, meaning if polling is enabled, the call is non-blocking. The new logic here appears to set blocking=True when polling is enabled, which is the opposite. This could be a bug.

src/a2a/client/transports/rest.py (366-392)

There is an inconsistency in how headers are handled between _execute_request and _send_stream_request. The _send_stream_request method correctly merges headers from self.httpx_client with request-specific headers. However, _execute_request only uses headers from the context (via _get_http_args) and does not include the base client headers from self.httpx_client.headers. This could lead to missing headers (e.g., A2A-Version) in non-streaming REST requests.

To ensure consistency, _execute_request should also merge the base client headers.

    async def _execute_request(
        self,
        method: str,
        target: str,
        tenant: str,
        payload: dict[str, Any] | None = None,
        context: ClientCallContext | None = None,
    ) -> dict[str, Any]:
        path = self._get_path(target, tenant)
        http_kwargs = self._get_http_args(context)
        payload = payload or {}

        headers = dict(self.httpx_client.headers.items())
        headers.update(http_kwargs.get('headers', {}))
        timeout = http_kwargs.get('timeout', httpx.USE_CLIENT_DEFAULT)

        json_payload = payload if method == 'POST' else None
        params = payload if method != 'POST' else None

        request = self.httpx_client.build_request(
            method,
            f'{self.url}{path}',
            json=json_payload,
            params=params,
            headers=headers,  # type: ignore[arg-type]
            timeout=timeout,  # type: ignore[arg-type]
        )
        return await self._send_request(request)

tests/client/transports/test_jsonrpc_client.py (651-688)

This test contains a large block of commented-out code which seems to be from a previous implementation or debugging session. It would be best to remove this commented-out code to improve test readability.