From b37a71a114934a4461113170912c55da9c035c90 Mon Sep 17 00:00:00 2001 From: Kazuhiro Sera Date: Tue, 25 Aug 2026 09:08:00 +0900 Subject: [PATCH] fix: preserve response IDs in redacted traces --- src/agents/models/openai_responses.py | 30 ++++ src/agents/tracing/span_data.py | 5 +- tests/test_responses_tracing.py | 221 ++++++++++++++++++++++++-- 3 files changed, 243 insertions(+), 13 deletions(-) diff --git a/src/agents/models/openai_responses.py b/src/agents/models/openai_responses.py index 75987c175d..87029d83e5 100644 --- a/src/agents/models/openai_responses.py +++ b/src/agents/models/openai_responses.py @@ -502,6 +502,9 @@ def __init__( def _non_null_or_omit(self, value: Any) -> Any: return value if value is not None else omit + def _uses_official_openai_endpoint(self) -> bool: + return is_official_openai_client(self._get_client()) + def _supports_default_prompt_cache_key(self) -> bool: return is_official_openai_client(self._get_client()) @@ -572,6 +575,11 @@ async def get_response( ) -> ModelResponse: with response_span(disabled=tracing.is_disabled()) as span_response: try: + redacted_response_id_endpoint_is_trusted = ( + not tracing.include_data() + and not tracing.is_disabled() + and self._uses_official_openai_endpoint() + ) response = await self._fetch_response( system_instructions, input, @@ -604,6 +612,11 @@ async def get_response( if tracing.include_data(): span_response.span_data.response = response span_response.span_data.input = input + elif ( + redacted_response_id_endpoint_is_trusted + and self._uses_official_openai_endpoint() + ): + span_response.span_data._response_id = response.id except asyncio.CancelledError: record_current_task_model_timeout_on_span( span_response, @@ -658,6 +671,11 @@ async def stream_response( """ with response_span(disabled=tracing.is_disabled()) as span_response: try: + redacted_response_id_endpoint_is_trusted = ( + not tracing.include_data() + and not tracing.is_disabled() + and self._uses_official_openai_endpoint() + ) stream = await self._fetch_response( system_instructions, input, @@ -680,6 +698,11 @@ async def stream_response( chunk_type = getattr(chunk, "type", None) if isinstance(chunk, ResponseCompletedEvent): final_response = chunk.response + if ( + redacted_response_id_endpoint_is_trusted + and self._uses_official_openai_endpoint() + ): + span_response.span_data._response_id = chunk.response.id if model_settings.preserve_raw_usage is True: _attach_raw_usage_snapshot(chunk.response, chunk.response.usage) usage = _usage_from_response(chunk.response) @@ -1111,6 +1134,13 @@ def __init__( ) self._ws_client_close_generation = 0 + def _uses_official_openai_endpoint(self) -> bool: + base_url = prepare_openai_client_websocket_base_url( + self._client, + context="Responses websocket", + ) + return is_official_openai_base_url(base_url, websocket=True) + def _supports_default_prompt_cache_key(self) -> bool: if self._client.websocket_base_url is not None: return is_official_openai_base_url(self._client.websocket_base_url, websocket=True) diff --git a/src/agents/tracing/span_data.py b/src/agents/tracing/span_data.py index 872388a736..57b7fe6226 100644 --- a/src/agents/tracing/span_data.py +++ b/src/agents/tracing/span_data.py @@ -215,7 +215,7 @@ class ResponseSpanData(SpanData): Includes response and input. """ - __slots__ = ("response", "input", "usage") + __slots__ = ("response", "input", "usage", "_response_id") def __init__( self, @@ -228,6 +228,7 @@ def __init__( # processor implementations self.input = input self.usage = usage + self._response_id: str | None = None @property def type(self) -> str: @@ -236,7 +237,7 @@ def type(self) -> str: def export(self) -> dict[str, Any]: return { "type": self.type, - "response_id": self.response.id if self.response is not None else None, + "response_id": (self.response.id if self.response is not None else self._response_id), "usage": self.usage, } diff --git a/tests/test_responses_tracing.py b/tests/test_responses_tracing.py index 71124d047d..c13253238a 100644 --- a/tests/test_responses_tracing.py +++ b/tests/test_responses_tracing.py @@ -1,10 +1,19 @@ +from typing import Any, cast + import pytest from inline_snapshot import snapshot from openai import AsyncOpenAI from openai.types.responses import ResponseCompletedEvent from openai.types.responses.response_usage import InputTokensDetails, OutputTokensDetails -from agents import ModelBehaviorError, ModelSettings, ModelTracing, OpenAIResponsesModel, trace +from agents import ( + ModelBehaviorError, + ModelSettings, + ModelTracing, + OpenAIResponsesModel, + OpenAIResponsesWSModel, + trace, +) from agents.tracing.span_data import ResponseSpanData from tests import model_test_helpers @@ -120,10 +129,13 @@ async def dummy_fetch_response( @pytest.mark.allow_call_model_methods @pytest.mark.asyncio -async def test_non_data_tracing_doesnt_set_response_id(monkeypatch): +async def test_non_data_tracing_preserves_response_id_without_response(monkeypatch): with trace(workflow_name="test"): # Create an instance of the model - model = OpenAIResponsesModel(model="test-model", openai_client=AsyncOpenAI(api_key="test")) + model = OpenAIResponsesModel( + model="test-model", + openai_client=AsyncOpenAI(api_key="test", base_url="https://api.openai.com/v1"), + ) # Mock _fetch_response to return a dummy response with a known id async def dummy_fetch_response( @@ -162,6 +174,7 @@ async def dummy_fetch_response( { "type": "response", "data": { + "response_id": "dummy-id", "usage": { "requests": 1, "input_tokens": 1, @@ -172,7 +185,7 @@ async def dummy_fetch_response( "cache_write_tokens": 0, }, "output_tokens_details": {"reasoning_tokens": 0}, - } + }, }, } ], @@ -182,6 +195,57 @@ async def dummy_fetch_response( [span] = fetch_ordered_spans() assert span.span_data.response is None + assert span.span_data.input is None + + +@pytest.mark.allow_call_model_methods +@pytest.mark.asyncio +async def test_non_data_tracing_omits_custom_endpoint_response_id(monkeypatch): + provider_response_id = "tenant-customer-123" + with trace(workflow_name="test"): + client = AsyncOpenAI(api_key="test", base_url="https://provider.example.test/v1") + model = OpenAIResponsesModel( + model="test-model", + openai_client=client, + ) + + async def dummy_fetch_response( + system_instructions, + input, + model_settings, + tools, + output_schema, + handoffs, + previous_response_id, + conversation_id, + stream, + prompt, + ): + response = DummyResponse() + response.id = provider_response_id + client.base_url = "https://api.openai.com/v1" + return response + + monkeypatch.setattr(model, "_fetch_response", dummy_fetch_response) + + model_response = await model.get_response( + "instr", + "input", + ModelSettings(), + [], + None, + [], + ModelTracing.ENABLED_WITHOUT_DATA, + previous_response_id=None, + ) + + assert model_response.response_id == provider_response_id + [span] = fetch_ordered_spans() + assert isinstance(span.span_data, ResponseSpanData) + assert span.span_data.export()["response_id"] is None + assert span.span_data.response is None + assert span.span_data.input is None + assert span.span_data.usage is not None @pytest.mark.allow_call_model_methods @@ -369,10 +433,16 @@ async def __aiter__(self): @pytest.mark.allow_call_model_methods @pytest.mark.asyncio -async def test_stream_non_data_tracing_doesnt_set_response_id(monkeypatch): +@pytest.mark.parametrize("close_at_completed", [False, True]) +async def test_stream_non_data_tracing_preserves_response_id_without_response( + monkeypatch, close_at_completed: bool +): with trace(workflow_name="test"): # Create an instance of the model - model = OpenAIResponsesModel(model="test-model", openai_client=AsyncOpenAI(api_key="test")) + model = OpenAIResponsesModel( + model="test-model", + openai_client=AsyncOpenAI(api_key="test", base_url="https://api.openai.com/v1"), + ) # Define a dummy fetch function that returns an async stream with a dummy response async def dummy_fetch_response( @@ -399,8 +469,7 @@ async def __aiter__(self): monkeypatch.setattr(model, "_fetch_response", dummy_fetch_response) - # Consume the stream to trigger processing of the final response - async for _ in model.stream_response( + stream = model.stream_response( "instr", "input", ModelSettings(), @@ -409,8 +478,15 @@ async def __aiter__(self): [], ModelTracing.ENABLED_WITHOUT_DATA, previous_response_id=None, - ): - pass + ) + if close_at_completed: + stream_agen = cast(Any, stream) + event = await stream_agen.__anext__() + assert event.type == "response.completed" + await stream_agen.aclose() + else: + async for _ in stream: + pass assert fetch_normalized_spans() == snapshot( [ @@ -420,6 +496,7 @@ async def __aiter__(self): { "type": "response", "data": { + "response_id": "dummy-id-123", "usage": { "requests": 1, "input_tokens": 0, @@ -430,7 +507,7 @@ async def __aiter__(self): "cache_write_tokens": 0, }, "output_tokens_details": {"reasoning_tokens": 0}, - } + }, }, } ], @@ -441,6 +518,128 @@ async def __aiter__(self): [span] = fetch_ordered_spans() assert isinstance(span.span_data, ResponseSpanData) assert span.span_data.response is None + assert span.span_data.input is None + + +@pytest.mark.allow_call_model_methods +@pytest.mark.asyncio +async def test_stream_non_data_tracing_omits_custom_endpoint_response_id(monkeypatch): + provider_response_id = "tenant-customer-123" + with trace(workflow_name="test"): + client = AsyncOpenAI(api_key="test", base_url="https://provider.example.test/v1") + model = OpenAIResponsesModel( + model="test-model", + openai_client=client, + ) + + async def dummy_fetch_response( + system_instructions, + input, + model_settings, + tools, + output_schema, + handoffs, + previous_response_id, + conversation_id, + stream, + prompt, + ): + class DummyStream: + async def __aiter__(self): + client.base_url = "https://api.openai.com/v1" + yield ResponseCompletedEvent( + type="response.completed", + response=model_test_helpers.get_response_obj([], provider_response_id), + sequence_number=0, + ) + + return DummyStream() + + monkeypatch.setattr(model, "_fetch_response", dummy_fetch_response) + + events = [ + event + async for event in model.stream_response( + "instr", + "input", + ModelSettings(), + [], + None, + [], + ModelTracing.ENABLED_WITHOUT_DATA, + previous_response_id=None, + ) + ] + + assert isinstance(events[-1], ResponseCompletedEvent) + assert events[-1].response.id == provider_response_id + [span] = fetch_ordered_spans() + assert isinstance(span.span_data, ResponseSpanData) + assert span.span_data.export()["response_id"] is None + assert span.span_data.response is None + assert span.span_data.input is None + assert span.span_data.usage is not None + + +@pytest.mark.allow_call_model_methods +@pytest.mark.asyncio +async def test_stream_non_data_tracing_preserves_id_for_https_official_websocket(monkeypatch): + provider_response_id = "resp-official-ws" + with trace(workflow_name="test"): + model = OpenAIResponsesWSModel( + model="test-model", + openai_client=AsyncOpenAI( + api_key="test", websocket_base_url="https://api.openai.com/v1" + ), + ) + assert model._supports_default_prompt_cache_key() is False + + async def dummy_fetch_response( + system_instructions, + input, + model_settings, + tools, + output_schema, + handoffs, + previous_response_id, + conversation_id, + stream, + prompt, + ): + class DummyStream: + async def __aiter__(self): + yield ResponseCompletedEvent( + type="response.completed", + response=model_test_helpers.get_response_obj([], provider_response_id), + sequence_number=0, + ) + + return DummyStream() + + monkeypatch.setattr(model, "_fetch_response", dummy_fetch_response) + + events = [ + event + async for event in model.stream_response( + "instr", + "input", + ModelSettings(), + [], + None, + [], + ModelTracing.ENABLED_WITHOUT_DATA, + previous_response_id=None, + ) + ] + + assert isinstance(events[-1], ResponseCompletedEvent) + assert events[-1].response.id == provider_response_id + [span] = fetch_ordered_spans() + assert isinstance(span.span_data, ResponseSpanData) + assert span.span_data.export()["response_id"] == provider_response_id + assert span.span_data.response is None + assert span.span_data.input is None + assert span.span_data.usage is not None @pytest.mark.allow_call_model_methods