From 85dcdb76d8e8f47c24fbb96b178ba2e0e436b59e Mon Sep 17 00:00:00 2001 From: Nikhil Pareek Date: Sun, 1 Nov 2026 21:55:13 -0800 Subject: [PATCH 01/10] fix(openai): protect Responses privacy and finalize interrupted calls Strip MCP authentication from telemetry-only tool copies, redact known credentials from attributes and exception telemetry, and keep echoed request fields out of response output. Honor failed terminal results, finalize explicit and high-level stream closes, and propagate cancellation after ending spans. Mark interrupted Parallel Task workflows as errors. Co-Authored-By: Claude Code --- .github/workflows/parallel-python-tests.yml | 2 +- .../openai/tests/test_responses_safety.py | 472 ++++++++++++++++++ .../_attributes/_responses_api.py | 5 +- .../openai/traceai_openai/_privacy.py | 74 +++ .../openai/traceai_openai/_request.py | 26 +- .../traceai_openai/_response_accumulator.py | 35 +- .../openai/traceai_openai/_span_io_handler.py | 11 +- .../openai/traceai_openai/_stream.py | 82 ++- .../openai/traceai_openai/_with_span.py | 25 +- python/frameworks/parallel/README.md | 2 +- .../parallel/examples/task_workflow.py | 4 +- .../parallel/tests/test_task_workflow.py | 31 ++ 12 files changed, 740 insertions(+), 29 deletions(-) create mode 100644 python/frameworks/openai/tests/test_responses_safety.py create mode 100644 python/frameworks/openai/traceai_openai/_privacy.py diff --git a/.github/workflows/parallel-python-tests.yml b/.github/workflows/parallel-python-tests.yml index 553945dd8..6ef2de5c6 100644 --- a/.github/workflows/parallel-python-tests.yml +++ b/.github/workflows/parallel-python-tests.yml @@ -37,7 +37,7 @@ jobs: distribution: wheel - python: '3.13' parallel: '1.3.7' - openai: '3.42.0' + openai: '3.42.1' distribution: sdist env: FI_BASE_URL: http://127.0.0.1:9 diff --git a/python/frameworks/openai/tests/test_responses_safety.py b/python/frameworks/openai/tests/test_responses_safety.py new file mode 100644 index 000000000..f6a5be333 --- /dev/null +++ b/python/frameworks/openai/tests/test_responses_safety.py @@ -0,0 +1,472 @@ +"""Responses privacy and lifecycle contracts with real SDKs, no vendor traffic.""" +import asyncio +import json + +import httpx +import openai +import pytest +from fi_instrumentation import TraceConfig, register +from fi_instrumentation.fi_types import ProjectType +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import StatusCode +from traceai_openai import OpenAIInstrumentor + +KEY = "synthetic-parallel-vendor-key-DO-NOT-EXPORT" +QUESTION = "synthetic private question" + +@pytest.fixture +def tracing(): + exporter = InMemorySpanExporter() + provider = register(project_name="recall-offline", project_type=ProjectType.OBSERVE, + span_exporter=exporter, batch=False, verbose=False) + try: + yield provider, exporter + finally: + provider.force_flush() + provider.shutdown() + + +@pytest.fixture +def openai_tracing(tracing): + provider, exporter = tracing + instrumentor = OpenAIInstrumentor() + instrumentor.instrument(tracer_provider=provider) + try: + yield provider, exporter + finally: + instrumentor.uninstrument() + + +def sdk_client(handler, async_=False): + http_client = (httpx.AsyncClient if async_ else httpx.Client)( + transport=httpx.MockTransport(handler), trust_env=False) + return (openai.AsyncOpenAI if async_ else openai.OpenAI)( + api_key=KEY, base_url="https://api.parallel.ai/v1", max_retries=0, http_client=http_client) + + +def response_body(status="completed"): + return { + "id": "resp_synthetic", "object": "response", "created_at": 1, "model": "parallel", + "status": status, "parallel_tool_calls": True, "tool_choice": "auto", "tools": [], + "output": [], "error": {"code": "server_error", "message": "synthetic upstream failed"} if status == "failed" else None, + "usage": {"input_tokens": 5, "output_tokens": 2, "total_tokens": 7, + "input_tokens_details": {"cached_tokens": 0}, "output_tokens_details": {"reasoning_tokens": 0}}, + } + + +def sse_event(kind, response): + return ("event: " + kind + "\ndata: " + json.dumps({"type": kind, "sequence_number": 1, "response": response}) + "\n\n").encode() + + +@pytest.mark.parametrize("async_", [False, True]) +@pytest.mark.parametrize("stream", [False, True]) +def test_failed_response_must_not_export_success(openai_tracing, async_, stream): + _, exporter = openai_tracing + def handler(request): + if stream: + return httpx.Response(200, headers={"content-type": "text/event-stream"}, + content=sse_event("response.failed", response_body("failed"))) + return httpx.Response(200, json=response_body("failed")) + async def run(): + async with sdk_client(handler, True) as client: + result = await client.responses.create(model="parallel", input=QUESTION, stream=stream) + if stream: + events = [event async for event in result] + assert events[0].type == "response.failed" + else: + assert result.status == "failed" + if async_: + asyncio.run(run()) + else: + with sdk_client(handler) as client: + result = client.responses.create(model="parallel", input=QUESTION, stream=stream) + if stream: + events = list(result) # Exhaustion is part of the lifecycle contract. + assert events[0].type == "response.failed" + else: + assert result.status == "failed" + spans = exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].attributes["gen_ai.provider.name"] == "parallel" + assert spans[0].status.status_code is StatusCode.ERROR, spans[0].to_json() + + +@pytest.mark.parametrize("async_", [False, True]) +def test_early_stream_close_must_finish_span(openai_tracing, async_): + provider, exporter = openai_tracing + def handler(request): + return httpx.Response(200, headers={"content-type": "text/event-stream"}, + content=sse_event("response.created", response_body("in_progress")) + + sse_event("response.completed", response_body())) + async def run(): + async with sdk_client(handler, True) as client: + stream = await client.responses.create(model="parallel", input=QUESTION, stream=True) + assert (await stream.__anext__()).type == "response.created" + await stream.close() + await stream.close() + assert stream.response.is_closed + if async_: + asyncio.run(run()) + else: + with sdk_client(handler) as client: + stream = client.responses.create(model="parallel", input=QUESTION, stream=True) + assert next(stream).type == "response.created" + stream.close() + stream.close() + assert stream.response.is_closed + provider.force_flush() + spans = exporter.get_finished_spans() + assert len(spans) == 1, f"closed SDK stream exported {len(spans)} spans" + assert spans[0].status.status_code is StatusCode.ERROR + + +@pytest.mark.parametrize("stage", ["request", "stream"]) +def test_async_cancellation_must_finish_span(openai_tracing, stage): + provider, exporter = openai_tracing + async def run(): + waiting = asyncio.Event() + class WaitingStream(httpx.AsyncByteStream): + async def __aiter__(self): + yield sse_event("response.created", response_body("in_progress")) + waiting.set() + await asyncio.Event().wait() + async def handler(request): + if stage == "request": + waiting.set() + await asyncio.Event().wait() + return httpx.Response(200, headers={"content-type": "text/event-stream"}, stream=WaitingStream()) + async with sdk_client(handler, True) as client: + stream = None + async def call(): + nonlocal stream + stream = await client.responses.create(model="parallel", input=QUESTION, stream=stage == "stream") + if stage == "stream": + async for _ in stream: + pass + task = asyncio.create_task(call()) + await asyncio.wait_for(waiting.wait(), timeout=3) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + if stream is not None: + await stream.close() + asyncio.run(run()) + provider.force_flush() + spans = exporter.get_finished_spans() + assert len(spans) == 1, f"cancelled SDK {stage} exported {len(spans)} spans" + assert spans[0].status.status_code is StatusCode.ERROR + + +@pytest.mark.parametrize("async_", [False, True]) +def test_api_key_in_error_must_be_redacted(openai_tracing, async_): + _, exporter = openai_tracing + def handler(request): + assert request.headers["authorization"] == "Bearer " + KEY + return httpx.Response(401, json={"error": {"code": "invalid_api_key", "message": "Invalid key: " + KEY}}) + async def run(): + async with sdk_client(handler, True) as client: + await client.responses.create(model="parallel", input=QUESTION) + with pytest.raises(openai.AuthenticationError) as caught: + if async_: + asyncio.run(run()) + else: + with sdk_client(handler) as client: + client.responses.create(model="parallel", input=QUESTION) + spans = exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].status.status_code is StatusCode.ERROR + assert KEY in str(caught.value) # The caller still receives the original SDK exception. + assert KEY not in spans[0].to_json(), "vendor credential exported in error status/event" + + +@pytest.mark.parametrize("async_", [False, True]) +def test_hidden_instructions_not_reexported_in_output(tracing, async_): + provider, exporter = tracing + instruction = "synthetic confidential system instruction DO-NOT-EXPORT" + instrumentor = OpenAIInstrumentor() + instrumentor.instrument(tracer_provider=provider, config=TraceConfig(hide_inputs=True)) + def handler(request): + assert json.loads(request.content)["instructions"] == instruction + body = response_body() + body["instructions"] = instruction + body["output"] = [{"type": "message", "id": "msg_1", "role": "assistant", "status": "completed", + "content": [{"type": "output_text", "text": "safe answer", "annotations": []}]}] + return httpx.Response(200, json=body) + async def run(): + async with sdk_client(handler, True) as client: + response = await client.responses.create(model="parallel", input=QUESTION, instructions=instruction) + assert response.instructions == instruction + try: + if async_: + asyncio.run(run()) + else: + with sdk_client(handler) as client: + response = client.responses.create(model="parallel", input=QUESTION, instructions=instruction) + assert response.instructions == instruction + span, = exporter.get_finished_spans() + assert span.attributes["input.value"] == "__REDACTED__" + assert "safe answer" in span.attributes["output.value"] + assert instruction not in span.to_json(), "FI_HIDE_INPUTS instruction leaked via output.value" + finally: + instrumentor.uninstrument() + + +@pytest.mark.parametrize("async_", [False, True]) +def test_mcp_credentials_not_captured_as_tool_definitions(openai_tracing, async_): + _, exporter = openai_tracing + oauth = "synthetic-oauth-credential-DO-NOT-EXPORT" + apikey = "synthetic-mcp-api-key-DO-NOT-EXPORT" + tool = {"type": "mcp", "server_label": "synthetic", "server_url": "https://synthetic.invalid/mcp", + "require_approval": "never", "authorization": oauth, "headers": {"X-Api-Key": apikey}} + def handler(request): + parsed = json.loads(request.content) + assert parsed["tools"][0]["authorization"] == oauth + assert parsed["tools"][0]["headers"]["X-Api-Key"] == apikey + # Parallel explicitly scrubs tool credentials in its response. + body = response_body() + body["tools"] = [{**tool, "authorization": None, "headers": None}] + return httpx.Response(200, json=body) + async def run(): + async with sdk_client(handler, True) as client: + await client.responses.create(model="parallel", input=QUESTION, tools=[tool]) + if async_: + asyncio.run(run()) + else: + with sdk_client(handler) as client: + client.responses.create(model="parallel", input=QUESTION, tools=[tool]) + span, = exporter.get_finished_spans() + assert oauth not in span.to_json(), "MCP OAuth token exported in gen_ai.tool.definitions" + assert apikey not in span.to_json(), "MCP API key exported in gen_ai.tool.definitions" + + +@pytest.mark.parametrize("async_", [False, True]) +def test_sdk_responses_stream_context_exit_finishes_span(openai_tracing, async_): + provider, exporter = openai_tracing + def handler(request): + return httpx.Response(200, headers={"content-type": "text/event-stream"}, + content=sse_event("response.created", response_body("in_progress")) + + sse_event("response.completed", response_body())) + async def run(): + async with sdk_client(handler, True) as client, client.responses.stream( + model="parallel", input=QUESTION + ) as stream: + async for event in stream: + assert event.type == "response.created" + break + if async_: + asyncio.run(run()) + else: + with sdk_client(handler) as client, client.responses.stream( + model="parallel", input=QUESTION + ) as stream: + for event in stream: + assert event.type == "response.created" + break + provider.force_flush() + spans = exporter.get_finished_spans() + assert len(spans) == 1, f"SDK context exit exported {len(spans)} spans" + assert spans[0].status.status_code is StatusCode.ERROR + + + +@pytest.mark.parametrize("async_", [False, True]) +@pytest.mark.parametrize("mode", ["complete", "failed", "incomplete", "error", "eof", "empty"]) +def test_terminal_status_and_close_are_consistent(openai_tracing, async_, mode): + _, exporter = openai_tracing + status = {"complete": "completed", "failed": "failed", "incomplete": "incomplete"}.get(mode) + body = response_body(status or "in_progress") + if mode == "failed": + body["error"]["message"] = "upstream rejected " + KEY + payload = sse_event("response.created", response_body("in_progress")) + if status: + payload += sse_event("response." + status, body) + elif mode == "error": + payload += ('data: ' + json.dumps({"type": "error", "code": "server_error", + "message": "upstream rejected " + KEY, "param": None, + "sequence_number": 2}) + '\n\n').encode() + + def handler(request): + return httpx.Response(200, headers={"content-type": "text/event-stream"}, content=payload) + + async def run(): + async with sdk_client(handler, True) as client: + stream = await client.responses.create(model="parallel", input=QUESTION, stream=True) + if mode != "empty": + async for event in stream: + if status and event.type == "response." + status: + break # Terminal event received, but no EOF read yet. + await stream.close() + await stream.close() + assert stream.response.is_closed + + if async_: + asyncio.run(run()) + else: + with sdk_client(handler) as client: + stream = client.responses.create(model="parallel", input=QUESTION, stream=True) + if mode != "empty": + for event in stream: + if status and event.type == "response." + status: + break + stream.close() + stream.close() + assert stream.response.is_closed + span, = exporter.get_finished_spans() + assert span.status.status_code is (StatusCode.OK if mode == "complete" else StatusCode.ERROR) + assert KEY not in span.to_json() + if status: + assert span.attributes["gen_ai.usage.input_tokens"] == 5 + + +@pytest.mark.parametrize("async_", [False, True]) +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("flag", ["hide_inputs", "hide_input_text", "hide_input_messages"]) +def test_echoed_request_fields_never_become_output(tracing, async_, stream, flag): + provider, exporter = tracing + instruction = "synthetic-private-system-text" + instrumentor = OpenAIInstrumentor() + instrumentor.instrument(tracer_provider=provider, config=TraceConfig(**{flag: True})) + body = response_body() + # Empty output exercises the accumulator's JSON fallback instead of output_text. + body.update(instructions=instruction, metadata={"prompt": instruction}, + input=instruction, text={"format": {"type": "json_schema", "name": instruction}}) + + def handler(request): + if stream: + return httpx.Response(200, headers={"content-type": "text/event-stream"}, + content=sse_event("response.completed", body)) + return httpx.Response(200, json=body) + + async def run(): + async with sdk_client(handler, True) as client: + result = await client.responses.create(model="parallel", input=QUESTION, + instructions=instruction, stream=stream) + if stream: + events = [event async for event in result] + assert events[0].response.instructions == instruction + else: + assert result.instructions == instruction + + try: + if async_: + asyncio.run(run()) + else: + with sdk_client(handler) as client: + result = client.responses.create(model="parallel", input=QUESTION, + instructions=instruction, stream=stream) + if stream: + events = list(result) + assert events[0].response.instructions == instruction + else: + assert result.instructions == instruction + span, = exporter.get_finished_spans() + assert instruction not in span.to_json() + output = json.loads(span.attributes["output.value"]) + assert output["status"] == "completed" + assert output["usage"]["input_tokens"] == 5 + finally: + instrumentor.uninstrument() + + +@pytest.mark.parametrize("async_", [False, True]) +def test_raw_failed_response_retains_sdk_contract(openai_tracing, async_): + _, exporter = openai_tracing + + def handler(request): + return httpx.Response(200, json=response_body("failed")) + + async def run(): + async with sdk_client(handler, True) as client: + raw = await client.responses.with_raw_response.create(model="parallel", input=QUESTION) + assert raw.status_code == 200 + assert raw.parse().status == "failed" + + if async_: + asyncio.run(run()) + else: + with sdk_client(handler) as client: + raw = client.responses.with_raw_response.create(model="parallel", input=QUESTION) + assert raw.status_code == 200 + assert raw.parse().status == "failed" + span, = exporter.get_finished_spans() + assert span.status.status_code is StatusCode.ERROR + + +@pytest.mark.parametrize("async_", [False, True]) +@pytest.mark.parametrize("stream", [False, True]) +def test_hidden_error_details_apply_to_failed_response_json(tracing, async_, stream): + provider, exporter = tracing + instrumentor = OpenAIInstrumentor() + instrumentor.instrument(tracer_provider=provider, config=TraceConfig(hide_error_details=True)) + message = "synthetic-private-error-details" + body = response_body("failed") + body["error"]["message"] = message + + def handler(request): + if stream: + return httpx.Response(200, headers={"content-type": "text/event-stream"}, + content=sse_event("response.failed", body)) + return httpx.Response(200, json=body) + + async def run(): + async with sdk_client(handler, True) as client: + result = await client.responses.create(model="parallel", input=QUESTION, stream=stream) + if stream: + events = [event async for event in result] + assert events[0].response.error.message == message + else: + assert result.error.message == message + + try: + if async_: + asyncio.run(run()) + else: + with sdk_client(handler) as client: + result = client.responses.create(model="parallel", input=QUESTION, stream=stream) + if stream: + events = list(result) + assert events[0].response.error.message == message + else: + assert result.error.message == message + span, = exporter.get_finished_spans() + assert span.status.status_code is StatusCode.ERROR + assert message not in span.to_json() + finally: + instrumentor.uninstrument() + + +@pytest.mark.parametrize("async_", [False, True]) +def test_mcp_wire_and_function_schema_are_not_mutated(openai_tracing, async_): + _, exporter = openai_tracing + token = 'synthetic-token-with-"quote"-and-\\slash' + mcp = {"type": "mcp", "server_label": "synthetic", "server_url": "https://synthetic.invalid/mcp", + "authorization": token, "headers": {"X-Custom-Credential": token}} + function = {"type": "function", "name": "lookup", "parameters": { + "type": "object", "properties": {"authorization": {"type": "string"}}}} + tools = [mcp, function] + original = json.loads(json.dumps(tools)) + + def handler(request): + assert json.loads(request.content)["tools"] == original + body = response_body() + body["tools"] = tools # Don't rely on a provider scrubbing its response. + return httpx.Response(200, json=body) + + async def run(): + async with sdk_client(handler, True) as client: + result = await client.responses.create(model="parallel", input=QUESTION, tools=tools) + assert result.tools[0].authorization == token + + if async_: + asyncio.run(run()) + else: + with sdk_client(handler) as client: + result = client.responses.create(model="parallel", input=QUESTION, tools=tools) + assert result.tools[0].authorization == token + assert tools == original + span, = exporter.get_finished_spans() + assert "synthetic-token-with" not in span.to_json() + definitions = json.loads(span.attributes["gen_ai.tool.definitions"]) + assert "authorization" not in definitions[0] + assert "headers" not in definitions[0] + assert definitions[1] == function diff --git a/python/frameworks/openai/traceai_openai/_attributes/_responses_api.py b/python/frameworks/openai/traceai_openai/_attributes/_responses_api.py index f92d6ec07..e4ff64ec5 100644 --- a/python/frameworks/openai/traceai_openai/_attributes/_responses_api.py +++ b/python/frameworks/openai/traceai_openai/_attributes/_responses_api.py @@ -6,7 +6,8 @@ from opentelemetry.util.types import AttributeValue from typing_extensions import assert_never -from fi_instrumentation import safe_json_dumps +from fi_instrumentation import safe_json_dumps +from traceai_openai._privacy import safe_tool from fi_instrumentation.fi_types import ( ImageAttributes, MessageAttributes, @@ -237,7 +238,7 @@ def _get_attributes_from_response_tool( obj: responses.tool.Tool, prefix: str = "", ) -> Iterator[Tuple[str, AttributeValue]]: - yield (f"{prefix}{ToolAttributes.TOOL_JSON_SCHEMA}", obj.model_dump_json()) + yield (f"{prefix}{ToolAttributes.TOOL_JSON_SCHEMA}", safe_json_dumps(safe_tool(obj.model_dump(mode="json")))) @classmethod @stop_on_exception diff --git a/python/frameworks/openai/traceai_openai/_privacy.py b/python/frameworks/openai/traceai_openai/_privacy.py new file mode 100644 index 000000000..ae4a66e9f --- /dev/null +++ b/python/frameworks/openai/traceai_openai/_privacy.py @@ -0,0 +1,74 @@ +"""Telemetry-only copies; never mutate SDK requests, responses or exceptions.""" + +from __future__ import annotations + +import json +from collections.abc import Iterable, Mapping +from typing import Any + +from fi_instrumentation import TraceConfig + +REDACTED = "__REDACTED__" + + +def redact(value: Any, secrets: Iterable[str]) -> Any: + if isinstance(value, str): + for secret in secrets: + value = value.replace(secret, REDACTED) + return value + if isinstance(value, (list, tuple)): + return type(value)(redact(item, secrets) for item in value) + return value + + +def request_credentials(instance: Any, options: Any, parameters: Mapping) -> tuple[str, ...]: + """Read static credentials only; never invoke a caller's token provider.""" + values = [getattr(instance, "api_key", None)] + for headers in (getattr(instance, "_custom_headers", None), getattr(options, "headers", None)): + if isinstance(headers, Mapping): + for name, value in headers.items(): + if str(name).lower() in {"authorization", "api-key", "x-api-key"}: + values.append(value) + if isinstance(value, str) and value.lower().startswith("bearer "): + values.append(value[7:]) + tools = parameters.get("tools") + if isinstance(tools, (list, tuple)): + for tool in tools: + if isinstance(tool, Mapping) and tool.get("type") == "mcp": + values.append(tool.get("authorization")) + headers = tool.get("headers") + if isinstance(headers, Mapping): + values.extend(headers.values()) + secrets = set() + for value in values: + if isinstance(value, str) and value: + secrets.update((value, json.dumps(value)[1:-1], repr(value)[1:-1])) + return tuple(sorted(secrets, key=len, reverse=True)) + + +def safe_tool(tool: Any) -> Any: + if isinstance(tool, Mapping) and tool.get("type") == "mcp": + return {key: value for key, value in tool.items() if key not in {"authorization", "headers"}} + return tool + + +def safe_request(parameters: Mapping) -> dict: + """MCP headers are arbitrary authentication material, not tool schemas.""" + result = dict(parameters) + tools = result.get("tools") + if isinstance(tools, (list, tuple)): + result["tools"] = [safe_tool(tool) for tool in tools] + return result + + +def response_output(response: Any, config: TraceConfig | None = None) -> dict: + """Only response-owned data belongs in output, not echoed instructions/tools. + + An allowlist also keeps new SDK request-echo fields from bypassing input + privacy flags. The SDK response returned to the application is unchanged. + """ + fields = {"id", "object", "created_at", "model", "status", "output", "usage", + "error", "incomplete_details"} + if config is not None and getattr(config, "hide_error_details", False): + fields.difference_update({"error", "incomplete_details"}) + return response.model_dump(include=fields, exclude_none=True, mode="json") diff --git a/python/frameworks/openai/traceai_openai/_request.py b/python/frameworks/openai/traceai_openai/_request.py index c6e8cf426..9e06d500e 100644 --- a/python/frameworks/openai/traceai_openai/_request.py +++ b/python/frameworks/openai/traceai_openai/_request.py @@ -18,11 +18,13 @@ from opentelemetry.util.types import AttributeValue from traceai_openai._client_provider import _get_client_provider from traceai_openai._providers import provider_from_host +from traceai_openai._privacy import redact, request_credentials, safe_request from traceai_openai._request_attributes_extractor import _RequestAttributesExtractor from traceai_openai._response_accumulator import ( _ChatCompletionAccumulator, _CompletionAccumulator, _ResponsesAccumulator, + _response_status, ) from traceai_openai._response_attributes_extractor import _ResponseAttributesExtractor from traceai_openai._span_io_handler import add_io_to_span_attributes @@ -57,12 +59,13 @@ def _start_as_current_span( attributes: Iterable[Tuple[str, AttributeValue]], context_attributes: Iterable[Tuple[str, AttributeValue]], extra_attributes: Iterable[Tuple[str, AttributeValue]], + secrets: Tuple[str, ...] = (), ) -> Iterator[_WithSpan]: # Keep primary attributes so finalization can restore them after optional # details: OTel's bounded attribute map evicts the oldest entries. initial_attributes = {} try: - initial_attributes = dict(attributes) + initial_attributes = {key: redact(value, secrets) for key, value in attributes} span = self._tracer.start_span(name=span_name, attributes=initial_attributes) except Exception: logger.exception("Failed to start span") @@ -75,6 +78,7 @@ def _start_as_current_span( ) as span: yield _WithSpan( span=span, + secrets=secrets, initial_attributes=initial_attributes, context_attributes=dict(context_attributes), extra_attributes=dict(extra_attributes), @@ -114,6 +118,7 @@ def responses_accumulator(request_parameters: _RequestParameters) -> Any: return _ResponsesAccumulator( request_parameters=request_parameters, chat_completion_type=openai.types.responses.response.Response, + config=self._config, response_attributes_extractor=self._response_attributes_extractor, ) @@ -206,9 +211,10 @@ def _finalize_response( Monkey-patch the response object to trace the stream, or finish tracing if the response is not a stream. """ + parsed_response = response if hasattr(response, "parse") and callable(response.parse): try: - response.parse() + parsed_response = response.parse() except Exception: logger.exception(f"Failed to parse response of type {type(response)}") @@ -240,7 +246,7 @@ def _finalize_response( else: _finish_tracing( - status=trace_api.Status(status_code=trace_api.StatusCode.OK), + status=_response_status(parsed_response), with_span=with_span, has_attributes=_ResponseAttributes( request_parameters=request_parameters, @@ -285,6 +291,8 @@ def __call__( return wrapped(*args, **kwargs) try: cast_to, request_parameters = _parse_request_args(args) + secrets = request_credentials(instance, args[1], request_parameters) + request_parameters = safe_request(request_parameters) # Extract input data more robustly input_data = ( request_parameters.get("messages") @@ -299,6 +307,7 @@ def __call__( with self._start_as_current_span( span_name=span_name, + secrets=secrets, attributes=chain( self._get_attributes_from_instance(instance), self._get_attributes_from_request( @@ -333,6 +342,7 @@ def __call__( None, response, is_streaming=self._is_streaming(response), + config=self._config, ) return self._finalize_response( response=response, @@ -340,7 +350,7 @@ def __call__( cast_to=cast_to, request_parameters=request_parameters, ) - except Exception as e: + except BaseException as e: status = trace_api.Status( status_code=trace_api.StatusCode.ERROR, description=f"{type(e).__name__}: {e}", @@ -362,6 +372,8 @@ async def __call__( return await wrapped(*args, **kwargs) try: cast_to, request_parameters = _parse_request_args(args) + secrets = request_credentials(instance, args[1], request_parameters) + request_parameters = safe_request(request_parameters) # Extract input data more robustly input_data = ( request_parameters.get("messages") @@ -376,6 +388,7 @@ async def __call__( return await wrapped(*args, **kwargs) with self._start_as_current_span( span_name=span_name, + secrets=secrets, attributes=chain( self._get_attributes_from_instance(instance), self._get_attributes_from_request( @@ -404,9 +417,10 @@ async def __call__( ) add_io_to_span_attributes( - with_span, None, response, is_streaming=self._is_streaming(response) + with_span, None, response, is_streaming=self._is_streaming(response), + config=self._config, ) - except Exception as exception: + except BaseException as exception: with_span.record_exception(exception) status = trace_api.Status( status_code=trace_api.StatusCode.ERROR, diff --git a/python/frameworks/openai/traceai_openai/_response_accumulator.py b/python/frameworks/openai/traceai_openai/_response_accumulator.py index 17417d213..3b73efdeb 100644 --- a/python/frameworks/openai/traceai_openai/_response_accumulator.py +++ b/python/frameworks/openai/traceai_openai/_response_accumulator.py @@ -17,16 +17,18 @@ Type, ) -from fi_instrumentation import safe_json_dumps +from fi_instrumentation import TraceConfig, safe_json_dumps from fi_instrumentation.fi_types import FiMimeTypeValues, SpanAttributes from opentelemetry.util.types import AttributeValue +from opentelemetry.trace import Status, StatusCode +from traceai_openai._privacy import response_output from traceai_openai._utils import _as_output_attributes, _to_dict, _ValueAndType, _io_value_and_type if TYPE_CHECKING: from openai.types import Completion from openai.types.chat import ChatCompletion from openai.types.responses.response import Response - from openai.types.responses.response_completed_event import ResponseCompletedEvent + from openai.types.responses import ResponseStreamEvent __all__ = ( "_CompletionAccumulator", @@ -43,8 +45,16 @@ def get_attributes_from_response( ) -> Iterator[Tuple[str, AttributeValue]]: ... +def _response_status(response: Any) -> Status: + if getattr(response, "status", None) in {"failed", "cancelled", "incomplete"}: + # The vendor's error message can contain credentials or request content. + return Status(StatusCode.ERROR, "Response did not complete") + return Status(StatusCode.OK) + + class _ResponsesAccumulator: __slots__ = ( + "_config", "_is_null", "_values", "_cached_result", @@ -58,16 +68,20 @@ def __init__( request_parameters: Mapping[str, Any], chat_completion_type: Type["Response"], response_attributes_extractor: Optional[_CanGetAttributesFromResponse] = None, + config: Optional[TraceConfig] = None, ) -> None: + self._config = config self._chat_completion_type = chat_completion_type self._request_parameters = request_parameters self._response_attributes_extractor = response_attributes_extractor self._is_null = True self._cached_result: Optional[Dict[str, Any]] = None - self._values: Optional["ResponseCompletedEvent"] = None + self._values: Optional["ResponseStreamEvent"] = None def process_chunk(self, chunk: Any) -> None: - if type(chunk).__name__ == "ResponseCompletedEvent": + if getattr(chunk, "type", None) in { + "response.completed", "response.failed", "response.incomplete", "error" + }: self._is_null = False self._cached_result = None self._values = chunk @@ -76,7 +90,14 @@ def _result(self) -> Any: if self._is_null: return None if self._values: - return self._values.response + return getattr(self._values, "response", None) + + def get_status(self) -> Status: + if self._values is None: + return Status(StatusCode.ERROR, "Response stream ended before a terminal event") + if self._values.type != "response.completed": + return Status(StatusCode.ERROR, "Response stream did not complete") + return _response_status(self._result()) def get_attributes(self) -> Iterator[Tuple[str, AttributeValue]]: if not (result := self._result()): @@ -97,9 +118,9 @@ def get_attributes(self) -> Iterator[Tuple[str, AttributeValue]]: if output_content: yield SpanAttributes.OUTPUT_VALUE, output_content else: - # Fallback to full response if we can't extract content + # Keep response-owned data, never echoed request fields. yield from _as_output_attributes( - _io_value_and_type(result), + _io_value_and_type(response_output(result, self._config)), ) def get_extra_attributes(self) -> Iterator[Tuple[str, AttributeValue]]: diff --git a/python/frameworks/openai/traceai_openai/_span_io_handler.py b/python/frameworks/openai/traceai_openai/_span_io_handler.py index 028d2faa5..772db2da3 100644 --- a/python/frameworks/openai/traceai_openai/_span_io_handler.py +++ b/python/frameworks/openai/traceai_openai/_span_io_handler.py @@ -2,8 +2,10 @@ import logging from typing import Any, Optional +from fi_instrumentation import TraceConfig from fi_instrumentation.fi_types import FiMimeTypeValues, SpanAttributes from traceai_openai._with_span import _WithSpan +from traceai_openai._privacy import response_output logger = logging.getLogger(__name__) logger.addHandler(logging.NullHandler()) @@ -88,7 +90,7 @@ def _set_mime_type(span: _WithSpan, key: str, value: Any) -> None: span.set_attribute(key, mime) -def _json_output_for(output_data: Any) -> Optional[str]: +def _json_output_for(output_data: Any, config: Optional[TraceConfig] = None) -> Optional[str]: """JSON for responses whose output is not a chat message: the Responses API object as JSON, and a compact summary for embeddings (never the vectors).""" try: @@ -98,7 +100,7 @@ def _json_output_for(output_data: Any) -> Optional[str]: return None if isinstance(output_data, Response): # Lead with the answer text (the SDK's `output_text`) so previews show it. - body = json.loads(output_data.model_dump_json(exclude_none=True)) + body = response_output(output_data, config) output_text = getattr(output_data, "output_text", None) if output_text: body = {"output_text": output_text, **body} @@ -116,7 +118,8 @@ def _json_output_for(output_data: Any) -> Optional[str]: def add_io_to_span_attributes( - span: _WithSpan, input_data: Any, output_data: Any, is_streaming: bool = False + span: _WithSpan, input_data: Any, output_data: Any, is_streaming: bool = False, + config: Optional[TraceConfig] = None ) -> None: """Add input/output data to span attributes. @@ -135,7 +138,7 @@ def add_io_to_span_attributes( if output_data is not None and not is_streaming: if hasattr(output_data, "parse") and not hasattr(output_data, "model_dump"): output_data = output_data.parse() - if (json_output := _json_output_for(output_data)) is not None: + if (json_output := _json_output_for(output_data, config)) is not None: span.set_attribute(SpanAttributes.OUTPUT_VALUE, json_output) span.set_attribute( SpanAttributes.OUTPUT_MIME_TYPE, FiMimeTypeValues.JSON.value diff --git a/python/frameworks/openai/traceai_openai/_stream.py b/python/frameworks/openai/traceai_openai/_stream.py index a3a60042e..20dd5a5c4 100644 --- a/python/frameworks/openai/traceai_openai/_stream.py +++ b/python/frameworks/openai/traceai_openai/_stream.py @@ -40,6 +40,38 @@ def _finish_on_exit(stream: Any, error: Optional[BaseException]) -> None: stream._finish_tracing(status=status) +class _ResponseCloseProxy(ObjectProxy): + """Responses.stream() closes the HTTP response, not the raw SDK stream. + + Expose a per-call proxy instead of patching HTTPX classes or the client's + transport. Internal SDK reads still use the original response unchanged. + """ + + def __init__(self, response: Any, stream: Any) -> None: + super().__init__(response) + self._self_stream = stream + + def close(self) -> None: + error = None + try: + self.__wrapped__.close() + except BaseException as exc: + error = exc + raise + finally: + _finish_on_exit(self._self_stream, error) + + async def aclose(self) -> None: + error = None + try: + await self.__wrapped__.aclose() + except BaseException as exc: + error = exc + raise + finally: + _finish_on_exit(self._self_stream, error) + + class _Stream(ObjectProxy): # type: ignore __slots__ = ( "_self_with_span", @@ -60,6 +92,13 @@ def __init__( self._self_is_finished = with_span.is_finished self._self_response_accumulator = response_accumulator + @property + def response(self) -> Any: + response = self.__wrapped__.response + if getattr(self._self_response_accumulator, "get_status", None) is not None: + return _ResponseCloseProxy(response, self) + return response + def __enter__(self) -> "_Stream": # ObjectProxy otherwise returns the SDK stream, bypassing our iterator. self.__wrapped__.__enter__() @@ -74,6 +113,16 @@ def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> Any: finally: _finish_on_exit(self, exc) + def close(self) -> None: + error = None + try: + self.__wrapped__.close() + except BaseException as exc: + error = exc + raise + finally: + _finish_on_exit(self, error) + def __iter__(self) -> Iterator[Any]: return self @@ -85,7 +134,7 @@ def __next__(self) -> Any: status = trace_api.Status(status_code=trace_api.StatusCode.OK) self._finish_tracing(status=status) raise - except Exception as exception: + except BaseException as exception: if not self._self_is_finished: status = trace_api.Status( status_code=trace_api.StatusCode.ERROR, @@ -115,6 +164,12 @@ def _finish_tracing( self, status: Optional[trace_api.Status] = None, ) -> None: + if self._self_is_finished: + return + if status is None or status.status_code is not trace_api.StatusCode.ERROR: + get_status = getattr(self._self_response_accumulator, "get_status", None) + if get_status is not None: + status = get_status() _finish_tracing( status=status, with_span=self._self_with_span, @@ -151,6 +206,13 @@ def __init__( self._self_is_finished = with_span.is_finished self._self_response_accumulator = response_accumulator + @property + def response(self) -> Any: + response = self.__wrapped__.response + if getattr(self._self_response_accumulator, "get_status", None) is not None: + return _ResponseCloseProxy(response, self) + return response + async def __aenter__(self) -> "_AsyncStream": await self.__wrapped__.__aenter__() return self @@ -164,6 +226,16 @@ async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> Any: finally: _finish_on_exit(self, exc) + async def close(self) -> None: + error = None + try: + await self.__wrapped__.close() + except BaseException as exc: + error = exc + raise + finally: + _finish_on_exit(self, error) + def __aiter__(self): return self @@ -175,7 +247,7 @@ async def __anext__(self): status = trace_api.Status(status_code=trace_api.StatusCode.OK) self._finish_tracing(status=status) raise - except Exception as exception: + except BaseException as exception: if not self._self_is_finished: status = trace_api.Status( status_code=trace_api.StatusCode.ERROR, @@ -205,6 +277,12 @@ def _finish_tracing( self, status: Optional[trace_api.Status] = None, ) -> None: + if self._self_is_finished: + return + if status is None or status.status_code is not trace_api.StatusCode.ERROR: + get_status = getattr(self._self_response_accumulator, "get_status", None) + if get_status is not None: + status = get_status() _finish_tracing( status=status, with_span=self._self_with_span, diff --git a/python/frameworks/openai/traceai_openai/_with_span.py b/python/frameworks/openai/traceai_openai/_with_span.py index 5f5419ad0..8efb36318 100644 --- a/python/frameworks/openai/traceai_openai/_with_span.py +++ b/python/frameworks/openai/traceai_openai/_with_span.py @@ -1,9 +1,11 @@ import logging -from typing import Any, Optional +import traceback +from typing import Any, Optional, Tuple from fi_instrumentation.fi_types import SpanAttributes from opentelemetry import trace as trace_api from opentelemetry.util.types import Attributes +from traceai_openai._privacy import redact logger = logging.getLogger(__name__) logger.addHandler(logging.NullHandler()) @@ -22,6 +24,7 @@ class _WithSpan: __slots__ = ( "_span", + "_secrets", "_primary_attributes", "_context_attributes", "_extra_attributes", @@ -34,8 +37,10 @@ def __init__( context_attributes: Attributes = None, extra_attributes: Attributes = None, initial_attributes: Attributes = None, + secrets: Tuple[str, ...] = (), ) -> None: self._span = span + self._secrets = secrets self._primary_attributes = dict(initial_attributes or {}) self._context_attributes = context_attributes self._extra_attributes = extra_attributes @@ -49,11 +54,20 @@ def __init__( def is_finished(self) -> bool: return self._is_finished - def record_exception(self, exception: Exception) -> None: + def record_exception(self, exception: BaseException) -> None: if self._is_finished: return try: - self._span.record_exception(exception) + # Emit the standard event through FITracer's masking API, without + # handing raw credential-bearing exception text to an exporter. + cls = type(exception) + self._span.add_event("exception", attributes={ + "exception.type": f"{cls.__module__}.{cls.__qualname__}", + "exception.message": redact(str(exception), self._secrets), + "exception.stacktrace": redact("".join(traceback.format_exception( + cls, exception, exception.__traceback__)), self._secrets), + "exception.escaped": "False", + }) except Exception: logger.exception("Failed to record exception on span") @@ -96,7 +110,9 @@ def finish_tracing( self.set_attribute(key, value) if status is not None: try: - self._span.set_status(status=status) + self._span.set_status(status=trace_api.Status( + status.status_code, redact(status.description, self._secrets) + )) except Exception: logger.exception("Failed to set status code on span") try: @@ -110,6 +126,7 @@ def set_attribute(self, key: str, value: Any) -> None: if self._is_finished: return try: + value = redact(value, self._secrets) self._span.set_attribute(key, value) if key in self._primary_attributes and value is not None: self._primary_attributes[key] = value diff --git a/python/frameworks/parallel/README.md b/python/frameworks/parallel/README.md index ab85d79d7..f6e9448fa 100644 --- a/python/frameworks/parallel/README.md +++ b/python/frameworks/parallel/README.md @@ -234,7 +234,7 @@ def research(client: Parallel, tracer: FITracer, question: str) -> TaskRunResult span.set_attribute("parallel.run_id", run.run_id) span.set_attribute("parallel.processor", "base") result = client.task_run.result(run.run_id, api_timeout=600, timeout=610) - except Exception: + except BaseException: span.set_status(Status(StatusCode.ERROR, "Task request failed")) raise span.set_attribute("parallel.task.status", result.run.status) diff --git a/python/frameworks/parallel/examples/task_workflow.py b/python/frameworks/parallel/examples/task_workflow.py index a3c32c20f..cbaf2b09f 100644 --- a/python/frameworks/parallel/examples/task_workflow.py +++ b/python/frameworks/parallel/examples/task_workflow.py @@ -25,7 +25,7 @@ def research(client: Parallel, tracer: FITracer, question: str) -> TaskRunResult span.set_attribute("parallel.run_id", run.run_id) span.set_attribute("parallel.processor", "base") result = client.task_run.result(run.run_id, api_timeout=600, timeout=610) - except Exception: + except BaseException: span.set_status(Status(StatusCode.ERROR, "Task request failed")) raise span.set_attribute("parallel.task.status", result.run.status) @@ -48,7 +48,7 @@ async def research_async(client: AsyncParallel, tracer: FITracer, question: str) span.set_attribute("parallel.cancelled", True) span.set_status(Status(StatusCode.ERROR, "Task wait cancelled")) raise - except Exception: + except BaseException: span.set_status(Status(StatusCode.ERROR, "Task request failed")) raise span.set_attribute("parallel.task.status", result.run.status) diff --git a/python/frameworks/parallel/tests/test_task_workflow.py b/python/frameworks/parallel/tests/test_task_workflow.py index 52d40bdb4..1e6e09b93 100644 --- a/python/frameworks/parallel/tests/test_task_workflow.py +++ b/python/frameworks/parallel/tests/test_task_workflow.py @@ -205,3 +205,34 @@ def log_message(self, *_): server.shutdown() server.server_close() thread.join(timeout=3) + + +@pytest.mark.parametrize("stage", ["create", "wait"]) +@pytest.mark.parametrize("error_cls", [KeyboardInterrupt, SystemExit]) +def test_task_interrupt_is_error_with_registered_processor(stage, error_cls): + from fi_instrumentation import register + from fi_instrumentation.fi_types import ProjectType + + exporter = InMemorySpanExporter() + provider = register(project_name="task-interrupt", project_type=ProjectType.OBSERVE, + span_exporter=exporter, batch=False, verbose=False) + abort = error_cls("synthetic local abort " + KEY) + + def handler(request): + if stage == "create" or request.method == "GET": + raise abort + return httpx.Response(200, json=run_body("queued")) + + try: + with Parallel(api_key=KEY, max_retries=0, http_client=httpx.Client( + transport=httpx.MockTransport(handler), trust_env=False, + )) as client: + with pytest.raises(error_cls) as caught: + example.research(client, FITracer(provider.get_tracer("task-interrupt")), QUESTION) + assert caught.value is abort + span, = exporter.get_finished_spans() + assert span.status.status_code is StatusCode.ERROR + assert KEY not in span.to_json() + assert span.events == () + finally: + provider.shutdown() From d4595dcdc24269c6f2a6ab59ba2c7c7b046dcc2e Mon Sep 17 00:00:00 2001 From: Nikhil Pareek Date: Sun, 1 Nov 2026 22:41:02 -0800 Subject: [PATCH 02/10] fix(openai): close Responses privacy and raw-stream gaps Redact resolved request credentials, prevent core traceback reconstruction from restoring secrets, mask nested MCP errors, and trace parsed raw streams. Keep legacy core and published recipe contracts covered. Co-Authored-By: Claude Code --- .github/workflows/parallel-python-tests.yml | 3 +- .../tests/test_perplexity_recipe.py | 36 +- .../test_responses_privacy_and_raw_streams.py | 411 ++++++++++++++++++ .../openai/tests/test_responses_safety.py | 5 + .../openai/traceai_openai/_privacy.py | 32 +- .../openai/traceai_openai/_request.py | 22 +- .../openai/traceai_openai/_stream.py | 17 + .../openai/traceai_openai/_with_span.py | 17 +- 8 files changed, 525 insertions(+), 18 deletions(-) create mode 100644 python/frameworks/openai/tests/test_responses_privacy_and_raw_streams.py diff --git a/.github/workflows/parallel-python-tests.yml b/.github/workflows/parallel-python-tests.yml index 6ef2de5c6..97445cf90 100644 --- a/.github/workflows/parallel-python-tests.yml +++ b/.github/workflows/parallel-python-tests.yml @@ -65,7 +65,8 @@ jobs: packages=(python/frameworks/{parallel,openai}/dist/*.tar.gz) fi .parallel-test-venv/bin/pip install "${packages[@]}" \ - "parallel-web==$PARALLEL_VERSION" "openai==$OPENAI_VERSION" pytest pytest-asyncio + "parallel-web==$PARALLEL_VERSION" "openai==$OPENAI_VERSION" \ + 'fi-instrumentation-otel>=1.2.2' pytest pytest-asyncio - name: Test with only loopback networking shell: bash run: | diff --git a/python/examples/perplexity/tests/test_perplexity_recipe.py b/python/examples/perplexity/tests/test_perplexity_recipe.py index 6ab2c93a2..b93530b1a 100644 --- a/python/examples/perplexity/tests/test_perplexity_recipe.py +++ b/python/examples/perplexity/tests/test_perplexity_recipe.py @@ -37,6 +37,19 @@ == Path(CORE_DISTRIBUTION.locate_file("fi_instrumentation/__init__.py")).resolve() ) +# The recipe CI deliberately retains the old published SDK as a compatibility +# control. Only that exact installed distribution exports echoed instructions; +# current source must not reintroduce them to satisfy legacy expectations. +try: + OPENAI_DISTRIBUTION = distribution("traceAI-openai") +except PackageNotFoundError: + LEGACY_OUTPUT = False +else: + LEGACY_OUTPUT = OPENAI_DISTRIBUTION.version == "0.1.10" and ( + Path(traceai_openai.__file__).resolve() + == Path(OPENAI_DISTRIBUTION.locate_file("traceai_openai/__init__.py")).resolve() + ) + TEST_DIR = Path(__file__).resolve().parent RECIPE_DIR = TEST_DIR.parent VENDOR_KEY = "placeholder-perplexity-key" @@ -254,7 +267,10 @@ def test_responses_at_documented_host(tracing): assert raw["usage"]["input_tokens_details"][key] == value assert raw["usage"]["output_tokens_details"] == fake.USAGE["output_tokens_details"] assert raw["usage"]["tool_calls_details"] == fake.USAGE["tool_calls_details"] - assert fake.INSTRUCTIONS_MARKER in raw["instructions"] + if LEGACY_OUTPUT: + assert fake.INSTRUCTIONS_MARKER in raw["instructions"] + else: + assert "instructions" not in raw assert fake.SEARCH_URL in attrs["output.value"] assert fake.SEARCH_SNIPPET in attrs["output.value"] @@ -349,9 +365,11 @@ def handler(request): assert _call(handler, prompt=marker) == fake.ANSWER span, _ = _assert_export(receiver, provider, project, hide_inputs=hidden) assert (marker in json.dumps(span)) is (not hidden) - # Output echoes remain visible even when the input itself is hidden. - for output_marker in (fake.QUERY_MARKER, fake.INSTRUCTIONS_MARKER): - assert output_marker in _attributes(span)["output.value"] + # Search queries are response-owned output; instructions are input. + output = _attributes(span)["output.value"] + assert fake.QUERY_MARKER in output + assert (fake.INSTRUCTIONS_MARKER in output) is LEGACY_OUTPUT + assert fake.ANSWER in output @pytest.mark.parametrize("stream", [False, True], ids=["nonstream", "stream"]) @@ -363,8 +381,11 @@ def test_hide_outputs_with_control(tracing, monkeypatch, stream): span, attrs = _assert_export(receiver, provider, project) _assert_usage(attrs) assert (fake.ANSWER in json.dumps(span)) is (not hidden) - for marker in (fake.SEARCH_URL, fake.SEARCH_SNIPPET, fake.QUERY_MARKER, fake.INSTRUCTIONS_MARKER): + for marker in (fake.SEARCH_URL, fake.SEARCH_SNIPPET, fake.QUERY_MARKER): assert (marker in json.dumps(span)) is (not hidden and not stream) + assert (fake.INSTRUCTIONS_MARKER in json.dumps(span)) is ( + LEGACY_OUTPUT and not hidden and not stream + ) @pytest.mark.parametrize("hidden", [False, True], ids=["visible", "hidden"]) @@ -381,8 +402,9 @@ def test_empty_text_stream_raw_fallback(tracing, monkeypatch, hidden): list(client.responses.create(model=fake.REQUEST_MODEL, input="empty stream", stream=True)) span, attrs = _assert_export(receiver, provider, project) _assert_usage(attrs) - for marker in (fake.SEARCH_URL, fake.QUERY_MARKER, fake.INSTRUCTIONS_MARKER): + for marker in (fake.SEARCH_URL, fake.QUERY_MARKER): assert (marker in json.dumps(span)) is (not hidden) + assert (fake.INSTRUCTIONS_MARKER in json.dumps(span)) is (LEGACY_OUTPUT and not hidden) def test_cache_details_require_output_details(tracing): @@ -457,7 +479,7 @@ def test_subprocess_app(tmp_path, stream, cli_model): assert (fake.SEARCH_URL in json.dumps(attrs)) is (not stream) assert (fake.SEARCH_SNIPPET in json.dumps(attrs)) is (not stream) assert (fake.QUERY_MARKER in json.dumps(attrs)) is (not stream) - assert (fake.INSTRUCTIONS_MARKER in json.dumps(attrs)) is (not stream) + assert (fake.INSTRUCTIONS_MARKER in json.dumps(attrs)) is (LEGACY_OUTPUT and not stream) [export] = receiver.requests() assert export["path"] == "/tracer/v1/traces" assert export["headers"]["x-api-key"] == FI_KEY diff --git a/python/frameworks/openai/tests/test_responses_privacy_and_raw_streams.py b/python/frameworks/openai/tests/test_responses_privacy_and_raw_streams.py new file mode 100644 index 000000000..f9a83ae2e --- /dev/null +++ b/python/frameworks/openai/tests/test_responses_privacy_and_raw_streams.py @@ -0,0 +1,411 @@ +"""Responses privacy and raw-envelope contracts using synthetic SDK transports.""" + +import asyncio +import inspect +import json + +import httpx +import openai +import pytest +from fi_instrumentation import TraceConfig, register +from fi_instrumentation.fi_types import ProjectType +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import StatusCode +from traceai_openai import OpenAIInstrumentor + +KEY = "synthetic-review-key-not-for-export" +ERROR = "synthetic-private-tool-error" + + +@pytest.fixture +def tracing(): + exporter = InMemorySpanExporter() + provider = register( + project_name="independent-review", + project_type=ProjectType.OBSERVE, + span_exporter=exporter, + batch=False, + verbose=False, + ) + try: + yield provider, exporter + finally: + OpenAIInstrumentor().uninstrument() + provider.force_flush() + provider.shutdown() + + +def client(handler, async_=False, key=KEY): + http = (httpx.AsyncClient if async_ else httpx.Client)( + transport=httpx.MockTransport(handler), trust_env=False + ) + return (openai.AsyncOpenAI if async_ else openai.OpenAI)( + api_key=key, + base_url="https://api.parallel.ai/v1", + max_retries=0, + http_client=http, + ) + + +def body(status="completed"): + return { + "id": "resp_review", + "object": "response", + "created_at": 1, + "model": "parallel", + "status": status, + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "output": [], + "error": None, + "usage": { + "input_tokens": 5, + "output_tokens": 2, + "total_tokens": 7, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens_details": {"reasoning_tokens": 0}, + }, + } + + +def sse(kind, response): + return ( + "event: " + + kind + + "\ndata: " + + json.dumps({"type": kind, "sequence_number": 1, "response": response}) + + "\n\n" + ).encode() + + +@pytest.mark.parametrize("async_", [False, True]) +@pytest.mark.parametrize("traceback_marker", [False, True]) +def test_key_redaction_survives_core_status_cleaning(tracing, async_, traceback_marker): + provider, exporter = tracing + OpenAIInstrumentor().instrument(tracer_provider=provider) + message = "Invalid credential: " + KEY + if traceback_marker: + message += "\nTraceback (most recent call last):\n synthetic remote failure" + + def handler(request): + assert request.headers["authorization"] == "Bearer " + KEY + return httpx.Response( + 401, json={"error": {"code": "invalid_api_key", "message": message}} + ) + + async def run(): + async with client(handler, True) as c: + await c.responses.create(model="parallel", input="synthetic question") + + with pytest.raises(openai.AuthenticationError) as caught: + if async_: + asyncio.run(run()) + else: + with client(handler) as c: + c.responses.create(model="parallel", input="synthetic question") + (span,) = exporter.get_finished_spans() + assert KEY in str(caught.value) + assert KEY not in span.to_json(), span.to_json() + + +@pytest.mark.parametrize("async_", [False, True]) +def test_callable_credentials_are_scrubbed_without_extra_provider_calls( + tracing, async_ +): + if tuple(map(int, openai.__version__.split(".")[:2])) < (3, 0): + pytest.skip("callable OpenAI credentials are a latest-SDK feature") + provider, exporter = tracing + OpenAIInstrumentor().instrument(tracer_provider=provider) + calls = [] + + def token(): + calls.append(1) + return KEY + + async def atoken(): + return token() + + def handler(request): + assert request.headers["authorization"] == "Bearer " + KEY + return httpx.Response( + 401, + json={ + "error": {"code": "invalid_api_key", "message": "Invalid key: " + KEY} + }, + ) + + async def run(): + async with client(handler, True, atoken) as c: + await c.responses.create(model="parallel", input="synthetic question") + + with pytest.raises(openai.AuthenticationError) as caught: + if async_: + asyncio.run(run()) + else: + with client(handler, key=token) as c: + c.responses.create(model="parallel", input="synthetic question") + (span,) = exporter.get_finished_spans() + assert len(calls) == 1 + assert KEY in str(caught.value) + assert caught.value.request.headers["authorization"] == "Bearer " + KEY + assert KEY not in span.to_json(), span.to_json() + + +@pytest.mark.parametrize("async_", [False, True]) +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("hidden", [False, True]) +@pytest.mark.parametrize( + "kind", + ["http_error", "mcp_protocol_error", "mcp_tool_execution_error", "mcp_list_tools"], +) +@pytest.mark.skipif( + "hide_error_details" not in inspect.signature(TraceConfig).parameters, + reason="Error-detail masking requires a core with hide_error_details support", +) +def test_hidden_error_details_cover_mcp_failure_output( + tracing, async_, stream, hidden, kind +): + provider, exporter = tracing + OpenAIInstrumentor().instrument( + tracer_provider=provider, config=TraceConfig(hide_error_details=hidden) + ) + response = body() + response["output"] = [ + { + "id": "mcp_review", + "type": "mcp_call", + "name": "lookup", + "arguments": "{}", + "server_label": "test", + "error": {"type": "http_error", "code": 500, "message": ERROR}, + "output": "ordinary tool result", + } + ] + if kind == "mcp_list_tools": + response["output"] = [ + { + "id": "mcp_list", + "type": kind, + "server_label": "test", + "tools": [], + "error": ERROR, + } + ] + elif kind == "mcp_tool_execution_error": + response["output"][0]["error"] = {"type": kind, "content": {"text": ERROR}} + else: + response["output"][0]["error"]["type"] = kind + + def handler(request): + if stream: + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=sse("response.completed", response), + ) + return httpx.Response(200, json=response) + + async def run(): + async with client(handler, True) as c: + result = await c.responses.create( + model="parallel", input="question", stream=stream + ) + if stream: + events = [event async for event in result] + assert ERROR in events[-1].response.model_dump_json() + else: + assert ERROR in result.model_dump_json() + + if async_: + asyncio.run(run()) + else: + with client(handler) as c: + result = c.responses.create( + model="parallel", input="question", stream=stream + ) + if stream: + assert ERROR in list(result)[-1].response.model_dump_json() + else: + assert ERROR in result.model_dump_json() + (span,) = exporter.get_finished_spans() + assert (ERROR in span.to_json()) is (not hidden), span.to_json() + output = json.loads(span.attributes["output.value"])["output"][0] + if kind != "mcp_list_tools": + assert output["output"] == "ordinary tool result" + assert output["error"]["type"] == kind + if kind != "mcp_tool_execution_error": + assert output["error"]["code"] == 500 + elif hidden: + assert "error" not in output + + +@pytest.mark.parametrize("async_", [False, True]) +@pytest.mark.parametrize("status", ["failed", "incomplete", "completed", "in_progress"]) +@pytest.mark.parametrize("close_http", [False, True]) +def test_raw_stream_lifecycle_retains_sdk_contract(tracing, async_, status, close_http): + provider, exporter = tracing + OpenAIInstrumentor().instrument(tracer_provider=provider) + response = body(status) + if status == "failed": + response["error"] = {"code": "server_error", "message": "synthetic failure"} + kind = "response.created" if status == "in_progress" else "response." + status + + def handler(request): + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=sse(kind, response), + ) + + async def run(): + async with client(handler, True) as c: + raw = await c.responses.with_raw_response.create( + model="parallel", input="question", stream=True + ) + assert raw.status_code == 200 + assert raw.headers["content-type"] == "text/event-stream" + assert isinstance(raw, openai._legacy_response.LegacyAPIResponse) + assert not exporter.get_finished_spans() + stream = raw.parse() + assert raw.parse() is stream + assert raw.parse(to=openai.types.responses.Response) is stream + assert (await stream.__anext__()).response.status == status + if close_http: + await raw.http_response.aclose() + else: + await stream.close() + await stream.close() + assert raw.http_response.is_closed + + if async_: + asyncio.run(run()) + else: + with client(handler) as c: + raw = c.responses.with_raw_response.create( + model="parallel", input="question", stream=True + ) + assert raw.status_code == 200 + assert raw.headers["content-type"] == "text/event-stream" + assert isinstance(raw, openai._legacy_response.LegacyAPIResponse) + assert not exporter.get_finished_spans() + stream = raw.parse() + assert raw.parse() is stream + assert raw.parse(to=openai.types.responses.Response) is stream + assert next(stream).response.status == status + if close_http: + raw.http_response.close() + else: + stream.close() + stream.close() + assert raw.http_response.is_closed + (span,) = exporter.get_finished_spans() + assert span.status.status_code is ( + StatusCode.OK if status == "completed" else StatusCode.ERROR + ) + if status != "in_progress": + assert span.attributes["gen_ai.usage.input_tokens"] == 5 + assert span.attributes["gen_ai.usage.output_tokens"] == 2 + assert span.attributes["gen_ai.usage.total_tokens"] == 7 + + +def test_concurrent_callable_credentials_use_each_failed_request(tracing): + if tuple(map(int, openai.__version__.split(".")[:2])) < (3, 0): + pytest.skip("callable OpenAI credentials are a latest-SDK feature") + provider, exporter = tracing + OpenAIInstrumentor().instrument(tracer_provider=provider) + keys = [KEY + "-first", KEY + "-second"] + issued = [] + + async def run(): + both_sent = asyncio.Event() + received = [] + + async def token(): + value = keys[len(issued)] + issued.append(value) + return value + + async def handler(request): + credential = request.headers["authorization"][7:] + received.append(credential) + if len(received) == 2: + both_sent.set() + await asyncio.wait_for(both_sent.wait(), 3) + return httpx.Response( + 401, + json={ + "error": { + "code": "invalid_api_key", + "message": "Invalid key: " + credential, + } + }, + ) + + async with client(handler, True, token) as c: + results = await asyncio.gather( + *[c.responses.create(model="parallel", input="question") for _ in keys], + return_exceptions=True, + ) + assert set(received) == set(keys) + assert all(isinstance(result, openai.AuthenticationError) for result in results) + for result in results: + assert result.request.headers["authorization"][7:] in str(result) + + asyncio.run(run()) + assert issued == keys + spans = exporter.get_finished_spans() + assert len(spans) == 2 + for span in spans: + assert span.status.status_code is StatusCode.ERROR + assert all(key not in span.to_json() for key in keys) + + +@pytest.mark.parametrize("async_", [False, True]) +@pytest.mark.parametrize("raw", [False, True]) +def test_callable_credentials_survive_stream_creation(tracing, async_, raw): + if tuple(map(int, openai.__version__.split(".")[:2])) < (3, 0): + pytest.skip("callable OpenAI credentials are a latest-SDK feature") + provider, exporter = tracing + OpenAIInstrumentor().instrument(tracer_provider=provider) + calls = [] + + def token(): + calls.append(1) + return KEY + + async def atoken(): + return token() + + response = body("failed") + response["error"] = {"code": "server_error", "message": "Rejected " + KEY} + + def handler(request): + assert request.headers["authorization"] == "Bearer " + KEY + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=sse("response.failed", response), + ) + + async def run(): + async with client(handler, True, atoken) as c: + api = c.responses.with_raw_response if raw else c.responses + result = await api.create(model="parallel", input="question", stream=True) + stream = result.parse() if raw else result + assert KEY in (await stream.__anext__()).response.error.message + await stream.close() + + if async_: + asyncio.run(run()) + else: + with client(handler, key=token) as c: + api = c.responses.with_raw_response if raw else c.responses + result = api.create(model="parallel", input="question", stream=True) + stream = result.parse() if raw else result + assert KEY in next(stream).response.error.message + stream.close() + assert len(calls) == 1 + (span,) = exporter.get_finished_spans() + assert span.status.status_code is StatusCode.ERROR + assert KEY not in span.to_json() diff --git a/python/frameworks/openai/tests/test_responses_safety.py b/python/frameworks/openai/tests/test_responses_safety.py index f6a5be333..fb2d773ff 100644 --- a/python/frameworks/openai/tests/test_responses_safety.py +++ b/python/frameworks/openai/tests/test_responses_safety.py @@ -1,5 +1,6 @@ """Responses privacy and lifecycle contracts with real SDKs, no vendor traffic.""" import asyncio +import inspect import json import httpx @@ -394,6 +395,10 @@ async def run(): @pytest.mark.parametrize("async_", [False, True]) @pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.skipif( + "hide_error_details" not in inspect.signature(TraceConfig).parameters, + reason="Error-detail masking requires a core with hide_error_details support", +) def test_hidden_error_details_apply_to_failed_response_json(tracing, async_, stream): provider, exporter = tracing instrumentor = OpenAIInstrumentor() diff --git a/python/frameworks/openai/traceai_openai/_privacy.py b/python/frameworks/openai/traceai_openai/_privacy.py index ae4a66e9f..98521c73e 100644 --- a/python/frameworks/openai/traceai_openai/_privacy.py +++ b/python/frameworks/openai/traceai_openai/_privacy.py @@ -46,6 +46,19 @@ def request_credentials(instance: Any, options: Any, parameters: Mapping) -> tup return tuple(sorted(secrets, key=len, reverse=True)) +def response_credentials(response: Any) -> tuple[str, ...]: + """Use the request actually sent, not a shared client's last refreshed token. + + Raw responses and streams retain their HTTP response. Reading its request + headers never calls the user's token provider or modifies authentication. + """ + http_response = getattr(response, "http_response", None) + if http_response is None: + http_response = getattr(response, "response", None) + request = getattr(http_response, "request", None) + return request_credentials(None, request, {}) + + def safe_tool(tool: Any) -> Any: if isinstance(tool, Mapping) and tool.get("type") == "mcp": return {key: value for key, value in tool.items() if key not in {"authorization", "headers"}} @@ -71,4 +84,21 @@ def response_output(response: Any, config: TraceConfig | None = None) -> dict: "error", "incomplete_details"} if config is not None and getattr(config, "hide_error_details", False): fields.difference_update({"error", "incomplete_details"}) - return response.model_dump(include=fields, exclude_none=True, mode="json") + output = response.model_dump(include=fields, exclude_none=True, mode="json") + if config is not None and getattr(config, "hide_error_details", False): + # These provider-owned error fields are not ordinary tool results or + # assistant text. Work on the dump only, leaving SDK objects untouched. + for item in output.get("output", []): + if not isinstance(item, dict) or item.get("type") not in { + "mcp_call", "mcp_list_tools", "image_generation_call", + }: + continue + error = item.pop("error", None) + if isinstance(error, dict) and error.get("type") in { + "http_error", "mcp_protocol_error", "mcp_tool_execution_error", + }: + classification = {"type": error["type"]} + if type(error.get("code")) is int: + classification["code"] = error["code"] + item["error"] = classification + return output diff --git a/python/frameworks/openai/traceai_openai/_request.py b/python/frameworks/openai/traceai_openai/_request.py index 9e06d500e..44d59b227 100644 --- a/python/frameworks/openai/traceai_openai/_request.py +++ b/python/frameworks/openai/traceai_openai/_request.py @@ -18,7 +18,7 @@ from opentelemetry.util.types import AttributeValue from traceai_openai._client_provider import _get_client_provider from traceai_openai._providers import provider_from_host -from traceai_openai._privacy import redact, request_credentials, safe_request +from traceai_openai._privacy import redact, request_credentials, response_credentials, safe_request from traceai_openai._request_attributes_extractor import _RequestAttributesExtractor from traceai_openai._response_accumulator import ( _ChatCompletionAccumulator, @@ -28,7 +28,7 @@ ) from traceai_openai._response_attributes_extractor import _ResponseAttributesExtractor from traceai_openai._span_io_handler import add_io_to_span_attributes -from traceai_openai._stream import _AsyncStream, _ResponseAccumulator, _Stream +from traceai_openai._stream import _AsyncStream, _RawResponse, _ResponseAccumulator, _Stream from traceai_openai._utils import ( _as_input_attributes, _finish_tracing, @@ -218,9 +218,9 @@ def _finalize_response( except Exception: logger.exception(f"Failed to parse response of type {type(response)}") - if self._is_streaming(response): - # Determine if the response is asynchronous - if self._is_async_stream(response): + if self._is_streaming(parsed_response): + # Raw-response parse() returns the stream, not the outer envelope. + if self._is_async_stream(parsed_response): stream_wrapper = _AsyncStream else: stream_wrapper = _Stream @@ -237,11 +237,13 @@ def _finalize_response( # Wrap the response with the appropriate stream wrapper wrapped_stream = stream_wrapper( - stream=response, + stream=parsed_response, with_span=with_span, response_accumulator=response_accumulator, ) + if parsed_response is not response: + return _RawResponse(response, parsed_response, wrapped_stream) return wrapped_stream else: @@ -328,6 +330,8 @@ def __call__( add_io_to_span_attributes(with_span, input_data, None) try: response = wrapped(*args, **kwargs) + with_span.add_secrets(request_credentials(instance, None, {})) + with_span.add_secrets(response_credentials(response)) if isinstance(response, CreateEmbeddingResponse): with_span.set_attribute( @@ -351,6 +355,7 @@ def __call__( request_parameters=request_parameters, ) except BaseException as e: + with_span.add_secrets(request_credentials(instance, getattr(e, "request", None), {})) status = trace_api.Status( status_code=trace_api.StatusCode.ERROR, description=f"{type(e).__name__}: {e}", @@ -409,6 +414,8 @@ async def __call__( add_io_to_span_attributes(with_span, input_data, None) try: response = await wrapped(*args, **kwargs) + with_span.add_secrets(request_credentials(instance, None, {})) + with_span.add_secrets(response_credentials(response)) if isinstance(response, CreateEmbeddingResponse): with_span.set_attribute( @@ -421,6 +428,9 @@ async def __call__( config=self._config, ) except BaseException as exception: + with_span.add_secrets(request_credentials( + instance, getattr(exception, "request", None), {} + )) with_span.record_exception(exception) status = trace_api.Status( status_code=trace_api.StatusCode.ERROR, diff --git a/python/frameworks/openai/traceai_openai/_stream.py b/python/frameworks/openai/traceai_openai/_stream.py index 20dd5a5c4..393977e46 100644 --- a/python/frameworks/openai/traceai_openai/_stream.py +++ b/python/frameworks/openai/traceai_openai/_stream.py @@ -72,6 +72,23 @@ async def aclose(self) -> None: _finish_on_exit(self._self_stream, error) +class _RawResponse(ObjectProxy): + """Preserve the raw SDK envelope and parse cache while tracing its stream.""" + + def __init__(self, response: Any, parsed_stream: Any, stream: Any) -> None: + super().__init__(response) + self._self_parsed_stream = parsed_stream + self._self_stream = stream + + def parse(self, *args: Any, **kwargs: Any) -> Any: + parsed = self.__wrapped__.parse(*args, **kwargs) + return self._self_stream if parsed is self._self_parsed_stream else parsed + + @property + def http_response(self) -> Any: + return self._self_stream.response + + class _Stream(ObjectProxy): # type: ignore __slots__ = ( "_self_with_span", diff --git a/python/frameworks/openai/traceai_openai/_with_span.py b/python/frameworks/openai/traceai_openai/_with_span.py index 8efb36318..c3b336328 100644 --- a/python/frameworks/openai/traceai_openai/_with_span.py +++ b/python/frameworks/openai/traceai_openai/_with_span.py @@ -54,6 +54,11 @@ def __init__( def is_finished(self) -> bool: return self._is_finished + def add_secrets(self, secrets: Tuple[str, ...]) -> None: + # Credentials resolved by an SDK token provider are available only after + # sending the request. Retain all known values for this span's lifetime. + self._secrets = tuple(sorted(set((*self._secrets, *secrets)), key=len, reverse=True)) + def record_exception(self, exception: BaseException) -> None: if self._is_finished: return @@ -110,9 +115,15 @@ def finish_tracing( self.set_attribute(key, value) if status is not None: try: - self._span.set_status(status=trace_api.Status( - status.status_code, redact(status.description, self._secrets) - )) + description = redact(status.description, self._secrets) + if description: + # Core removes tracebacks by consulting the in-flight SDK + # exception. Its original message may contain credentials. + # Remove the traceback here so core cannot reconstruct that + # unsanitized message; the redacted event keeps the stack. + description = description.partition("Traceback (most recent call last)")[0] + description = description.rstrip() or None + self._span.set_status(status=trace_api.Status(status.status_code, description)) except Exception: logger.exception("Failed to set status code on span") try: From a37ce8bf66fd4e42e4d7c4f0a585b8b9e09b1847 Mon Sep 17 00:00:00 2001 From: Nikhil Pareek Date: Sun, 1 Nov 2026 23:29:58 -0800 Subject: [PATCH 03/10] fix(openai): preserve request credentials and lazy response lifecycles Capture actual HTTP request credentials in a per-call context before SDK parsing. Preserve lazy async envelope parsing, finalize envelope closes, and trace every explicit stream parse cache entry with a shared terminal result. Co-Authored-By: Claude Code --- .../openai/tests/test_response_envelopes.py | 455 ++++++++++++++++++ .../openai/traceai_openai/__init__.py | 13 + .../openai/traceai_openai/_request.py | 123 ++--- .../openai/traceai_openai/_request_context.py | 41 ++ .../openai/traceai_openai/_stream.py | 90 +++- 5 files changed, 641 insertions(+), 81 deletions(-) create mode 100644 python/frameworks/openai/tests/test_response_envelopes.py create mode 100644 python/frameworks/openai/traceai_openai/_request_context.py diff --git a/python/frameworks/openai/tests/test_response_envelopes.py b/python/frameworks/openai/tests/test_response_envelopes.py new file mode 100644 index 000000000..593bd62c6 --- /dev/null +++ b/python/frameworks/openai/tests/test_response_envelopes.py @@ -0,0 +1,455 @@ +"""Independent review probes; all credentials and transports are synthetic.""" + +import asyncio +import json + +import httpx +import openai +import pytest +from fi_instrumentation import register +from fi_instrumentation.fi_types import ProjectType +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import StatusCode +from traceai_openai import OpenAIInstrumentor + + +@pytest.fixture +def tracing(): + exporter = InMemorySpanExporter() + provider = register( + project_name="review-edges", + project_type=ProjectType.OBSERVE, + span_exporter=exporter, + batch=False, + verbose=False, + ) + OpenAIInstrumentor().instrument(tracer_provider=provider) + try: + yield provider, exporter + finally: + OpenAIInstrumentor().uninstrument() + provider.force_flush() + provider.shutdown() + + +def client(handler, async_=False, key="synthetic-static-key"): + http = (httpx.AsyncClient if async_ else httpx.Client)( + transport=httpx.MockTransport(handler), trust_env=False + ) + return (openai.AsyncOpenAI if async_ else openai.OpenAI)( + api_key=key, + base_url="https://api.parallel.ai/v1", + max_retries=0, + http_client=http, + ) + + +def body(status="completed"): + return { + "id": "resp_edge", + "object": "response", + "created_at": 1, + "model": "parallel", + "status": status, + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "output": [], + "error": None, + "usage": { + "input_tokens": 5, + "output_tokens": 2, + "total_tokens": 7, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens_details": {"reasoning_tokens": 0}, + }, + } + + +def sse(status): + response = body(status) + if status == "failed": + response["error"] = {"code": "server_error", "message": "synthetic failure"} + kind = "response.created" if status == "in_progress" else "response." + status + return ( + "event: " + + kind + + "\ndata: " + + json.dumps({"type": kind, "sequence_number": 1, "response": response}) + + "\n\n" + ).encode() + + +def test_concurrent_callable_nonstream_credentials(tracing): + if int(openai.__version__.split(".")[0]) < 3: + pytest.skip("callable auth not available at floor") + _provider, exporter = tracing + keys = ["synthetic-alpha-credential-198273", "synthetic-bravo-credential-543215"] + issued = [] + + async def run(): + received = [] + both_sent = asyncio.Event() + + async def token(): + key = keys[len(issued)] + issued.append(key) + return key + + async def handler(request): + key = request.headers["authorization"][7:] + received.append(key) + if len(received) == 2: + both_sent.set() + await asyncio.wait_for(both_sent.wait(), 3) + response = body("failed") + response["error"] = { + "code": "server_error", + "message": "Rejected token " + key, + } + return httpx.Response(200, json=response) + + async with client(handler, True, token) as c: + responses = await asyncio.gather( + *[c.responses.create(model="parallel", input="question") for _ in keys] + ) + assert set(received) == set(keys) + assert {r.error.message for r in responses} == { + "Rejected token " + k for k in keys + } + + asyncio.run(run()) + assert issued == keys + spans = exporter.get_finished_spans() + assert len(spans) == 2 + assert all(s.status.status_code == StatusCode.ERROR for s in spans) + leaked = { + key: [ + name + for span in spans + for name, value in span.attributes.items() + if key in str(value) + ] + for key in keys + } + assert not any(leaked.values()), leaked + + +@pytest.mark.parametrize("async_", [False, True]) +@pytest.mark.parametrize("status", ["failed", "completed", "in_progress"]) +@pytest.mark.parametrize("exhaust", [False, True]) +def test_streaming_response_envelope_lifecycle(tracing, async_, status, exhaust): + provider, exporter = tracing + + def handler(request): + return httpx.Response( + 200, headers={"content-type": "text/event-stream"}, content=sse(status) + ) + + async def run(): + async with client(handler, True) as c: + async with c.responses.with_streaming_response.create( + model="parallel", input="question", stream=True + ) as raw: + stream = await raw.parse() + if exhaust: + events = [chunk async for chunk in stream] + assert events[-1].response.status == status + else: + assert (await stream.__anext__()).response.status == status + + if async_: + asyncio.run(run()) + else: + with client(handler) as c: + with c.responses.with_streaming_response.create( + model="parallel", input="question", stream=True + ) as raw: + stream = raw.parse() + if exhaust: + events = list(stream) + assert events[-1].response.status == status + else: + assert next(stream).response.status == status + provider.force_flush() + spans = exporter.get_finished_spans() + assert len(spans) == 1, f"Expected one finalized span, got {len(spans)}" + (span,) = spans + assert span.status.status_code == ( + StatusCode.OK if status == "completed" else StatusCode.ERROR + ), span.to_json() + if status != "in_progress": + assert span.attributes["gen_ai.usage.total_tokens"] == 7, span.to_json() + + +@pytest.mark.parametrize("async_", [False, True]) +@pytest.mark.parametrize("instrumented", [False, True]) +def test_raw_explicit_stream_parse(tracing, async_, instrumented): + provider, exporter = tracing + if not instrumented: + OpenAIInstrumentor().uninstrument() + from openai.types.responses import ResponseStreamEvent + + def handler(request): + return httpx.Response( + 200, headers={"content-type": "text/event-stream"}, content=sse("failed") + ) + + async def run(): + async with client(handler, True) as c: + raw = await c.responses.with_raw_response.create( + model="parallel", input="question", stream=True + ) + stream = raw.parse(to=openai.AsyncStream[ResponseStreamEvent]) + assert raw.parse(to=openai.AsyncStream[ResponseStreamEvent]) is stream + events = [event async for event in stream] + assert events[-1].response.status == "failed" + await stream.close() + + if async_: + asyncio.run(run()) + else: + with client(handler) as c: + raw = c.responses.with_raw_response.create( + model="parallel", input="question", stream=True + ) + stream = raw.parse(to=openai.Stream[ResponseStreamEvent]) + assert raw.parse(to=openai.Stream[ResponseStreamEvent]) is stream + events = list(stream) + assert events[-1].response.status == "failed" + stream.close() + provider.force_flush() + spans = exporter.get_finished_spans() + if instrumented: + assert len(spans) == 1, f"Expected one finalized span, got {len(spans)}" + assert spans[0].status.status_code is StatusCode.ERROR + assert spans[0].attributes["gen_ai.usage.total_tokens"] == 7 + else: + assert not spans + + +@pytest.mark.parametrize("async_", [False, True]) +@pytest.mark.parametrize("streaming", [False, True]) +@pytest.mark.parametrize("parse", [False, True]) +def test_lazy_envelope_does_not_read_until_requested(tracing, async_, streaming, parse): + provider, exporter = tracing + reads = [] + closes = [] + payload = sse("completed") if streaming else json.dumps(body()).encode() + + class SyncBody(httpx.SyncByteStream): + def __iter__(self): + reads.append("read") + yield payload + + def close(self): + closes.append("close") + + class AsyncBody(httpx.AsyncByteStream): + async def __aiter__(self): + reads.append("read") + yield payload + + async def aclose(self): + closes.append("close") + + def handler(request): + return httpx.Response( + 200, + headers={ + "content-type": "text/event-stream" if streaming else "application/json" + }, + stream=AsyncBody() if async_ else SyncBody(), + ) + + async def run(): + async with client(handler, True) as c: + async with c.responses.with_streaming_response.create( + model="parallel", input="question", stream=streaming + ) as raw: + assert isinstance(raw, openai.AsyncAPIResponse) + assert not reads + assert not exporter.get_finished_spans() + if parse: + result = await raw.parse() + assert await raw.parse() is result + if streaming: + assert not reads + assert (await result.__anext__()).response.status == "completed" + else: + assert result.status == "completed" + + if async_: + asyncio.run(run()) + else: + with client(handler) as c: + with c.responses.with_streaming_response.create( + model="parallel", input="question", stream=streaming + ) as raw: + assert isinstance(raw, openai.APIResponse) + assert not reads + assert not exporter.get_finished_spans() + if parse: + result = raw.parse() + assert raw.parse() is result + if streaming: + assert not reads + assert next(result).response.status == "completed" + else: + assert result.status == "completed" + assert reads == (["read"] if parse else []) + assert closes == ["close"] + provider.force_flush() + (span,) = exporter.get_finished_spans() + assert span.status.status_code == (StatusCode.OK if parse else StatusCode.ERROR) + if parse: + assert span.attributes["gen_ai.usage.total_tokens"] == 7 + + +@pytest.mark.parametrize("async_", [False, True]) +@pytest.mark.parametrize("modern", [False, True]) +@pytest.mark.parametrize("close_http", [False, True]) +def test_explicit_stream_cache_and_envelope_close(tracing, async_, modern, close_http): + provider, exporter = tracing + from openai.types.responses import ResponseStreamEvent + + def handler(request): + return httpx.Response( + 200, headers={"content-type": "text/event-stream"}, content=sse("completed") + ) + + async def consume(raw): + if modern: + default = await raw.parse() + typed = await raw.parse(to=openai.AsyncStream[ResponseStreamEvent]) + assert await raw.parse(to=openai.AsyncStream[ResponseStreamEvent]) is typed + assert await raw.parse() is default + else: + default = raw.parse() + typed = raw.parse(to=openai.AsyncStream[ResponseStreamEvent]) + assert raw.parse(to=openai.AsyncStream[ResponseStreamEvent]) is typed + assert raw.parse() is default + assert typed is not default + assert (await typed.__anext__()).response.status == "completed" + assert not exporter.get_finished_spans() + if close_http: + await raw.http_response.aclose() + elif modern: + await raw.close() + else: + await typed.close() + # Both cache entries and every close path share one terminal result. + await default.close() + await typed.close() + + async def run(): + async with client(handler, True) as c: + if modern: + async with c.responses.with_streaming_response.create( + model="parallel", input="question", stream=True + ) as raw: + await consume(raw) + else: + raw = await c.responses.with_raw_response.create( + model="parallel", input="question", stream=True + ) + await consume(raw) + + def consume_sync(raw): + default = raw.parse() + typed = raw.parse(to=openai.Stream[ResponseStreamEvent]) + assert raw.parse(to=openai.Stream[ResponseStreamEvent]) is typed + assert raw.parse() is default + assert typed is not default + assert next(typed).response.status == "completed" + assert not exporter.get_finished_spans() + if close_http: + raw.http_response.close() + elif modern: + raw.close() + else: + typed.close() + default.close() + typed.close() + + if async_: + asyncio.run(run()) + else: + with client(handler) as c: + if modern: + with c.responses.with_streaming_response.create( + model="parallel", input="question", stream=True + ) as raw: + consume_sync(raw) + else: + raw = c.responses.with_raw_response.create( + model="parallel", input="question", stream=True + ) + consume_sync(raw) + provider.force_flush() + (span,) = exporter.get_finished_spans() + assert span.status.status_code is StatusCode.OK + assert span.attributes["gen_ai.usage.total_tokens"] == 7 + + +@pytest.mark.parametrize("cancel", [False, True]) +def test_async_lazy_parse_failure_finishes_once(tracing, cancel): + provider, exporter = tracing + entered = None + + class BrokenBody(httpx.AsyncByteStream): + async def __aiter__(self): + entered.set() + if cancel: + await asyncio.Event().wait() + raise httpx.ReadError("synthetic unreadable response") + yield b"" # Make the method an async generator. + + def handler(request): + return httpx.Response( + 200, headers={"content-type": "application/json"}, stream=BrokenBody() + ) + + async def run(): + nonlocal entered + entered = asyncio.Event() + async with client(handler, True) as c: + async with c.responses.with_streaming_response.create( + model="parallel", input="question" + ) as raw: + task = asyncio.create_task(raw.parse()) + await asyncio.wait_for(entered.wait(), 3) + if cancel: + task.cancel() + with pytest.raises( + asyncio.CancelledError if cancel else httpx.ReadError + ): + await task + assert len(exporter.get_finished_spans()) == 1 + + asyncio.run(run()) + provider.force_flush() + (span,) = exporter.get_finished_spans() + assert span.status.status_code is StatusCode.ERROR + + +def test_credential_context_is_reset_and_hooks_restore(tracing): + from traceai_openai._request_context import _active_request + + original_sync = OpenAIInstrumentor()._original_process_response + original_async = OpenAIInstrumentor()._original_async_process_response + assert _active_request.get() is None + with client(lambda request: httpx.Response(200, json=body())) as c: + c.responses.create(model="parallel", input="question") + assert _active_request.get() is None + with client( + lambda request: httpx.Response( + 401, json={"error": {"message": "synthetic error"}} + ) + ) as c: + with pytest.raises(openai.AuthenticationError): + c.responses.create(model="parallel", input="question") + assert _active_request.get() is None + OpenAIInstrumentor().uninstrument() + assert openai.OpenAI._process_response is original_sync + assert openai.AsyncOpenAI._process_response is original_async diff --git a/python/frameworks/openai/traceai_openai/__init__.py b/python/frameworks/openai/traceai_openai/__init__.py index 891928712..e7633a2a3 100644 --- a/python/frameworks/openai/traceai_openai/__init__.py +++ b/python/frameworks/openai/traceai_openai/__init__.py @@ -16,6 +16,7 @@ from opentelemetry.instrumentation.instrumentor import BaseInstrumentor # type: ignore from traceai_openai._client_provider import set_client_provider as set_client_provider from traceai_openai._request import _AsyncRequest, _Request +from traceai_openai._request_context import capture_response_credentials from traceai_openai.package import _instruments from traceai_openai.version import __version__ from wrapt import wrap_function_wrapper @@ -34,6 +35,8 @@ class OpenAIInstrumentor(BaseInstrumentor): # type: ignore __slots__ = ( "_original_request", "_original_async_request", + "_original_process_response", + "_original_async_process_response", ) def instrumentation_dependencies(self) -> Collection[str]: @@ -63,6 +66,14 @@ def _instrument(self, **kwargs: Any) -> None: name="AsyncOpenAI.request", wrapper=_AsyncRequest(tracer=tracer, openai=openai, config=config), ) + self._original_process_response = openai.OpenAI._process_response + self._original_async_process_response = openai.AsyncOpenAI._process_response + for client in ("OpenAI", "AsyncOpenAI"): + wrap_function_wrapper( + module=_MODULE, + name=f"{client}._process_response", + wrapper=capture_response_credentials, + ) if Protect is not None: self._original_protect = Protect.protect wrap_function_wrapper( @@ -77,3 +88,5 @@ def _uninstrument(self, **kwargs: Any) -> None: openai = import_module(_MODULE) openai.OpenAI.request = self._original_request openai.AsyncOpenAI.request = self._original_async_request + openai.OpenAI._process_response = self._original_process_response + openai.AsyncOpenAI._process_response = self._original_async_process_response diff --git a/python/frameworks/openai/traceai_openai/_request.py b/python/frameworks/openai/traceai_openai/_request.py index 44d59b227..694db0276 100644 --- a/python/frameworks/openai/traceai_openai/_request.py +++ b/python/frameworks/openai/traceai_openai/_request.py @@ -1,6 +1,7 @@ import logging from abc import ABC from contextlib import contextmanager +from inspect import iscoroutinefunction from itertools import chain from types import ModuleType from typing import Any, Awaitable, Callable, Iterable, Iterator, Mapping, Optional, Tuple @@ -10,6 +11,7 @@ FiSpanKindValues, SpanAttributes, ) +from openai._legacy_response import LegacyAPIResponse from openai.types.create_embedding_response import CreateEmbeddingResponse from opentelemetry import context as context_api from opentelemetry import trace as trace_api @@ -19,6 +21,7 @@ from traceai_openai._client_provider import _get_client_provider from traceai_openai._providers import provider_from_host from traceai_openai._privacy import redact, request_credentials, response_credentials, safe_request +from traceai_openai._request_context import capture_request_credentials from traceai_openai._request_attributes_extractor import _RequestAttributesExtractor from traceai_openai._response_accumulator import ( _ChatCompletionAccumulator, @@ -28,7 +31,13 @@ ) from traceai_openai._response_attributes_extractor import _ResponseAttributesExtractor from traceai_openai._span_io_handler import add_io_to_span_attributes -from traceai_openai._stream import _AsyncStream, _RawResponse, _ResponseAccumulator, _Stream +from traceai_openai._stream import ( + _AsyncRawResponse, + _AsyncStream, + _RawResponse, + _ResponseAccumulator, + _Stream, +) from traceai_openai._utils import ( _as_input_attributes, _finish_tracing, @@ -211,52 +220,50 @@ def _finalize_response( Monkey-patch the response object to trace the stream, or finish tracing if the response is not a stream. """ - parsed_response = response - if hasattr(response, "parse") and callable(response.parse): - try: - parsed_response = response.parse() - except Exception: - logger.exception(f"Failed to parse response of type {type(response)}") - - if self._is_streaming(parsed_response): - # Raw-response parse() returns the stream, not the outer envelope. - if self._is_async_stream(parsed_response): - stream_wrapper = _AsyncStream - else: - stream_wrapper = _Stream - - # Get the response accumulator if available - response_accumulator_factory = self._response_accumulator_factories.get( - cast_to - ) - response_accumulator = ( - response_accumulator_factory(request_parameters) - if response_accumulator_factory - else None - ) - - # Wrap the response with the appropriate stream wrapper - wrapped_stream = stream_wrapper( - stream=parsed_response, - with_span=with_span, - response_accumulator=response_accumulator, - ) - - if parsed_response is not response: - return _RawResponse(response, parsed_response, wrapped_stream) - return wrapped_stream + response_accumulator_factory = self._response_accumulator_factories.get(cast_to) + response_accumulator = ( + response_accumulator_factory(request_parameters) + if response_accumulator_factory + else None + ) - else: + def on_parse(parsed: Any) -> Any: + if self._is_streaming(parsed): + stream_wrapper = _AsyncStream if self._is_async_stream(parsed) else _Stream + return stream_wrapper( + stream=parsed, + with_span=with_span, + response_accumulator=response_accumulator, + ) + if isinstance(parsed, CreateEmbeddingResponse): + with_span.set_attribute( + SpanAttributes.EMBEDDING_EMBEDDINGS, self._raw_embeddings_json(parsed) + ) + add_io_to_span_attributes(with_span, None, parsed, config=self._config) _finish_tracing( - status=_response_status(parsed_response), + status=_response_status(parsed), with_span=with_span, has_attributes=_ResponseAttributes( request_parameters=request_parameters, - response=response, + response=parsed, response_attributes_extractor=self._response_attributes_extractor, ), ) - return response + return parsed + + if hasattr(response, "parse") and callable(response.parse): + # Modern envelopes are deliberately lazy, even for non-SSE responses. + # Do not call async parse from sync I/O/attribute extraction paths. + raw_wrapper = ( + _AsyncRawResponse if iscoroutinefunction(response.parse) else _RawResponse + ) + raw = raw_wrapper(response, with_span, response_accumulator, on_parse) + if isinstance(response, LegacyAPIResponse): + # Preserve eager legacy non-stream telemetry. Stream parse caches + # remain SDK-owned, and every later explicit parse type is traced. + raw.parse() + return raw + return on_parse(response) def _raw_embeddings_json(self, response: CreateEmbeddingResponse) -> str: """The raw embeddings response; vectors are replaced by their dimension @@ -329,25 +336,11 @@ def __call__( ) add_io_to_span_attributes(with_span, input_data, None) try: - response = wrapped(*args, **kwargs) + with capture_request_credentials(instance, with_span): + response = wrapped(*args, **kwargs) with_span.add_secrets(request_credentials(instance, None, {})) with_span.add_secrets(response_credentials(response)) - if isinstance(response, CreateEmbeddingResponse): - with_span.set_attribute( - SpanAttributes.EMBEDDING_EMBEDDINGS, - self._raw_embeddings_json(response), - ) - - # Add output data to span attributes after getting response - if not self._is_streaming(response): - add_io_to_span_attributes( - with_span, - None, - response, - is_streaming=self._is_streaming(response), - config=self._config, - ) return self._finalize_response( response=response, with_span=with_span, @@ -413,20 +406,10 @@ async def __call__( ) add_io_to_span_attributes(with_span, input_data, None) try: - response = await wrapped(*args, **kwargs) + with capture_request_credentials(instance, with_span): + response = await wrapped(*args, **kwargs) with_span.add_secrets(request_credentials(instance, None, {})) with_span.add_secrets(response_credentials(response)) - - if isinstance(response, CreateEmbeddingResponse): - with_span.set_attribute( - SpanAttributes.EMBEDDING_EMBEDDINGS, - self._raw_embeddings_json(response), - ) - - add_io_to_span_attributes( - with_span, None, response, is_streaming=self._is_streaming(response), - config=self._config, - ) except BaseException as exception: with_span.add_secrets(request_credentials( instance, getattr(exception, "request", None), {} @@ -491,12 +474,6 @@ def __init__( request_parameters: Mapping[str, Any], response_attributes_extractor: _ResponseAttributesExtractor, ) -> None: - if hasattr(response, "parse") and callable(response.parse): - # E.g. see https://github.com/openai/openai-python/blob/f1c7d714914e3321ca2e72839fe2d132a8646e7f/src/openai/_base_client.py#L518 # noqa: E501 - try: - response = response.parse() - except Exception: - logger.exception(f"Failed to parse response of type {type(response)}") self._request_parameters = request_parameters self._response = response self._response_attributes_extractor = response_attributes_extractor diff --git a/python/frameworks/openai/traceai_openai/_request_context.py b/python/frameworks/openai/traceai_openai/_request_context.py new file mode 100644 index 000000000..c81131cf7 --- /dev/null +++ b/python/frameworks/openai/traceai_openai/_request_context.py @@ -0,0 +1,41 @@ +"""Capture request-local credentials before the SDK discards its HTTP envelope.""" + +import logging +from contextlib import contextmanager +from contextvars import ContextVar +from typing import Any, Iterator, Optional, Tuple + +from traceai_openai._privacy import request_credentials +from traceai_openai._with_span import _WithSpan + +logger = logging.getLogger(__name__) +_active_request: ContextVar[Optional[Tuple[Any, _WithSpan]]] = ContextVar( + "traceai_openai_request", default=None +) + + +@contextmanager +def capture_request_credentials(client: Any, span: _WithSpan) -> Iterator[None]: + token = _active_request.set((client, span)) + try: + yield + finally: + _active_request.reset(token) + + +def capture_response_credentials( + wrapped: Any, instance: Any, args: Any, kwargs: Any +) -> Any: + # Both SDK bounds pass response by keyword to _process_response, before + # parsing removes the request. ContextVars isolate concurrent calls without + # modifying shared clients/transports or invoking a token provider again. + active = _active_request.get() + if active is not None and active[0] is instance: + try: + response = kwargs.get("response") + active[1].add_secrets( + request_credentials(None, getattr(response, "request", None), {}) + ) + except Exception: + logger.warning("Failed to capture request credentials") + return wrapped(*args, **kwargs) diff --git a/python/frameworks/openai/traceai_openai/_stream.py b/python/frameworks/openai/traceai_openai/_stream.py index 393977e46..9829860a3 100644 --- a/python/frameworks/openai/traceai_openai/_stream.py +++ b/python/frameworks/openai/traceai_openai/_stream.py @@ -1,5 +1,5 @@ import logging -from typing import Any, AsyncIterator, Generator, Iterator, Optional, Protocol, Tuple +from typing import Any, AsyncIterator, Callable, Generator, Iterator, Optional, Protocol, Tuple from openai.types.chat import ChatCompletionChunk from opentelemetry import trace as trace_api @@ -73,20 +73,94 @@ async def aclose(self) -> None: class _RawResponse(ObjectProxy): - """Preserve the raw SDK envelope and parse cache while tracing its stream.""" + """Trace every parse-cache entry without eagerly reading a lazy envelope.""" - def __init__(self, response: Any, parsed_stream: Any, stream: Any) -> None: + def __init__( + self, + response: Any, + with_span: _WithSpan, + response_accumulator: Optional[_ResponseAccumulator], + on_parse: Callable[[Any], Any], + ) -> None: super().__init__(response) - self._self_parsed_stream = parsed_stream - self._self_stream = stream + self._self_with_span = with_span + self._self_response_accumulator = response_accumulator + self._self_on_parse = on_parse + self._self_parsed = {} + + @property + def _self_is_finished(self) -> bool: + return self._self_with_span.is_finished + + def _trace_parsed(self, parsed: Any) -> Any: + # The SDK owns cache keys (including generic stream types). Cache by + # identity and retain the original through the proxy, so repeated parse + # calls return the same wrapper without collapsing distinct SDK entries. + key = id(parsed) + if key not in self._self_parsed: + self._self_parsed[key] = self._self_on_parse(parsed) + return self._self_parsed[key] def parse(self, *args: Any, **kwargs: Any) -> Any: - parsed = self.__wrapped__.parse(*args, **kwargs) - return self._self_stream if parsed is self._self_parsed_stream else parsed + try: + parsed = self.__wrapped__.parse(*args, **kwargs) + except BaseException as exc: + _finish_on_exit(self, exc) + raise + return self._trace_parsed(parsed) @property def http_response(self) -> Any: - return self._self_stream.response + return _ResponseCloseProxy(self.__wrapped__.http_response, self) + + def close(self) -> None: + error = None + try: + self.__wrapped__.close() + except BaseException as exc: + error = exc + raise + finally: + _finish_on_exit(self, error) + + def _finish_tracing(self, status: Optional[trace_api.Status] = None) -> None: + if self._self_is_finished: + return + accumulator = self._self_response_accumulator + if status is None or status.status_code is not trace_api.StatusCode.ERROR: + get_status = getattr(accumulator, "get_status", None) + if get_status is not None: + status = get_status() + elif not self._self_parsed: + status = trace_api.Status( + trace_api.StatusCode.ERROR, "Response closed before parsing" + ) + if accumulator is not None: + _finish_tracing(self._self_with_span, accumulator, status) + else: + self._self_with_span.finish_tracing(status=status) + + +class _AsyncRawResponse(_RawResponse): + """AsyncAPIResponse has awaitable parse/close; LegacyAPIResponse does not.""" + + async def parse(self, *args: Any, **kwargs: Any) -> Any: + try: + parsed = await self.__wrapped__.parse(*args, **kwargs) + except BaseException as exc: + _finish_on_exit(self, exc) + raise + return self._trace_parsed(parsed) + + async def close(self) -> None: + error = None + try: + await self.__wrapped__.close() + except BaseException as exc: + error = exc + raise + finally: + _finish_on_exit(self, error) class _Stream(ObjectProxy): # type: ignore From f180afe58f6ec3c1394bcb77368e14d5ef565dc5 Mon Sep 17 00:00:00 2001 From: Nikhil Pareek Date: Mon, 2 Nov 2026 03:35:52 -0800 Subject: [PATCH 04/10] fix(openai): mark interrupted agents and truncated chat streams as errors Handle non-Exception aborts through the privacy-aware core tracer. Require terminal finish reasons for every requested chat choice before marking a stream successful. Reuse the stream status hook from PR #520 without duplicating its cancellation and close fixes. Co-Authored-By: Claude Code --- .../instrumentation/_tracers.py | 16 ++- .../tests/test_chat_stream_completion.py | 121 ++++++++++++++++++ .../tests/test_stream_context_manager.py | 4 +- .../traceai_openai/_response_accumulator.py | 25 ++++ .../openai/traceai_openai/_stream.py | 12 ++ python/tests/test_tracer_cancellation.py | 79 ++++++++++++ 6 files changed, 254 insertions(+), 3 deletions(-) create mode 100644 python/frameworks/openai/tests/test_chat_stream_completion.py create mode 100644 python/tests/test_tracer_cancellation.py diff --git a/python/fi_instrumentation/instrumentation/_tracers.py b/python/fi_instrumentation/instrumentation/_tracers.py index 779854676..24f2db811 100644 --- a/python/fi_instrumentation/instrumentation/_tracers.py +++ b/python/fi_instrumentation/instrumentation/_tracers.py @@ -203,7 +203,21 @@ def start_as_current_span( record_exception=record_exception, set_status_on_exception=set_status_on_exception, ) as current_span: - yield cast(FiSpan, current_span) + try: + yield cast(FiSpan, current_span) + except BaseException as error: + # OTel use_span handles Exception, but deliberately excludes + # cancellation, KeyboardInterrupt and other local aborts. Record + # those before use_span ends the span; leave ordinary errors to + # OTel so they are not recorded twice. FiSpan applies privacy. + if not isinstance(error, Exception) and span.is_recording(): + if record_exception: + span.record_exception(error, escaped=True) + if set_status_on_exception: + span.set_status(Status( + StatusCode.ERROR, f"{type(error).__name__}: {error}" + )) + raise def start_span( self, diff --git a/python/frameworks/openai/tests/test_chat_stream_completion.py b/python/frameworks/openai/tests/test_chat_stream_completion.py new file mode 100644 index 000000000..908b1e00b --- /dev/null +++ b/python/frameworks/openai/tests/test_chat_stream_completion.py @@ -0,0 +1,121 @@ +"""HTTP EOF is not evidence that every requested chat choice completed.""" +import asyncio +import json + +import httpx +import pytest +from fi_instrumentation import TraceConfig +from opentelemetry.trace import StatusCode + +from .test_openai_compatible_spans import ( + REQUEST_MODEL, async_client, client, make_tracer, only_span, +) + + +def chunk(index=0, finish_reason=None, content=None): + return { + "id": "synthetic", "object": "chat.completion.chunk", "created": 1, + "model": REQUEST_MODEL, + "choices": [{"index": index, "delta": {"content": content} if content else {}, + "finish_reason": finish_reason}], + } + + +def stream_response(chunks, done): + body = "".join(f"data: {json.dumps(c)}\n\n" for c in chunks) + if done: + body += "data: [DONE]\n\n" + return httpx.Response(200, headers={"content-type": "text/event-stream"}, content=body.encode()) + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("done", [False, True]) +@pytest.mark.parametrize("case", [ + "empty", "partial", "usage_only", "stop", "length", "tool_calls", "content_filter", + "function_call", "two_complete", "second_partial", "second_missing", "usage_tail", +]) +def test_chat_completion_requires_terminal_choices(make_tracer, asynchronous, done, case): + exporter = make_tracer(TraceConfig(hide_inputs=True, hide_outputs=True, hide_error_details=True)) + chunks = [chunk(content="synthetic-private-output")] + n = 1 + complete = case in {"stop", "length", "tool_calls", "content_filter", "function_call", "two_complete", "usage_tail"} + usage = {**chunk(), "choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3, "total_tokens": 5}} + if case == "empty": + chunks = [] + elif case == "usage_only": + chunks = [usage] + elif case in {"two_complete", "second_partial", "second_missing"}: + n = 2 + chunks.append(chunk(finish_reason="stop")) + if case != "second_missing": + chunks.append(chunk(index=1, content="second")) + if case == "two_complete": + chunks.append(chunk(index=1, finish_reason="stop")) + elif case == "usage_tail": + chunks.extend([chunk(finish_reason="stop"), usage]) + elif complete: + chunks.append(chunk(finish_reason=case)) + params = dict(model=REQUEST_MODEL, messages=[{"role": "user", "content": "synthetic-input"}], stream=True, n=n) + + async def run_async(): + async with async_client(lambda _: stream_response(chunks, done)) as sdk: + stream = await sdk.chat.completions.create(**params) + actual = [c async for c in stream] + assert len(actual) == len(chunks) + await stream.close() + + if asynchronous: + asyncio.run(run_async()) + else: + with client(lambda _: stream_response(chunks, done)) as sdk: + stream = sdk.chat.completions.create(**params) + actual = list(stream) + assert len(actual) == len(chunks) + stream.close() + span, attrs = only_span(exporter) + assert span.status.status_code == (StatusCode.OK if complete else StatusCode.ERROR) + assert "synthetic-private-output" not in str(attrs) + if case == "usage_tail": + assert attrs["gen_ai.usage.total_tokens"] == 5 + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("terminal", [False, True]) +@pytest.mark.parametrize("consumer_error", [False, True]) +def test_context_exit_uses_observed_terminal_state(make_tracer, asynchronous, terminal, consumer_error): + exporter = make_tracer() + chunks = [chunk(content="partial", finish_reason="stop" if terminal else None)] + params = dict(model=REQUEST_MODEL, messages=[{"role": "user", "content": "hi"}], stream=True) + error = ValueError("synthetic consumer error") + + async def run_async(): + async with async_client(lambda _: stream_response(chunks, True)) as sdk: + stream = await sdk.chat.completions.create(**params) + async with stream: + await stream.__anext__() + if consumer_error: + raise error + + def run_sync(): + with client(lambda _: stream_response(chunks, True)) as sdk: + stream = sdk.chat.completions.create(**params) + with stream: + next(stream) + if consumer_error: + raise error + + def run(): + if asynchronous: + asyncio.run(run_async()) + else: + run_sync() + + if consumer_error: + with pytest.raises(ValueError) as caught: + run() + assert caught.value is error + else: + run() + span, _ = only_span(exporter) + assert span.status.status_code == (StatusCode.OK if terminal and not consumer_error else StatusCode.ERROR) + assert len([e for e in span.events if e.name == "exception"]) == int(consumer_error) diff --git a/python/frameworks/openai/tests/test_stream_context_manager.py b/python/frameworks/openai/tests/test_stream_context_manager.py index 684f57dd9..388024ec9 100644 --- a/python/frameworks/openai/tests/test_stream_context_manager.py +++ b/python/frameworks/openai/tests/test_stream_context_manager.py @@ -31,7 +31,7 @@ def test_sync_stream_context_finishes_once(make_tracer, mode): span, attrs = only_span(exporter) assert attrs["gen_ai.request.model"] == REQUEST_MODEL assert attrs["gen_ai.provider.name"] == "baseten" - assert span.status.status_code == (StatusCode.ERROR if mode == "error" else StatusCode.OK) + assert span.status.status_code == (StatusCode.OK if mode == "full" else StatusCode.ERROR) if mode == "full": assert attrs["output.value"] == "Tokyo" if mode == "early": @@ -66,7 +66,7 @@ async def run(): span, attrs = only_span(exporter) assert attrs["gen_ai.request.model"] == REQUEST_MODEL assert attrs["gen_ai.provider.name"] == "baseten" - assert span.status.status_code == (StatusCode.ERROR if mode == "error" else StatusCode.OK) + assert span.status.status_code == (StatusCode.OK if mode == "full" else StatusCode.ERROR) if mode == "full": assert attrs["output.value"] == "Tokyo" if mode == "early": diff --git a/python/frameworks/openai/traceai_openai/_response_accumulator.py b/python/frameworks/openai/traceai_openai/_response_accumulator.py index 17417d213..743aa60e5 100644 --- a/python/frameworks/openai/traceai_openai/_response_accumulator.py +++ b/python/frameworks/openai/traceai_openai/_response_accumulator.py @@ -19,6 +19,7 @@ from fi_instrumentation import safe_json_dumps from fi_instrumentation.fi_types import FiMimeTypeValues, SpanAttributes +from opentelemetry import trace as trace_api from opentelemetry.util.types import AttributeValue from traceai_openai._utils import _as_output_attributes, _to_dict, _ValueAndType, _io_value_and_type @@ -122,6 +123,7 @@ class _ChatCompletionAccumulator: "_chat_completion_type", "_content", "_raw_data", + "_choice_completion", ) def __init__( @@ -136,6 +138,7 @@ def __init__( self._is_null = True self._content: List[str] = [] self._raw_data: List[Dict[str, Any]] = [] + self._choice_completion: Dict[int, bool] = {} self._cached_result: Optional[Dict[str, Any]] = None self._values = _ValuesAccumulator( choices=_IndexedAccumulator( @@ -162,6 +165,13 @@ def __init__( def process_chunk(self, chunk: Any) -> None: raw_chunk = _to_dict(chunk) self._raw_data.append(raw_chunk) + for choice in raw_chunk.get("choices", []): + index = choice.get("index") + if isinstance(index, int): + self._choice_completion[index] = ( + self._choice_completion.get(index, False) + or bool(choice.get("finish_reason")) + ) self._is_null = False self._cached_result = None @@ -185,6 +195,21 @@ def process_chunk(self, chunk: Any) -> None: if content is not None: self._content.append(content) + def get_status(self) -> trace_api.Status: + # The SDK consumes [DONE] internally and treats transport EOF as normal + # exhaustion. Only per-choice finish reasons prove model completion. + # Usage-only tail chunks must not clear an already completed choice. + expected = self._request_parameters.get("n", 1) + if not isinstance(expected, int) or expected < 1: + expected = 1 + if (len(self._choice_completion) >= expected + and all(self._choice_completion.values())): + return trace_api.Status(trace_api.StatusCode.OK) + return trace_api.Status( + trace_api.StatusCode.ERROR, + "Chat stream ended before all choices completed", + ) + def _result(self) -> Optional[Dict[str, Any]]: if self._is_null: return None diff --git a/python/frameworks/openai/traceai_openai/_stream.py b/python/frameworks/openai/traceai_openai/_stream.py index a3a60042e..716a3a18e 100644 --- a/python/frameworks/openai/traceai_openai/_stream.py +++ b/python/frameworks/openai/traceai_openai/_stream.py @@ -115,6 +115,12 @@ def _finish_tracing( self, status: Optional[trace_api.Status] = None, ) -> None: + if self._self_is_finished: + return + if status is None or status.status_code is not trace_api.StatusCode.ERROR: + get_status = getattr(self._self_response_accumulator, "get_status", None) + if get_status is not None: + status = get_status() _finish_tracing( status=status, with_span=self._self_with_span, @@ -205,6 +211,12 @@ def _finish_tracing( self, status: Optional[trace_api.Status] = None, ) -> None: + if self._self_is_finished: + return + if status is None or status.status_code is not trace_api.StatusCode.ERROR: + get_status = getattr(self._self_response_accumulator, "get_status", None) + if get_status is not None: + status = get_status() _finish_tracing( status=status, with_span=self._self_with_span, diff --git a/python/tests/test_tracer_cancellation.py b/python/tests/test_tracer_cancellation.py new file mode 100644 index 000000000..56be5a513 --- /dev/null +++ b/python/tests/test_tracer_cancellation.py @@ -0,0 +1,79 @@ +"""Local aborts must not become successful spans at the FI processor.""" +import asyncio + +import pytest +from fi_instrumentation import TraceConfig, register +from fi_instrumentation.fi_types import ProjectType +from fi_instrumentation.instrumentation._tracers import FITracer +from opentelemetry import trace +from opentelemetry.exporter.otlp.proto.common.trace_encoder import encode_spans +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import StatusCode + + +@pytest.mark.parametrize("error_type", [KeyboardInterrupt, asyncio.CancelledError, GeneratorExit, ValueError]) +@pytest.mark.parametrize("record_exception", [False, True]) +@pytest.mark.parametrize("set_status", [False, True]) +@pytest.mark.parametrize("hide_details", [False, True]) +@pytest.mark.parametrize("end_on_exit", [False, True]) +def test_context_abort_preserves_flags_and_exception( + error_type, record_exception, set_status, hide_details, end_on_exit, +): + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + tracer = FITracer(provider.get_tracer(__name__), TraceConfig(hide_error_details=hide_details)) + error = error_type("synthetic-private-abort") + try: + with tracer.start_as_current_span("parent") as parent: + with pytest.raises(error_type) as caught: + with tracer.start_as_current_span( + "child", record_exception=record_exception, + set_status_on_exception=set_status, end_on_exit=end_on_exit, + ) as child: + assert trace.get_current_span() is child + raise error + assert caught.value is error + assert trace.get_current_span() is parent + if not end_on_exit: + assert child.is_recording() + assert not exporter.get_finished_spans() + child.end() + spans = exporter.get_finished_spans() + assert len(spans) == 2 + span = next(s for s in spans if s.name == "child") + assert span.parent.span_id == parent.get_span_context().span_id + assert span.status.status_code == (StatusCode.ERROR if set_status else StatusCode.UNSET) + exceptions = [e for e in span.events if e.name == "exception"] + assert len(exceptions) == int(record_exception) + wire = encode_spans([span]).SerializeToString() + if hide_details: + assert b"synthetic-private-abort" not in wire + elif record_exception or set_status: + assert b"synthetic-private-abort" in wire + finally: + provider.shutdown() + + +@pytest.mark.parametrize("batch", [False, True]) +@pytest.mark.parametrize("error_type", [KeyboardInterrupt, asyncio.CancelledError]) +def test_registered_processor_exports_interrupted_agent_as_error(batch, error_type): + exporter = InMemorySpanExporter() + provider = register( + project_name="synthetic-cancellation", project_type=ProjectType.OBSERVE, + span_exporter=exporter, batch=batch, verbose=False, + ) + tracer = FITracer(provider.get_tracer(__name__), TraceConfig(hide_error_details=True)) + try: + with pytest.raises(error_type): + with tracer.start_as_current_span("order-assistant", fi_span_kind="agent"): + raise error_type("synthetic-private-abort") + provider.force_flush() + (span,) = exporter.get_finished_spans() + assert span.status.status_code == StatusCode.ERROR + assert len([e for e in span.events if e.name == "exception"]) == 1 + assert b"synthetic-private-abort" not in encode_spans([span]).SerializeToString() + finally: + provider.shutdown() From 0c16bad668073d12e12d382acd1feee68266a883 Mon Sep 17 00:00:00 2001 From: Nikhil Pareek Date: Mon, 2 Nov 2026 03:43:04 -0800 Subject: [PATCH 05/10] test(openai): complete the Fireworks success stream fixture Include a terminal finish reason before the usage tail. Preserve strict truncated-stream detection and run both optional migration-route suites. Co-Authored-By: Claude Code --- .../openai/tests/test_fireworks_route.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/python/frameworks/openai/tests/test_fireworks_route.py b/python/frameworks/openai/tests/test_fireworks_route.py index d5698b981..a09194047 100644 --- a/python/frameworks/openai/tests/test_fireworks_route.py +++ b/python/frameworks/openai/tests/test_fireworks_route.py @@ -90,6 +90,21 @@ def response(outcome): } for delta in deltas ] + # A successful generation finishes its choice before the usage-only + # tail. [DONE] alone is transport termination, not model completion. + chunks.append( + { + "id": "c", + "object": "chat.completion.chunk", + "created": 1, + "model": MODEL, + "choices": [{ + "index": 0, + "delta": {}, + "finish_reason": "tool_calls" if outcome == "tool_stream" else "stop", + }], + } + ) chunks.append( { "id": "c", From 2eae211d73fc327bfe55b45e5372f0a049f75d77 Mon Sep 17 00:00:00 2001 From: Nikhil Pareek Date: Mon, 2 Nov 2026 04:28:13 -0800 Subject: [PATCH 06/10] fix(openai): honor effective body in stream completion checks Apply extra_body overrides and Omit removals before extracting request metadata and expected streamed choices. Cover sync/async on-wire override directions, removal, null/default counts, and non-mutating shallow merge behavior. Co-Authored-By: Claude Code --- .../tests/test_chat_stream_completion.py | 55 +++++++++++++++++++ .../openai/tests/test_request_body.py | 41 ++++++++++++++ .../openai/traceai_openai/_request.py | 20 +++++-- 3 files changed, 111 insertions(+), 5 deletions(-) create mode 100644 python/frameworks/openai/tests/test_request_body.py diff --git a/python/frameworks/openai/tests/test_chat_stream_completion.py b/python/frameworks/openai/tests/test_chat_stream_completion.py index 908b1e00b..bd1f75344 100644 --- a/python/frameworks/openai/tests/test_chat_stream_completion.py +++ b/python/frameworks/openai/tests/test_chat_stream_completion.py @@ -3,6 +3,7 @@ import json import httpx +import openai import pytest from fi_instrumentation import TraceConfig from opentelemetry.trace import StatusCode @@ -79,6 +80,60 @@ async def run_async(): assert attrs["gen_ai.usage.total_tokens"] == 5 +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("declared_n,override_n,completed_choices,expected_status", [ + pytest.param(2, 1, 1, StatusCode.OK, id="override-down"), + pytest.param(1, 2, 1, StatusCode.ERROR, id="override-up-incomplete"), + pytest.param(1, 2, 2, StatusCode.OK, id="override-up-complete"), + pytest.param(2, openai.Omit(), 1, StatusCode.OK, id="remove-n"), + pytest.param(2, None, 1, StatusCode.OK, id="null-n"), + pytest.param(openai.NOT_GIVEN, 2, 1, StatusCode.ERROR, id="extra-body-only-n"), +]) +def test_completion_count_uses_effective_body( + make_tracer, asynchronous, declared_n, override_n, completed_choices, expected_status, +): + exporter = make_tracer() + chunks = [chunk(index=i, finish_reason="stop") for i in range(completed_choices)] + wire_bodies = [] + + def handler(request): + body = json.loads(request.content) + wire_bodies.append(body) + if isinstance(override_n, openai.Omit): + assert "n" not in body + else: + assert body["n"] == override_n + return stream_response(chunks, done=True) + + extra_body = {"n": override_n} + params = dict( + model=REQUEST_MODEL, messages=[{"role": "user", "content": "synthetic-input"}], + stream=True, n=declared_n, extra_body=extra_body, + ) + + async def run_async(): + async with async_client(handler) as sdk: + stream = await sdk.chat.completions.create(**params) + assert len([c async for c in stream]) == completed_choices + await stream.close() + + if asynchronous: + asyncio.run(run_async()) + else: + with client(handler) as sdk: + stream = sdk.chat.completions.create(**params) + assert len(list(stream)) == completed_choices + stream.close() + + assert len(wire_bodies) == 1 + assert extra_body == {"n": override_n} + span, attrs = only_span(exporter) + assert span.status.status_code == expected_status + # Request metadata and completion detection must describe the same body. + expected_parameters = {k: v for k, v in wire_bodies[0].items() if k != "messages"} + assert json.loads(attrs["gen_ai.request.parameters"]) == expected_parameters + + @pytest.mark.parametrize("asynchronous", [False, True]) @pytest.mark.parametrize("terminal", [False, True]) @pytest.mark.parametrize("consumer_error", [False, True]) diff --git a/python/frameworks/openai/tests/test_request_body.py b/python/frameworks/openai/tests/test_request_body.py new file mode 100644 index 000000000..70ec23114 --- /dev/null +++ b/python/frameworks/openai/tests/test_request_body.py @@ -0,0 +1,41 @@ +"""Request parsing follows the SDK's shallow extra_body merge without mutation.""" +from types import SimpleNamespace + +import openai +import pytest + +from traceai_openai._request import _parse_request_args + + +@pytest.mark.parametrize("json_data,extra_json,expected", [ + pytest.param({"n": 2}, {"n": 1}, {"n": 1}, id="override"), + pytest.param({"n": 2}, {"n": openai.Omit()}, {}, id="remove"), + pytest.param({"n": 2}, {"n": None}, {"n": None}, id="preserve-null"), + pytest.param(None, {"n": 2}, {"n": 2}, id="extra-only"), + pytest.param({"n": 2}, None, {"n": 2}, id="no-extra"), + pytest.param(None, None, {}, id="no-body"), + pytest.param( + {"metadata": {"old": "value", "shared": "old"}}, + {"metadata": {"shared": "new"}}, + {"metadata": {"shared": "new"}}, + id="shallow-replacement", + ), +]) +def test_parse_effective_body(json_data, extra_json, expected): + original_json = dict(json_data) if json_data is not None else None + original_extra = dict(extra_json) if extra_json is not None else None + options = SimpleNamespace(json_data=json_data, extra_json=extra_json) + cast_to = openai.types.chat.ChatCompletion + + actual_type, actual = _parse_request_args((cast_to, options)) + + assert actual_type is cast_to + assert actual == expected + assert options.json_data == original_json + assert options.extra_json == original_extra + + +def test_options_without_extra_json_remain_supported(): + options = SimpleNamespace(json_data={"n": 2}) + _, actual = _parse_request_args((openai.types.chat.ChatCompletion, options)) + assert actual == {"n": 2} diff --git a/python/frameworks/openai/traceai_openai/_request.py b/python/frameworks/openai/traceai_openai/_request.py index c6e8cf426..fba03fe83 100644 --- a/python/frameworks/openai/traceai_openai/_request.py +++ b/python/frameworks/openai/traceai_openai/_request.py @@ -10,6 +10,7 @@ FiSpanKindValues, SpanAttributes, ) +from openai import Omit from openai.types.create_embedding_response import CreateEmbeddingResponse from opentelemetry import context as context_api from opentelemetry import trace as trace_api @@ -437,12 +438,21 @@ def _parse_request_args(args: Tuple[type, Any]) -> Tuple[type, Mapping[str, Any] # The targeted signature of `request` is here: # https://github.com/openai/openai-python/blob/f1c7d714914e3321ca2e72839fe2d132a8646e7f/src/openai/_base_client.py#L846-L847 # noqa: E501 cast_to: type = args[0] + json_data = getattr(args[1], "json_data", None) + extra_json = getattr(args[1], "extra_json", None) + # Mirror the SDK's shallow extra_body merge without mutating its options. + # Overrides (including Omit removals) affect both request metadata and the + # number of terminal choices required for a successful streamed completion. + if extra_json is not None: + if json_data is None: + json_data = extra_json + elif isinstance(json_data, Mapping) and isinstance(extra_json, Mapping): + json_data = { + key: value for key, value in {**json_data, **extra_json}.items() + if not isinstance(value, Omit) + } request_parameters: Mapping[str, Any] = ( - json_data - # See https://github.com/openai/openai-python/blob/f1c7d714914e3321ca2e72839fe2d132a8646e7f/src/openai/_models.py#L427 # noqa: E501 - if hasattr(args[1], "json_data") - and isinstance(json_data := args[1].json_data, Mapping) - else {} + json_data if isinstance(json_data, Mapping) else {} ) # FIXME: Because request parameters is just a Mapping, it can contain any value as long as it # serializes correctly in an HTTP request body. For example, Enum values may be present if a From 6dfc7b1ca46cce0f5535cf6094f8debda59467c5 Mon Sep 17 00:00:00 2001 From: Nikhil Pareek Date: Mon, 2 Nov 2026 04:41:10 -0800 Subject: [PATCH 07/10] fix(openai): scope effective body to chat completion validation Preserve existing exported request metadata and recipe contracts; only the chat stream accumulator consumes the effective body. Verify arbitrary extra_body values are not newly exported while overridden and removed choice counts determine completion status. Co-Authored-By: Claude Code --- .../tests/test_chat_stream_completion.py | 15 ++++++++----- .../openai/tests/test_request_body.py | 4 +++- .../openai/traceai_openai/_request.py | 22 ++++++++++++++----- 3 files changed, 30 insertions(+), 11 deletions(-) diff --git a/python/frameworks/openai/tests/test_chat_stream_completion.py b/python/frameworks/openai/tests/test_chat_stream_completion.py index bd1f75344..695bfe8f8 100644 --- a/python/frameworks/openai/tests/test_chat_stream_completion.py +++ b/python/frameworks/openai/tests/test_chat_stream_completion.py @@ -105,7 +105,7 @@ def handler(request): assert body["n"] == override_n return stream_response(chunks, done=True) - extra_body = {"n": override_n} + extra_body = {"n": override_n, "synthetic_extra": "not-exported"} params = dict( model=REQUEST_MODEL, messages=[{"role": "user", "content": "synthetic-input"}], stream=True, n=declared_n, extra_body=extra_body, @@ -126,12 +126,17 @@ async def run_async(): stream.close() assert len(wire_bodies) == 1 - assert extra_body == {"n": override_n} + assert extra_body == {"n": override_n, "synthetic_extra": "not-exported"} span, attrs = only_span(exporter) assert span.status.status_code == expected_status - # Request metadata and completion detection must describe the same body. - expected_parameters = {k: v for k, v in wire_bodies[0].items() if k != "messages"} - assert json.loads(attrs["gen_ai.request.parameters"]) == expected_parameters + # This fix must not expand exported metadata to arbitrary extra_body fields. + parameters = json.loads(attrs["gen_ai.request.parameters"]) + if declared_n is openai.NOT_GIVEN: + assert "n" not in parameters + else: + assert parameters["n"] == declared_n + assert "synthetic_extra" not in str(attrs) + assert "not-exported" not in str(attrs) @pytest.mark.parametrize("asynchronous", [False, True]) diff --git a/python/frameworks/openai/tests/test_request_body.py b/python/frameworks/openai/tests/test_request_body.py index 70ec23114..1351fa537 100644 --- a/python/frameworks/openai/tests/test_request_body.py +++ b/python/frameworks/openai/tests/test_request_body.py @@ -27,7 +27,9 @@ def test_parse_effective_body(json_data, extra_json, expected): options = SimpleNamespace(json_data=json_data, extra_json=extra_json) cast_to = openai.types.chat.ChatCompletion - actual_type, actual = _parse_request_args((cast_to, options)) + _, metadata_parameters = _parse_request_args((cast_to, options)) + assert metadata_parameters == (json_data or {}) + actual_type, actual = _parse_request_args((cast_to, options), include_extra_body=True) assert actual_type is cast_to assert actual == expected diff --git a/python/frameworks/openai/traceai_openai/_request.py b/python/frameworks/openai/traceai_openai/_request.py index fba03fe83..4e20a0414 100644 --- a/python/frameworks/openai/traceai_openai/_request.py +++ b/python/frameworks/openai/traceai_openai/_request.py @@ -202,6 +202,7 @@ def _finalize_response( with_span: _WithSpan, cast_to: type, request_parameters: Mapping[str, Any], + request_options: Any = None, ) -> Any: """ Monkey-patch the response object to trace the stream, or finish tracing if the response is @@ -224,8 +225,15 @@ def _finalize_response( response_accumulator_factory = self._response_accumulator_factories.get( cast_to ) + accumulator_parameters = request_parameters + if cast_to is self._openai.types.chat.ChatCompletion and request_options is not None: + # Completion validation needs the on-wire choice count. Keep + # extra_body out of the existing exported request metadata. + _, accumulator_parameters = _parse_request_args( + (cast_to, request_options), include_extra_body=True, + ) response_accumulator = ( - response_accumulator_factory(request_parameters) + response_accumulator_factory(accumulator_parameters) if response_accumulator_factory else None ) @@ -340,6 +348,7 @@ def __call__( with_span=with_span, cast_to=cast_to, request_parameters=request_parameters, + request_options=args[1], ) except Exception as e: status = trace_api.Status( @@ -423,6 +432,7 @@ async def __call__( with_span=with_span, cast_to=cast_to, request_parameters=request_parameters, + request_options=args[1], ) except Exception: logger.exception( @@ -432,7 +442,9 @@ async def __call__( return response -def _parse_request_args(args: Tuple[type, Any]) -> Tuple[type, Mapping[str, Any]]: +def _parse_request_args( + args: Tuple[type, Any], *, include_extra_body: bool = False, +) -> Tuple[type, Mapping[str, Any]]: # We don't use `signature(request).bind()` because `request` could have been monkey-patched # (incorrectly) by others and the signature at runtime may not match the original. # The targeted signature of `request` is here: @@ -441,9 +453,9 @@ def _parse_request_args(args: Tuple[type, Any]) -> Tuple[type, Mapping[str, Any] json_data = getattr(args[1], "json_data", None) extra_json = getattr(args[1], "extra_json", None) # Mirror the SDK's shallow extra_body merge without mutating its options. - # Overrides (including Omit removals) affect both request metadata and the - # number of terminal choices required for a successful streamed completion. - if extra_json is not None: + # Overrides (including Omit removals) affect the number of terminal choices + # required for a successful stream, not the existing metadata contract. + if include_extra_body and extra_json is not None: if json_data is None: json_data = extra_json elif isinstance(json_data, Mapping) and isinstance(extra_json, Mapping): From 3e19fa3b2333f812228b62fade8450a57a0fa798 Mon Sep 17 00:00:00 2001 From: Nikhil Pareek Date: Mon, 2 Nov 2026 05:36:03 -0800 Subject: [PATCH 08/10] fix(openai): redact effective MCP tools and track raw body consumption Co-Authored-By: Claude Code --- ...test_effective_body_and_raw_consumption.py | 260 ++++++++++++++++++ .../openai/traceai_openai/_privacy.py | 9 +- .../traceai_openai/_response_accumulator.py | 7 +- .../openai/traceai_openai/_stream.py | 58 +++- 4 files changed, 331 insertions(+), 3 deletions(-) create mode 100644 python/frameworks/openai/tests/test_effective_body_and_raw_consumption.py diff --git a/python/frameworks/openai/tests/test_effective_body_and_raw_consumption.py b/python/frameworks/openai/tests/test_effective_body_and_raw_consumption.py new file mode 100644 index 000000000..74a5c053d --- /dev/null +++ b/python/frameworks/openai/tests/test_effective_body_and_raw_consumption.py @@ -0,0 +1,260 @@ +"""Effective credentials and lazy body lifecycles through real SDK transports.""" + +import asyncio +from copy import deepcopy +import json +from types import SimpleNamespace + +try: + import httpx2 as httpx +except ImportError: + import httpx +import openai +import pytest +from fi_instrumentation import TraceConfig, register +from fi_instrumentation.fi_types import ProjectType +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import StatusCode +from traceai_openai import OpenAIInstrumentor +from traceai_openai._privacy import request_credentials + + +@pytest.fixture +def tracing(): + exporter = InMemorySpanExporter() + provider = register( + project_name="body-consumption", + project_type=ProjectType.OBSERVE, + span_exporter=exporter, + batch=False, + verbose=False, + ) + OpenAIInstrumentor().instrument(tracer_provider=provider, config=TraceConfig()) + try: + yield provider, exporter + finally: + OpenAIInstrumentor().uninstrument() + provider.force_flush() + provider.shutdown() + + +def sdk(handler, async_=False): + http = (httpx.AsyncClient if async_ else httpx.Client)( + transport=httpx.MockTransport(handler), + trust_env=False, + ) + return (openai.AsyncOpenAI if async_ else openai.OpenAI)( + api_key="synthetic-static-key", + max_retries=0, + http_client=http, + ) + + +def response_body(status="completed"): + return { + "id": "resp_review", + "object": "response", + "created_at": 1, + "model": "synthetic", + "status": status, + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "output": [], + "error": None, + "usage": { + "input_tokens": 5, + "output_tokens": 2, + "total_tokens": 7, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens_details": {"reasoning_tokens": 0}, + }, + } + + +@pytest.mark.parametrize("async_", [False, True]) +@pytest.mark.parametrize("extra", [False, True]) +@pytest.mark.parametrize("credential", ["authorization", "header"]) +def test_mcp_credentials_in_effective_body(tracing, async_, extra, credential): + _, exporter = tracing + secret = 'synthetic-extra-mcp-credential-17498-"quoted"' + tool = { + "type": "mcp", + "server_label": "synthetic", + "server_url": "https://synthetic.invalid/mcp", + **( + {"authorization": secret} + if credential == "authorization" + else {"headers": {"X-Private-Token": secret}} + ), + } + original_tool = deepcopy(tool) + captured = [] + + def handler(request): + captured.append(json.loads(request.content)) + assert captured[-1]["tools"][0] == tool + return httpx.Response( + 400, + json={ + "error": { + "code": "invalid_request", + "message": "MCP authentication failed for " + secret, + } + }, + ) + + params = {"model": "synthetic", "input": "question"} + params.update( + {"tools": [], "extra_body": {"tools": [tool]}} if extra else {"tools": [tool]} + ) + + async def run(): + async with sdk(handler, True) as client: + await client.responses.create(**params) + + with pytest.raises(openai.BadRequestError) as caught: + if async_: + asyncio.run(run()) + else: + with sdk(handler) as client: + client.responses.create(**params) + assert captured and tool == original_tool + assert secret in caught.value.body["message"] + (span,) = exporter.get_finished_spans() + assert span.status.status_code is StatusCode.ERROR + assert "synthetic-extra-mcp-credential" not in span.to_json() + if extra: + assert "synthetic.invalid/mcp" not in span.to_json() + + +@pytest.mark.parametrize("override", [None, [], openai.Omit()]) +def test_effective_tools_override_original_without_mutation(override): + original = {"tools": [{"type": "mcp", "authorization": "synthetic-original"}]} + extra = {"tools": override} + secrets = request_credentials(None, SimpleNamespace(extra_json=extra), original) + assert "synthetic-original" not in secrets + assert original["tools"][0]["authorization"] == "synthetic-original" + assert extra["tools"] is override + + +@pytest.mark.parametrize("async_", [False, True]) +@pytest.mark.parametrize("mode", ["parse", "json", "text", "read"]) +@pytest.mark.parametrize("status", ["completed", "failed", "incomplete", "cancelled"]) +def test_raw_body_consumption(tracing, async_, mode, status): + _, exporter = tracing + body = response_body(status) + + def handler(request): + return httpx.Response(200, json=body) + + def check(result): + if mode == "parse": + assert result.status == status + elif mode == "json": + assert result == body + else: + assert json.loads(result) == body + + async def run(): + async with sdk(handler, True) as client: + async with client.responses.with_streaming_response.create( + model="synthetic", + input="question", + ) as raw: + check(await getattr(raw, mode)()) + check(await getattr(raw, mode)()) + await raw.close() + + if async_: + asyncio.run(run()) + else: + with sdk(handler) as client: + with client.responses.with_streaming_response.create( + model="synthetic", + input="question", + ) as raw: + check(getattr(raw, mode)()) + check(getattr(raw, mode)()) + raw.close() + (span,) = exporter.get_finished_spans() + assert span.status.status_code is ( + StatusCode.OK if status == "completed" else StatusCode.ERROR + ) + + +@pytest.mark.parametrize("async_", [False, True]) +@pytest.mark.parametrize("mode", ["read", "text", "json"]) +@pytest.mark.parametrize("outcome", ["completed", "unread", "error", "cancel", "sse"]) +def test_raw_consumption_is_lazy_and_errors_finish_once(tracing, async_, mode, outcome): + _, exporter = tracing + reads = [] + closes = [] + error = ( + KeyboardInterrupt("synthetic abort") + if outcome == "cancel" + else RuntimeError("synthetic read failure") + ) + body = json.dumps(response_body()).encode() + + class SyncBody(httpx.SyncByteStream): + def __iter__(self): + reads.append(1) + if outcome in {"error", "cancel"}: + raise error + yield body + + def close(self): + closes.append(1) + + class AsyncBody(httpx.AsyncByteStream): + async def __aiter__(self): + reads.append(1) + if outcome in {"error", "cancel"}: + raise error + yield body + + async def aclose(self): + closes.append(1) + + def handler(request): + return httpx.Response( + 200, + headers={"content-type": "application/json"}, + stream=AsyncBody() if async_ else SyncBody(), + ) + + params = {"model": "synthetic", "input": "question", "stream": outcome == "sse"} + + async def run(): + async with sdk(handler, True) as client: + async with client.responses.with_streaming_response.create(**params) as raw: + assert not reads + assert not exporter.get_finished_spans() + if outcome != "unread": + await getattr(raw, mode)() + await raw.close() + + def run_sync(): + with sdk(handler) as client: + with client.responses.with_streaming_response.create(**params) as raw: + assert not reads + assert not exporter.get_finished_spans() + if outcome != "unread": + getattr(raw, mode)() + raw.close() + + if outcome in {"error", "cancel"}: + with pytest.raises(type(error)) as caught: + asyncio.run(run()) if async_ else run_sync() + assert caught.value is error + else: + asyncio.run(run()) if async_ else run_sync() + assert len(reads) == (0 if outcome == "unread" else 1) + assert closes + (span,) = exporter.get_finished_spans() + assert span.status.status_code is ( + StatusCode.OK if outcome == "completed" else StatusCode.ERROR + ) + if outcome in {"error", "cancel"}: + assert len([event for event in span.events if event.name == "exception"]) == 1 diff --git a/python/frameworks/openai/traceai_openai/_privacy.py b/python/frameworks/openai/traceai_openai/_privacy.py index 98521c73e..cdeafe6b6 100644 --- a/python/frameworks/openai/traceai_openai/_privacy.py +++ b/python/frameworks/openai/traceai_openai/_privacy.py @@ -31,7 +31,14 @@ def request_credentials(instance: Any, options: Any, parameters: Mapping) -> tup values.append(value) if isinstance(value, str) and value.lower().startswith("bearer "): values.append(value[7:]) - tools = parameters.get("tools") + # extra_body is a shallow SDK override. Inspect its effective tools only + # for private redaction; never add extra_body to exported request metadata. + extra_json = getattr(options, "extra_json", None) + tools = ( + extra_json["tools"] + if isinstance(extra_json, Mapping) and "tools" in extra_json + else parameters.get("tools") + ) if isinstance(tools, (list, tuple)): for tool in tools: if isinstance(tool, Mapping) and tool.get("type") == "mcp": diff --git a/python/frameworks/openai/traceai_openai/_response_accumulator.py b/python/frameworks/openai/traceai_openai/_response_accumulator.py index 3b73efdeb..8782d40c7 100644 --- a/python/frameworks/openai/traceai_openai/_response_accumulator.py +++ b/python/frameworks/openai/traceai_openai/_response_accumulator.py @@ -46,7 +46,12 @@ def get_attributes_from_response( def _response_status(response: Any) -> Status: - if getattr(response, "status", None) in {"failed", "cancelled", "incomplete"}: + status = ( + response.get("status") + if isinstance(response, Mapping) + else getattr(response, "status", None) + ) + if isinstance(status, str) and status in {"failed", "cancelled", "incomplete"}: # The vendor's error message can contain credentials or request content. return Status(StatusCode.ERROR, "Response did not complete") return Status(StatusCode.OK) diff --git a/python/frameworks/openai/traceai_openai/_stream.py b/python/frameworks/openai/traceai_openai/_stream.py index 9829860a3..7fc65c9f9 100644 --- a/python/frameworks/openai/traceai_openai/_stream.py +++ b/python/frameworks/openai/traceai_openai/_stream.py @@ -4,6 +4,7 @@ from openai.types.chat import ChatCompletionChunk from opentelemetry import trace as trace_api from opentelemetry.util.types import AttributeValue +from traceai_openai._response_accumulator import _response_status from traceai_openai._utils import _finish_tracing from traceai_openai._with_span import _WithSpan from wrapt import ObjectProxy @@ -87,6 +88,41 @@ def __init__( self._self_response_accumulator = response_accumulator self._self_on_parse = on_parse self._self_parsed = {} + self._self_body_status: Optional[trace_api.Status] = None + + def _record_body_consumption(self) -> None: + # read/text/json consume a non-SSE envelope without calling parse(). + # Do not force SDK parsing (or its post_parser) on a raw-body caller. + # SSE still needs a terminal event from the traced iterator. + if ( + self._self_is_finished + or self._self_body_status is not None + or getattr(self.__wrapped__, "_is_sse_stream", False) + ): + return + try: + body = self.__wrapped__.http_response.json() + except Exception: + body = None # Successful raw reads need not contain JSON. + self._self_body_status = _response_status(body) + + def _consume_body(self, method: str, *args: Any, **kwargs: Any) -> Any: + try: + result = getattr(self.__wrapped__, method)(*args, **kwargs) + except BaseException as exc: + _finish_on_exit(self, exc) + raise + self._record_body_consumption() + return result + + def read(self, *args: Any, **kwargs: Any) -> Any: + return self._consume_body("read", *args, **kwargs) + + def text(self, *args: Any, **kwargs: Any) -> Any: + return self._consume_body("text", *args, **kwargs) + + def json(self, *args: Any, **kwargs: Any) -> Any: + return self._consume_body("json", *args, **kwargs) @property def _self_is_finished(self) -> bool: @@ -129,7 +165,9 @@ def _finish_tracing(self, status: Optional[trace_api.Status] = None) -> None: accumulator = self._self_response_accumulator if status is None or status.status_code is not trace_api.StatusCode.ERROR: get_status = getattr(accumulator, "get_status", None) - if get_status is not None: + if self._self_body_status is not None: + status = self._self_body_status + elif get_status is not None: status = get_status() elif not self._self_parsed: status = trace_api.Status( @@ -144,6 +182,24 @@ def _finish_tracing(self, status: Optional[trace_api.Status] = None) -> None: class _AsyncRawResponse(_RawResponse): """AsyncAPIResponse has awaitable parse/close; LegacyAPIResponse does not.""" + async def _consume_body(self, method: str, *args: Any, **kwargs: Any) -> Any: + try: + result = await getattr(self.__wrapped__, method)(*args, **kwargs) + except BaseException as exc: + _finish_on_exit(self, exc) + raise + self._record_body_consumption() + return result + + async def read(self, *args: Any, **kwargs: Any) -> Any: + return await self._consume_body("read", *args, **kwargs) + + async def text(self, *args: Any, **kwargs: Any) -> Any: + return await self._consume_body("text", *args, **kwargs) + + async def json(self, *args: Any, **kwargs: Any) -> Any: + return await self._consume_body("json", *args, **kwargs) + async def parse(self, *args: Any, **kwargs: Any) -> Any: try: parsed = await self.__wrapped__.parse(*args, **kwargs) From ff3feb75ca6564d2f5da8c7f9dd51b9064dbb75f Mon Sep 17 00:00:00 2001 From: Nikhil Pareek Date: Mon, 2 Nov 2026 13:39:26 -0800 Subject: [PATCH 09/10] fix(openai): preserve raw envelopes and response privacy across parse modes Co-Authored-By: Claude Code --- python/frameworks/openai/pyproject.toml | 2 +- .../tests/test_core_dependency_privacy.py | 42 ++++ .../openai/tests/test_response_edge_cases.py | 185 ++++++++++++++++++ .../openai/traceai_openai/_privacy.py | 7 +- .../openai/traceai_openai/_request.py | 18 +- .../openai/traceai_openai/_stream.py | 48 +++++ 6 files changed, 297 insertions(+), 5 deletions(-) create mode 100644 python/frameworks/openai/tests/test_core_dependency_privacy.py create mode 100644 python/frameworks/openai/tests/test_response_edge_cases.py diff --git a/python/frameworks/openai/pyproject.toml b/python/frameworks/openai/pyproject.toml index 3a189f628..690651844 100644 --- a/python/frameworks/openai/pyproject.toml +++ b/python/frameworks/openai/pyproject.toml @@ -28,7 +28,7 @@ classifiers = [ [tool.poetry.dependencies] python = ">3.9,<3.14" -fi-instrumentation-otel = ">=0.1.11" +fi-instrumentation-otel = ">=1.2.2" openai = ">=1.69.0" [tool.poetry.urls] diff --git a/python/frameworks/openai/tests/test_core_dependency_privacy.py b/python/frameworks/openai/tests/test_core_dependency_privacy.py new file mode 100644 index 000000000..160b0a30b --- /dev/null +++ b/python/frameworks/openai/tests/test_core_dependency_privacy.py @@ -0,0 +1,42 @@ +"""Synthetic fixture to verify declared core compatibility and hidden error details.""" +import asyncio +import httpx +import openai +import pytest +from fi_instrumentation import TraceConfig, register +from fi_instrumentation.fi_types import ProjectType +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from traceai_openai import OpenAIInstrumentor + +MESSAGE = 'synthetic-private-upstream-error-483129' + +@pytest.mark.parametrize('async_', [False, True]) +def test_hidden_failed_response(async_, monkeypatch): + monkeypatch.setenv('FI_HIDE_ERROR_DETAILS', 'true') + exporter = InMemorySpanExporter() + provider = register(project_name='core-floor-probe', project_type=ProjectType.OBSERVE, + span_exporter=exporter, batch=False, verbose=False) + instrumentor = OpenAIInstrumentor() + instrumentor.instrument(tracer_provider=provider, config=TraceConfig(hide_inputs=True)) + body = {'id': 'resp_probe', 'object': 'response', 'created_at': 1, 'model': 'synthetic', + 'status': 'failed', 'output': [], 'error': {'code': 'server_error', 'message': MESSAGE}, + 'usage': {'input_tokens': 1, 'output_tokens': 0, 'total_tokens': 1}} + def handler(req): return httpx.Response(200, json=body) + def make_client(): + http = (httpx.AsyncClient if async_ else httpx.Client)(transport=httpx.MockTransport(handler), trust_env=False) + return (openai.AsyncOpenAI if async_ else openai.OpenAI)(api_key='synthetic-key', base_url='https://api.parallel.ai/v1', max_retries=0, http_client=http) + try: + if async_: + async def run(): + async with make_client() as client: + assert (await client.responses.create(model='synthetic', input='private input')).status == 'failed' + asyncio.run(run()) + else: + with make_client() as client: + assert client.responses.create(model='synthetic', input='private input').status == 'failed' + spans = exporter.get_finished_spans() + assert len(spans) == 1 + assert MESSAGE not in spans[0].to_json() + finally: + instrumentor.uninstrument() + provider.shutdown() diff --git a/python/frameworks/openai/tests/test_response_edge_cases.py b/python/frameworks/openai/tests/test_response_edge_cases.py new file mode 100644 index 000000000..de31ae8ed --- /dev/null +++ b/python/frameworks/openai/tests/test_response_edge_cases.py @@ -0,0 +1,185 @@ +"""Round 3 independent real SDK probes; synthetic credentials and MockTransport only.""" +import asyncio +import json +import httpx +import openai +from typing import Any, Dict +import pytest +from fi_instrumentation import TraceConfig, register +from fi_instrumentation.fi_types import ProjectType +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import StatusCode +from traceai_openai import OpenAIInstrumentor + +SECRET_INPUT = 'synthetic-private-system-instruction-827391' + +def body(status='completed'): + return {'id': 'resp_independent', 'object': 'response', 'created_at': 1, + 'model': 'parallel', 'status': status, 'parallel_tool_calls': True, + 'tool_choice': 'auto', 'tools': [], 'instructions': SECRET_INPUT, + 'output': [{'id': 'msg_1', 'type': 'message', 'role': 'assistant', + 'status': 'completed', 'content': [{'type': 'output_text', + 'text': 'safe public answer', 'annotations': []}]}], + 'error': None, 'usage': {'input_tokens': 5, 'output_tokens': 2, + 'total_tokens': 7, 'input_tokens_details': {'cached_tokens': 0}, + 'output_tokens_details': {'reasoning_tokens': 0}}} + +@pytest.fixture +def tracing(): + exporter = InMemorySpanExporter() + provider = register(project_name='review3', project_type=ProjectType.OBSERVE, + span_exporter=exporter, batch=False, verbose=False) + OpenAIInstrumentor().instrument(tracer_provider=provider, config=TraceConfig(hide_inputs=True)) + try: + yield provider, exporter + finally: + OpenAIInstrumentor().uninstrument() + provider.force_flush() + provider.shutdown() + +def client(handler, async_=False, base_url='https://api.parallel.ai/v1'): + http = (httpx.AsyncClient if async_ else httpx.Client)( + transport=httpx.MockTransport(handler), trust_env=False) + return (openai.AsyncOpenAI if async_ else openai.OpenAI)( + api_key='synthetic-api-key-938472', base_url=base_url, + max_retries=0, http_client=http) + +@pytest.mark.parametrize('async_', [False, True]) +@pytest.mark.parametrize('to_dict', [False, True]) +def test_explicit_dict_parse_keeps_inputs_hidden(tracing, async_, to_dict): + provider, exporter = tracing + def handler(request): + return httpx.Response(200, json=body()) + async def run(): + async with client(handler, True) as c: + async with c.responses.with_streaming_response.create( + model='parallel', instructions=SECRET_INPUT, input='question') as raw: + result = await raw.parse(**({'to': Dict[str, Any] if openai.__version__.startswith('1.') else dict} if to_dict else {})) + assert (result['instructions'] if to_dict else result.instructions) == SECRET_INPUT + if async_: + asyncio.run(run()) + else: + with client(handler) as c: + with c.responses.with_streaming_response.create( + model='parallel', instructions=SECRET_INPUT, input='question') as raw: + result = raw.parse(**({'to': Dict[str, Any] if openai.__version__.startswith('1.') else dict} if to_dict else {})) + assert (result['instructions'] if to_dict else result.instructions) == SECRET_INPUT + provider.force_flush() + (span,) = exporter.get_finished_spans() + assert 'safe public answer' in span.attributes['output.value'] + leaked = {k: v for k, v in span.attributes.items() if SECRET_INPUT in str(v)} + assert not leaked, leaked + +@pytest.mark.parametrize('async_', [False, True]) +@pytest.mark.parametrize('method', ['parse', 'json', 'read', 'text', 'iter_lines']) +def test_successful_raw_body_consumption_is_not_error(tracing, async_, method): + provider, exporter = tracing + payload = json.dumps(body()).encode() + class SyncBody(httpx.SyncByteStream): + def __iter__(self): + yield payload + class AsyncBody(httpx.AsyncByteStream): + async def __aiter__(self): + yield payload + def handler(request): + return httpx.Response(200, headers={'content-type':'application/json'}, + stream=AsyncBody() if async_ else SyncBody()) + async def run(): + async with client(handler, True) as c: + async with c.responses.with_streaming_response.create(model='parallel', input='question') as raw: + if method == 'iter_lines': + result = [line async for line in raw.iter_lines()] + else: + result = await getattr(raw, method)() + assert result + if async_: + asyncio.run(run()) + else: + with client(handler) as c: + with c.responses.with_streaming_response.create(model='parallel', input='question') as raw: + result = list(raw.iter_lines()) if method == 'iter_lines' else getattr(raw, method)() + assert result + provider.force_flush() + (span,) = exporter.get_finished_spans() + assert span.status.status_code == StatusCode.OK, span.to_json() + +@pytest.mark.parametrize('async_', [False, True]) +@pytest.mark.parametrize('instrumented', [False, True]) +def test_raw_malformed_body_remains_inspectable(tracing, async_, instrumented): + provider, exporter = tracing + if not instrumented: + OpenAIInstrumentor().uninstrument() + def handler(request): + return httpx.Response(200, headers={'content-type': 'application/json'}, content=b'{"truncated":') + async def run(): + async with client(handler, True) as c: + raw = await c.responses.with_raw_response.create(model='parallel', input='question') + assert raw.status_code == 200 + assert raw.content == b'{"truncated":' + with pytest.raises(json.JSONDecodeError): + raw.parse() + if async_: + asyncio.run(run()) + else: + with client(handler) as c: + raw = c.responses.with_raw_response.create(model='parallel', input='question') + assert raw.status_code == 200 + assert raw.content == b'{"truncated":' + with pytest.raises(json.JSONDecodeError): + raw.parse() + +@pytest.mark.parametrize('async_', [False, True]) +def test_speech_stream_to_file_not_marked_failed(tracing, async_, tmp_path): + from pathlib import Path + provider, exporter = tracing + payload = b'synthetic-audio-data' + target = tmp_path / ('audio-async.bin' if async_ else 'audio-sync.bin') + class SyncBody(httpx.SyncByteStream): + def __iter__(self): + yield payload + class AsyncBody(httpx.AsyncByteStream): + async def __aiter__(self): + yield payload + def handler(request): + return httpx.Response(200, headers={'content-type':'audio/mpeg'}, + stream=AsyncBody() if async_ else SyncBody()) + async def run(): + async with client(handler, True, base_url='https://api.openai.com/v1') as c: + async with c.audio.speech.with_streaming_response.create( + model='tts-1', voice='alloy', input='Hello') as raw: + await raw.stream_to_file(target) + if async_: + asyncio.run(run()) + else: + with client(handler, base_url='https://api.openai.com/v1') as c: + with c.audio.speech.with_streaming_response.create( + model='tts-1', voice='alloy', input='Hello') as raw: + raw.stream_to_file(target) + assert target.read_bytes() == payload + provider.force_flush() + (span,) = exporter.get_finished_spans() + assert span.status.status_code == StatusCode.OK, span.to_json() + + +@pytest.mark.parametrize('async_', [False, True]) +def test_early_aborted_raw_iterator_is_error(tracing, async_): + provider, exporter = tracing + def handler(request): + return httpx.Response(200, headers={'content-type': 'text/plain'}, content=b'line one\nline two\n') + async def run(): + async with client(handler, True) as c: + async with c.responses.with_streaming_response.create(model='parallel', input='question') as raw: + iterator = raw.iter_lines() + assert await iterator.__anext__() == 'line one' + await iterator.aclose() + if async_: + asyncio.run(run()) + else: + with client(handler) as c: + with c.responses.with_streaming_response.create(model='parallel', input='question') as raw: + iterator = raw.iter_lines() + assert next(iterator) == 'line one' + iterator.close() + provider.force_flush() + (span,) = exporter.get_finished_spans() + assert span.status.status_code == StatusCode.ERROR diff --git a/python/frameworks/openai/traceai_openai/_privacy.py b/python/frameworks/openai/traceai_openai/_privacy.py index cdeafe6b6..3ff5d3ade 100644 --- a/python/frameworks/openai/traceai_openai/_privacy.py +++ b/python/frameworks/openai/traceai_openai/_privacy.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +from copy import deepcopy from collections.abc import Iterable, Mapping from typing import Any @@ -91,7 +92,11 @@ def response_output(response: Any, config: TraceConfig | None = None) -> dict: "error", "incomplete_details"} if config is not None and getattr(config, "hide_error_details", False): fields.difference_update({"error", "incomplete_details"}) - output = response.model_dump(include=fields, exclude_none=True, mode="json") + if isinstance(response, Mapping): + output = {key: deepcopy(value) for key, value in response.items() + if key in fields and value is not None} + else: + output = response.model_dump(include=fields, exclude_none=True, mode="json") if config is not None and getattr(config, "hide_error_details", False): # These provider-owned error fields are not ordinary tool results or # assistant text. Work on the dump only, leaving SDK objects untouched. diff --git a/python/frameworks/openai/traceai_openai/_request.py b/python/frameworks/openai/traceai_openai/_request.py index 694db0276..feabafdb5 100644 --- a/python/frameworks/openai/traceai_openai/_request.py +++ b/python/frameworks/openai/traceai_openai/_request.py @@ -20,7 +20,7 @@ from opentelemetry.util.types import AttributeValue from traceai_openai._client_provider import _get_client_provider from traceai_openai._providers import provider_from_host -from traceai_openai._privacy import redact, request_credentials, response_credentials, safe_request +from traceai_openai._privacy import redact, request_credentials, response_credentials, response_output, safe_request from traceai_openai._request_context import capture_request_credentials from traceai_openai._request_attributes_extractor import _RequestAttributesExtractor from traceai_openai._response_accumulator import ( @@ -239,7 +239,12 @@ def on_parse(parsed: Any) -> Any: with_span.set_attribute( SpanAttributes.EMBEDDING_EMBEDDINGS, self._raw_embeddings_json(parsed) ) - add_io_to_span_attributes(with_span, None, parsed, config=self._config) + telemetry_output = ( + response_output(parsed, self._config) + if cast_to is self._openai.types.responses.response.Response + and isinstance(parsed, dict) else parsed + ) + add_io_to_span_attributes(with_span, None, telemetry_output, config=self._config) _finish_tracing( status=_response_status(parsed), with_span=with_span, @@ -261,7 +266,14 @@ def on_parse(parsed: Any) -> Any: if isinstance(response, LegacyAPIResponse): # Preserve eager legacy non-stream telemetry. Stream parse caches # remain SDK-owned, and every later explicit parse type is traced. - raw.parse() + try: + parsed = response.parse() + except Exception: + # Optional eager telemetry must not change SDK behavior: callers + # still receive the envelope and may inspect malformed bodies. + logger.debug("Could not eagerly parse raw response for telemetry", exc_info=True) + else: + raw._trace_parsed(parsed) return raw return on_parse(response) diff --git a/python/frameworks/openai/traceai_openai/_stream.py b/python/frameworks/openai/traceai_openai/_stream.py index 7fc65c9f9..650f81182 100644 --- a/python/frameworks/openai/traceai_openai/_stream.py +++ b/python/frameworks/openai/traceai_openai/_stream.py @@ -124,6 +124,29 @@ def text(self, *args: Any, **kwargs: Any) -> Any: def json(self, *args: Any, **kwargs: Any) -> Any: return self._consume_body("json", *args, **kwargs) + def _consume_iterator(self, method: str, *args: Any, **kwargs: Any) -> Iterator[Any]: + try: + for chunk in getattr(self.__wrapped__, method)(*args, **kwargs): + yield chunk + except BaseException as exc: + _finish_on_exit(self, exc) + raise + else: + # The caller may stop iteration early; only exhaustion consumes it. + self._record_body_consumption() + + def iter_bytes(self, *args: Any, **kwargs: Any) -> Iterator[bytes]: + return self._consume_iterator("iter_bytes", *args, **kwargs) + + def iter_text(self, *args: Any, **kwargs: Any) -> Iterator[str]: + return self._consume_iterator("iter_text", *args, **kwargs) + + def iter_lines(self, *args: Any, **kwargs: Any) -> Iterator[str]: + return self._consume_iterator("iter_lines", *args, **kwargs) + + def stream_to_file(self, *args: Any, **kwargs: Any) -> None: + self._consume_body("stream_to_file", *args, **kwargs) + @property def _self_is_finished(self) -> bool: return self._self_with_span.is_finished @@ -200,6 +223,31 @@ async def text(self, *args: Any, **kwargs: Any) -> Any: async def json(self, *args: Any, **kwargs: Any) -> Any: return await self._consume_body("json", *args, **kwargs) + async def _consume_iterator(self, method: str, *args: Any, **kwargs: Any) -> AsyncIterator[Any]: + try: + async for chunk in getattr(self.__wrapped__, method)(*args, **kwargs): + yield chunk + except BaseException as exc: + _finish_on_exit(self, exc) + raise + else: + self._record_body_consumption() + + async def iter_bytes(self, *args: Any, **kwargs: Any) -> AsyncIterator[bytes]: + async for chunk in self._consume_iterator("iter_bytes", *args, **kwargs): + yield chunk + + async def iter_text(self, *args: Any, **kwargs: Any) -> AsyncIterator[str]: + async for chunk in self._consume_iterator("iter_text", *args, **kwargs): + yield chunk + + async def iter_lines(self, *args: Any, **kwargs: Any) -> AsyncIterator[str]: + async for chunk in self._consume_iterator("iter_lines", *args, **kwargs): + yield chunk + + async def stream_to_file(self, *args: Any, **kwargs: Any) -> None: + await self._consume_body("stream_to_file", *args, **kwargs) + async def parse(self, *args: Any, **kwargs: Any) -> Any: try: parsed = await self.__wrapped__.parse(*args, **kwargs) From ae1e21b0150d2440ddda946c4a452fcbbc36aeae Mon Sep 17 00:00:00 2001 From: Nikhil Pareek Date: Mon, 2 Nov 2026 13:44:54 -0800 Subject: [PATCH 10/10] fix(openai): preserve raw Responses privacy and SDK contracts Co-Authored-By: Claude Code --- .../tests/test_raw_envelope_contract.py | 189 ++++++++++++++++++ .../openai/traceai_openai/_request.py | 35 +++- .../openai/traceai_openai/_stream.py | 93 ++++++++- 3 files changed, 308 insertions(+), 9 deletions(-) create mode 100644 python/frameworks/openai/tests/test_raw_envelope_contract.py diff --git a/python/frameworks/openai/tests/test_raw_envelope_contract.py b/python/frameworks/openai/tests/test_raw_envelope_contract.py new file mode 100644 index 000000000..7f8bba9d0 --- /dev/null +++ b/python/frameworks/openai/tests/test_raw_envelope_contract.py @@ -0,0 +1,189 @@ +"""Raw OpenAI SDK response contracts using synthetic in-memory transport.""" + +import asyncio +import json + +import httpx +import openai +import pytest +from fi_instrumentation import TraceConfig, register +from fi_instrumentation.fi_types import ProjectType +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import StatusCode +from traceai_openai import OpenAIInstrumentor + + +@pytest.fixture +def tracing(): + exporter = InMemorySpanExporter() + provider = register( + project_name="raw-review", project_type=ProjectType.OBSERVE, + span_exporter=exporter, batch=False, verbose=False, + ) + config = TraceConfig( + hide_inputs=True, hide_input_messages=True, + hide_llm_invocation_parameters=True, hide_error_details=True, + ) + OpenAIInstrumentor().instrument(tracer_provider=provider, config=config) + try: + yield exporter + finally: + OpenAIInstrumentor().uninstrument() + provider.shutdown() + + +def body(status="completed"): + return { + "id": "resp_review", "object": "response", "created_at": 1, + "model": "synthetic", "status": status, "parallel_tool_calls": True, + "tool_choice": "auto", "tools": [], "output": [], "error": None, + "usage": { + "input_tokens": 5, "output_tokens": 2, "total_tokens": 7, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens_details": {"reasoning_tokens": 0}, + }, + } + + +def client(payload, async_): + class SyncBody(httpx.SyncByteStream): + def __iter__(self): + yield payload + + class AsyncBody(httpx.AsyncByteStream): + async def __aiter__(self): + yield payload + + def handler(request): + return httpx.Response( + 200, headers={"content-type": "application/json"}, + stream=AsyncBody() if async_ else SyncBody(), + ) + + http_client = (httpx.AsyncClient if async_ else httpx.Client)( + transport=httpx.MockTransport(handler), trust_env=False, + ) + return (openai.AsyncOpenAI if async_ else openai.OpenAI)( + api_key="synthetic-credential", max_retries=0, http_client=http_client, + ) + + +@pytest.mark.parametrize("async_", [False, True]) +@pytest.mark.parametrize("parse_to", ["default", "dict", "str"]) +def test_raw_custom_parse_respects_input_and_error_privacy(tracing, async_, parse_to): + hidden_input = "synthetic-private-input-24193" + hidden_error = "synthetic-private-provider-error-88123" + response = body("failed") + response.update( + instructions=hidden_input, metadata={"private": hidden_input}, + input=hidden_input, error={"code": "server_error", "message": hidden_error}, + ) + payload = json.dumps(response).encode() + kwargs = {} if parse_to == "default" else { + "to": dict[str, object] if parse_to == "dict" else str, + } + + async def run_async(): + async with client(payload, True) as c: + async with c.responses.with_streaming_response.create( + model="synthetic", input=hidden_input, instructions=hidden_input, + ) as raw: + return await raw.parse(**kwargs) + + if async_: + result = asyncio.run(run_async()) + else: + with client(payload, False) as c: + with c.responses.with_streaming_response.create( + model="synthetic", input=hidden_input, instructions=hidden_input, + ) as raw: + result = raw.parse(**kwargs) + + # The application still receives the original unmodified response. + application_data = ( + result if isinstance(result, str) else json.dumps(result) if isinstance(result, dict) + else result.model_dump_json() + ) + assert hidden_error in application_data + span, = tracing.get_finished_spans() + assert hidden_input not in span.to_json() + assert hidden_error not in span.to_json() + assert span.status.status_code is StatusCode.ERROR + + +@pytest.mark.parametrize("async_", [False, True]) +@pytest.mark.parametrize("status", ["completed", "failed"]) +@pytest.mark.parametrize("mode", ["read", "iter_bytes", "iter_text", "iter_lines"]) +def test_non_sse_raw_iteration_uses_body_status(tracing, async_, status, mode): + payload = json.dumps(body(status)).encode() + + async def run_async(): + async with client(payload, True) as c: + async with c.responses.with_streaming_response.create( + model="synthetic", input="question", + ) as raw: + if mode == "read": + assert await raw.read() == payload + else: + chunks = [part async for part in getattr(raw, mode)()] + actual = (b"" if mode == "iter_bytes" else "").join(chunks) + assert actual == (payload if mode == "iter_bytes" else payload.decode()) + + if async_: + asyncio.run(run_async()) + else: + with client(payload, False) as c: + with c.responses.with_streaming_response.create( + model="synthetic", input="question", + ) as raw: + if mode == "read": + assert raw.read() == payload + else: + chunks = list(getattr(raw, mode)()) + actual = (b"" if mode == "iter_bytes" else "").join(chunks) + assert actual == (payload if mode == "iter_bytes" else payload.decode()) + + span, = tracing.get_finished_spans() + expected = StatusCode.OK if status == "completed" else StatusCode.ERROR + assert span.status.status_code is expected + + +@pytest.mark.parametrize("async_", [False, True]) +@pytest.mark.parametrize("instrumented", [False, True]) +@pytest.mark.parametrize("endpoint", ["responses", "chat"]) +def test_legacy_raw_text_remains_string(tracing, async_, instrumented, endpoint): + if not instrumented: + OpenAIInstrumentor().uninstrument() + response = body() + if endpoint == "chat": + response = { + "id": "chat_review", "object": "chat.completion", "created": 1, + "model": "synthetic", "choices": [{ + "index": 0, "message": {"role": "assistant", "content": "answer"}, + "finish_reason": "stop", + }], + } + payload = json.dumps(response).encode() + + def params(c): + if endpoint == "responses": + return c.responses, {"model": "synthetic", "input": "question"} + return c.chat.completions, { + "model": "synthetic", "messages": [{"role": "user", "content": "question"}], + } + + async def run_async(): + async with client(payload, True) as c: + resource, kwargs = params(c) + raw = await resource.with_raw_response.create(**kwargs) + assert raw.content == payload + assert json.loads(raw.text) == response + + if async_: + asyncio.run(run_async()) + else: + with client(payload, False) as c: + resource, kwargs = params(c) + raw = resource.with_raw_response.create(**kwargs) + assert raw.content == payload + assert json.loads(raw.text) == response diff --git a/python/frameworks/openai/traceai_openai/_request.py b/python/frameworks/openai/traceai_openai/_request.py index eec002459..98aec1227 100644 --- a/python/frameworks/openai/traceai_openai/_request.py +++ b/python/frameworks/openai/traceai_openai/_request.py @@ -1,3 +1,4 @@ +import json import logging from abc import ABC from contextlib import contextmanager @@ -35,6 +36,7 @@ from traceai_openai._stream import ( _AsyncRawResponse, _AsyncStream, + _LegacyRawResponse, _RawResponse, _ResponseAccumulator, _Stream, @@ -248,14 +250,35 @@ def on_parse(parsed: Any) -> Any: with_span.set_attribute( SpanAttributes.EMBEDDING_EMBEDDINGS, self._raw_embeddings_json(parsed) ) - add_io_to_span_attributes(with_span, None, parsed, config=self._config) - + telemetry_response = parsed + status = _response_status(parsed) + if cast_to is self._openai.types.responses.response.Response and not isinstance( + parsed, cast_to + ): + # parse(to=dict/str) returns the original SDK value to the caller, + # but telemetry must pass through the same Response allowlist as + # the default parse. Never serialize the arbitrary raw value. + try: + body = json.loads(parsed) if isinstance(parsed, str) else parsed + if not isinstance(body, Mapping): + raise ValueError("Responses body is not an object") + status = _response_status(body) + telemetry_response = cast_to.model_validate(body) + except (TypeError, ValueError): + telemetry_response = None + status = trace_api.Status( + trace_api.StatusCode.ERROR, "Response body could not be inspected" + ) + if telemetry_response is not None: + add_io_to_span_attributes( + with_span, None, telemetry_response, config=self._config + ) _finish_tracing( - status=_response_status(parsed), + status=status, with_span=with_span, has_attributes=_ResponseAttributes( request_parameters=request_parameters, - response=parsed, + response=telemetry_response, response_attributes_extractor=self._response_attributes_extractor, ), ) @@ -265,7 +288,9 @@ def on_parse(parsed: Any) -> Any: # Modern envelopes are deliberately lazy, even for non-SSE responses. # Do not call async parse from sync I/O/attribute extraction paths. raw_wrapper = ( - _AsyncRawResponse if iscoroutinefunction(response.parse) else _RawResponse + _LegacyRawResponse if isinstance(response, LegacyAPIResponse) + else _AsyncRawResponse if iscoroutinefunction(response.parse) + else _RawResponse ) raw = raw_wrapper(response, with_span, response_accumulator, on_parse) if isinstance(response, LegacyAPIResponse): diff --git a/python/frameworks/openai/traceai_openai/_stream.py b/python/frameworks/openai/traceai_openai/_stream.py index 7fc65c9f9..3629fa3ab 100644 --- a/python/frameworks/openai/traceai_openai/_stream.py +++ b/python/frameworks/openai/traceai_openai/_stream.py @@ -1,3 +1,4 @@ +import json import logging from typing import Any, AsyncIterator, Callable, Generator, Iterator, Optional, Protocol, Tuple @@ -90,10 +91,10 @@ def __init__( self._self_parsed = {} self._self_body_status: Optional[trace_api.Status] = None - def _record_body_consumption(self) -> None: + def _record_body_consumption(self, body_bytes: Optional[bytes] = None) -> None: # read/text/json consume a non-SSE envelope without calling parse(). - # Do not force SDK parsing (or its post_parser) on a raw-body caller. - # SSE still needs a terminal event from the traced iterator. + # Iterators do not cache HTTPX content, so inspect their collected bytes. + # Never invoke the SDK post_parser or export the raw body. if ( self._self_is_finished or self._self_body_status is not None @@ -101,7 +102,10 @@ def _record_body_consumption(self) -> None: ): return try: - body = self.__wrapped__.http_response.json() + body = ( + json.loads(body_bytes) if body_bytes is not None + else self.__wrapped__.http_response.json() + ) except Exception: body = None # Successful raw reads need not contain JSON. self._self_body_status = _response_status(body) @@ -124,6 +128,43 @@ def text(self, *args: Any, **kwargs: Any) -> Any: def json(self, *args: Any, **kwargs: Any) -> Any: return self._consume_body("json", *args, **kwargs) + def _iterate_body(self, method: str, *args: Any, **kwargs: Any) -> Iterator[Any]: + is_sse = getattr(self.__wrapped__, "_is_sse_stream", False) + body = bytearray() + too_large = False + try: + for chunk in getattr(self.__wrapped__, method)(*args, **kwargs): + if not is_sse and not too_large: + encoded = chunk if isinstance(chunk, bytes) else chunk.encode("utf-8") + if method == "iter_lines": + encoded += b"\n" + if len(body) + len(encoded) <= 8 * 1024 * 1024: + body.extend(encoded) + else: + too_large = True + body.clear() + yield chunk + except BaseException as exc: + _finish_on_exit(self, exc) + raise + else: + # Partial iteration and SSE must not report a completed response. + if too_large: + self._self_body_status = trace_api.Status( + trace_api.StatusCode.ERROR, "Response body too large to inspect" + ) + elif not is_sse: + self._record_body_consumption(bytes(body)) + + def iter_bytes(self, *args: Any, **kwargs: Any) -> Iterator[bytes]: + return self._iterate_body("iter_bytes", *args, **kwargs) + + def iter_text(self, *args: Any, **kwargs: Any) -> Iterator[str]: + return self._iterate_body("iter_text", *args, **kwargs) + + def iter_lines(self, *args: Any, **kwargs: Any) -> Iterator[str]: + return self._iterate_body("iter_lines", *args, **kwargs) + @property def _self_is_finished(self) -> bool: return self._self_with_span.is_finished @@ -179,6 +220,14 @@ def _finish_tracing(self, status: Optional[trace_api.Status] = None) -> None: self._self_with_span.finish_tracing(status=status) +class _LegacyRawResponse(_RawResponse): + """LegacyAPIResponse.text is a string property, unlike APIResponse.text().""" + + @property + def text(self) -> str: + return self.__wrapped__.text + + class _AsyncRawResponse(_RawResponse): """AsyncAPIResponse has awaitable parse/close; LegacyAPIResponse does not.""" @@ -200,6 +249,42 @@ async def text(self, *args: Any, **kwargs: Any) -> Any: async def json(self, *args: Any, **kwargs: Any) -> Any: return await self._consume_body("json", *args, **kwargs) + async def _aiterate_body(self, method: str, *args: Any, **kwargs: Any) -> AsyncIterator[Any]: + is_sse = getattr(self.__wrapped__, "_is_sse_stream", False) + body = bytearray() + too_large = False + try: + async for chunk in getattr(self.__wrapped__, method)(*args, **kwargs): + if not is_sse and not too_large: + encoded = chunk if isinstance(chunk, bytes) else chunk.encode("utf-8") + if method == "iter_lines": + encoded += b"\n" + if len(body) + len(encoded) <= 8 * 1024 * 1024: + body.extend(encoded) + else: + too_large = True + body.clear() + yield chunk + except BaseException as exc: + _finish_on_exit(self, exc) + raise + else: + if too_large: + self._self_body_status = trace_api.Status( + trace_api.StatusCode.ERROR, "Response body too large to inspect" + ) + elif not is_sse: + self._record_body_consumption(bytes(body)) + + def iter_bytes(self, *args: Any, **kwargs: Any) -> AsyncIterator[bytes]: + return self._aiterate_body("iter_bytes", *args, **kwargs) + + def iter_text(self, *args: Any, **kwargs: Any) -> AsyncIterator[str]: + return self._aiterate_body("iter_text", *args, **kwargs) + + def iter_lines(self, *args: Any, **kwargs: Any) -> AsyncIterator[str]: + return self._aiterate_body("iter_lines", *args, **kwargs) + async def parse(self, *args: Any, **kwargs: Any) -> Any: try: parsed = await self.__wrapped__.parse(*args, **kwargs)