CopyingTaskStoreAdapter.

Author: bartek-w

src/a2a/server/tasks/copying_task_store.py

@@ -0,0 +1,61 @@
+from __future__ import annotations
+
+import logging
+
+from typing import TYPE_CHECKING
+
+
+if TYPE_CHECKING:
+    from a2a.server.context import ServerCallContext
+from a2a.server.tasks.task_store import TaskStore
+from a2a.types.a2a_pb2 import ListTasksRequest, ListTasksResponse, Task
+
+
+logger = logging.getLogger(__name__)
+
+
+class CopyingTaskStoreAdapter(TaskStore):
+    """An adapter that ensures deep copies of tasks are passed to and returned from the underlying TaskStore.
+
+    This prevents accidental shared mutable state bugs where code modifies a Task object
+    retrieved from the store without explicitly saving it, which hides missing save calls.
+    """
+
+    def __init__(self, underlying_store: TaskStore):
+        self._store = underlying_store
+
+    async def save(
+        self, task: Task, context: ServerCallContext | None = None
+    ) -> None:
+        """Saves a copy of the task to the underlying store."""
+        task_copy = Task()
+        task_copy.CopyFrom(task)
+        await self._store.save(task_copy, context)
+
+    async def get(
+        self, task_id: str, context: ServerCallContext | None = None
+    ) -> Task | None:
+        """Retrieves a task from the underlying store and returns a copy."""
+        task = await self._store.get(task_id, context)
+        if task is None:
+            return None
+        task_copy = Task()
+        task_copy.CopyFrom(task)
+        return task_copy
+
+    async def list(
+        self,
+        params: ListTasksRequest,
+        context: ServerCallContext | None = None,
+    ) -> ListTasksResponse:
+        """Retrieves a list of tasks from the underlying store and returns a copy."""
+        response = await self._store.list(params, context)
+        response_copy = ListTasksResponse()
+        response_copy.CopyFrom(response)
+        return response_copy
+
+    async def delete(
+        self, task_id: str, context: ServerCallContext | None = None
+    ) -> None:
+        """Deletes a task from the underlying store."""
+        await self._store.delete(task_id, context)

src/a2a/server/tasks/inmemory_task_store.py

@@ -3,6 +3,7 @@
 
 from a2a.server.context import ServerCallContext
 from a2a.server.owner_resolver import OwnerResolver, resolve_user_scope
+from a2a.server.tasks.copying_task_store import CopyingTaskStoreAdapter
 from a2a.server.tasks.task_store import TaskStore
 from a2a.types import a2a_pb2
 from a2a.types.a2a_pb2 import Task
@@ -14,8 +15,8 @@
 logger = logging.getLogger(__name__)
 
 
-class InMemoryTaskStore(TaskStore):
-    """In-memory implementation of TaskStore.
+class _InMemoryTaskStoreImpl(TaskStore):
+    """Internal In-memory implementation of TaskStore.
 
     Stores task objects in a nested dictionary in memory, keyed by owner then task_id.
     Task data is lost when the server process stops.
@@ -25,8 +26,8 @@ def __init__(
         self,
         owner_resolver: OwnerResolver = resolve_user_scope,
     ) -> None:
-        """Initializes the InMemoryTaskStore."""
-        logger.debug('Initializing InMemoryTaskStore')
+        """Initializes the internal _InMemoryTaskStoreImpl."""
+        logger.debug('Initializing _InMemoryTaskStoreImpl')
         self.tasks: dict[str, dict[str, Task]] = {}
         self.lock = asyncio.Lock()
         self.owner_resolver = owner_resolver
@@ -183,3 +184,55 @@ async def delete(
             if not owner_tasks:
                 del self.tasks[owner]
                 logger.debug('Removed empty owner %s from store.', owner)
+
+
+class InMemoryTaskStore(TaskStore):
+    """In-memory implementation of TaskStore.
+
+    Can optionally use CopyingTaskStoreAdapter to wrap the internal dictionary-based
+    implementation, preventing shared mutable state issues by always returning and
+    storing deep copies.
+    """
+
+    def __init__(
+        self,
+        owner_resolver: OwnerResolver = resolve_user_scope,
+        use_copying: bool = True,
+    ) -> None:
+        """Initializes the InMemoryTaskStore.
+
+        Args:
+            owner_resolver: Resolver for task owners.
+            use_copying: If True, the store will return and save deep copies of tasks.
+              Copying behavior is consistent with database task stores.
+        """
+        self._impl = _InMemoryTaskStoreImpl(owner_resolver=owner_resolver)
+        self._store: TaskStore = (
+            CopyingTaskStoreAdapter(self._impl) if use_copying else self._impl
+        )
+
+    async def save(
+        self, task: Task, context: ServerCallContext | None = None
+    ) -> None:
+        """Saves or updates a task in the store."""
+        await self._store.save(task, context)
+
+    async def get(
+        self, task_id: str, context: ServerCallContext | None = None
+    ) -> Task | None:
+        """Retrieves a task from the store by ID."""
+        return await self._store.get(task_id, context)
+
+    async def list(
+        self,
+        params: a2a_pb2.ListTasksRequest,
+        context: ServerCallContext | None = None,
+    ) -> a2a_pb2.ListTasksResponse:
+        """Retrieves a list of tasks from the store."""
+        return await self._store.list(params, context)
+
+    async def delete(
+        self, task_id: str, context: ServerCallContext | None = None
+    ) -> None:
+        """Deletes a task from the store by ID."""
+        await self._store.delete(task_id, context)

tests/integration/test_copying_observability.py

@@ -0,0 +1,175 @@
+import httpx
+import pytest
+from typing import NamedTuple
+
+from a2a.client.client import Client, ClientConfig
+from a2a.client.client_factory import ClientFactory
+from a2a.server.agent_execution import AgentExecutor, RequestContext
+from a2a.server.apps import A2AFastAPIApplication
+from a2a.server.events import EventQueue
+from a2a.server.events.in_memory_queue_manager import InMemoryQueueManager
+from a2a.server.request_handlers import DefaultRequestHandler
+from a2a.server.tasks import TaskUpdater
+from a2a.server.tasks.inmemory_task_store import InMemoryTaskStore
+from a2a.types import (
+    AgentCapabilities,
+    AgentCard,
+    AgentInterface,
+    Artifact,
+    GetTaskRequest,
+    Message,
+    Part,
+    Role,
+    SendMessageRequest,
+    TaskState,
+)
+from a2a.utils import TransportProtocol
+
+
+class MockMutatingAgentExecutor(AgentExecutor):
+    async def execute(self, context: RequestContext, event_queue: EventQueue):
+        assert context.task_id is not None
+        assert context.context_id is not None
+        task_updater = TaskUpdater(
+            event_queue,
+            context.task_id,
+            context.context_id,
+        )
+
+        user_input = context.get_user_input()
+
+        if user_input == 'Init task':
+            # Explicitly save status change to ensure task exists with some state
+            await task_updater.update_status(
+                TaskState.TASK_STATE_WORKING,
+                message=task_updater.new_agent_message(
+                    [Part(text='task working')]
+                ),
+            )
+        else:
+            # Mutate the task WITHOUT saving it properly
+            context.current_task.artifacts.append(
+                Artifact(
+                    name='leaked-artifact',
+                    parts=[Part(text='leaked artifact')],
+                )
+            )
+
+    async def cancel(self, context: RequestContext, event_queue: EventQueue):
+        raise NotImplementedError('Cancellation is not supported')
+
+
+@pytest.fixture
+def agent_card() -> AgentCard:
+    return AgentCard(
+        name='Mutating Agent',
+        description='Real in-memory integration testing.',
+        version='1.0.0',
+        capabilities=AgentCapabilities(
+            streaming=True, push_notifications=False
+        ),
+        skills=[],
+        default_input_modes=['text/plain'],
+        default_output_modes=['text/plain'],
+        supported_interfaces=[
+            AgentInterface(
+                protocol_binding=TransportProtocol.JSONRPC,
+                url='http://testserver',
+            ),
+        ],
+    )
+
+
+class ClientSetup(NamedTuple):
+    client: Client
+    task_store: InMemoryTaskStore
+    use_copying: bool
+
+
+def setup_client(agent_card: AgentCard, use_copying: bool) -> ClientSetup:
+    task_store = InMemoryTaskStore(use_copying=use_copying)
+    handler = DefaultRequestHandler(
+        agent_executor=MockMutatingAgentExecutor(),
+        task_store=task_store,
+        queue_manager=InMemoryQueueManager(),
+    )
+    app_builder = A2AFastAPIApplication(
+        agent_card, handler, extended_agent_card=agent_card
+    )
+    app = app_builder.build()
+    httpx_client = httpx.AsyncClient(
+        transport=httpx.ASGITransport(app=app), base_url='http://testserver'
+    )
+    factory = ClientFactory(
+        config=ClientConfig(
+            httpx_client=httpx_client,
+            supported_protocol_bindings=[TransportProtocol.JSONRPC],
+        )
+    )
+    client = factory.create(agent_card)
+    return ClientSetup(
+        client=client,
+        task_store=task_store,
+        use_copying=use_copying,
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize('use_copying', [True, False])
+async def test_mutation_observability(agent_card: AgentCard, use_copying: bool):
+    """Tests that task mutations are observable when copying is disabled.
+
+    When copying is disabled, the agent mutates the task in-place and the
+    changes are observable by the client. When copying is enabled, the agent
+    mutates a copy of the task and the changes are not observable by the client.
+
+    It is ok to remove the `use_copying` parameter from the system in the future
+    to make InMemoryTaskStore consistent with other task stores.
+    """
+    client_setup = setup_client(agent_card, use_copying)
+    client = client_setup.client
+
+    # 1. First message to create the task
+    message_to_send = Message(
+        role=Role.ROLE_USER,
+        message_id='msg-mut-init',
+        parts=[Part(text='Init task')],
+    )
+
+    events = [
+        event
+        async for event in client.send_message(
+            request=SendMessageRequest(message=message_to_send)
+        )
+    ]
+
+    task = events[-1][1]
+    assert task is not None
+    task_id = task.id
+
+    # 2. Second message to mutate it
+    message_to_send_2 = Message(
+        role=Role.ROLE_USER,
+        message_id='msg-mut-do',
+        task_id=task_id,
+        parts=[Part(text='Update task without saving it')],
+    )
+
+    _ = [
+        event
+        async for event in client.send_message(
+            request=SendMessageRequest(message=message_to_send_2)
+        )
+    ]
+
+    # 3. Get task via client
+    retrieved_task = await client.get_task(request=GetTaskRequest(id=task_id))
+
+    # 4. Assert behavior based on `use_copying`
+    if use_copying:
+        # The un-saved artifact IS NOT leaked to the client
+        assert len(retrieved_task.artifacts) == 0
+    else:
+        # The un-saved artifact IS leaked to the client
+        assert len(retrieved_task.artifacts) == 1
+        assert retrieved_task.artifacts[0].name == 'leaked-artifact'