import logging import os import subprocess import sys import textwrap import threading import time from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer 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_shutdown_timeout_returns_when_exporter_blocks( caplog: pytest.LogCaptureFixture, ) -> None: export_started = threading.Event() release_export = threading.Event() class BlockingExporter(TracingExporter): def export(self, items: list[Trace | Span[Any]]) -> None: export_started.set() release_export.wait(timeout=5.0) processor = BatchTraceProcessor( exporter=BlockingExporter(), max_queue_size=1, schedule_delay=60.0, export_trigger_ratio=1.0, ) processor.on_span_end(get_span(processor)) assert export_started.wait(timeout=2.0) start = time.monotonic() with caplog.at_level(logging.WARNING): processor.shutdown(timeout=0.05) elapsed = time.monotonic() - start assert elapsed < 0.5 assert "shutdown timeout reached" in caplog.text release_export.set() if processor._worker_thread: processor._worker_thread.join(timeout=2.0) def test_batch_trace_processor_shutdown_passes_deadline_to_exporter() -> None: seen_deadlines: list[float | None] = [] class DeadlineExporter(TracingExporter): def export(self, items: list[Trace | Span[Any]]) -> None: raise AssertionError("shutdown should use the deadline-aware exporter path") def _export_with_deadline( self, items: list[Trace | Span[Any]], deadline: float | None ) -> None: seen_deadlines.append(deadline) processor = BatchTraceProcessor(exporter=DeadlineExporter()) processor._queue.put_nowait(get_span(processor)) processor.shutdown(timeout=1.0) assert len(seen_deadlines) == 1 assert seen_deadlines[0] is not None def test_batch_trace_processor_survives_exporter_exception(): """A failing exporter must not kill the background worker thread. Previously, an exception raised inside ``exporter.export`` propagated out of ``_export_batches`` and killed the ``_run`` thread, causing all subsequent spans to silently accumulate in the queue until it filled up. """ class FlakyExporter(TracingExporter): def __init__(self) -> None: self.call_count = 0 self.exported: list[Trace | Span[Any]] = [] def export(self, items: list[Trace | Span[Any]]) -> None: self.call_count += 1 if self.call_count == 1: raise RuntimeError("simulated exporter failure") self.exported.extend(items) exporter = FlakyExporter() processor = BatchTraceProcessor(exporter, schedule_delay=0.05, max_batch_size=1) processor.on_span_end(get_span(processor)) processor.on_span_end(get_span(processor)) processor.on_span_end(get_span(processor)) # Give the worker time to encounter the failure and continue processing. time.sleep(0.3) assert processor._worker_thread is not None assert processor._worker_thread.is_alive(), "Worker thread must survive an exporter exception" processor.shutdown(timeout=2.0) # First batch raised; the remaining two items must still have been exported. assert len(exporter.exported) == 2 assert exporter.call_count >= 3 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() 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): 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) with patch.object(exporter._shutdown_event, "wait", return_value=False) as wait_for_retry: exporter.export([get_span(mock_processor())]) # Should retry up to max_retries times assert mock_client.return_value.post.call_count == 3 assert wait_for_retry.call_count == 2 exporter.close() @patch("httpx.Client") def test_backend_span_exporter_deadline_stops_during_5xx_retry_backoff(mock_client): mock_response = MagicMock() mock_response.status_code = 504 mock_client.return_value.post.return_value = mock_response exporter = BackendSpanExporter(api_key="test_key", max_retries=3, base_delay=1.0) with patch("time.sleep") as sleep_for_retry: exporter._export_with_deadline( [get_span(mock_processor())], deadline=time.monotonic() + 0.01 ) assert mock_client.return_value.post.call_count == 1 sleep_for_retry.assert_called_once() assert sleep_for_retry.call_args.args[0] <= 0.1 exporter.close() @patch("httpx.Client") def test_batch_trace_processor_shutdown_interrupts_exporter_retry_backoff(mock_client): post_called = threading.Event() mock_response = MagicMock() mock_response.status_code = 504 def post(**kwargs: Any) -> Any: post_called.set() return mock_response mock_client.return_value.post.side_effect = post exporter = BackendSpanExporter( api_key="test_key", max_retries=100, base_delay=10.0, max_delay=10.0, ) processor = BatchTraceProcessor( exporter=exporter, max_queue_size=1, max_batch_size=1, schedule_delay=60.0, export_trigger_ratio=1.0, ) processor.on_span_end(get_span(processor)) assert post_called.wait(timeout=2.0) start = time.monotonic() processor.shutdown(timeout=1.0) elapsed = time.monotonic() - start assert elapsed < 0.5 assert processor._worker_thread is not None assert not processor._worker_thread.is_alive() assert mock_client.return_value.post.call_count == 1 exporter.close() @patch("httpx.Client") def test_batch_trace_processor_shutdown_without_timeout_preserves_export_retries(mock_client): mock_response = MagicMock() mock_response.status_code = 504 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, ) processor = BatchTraceProcessor(exporter=exporter) processor._queue.put_nowait(get_span(processor)) with patch.object(exporter._shutdown_event, "wait", return_value=False) as wait_for_retry: processor.shutdown(timeout=None) assert mock_client.return_value.post.call_count == 3 assert wait_for_retry.call_count == 2 exporter.close() @pytest.mark.serial def test_tracing_atexit_cleanup_timeout_preserves_process_exit_code_on_504() -> None: request_seen = threading.Event() class Always504Handler(BaseHTTPRequestHandler): def do_POST(self) -> None: request_seen.set() self.send_response(504) self.end_headers() self.wfile.write(b"gateway timeout") def log_message(self, format: str, *args: Any) -> None: return server = ThreadingHTTPServer(("127.0.0.1", 0), Always504Handler) server_thread = threading.Thread(target=server.serve_forever, daemon=True) server_thread.start() script = textwrap.dedent( f""" import sys import time from agents.tracing import custom_span, trace from agents.tracing.processors import BackendSpanExporter, BatchTraceProcessor from agents.tracing.provider import DefaultTraceProvider from agents.tracing import setup as tracing_setup tracing_setup._DEFAULT_SHUTDOWN_TIMEOUT = 0.2 exporter = BackendSpanExporter( api_key="test_key", endpoint="http://127.0.0.1:{server.server_port}/traces/ingest", max_retries=100, base_delay=10.0, max_delay=10.0, ) processor = BatchTraceProcessor( exporter=exporter, max_queue_size=1, max_batch_size=1, schedule_delay=60.0, export_trigger_ratio=1.0, ) provider = DefaultTraceProvider() provider.register_processor(processor) original_shutdown = provider.shutdown def timed_shutdown(*args, **kwargs): shutdown_started = time.monotonic() try: return original_shutdown(*args, **kwargs) finally: print( f"shutdown_elapsed={{time.monotonic() - shutdown_started:.6f}}", flush=True, ) provider.shutdown = timed_shutdown tracing_setup.set_trace_provider(provider) with trace("probe"): with custom_span("probe-span"): pass time.sleep(0.3) sys.exit(7) """ ) try: result = subprocess.run( [sys.executable, "-c", script], check=False, capture_output=True, text=True, timeout=10.0, ) finally: server.shutdown() server.server_close() assert request_seen.is_set() assert result.returncode == 7 shutdown_elapsed_prefix = "shutdown_elapsed=" shutdown_elapsed_lines = [ line for line in result.stdout.splitlines() if line.startswith(shutdown_elapsed_prefix) ] assert len(shutdown_elapsed_lines) == 1 assert float(shutdown_elapsed_lines[0][len(shutdown_elapsed_prefix) :]) < 0.5 @patch("httpx.Client") def test_backend_span_exporter_request_error(mock_client): # 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) with patch.object(exporter._shutdown_event, "wait", return_value=False) as wait_for_retry: exporter.export([get_span(mock_processor())]) # Should retry up to max_retries times assert mock_client.return_value.post.call_count == 2 wait_for_retry.assert_called_once() 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_drops_non_generation_usage_for_openai_endpoint(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 "usage" not in sent_payload["span_data"] exporter.close() @patch("httpx.Client") def test_backend_span_exporter_keeps_non_generation_usage_for_custom_endpoint(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", 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} 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": "", } 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()