Refactor database stores for robust custom conversion handling and improve v0.3 compatibility.

Author: sokoliva

src/a2a/compat/v0_3/conversions.py

@@ -1395,9 +1395,7 @@ def compat_task_model_to_core(task_model: TaskModel) -> pb2_v10.Task:
     compat_task = types_v03.Task(
         id=task_model.id,
         context_id=task_model.context_id,
-        status=types_v03.TaskStatus.model_validate(task_model.status)
-        if task_model.status
-        else None,
+        status=types_v03.TaskStatus.model_validate(task_model.status),
         artifacts=(
             [types_v03.Artifact.model_validate(a) for a in task_model.artifacts]
             if task_model.artifacts

src/a2a/server/tasks/database_push_notification_config_store.py

@@ -1,4 +1,5 @@
 # ruff: noqa: PLC0415
+import inspect
 import logging
 
 from typing import TYPE_CHECKING
@@ -13,9 +14,7 @@
         AsyncSession,
         async_sessionmaker,
     )
-    from sqlalchemy.orm import (
-        class_mapper,
-    )
+    from sqlalchemy.orm import class_mapper
 except ImportError as e:
     raise ImportError(
         'DatabasePushNotificationConfigStore requires SQLAlchemy and a database driver. '
@@ -157,12 +156,11 @@ def _to_orm(
 
         The config data is serialized to JSON bytes, and encrypted if a key is configured.
         """
-        if self.core_to_model_conversion:
-            conversion = self.core_to_model_conversion
-            # bound method
-            if hasattr(conversion, '__func__'):
+        if conversion := self.core_to_model_conversion:
+            # If it's a bound method of this instance, call the underlying function
+            # to avoid passing 'self' twice.
+            if inspect.ismethod(conversion):
                 return conversion.__func__(task_id, config, owner, self._fernet)
-            # instance method
             return conversion(task_id, config, owner, self._fernet)
 
         json_payload = MessageToJson(config).encode('utf-8')
@@ -187,12 +185,11 @@ def _from_orm(
 
         Handles decryption if a key is configured, with a fallback to plain JSON.
         """
-        if self.model_to_core_conversion:
-            conversion = self.model_to_core_conversion
-            # bound method
-            if hasattr(conversion, '__func__'):
+        if conversion := self.model_to_core_conversion:
+            # If it's a bound method of this instance, call the underlying function
+            # to avoid passing 'self' twice.
+            if inspect.ismethod(conversion):
                 return conversion.__func__(model_instance)
-            # instance method
             return conversion(model_instance)
 
         payload = model_instance.config_data

src/a2a/server/tasks/database_task_store.py

@@ -1,26 +1,18 @@
+import inspect
 import logging
 
 from datetime import datetime, timezone
 from typing import TYPE_CHECKING
 
 
 try:
-    from sqlalchemy import (
-        Table,
-        and_,
-        delete,
-        func,
-        or_,
-        select,
-    )
+    from sqlalchemy import Table, and_, delete, func, or_, select
     from sqlalchemy.ext.asyncio import (
         AsyncEngine,
         AsyncSession,
         async_sessionmaker,
     )
-    from sqlalchemy.orm import (
-        class_mapper,
-    )
+    from sqlalchemy.orm import class_mapper
 except ImportError as e:
     raise ImportError(
         'DatabaseTaskStore requires SQLAlchemy and a database driver. '
@@ -125,12 +117,11 @@ async def _ensure_initialized(self) -> None:
 
     def _to_orm(self, task: Task, owner: str) -> TaskModel:
         """Maps a Proto Task to a SQLAlchemy TaskModel instance."""
-        if self.core_to_model_conversion:
-            conversion = self.core_to_model_conversion
-            # bound method
-            if hasattr(conversion, '__func__'):
+        if conversion := self.core_to_model_conversion:
+            # If it's a bound method of this instance, call the underlying function
+            # to avoid passing 'self' twice.
+            if inspect.ismethod(conversion):
                 return conversion.__func__(task, owner)
-            # instance method
             return conversion(task, owner)
 
         return self.task_model(
@@ -154,12 +145,11 @@ def _to_orm(self, task: Task, owner: str) -> TaskModel:
 
     def _from_orm(self, task_model: TaskModel) -> Task:
         """Maps a SQLAlchemy TaskModel to a Proto Task instance."""
-        if self.model_to_core_conversion:
-            conversion = self.model_to_core_conversion
-            # bound method
-            if hasattr(conversion, '__func__'):
+        if conversion := self.model_to_core_conversion:
+            # If it's a bound method of this instance, call the underlying function
+            # to avoid passing 'self' twice.
+            if inspect.ismethod(conversion):
                 return conversion.__func__(task_model)
-            # instance method
             return conversion(task_model)
 
         if task_model.protocol_version == '1.0':

tests/server/tasks/test_database_push_notification_config_store.py

@@ -820,21 +820,29 @@ async def test_custom_conversion():
         DatabasePushNotificationConfigStore.model_to_core_conversion = None
 
 
+@pytest.mark.parametrize('assignment_type', ['class', 'instance'])
 @pytest.mark.asyncio
 async def test_core_to_0_3_model_conversion(
     db_store_parameterized: DatabasePushNotificationConfigStore,
+    assignment_type: str,
 ) -> None:
     """Test storing and retrieving push notification configs in v0.3 format using conversion utilities.
 
+    Tests both class-level and instance-level assignment of the conversion function.
     Setting the model_to_core_conversion to compat_push_notification_config_model_to_core would be redundant as
     it is always called when retrieving 0.3 PushNotificationConfigs.
     """
     store = db_store_parameterized
 
     # Set the v0.3 persistence utilities
-    DatabasePushNotificationConfigStore.core_to_model_conversion = (
-        core_to_compat_push_notification_config_model
-    )
+    if assignment_type == 'class':
+        DatabasePushNotificationConfigStore.core_to_model_conversion = (
+            core_to_compat_push_notification_config_model
+        )
+    else:
+        store.core_to_model_conversion = (
+            core_to_compat_push_notification_config_model
+        )
 
     try:
         task_id = 'v03-persistence-task'
@@ -877,6 +885,10 @@ async def test_core_to_0_3_model_conversion(
         assert retrieved.token == original_config.token
 
     finally:
-        # Reset class variables
-        DatabasePushNotificationConfigStore.core_to_model_conversion = None
+        # Reset conversion attributes
+        if assignment_type == 'class':
+            DatabasePushNotificationConfigStore.core_to_model_conversion = None
+        else:
+            store.core_to_model_conversion = None
+
         await store.delete_info(task_id, MINIMAL_CALL_CONTEXT)

tests/server/tasks/test_database_task_store.py

@@ -830,9 +830,6 @@ async def test_get_0_3_task_detailed(
     await db_store_parameterized.delete(task_id, context_user)
 
 
-# Ensure aiosqlite, asyncpg, and aiomysql are installed in the test environment (added to pyproject.toml).
-
-
 @pytest.mark.asyncio
 async def test_custom_conversion():
     engine = MagicMock()
@@ -863,19 +860,25 @@ async def test_custom_conversion():
         DatabaseTaskStore.model_to_core_conversion = None
 
 
+@pytest.mark.parametrize('assignment_type', ['class', 'instance'])
 @pytest.mark.asyncio
 async def test_core_to_0_3_model_conversion(
     db_store_parameterized: DatabaseTaskStore,
+    assignment_type: str,
 ) -> None:
     """Test storing and retrieving tasks in v0.3 format using conversion utilities.
 
-    Setting the model_to_core_conversion class variables to compat_task_model_to_core is redundant
+    Tests both class-level and instance-level assignment of the conversion function.
+    Setting the model_to_core_conversion class variables to compat_task_model_to_core would be redundant
     as it is always called when retrieving 0.3 tasks.
     """
     store = db_store_parameterized
 
     # Set the v0.3 persistence utilities
-    DatabaseTaskStore.core_to_model_conversion = core_to_compat_task_model
+    if assignment_type == 'class':
+        DatabaseTaskStore.core_to_model_conversion = core_to_compat_task_model
+    else:
+        store.core_to_model_conversion = core_to_compat_task_model
 
     try:
         task_id = 'v03-persistence-task'
@@ -905,7 +908,13 @@ async def test_core_to_0_3_model_conversion(
         assert dict(retrieved_task.metadata) == {'key': 'value'}
 
     finally:
-        # Reset class variables
-        DatabaseTaskStore.core_to_model_conversion = None
-        DatabaseTaskStore.model_to_core_conversion = None
+        # Reset conversion attributes
+        if assignment_type == 'class':
+            DatabaseTaskStore.core_to_model_conversion = None
+        else:
+            store.core_to_model_conversion = None
+
         await store.delete('v03-persistence-task')
+
+
+# Ensure aiosqlite, asyncpg, and aiomysql are installed in the test environment (added to pyproject.toml).