feat: Implement a vertex based task store by gaborfeher · Pull Request #752 · a2aproject/a2a-python

gaborfeher · 2026-03-02T14:57:30Z

Create a task store adapter that uses the Vertex Managed Task Store as a backend.

Fixes #751

gemini-code-assist · 2026-03-02T14:57:55Z

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 significant new feature by integrating Google Cloud Vertex AI as a backend for task storage. This change enhances the system's persistence capabilities, moving beyond the existing in-memory solution to a robust, cloud-based task management system. The core implementation involves a new VertexTaskStore and a dedicated conversion layer to handle data mapping between the SDK's task models and Vertex AI's proto types. This also includes an architectural improvement to make the task store pluggable, allowing for greater flexibility in choosing persistence mechanisms.

Highlights

New Vertex AI Task Store Implementation: Introduced a VertexTaskStore that leverages Google Cloud Vertex AI's Agent Engine to persist and manage tasks, providing an alternative to the in-memory store.
Task Conversion Utilities: Developed vertex_task_converter.py to facilitate seamless conversion between the SDK's internal Task and related types (e.g., Part, Artifact, TaskState) and Vertex AI's corresponding proto types.
Pluggable Task Store Architecture: Refactored the sut_agent.py to accept a TaskStore instance via dependency injection, making the task persistence layer configurable and allowing different backends (like Vertex AI) to be used.
Comprehensive Testing and Setup: Added extensive unit tests for the task converter and integration tests for the VertexTaskStore, covering task creation, retrieval, and updates (status, metadata, artifacts). A new shell script (run_vertex_tests.sh) was also added to simplify running these tests.
Dependency Management: Updated pyproject.toml to include the google-cloud-aiplatform dependency, specified as a direct Git reference, and configured hatch to allow direct references for metadata.

Changelog

pyproject.toml
- Added 'vertex' extra dependency for google-cloud-aiplatform with a direct Git reference.
- Configured hatch.metadata to allow direct references for package building.
scripts/run_vertex_tests.sh
- Added a new script to run pytest for Vertex AI task store and converter tests.
- Included environment variable checks for VERTEX_PROJECT, VERTEX_LOCATION, VERTEX_BASE_URL, and VERTEX_API_VERSION.
src/a2a/server/tasks/vertex_task_converter.py
- Added a new module for converting between SDK task types and Vertex AI proto types.
- Implemented functions for converting TaskState, Part (Text, Data, FileWithBytes, FileWithUri), Artifact, and Task objects in both directions.
src/a2a/server/tasks/vertex_task_store.py
- Added a new VertexTaskStore class, implementing the TaskStore interface.
- Implemented save, get, _create, and _update methods to interact with Vertex AI's Agent Engine for task persistence.
- Included logic for handling state, metadata, and artifact changes during task updates.
tck/sut_agent.py
- Imported TaskStore for type hinting.
- Refactored the main function into a serve function that accepts a TaskStore instance.
- Updated the DefaultRequestHandler initialization to use the injected task_store.
tck/sut_agent_with_vertex_task_store.py
- Added a new script to demonstrate running the sut_agent with the VertexTaskStore.
- Initialized a vertexai.Client and VertexTaskStore using environment variables and a newly created agent engine.
tests/server/tasks/test_vertex_task_converter.py
- Added a new test file for vertex_task_converter.py.
- Included tests for to_sdk_task_state, to_stored_task_state, to_stored_part, to_sdk_part, to_stored_artifact, to_sdk_artifact, to_stored_task, and to_sdk_task functions.
tests/server/tasks/test_vertex_task_store.py
- Added a new test file for VertexTaskStore.
- Implemented asynchronous tests for save, get, update (status, metadata, adding/updating/deleting artifacts), and metadata field mapping.
- Created a pytest_asyncio fixture to set up and tear down the VertexTaskStore with a Vertex AI Agent Engine.

Activity

The pull request is currently marked as 'Work In Progress' (WIP).
This PR aims to address issue [Feat]: add support for the vertex managed task store service #751, indicating it's part of a planned feature or bug fix.

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.

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 a new TaskStore implementation based on Google Cloud Vertex AI, along with the necessary type converters and integration tests. The changes are well-structured, with a clear separation of concerns between the store and the data conversion logic. The addition of comprehensive integration tests that run against a sandbox environment is a great way to ensure correctness.

I've identified a few areas for improvement, including removing hardcoded test configurations, addressing some copy-paste errors in docstrings, removing debug print statements, and clarifying the handling of task history. There's also a dependency on a non-standard git reference in pyproject.toml that might cause issues in the long run. All comments align with the provided rules or are not covered by them, and thus no comments were dropped or modified.

pyproject.toml

src/a2a/server/tasks/vertex_task_converter.py

src/a2a/server/tasks/vertex_task_store.py

src/a2a/contrib/tasks/vertex_task_store.py

tck/sut_agent_with_vertex_task_store.py

scripts/run_vertex_tests.sh

src/a2a/server/tasks/vertex_task_converter.py

src/a2a/server/tasks/vertex_task_store.py

ishymko

Nice!

Added a few non-critical double-checking comments around conversion logic and testing.

pyproject.toml

src/a2a/contrib/tasks/vertex_task_converter.py

uv.lock

ishymko · 2026-03-05T11:32:13Z

src/a2a/contrib/tasks/vertex_task_converter.py

+    if isinstance(part.root, DataPart):
+        data_bytes = json.dumps(part.root.data).encode('utf-8')
+        return types.Part(
+            inline_data=types.Blob(
+                mime_type='application/json', data=data_bytes
+            )
+        )


I believe this isn't going to roundtrip back to the A2A DataPart? Maybe application/json MIME type should be used as a signal when converting back (one can have a FileWithBytes representing the same though but it feels that this use case is going to be less popular)?

Yes, that's correct. So are you saying that it's better to turn FileWithBytes of type application/json into DataPart than vice versa?

src/a2a/contrib/tasks/vertex_task_converter.py

ishymko · 2026-03-05T11:50:45Z

tests/contrib/tasks/test_vertex_task_converter.py

Maybe some roundtrip tests can be added here (or existing replaced with roundtrip)?

I see that test_vertex_task_store.py does save-read tests but as it requires real Vertex it's not going to be executed in CI.

tests/contrib/tasks/test_vertex_task_store.py

…ssist suggestion)

gaborfeher requested a review from a team as a code owner March 2, 2026 14:57

gemini-code-assist bot reviewed Mar 2, 2026

View reviewed changes

gaborfeher force-pushed the main branch 11 times, most recently from ad1e1ce to fc58455 Compare March 4, 2026 17:04

Implement a vertex based task store

708cd9f

gaborfeher force-pushed the main branch from fc58455 to 708cd9f Compare March 4, 2026 17:17

ishymko reviewed Mar 5, 2026

View reviewed changes

gaborfeher added 4 commits March 5, 2026 12:55

change to the base directory before invoking test runs (gemini-code-a…

78e9da3

…ssist suggestion)

Address review comments

95a3887

Fix test comments and script location

7fdafe7

apply the ruff reformatter

18ce0c4

gaborfeher force-pushed the main branch from f7fabe9 to 18ce0c4 Compare March 5, 2026 13:46

fix the handling of unknown part types

0490e99

gaborfeher force-pushed the main branch from 1196db8 to 0490e99 Compare March 5, 2026 15:03

Create a fake vertex backend for testing

c665e3c

Conversation

gaborfeher commented Mar 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

gemini-code-assist bot commented Mar 2, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ishymko left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ishymko Mar 5, 2026

Choose a reason for hiding this comment

Uh oh!

gaborfeher Mar 5, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

ishymko Mar 5, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

gaborfeher commented Mar 2, 2026 •

edited

Loading