feat: handle tenant in Client

[sokoliva]

Description

Thank you for opening a Pull Request!
Before submitting your PR, there are a few things you can do to make sure it goes smoothly:

• Follow the CONTRIBUTING Guide.
• Make your Pull Request title in the https://www.conventionalcommits.org/ specification.
  • Important Prefixes for release-please:
    • fix: which represents bug fixes, and correlates to a SemVer patch.
    • feat: represents a new feature, and correlates to a SemVer minor.
    • feat!:, or fix!:, refactor!:, etc., which represent a breaking change (indicated by the !) and will result in a SemVer major.
• Ensure the tests and linter pass (Run bash scripts/format.sh from the repository root to format)
• Appropriate docs were updated (if necessary)

Fixes #<issue_number_goes_here> 🦕

[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 introduces a robust mechanism for handling multi-tenancy within the client by implementing a transport decorator pattern. This change ensures that tenant information, when present in an agent's supported interface, is automatically and consistently applied to all outgoing requests, simplifying client-side tenant management and enabling seamless interaction with tenant-aware services.

Highlights

• Introduced Tenant Transport Decorator: A new TenantTransportDecorator class was added, implementing the decorator pattern to automatically attach tenant information to outgoing requests for various transport types (REST, JSONRPC, gRPC).
• Client Factory Integration: The ClientFactory was updated to detect tenant information from the AgentInterface and wrap the created transport with the TenantTransportDecorator if a tenant is specified.
• REST Transport Tenant Handling: The RestTransport was modified to prepend the tenant ID to the URL path for all relevant API endpoints, ensuring tenant-specific routing for REST calls. It also includes logic to extract the tenant from the request payload.
• Comprehensive Testing: New unit and integration tests were added to validate the functionality of the TenantTransportDecorator, its application by the ClientFactory, and correct tenant-prepending behavior in the RestTransport across different request types.

Changelog

• src/a2a/client/client_factory.py
  • Imported TenantTransportDecorator for use in client creation.
  • Modified the create method to use selected_interface instead of transport_url for better clarity and access to tenant information.
  • Implemented logic to wrap the instantiated transport with TenantTransportDecorator if a tenant is specified in the selected_interface.
• src/a2a/client/transports/init.py
  • Imported TenantTransportDecorator from base.py.
  • Exported TenantTransportDecorator in the __all__ list to make it publicly accessible.
• src/a2a/client/transports/base.py
  • Imported Message from google.protobuf.message.
  • Added the TenantTransportDecorator class, which extends ClientTransport.
  • Implemented __init__ to store the base transport and tenant ID.
  • Added an update_tenant method to set the tenant on a request message if not already present.
  • Overrode all ClientTransport methods (send_message, send_message_streaming, get_task, etc.) to apply the update_tenant logic before delegating to the base transport.
  • Overrode the close method to delegate to the base transport's close method.
• src/a2a/client/transports/rest.py
  • Added a private helper method _get_path to prepend the tenant ID to a base path.
  • Added a private helper method _pop_tenant to extract and remove the tenant from a dictionary.
  • Modified send_message to extract the tenant and construct a tenant-prefixed URL path.
  • Modified send_message_streaming to extract the tenant and construct a tenant-prefixed URL path.
  • Modified get_task to extract the tenant and construct a tenant-prefixed URL path.
  • Modified list_tasks to extract the tenant from the payload and construct a tenant-prefixed URL path.
  • Modified cancel_task to extract the tenant and construct a tenant-prefixed URL path.
  • Modified set_task_callback to extract the tenant and construct a tenant-prefixed URL path.
  • Modified get_task_callback to extract the tenant and construct a tenant-prefixed URL path.
  • Modified list_task_callback to extract the tenant and construct a tenant-prefixed URL path.
  • Modified delete_task_callback to extract the tenant and construct a tenant-prefixed URL path.
  • Modified subscribe to use the tenant from the request and construct a tenant-prefixed URL path.
  • Removed the unused _model_to_query_params helper function.
• tests/client/test_base_client.py
  • Imported AsyncGenerator, SendMessageRequest, and TenantTransportDecorator.
  • Added a new asynchronous test test_tenant_transport_decorator to verify that the decorator correctly sets the tenant on a SendMessageRequest.
• tests/client/test_client_factory.py
  • Imported AsyncGenerator and TenantTransportDecorator.
  • Added a new test test_client_factory_applies_tenant_decorator to confirm that the ClientFactory correctly wraps the transport with TenantTransportDecorator when a tenant is present in the agent card.
• tests/client/transports/test_rest_client.py
  • Imported various request types (CancelTaskRequest, CreateTaskPushNotificationConfigRequest, etc.) for comprehensive testing.
  • Added TestRestTransportTenant class with parameterized tests to verify that all REST methods correctly prepend the tenant to the URL path when a tenant is provided in the request.
  • Added parameterized tests for streaming REST methods to ensure tenant-prepending in aconnect_sse calls.
• tests/integration/test_tenant.py
  • Added a new integration test file to verify tenant handling across different transport types (REST, JSONRPC, gRPC).
  • Included fixtures for agent_card with tenant information.
  • Added tests for TenantTransportDecorator with RestTransport, JsonRpcTransport, and GrpcTransport to ensure tenant propagation and URL modification.
  • Added a test to confirm that an explicitly provided tenant in a request overrides the tenant from the agent card.

Activity

• The pull request was opened by sokoliva, introducing a new feature to handle tenant information within the client.
• The changes adhere to the conventional commits specification, indicating a new feature (feat:).
• The pull request description includes a checklist for contributing guidelines, testing, and documentation, suggesting a structured approach to development.

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 introduces multi-tenancy support to the client. The implementation is clean, using a TenantTransportDecorator to inject the tenant information into requests. This is a good use of the decorator pattern that keeps the transport logic separate. The changes are well-tested, with new unit and integration tests covering various transports (REST, JSON-RPC, gRPC) and edge cases like explicit tenant overrides. Overall, this is a solid feature addition.

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