test_tracing.py

from types import SimpleNamespace
from typing import Any, cast
from unittest.mock import AsyncMock, Mock, patch

import pytest
from openai.types.realtime.realtime_session_create_request import (
    RealtimeSessionCreateRequest,
)
from openai.types.realtime.realtime_tracing_config import TracingConfiguration

from agents.realtime.agent import RealtimeAgent
from agents.realtime.model import RealtimeModel, RealtimeModelConfig
from agents.realtime.openai_realtime import OpenAIRealtimeWebSocketModel
from agents.realtime.session import RealtimeSession


@pytest.fixture(autouse=True)
def mock_client_secret_request(monkeypatch):
    records: dict[str, list[dict[str, Any]]] = {"init_kwargs": [], "sessions": []}

    class DummySecrets:
        async def create(self, *, session: dict[str, Any]) -> SimpleNamespace:
            records["sessions"].append(session)
            return SimpleNamespace(value="ek_test")

    class DummyRealtime:
        def __init__(self):
            self.client_secrets = DummySecrets()

    class DummyClient:
        def __init__(self, *args, **kwargs):
            records["init_kwargs"].append(kwargs)
            self.realtime = DummyRealtime()

        async def close(self) -> None:
            return None

    monkeypatch.setattr(
        "agents.realtime.openai_realtime.AsyncOpenAI",
        DummyClient,
    )

    return records


class TestRealtimeTracingIntegration:
    """Test tracing configuration and session.update integration."""

    @pytest.fixture
    def model(self):
        """Create a fresh model instance for each test."""
        return OpenAIRealtimeWebSocketModel()

    @pytest.fixture
    def mock_websocket(self):
        """Create a mock websocket connection."""
        mock_ws = AsyncMock()
        mock_ws.send = AsyncMock()
        mock_ws.close = AsyncMock()
        return mock_ws

    @pytest.mark.asyncio
    async def test_tracing_config_storage_and_defaults(self, model, mock_websocket):
        """Test that tracing config is stored correctly and defaults to 'auto'."""
        # Test with explicit tracing config
        config_with_tracing = {
            "api_key": "test-key",
            "initial_model_settings": {
                "tracing": {
                    "workflow_name": "test_workflow",
                    "group_id": "group_123",
                    "metadata": {"version": "1.0"},
                }
            },
        }

        async def async_websocket(*args, **kwargs):
            return mock_websocket

        with patch("websockets.connect", side_effect=async_websocket):
            with patch("asyncio.create_task") as mock_create_task:
                mock_task = AsyncMock()
                mock_create_task.return_value = mock_task
                mock_create_task.side_effect = lambda coro: (coro.close(), mock_task)[1]

                await model.connect(config_with_tracing)

                # Should store the tracing config
                assert model._tracing_config == {
                    "workflow_name": "test_workflow",
                    "group_id": "group_123",
                    "metadata": {"version": "1.0"},
                }


        # Test without tracing config - should default to "auto"
        model2 = OpenAIRealtimeWebSocketModel()
        config_no_tracing = {
            "api_key": "test-key",
            "initial_model_settings": {},
        }

        with patch("websockets.connect", side_effect=async_websocket):
            with patch("asyncio.create_task") as mock_create_task:
                mock_create_task.side_effect = lambda coro: (coro.close(), mock_task)[1]

                await model2.connect(config_no_tracing)  # type: ignore[arg-type]
                assert model2._tracing_config == "auto"

    @pytest.mark.asyncio
    async def test_send_tracing_config_on_session_created(self, model, mock_websocket):
        """Test that tracing config is sent when session.created event is received."""
        config = {
            "api_key": "test-key",
            "initial_model_settings": {
                "tracing": {"workflow_name": "test_workflow", "group_id": "group_123"}
            },
        }

        async def async_websocket(*args, **kwargs):
            return mock_websocket

        with patch("websockets.connect", side_effect=async_websocket):
            with patch("asyncio.create_task") as mock_create_task:
                mock_task = AsyncMock()
                mock_create_task.side_effect = lambda coro: (coro.close(), mock_task)[1]

                await model.connect(config)

                # Simulate session.created event
                session_created_event = {
                    "type": "session.created",
                    "event_id": "event_123",
                    "session": {"id": "session_456", "type": "realtime", "model": "gpt-realtime"},
                }

                with patch.object(model, "_send_raw_message") as mock_send_raw_message:
                    await model._handle_ws_event(session_created_event)

                    # Should send session.update with tracing config
                    from openai.types.realtime.session_update_event import (
                        SessionUpdateEvent,
                    )

                    mock_send_raw_message.assert_called_once()
                    call_args = mock_send_raw_message.call_args[0][0]
                    assert isinstance(call_args, SessionUpdateEvent)
                    assert call_args.type == "session.update"
                    session_req = cast(RealtimeSessionCreateRequest, call_args.session)
                    assert isinstance(session_req.tracing, TracingConfiguration)
                    assert session_req.tracing.workflow_name == "test_workflow"
                    assert session_req.tracing.group_id == "group_123"

    @pytest.mark.asyncio
    async def test_send_tracing_config_auto_mode(self, model, mock_websocket):
        """Test that 'auto' tracing config is sent correctly."""
        config = {
            "api_key": "test-key",
            "initial_model_settings": {},  # No tracing config - defaults to "auto"
        }

        async def async_websocket(*args, **kwargs):
            return mock_websocket

        with patch("websockets.connect", side_effect=async_websocket):
            with patch("asyncio.create_task") as mock_create_task:
                mock_task = AsyncMock()
                mock_create_task.side_effect = lambda coro: (coro.close(), mock_task)[1]

                await model.connect(config)

                session_created_event = {
                    "type": "session.created",
                    "event_id": "event_123",
                    "session": {"id": "session_456", "type": "realtime", "model": "gpt-realtime"},
                }

                with patch.object(model, "_send_raw_message") as mock_send_raw_message:
                    await model._handle_ws_event(session_created_event)

                    # Should send session.update with "auto"
                    from openai.types.realtime.session_update_event import SessionUpdateEvent

                    mock_send_raw_message.assert_called_once()
                    call_args = mock_send_raw_message.call_args[0][0]
                    assert isinstance(call_args, SessionUpdateEvent)
                    assert call_args.type == "session.update"
                    session_req = cast(RealtimeSessionCreateRequest, call_args.session)
                    assert session_req.tracing == "auto"

    @pytest.mark.asyncio
    async def test_tracing_config_none_skips_session_update(self, model, mock_websocket):
        """Test that None tracing config skips sending session.update."""
        # Manually set tracing config to None (this would happen if explicitly set)
        model._tracing_config = None

        session_created_event = {
            "type": "session.created",
            "event_id": "event_123",
            "session": {"id": "session_456", "type": "realtime", "model": "gpt-realtime"},
        }

        with patch.object(model, "send_event") as mock_send_event:
            await model._handle_ws_event(session_created_event)

            # Should not send any session.update
            mock_send_event.assert_not_called()

    @pytest.mark.asyncio
    async def test_tracing_config_with_metadata_serialization(self, model, mock_websocket):
        """Test that complex metadata in tracing config is handled correctly."""
        complex_metadata = {
            "user_id": "user_123",
            "session_type": "demo",
            "features": ["audio", "tools"],
            "config": {"timeout": 30, "retries": 3},
        }

        config = {
            "api_key": "test-key",
            "initial_model_settings": {
                "tracing": {"workflow_name": "complex_workflow", "metadata": complex_metadata}
            },
        }

        async def async_websocket(*args, **kwargs):
            return mock_websocket

        with patch("websockets.connect", side_effect=async_websocket):
            with patch("asyncio.create_task") as mock_create_task:
                mock_task = AsyncMock()
                mock_create_task.side_effect = lambda coro: (coro.close(), mock_task)[1]

                await model.connect(config)

                session_created_event = {
                    "type": "session.created",
                    "event_id": "event_123",
                    "session": {"id": "session_456", "type": "realtime", "model": "gpt-realtime"},
                }

                with patch.object(model, "_send_raw_message") as mock_send_raw_message:
                    await model._handle_ws_event(session_created_event)

                    # Should send session.update with complete tracing config including metadata
                    from openai.types.realtime.session_update_event import (
                        SessionUpdateEvent,
                    )

                    mock_send_raw_message.assert_called_once()
                    call_args = mock_send_raw_message.call_args[0][0]
                    assert isinstance(call_args, SessionUpdateEvent)
                    assert call_args.type == "session.update"
                    session_req = cast(RealtimeSessionCreateRequest, call_args.session)
                    assert isinstance(session_req.tracing, TracingConfiguration)
                    assert session_req.tracing.workflow_name == "complex_workflow"
                    assert session_req.tracing.metadata == complex_metadata

    @pytest.mark.asyncio
    async def test_tracing_disabled_prevents_tracing(self, mock_websocket):
        """Test that tracing_disabled=True prevents tracing configuration."""

        # Create a test agent and mock model
        agent = RealtimeAgent(name="test_agent", instructions="test")
        agent.handoffs = []

        mock_model = Mock(spec=RealtimeModel)

        # Create session with tracing disabled
        session = RealtimeSession(
            model=mock_model,
            agent=agent,
            context=None,
            model_config=None,
            run_config={"tracing_disabled": True},
        )

        # Test the _get_updated_model_settings_from_agent method directly
        model_settings = await session._get_updated_model_settings_from_agent(
            starting_settings=None, agent=agent
        )

        # When tracing is disabled, model settings should have tracing=None
        assert model_settings["tracing"] is None

    @pytest.mark.asyncio
    async def test_connect_sets_sdk_headers_and_subprotocols(
        self,
        mock_websocket,
        mock_client_secret_request,
    ):
        """Ensure websocket handshake mirrors Agents JS with client secrets."""
        model = OpenAIRealtimeWebSocketModel()
        config: RealtimeModelConfig = {
            "api_key": "sk-test",
            "initial_model_settings": {},
        }

        captured_kwargs: dict[str, Any] = {}

        async def async_websocket(*args, **kwargs):
            captured_kwargs.update(kwargs)
            return mock_websocket

        with patch("websockets.connect", side_effect=async_websocket):
            with patch("asyncio.create_task") as mock_create_task:
                mock_task = AsyncMock()
                mock_create_task.return_value = mock_task
                mock_create_task.side_effect = lambda coro: (coro.close(), mock_task)[1]

                await model.connect(config)

        headers = captured_kwargs["additional_headers"]
        assert "Authorization" not in headers

        subprotocols = captured_kwargs["subprotocols"]
        assert subprotocols[0] == "realtime"
        assert subprotocols[1].startswith("openai-insecure-api-key.ek_test")
        assert subprotocols[2].startswith("openai-agents-sdk.python.")
        # Ensure client secret API was called once
        assert mock_client_secret_request["init_kwargs"] == [{"api_key": "sk-test"}]
        assert mock_client_secret_request["sessions"] == [
            {"type": "realtime", "model": "gpt-realtime"}
        ]

    @pytest.mark.asyncio
    async def test_connect_with_ephemeral_key_skips_client_secret(
        self,
        mock_websocket,
        mock_client_secret_request,
    ):
        """Ensure pre-generated ek_ keys are used directly without calling the API."""
        model = OpenAIRealtimeWebSocketModel()
        config: RealtimeModelConfig = {
            "api_key": "ek_existing",
            "initial_model_settings": {},
        }

        captured_kwargs: dict[str, Any] = {}

        async def async_websocket(*args, **kwargs):
            captured_kwargs.update(kwargs)
            return mock_websocket

        with patch("websockets.connect", side_effect=async_websocket):
            with patch("asyncio.create_task") as mock_create_task:
                mock_task = AsyncMock()
                mock_create_task.return_value = mock_task
                mock_create_task.side_effect = lambda coro: (coro.close(), mock_task)[1]

                await model.connect(config)

        # No client secret API calls should have been made
        assert mock_client_secret_request["init_kwargs"] == []
        assert mock_client_secret_request["sessions"] == []

        subprotocols = captured_kwargs["subprotocols"]
        assert subprotocols[1] == "openai-insecure-api-key.ek_existing"