diff --git a/src/a2a/server/cluster/database_task_store.py b/src/a2a/server/cluster/database_task_store.py index fb9c8ffa0..dc0bf71b6 100644 --- a/src/a2a/server/cluster/database_task_store.py +++ b/src/a2a/server/cluster/database_task_store.py @@ -258,6 +258,7 @@ async def get( commits) rather than failing the read. """ await self._db._ensure_initialized() # noqa: SLF001 + await self._ensure_cluster_tables() owner = self._db.owner_resolver(context) attempts = 0 while True: @@ -355,6 +356,8 @@ async def list( async def delete(self, task_id: str, context: ServerCallContext) -> None: """Deletes a task and its version row via the underlying store.""" + await self._db._ensure_initialized() # noqa: SLF001 + await self._ensure_cluster_tables() owner = self._db.owner_resolver(context) await self._db.delete(task_id, context) async with self._db.async_session_maker.begin() as session: diff --git a/tests/server/cluster/test_database_versioned_task_store.py b/tests/server/cluster/test_database_versioned_task_store.py index 624dee448..a00dc3612 100644 --- a/tests/server/cluster/test_database_versioned_task_store.py +++ b/tests/server/cluster/test_database_versioned_task_store.py @@ -136,6 +136,24 @@ async def versioned_store( await engine.dispose() +@pytest_asyncio.fixture(params=DB_CONFIGS) +async def uninitialized_store( + request, +) -> AsyncGenerator[VersionedDatabaseTaskStore, None]: + """A store that creates its tables lazily, without `initialize()`.""" + db_url, dialect_name = request.param + if db_url is None: + pytest.skip(f'DSN for {dialect_name} not set in environment variables.') + engine = create_async_engine(db_url) + _enforce_sqlite_foreign_keys(engine) + try: + yield VersionedDatabaseTaskStore(engine=engine, create_table=True) + finally: + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.drop_all) + await engine.dispose() + + @pytest.mark.asyncio async def test_is_a_versioned_task_store( versioned_store: VersionedDatabaseTaskStore, @@ -654,6 +672,21 @@ async def test_delete_removes_version_row( assert stored is not None +@pytest.mark.asyncio +async def test_get_without_initialize_creates_tables( + uninitialized_store: VersionedDatabaseTaskStore, +) -> None: + assert await uninitialized_store.get('task-abc', TEST_CONTEXT) is None + + +@pytest.mark.asyncio +async def test_delete_without_initialize_creates_tables( + uninitialized_store: VersionedDatabaseTaskStore, +) -> None: + await uninitialized_store.delete('task-abc', TEST_CONTEXT) + assert await uninitialized_store.get('task-abc', TEST_CONTEXT) is None + + @pytest.mark.asyncio async def test_versions_count_writes( versioned_store: VersionedDatabaseTaskStore,