feat: Add support for custom task store in to_a2a function#3839
feat: Add support for custom task store in to_a2a function#3839secprog wants to merge 20 commits intogoogle:mainfrom
Conversation
Summary of ChangesHello @secprog, 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 significant enhancement to the Highlights
Using Gemini Code AssistThe 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
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 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
|
adk-bot
added
the
core
|
Response from ADK Triaging Agent Hello @secprog, thank you for creating this PR! To help us review this PR, could you please associate a GitHub issue with it? If there is no existing issue, could you please create one? In addition, could you please add a This information will help reviewers to review your PR more efficiently. Thanks! |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request enhances the to_a2a function by adding an optional task_store parameter. This allows users to inject a custom task store, providing more flexibility, while defaulting to InMemoryTaskStore for backward compatibility and ease of use. The implementation is clean and follows best practices. A corresponding unit test has been added to verify that a provided custom task store is used correctly and that the default is not instantiated in that case. The changes are well-executed and improve the utility of the to_a2a function. Overall, this is a great contribution.
rohityan
added
a2a
|
Hi @secprog ,Thank you for your contribution! We appreciate you taking the time to submit this pull request. |
rohityan
added
the
request clarification
|
Hi @seanzhougoogle , can you please review this. Looks like a good feature. |
rohityan
added
the
needs review
|
@ryanaiagent, @seanzhougoogle ready to review, i runned autoformat.sh for fix the lint errors. |
|
Hi @ryanaiagent, @seanzhougoogle any update on this? |
|
Hi @secprog , we appreciate your patience and support. Can you please fix the failing unit tests. |
rohityan
added
request clarification
|
Hi @ryanaiagent the tests are not the one's introduced with this PR, but i updated the current PR with the main branch maybe it is fixed now, let's give it a try :). |
rohityan
added
needs review
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request successfully introduces the ability to provide a custom task store to the to_a2a function, enhancing its flexibility. The changes include updating the function signature, adding comprehensive documentation for the new parameter, and providing an example of its usage. Crucially, a new unit test has been added to verify that the custom task store is correctly utilized and that the default InMemoryTaskStore is not instantiated when a custom one is supplied. The minor whitespace changes in other files are also cleanups. Overall, the implementation is clean, well-documented, and adequately tested.
|
Hi @secprog , Your PR has been received by the team and is currently under review. We will provide feedback as soon as we have an update to share. |
|
Hi @GWeale , can you please review this. |
…er.py ADK v1.27.x hardcodes InMemoryTaskStore in get_fast_api_app with no injection point (tracked in upstream PR #3839). This adds an interim workaround: - code_agent/a2a/stores.py: patch_adk_stores() monkey-patches a2a.server.tasks.InMemoryTaskStore (and InMemoryPushNotificationConfigStore) with DatabaseTaskStore-backed equivalents when DATABASE_URL is set. The patch must run before get_fast_api_app is called since ADK imports the class at function-call time inside the `if a2a:` block. - main.py: calls patch_adk_stores() before get_fast_api_app so tasks survive pod restarts and are shared across replicas in EKS. Falls back to InMemory when DATABASE_URL is absent (local/CI). - code_agent/a2a/server.py: marked DEPRECATED with full migration rationale. The custom JSON-RPC dispatcher is superseded by ADK's built-in A2A handler + the stores patch. File kept for reference until EKS rollout is confirmed. Remove stores.py + the patch_adk_stores() call in main.py once ADK exposes a native a2a_task_store parameter in get_fast_api_app. Track: google/adk-python#3839 https://claude.ai/code/session_01KKh2oYnyYKq4MM7JoDNtvY
…t_api_app ## Problem `get_fast_api_app()` unconditionally creates an `InMemoryTaskStore` and `InMemoryPushNotificationConfigStore`, making it impossible for callers to inject persistent or shared stores without patching ADK internals. This is especially painful in production deployments where: - Multiple replicas need a shared task store to route A2A callbacks correctly - Operators want task state to survive server restarts (e.g. SQLite/Postgres) ## Solution Adds two new optional keyword arguments to `get_fast_api_app()`: - `a2a_task_store: Optional[Any] = None` - `a2a_push_config_store: Optional[Any] = None` When `None` (the default), the existing `InMemory*` defaults are used — fully backward-compatible. When provided, the caller-supplied instances are forwarded directly to `DefaultRequestHandler`. This mirrors the pattern introduced for the lower-level `to_a2a()` helper in PR google#3839. ## Usage from a2a.server.tasks import DatabaseTaskStore from sqlalchemy.ext.asyncio import create_async_engine engine = create_async_engine("postgresql+asyncpg://user:pw@host/db") app = get_fast_api_app( agents_dir="agents/", web=True, a2a=True, a2a_task_store=DatabaseTaskStore(engine), ) ## Tests Two new test cases added to tests/unittests/cli/test_fast_api.py: - test_a2a_uses_in_memory_task_store_by_default - test_a2a_custom_task_store_bypasses_in_memory_default
…t_api_app ## Problem `get_fast_api_app()` unconditionally creates an `InMemoryTaskStore` and `InMemoryPushNotificationConfigStore`, making it impossible for callers to inject persistent or shared stores without patching ADK internals. This is especially painful in production deployments where: - Multiple replicas need a shared task store to route A2A callbacks correctly - Operators want task state to survive server restarts (e.g. SQLite/Postgres) ## Solution Adds two new optional keyword arguments to `get_fast_api_app()`: - `a2a_task_store: Optional[Any] = None` - `a2a_push_config_store: Optional[Any] = None` When `None` (the default), the existing `InMemory*` defaults are used — fully backward-compatible. When provided, the caller-supplied instances are forwarded directly to `DefaultRequestHandler`. This mirrors the pattern introduced for the lower-level `to_a2a()` helper in PR google#3839. ## Usage from a2a.server.tasks import DatabaseTaskStore from sqlalchemy.ext.asyncio import create_async_engine engine = create_async_engine("postgresql+asyncpg://user:pw@host/db") app = get_fast_api_app( agents_dir="agents/", web=True, a2a=True, a2a_task_store=DatabaseTaskStore(engine), ) ## Tests Two new test cases added to tests/unittests/cli/test_fast_api.py: - test_a2a_uses_in_memory_task_store_by_default - test_a2a_custom_task_store_bypasses_in_memory_default
…t_api_app ## Problem `get_fast_api_app()` unconditionally creates an `InMemoryTaskStore` and `InMemoryPushNotificationConfigStore`, making it impossible for callers to inject persistent or shared stores without patching ADK internals. This is especially painful in production deployments where: - Multiple replicas need a shared task store to route A2A callbacks correctly - Operators want task state to survive server restarts (e.g. SQLite/Postgres) ## Solution Adds two new optional keyword arguments to `get_fast_api_app()`: - `a2a_task_store: Optional[Any] = None` - `a2a_push_config_store: Optional[Any] = None` When `None` (the default), the existing `InMemory*` defaults are used — fully backward-compatible. When provided, the caller-supplied instances are forwarded directly to `DefaultRequestHandler`. This mirrors the pattern introduced for the lower-level `to_a2a()` helper in PR google#3839. ## Usage from a2a.server.tasks import DatabaseTaskStore from sqlalchemy.ext.asyncio import create_async_engine engine = create_async_engine("postgresql+asyncpg://user:pw@host/db") app = get_fast_api_app( agents_dir="agents/", web=True, a2a=True, a2a_task_store=DatabaseTaskStore(engine), ) ## Tests Two new test cases added to tests/unittests/cli/test_fast_api.py: - test_a2a_uses_in_memory_task_store_by_default - test_a2a_custom_task_store_bypasses_in_memory_default
3 participants
This update introduces an optional task_store parameter to the to_a2a function, allowing users to provide a custom task store instance. If not specified, the function defaults to using InMemoryTaskStore. Additionally, unit tests have been added to verify the functionality with a custom task store.
Replaces my old closed PR #3639.