diff --git a/src/a2a/server/agent_execution/active_task.py b/src/a2a/server/agent_execution/active_task.py index c5dcbb9dd..579a7fe4c 100644 --- a/src/a2a/server/agent_execution/active_task.py +++ b/src/a2a/server/agent_execution/active_task.py @@ -57,7 +57,10 @@ ) from a2a.server.tasks.task_manager import TaskManager -from a2a.server.cluster.event_stream import VersionedEvent +from a2a.server.cluster.event_stream import ( + VersionedEvent, + _stream_accepts_context, +) from a2a.server.cluster.task_store import ConcurrentTaskModificationError from a2a.server.events.event_queue_v2 import ( AsyncQueue, @@ -435,10 +438,17 @@ async def _enqueue_to_subscribers( stream = self.active_task._event_stream if stream is not None and isinstance(event, Event): version = self.active_task._task_manager.current_version - await stream.publish( - self.active_task._task_id, - VersionedEvent(event=event, version=version), - ) + versioned_event = VersionedEvent(event=event, version=version) + task_id = self.active_task._task_id + if _stream_accepts_context(stream, 'publish'): + # TODO(ISSUE #1374): remove this check once ServerCallContext is a required parameter on TaskEventStream. + await stream.publish( + task_id, + versioned_event, + context=self.active_task._task_manager._call_context, # ty:ignore[unknown-argument] + ) + else: + await stream.publish(task_id, versioned_event) self.active_task._event_queue_agent.task_done() diff --git a/src/a2a/server/cluster/database_event_stream.py b/src/a2a/server/cluster/database_event_stream.py index d77db1971..924a8e7e9 100644 --- a/src/a2a/server/cluster/database_event_stream.py +++ b/src/a2a/server/cluster/database_event_stream.py @@ -25,8 +25,14 @@ from a2a.server.cluster.event_stream import TaskEventStream, VersionedEvent from a2a.server.cluster.version import TaskVersion +from a2a.server.context import ServerCallContext from a2a.server.events.event_queue import Event -from a2a.server.models import Base, TaskEventModel, create_task_event_model +from a2a.server.models import ( + Base, + TaskEventModel, + create_task_event_model, +) +from a2a.server.owner_resolver import OwnerResolver from a2a.types.a2a_pb2 import StreamResponse @@ -64,6 +70,7 @@ class DatabaseTaskEventStream(TaskEventStream): """`TaskEventStream` backed by polling the shared ``task_events`` table.""" _event_model: type[TaskEventModel] + owner_resolver: OwnerResolver | None def __init__( self, @@ -71,6 +78,7 @@ def __init__( create_table: bool = True, table_name: str = 'task_events', poll_interval_s: float = _DEFAULT_POLL_INTERVAL_S, + owner_resolver: OwnerResolver | None = None, ) -> None: """Initializes the stream over an existing SQLAlchemy AsyncEngine.""" self._engine = engine @@ -78,6 +86,7 @@ def __init__( self._create_table = create_table self._poll_interval_s = poll_interval_s self._initialized = False + self.owner_resolver = owner_resolver self._event_model = ( # ty:ignore[invalid-assignment] TaskEventModel if table_name == 'task_events' @@ -99,24 +108,34 @@ async def _ensure_initialized(self) -> None: if not self._initialized: await self.initialize() - async def publish(self, task_id: str, event: VersionedEvent) -> None: + async def publish( + self, + task_id: str, + event: VersionedEvent, + context: ServerCallContext | None = None, + ) -> None: """No-op: events are persisted transactionally by the task store.""" - del task_id, event + del task_id, event, context async def _seq_at_or_before_version( - self, session: AsyncSession, task_id: str, after: TaskVersion + self, + session: AsyncSession, + task_id: str, + after: TaskVersion, + owner: str | None, ) -> int: """Seq to start polling from so nothing after `after` is missed.""" if after.is_missing: return 0 + conditions = [ + self._event_model.task_id == task_id, + self._event_model.task_version <= _as_int(after), + ] + if owner is not None: + conditions.append(self._event_model.owner == owner) stmt = ( select(self._event_model.seq) - .where( - and_( - self._event_model.task_id == task_id, - self._event_model.task_version <= _as_int(after), - ) - ) + .where(and_(*conditions)) .order_by(self._event_model.seq.desc()) .limit(1) ) @@ -124,30 +143,40 @@ async def _seq_at_or_before_version( return result or 0 async def subscribe( # type: ignore[override] - self, task_id: str, *, after: TaskVersion + self, + task_id: str, + *, + after: TaskVersion, + context: ServerCallContext | None = None, ) -> AsyncGenerator[VersionedEvent, None]: """Polls the log for events of `task_id` newer than `after`.""" await self._ensure_initialized() + owner = ( + self.owner_resolver(context or ServerCallContext()) + if self.owner_resolver is not None + else None + ) async with self._session_maker() as session: cursor = await self._seq_at_or_before_version( - session, task_id, after + session, task_id, after, owner ) while True: try: async with self._session_maker() as session: + conditions = [ + self._event_model.task_id == task_id, + self._event_model.seq > cursor, + ] + if owner is not None: + conditions.append(self._event_model.owner == owner) stmt = ( select( self._event_model.seq, self._event_model.task_version, self._event_model.event_data, ) - .where( - and_( - self._event_model.task_id == task_id, - self._event_model.seq > cursor, - ) - ) + .where(and_(*conditions)) .order_by(self._event_model.seq.asc()) ) rows = (await session.execute(stmt)).all() @@ -176,6 +205,8 @@ async def subscribe( # type: ignore[override] if not rows: await asyncio.sleep(self._poll_interval_s) - async def destroy(self, task_id: str) -> None: + async def destroy( + self, task_id: str, context: ServerCallContext | None = None + ) -> None: """No-op: the append-only log is retained; subscribers stop on their own.""" - del task_id + del task_id, context diff --git a/src/a2a/server/cluster/event_stream.py b/src/a2a/server/cluster/event_stream.py index d65838a39..924e047e2 100644 --- a/src/a2a/server/cluster/event_stream.py +++ b/src/a2a/server/cluster/event_stream.py @@ -1,3 +1,7 @@ +import inspect +import logging +import warnings + from abc import ABC, abstractmethod from collections.abc import AsyncGenerator from dataclasses import dataclass @@ -6,6 +10,9 @@ from a2a.server.events.event_queue import Event +logger = logging.getLogger(__name__) + + @dataclass(frozen=True) class VersionedEvent: """An event together with the task version produced by applying it.""" @@ -30,3 +37,38 @@ def subscribe( @abstractmethod async def destroy(self, task_id: str) -> None: """Releases resources for a task that has reached a terminal state.""" + + +_context_support_cache: dict[tuple[type, str], bool] = {} + + +def _stream_accepts_context(stream: TaskEventStream, method_name: str) -> bool: + """Temporary solution to support ServerCallContext being passed to TaskEventStream. + + Checks if given method of TaskEventStream accepts context parameter. + If not, it logs a warning and returns False. + + TODO(ISSUE #1374): remove this function once ServerCallContext is a required parameter on TaskEventStream. + """ + key = (type(stream), method_name) + cached = _context_support_cache.get(key) + if cached is not None: + return cached + + params = inspect.signature(getattr(stream, method_name)).parameters + accepts = 'context' in params or any( + p.kind is p.VAR_KEYWORD for p in params.values() + ) + _context_support_cache[key] = accepts + + if not accepts: + message = ( + f'{type(stream).__name__}.{method_name} does not accept a ' + "'context' parameter; ownership checks are skipped for this " + 'TaskEventStream implementation. This will become a required ' + 'parameter in a future major version.' + ) + warnings.warn(message, DeprecationWarning, stacklevel=3) + logger.warning(message) + + return accepts diff --git a/src/a2a/server/request_handlers/default_request_handler_v2.py b/src/a2a/server/request_handlers/default_request_handler_v2.py index 9b24c7971..5655fcfdc 100644 --- a/src/a2a/server/request_handlers/default_request_handler_v2.py +++ b/src/a2a/server/request_handlers/default_request_handler_v2.py @@ -19,7 +19,10 @@ from a2a.server.agent_execution.plugin_agent_executor import ( PluginAgentExecutor, ) -from a2a.server.cluster.event_stream import VersionedEvent +from a2a.server.cluster.event_stream import ( + VersionedEvent, + _stream_accepts_context, +) from a2a.server.cluster.task_store import ( ConcurrentTaskModificationError, LegacyTaskStoreAdapter, @@ -375,9 +378,16 @@ async def _write_cancel( ) ) if self._event_stream is not None: - await self._event_stream.publish( - task_id, VersionedEvent(event=event, version=new_version) - ) + versioned_event = VersionedEvent(event=event, version=new_version) + if _stream_accepts_context(self._event_stream, 'publish'): + # TODO(ISSUE #1374): remove this check once ServerCallContext is a required parameter on TaskEventStream. + await self._event_stream.publish( + task_id, + versioned_event, + context=context, # ty:ignore[unknown-argument] + ) + else: + await self._event_stream.publish(task_id, versioned_event) return cancelled def _validate_task_id_match(self, task_id: str, event_task_id: str) -> None: @@ -657,7 +667,7 @@ async def on_subscribe_to_task( # noqa: D102 # Not running here: serve the snapshot and tail the remote stream. async for event in self._subscribe_remote( - task_id, task, snapshot_version, stream + task_id, task, snapshot_version, stream, context ): yield event @@ -667,13 +677,22 @@ async def _subscribe_remote( task: Task, snapshot_version: TaskVersion, stream: TaskEventStream, + context: ServerCallContext, ) -> AsyncGenerator[Event, None]: """Serves a resubscription for a task running on another replica.""" yield task - async with aclosing( - stream.subscribe(task_id, after=snapshot_version) - ) as subscription: + if _stream_accepts_context(stream, 'subscribe'): + # TODO(ISSUE #1374): remove this check once ServerCallContext is a required parameter on TaskEventStream. + subscription_gen = stream.subscribe( + task_id, + after=snapshot_version, + context=context, # ty:ignore[unknown-argument] + ) + else: + subscription_gen = stream.subscribe(task_id, after=snapshot_version) + + async with aclosing(subscription_gen) as subscription: async for versioned in subscription: if not versioned.version.is_after(snapshot_version): continue diff --git a/tests/server/cluster/test_database_event_stream.py b/tests/server/cluster/test_database_event_stream.py index 98a3a6a59..5a0b025be 100644 --- a/tests/server/cluster/test_database_event_stream.py +++ b/tests/server/cluster/test_database_event_stream.py @@ -4,6 +4,7 @@ import contextlib import os import uuid +import warnings from collections.abc import AsyncGenerator @@ -22,8 +23,13 @@ stream_response_to_event, ) from a2a.server.cluster.database_task_store import VersionedDatabaseTaskStore +from a2a.server.cluster.event_stream import ( + TaskEventStream, + _stream_accepts_context, +) from a2a.server.context import ServerCallContext from a2a.server.models import Base +from a2a.server.owner_resolver import resolve_user_scope from a2a.types.a2a_pb2 import ( Message, Part, @@ -181,6 +187,32 @@ def test_as_int_rejects_non_integer_version() -> None: _as_int(TaskVersion('etag')) +class _LegacyStream(TaskEventStream): + """A `TaskEventStream` with no `context` support, for compat testing.""" + + async def publish(self, task_id: str, event: VersionedEvent) -> None: + pass + + async def subscribe( + self, task_id: str, *, after: TaskVersion + ) -> AsyncGenerator[VersionedEvent, None]: + return + yield # unreachable; makes this an async generator function + + async def destroy(self, task_id: str) -> None: + pass + + +def test_stream_accepts_context_warns_once_per_type() -> None: + """A non-updated stream is detected once and cached, not re-warned per call.""" + stream = _LegacyStream() + with pytest.warns(DeprecationWarning, match='does not accept'): + assert _stream_accepts_context(stream, 'publish') is False + with warnings.catch_warnings(): + warnings.simplefilter('error') + assert _stream_accepts_context(stream, 'publish') is False + + @pytest.mark.asyncio @pytest.mark.timeout(10) async def test_initialize_creates_table_and_is_idempotent( @@ -278,6 +310,104 @@ async def consume() -> None: await stream_engine.dispose() +@pytest.mark.asyncio +@pytest.mark.timeout(10) +async def test_owner_context_succeeds_on_all_methods(db_url: str) -> None: + """The owner's context sees its own events and is accepted everywhere.""" + task_id = 'task-owner-ctx' + store_engine = create_async_engine(db_url) + stream_engine = create_async_engine(db_url) + _enforce_sqlite_foreign_keys(store_engine) + _enforce_sqlite_foreign_keys(stream_engine) + try: + store = VersionedDatabaseTaskStore( + engine=store_engine, create_table=False + ) + stream = DatabaseTaskEventStream( + engine=stream_engine, + create_table=False, + owner_resolver=resolve_user_scope, + ) + await store.initialize() + await stream.initialize() + + v1 = await store.save( + create_task(task_id), + event=None, + prev=None, + prev_version=TaskVersion.MISSING, + context=TEST_CONTEXT, + ) + v2 = await store.save( + create_task(task_id), + event=status_event(task_id, state=TaskState.TASK_STATE_WORKING), + prev=None, + prev_version=v1, + context=TEST_CONTEXT, + ) + + agen = stream.subscribe(task_id, after=v1, context=TEST_CONTEXT) + received = await asyncio.wait_for(agen.__anext__(), timeout=5) + assert received.event.status.state == TaskState.TASK_STATE_WORKING + await agen.aclose() + + await stream.publish( + task_id, + VersionedEvent(create_task(task_id), v2), + context=TEST_CONTEXT, + ) + await stream.destroy(task_id, context=TEST_CONTEXT) + finally: + await store_engine.dispose() + await stream_engine.dispose() + + +@pytest.mark.asyncio +@pytest.mark.timeout(10) +async def test_other_owner_context_does_not_receive_events(db_url: str) -> None: + """A different owner does not see events belonging to another owner.""" + task_id = 'task-other-owner' + store_engine = create_async_engine(db_url) + stream_engine = create_async_engine(db_url) + _enforce_sqlite_foreign_keys(store_engine) + _enforce_sqlite_foreign_keys(stream_engine) + try: + store = VersionedDatabaseTaskStore( + engine=store_engine, create_table=False + ) + stream = DatabaseTaskEventStream( + engine=stream_engine, + create_table=False, + owner_resolver=resolve_user_scope, + ) + await store.initialize() + await stream.initialize() + + v1 = await store.save( + create_task(task_id), + event=None, + prev=None, + prev_version=TaskVersion.MISSING, + context=TEST_CONTEXT, + ) + await store.save( + create_task(task_id), + event=status_event(task_id, state=TaskState.TASK_STATE_WORKING), + prev=None, + prev_version=v1, + context=TEST_CONTEXT, + ) + other = ServerCallContext(user=SampleUser('someone_else')) + + agen = stream.subscribe(task_id, after=v1, context=other) + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait_for(agen.__anext__(), timeout=0.15) + await agen.aclose() + finally: + await store_engine.dispose() + await stream_engine.dispose() + + @pytest.mark.asyncio @pytest.mark.timeout(10) async def test_publish_is_noop(db_url: str) -> None: