Skip to content

feat(a2a): add a2a_task_store and a2a_push_config_store params to get_fast_api_app#4970

Open
STHITAPRAJNAS wants to merge 2 commits intogoogle:mainfrom
STHITAPRAJNAS:feat/fast-api-custom-task-store
Open

feat(a2a): add a2a_task_store and a2a_push_config_store params to get_fast_api_app#4970
STHITAPRAJNAS wants to merge 2 commits intogoogle:mainfrom
STHITAPRAJNAS:feat/fast-api-custom-task-store

Conversation

@STHITAPRAJNAS
Copy link

@STHITAPRAJNAS STHITAPRAJNAS commented Mar 24, 2026

Closes #4971.

Follows the same pattern as #3839 (to_a2a), but for the higher-level get_fast_api_app entry point.

Currently both stores are unconditionally created as in-memory instances inside the function, with no way for callers to inject their own. This makes it impossible to use a persistent backend (e.g. DatabaseTaskStore backed by SQLite or Postgres) without monkey-patching.

The fix adds two optional kwargs:

def get_fast_api_app(
    ...
    a2a_task_store: Optional[Any] = None,
    a2a_push_config_store: Optional[Any] = None,
    ...
)

When None (the default), the existing InMemory* instances are created as before, so there is no change in behaviour for existing callers.

Both stores are now initialised once before the per-agent loop rather than push_config_store being re-created on every iteration, which was likely unintentional in the original code.

Two tests added:

  • test_a2a_uses_in_memory_task_store_by_default — default path unchanged
  • test_a2a_custom_task_store_is_used — custom store forwarded, InMemoryTaskStore not constructed

…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
@google-cla
Copy link

google-cla bot commented Mar 24, 2026

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

@adk-bot
Copy link
Collaborator

adk-bot commented Mar 24, 2026

Response from ADK Triaging Agent

Hello @STHITAPRAJNAS, thank you for your contribution!

To help us track this new feature, could you please create a GitHub issue and associate it with this PR?

Additionally, it looks like the Contributor License Agreement (CLA) check has not passed. Please ensure you have signed the CLA, as it is a requirement for us to accept contributions. You can find more details in the "cla/google" check at the bottom of the pull request.

Taking care of these items will help us to review and merge your PR more efficiently. Thank you!

- hoist a2a_push_config_store init outside per-agent loop, matching
  the task_store pattern (both now use if-None guard at the top)
- remove redundant local push_config_store variable; use
  a2a_push_config_store directly in DefaultRequestHandler call
- trim docstring: drop inline Example block, shorten to match the
  one-liner style used by the other params in the same function
- rename test to test_a2a_custom_task_store_is_used (shorter, matches
  existing naming style)
- fix test mock setup to use an explicit mock_a2a_app_instance like
  the test_app_with_a2a fixture does
- simplify test assertion: drop intermediate call_kwargs variable
@STHITAPRAJNAS STHITAPRAJNAS changed the title feat(a2a): expose a2a_task_store and a2a_push_config_store in get_fast_api_app feat(a2a): add a2a_task_store and a2a_push_config_store params to get_fast_api_app Mar 24, 2026
@adk-bot adk-bot added the core [Component] This issue is related to the core interface and implementation label Mar 24, 2026
@rohityan rohityan self-assigned this Mar 24, 2026
@rohityan rohityan added the request clarification [Status] The maintainer need clarification or more information from the author label Mar 24, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

core [Component] This issue is related to the core interface and implementation request clarification [Status] The maintainer need clarification or more information from the author

Projects

None yet

Development

Successfully merging this pull request may close these issues.

get_fast_api_app does not allow injecting a custom A2A task store

3 participants