OpenAI realtime: a long chat context sync or reconnect replay is sent in one event loop turn

# /// script
# requires-python = ">=3.10"
# dependencies = ["livekit-agents", "livekit-plugins-openai"]
# ///
"""The OpenAI Realtime session sends a burst of client events in one event loop turn.

The script starts a local websocket server in a child process. The server speaks
enough of the OpenAI Realtime protocol to confirm chat items. No request goes to
OpenAI.

1. update_chat_ctx: the session syncs a chat context of 1000 messages.
2. reconnect: max_session_duration expires, and the session replays the chat
   context on a new connection.

A probe task runs once per event loop turn. It counts the client events that the
session sent since its previous turn ("openai_client_event_queued" events). The
time in parentheses is the time from the first to the last send of the largest
turn. It depends on the machine.

--fixed applies the proposed change to RealtimeSession._run_ws and
RealtimeSession._main_task in the process.
"""

from __future__ import annotations

import __future__
import asyncio
import inspect
import json
import logging
import sys
import textwrap
import time
from importlib.metadata import version

from aiohttp import WSMsgType, web
from livekit.agents.llm import ChatContext
from livekit.agents.llm._realtime import openai as realtime_openai
from livekit.plugins.openai.realtime import RealtimeModel

MESSAGES = 1000
MAX_SESSION_DURATION = 3.0
TIMEOUT = 30.0

FIX = {
    "_run_ws": (
        (
            "            except Exception:\n"
            '                logger.exception("failed to send event")\n'
            "\n"
            "        closing = True\n",
            "            except Exception:\n"
            '                logger.exception("failed to send event")\n'
            "\n"
            "            # one event per loop turn, so a long queue does not block other tasks\n"
            "            await asyncio.sleep(0)\n"
            "\n"
            "        closing = True\n",
        ),
    ),
    "_main_task": (
        (
            '                self.emit("openai_client_event_queued", ev)\n'
            "                await ws_conn.send_str(json.dumps(ev))\n",
            '                self.emit("openai_client_event_queued", ev)\n'
            "                await ws_conn.send_str(json.dumps(ev))\n"
            "                # one event per loop turn, so a long replay does not block other tasks\n"
            "                await asyncio.sleep(0)\n",
        ),
    ),
}


def apply_fix() -> None:
    for name, replacements in FIX.items():
        function = inspect.unwrap(getattr(realtime_openai.RealtimeSession, name))
        source = textwrap.dedent(inspect.getsource(function))
        for old, new in replacements:
            if source.count(old) != 1:
                sys.exit(f"--fixed: {name} does not contain the expected text:\n{old}")
            source = source.replace(old, new)
        namespace = dict(vars(realtime_openai))
        code = compile(
            source,
            realtime_openai.__file__,
            "exec",
            flags=__future__.annotations.compiler_flag,
            dont_inherit=True,
        )
        exec(code, namespace)
        function.__code__ = inspect.unwrap(namespace[name]).__code__


async def realtime_endpoint(request: web.Request) -> web.WebSocketResponse:
    """Confirm the session update and every conversation.item.create."""
    ws = web.WebSocketResponse()
    await ws.prepare(request)
    session = {"id": "sess_local", "object": "realtime.session", "model": "gpt-realtime"}
    await ws.send_str(json.dumps({"type": "session.created", "event_id": "ev_c", "session": session}))
    async for message in ws:
        if message.type is not WSMsgType.TEXT:
            continue
        event = json.loads(message.data)
        if event["type"] == "session.update":
            await ws.send_str(
                json.dumps({"type": "session.updated", "event_id": "ev_u", "session": session})
            )
        elif event["type"] == "conversation.item.create":
            previous_item_id = event.get("previous_item_id")
            await ws.send_str(
                json.dumps(
                    {
                        "type": "conversation.item.added",
                        "event_id": f"ev_{event['item']['id']}",
                        "previous_item_id": None if previous_item_id == "root" else previous_item_id,
                        "item": event["item"],
                    }
                )
            )
    return ws


async def serve() -> None:
    """Child process: serve until the parent closes stdin."""
    app = web.Application()
    app.router.add_get("/realtime", realtime_endpoint)
    runner = web.AppRunner(app, shutdown_timeout=1.0)
    await runner.setup()
    site = web.TCPSite(runner, "127.0.0.1", 0)
    await site.start()
    print(runner.addresses[0][1], flush=True)
    await asyncio.get_running_loop().run_in_executor(None, sys.stdin.read)
    await runner.cleanup()


class SendProbe:
    """Count the client events that the session sends in each event loop turn."""

    def __init__(self, session: realtime_openai.RealtimeSession) -> None:
        self.reset()
        session.on("openai_client_event_queued", self._on_event)

    def reset(self) -> None:
        self.sent = 0
        self.most_in_one_turn = 0
        self.largest_turn_ms = 0.0
        self._turn: list[float] = []

    def _on_event(self, _event: dict) -> None:
        self.sent += 1
        self._turn.append(time.perf_counter())

    def end_turn(self) -> None:
        if len(self._turn) > self.most_in_one_turn:
            self.most_in_one_turn = len(self._turn)
            self.largest_turn_ms = (self._turn[-1] - self._turn[0]) * 1000
        self._turn.clear()

    async def run(self) -> None:
        while True:
            self.end_turn()
            await asyncio.sleep(0)

    def report(self, phase: str) -> int:
        self.end_turn()
        print(
            f"{phase}: the session sent {self.sent} events, at most {self.most_in_one_turn} "
            f"in one event loop turn ({self.largest_turn_ms:.1f} ms)"
        )
        most = self.most_in_one_turn
        self.reset()
        return most


async def wait_for(condition, timeout: float = TIMEOUT) -> bool:
    for _ in range(int(timeout / 0.01)):
        if condition():
            return True
        await asyncio.sleep(0.01)
    return False


async def sync_then_reconnect(base_url: str) -> tuple[int, int]:
    model = RealtimeModel(
        api_key="local",
        base_url=base_url,
        model="gpt-realtime",
        max_session_duration=MAX_SESSION_DURATION,
    )
    session = model.session()
    probe = SendProbe(session)
    reconnected: list = []
    session.on("session_reconnected", reconnected.append)
    if not await wait_for(lambda: probe.sent == 1):
        sys.exit("the session did not send session.update")
    probe.reset()
    probe_task = asyncio.create_task(probe.run())

    chat_ctx = ChatContext()
    for i in range(MESSAGES):
        chat_ctx.add_message(role="user" if i % 2 == 0 else "assistant", content=f"message {i}")
    await asyncio.wait_for(session.update_chat_ctx(chat_ctx), TIMEOUT)
    if reconnected:
        sys.exit("max_session_duration expired during update_chat_ctx, raise MAX_SESSION_DURATION")
    sync = probe.report(f"update_chat_ctx with {MESSAGES} messages")

    if not await wait_for(lambda: reconnected):
        sys.exit("the session did not reconnect")
    replay = probe.report("reconnect after max_session_duration")

    probe_task.cancel()
    await session.aclose()
    await model.aclose()
    return sync, replay


async def main() -> int:
    server = await asyncio.create_subprocess_exec(
        sys.executable,
        __file__,
        "--serve",
        stdin=asyncio.subprocess.PIPE,
        stdout=asyncio.subprocess.PIPE,
    )
    try:
        port = int(await asyncio.wait_for(server.stdout.readline(), TIMEOUT))
        sync, replay = await sync_then_reconnect(f"http://127.0.0.1:{port}")
    finally:
        server.stdin.close()
        await asyncio.wait_for(server.wait(), TIMEOUT)

    most = max(sync, replay)
    if most > 1:
        print(f"PROBLEM PRESENT: the session sent up to {most} client events in one event loop turn")
        return 1
    print("PROBLEM ABSENT: the session sent at most 1 client event in one event loop turn")
    return 0


if __name__ == "__main__":
    if "--serve" in sys.argv:
        asyncio.run(serve())
        sys.exit(0)
    logging.getLogger("livekit").setLevel(logging.CRITICAL + 1)
    if "--fixed" in sys.argv:
        apply_fix()
    print(
        f"livekit-agents {version('livekit-agents')}, "
        f"livekit-plugins-openai {version('livekit-plugins-openai')}, "
        f"aiohttp {version('aiohttp')}, Python {sys.version.split()[0]}"
    )
    sys.exit(asyncio.run(main()))