refactor: replace PydanticType/PydanticListType with explicit JSON serialization/deserialization for protobuf models and add protocol versioning to task storage.

Author: sokoliva

src/a2a/server/models.py

@@ -1,5 +1,5 @@
 from datetime import datetime
-from typing import TYPE_CHECKING, Any, Generic, TypeVar
+from typing import TYPE_CHECKING, Any
 
 
 if TYPE_CHECKING:
@@ -11,26 +11,14 @@ def override(func):  # noqa: ANN001, ANN201
         return func
 
 
-from google.protobuf.json_format import MessageToDict, ParseDict, ParseError
-from google.protobuf.message import Message as ProtoMessage
-from pydantic import BaseModel, ValidationError
-
-from a2a.compat.v0_3 import conversions
-from a2a.compat.v0_3 import types as types_v03
-from a2a.types.a2a_pb2 import Artifact, Message, TaskStatus
-
-
 try:
-    from sqlalchemy import JSON, DateTime, Dialect, Index, LargeBinary, String
+    from sqlalchemy import JSON, DateTime, Index, LargeBinary, String
     from sqlalchemy.orm import (
         DeclarativeBase,
         Mapped,
         declared_attr,
         mapped_column,
     )
-    from sqlalchemy.types import (
-        TypeDecorator,
-    )
 except ImportError as e:
     raise ImportError(
         'Database models require SQLAlchemy. '
@@ -42,130 +30,6 @@ def override(func):  # noqa: ANN001, ANN201
     ) from e
 
 
-T = TypeVar('T')
-
-
-class PydanticType(TypeDecorator[T], Generic[T]):
-    """SQLAlchemy type that handles Pydantic model and Protobuf message serialization."""
-
-    impl = JSON
-    cache_ok = True
-
-    def __init__(self, pydantic_type: type[T], **kwargs: dict[str, Any]):
-        """Initialize the PydanticType.
-
-        Args:
-            pydantic_type: The Pydantic model or Protobuf message type to handle.
-            **kwargs: Additional arguments for TypeDecorator.
-        """
-        self.pydantic_type = pydantic_type
-        super().__init__(**kwargs)
-
-    def process_bind_param(
-        self, value: T | None, dialect: Dialect
-    ) -> dict[str, Any] | None:
-        """Convert Pydantic model or Protobuf message to a JSON-serializable dictionary for the database."""
-        if value is None:
-            return None
-        if isinstance(value, ProtoMessage):
-            return MessageToDict(value, preserving_proto_field_name=False)
-        if isinstance(value, BaseModel):
-            return value.model_dump(mode='json')
-        return value  # type: ignore[return-value]
-
-    def process_result_value(
-        self, value: dict[str, Any] | None, dialect: Dialect
-    ) -> T | None:
-        """Convert a JSON-like dictionary from the database back to a Pydantic model or Protobuf message."""
-        if value is None:
-            return None
-        # Check if it's a protobuf message class
-        if isinstance(self.pydantic_type, type) and issubclass(
-            self.pydantic_type, ProtoMessage
-        ):
-            try:
-                return ParseDict(value, self.pydantic_type())  # type: ignore[return-value]
-            except (ParseError, ValueError):
-                # Try legacy conversion
-                legacy_map = _get_legacy_conversions()
-                if self.pydantic_type in legacy_map:
-                    legacy_type, convert_func = legacy_map[self.pydantic_type]
-                    try:
-                        legacy_instance = legacy_type.model_validate(value)
-                        return convert_func(legacy_instance)
-                    except ValidationError:
-                        pass
-                raise
-        # Assume it's a Pydantic model
-        return self.pydantic_type.model_validate(value)  # type: ignore[attr-defined]
-
-
-class PydanticListType(TypeDecorator, Generic[T]):
-    """SQLAlchemy type that handles lists of Pydantic models or Protobuf messages."""
-
-    impl = JSON
-    cache_ok = True
-
-    def __init__(self, pydantic_type: type[T], **kwargs: dict[str, Any]):
-        """Initialize the PydanticListType.
-
-        Args:
-            pydantic_type: The Pydantic model or Protobuf message type for items in the list.
-            **kwargs: Additional arguments for TypeDecorator.
-        """
-        self.pydantic_type = pydantic_type
-        super().__init__(**kwargs)
-
-    def process_bind_param(
-        self, value: list[T] | None, dialect: Dialect
-    ) -> list[dict[str, Any]] | None:
-        """Convert a list of Pydantic models or Protobuf messages to a JSON-serializable list for the DB."""
-        if value is None:
-            return None
-        result: list[dict[str, Any]] = []
-        for item in value:
-            if isinstance(item, ProtoMessage):
-                result.append(
-                    MessageToDict(item, preserving_proto_field_name=False)
-                )
-            elif isinstance(item, BaseModel):
-                result.append(item.model_dump(mode='json'))
-            else:
-                result.append(item)  # type: ignore[arg-type]
-        return result
-
-    def process_result_value(
-        self, value: list[dict[str, Any]] | None, dialect: Dialect
-    ) -> list[T] | None:
-        """Convert a JSON-like list from the DB back to a list of Pydantic models or Protobuf messages."""
-        if value is None:
-            return None
-        # Check if it's a protobuf message class
-        if isinstance(self.pydantic_type, type) and issubclass(
-            self.pydantic_type, ProtoMessage
-        ):
-            result = []
-            legacy_map = _get_legacy_conversions()
-            legacy_info = legacy_map.get(self.pydantic_type)
-
-            for item in value:
-                try:
-                    result.append(ParseDict(item, self.pydantic_type()))
-                except (ParseError, ValueError):  # noqa: PERF203
-                    if legacy_info:
-                        legacy_type, convert_func = legacy_info
-                        try:
-                            legacy_instance = legacy_type.model_validate(item)
-                            result.append(convert_func(legacy_instance))
-                            continue
-                        except ValidationError:
-                            pass
-                    raise
-            return result  # type: ignore[return-value]
-        # Assume it's a Pydantic model
-        return [self.pydantic_type.model_validate(item) for item in value]  # type: ignore[attr-defined]
-
-
 # Base class for all database models
 class Base(DeclarativeBase):
     """Base class for declarative models in A2A SDK."""
@@ -184,25 +48,17 @@ class TaskMixin:
     last_updated: Mapped[datetime | None] = mapped_column(
         DateTime, nullable=True
     )
-
-    # Properly typed Pydantic fields with automatic serialization
-    status: Mapped[TaskStatus] = mapped_column(PydanticType(TaskStatus))
-    artifacts: Mapped[list[Artifact] | None] = mapped_column(
-        PydanticListType(Artifact), nullable=True
-    )
-    history: Mapped[list[Message] | None] = mapped_column(
-        PydanticListType(Message), nullable=True
-    )
+    status: Mapped[Any] = mapped_column(JSON)
+    artifacts: Mapped[list[Any] | None] = mapped_column(JSON, nullable=True)
+    history: Mapped[list[Any] | None] = mapped_column(JSON, nullable=True)
     protocol_version: Mapped[str | None] = mapped_column(
         String(16), nullable=True
     )
 
-    # Using declared_attr to avoid conflict with Pydantic's metadata
-    @declared_attr
-    @classmethod
-    def task_metadata(cls) -> Mapped[dict[str, Any] | None]:
-        """Define the 'metadata' column, avoiding name conflicts with Pydantic."""
-        return mapped_column(JSON, nullable=True, name='metadata')
+    # Using 'task_metadata' to avoid conflict with SQLAlchemy's 'Base.metadata'
+    task_metadata: Mapped[dict[str, Any] | None] = mapped_column(
+        JSON, nullable=True, name='metadata'
+    )
 
     @override
     def __repr__(self) -> str:
@@ -329,15 +185,3 @@ class PushNotificationConfigModel(PushNotificationConfigMixin, Base):
     """Default push notification config model with standard table name."""
 
     __tablename__ = 'push_notification_configs'
-
-
-def _get_legacy_conversions() -> dict[type, tuple[type, Any]]:
-    """Get the mapping of current types to their legacy counterparts and conversion functions."""
-    return {
-        TaskStatus: (
-            types_v03.TaskStatus,
-            conversions.to_core_task_status,
-        ),
-        Message: (types_v03.Message, conversions.to_core_message),
-        Artifact: (types_v03.Artifact, conversions.to_core_artifact),
-    }

src/a2a/server/tasks/database_push_notification_config_store.py

@@ -167,22 +167,6 @@ def _to_orm(
             protocol_version='1.0',
         )
 
-    def _parse_config(
-        self, json_payload: str, protocol_version: str | None = None
-    ) -> PushNotificationConfig:
-        """Parses a JSON payload into a PushNotificationConfig proto.
-
-        Uses protocol_version to decide between modern parsing and legacy fallback.
-        """
-        if protocol_version == '1.0':
-            return Parse(json_payload, PushNotificationConfig())
-
-        # Legacy case: no version or older
-        legacy_instance = types_v03.PushNotificationConfig.model_validate_json(
-            json_payload
-        )
-        return conversions.to_core_push_notification_config(legacy_instance)
-
     def _from_orm(
         self, model_instance: PushNotificationConfigModel
     ) -> TaskPushNotificationConfig:
@@ -355,3 +339,22 @@ async def delete_info(
                     owner,
                     config_id,
                 )
+
+    def _parse_config(
+        self, json_payload: str, protocol_version: str | None = None
+    ) -> TaskPushNotificationConfig:
+        """Parses a JSON payload into a TaskPushNotificationConfig proto.
+
+        Uses protocol_version to decide between modern parsing and legacy fallback.
+        """
+        if protocol_version == '1.0':
+            return Parse(json_payload, TaskPushNotificationConfig())
+
+        legacy_instance = (
+            types_v03.TaskPushNotificationConfig.model_validate_json(
+                json_payload
+            )
+        )
+        return conversions.to_core_task_push_notification_config(
+            legacy_instance
+        )

src/a2a/server/tasks/database_task_store.py

@@ -31,8 +31,10 @@
         "or 'pip install a2a-sdk[sql]'"
     ) from e
 
-from google.protobuf.json_format import MessageToDict
+from google.protobuf.json_format import MessageToDict, ParseDict
 
+from a2a.compat.v0_3 import conversions
+from a2a.compat.v0_3 import types as types_v03
 from a2a.server.context import ServerCallContext
 from a2a.server.models import Base, TaskModel, create_task_model
 from a2a.server.owner_resolver import OwnerResolver, resolve_user_scope
@@ -117,8 +119,7 @@ async def _ensure_initialized(self) -> None:
 
     def _to_orm(self, task: Task, owner: str) -> TaskModel:
         """Maps a Proto Task to a SQLAlchemy TaskModel instance."""
-        # Pass proto objects directly - PydanticType/PydanticListType
-        # handle serialization via process_bind_param
+        task_dict = MessageToDict(task)
         return self.task_model(
             id=task.id,
             context_id=task.context_id,
@@ -129,33 +130,52 @@ def _to_orm(self, task: Task, owner: str) -> TaskModel:
                 if task.HasField('status') and task.status.HasField('timestamp')
                 else None
             ),
-            status=task.status if task.HasField('status') else None,
-            artifacts=list(task.artifacts) if task.artifacts else [],
-            history=list(task.history) if task.history else [],
-            task_metadata=(
-                MessageToDict(task.metadata) if task.metadata.fields else None
-            ),
+            status=task_dict.get('status'),
+            artifacts=task_dict.get('artifacts', []),
+            history=task_dict.get('history', []),
+            task_metadata=task_dict.get('metadata'),
+            protocol_version='1.0',
         )
 
     def _from_orm(self, task_model: TaskModel) -> Task:
         """Maps a SQLAlchemy TaskModel to a Proto Task instance."""
-        # PydanticType/PydanticListType already deserialize to proto objects
-        # via process_result_value, so we can construct the Task directly
+        # Data is stored as raw JSON (dicts/lists), so we parse it manually
         task = Task(
             id=task_model.id,
             context_id=task_model.context_id,
         )
-        if task_model.status:
-            task.status.CopyFrom(task_model.status)
-        if task_model.artifacts:
-            task.artifacts.extend(task_model.artifacts)
-        if task_model.history:
-            task.history.extend(task_model.history)
-        if task_model.task_metadata:
-            task.metadata.update(
-                cast('dict[str, Any]', task_model.task_metadata)
-            )
-        return task
+        if task_model.protocol_version == '1.0':
+            if task_model.status:
+                ParseDict(
+                    cast('dict[str, Any]', task_model.status), task.status
+                )
+            if task_model.artifacts:
+                for art_dict in cast(
+                    'list[dict[str, Any]]', task_model.artifacts
+                ):
+                    art = task.artifacts.add()
+                    ParseDict(art_dict, art)
+            if task_model.history:
+                for msg_dict in cast(
+                    'list[dict[str, Any]]', task_model.history
+                ):
+                    msg = task.history.add()
+                    ParseDict(msg_dict, msg)
+            if task_model.task_metadata:
+                task.metadata.update(
+                    cast('dict[str, Any]', task_model.task_metadata)
+                )
+            return task
+        # Reconstruct legacy task from raw columns (which are dicts/lists here)
+        legacy_task = types_v03.Task(
+            id=task_model.id,
+            context_id=task_model.context_id,
+            status=cast('dict[str, Any]', task_model.status),
+            artifacts=cast('list[dict[str, Any]]', task_model.artifacts),
+            history=cast('list[dict[str, Any]]', task_model.history),
+            metadata=cast('dict[str, Any]', task_model.task_metadata),
+        )
+        return conversions.to_core_task(legacy_task)
 
     async def save(
         self, task: Task, context: ServerCallContext | None = None