Refactoring and bug fixes

Author: kthota-g

.github/workflows/unit-tests.yml

@@ -35,7 +35,7 @@ jobs:
         run: |
           sudo apt-get update && sudo apt-get install -y postgresql-client
           PGPASSWORD=postgres psql -h localhost -p 5432 -U postgres -d a2a_test -f ${{ github.workspace }}/docker/postgres/init.sql
-          export POSTGRES_TEST_DSN="postgresql://postgres:postgres@localhost:5432/a2a_test"
+          export POSTGRES_TEST_DSN="postgresql+asyncpg://postgres:postgres@localhost:5432/a2a_test"
 
       - name: Install uv
         uses: astral-sh/setup-uv@v6

pyproject.toml

@@ -20,7 +20,7 @@ dependencies = [
   "grpcio>=1.60",
   "grpcio-tools>=1.60",
   "grpcio_reflection>=1.7.0",
-  "protobuf==5.29.5",
+  "protobuf==5.29.5"
 ]
 
 classifiers = [
@@ -38,22 +38,16 @@ classifiers = [
 
 [project.optional-dependencies]
 postgresql = [
-    "sqlalchemy>=2.0.0",
-    "asyncpg>=0.30.0",
+    "sqlalchemy[asyncio,postgresql-asyncpg]>=2.0.0"
 ]
 mysql = [
-    "sqlalchemy>=2.0.0",
-    "aiomysql>=0.2.0",
+    "sqlalchemy[asyncio,aiomysql]>=2.0.0"
 ]
 sqlite = [
-    "sqlalchemy>=2.0.0",
-    "aiosqlite>=0.19.0",
+    "sqlalchemy[asyncio,aiosqlite]>=2.0.0"
 ]
 sql = [
-    "sqlalchemy>=2.0.0",
-    "asyncpg>=0.30.0",
-    "aiomysql>=0.2.0",
-    "aiosqlite>=0.19.0",
+    "sqlalchemy[asyncio,postgresql-asyncpg,aiomysql,aiosqlite]>=2.0.0"
 ]
 
 [project.urls]
@@ -104,8 +98,7 @@ dev = [
   "uv-dynamic-versioning>=0.8.2",
   "types-protobuf",
   "types-requests",
-  "pre-commit",
-  "asyncpg-stubs>=0.30.1",
+  "pre-commit"
 ]
 
 [[tool.uv.index]]

src/a2a/server/models.py

@@ -109,9 +109,11 @@ class Base(DeclarativeBase):
 class TaskMixin:
     """Mixin providing standard task columns with proper type handling."""
 
-    id: Mapped[str] = mapped_column(String, primary_key=True, index=True)
-    contextId: Mapped[str] = mapped_column(String, nullable=False)  # noqa: N815
-    kind: Mapped[str] = mapped_column(String, nullable=False, default='task')
+    id: Mapped[str] = mapped_column(String(36), primary_key=True, index=True)
+    contextId: Mapped[str] = mapped_column(String(36), nullable=False)  # noqa: N815
+    kind: Mapped[str] = mapped_column(
+        String(16), nullable=False, default='task'
+    )
 
     # Properly typed Pydantic fields with automatic serialization
     status: Mapped[TaskStatus] = mapped_column(PydanticType(TaskStatus))

src/a2a/server/tasks/database_task_store.py

@@ -124,6 +124,7 @@ async def save(self, task: Task) -> None:
                     .values(**update_data)
                 )
                 await session.execute(stmt_update)
+                logger.debug(f'Task {task.id} updated successfully.')
             else:
                 logger.debug(f'Saving new task {task.id} to the database.')
                 # Map Pydantic fields to database columns
@@ -139,8 +140,7 @@ async def save(self, task: Task) -> None:
                     ),  # Map metadata field to task_metadata column
                 )
                 session.add(new_task_model)
-            # Commit is automatic when using session.begin()
-        logger.info(f'Task {task.id} saved successfully.')
+                logger.info(f'Task {task.id} created successfully.')
 
     async def get(self, task_id: str) -> Task | None:
         """Retrieves a task from the database by ID."""

tests/server/tasks/test_database_task_store.py

@@ -125,17 +125,13 @@ async def db_store_parameterized(
             await conn.run_sync(Base.metadata.create_all)
 
         # create_table=False as we've explicitly created tables above.
-        store = DatabaseTaskStore(db_url, create_table=False)
+        store = DatabaseTaskStore(engine=engine, create_table=False)
         # Initialize the store (connects, etc.). Safe to call even if tables exist.
         await store.initialize()
 
         yield store
 
     finally:
-        # Teardown
-        if store:  # If store was successfully created
-            await store.close()  # Closes the store's own engine
-
         if engine:  # If engine was created for setup/teardown
             # Drop tables using the fixture's engine
             async with engine.begin() as conn: