feat: add traceability extension support #387
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @mindpower, I'm Gemini Code Assist1! 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 foundational extension mechanism for the A2A client and server. This new framework enables developers to inject custom logic at critical points in the message processing flow, significantly enhancing the system's modularity and extensibility. A practical example, the TraceExtension, is included to demonstrate how this mechanism can be leveraged for features like end-to-end traceability.
Highlights
- Core Extension Framework: A new
Extensionbase class has been introduced, providing a standardized interface for custom logic to be installed and executed within the A2A client and server components. This promotes modularity and extensibility. - Client-Side Extension Integration: The
BaseClientnow supports installing extensions via a newinstall_extensionmethod. Installed extensions can intercept and process messages sent from the client through a newon_client_messagehook, called before the message is sent (lines 71-72 insrc/a2a/client/base_client.py). - Server-Side Extension Integration: Similarly, the
DefaultRequestHandleron the server side now allows extensions to be installed via itsinstall_extensionmethod. These extensions can process incoming messages through anon_server_messagehook, called when a message is received for execution (lines 192-193 insrc/a2a/server/request_handlers/default_request_handler.py). - Traceability Extension Example: A concrete
TraceExtensionhas been added, demonstrating how the new framework can be used to implement tracing capabilities. This includes new Pydantic models (ResponseTrace,Step,AgentInvocation,ToolInvocation) for capturing detailed trace information such as call types, costs, tokens, and latency. - Dependency Updates: The
uv.lockfile has been updated to relax the version constraints forfastapi(from>=0.116.1to>=0.95.0) andprotobuf(from==5.29.5to>=5.29.5), potentially broadening compatibility with other packages.
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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| 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 issue 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 or fill out our survey to provide feedback.
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
-
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. ↩
There was a problem hiding this comment.
Code Review
This pull request introduces a new extension system to support traceability. The changes include a base Extension class, a TraceExtension implementation, and integration points in the client and server handlers. The models for trace data and corresponding tests are also added.
src/a2a/extensions/trace.py
Outdated
| if hasattr(message, 'metadata') and 'trace' in message.metadata: | ||
| print(f"Received trace: {message.metadata['trace']}") |
There was a problem hiding this comment.
There's a potential TypeError in on_server_message. The expression 'trace' in message.metadata will raise an exception if message.metadata is None. You should ensure message.metadata is not None before checking for a key within it.
| if hasattr(message, 'metadata') and 'trace' in message.metadata: | |
| print(f"Received trace: {message.metadata['trace']}") | |
| if hasattr(message, 'metadata') and message.metadata and 'trace' in message.metadata: | |
| print(f"Received trace: {message.metadata['trace']}") |
| def on_client_message(self, message: Any) -> None: | ||
| """Called when a message is sent from the client.""" |
There was a problem hiding this comment.
For better type safety, the message parameter in on_client_message should be typed as Message instead of Any. Based on the usage in base_client.py and default_request_handler.py, this parameter is always an instance of a2a.types.Message.
To make this work, you'll also need to add from a2a.types import Message inside the if TYPE_CHECKING: block at the top of the file to avoid circular imports.
| def on_client_message(self, message: Any) -> None: | |
| """Called when a message is sent from the client.""" | |
| def on_client_message(self, message: Message) -> None: |
| def on_server_message(self, message: Any) -> None: | ||
| """Called when a message is received by the server.""" |
There was a problem hiding this comment.
For better type safety, the message parameter in on_server_message should be typed as Message instead of Any. Based on the usage in base_client.py and default_request_handler.py, this parameter is always an instance of a2a.types.Message.
To make this work, you'll also need to add from a2a.types import Message inside the if TYPE_CHECKING: block at the top of the file to avoid circular imports.
| def on_server_message(self, message: Any) -> None: | |
| """Called when a message is received by the server.""" | |
| def on_server_message(self, message: Message) -> None: |
| self._running_agents_lock = asyncio.Lock() | ||
| self._extensions: list[Extension] = [] | ||
|
|
||
| def install_extension(self, extension: Extension, server: Any) -> None: |
There was a problem hiding this comment.
The server parameter is typed as Any. For better type safety, it should be typed more specifically. Based on the Extension.install method signature, it should be A2AServer. To avoid circular imports, you can use a string forward reference: server: "A2AServer". You would then need to add from a2a.server.server import A2AServer inside a TYPE_CHECKING block at the top of the file.
| def install_extension(self, extension: Extension, server: Any) -> None: | |
| def install_extension(self, extension: Extension, server: "A2AServer") -> None: |
| steps=[ | ||
| Step( | ||
| step_id='step-1-agent', | ||
| trace_id='trace-example-1234p', |
There was a problem hiding this comment.
The trace_id for this Step is 'trace-example-1234p', which is different from the ResponseTrace's trace_id ('trace-example-12345'). While this doesn't break the current serialization test, it represents logically inconsistent test data which might be confusing for future readers. A step's trace_id should typically match the ID of the trace it belongs to.
| trace_id='trace-example-1234p', | |
| trace_id='trace-example-12345', |
| for extension in server_handler._extensions: | ||
| extension.on_server_message(message) | ||
|
|
||
| # Check that the server-side handler was called | ||
| # (in this case, it just prints a message) | ||
| # We can't easily check the output of print, so we'll just | ||
| # assume it worked if no exceptions were raised. |
There was a problem hiding this comment.
The test for the server-side message handling is currently implicit, assuming it works if no exception is raised. You can make this test more explicit and robust by mocking the print function and asserting that it's called with the expected arguments. You will need to add import unittest.mock or similar to the top of the file.
| for extension in server_handler._extensions: | |
| extension.on_server_message(message) | |
| # Check that the server-side handler was called | |
| # (in this case, it just prints a message) | |
| # We can't easily check the output of print, so we'll just | |
| # assume it worked if no exceptions were raised. | |
| with unittest.mock.patch('builtins.print') as mock_print: | |
| for extension in server_handler._extensions: | |
| extension.on_server_message(message) | |
| # Check that the server-side handler was called | |
| mock_print.assert_called_once_with("Received trace: client-trace") |
holtskinner
left a comment
holtskinner
left a comment
There was a problem hiding this comment.
@mikeas1 Can you review since this is an extension?
Also, general question. Should extensions be going in the core SDK like this? Or should they exist as samples/independent libraries.
Or if we have them in the core SDK, should they be optional dependencies or anything like that?
|
Please hold on review on this. I am implementing the traceability extension in the a2a-samples repo currently but keep this in case some additional change is needed for the extension framework. |
Ok, thanks for the update. I'm going to mark this as a draft for now. |
|
@mindpower Is this PR/branch still needed? |
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:
CONTRIBUTINGGuide.fix:which represents bug fixes, and correlates to a SemVer patch.feat:represents a new feature, and correlates to a SemVer minor.feat!:, orfix!:,refactor!:, etc., which represent a breaking change (indicated by the!) and will result in a SemVer major.bash scripts/format.shfrom the repository root to format)Fixes #<issue_number_goes_here> 🦕