Merge pull request #81 from GetStream/hotfix/agent-finish-events-merge

hotfix: events, agent finish

Author: yarikdevcom

agents-core/vision_agents/core/agents/agents.py

@@ -12,7 +12,7 @@
 from opentelemetry.trace import Tracer
 
 from ..edge import sfu_events
-from ..edge.events import AudioReceivedEvent, TrackAddedEvent
+from ..edge.events import AudioReceivedEvent, TrackAddedEvent, CallEndedEvent
 from ..edge.types import Connection, Participant, PcmData, TrackType, User
 from ..events.manager import EventManager
 from ..llm.events import (
@@ -267,6 +267,28 @@ async def join(self, call: Call) -> "AgentSessionContextManager":
 
             return AgentSessionContextManager(self, self._connection)
 
+    async def finish(self):
+        """Wait for the call to end gracefully.
+        Subscribes to the edge transport's `call_ended` event and awaits it. If
+        no connection is active, returns immediately.
+        """
+        # If connection is None or already closed, return immediately
+        if not self._connection:
+            logging.info("🔚 Agent connection already closed, finishing immediately")
+            return
+
+        @self.edge.events.subscribe
+        async def on_ended(event: CallEndedEvent):
+            self._is_running = False
+
+        while self._is_running:
+            try:
+                await asyncio.sleep(0.0001)
+            except asyncio.CancelledError:
+                self._is_running = False
+
+        await asyncio.shield(self.close())
+
     async def close(self):
         """Clean up all connections and resources.
 

agents-core/vision_agents/core/events/manager.py

@@ -1,4 +1,5 @@
 import asyncio
+import uuid
 import collections
 import logging
 import types
@@ -141,6 +142,7 @@ def __init__(self, ignore_unknown_events: bool = True):
         self._processing_task: Optional[asyncio.Task[Any]] = None
         self._shutdown = False
         self._silent_events: set[type] = set()
+        self._handler_tasks: Dict[uuid.UUID, asyncio.Task[Any]] = {}
 
         self.register(ExceptionEvent)
         self.register(HealthCheckEvent)
@@ -182,6 +184,8 @@ def register(self, event_class, ignore_not_compatible=False):
             #    raise KeyError(f"{event_class.type} is already registered.")
             self._events[event_class.type] = event_class
             logger.info(f"Registered new event {event_class} - {event_class.type}")
+        elif event_class.__name__.endswith('BaseEvent'):
+            return
         elif not ignore_not_compatible:
             raise ValueError(f"Provide valid class that ends on '*Event' and 'type' attribute: {event_class}")
         else:
@@ -458,6 +462,9 @@ async def wait(self, timeout: float = 10.0):
         start_time = asyncio.get_event_loop().time()
         while self._queue and (asyncio.get_event_loop().time() - start_time) < timeout:
             await asyncio.sleep(0.01)
+
+        if self._handler_tasks:
+            await asyncio.wait(list(self._handler_tasks.values()))
 
     def _start_processing_task(self):
         """Start the background event processing task."""
@@ -488,18 +495,27 @@ async def _process_events_loop(self):
             elif cancelled_exc:
                 raise cancelled_exc
             else:
+                cleanup_ids = set(task_id for task_id, task in self._handler_tasks.items() if task.done())
+                for task_id in cleanup_ids:
+                    self._handler_tasks.pop(task_id)
                 await asyncio.sleep(0.0001)
 
+    async def _run_handler(self, handler, event):
+        try:
+            return await handler(event)
+        except Exception as exc:
+            self._queue.appendleft(ExceptionEvent(exc, handler))  # type: ignore[arg-type]
+            module_name = getattr(handler, '__module__', 'unknown')
+            logger.exception(f"Error calling handler {handler.__name__} from {module_name} for event {event.type}")
+
     async def _process_single_event(self, event):
         """Process a single event."""
         for handler in self._handlers.get(event.type, []):
-            try:
-                module_name = getattr(handler, '__module__', 'unknown')
-                if event.type not in self._silent_events:
-                    logger.info(f"Called handler {handler.__name__} from {module_name} for event {event.type}")
-                await handler(event)
-            except Exception as exc:
-                self._queue.appendleft(ExceptionEvent(exc, handler))  # type: ignore[arg-type]
-                module_name = getattr(handler, '__module__', 'unknown')
-                logger.exception(f"Error calling handler {handler.__name__} from {module_name} for event {event.type}")
+            module_name = getattr(handler, '__module__', 'unknown')
+            if event.type not in self._silent_events:
+                logger.info(f"Called handler {handler.__name__} from {module_name} for event {event.type}")
+
+            loop = asyncio.get_running_loop()
+            handler_task = loop.create_task(self._run_handler(handler, event))
+            self._handler_tasks[uuid.uuid4()] = handler_task
 

examples/01_simple_agent_example/pyproject.toml

@@ -25,5 +25,6 @@ dependencies = [
 "vision-agents-plugins-anthropic" = {path = "../../plugins/anthropic", editable=true}
 "vision-agents-plugins-getstream" = {path = "../../plugins/getstream", editable=true}
 "vision-agents-plugins-openai" = {path = "../../plugins/openai", editable=true}
+"vision-agents-plugins-smart-turn" = {path = "../../plugins/smart_turn", editable=true}
 
 "vision-agents" = {path = "../../agents-core", editable=true}

plugins/getstream/vision_agents/plugins/getstream/stream_edge_transport.py

@@ -102,7 +102,7 @@ async def _on_track_published(self, event: sfu_events.TrackPublishedEvent):
         # SFU might send TrackPublishedEvent before WebRTC processes track_added
         track_id = None
         timeout = 10.0
-        poll_interval = 0.01  # 50ms
+        poll_interval = 0.01  # 10ms
         elapsed = 0.0
 
         while elapsed < timeout:
@@ -125,7 +125,7 @@ async def _on_track_published(self, event: sfu_events.TrackPublishedEvent):
         if track_id:
             # Store with correct type from SFU
             self._track_map[track_key] = {"track_id": track_id, "published": True}
-            self.logger.info(f"Track published: {track_type_name} from {user_id}, track_id: {track_id} (waited {elapsed:.2f}s)")
+            self.logger.info(f"Trackmap published: {track_type_name} from {user_id}, track_id: {track_id} (waited {elapsed:.2f}s)")
 
             # NOW spawn TrackAddedEvent with correct type
             self.events.send(events.TrackAddedEvent(
@@ -139,6 +139,9 @@ async def _on_track_published(self, event: sfu_events.TrackPublishedEvent):
             raise TimeoutError(
                 f"Timeout waiting for pending track: {track_type_name} ({expected_kind}) from user {user_id}, "
                 f"session {session_id}. Waited {timeout}s but WebRTC track_added with matching kind was never received."
+                f"Pending tracks: {self._pending_tracks}\n"
+                f"Key: {track_key}\n"
+                f"Track map: {self._track_map}\n"
             )
 
     async def _on_track_removed(self, event: sfu_events.ParticipantLeftEvent | sfu_events.TrackUnpublishedEvent):
@@ -230,16 +233,6 @@ async def join(self, agent: "Agent", call: Call) -> StreamConnection:
 
         self._connection = connection
 
-        original_on_subscriber_offer = self._connection._on_subscriber_offer  # type: ignore[attr-defined]
-
-        async def _safe_on_subscriber_offer(event):
-            if self._connection.subscriber_pc is None:  # type: ignore[attr-defined]
-                self.logger.debug("Ignoring subscriber offer after subscriber_pc teardown")
-                return
-            await original_on_subscriber_offer(event)
-
-        self._connection._on_subscriber_offer = _safe_on_subscriber_offer  # type: ignore[attr-defined]
-
         @self._connection.on("audio")
         async def on_audio_received(pcm: PcmData, participant: Participant):
             self.events.send(events.AudioReceivedEvent(