Do not subscribe to WS events inside Agent.

Author: dangusev

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

@@ -538,15 +538,6 @@ async def join(
             with self.span("edge.publish_tracks"):
                 await self.edge.publish_tracks(audio_track, video_track)
 
-        connection._connection._coordinator_ws_client.on_wildcard(
-            "*",
-            lambda event_name, event: self.events.send(event),
-        )
-
-        connection._connection._ws_client.on_wildcard(
-            "*",
-            lambda event_name, event: self.events.send(event),
-        )
 
         from .agent_session import AgentSessionContextManager
 

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

@@ -43,9 +43,9 @@ def __init__(self, connection: ConnectionManager):
     def participants(self) -> ParticipantsState:
         return self._connection.participants_state
 
-    async def close(self):
+    async def close(self, timeout: float = 2.0):
         try:
-            await asyncio.wait_for(self._connection.leave(), timeout=2.0)
+            await asyncio.wait_for(self._connection.leave(), timeout=timeout)
         except asyncio.TimeoutError:
             logger.warning("Connection leave timed out during close")
         except RuntimeError as e:
@@ -83,11 +83,19 @@ def __init__(self, **kwargs):
         # track_id -> (user_id, session_id, webrtc_type_string)
         self._pending_tracks: dict = {}
 
+        self._real_connection: Optional[ConnectionManager] = None
+
         # Register event handlers
         self.events.subscribe(self._on_track_published)
         self.events.subscribe(self._on_track_removed)
         self.events.subscribe(self._on_call_ended)
 
+    @property
+    def _connection(self) -> ConnectionManager:
+        if self._real_connection is None:
+            raise ValueError("Edge connection is not set")
+        return self._real_connection
+
     def _get_webrtc_kind(self, track_type_int: int) -> str:
         """Get the expected WebRTC kind (audio/video) for a SFU track type."""
         # Map SFU track types to WebRTC kinds
@@ -105,96 +113,6 @@ def _get_webrtc_kind(self, track_type_int: int) -> str:
             # Default to video for unknown types
             return "video"
 
-    async def _subscribe_to_existing_tracks(
-        self, connection: ConnectionManager
-    ) -> None:
-        """Subscribe to tracks from participants who joined before the agent."""
-        from vision_agents.core.edge.sfu_events import Participant as SfuParticipant
-
-        participants = connection.participants_state.get_participants()
-        subscription_manager = connection._subscription_manager
-        tracks_to_subscribe = []
-
-        for participant in participants:
-            if participant.user_id == self.agent_user_id:
-                continue
-
-            for track_type_int in participant.published_tracks:
-                # Create a mock event for the subscription manager
-                class MockTrackPublishedEvent:
-                    def __init__(self, p, track_type):
-                        self.user_id = p.user_id
-                        self.session_id = p.session_id
-                        self.type = track_type
-                        self.participant = p
-
-                mock_event = MockTrackPublishedEvent(participant, track_type_int)
-
-                try:
-                    await subscription_manager.handle_track_published(mock_event)
-                    tracks_to_subscribe.append((participant, track_type_int))
-                except Exception as e:
-                    logger.error(f"Failed to subscribe to existing track: {e}")
-
-        # Poll for WebRTC tracks to arrive after subscription
-        for participant, track_type_int in tracks_to_subscribe:
-            expected_kind = self._get_webrtc_kind(track_type_int)
-            track_key = (
-                participant.user_id,
-                participant.session_id,
-                track_type_int,
-            )
-
-            if track_key in self._track_map:
-                continue
-
-            # Poll for WebRTC track ID with timeout (same pattern as _on_track_published)
-            track_id = None
-            timeout = 10.0
-            poll_interval = 0.01
-            elapsed = 0.0
-
-            while elapsed < timeout:
-                for tid, (pending_user, pending_session, pending_kind) in list(
-                    self._pending_tracks.items()
-                ):
-                    if (
-                        pending_user == participant.user_id
-                        and pending_session == participant.session_id
-                        and pending_kind == expected_kind
-                    ):
-                        track_id = tid
-                        del self._pending_tracks[tid]
-                        break
-
-                if track_id:
-                    break
-
-                await asyncio.sleep(poll_interval)
-                elapsed += poll_interval
-
-            if track_id:
-                self._track_map[track_key] = {
-                    "track_id": track_id,
-                    "published": True,
-                }
-                sfu_participant = SfuParticipant.from_proto(participant)
-
-                self.events.send(
-                    events.TrackAddedEvent(
-                        plugin_name="getstream",
-                        track_id=track_id,
-                        track_type=track_type_int,
-                        user=sfu_participant,
-                        participant=sfu_participant,
-                    )
-                )
-            else:
-                logger.warning(
-                    f"No pending track for existing participant: "
-                    f"user={participant.user_id}, type={TrackType.Name(track_type_int)}"
-                )
-
     async def _on_track_published(self, event: sfu_events.TrackPublishedEvent):
         """Handle track published events from SFU - spawn TrackAddedEvent with correct type."""
         if not event.payload:
@@ -366,13 +284,10 @@ async def join(self, agent: "Agent", call: Call) -> StreamConnection:
         This function
         - initializes the chat channel
         - has the agent.agent_user join the call
-        - connect incoming audio/video to the agent
+        - connects incoming audio/video to the agent
         - connecting agent's outgoing audio/video to the call
-
-        TODO:
-        - process track flow
-
         """
+
         # Traditional mode - use WebRTC connection
         # Configure subscription for audio and video
         subscription_config = SubscriptionConfig(
@@ -401,13 +316,16 @@ async def on_audio_received(pcm: PcmData):
                 )
             )
 
-        await (
-            connection.__aenter__()
-        )  # TODO: weird API? there should be a manual version
-        self._connection = connection
+        # Re-emit certain events from the underlying RTC stack
+        # for the Agent to subscribe.
+        connection.on("participant_joined", self.events.send)
+        connection.on("participant_left", self.events.send)
+        connection.on("track_published", self.events.send)
+        connection.on("track_unpublished", self.events.send)
 
-        # Subscribe to tracks from participants who joined before the agent
-        await self._subscribe_to_existing_tracks(connection)
+        # Start the connection
+        await connection.__aenter__()
+        self._real_connection = connection
 
         standardize_connection = StreamConnection(connection)
         return standardize_connection