test_trace_processor.py

import os
import threading
import time
from typing import Any, cast
from unittest.mock import MagicMock, patch

import httpx
import pytest

from agents.tracing import flush_traces, get_trace_provider
from agents.tracing.processor_interface import TracingExporter, TracingProcessor
from agents.tracing.processors import BackendSpanExporter, BatchTraceProcessor
from agents.tracing.provider import DefaultTraceProvider, TraceProvider
from agents.tracing.span_data import AgentSpanData
from agents.tracing.spans import Span, SpanImpl
from agents.tracing.traces import Trace, TraceImpl


def get_span(processor: TracingProcessor) -> SpanImpl[AgentSpanData]:
    """Create a minimal agent span for testing processors."""
    return SpanImpl(
        trace_id="test_trace_id",
        span_id="test_span_id",
        parent_id=None,
        processor=processor,
        span_data=AgentSpanData(name="test_agent"),
        tracing_api_key=None,
    )


def get_trace(processor: TracingProcessor) -> TraceImpl:
    """Create a minimal trace."""
    return TraceImpl(
        name="test_trace",
        trace_id="test_trace_id",
        group_id="test_session_id",
        metadata={},
        processor=processor,
        tracing_api_key=None,
    )


@pytest.fixture
def mocked_exporter():
    exporter = MagicMock()
    exporter.export = MagicMock()
    return exporter


def test_batch_trace_processor_on_trace_start(mocked_exporter):
    processor = BatchTraceProcessor(exporter=mocked_exporter, schedule_delay=0.1)
    test_trace = get_trace(processor)

    processor.on_trace_start(test_trace)
    assert processor._queue.qsize() == 1, "Trace should be added to the queue"

    # Shutdown to clean up the worker thread
    processor.shutdown()


def test_batch_trace_processor_on_span_end(mocked_exporter):
    processor = BatchTraceProcessor(exporter=mocked_exporter, schedule_delay=0.1)
    test_span = get_span(processor)

    processor.on_span_end(test_span)
    assert processor._queue.qsize() == 1, "Span should be added to the queue"

    # Shutdown to clean up the worker thread
    processor.shutdown()


def test_batch_trace_processor_queue_full(mocked_exporter):
    processor = BatchTraceProcessor(exporter=mocked_exporter, max_queue_size=2, schedule_delay=0.1)
    # Fill the queue
    processor.on_trace_start(get_trace(processor))
    processor.on_trace_start(get_trace(processor))
    assert processor._queue.full() is True

    # Next item should not be queued
    processor.on_trace_start(get_trace(processor))
    assert processor._queue.qsize() == 2, "Queue should not exceed max_queue_size"

    processor.on_span_end(get_span(processor))
    assert processor._queue.qsize() == 2, "Queue should not exceed max_queue_size"

    processor.shutdown()


def test_batch_processor_doesnt_enqueue_on_trace_end_or_span_start(mocked_exporter):
    processor = BatchTraceProcessor(exporter=mocked_exporter)

    processor.on_trace_start(get_trace(processor))
    assert processor._queue.qsize() == 1, "Trace should be queued"

    processor.on_span_start(get_span(processor))
    assert processor._queue.qsize() == 1, "Span should not be queued"

    processor.on_span_end(get_span(processor))
    assert processor._queue.qsize() == 2, "Span should be queued"

    processor.on_trace_end(get_trace(processor))
    assert processor._queue.qsize() == 2, "Nothing new should be queued"

    processor.shutdown()


def test_batch_trace_processor_force_flush(mocked_exporter):
    processor = BatchTraceProcessor(exporter=mocked_exporter, max_batch_size=2, schedule_delay=5.0)

    processor.on_trace_start(get_trace(processor))
    processor.on_span_end(get_span(processor))
    processor.on_span_end(get_span(processor))

    processor.force_flush()

    # Ensure exporter.export was called with all items
    # Because max_batch_size=2, it may have been called multiple times
    total_exported = 0
    for call_args in mocked_exporter.export.call_args_list:
        batch = call_args[0][0]  # first positional arg to export() is the items list
        total_exported += len(batch)

    # We pushed 3 items; ensure they all got exported
    assert total_exported == 3

    processor.shutdown()


def test_batch_trace_processor_force_flush_waits_for_in_flight_background_export():
    export_started = threading.Event()
    export_continue = threading.Event()

    class BlockingExporter(TracingExporter):
        def export(self, items: list[Trace | Span[Any]]) -> None:
            export_started.set()
            assert export_continue.wait(timeout=2.0)

    processor = BatchTraceProcessor(exporter=BlockingExporter(), schedule_delay=0.01)
    processor.on_trace_start(get_trace(processor))

    assert export_started.wait(timeout=2.0)

    flush_thread = threading.Thread(target=processor.force_flush)
    flush_thread.start()

    time.sleep(0.1)
    assert flush_thread.is_alive(), "force_flush() should wait for an in-flight export"

    export_continue.set()
    flush_thread.join(timeout=2.0)

    assert not flush_thread.is_alive()

    processor.shutdown()


def test_batch_trace_processor_shutdown_flushes(mocked_exporter):
    processor = BatchTraceProcessor(exporter=mocked_exporter, schedule_delay=5.0)
    processor.on_trace_start(get_trace(processor))
    processor.on_span_end(get_span(processor))
    qsize_before = processor._queue.qsize()
    assert qsize_before == 2

    processor.shutdown()

    # Ensure everything was exported after shutdown
    total_exported = 0
    for call_args in mocked_exporter.export.call_args_list:
        batch = call_args[0][0]
        total_exported += len(batch)

    assert total_exported == 2, "All items in the queue should be exported upon shutdown"


def test_batch_trace_processor_scheduled_export(mocked_exporter):
    """
    Tests that items are automatically exported when the schedule_delay expires.
    We mock time.time() so we can trigger the condition without waiting in real time.
    """
    with patch("time.time") as mock_time:
        base_time = 1000.0
        mock_time.return_value = base_time

        processor = BatchTraceProcessor(exporter=mocked_exporter, schedule_delay=1.0)

        processor.on_span_end(get_span(processor))  # queue size = 1

        # Now artificially advance time beyond the next export time
        mock_time.return_value = base_time + 2.0  # > base_time + schedule_delay
        # Let the background thread run a bit
        time.sleep(0.3)

        # Check that exporter.export was eventually called
        # Because the background thread runs, we might need a small sleep
        processor.shutdown()

    total_exported = 0
    for call_args in mocked_exporter.export.call_args_list:
        batch = call_args[0][0]
        total_exported += len(batch)

    assert total_exported == 1, "Item should be exported after scheduled delay"


def test_flush_traces_delegates_to_default_trace_provider():
    provider = DefaultTraceProvider()
    mock_processor = MagicMock()
    provider.register_processor(mock_processor)

    with patch("agents.tracing.setup.GLOBAL_TRACE_PROVIDER", provider):
        flush_traces()

    mock_processor.force_flush.assert_called_once()


def test_flush_traces_is_importable_from_top_level_agents_package():
    from agents import flush_traces as top_level_flush_traces

    assert top_level_flush_traces is flush_traces


def test_default_trace_provider_force_flush_respects_disabled_flag():
    provider = DefaultTraceProvider()
    mock_processor = MagicMock()
    provider.register_processor(mock_processor)

    provider.set_disabled(True)
    provider.force_flush()

    mock_processor.force_flush.assert_not_called()


def test_trace_provider_force_flush_and_shutdown_default_to_noops():
    class MinimalProvider(TraceProvider):
        def register_processor(self, processor: TracingProcessor) -> None:
            pass

        def set_processors(self, processors: list[TracingProcessor]) -> None:
            pass

        def get_current_trace(self):
            return None

        def get_current_span(self):
            return None

        def set_disabled(self, disabled: bool) -> None:
            pass

        def time_iso(self) -> str:
            return ""

        def gen_trace_id(self) -> str:
            return "trace_123"

        def gen_span_id(self) -> str:
            return "span_123"

        def gen_group_id(self) -> str:
            return "group_123"

        def create_trace(
            self,
            name,
            trace_id=None,
            group_id=None,
            metadata=None,
            disabled=False,
            tracing=None,
        ):
            raise NotImplementedError

        def create_span(self, span_data, span_id=None, parent=None, disabled=False):
            raise NotImplementedError

    provider = MinimalProvider()
    provider.force_flush()
    provider.shutdown()


def test_get_trace_provider_force_flush_flushes_default_processor(mocked_exporter):
    provider = DefaultTraceProvider()
    processor = BatchTraceProcessor(exporter=mocked_exporter, schedule_delay=60.0)
    provider.register_processor(processor)

    with patch("agents.tracing.setup.GLOBAL_TRACE_PROVIDER", provider):
        processor.on_trace_start(get_trace(processor))
        processor.on_span_end(get_span(processor))

        get_trace_provider().force_flush()

    total_exported = sum(
        len(call_args[0][0]) for call_args in mocked_exporter.export.call_args_list
    )
    assert total_exported == 2
    processor.shutdown()


@pytest.fixture
def patched_time_sleep():
    """
    Fixture to replace time.sleep with a no-op to speed up tests
    that rely on retry/backoff logic.
    """
    with patch("time.sleep") as mock_sleep:
        yield mock_sleep


def mock_processor():
    processor = MagicMock()
    processor.on_trace_start = MagicMock()
    processor.on_span_end = MagicMock()
    return processor


@patch("httpx.Client")
def test_backend_span_exporter_no_items(mock_client):
    exporter = BackendSpanExporter(api_key="test_key")
    exporter.export([])
    # No calls should be made if there are no items
    mock_client.return_value.post.assert_not_called()
    exporter.close()


@patch("httpx.Client")
def test_backend_span_exporter_no_api_key(mock_client):
    # Ensure that os.environ is empty (sometimes devs have the openai api key set in their env)

    with patch.dict(os.environ, {}, clear=True):
        exporter = BackendSpanExporter(api_key=None)
        exporter.export([get_span(mock_processor())])

        # Should log an error and return without calling post
        mock_client.return_value.post.assert_not_called()
        exporter.close()


@patch("httpx.Client")
def test_backend_span_exporter_2xx_success(mock_client):
    mock_response = MagicMock()
    mock_response.status_code = 200
    mock_client.return_value.post.return_value = mock_response

    exporter = BackendSpanExporter(api_key="test_key")
    exporter.export([get_span(mock_processor()), get_trace(mock_processor())])

    # Should have called post exactly once
    mock_client.return_value.post.assert_called_once()
    exporter.close()


@patch("httpx.Client")
def test_backend_span_exporter_4xx_client_error(mock_client):
    mock_response = MagicMock()
    mock_response.status_code = 400
    mock_response.text = "Bad Request"
    mock_client.return_value.post.return_value = mock_response

    exporter = BackendSpanExporter(api_key="test_key")
    exporter.export([get_span(mock_processor())])

    # 4xx should not be retried
    mock_client.return_value.post.assert_called_once()
    exporter.close()


@patch("httpx.Client")
def test_backend_span_exporter_5xx_retry(mock_client, patched_time_sleep):
    mock_response = MagicMock()
    mock_response.status_code = 500

    # Make post() return 500 every time
    mock_client.return_value.post.return_value = mock_response

    exporter = BackendSpanExporter(api_key="test_key", max_retries=3, base_delay=0.1, max_delay=0.2)
    exporter.export([get_span(mock_processor())])

    # Should retry up to max_retries times
    assert mock_client.return_value.post.call_count == 3

    exporter.close()


@patch("httpx.Client")
def test_backend_span_exporter_request_error(mock_client, patched_time_sleep):
    # Make post() raise a RequestError each time
    mock_client.return_value.post.side_effect = httpx.RequestError("Network error")

    exporter = BackendSpanExporter(api_key="test_key", max_retries=2, base_delay=0.1, max_delay=0.2)
    exporter.export([get_span(mock_processor())])

    # Should retry up to max_retries times
    assert mock_client.return_value.post.call_count == 2

    exporter.close()


@patch("httpx.Client")
def test_backend_span_exporter_close(mock_client):
    exporter = BackendSpanExporter(api_key="test_key")
    exporter.close()

    # Ensure underlying http client is closed
    mock_client.return_value.close.assert_called_once()


@patch("httpx.Client")
def test_backend_span_exporter_sanitizes_generation_usage_for_openai_tracing(mock_client):
    """Unsupported usage keys should be stripped before POSTing to OpenAI tracing."""

    class DummyItem:
        tracing_api_key = None

        def __init__(self):
            self.exported_payload: dict[str, Any] = {
                "object": "trace.span",
                "span_data": {
                    "type": "generation",
                    "usage": {
                        "requests": 1,
                        "input_tokens": 10,
                        "output_tokens": 5,
                        "total_tokens": 15,
                        "input_tokens_details": {"cached_tokens": 1},
                        "output_tokens_details": {"reasoning_tokens": 2},
                    },
                },
            }

        def export(self):
            return self.exported_payload

    mock_response = MagicMock()
    mock_response.status_code = 200
    mock_client.return_value.post.return_value = mock_response

    exporter = BackendSpanExporter(api_key="test_key")
    item = DummyItem()
    exporter.export([cast(Any, item)])

    sent_payload = mock_client.return_value.post.call_args.kwargs["json"]["data"][0]
    sent_usage = sent_payload["span_data"]["usage"]
    assert "requests" not in sent_usage
    assert "total_tokens" not in sent_usage
    assert "input_tokens_details" not in sent_usage
    assert "output_tokens_details" not in sent_usage
    assert sent_usage["input_tokens"] == 10
    assert sent_usage["output_tokens"] == 5
    assert sent_usage["details"] == {
        "requests": 1,
        "total_tokens": 15,
        "input_tokens_details": {"cached_tokens": 1},
        "output_tokens_details": {"reasoning_tokens": 2},
    }

    # Ensure the original exported object has not been mutated.
    assert "requests" in item.exported_payload["span_data"]["usage"]
    assert item.exported_payload["span_data"]["usage"]["total_tokens"] == 15
    exporter.close()


@patch("httpx.Client")
def test_backend_span_exporter_truncates_large_input_for_openai_tracing(mock_client):
    class DummyItem:
        tracing_api_key = None

        def __init__(self):
            self.exported_payload: dict[str, Any] = {
                "object": "trace.span",
                "span_data": {
                    "type": "generation",
                    "input": "x" * (BackendSpanExporter._OPENAI_TRACING_MAX_FIELD_BYTES + 5_000),
                },
            }

        def export(self):
            return self.exported_payload

    mock_response = MagicMock()
    mock_response.status_code = 200
    mock_client.return_value.post.return_value = mock_response

    exporter = BackendSpanExporter(api_key="test_key")
    item = DummyItem()
    exporter.export([cast(Any, item)])

    sent_payload = mock_client.return_value.post.call_args.kwargs["json"]["data"][0]
    sent_input = sent_payload["span_data"]["input"]
    assert isinstance(sent_input, str)
    assert sent_input.endswith(exporter._OPENAI_TRACING_STRING_TRUNCATION_SUFFIX)
    assert exporter._value_json_size_bytes(sent_input) <= exporter._OPENAI_TRACING_MAX_FIELD_BYTES
    assert item.exported_payload["span_data"]["input"] != sent_input
    exporter.close()


@patch("httpx.Client")
def test_backend_span_exporter_truncates_large_structured_input_without_stringifying(mock_client):
    class NoStringifyDict(dict[str, Any]):
        def __str__(self) -> str:
            raise AssertionError("__str__ should not be called for oversized non-string previews")

    class DummyItem:
        tracing_api_key = None

        def __init__(self):
            payload_input = NoStringifyDict(
                blob="x" * (BackendSpanExporter._OPENAI_TRACING_MAX_FIELD_BYTES + 5_000)
            )
            self.exported_payload: dict[str, Any] = {
                "object": "trace.span",
                "span_data": {
                    "type": "generation",
                    "input": payload_input,
                },
            }

        def export(self):
            return self.exported_payload

    mock_response = MagicMock()
    mock_response.status_code = 200
    mock_client.return_value.post.return_value = mock_response

    exporter = BackendSpanExporter(api_key="test_key")
    exporter.export([cast(Any, DummyItem())])

    sent_payload = mock_client.return_value.post.call_args.kwargs["json"]["data"][0]
    sent_input = sent_payload["span_data"]["input"]
    assert isinstance(sent_input, dict)
    assert isinstance(sent_input["blob"], str)
    assert sent_input["blob"].endswith(exporter._OPENAI_TRACING_STRING_TRUNCATION_SUFFIX)
    assert exporter._value_json_size_bytes(sent_input) <= exporter._OPENAI_TRACING_MAX_FIELD_BYTES
    exporter.close()


@patch("httpx.Client")
def test_backend_span_exporter_keeps_generation_usage_for_custom_endpoint(mock_client):
    class DummyItem:
        tracing_api_key = None

        def __init__(self):
            self.exported_payload = {
                "object": "trace.span",
                "span_data": {
                    "type": "generation",
                    "usage": {
                        "requests": 1,
                        "input_tokens": 10,
                        "output_tokens": 5,
                    },
                },
            }

        def export(self):
            return self.exported_payload

    mock_response = MagicMock()
    mock_response.status_code = 200
    mock_client.return_value.post.return_value = mock_response

    exporter = BackendSpanExporter(
        api_key="test_key",
        endpoint="https://example.com/v1/traces/ingest",
    )
    exporter.export([cast(Any, DummyItem())])

    sent_payload = mock_client.return_value.post.call_args.kwargs["json"]["data"][0]
    assert sent_payload["span_data"]["usage"]["requests"] == 1
    assert sent_payload["span_data"]["usage"]["input_tokens"] == 10
    assert sent_payload["span_data"]["usage"]["output_tokens"] == 5
    exporter.close()


@patch("httpx.Client")
def test_backend_span_exporter_does_not_modify_non_generation_usage(mock_client):
    class DummyItem:
        tracing_api_key = None

        def export(self):
            return {
                "object": "trace.span",
                "span_data": {
                    "type": "function",
                    "usage": {"requests": 1},
                },
            }

    mock_response = MagicMock()
    mock_response.status_code = 200
    mock_client.return_value.post.return_value = mock_response

    exporter = BackendSpanExporter(api_key="test_key")
    exporter.export([cast(Any, DummyItem())])

    sent_payload = mock_client.return_value.post.call_args.kwargs["json"]["data"][0]
    assert sent_payload["span_data"]["usage"] == {"requests": 1}
    exporter.close()


def test_sanitize_for_openai_tracing_api_keeps_allowed_generation_usage():
    exporter = BackendSpanExporter(api_key="test_key")
    payload = {
        "object": "trace.span",
        "span_data": {
            "type": "generation",
            "usage": {
                "input_tokens": 1,
                "output_tokens": 2,
            },
        },
    }
    assert exporter._sanitize_for_openai_tracing_api(payload) is payload
    exporter.close()


@patch("httpx.Client")
def test_backend_span_exporter_keeps_large_input_for_custom_endpoint(mock_client):
    class DummyItem:
        tracing_api_key = None

        def __init__(self):
            self.exported_payload: dict[str, Any] = {
                "object": "trace.span",
                "span_data": {
                    "type": "generation",
                    "input": "x" * (BackendSpanExporter._OPENAI_TRACING_MAX_FIELD_BYTES + 5_000),
                },
            }

        def export(self):
            return self.exported_payload

    mock_response = MagicMock()
    mock_response.status_code = 200
    mock_client.return_value.post.return_value = mock_response

    exporter = BackendSpanExporter(
        api_key="test_key",
        endpoint="https://example.com/v1/traces/ingest",
    )
    item = DummyItem()
    exporter.export([cast(Any, item)])

    sent_payload: dict[str, Any] = mock_client.return_value.post.call_args.kwargs["json"]["data"][0]
    assert sent_payload["span_data"]["input"] == item.exported_payload["span_data"]["input"]
    exporter.close()


def test_sanitize_for_openai_tracing_api_moves_unsupported_generation_usage_to_details():
    exporter = BackendSpanExporter(api_key="test_key")
    payload = {
        "object": "trace.span",
        "span_data": {
            "type": "generation",
            "usage": {
                "input_tokens": 1,
                "output_tokens": 2,
                "total_tokens": 3,
                "input_tokens_details": {"cached_tokens": 0},
                "output_tokens_details": {"reasoning_tokens": 0},
                "details": {"provider": "litellm"},
            },
        },
    }
    sanitized = exporter._sanitize_for_openai_tracing_api(payload)
    assert sanitized["span_data"]["usage"] == {
        "input_tokens": 1,
        "output_tokens": 2,
        "details": {
            "provider": "litellm",
            "total_tokens": 3,
            "input_tokens_details": {"cached_tokens": 0},
            "output_tokens_details": {"reasoning_tokens": 0},
        },
    }
    exporter.close()


def test_sanitize_for_openai_tracing_api_filters_non_json_values_in_usage_details():
    exporter = BackendSpanExporter(api_key="test_key")
    non_json = object()
    payload = {
        "object": "trace.span",
        "span_data": {
            "type": "generation",
            "usage": {
                "input_tokens": 1,
                "output_tokens": 2,
                "input_tokens_details": {
                    "cached_tokens": 0,
                    "bad": non_json,
                },
                "output_tokens_details": {"reasoning_tokens": 0},
                "provider_usage": [1, non_json, {"ok": True, "bad": non_json}],
                "details": {
                    "provider": "litellm",
                    "bad": non_json,
                    "nested": {"keep": 1, "bad": non_json},
                },
            },
        },
    }
    sanitized = exporter._sanitize_for_openai_tracing_api(payload)
    assert sanitized["span_data"]["usage"] == {
        "input_tokens": 1,
        "output_tokens": 2,
        "details": {
            "provider": "litellm",
            "nested": {"keep": 1},
            "input_tokens_details": {"cached_tokens": 0},
            "output_tokens_details": {"reasoning_tokens": 0},
            "provider_usage": [1, {"ok": True}],
        },
    }
    exporter.close()


def test_sanitize_for_openai_tracing_api_handles_cyclic_usage_values():
    exporter = BackendSpanExporter(api_key="test_key")
    cyclic_dict: dict[str, Any] = {}
    cyclic_dict["self"] = cyclic_dict
    cyclic_list: list[Any] = []
    cyclic_list.append(cyclic_list)

    payload = {
        "object": "trace.span",
        "span_data": {
            "type": "generation",
            "usage": {
                "input_tokens": 1,
                "output_tokens": 2,
                "input_tokens_details": cyclic_dict,
                "details": {
                    "provider": "litellm",
                    "cycle": cyclic_list,
                },
            },
        },
    }

    sanitized = exporter._sanitize_for_openai_tracing_api(payload)
    assert sanitized["span_data"]["usage"] == {
        "input_tokens": 1,
        "output_tokens": 2,
        "details": {
            "provider": "litellm",
            "cycle": [],
            "input_tokens_details": {},
        },
    }
    exporter.close()


def test_sanitize_for_openai_tracing_api_drops_non_dict_generation_usage_details():
    exporter = BackendSpanExporter(api_key="test_key")
    payload = {
        "object": "trace.span",
        "span_data": {
            "type": "generation",
            "usage": {
                "input_tokens": 1,
                "output_tokens": 2,
                "details": "invalid",
            },
        },
    }
    sanitized = exporter._sanitize_for_openai_tracing_api(payload)
    assert sanitized["span_data"]["usage"] == {
        "input_tokens": 1,
        "output_tokens": 2,
    }
    exporter.close()


def test_sanitize_for_openai_tracing_api_drops_generation_usage_missing_required_tokens():
    exporter = BackendSpanExporter(api_key="test_key")
    payload = {
        "object": "trace.span",
        "span_data": {
            "type": "generation",
            "usage": {
                "input_tokens": 1,
                "total_tokens": 3,
                "input_tokens_details": {"cached_tokens": 0},
                "output_tokens_details": {"reasoning_tokens": 0},
            },
        },
    }
    sanitized = exporter._sanitize_for_openai_tracing_api(payload)
    assert sanitized["span_data"] == {
        "type": "generation",
    }
    exporter.close()


def test_sanitize_for_openai_tracing_api_rejects_boolean_token_counts():
    exporter = BackendSpanExporter(api_key="test_key")
    payload = {
        "object": "trace.span",
        "span_data": {
            "type": "generation",
            "usage": {
                "input_tokens": True,
                "output_tokens": False,
                "input_tokens_details": {"cached_tokens": 0},
                "output_tokens_details": {"reasoning_tokens": 0},
            },
        },
    }
    sanitized = exporter._sanitize_for_openai_tracing_api(payload)
    assert sanitized["span_data"] == {
        "type": "generation",
    }
    exporter.close()


def test_sanitize_for_openai_tracing_api_skips_non_dict_generation_usage():
    exporter = BackendSpanExporter(api_key="test_key")
    payload = {
        "object": "trace.span",
        "span_data": {
            "type": "generation",
            "usage": None,
        },
    }
    assert exporter._sanitize_for_openai_tracing_api(payload) is payload
    exporter.close()


def test_sanitize_for_openai_tracing_api_keeps_small_input_without_mutation():
    exporter = BackendSpanExporter(api_key="test_key")
    payload = {
        "object": "trace.span",
        "span_data": {
            "type": "generation",
            "input": "short input",
            "usage": {"input_tokens": 1, "output_tokens": 2},
        },
    }

    assert exporter._sanitize_for_openai_tracing_api(payload) is payload
    exporter.close()


def test_sanitize_for_openai_tracing_api_truncates_oversized_output():
    exporter = BackendSpanExporter(api_key="test_key")
    payload: dict[str, Any] = {
        "object": "trace.span",
        "span_data": {
            "type": "function",
            "output": "x" * (BackendSpanExporter._OPENAI_TRACING_MAX_FIELD_BYTES + 5_000),
        },
    }

    sanitized = exporter._sanitize_for_openai_tracing_api(payload)
    assert sanitized is not payload
    assert sanitized["span_data"]["output"].endswith(
        exporter._OPENAI_TRACING_STRING_TRUNCATION_SUFFIX
    )
    assert (
        exporter._value_json_size_bytes(sanitized["span_data"]["output"])
        <= exporter._OPENAI_TRACING_MAX_FIELD_BYTES
    )
    assert payload["span_data"]["output"] != sanitized["span_data"]["output"]
    exporter.close()


def test_sanitize_for_openai_tracing_api_preserves_generation_input_list_shape():
    exporter = BackendSpanExporter(api_key="test_key")
    payload = {
        "object": "trace.span",
        "span_data": {
            "type": "generation",
            "input": [
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "input_audio",
                            "input_audio": {
                                "data": "x"
                                * (BackendSpanExporter._OPENAI_TRACING_MAX_FIELD_BYTES + 5_000),
                                "format": "wav",
                            },
                        }
                    ],
                }
            ],
            "usage": {"input_tokens": 1, "output_tokens": 1},
        },
    }

    sanitized = exporter._sanitize_for_openai_tracing_api(payload)
    sanitized_input = sanitized["span_data"]["input"]
    assert isinstance(sanitized_input, list)
    assert isinstance(sanitized_input[0], dict)
    assert sanitized_input[0]["role"] == "user"
    assert (
        exporter._value_json_size_bytes(sanitized_input) <= exporter._OPENAI_TRACING_MAX_FIELD_BYTES
    )
    exporter.close()


def test_sanitize_for_openai_tracing_api_replaces_unserializable_output():
    exporter = BackendSpanExporter(api_key="test_key")
    payload: dict[str, Any] = {
        "object": "trace.span",
        "span_data": {
            "type": "function",
            "output": b"x" * 10,
        },
    }

    sanitized = exporter._sanitize_for_openai_tracing_api(payload)
    assert sanitized["span_data"]["output"] == {
        "truncated": True,
        "original_type": "bytes",
        "preview": "<bytes bytes=10 truncated>",
    }
    exporter.close()


def test_truncate_json_value_for_limit_terminates_preview_dict_under_zero_budget():
    exporter = BackendSpanExporter(api_key="test_key")
    preview = exporter._truncated_preview(None)

    truncated = exporter._truncate_json_value_for_limit(preview, 0)

    assert truncated == {}
    exporter.close()


def test_sanitize_for_openai_tracing_api_handles_none_content_under_tight_budget():
    exporter = BackendSpanExporter(api_key="test_key")
    payload: dict[str, Any] = {
        "object": "trace.span",
        "span_data": {
            "type": "generation",
            "output": [
                {
                    "role": "assistant",
                    "content": None,
                    "name": "a" * 25_000,
                    "tool_calls": [],
                }
                for _ in range(8)
            ],
            "usage": {"input_tokens": 1, "output_tokens": 1},
        },
    }

    sanitized = exporter._sanitize_for_openai_tracing_api(payload)
    sanitized_output = cast(list[Any], sanitized["span_data"]["output"])

    assert isinstance(sanitized_output, list)
    assert sanitized_output != payload["span_data"]["output"]
    assert (
        exporter._value_json_size_bytes(sanitized_output)
        <= exporter._OPENAI_TRACING_MAX_FIELD_BYTES
    )
    assert any(item == {} for item in sanitized_output)
    exporter.close()


def test_truncate_string_for_json_limit_returns_original_when_within_limit():
    exporter = BackendSpanExporter(api_key="test_key")
    value = "hello"
    max_bytes = exporter._value_json_size_bytes(value)

    assert exporter._truncate_string_for_json_limit(value, max_bytes) == value
    exporter.close()


def test_truncate_string_for_json_limit_returns_suffix_when_limit_equals_suffix():
    exporter = BackendSpanExporter(api_key="test_key")
    max_bytes = exporter._value_json_size_bytes(exporter._OPENAI_TRACING_STRING_TRUNCATION_SUFFIX)

    assert (
        exporter._truncate_string_for_json_limit("x" * 100, max_bytes)
        == exporter._OPENAI_TRACING_STRING_TRUNCATION_SUFFIX
    )
    exporter.close()


def test_truncate_string_for_json_limit_returns_empty_when_suffix_too_large():
    exporter = BackendSpanExporter(api_key="test_key")
    max_bytes = (
        exporter._value_json_size_bytes(exporter._OPENAI_TRACING_STRING_TRUNCATION_SUFFIX) - 1
    )

    assert exporter._truncate_string_for_json_limit("x" * 100, max_bytes) == ""
    exporter.close()


def test_truncate_string_for_json_limit_handles_escape_heavy_input():
    exporter = BackendSpanExporter(api_key="test_key")
    value = ('\\"' * 40_000) + "tail"
    max_bytes = exporter._OPENAI_TRACING_MAX_FIELD_BYTES

    truncated = exporter._truncate_string_for_json_limit(value, max_bytes)

    assert truncated.endswith(exporter._OPENAI_TRACING_STRING_TRUNCATION_SUFFIX)
    assert exporter._value_json_size_bytes(truncated) <= max_bytes
    exporter.close()