feat(cluster): Introduce ServerCallContext to DatabaseTaskEventStream

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()
 

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,20 +70,23 @@ class DatabaseTaskEventStream(TaskEventStream):
     """`TaskEventStream` backed by polling the shared ``task_events`` table."""
 
     _event_model: type[TaskEventModel]
+    owner_resolver: OwnerResolver | None
 
     def __init__(
         self,
         engine: AsyncEngine,
         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
         self._session_maker = async_sessionmaker(engine, expire_on_commit=False)
         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,55 +108,75 @@ 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)
         )
         result = (await session.execute(stmt)).scalar_one_or_none()
         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

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

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