fix(llm): keep plugin-specific realtime handlers across a session swap

livekit-agents/livekit/agents/llm/realtime_fallback_adapter.py

@@ -6,7 +6,7 @@
 import weakref
 from collections.abc import AsyncIterable, Callable
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Literal
+from typing import TYPE_CHECKING, Any, Literal
 
 from livekit import rtc
 
@@ -55,7 +55,9 @@ class RealtimeAvailabilityChangedEvent:
     "can_disable_turn_detection",
 )
 
-# child events re-emitted on the wrapper
+# child events re-emitted on the wrapper. plugin-specific events (e.g. openai's
+# ``openai_server_event_received``) are not listed here -- the adapter cannot know them
+# ahead of time, so ``on()`` adds a forwarder for whatever a caller subscribes to.
 _FORWARDED_EVENTS: tuple[EventTypes, ...] = (
     "input_speech_started",
     "input_speech_stopped",
@@ -148,7 +150,9 @@ async def aclose(self) -> None:
             await model.aclose()
 
 
-class _FallbackRealtimeSession(RealtimeSession[Literal["realtime_availability_changed"]]):
+# ``str`` rather than a Literal: the wrapper stands in for any provider session, and a
+# plugin's own event names are only known to that plugin.
+class _FallbackRealtimeSession(RealtimeSession[str]):
     """Bound once by AgentActivity; swaps the inner child session internally."""
 
     def __init__(
@@ -164,17 +168,15 @@ def __init__(
         self._tools: NotGivenOr[list[Tool]] = NOT_GIVEN
         self._tool_choice: NotGivenOr[ToolChoice | None] = NOT_GIVEN
 
-        # stable per-event forwarders so they can be detached on swap
-        def _make_forwarder(event: EventTypes) -> Callable[[object], None]:
-            # callbacks receive only the payload, so bind the event name per forwarder
-            def _forward(ev: object) -> None:
-                self.emit(event, ev)
-
-            return _forward
-
-        self._forwarders: dict[EventTypes, Callable[[object], None]] = {
-            event: _make_forwarder(event) for event in _FORWARDED_EVENTS
-        }
+        # stable per-event forwarders so they can be detached on swap. ``_bind`` replays
+        # this dict onto every new child, so anything in it survives a swap.
+        self._forwarders: dict[str, Callable[..., None]] = {}
+        # whether ``_forwarders`` are currently attached to ``self._active``. a forwarder
+        # added while unbound is picked up by the next ``_bind`` instead of being attached
+        # to the child that is on its way out.
+        self._child_bound = False
+        for event in _FORWARDED_EVENTS:
+            self._add_forwarder(event)
 
         # per-model availability, with a cooldown after a failure
         self._available = [True] * len(adapter._models)
@@ -194,15 +196,55 @@ def _forward(ev: object) -> None:
         )
         self._bind(self._active)
 
+    def _add_forwarder(self, event: str) -> None:
+        """Start re-emitting ``event`` from the child, attaching to one if it is bound.
+
+        ``error`` is excluded: ``_on_child_error`` already re-emits it, sometimes re-stamped
+        as recoverable, and a second forwarder would duplicate every error.
+        """
+        if event in self._forwarders or event == "error":
+            return
+
+        # varargs, since a plugin event may carry any number of payload arguments
+        def _forward(*args: object) -> None:
+            self.emit(event, *args)
+
+        self._forwarders[event] = _forward
+        if self._child_bound:
+            self._active.on(event, _forward)
+
+    def on(self, event: str, callback: Callable[..., Any] | None = None) -> Callable[..., Any]:
+        """Subscribe to a child event on the wrapper, plugin-specific ones included.
+
+        Handlers registered here survive a swap. The adapter keeps a forwarder for every
+        subscribed event and re-attaches it to each new child, so a ``restart_session()`` or
+        a fallback does not silently drop them -- unlike subscribing on ``_active``, which
+        binds to the one child that a swap then discards.
+        """
+        self._add_forwarder(event)
+        return super().on(event, callback)
+
+    def once(self, event: str, callback: Callable[..., Any] | None = None) -> Callable[..., Any]:
+        """One-shot :meth:`on`, registering the forwarder the same way.
+
+        ``EventEmitter.once`` currently routes through ``on``, so this would work without the
+        override. Registering here too keeps one-shot plugin subscriptions from silently
+        never firing if that ever stops being true.
+        """
+        self._add_forwarder(event)
+        return super().once(event, callback)
+
     def _bind(self, child: RealtimeSession) -> None:
         for event, forwarder in self._forwarders.items():
             child.on(event, forwarder)
         child.on("error", self._on_child_error)
+        self._child_bound = True
 
     def _unbind(self, child: RealtimeSession) -> None:
         for event, forwarder in self._forwarders.items():
             child.off(event, forwarder)
         child.off("error", self._on_child_error)
+        self._child_bound = False
 
     def _set_available(self, index: int, available: bool) -> None:
         if self._available[index] == available:

tests/fake_realtime.py

@@ -167,12 +167,16 @@ def __init__(
         self.closed = False
         # when set, every session this model creates fails to bring up
         self.bring_up_error: Exception | None = None
+        # when set, session() itself raises, before any child exists to bind
+        self.session_error: Exception | None = None
 
     @property
     def active_session(self) -> FakeRealtimeSession:
         return self.created_sessions[-1]
 
     def session(self, *, turn_detection_disabled: bool = False) -> FakeRealtimeSession:
+        if self.session_error is not None:
+            raise self.session_error
         sess = FakeRealtimeSession(self, turn_detection_disabled=turn_detection_disabled)
         sess.update_error = self.bring_up_error
         self.created_sessions.append(sess)

tests/test_realtime_fallback.py

@@ -201,6 +201,159 @@ async def test_restart_preserves_wrapper_subscribers() -> None:
     assert received == ["after-restart"]
 
 
+# a provider-specific event the adapter knows nothing about ahead of time, standing in for
+# e.g. openai's "openai_server_event_received"
+_PLUGIN_EVENT = "provider_specific_event"
+
+
+async def test_restart_preserves_plugin_event_subscribers() -> None:
+    primary = FakeRealtimeModel()
+    adapter = RealtimeModelFallbackAdapter([primary])
+    session = adapter.session()
+    received: list[object] = []
+    session.on(_PLUGIN_EVENT, lambda ev: received.append(ev))
+
+    await adapter.restart_session()
+
+    # plugin-specific subscribers are re-attached to the new child, same as the generic ones
+    session._active.emit(_PLUGIN_EVENT, "after-restart")
+    assert received == ["after-restart"]
+
+
+async def test_fallback_preserves_plugin_event_subscribers() -> None:
+    primary = FakeRealtimeModel()
+    backup = FakeRealtimeModel()
+    adapter = RealtimeModelFallbackAdapter([primary, backup])
+    session = adapter.session()
+    received: list[object] = []
+    session.on(_PLUGIN_EVENT, lambda ev: received.append(ev))
+
+    primary.active_session.emit_error(recoverable=False)
+    await session._swap_task
+
+    # the subscriber follows an availability swap too, not just an explicit restart
+    assert session._active is backup.active_session
+    session._active.emit(_PLUGIN_EVENT, "after-fallback")
+    assert received == ["after-fallback"]
+
+
+async def test_plugin_event_detached_from_discarded_child() -> None:
+    primary = FakeRealtimeModel()
+    adapter = RealtimeModelFallbackAdapter([primary])
+    session = adapter.session()
+    received: list[object] = []
+    session.on(_PLUGIN_EVENT, lambda ev: received.append(ev))
+    old_child = primary.active_session
+
+    await adapter.restart_session()
+
+    # the discarded child is fully unbound, so a late event from it is not re-emitted
+    old_child.emit(_PLUGIN_EVENT, "from-dead-child")
+    assert received == []
+
+
+async def test_plugin_event_subscribed_mid_swap_binds_to_new_child() -> None:
+    primary = FakeRealtimeModel()
+    adapter = RealtimeModelFallbackAdapter([primary])
+    session = adapter.session()
+    old_child = primary.active_session
+
+    gate = asyncio.Event()
+    old_child.block_aclose = gate
+
+    restart = asyncio.create_task(adapter.restart_session())
+    await old_child.aclose_entered.wait()  # swap is mid-flight, no child bound
+
+    received: list[object] = []
+    session.on(_PLUGIN_EVENT, lambda ev: received.append(ev))
+    old_child.emit(_PLUGIN_EVENT, "from-dead-child")
+
+    gate.set()
+    await restart
+
+    # subscribing mid-swap attaches to the incoming child, never the one being discarded
+    session._active.emit(_PLUGIN_EVENT, "from-new-child")
+    assert received == ["from-new-child"]
+
+
+async def test_swap_survives_bring_up_raising_before_a_child_exists() -> None:
+    primary = FakeRealtimeModel()
+    backup1 = FakeRealtimeModel()
+    backup2 = FakeRealtimeModel()
+    adapter = RealtimeModelFallbackAdapter([primary, backup1, backup2])
+    session = adapter.session()
+    errors: list = []
+    session.on("error", lambda e: errors.append(e))
+
+    # session() raises, so _bring_up's failure path unbinds the *outgoing* child a second
+    # time -- and by then _forwarders holds an entry that was never attached to it
+    backup1.session_error = RuntimeError("cannot construct session")
+    old_child = primary.active_session
+    gate = asyncio.Event()
+    old_child.block_aclose = gate
+
+    received: list[object] = []
+    old_child.emit_error(recoverable=False)
+    await old_child.aclose_entered.wait()
+    # registered mid-swap, so it is never bound to old_child
+    session.on(_PLUGIN_EVENT, lambda ev: received.append(ev))
+    gate.set()
+    await session._swap_task
+
+    # detaching a forwarder the child never had is a no-op, so the swap still cascades
+    assert session._active_index == 2
+    assert session._active is backup2.active_session
+    assert all(e.recoverable for e in errors)
+    # ...and the mid-swap subscriber lands on the model the cascade settled on
+    session._active.emit(_PLUGIN_EVENT, "still-forwarded")
+    assert received == ["still-forwarded"]
+
+
+def test_forwards_multi_arg_plugin_events() -> None:
+    primary = FakeRealtimeModel()
+    session = RealtimeModelFallbackAdapter([primary]).session()
+    received: list[tuple[object, ...]] = []
+    session.on(_PLUGIN_EVENT, lambda *args: received.append(args))
+
+    primary.active_session.emit(_PLUGIN_EVENT, "a", "b")
+
+    # forwarders are varargs: a plugin event carrying several payload args keeps them all
+    assert received == [("a", "b")]
+
+
+async def test_once_plugin_subscription_survives_a_swap_and_fires_once() -> None:
+    primary = FakeRealtimeModel()
+    adapter = RealtimeModelFallbackAdapter([primary])
+    session = adapter.session()
+    received: list[object] = []
+    session.once(_PLUGIN_EVENT, lambda ev: received.append(ev))
+
+    await adapter.restart_session()
+
+    session._active.emit(_PLUGIN_EVENT, "first")
+    session._active.emit(_PLUGIN_EVENT, "second")
+
+    # a one-shot subscriber gets a forwarder too, is re-attached to the new child, and
+    # still fires exactly once
+    assert received == ["first"]
+
+
+def test_multi_arg_forwarding_is_independent_per_subscriber() -> None:
+    primary = FakeRealtimeModel()
+    session = RealtimeModelFallbackAdapter([primary]).session()
+    all_args: list[tuple[object, ...]] = []
+    first_only: list[object] = []
+    session.on(_PLUGIN_EVENT, lambda *args: all_args.append(args))
+    session.on(_PLUGIN_EVENT, lambda ev: first_only.append(ev))
+
+    primary.active_session.emit(_PLUGIN_EVENT, "a", "b")
+
+    # EventEmitter trims to each callback's own arity, so a single-arg subscriber on the
+    # same event doesn't shorten what the varargs one receives
+    assert all_args == [("a", "b")]
+    assert first_only == ["a"]
+
+
 async def test_restart_emits_no_error() -> None:
     primary = FakeRealtimeModel()
     adapter = RealtimeModelFallbackAdapter([primary])