Merge remote-tracking branch 'upstream/1.0-dev' into guglielmoc/refactor_rest_server

Author: guglielmo-san

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_client_server_integration.py

@@ -1,5 +1,4 @@
 import asyncio
-
 from collections.abc import AsyncGenerator
 from typing import Any, NamedTuple
 from unittest.mock import ANY, AsyncMock, patch
@@ -8,22 +7,25 @@
 import httpx
 import pytest
 import pytest_asyncio
-
 from cryptography.hazmat.primitives.asymmetric import ec
 from google.protobuf.json_format import MessageToDict
 from google.protobuf.timestamp_pb2 import Timestamp
 
 from a2a.client import Client, ClientConfig
 from a2a.client.base_client import BaseClient
 from a2a.client.card_resolver import A2ACardResolver
-from a2a.client.client_factory import ClientFactory
 from a2a.client.client import ClientCallContext
+from a2a.client.client_factory import ClientFactory
 from a2a.client.service_parameters import (
     ServiceParametersFactory,
     with_a2a_extensions,
 )
 from a2a.client.transports import JsonRpcTransport, RestTransport
 from starlette.applications import Starlette
+
+# Compat v0.3 imports for dedicated tests
+from a2a.compat.v0_3 import a2a_v0_3_pb2, a2a_v0_3_pb2_grpc
+from a2a.compat.v0_3.grpc_handler import CompatGrpcHandler
 from a2a.server.routes import (
     create_agent_card_routes,
     create_jsonrpc_routes,
@@ -55,12 +57,10 @@
     TaskStatus,
     TaskStatusUpdateEvent,
 )
-from a2a.utils.constants import (
-    TransportProtocol,
-)
+from a2a.utils.constants import TransportProtocol
 from a2a.utils.errors import (
-    ExtendedAgentCardNotConfiguredError,
     ContentTypeNotSupportedError,
+    ExtendedAgentCardNotConfiguredError,
     ExtensionSupportRequiredError,
     InternalError,
     InvalidAgentResponseError,
@@ -78,11 +78,6 @@
     create_signature_verifier,
 )
 
-# Compat v0.3 imports for dedicated tests
-from a2a.compat.v0_3 import a2a_v0_3_pb2, a2a_v0_3_pb2_grpc
-from a2a.compat.v0_3.grpc_handler import CompatGrpcHandler
-
-
 # --- Test Constants ---
 
 TASK_FROM_STREAM = Task(
@@ -374,9 +369,9 @@ def grpc_03_setup(
 ) -> TransportSetup:
     """Sets up the CompatGrpcTransport and in-process 0.3 server."""
     server_address, handler = grpc_03_server_and_handler
-    from a2a.compat.v0_3.grpc_transport import CompatGrpcTransport
     from a2a.client.base_client import BaseClient
     from a2a.client.client import ClientConfig
+    from a2a.compat.v0_3.grpc_transport import CompatGrpcTransport
 
     channel = grpc.aio.insecure_channel(server_address)
     transport = CompatGrpcTransport(channel=channel, agent_card=agent_card)
@@ -932,6 +927,73 @@ async def test_client_handles_a2a_errors(transport_setups, error_cls) -> None:
     await client.close()
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    'error_cls',
+    [
+        TaskNotFoundError,
+        TaskNotCancelableError,
+        PushNotificationNotSupportedError,
+        UnsupportedOperationError,
+        ContentTypeNotSupportedError,
+        InvalidAgentResponseError,
+        ExtendedAgentCardNotConfiguredError,
+        ExtensionSupportRequiredError,
+        VersionNotSupportedError,
+    ],
+)
+@pytest.mark.parametrize(
+    'handler_attr, client_method, request_params',
+    [
+        pytest.param(
+            'on_message_send_stream',
+            'send_message',
+            SendMessageRequest(
+                message=Message(
+                    role=Role.ROLE_USER,
+                    message_id='msg-integration-test',
+                    parts=[Part(text='Hello, integration test!')],
+                )
+            ),
+            id='stream',
+        ),
+        pytest.param(
+            'on_subscribe_to_task',
+            'subscribe',
+            SubscribeToTaskRequest(id='some-id'),
+            id='subscribe',
+        ),
+    ],
+)
+async def test_client_handles_a2a_errors_streaming(
+    transport_setups, error_cls, handler_attr, client_method, request_params
+) -> None:
+    """Integration test to verify error propagation from streaming handlers to client.
+
+    The handler raises an A2AError before yielding any events. All transports
+    must propagate this as the exact error_cls, not wrapped in an ExceptionGroup
+    or converted to a generic client error.
+    """
+    client = transport_setups.client
+    handler = transport_setups.handler
+
+    async def mock_generator(*args, **kwargs):
+        raise error_cls('Test error message')
+        yield
+
+    getattr(handler, handler_attr).side_effect = mock_generator
+
+    with pytest.raises(error_cls) as exc_info:
+        async for _ in getattr(client, client_method)(request=request_params):
+            pass
+
+    assert 'Test error message' in str(exc_info.value)
+
+    getattr(handler, handler_attr).side_effect = None
+
+    await client.close()
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     'request_kwargs, expected_error_code',