Revert not needed

Author: ishymko

src/a2a/server/models.py

@@ -141,10 +141,6 @@ class Base(DeclarativeBase):
 
 
 # TaskMixin that can be used with any table name
-_task_model_cache: dict[tuple[str, type], type] = {}
-_push_notification_config_model_cache: dict[tuple[str, type], type] = {}
-
-
 class TaskMixin:
     """Mixin providing standard task columns with proper type handling."""
 
@@ -220,9 +216,6 @@ def create_task_model(
 
             TaskModel = create_task_model('tasks', MyBase)
     """
-    cache_key = (table_name, base)
-    if cache_key in _task_model_cache:
-        return _task_model_cache[cache_key]
 
     class TaskModel(TaskMixin, base):  # type: ignore
         __tablename__ = table_name
@@ -239,10 +232,16 @@ def __repr__(self) -> str:
     TaskModel.__name__ = f'TaskModel_{table_name}'
     TaskModel.__qualname__ = f'TaskModel_{table_name}'
 
-    _task_model_cache[cache_key] = TaskModel
     return TaskModel
 
 
+# Default TaskModel for backward compatibility
+class TaskModel(TaskMixin, Base):
+    """Default task model with standard table name."""
+
+    __tablename__ = 'tasks'
+
+
 # PushNotificationConfigMixin that can be used with any table name
 class PushNotificationConfigMixin:
     """Mixin providing standard push notification config columns."""
@@ -266,9 +265,6 @@ def create_push_notification_config_model(
     base: type[DeclarativeBase] = Base,
 ) -> type:
     """Create a PushNotificationConfigModel class with a configurable table name."""
-    cache_key = (table_name, base)
-    if cache_key in _push_notification_config_model_cache:
-        return _push_notification_config_model_cache[cache_key]
 
     class PushNotificationConfigModel(PushNotificationConfigMixin, base):  # type: ignore
         __tablename__ = table_name
@@ -277,19 +273,22 @@ class PushNotificationConfigModel(PushNotificationConfigMixin, base):  # type: i
         def __repr__(self) -> str:
             """Return a string representation of the push notification config."""
             return (
-                f'<PushNotificationConfigModel[{table_name}](task_id="{self.task_id}", '
-                f'config_id="{self.config_id}")>'
+                f'<PushNotificationConfigModel[{table_name}]('
+                f'task_id="{self.task_id}", config_id="{self.config_id}")>'
             )
 
-    # Set a dynamic name for better debugging
     PushNotificationConfigModel.__name__ = (
         f'PushNotificationConfigModel_{table_name}'
     )
     PushNotificationConfigModel.__qualname__ = (
         f'PushNotificationConfigModel_{table_name}'
     )
 
-    _push_notification_config_model_cache[cache_key] = (
-        PushNotificationConfigModel
-    )
     return PushNotificationConfigModel
+
+
+# Default PushNotificationConfigModel for backward compatibility
+class PushNotificationConfigModel(PushNotificationConfigMixin, Base):
+    """Default push notification config model with standard table name."""
+
+    __tablename__ = 'push_notification_configs'

src/a2a/server/tasks/database_push_notification_config_store.py

@@ -2,7 +2,7 @@
 import json
 import logging
 
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 from google.protobuf.json_format import MessageToJson, Parse
 
@@ -30,6 +30,7 @@
 from a2a.server.context import ServerCallContext
 from a2a.server.models import (
     Base,
+    PushNotificationConfigModel,
     create_push_notification_config_model,
 )
 from a2a.server.owner_resolver import OwnerResolver, resolve_user_scope
@@ -56,7 +57,7 @@ class DatabasePushNotificationConfigStore(PushNotificationConfigStore):
     async_session_maker: async_sessionmaker[AsyncSession]
     create_table: bool
     _initialized: bool
-    config_model: Any
+    config_model: type[PushNotificationConfigModel]
     _fernet: 'Fernet | None'
     owner_resolver: OwnerResolver
 
@@ -90,7 +91,11 @@ def __init__(
         self.create_table = create_table
         self._initialized = False
         self.owner_resolver = owner_resolver
-        self.config_model = create_push_notification_config_model(table_name)
+        self.config_model = (
+            PushNotificationConfigModel
+            if table_name == 'push_notification_configs'
+            else create_push_notification_config_model(table_name)
+        )
         self._fernet = None
 
         if encryption_key:
@@ -141,7 +146,7 @@ async def _ensure_initialized(self) -> None:
 
     def _to_orm(
         self, task_id: str, config: PushNotificationConfig, owner: str
-    ) -> Any:
+    ) -> PushNotificationConfigModel:
         """Maps a PushNotificationConfig proto to a SQLAlchemy model instance.
 
         The config data is serialized to JSON bytes, and encrypted if a key is configured.
@@ -160,7 +165,9 @@ def _to_orm(
             config_data=data_to_store,
         )
 
-    def _from_orm(self, model_instance: Any) -> PushNotificationConfig:
+    def _from_orm(
+        self, model_instance: PushNotificationConfigModel
+    ) -> PushNotificationConfig:
         """Maps a SQLAlchemy model instance to a PushNotificationConfig proto.
 
         Handles decryption if a key is configured, with a fallback to plain JSON.

src/a2a/server/tasks/database_task_store.py

@@ -34,7 +34,7 @@
 from google.protobuf.json_format import MessageToDict
 
 from a2a.server.context import ServerCallContext
-from a2a.server.models import Base, create_task_model
+from a2a.server.models import Base, TaskModel, create_task_model
 from a2a.server.owner_resolver import OwnerResolver, resolve_user_scope
 from a2a.server.tasks.task_store import TaskStore
 from a2a.types import a2a_pb2
@@ -56,7 +56,7 @@ class DatabaseTaskStore(TaskStore):
     async_session_maker: async_sessionmaker[AsyncSession]
     create_table: bool
     _initialized: bool
-    task_model: type[Any]
+    task_model: type[TaskModel]
     owner_resolver: OwnerResolver
 
     def __init__(
@@ -86,7 +86,11 @@ def __init__(
         self._initialized = False
         self.owner_resolver = owner_resolver
 
-        self.task_model = create_task_model(table_name)
+        self.task_model = (
+            TaskModel
+            if table_name == 'tasks'
+            else create_task_model(table_name)
+        )
 
     async def initialize(self) -> None:
         """Initialize the database and create the table if needed."""
@@ -111,7 +115,7 @@ async def _ensure_initialized(self) -> None:
         if not self._initialized:
             await self.initialize()
 
-    def _to_orm(self, task: Task, owner: str) -> Any:
+    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
@@ -133,7 +137,7 @@ def _to_orm(self, task: Task, owner: str) -> Any:
             ),
         )
 
-    def _from_orm(self, task_model: Any) -> Task:
+    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

tests/server/tasks/test_database_push_notification_config_store.py

@@ -16,7 +16,6 @@
 
 import pytest_asyncio
 
-from typing import TYPE_CHECKING, Any
 from _pytest.mark.structures import ParameterSet
 
 # Now safe to import SQLAlchemy-dependent modules
@@ -33,10 +32,8 @@
 
 from a2a.server.models import (
     Base,
-    create_push_notification_config_model,
-)
-
-PushNotificationConfigModel: Any = create_push_notification_config_model()
+    PushNotificationConfigModel,
+)  # Important: To get Base.metadata
 from a2a.server.tasks import DatabasePushNotificationConfigStore
 from a2a.types.a2a_pb2 import (
     PushNotificationConfig,

tests/server/tasks/test_database_task_store.py

@@ -19,10 +19,7 @@
 
 from google.protobuf.json_format import MessageToDict
 
-from a2a.server.models import Base, create_task_model
-from typing import Any
-
-TaskModel: Any = create_task_model()
+from a2a.server.models import Base, TaskModel  # Important: To get Base.metadata
 from a2a.server.tasks.database_task_store import DatabaseTaskStore
 from a2a.types.a2a_pb2 import (
     Artifact,