From f39a54ee03775a8c9646d20540adbeb1d28390de Mon Sep 17 00:00:00 2001 From: Nikhil Pareek Date: Mon, 2 Nov 2026 08:53:48 -0800 Subject: [PATCH 1/7] fix(openai): attribute dedicated Baseten Python endpoints Recognize bounded dedicated and regional model hostnames without treating arbitrary Baseten subdomains as model endpoints. Pin host boundaries and real SDK-to-OTLP contracts on the declared OpenAI floor and registry latest. Co-Authored-By: Claude Code --- .../workflows/baseten-python-contracts.yml | 61 +++++ python/frameworks/openai/README.md | 19 ++ .../openai/tests/test_baseten_provider.py | 224 ++++++++++++++++++ .../openai/traceai_openai/_providers.py | 5 + 4 files changed, 309 insertions(+) create mode 100644 .github/workflows/baseten-python-contracts.yml create mode 100644 python/frameworks/openai/tests/test_baseten_provider.py diff --git a/.github/workflows/baseten-python-contracts.yml b/.github/workflows/baseten-python-contracts.yml new file mode 100644 index 000000000..d9a2a411d --- /dev/null +++ b/.github/workflows/baseten-python-contracts.yml @@ -0,0 +1,61 @@ +name: Baseten Python contracts + +on: + pull_request: + branches: [dev, main] + paths: + - 'python/frameworks/openai/**' + - 'python/fi_instrumentation/**' + - 'python/examples/baseten/**' + - 'python/tests/harness/**' + - '.github/workflows/baseten-python-contracts.yml' + push: + branches: [dev] + paths: + - 'python/frameworks/openai/**' + - 'python/fi_instrumentation/**' + - 'python/examples/baseten/**' + - 'python/tests/harness/**' + - '.github/workflows/baseten-python-contracts.yml' + +permissions: + contents: read + +jobs: + contracts: + name: Baseten OpenAI ${{ matrix.sdk }} + runs-on: ubuntu-24.04 + timeout-minutes: 15 + strategy: + fail-fast: false + matrix: + sdk: ['1.69.0', 'latest'] + env: + PYTHONDONTWRITEBYTECODE: '1' + SDK_VERSION: ${{ matrix.sdk }} + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + ref: ${{ github.event.pull_request.head.sha || github.sha }} + persist-credentials: false + - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: '3.12' + - name: Install source packages and the selected SDK + run: | + sdk=openai + if [ "$SDK_VERSION" != latest ]; then sdk="openai==$SDK_VERSION"; fi + python -m pip install -e python -e python/frameworks/openai "$sdk" httpx pytest + python -c 'import openai; print("Resolved OpenAI SDK:", openai.__version__)' + - name: Run network-isolated provider and recipe contracts + run: | + sudo unshare --net -- bash -c 'ip link set lo up; exec "$@"' bash \ + env PATH="$PATH" PYTHONDONTWRITEBYTECODE=1 \ + python -m pytest python/frameworks/openai/tests/test_baseten_provider.py \ + python/frameworks/openai/tests/test_openai_compatible_spans.py \ + --noconftest -p no:cacheprovider -o addopts='' -q + sudo unshare --net -- bash -c 'ip link set lo up; exec "$@"' bash \ + env PATH="$PATH" PYTHONDONTWRITEBYTECODE=1 \ + PYTHONPATH="python/examples/baseten/src:python/tests" \ + python -m pytest python/examples/baseten/tests \ + --noconftest -p no:cacheprovider -o addopts='' -q diff --git a/python/frameworks/openai/README.md b/python/frameworks/openai/README.md index ea7036cec..319f658d7 100644 --- a/python/frameworks/openai/README.md +++ b/python/frameworks/openai/README.md @@ -100,6 +100,25 @@ with OpenAI() as client: trace_provider.force_flush() # Flush before a short-lived script exits. ``` +## Baseten + +Use the same `OpenAIInstrumentor` with Baseten's OpenAI-compatible clients; no +second instrumentor is needed. After the Quickstart tracer setup, create your +client with `api_key=os.environ["BASETEN_API_KEY"]` and the appropriate base URL: + +- Model APIs: `https://inference.baseten.co/v1`. +- Dedicated deployments: your deployment's OpenAI-compatible base URL, such as + `https://model-abc123.api.baseten.co/environments/production/sync/v1`. + +The hosted endpoint and dedicated `model-*.api.baseten.co` endpoints, including +regional model endpoints, are attributed to `baseten`. Unknown proxies retain +`openai` unless you set an explicit client provider override. Supported API +routes and model names depend on the deployment; use the values in +[Baseten's inference API reference](https://docs.baseten.co/reference/inference-api/overview). +The standalone Model APIs recipe in `python/examples/baseten` accepts only Model +APIs URLs; use the OpenAI client directly for a dedicated deployment. Baseten's +separate Performance Client and non-OpenAI protocols need their own tracing setup. + ## Vercel AI Gateway Use the same `OpenAIInstrumentor` for Chat Completions sent to Vercel AI Gateway; diff --git a/python/frameworks/openai/tests/test_baseten_provider.py b/python/frameworks/openai/tests/test_baseten_provider.py new file mode 100644 index 000000000..0ffe6b817 --- /dev/null +++ b/python/frameworks/openai/tests/test_baseten_provider.py @@ -0,0 +1,224 @@ +"""Baseten attribution through the real OpenAI SDK and a loopback OTLP receiver. + +HTTP model responses and credentials are synthetic. Resource coverage asserts +instrumentor behavior, not that every Baseten deployment serves every API. +""" + +import asyncio +import json +import sys +from pathlib import Path +from urllib.parse import urlsplit + +import httpx +import openai +import pytest +from fi_instrumentation import TraceConfig, register +from fi_instrumentation.fi_types import ProjectType + +from traceai_openai import OpenAIInstrumentor +from traceai_openai._providers import provider_from_host + +from .test_openai_compatible_spans import ( + ANSWER, REQUEST_MODEL, RESPONSE_MODEL, async_client, chat_completion, + chat_stream, client, embeddings_body, responses_body, responses_stream, unauthorized, +) + +# Use the repository's collector harness instead of a second telemetry implementation. +sys.path.insert(0, str(Path(__file__).resolve().parents[3] / "tests")) +from harness import Receiver, _flatten_attributes + +BASE_URLS = [ + "https://inference.baseten.co/v1", + "https://model-abc123.api.baseten.co/environments/production/sync/v1", + "https://model-abc123-us-west.api.baseten.co/sync/v1", +] +PROMPT = "synthetic-baseten-private-prompt" + + +@pytest.mark.parametrize("host,expected", [ + ("inference.baseten.co", "baseten"), + ("model-abc123.api.baseten.co", "baseten"), + ("model-abc123-us-west.api.baseten.co", "baseten"), + ("MODEL-Abc123.API.BASETEN.CO.", "baseten"), + ("model-a.api.baseten.co", "baseten"), + ("model-" + "a" * 57 + ".api.baseten.co", "baseten"), + ("model-" + "a" * 58 + ".api.baseten.co", "openai"), + ("model-.api.baseten.co", "openai"), + ("model--abc.api.baseten.co", "openai"), + ("model-abc-.api.baseten.co", "openai"), + ("model-abc_123.api.baseten.co", "openai"), + ("model-abc.123.api.baseten.co", "openai"), + ("model-abc.api.baseten.co.customer.example", "openai"), + ("proxy.model-abc.api.baseten.co", "openai"), + ("notmodel-abc.api.baseten.co", "openai"), + ("chain-abc.api.baseten.co", "openai"), + ("api.baseten.co", "openai"), + ("baseten.co", "openai"), + ("inference.baseten.co.customer.example", "openai"), + ("api.openai.com", "openai"), + (None, "openai"), +]) +def test_baseten_host_boundaries(host, expected): + assert provider_from_host(host) == expected + + +@pytest.fixture(scope="module") +def receiver(): + with Receiver() as sink: + yield sink + + +@pytest.fixture +def wire_tracer(receiver, monkeypatch): + # No inherited telemetry destinations, auth, privacy settings or HTTP proxies. + import os + for name in tuple(os.environ): + if name.startswith(("FI_", "OTEL_", "OPENAI_")) or name.lower() in ( + "http_proxy", "https_proxy", "all_proxy", "no_proxy" + ): + monkeypatch.delenv(name) + monkeypatch.setenv("NO_PROXY", "*") + monkeypatch.setenv("FI_API_KEY", "synthetic-fi-key") + monkeypatch.setenv("FI_SECRET_KEY", "synthetic-fi-secret") + monkeypatch.setenv("FI_BASE_URL", receiver.origin) + receiver.clear() + provider = register( + project_name="baseten-contract", project_type=ProjectType.OBSERVE, + set_global_tracer_provider=False, verbose=False, + ) + instrumentor = OpenAIInstrumentor() + + def configure(hidden=False): + instrumentor.instrument(tracer_provider=provider, config=TraceConfig( + hide_inputs=hidden, hide_outputs=hidden, + hide_llm_invocation_parameters=hidden, + )) + return provider + + try: + yield configure + finally: + instrumentor.uninstrument() + provider.force_flush() + provider.shutdown() + + +def only_span(receiver, provider): + assert provider.force_flush() + spans = receiver.spans() + assert len(spans) == 1 + exports = receiver.requests() + assert len(exports) == 1 + assert exports[0]["path"] == "/tracer/v1/traces" + assert exports[0]["headers"]["x-api-key"] == "synthetic-fi-key" + assert exports[0]["headers"]["x-secret-key"] == "synthetic-fi-secret" + assert exports[0]["resource_attributes"][0]["project_name"] == "baseten-contract" + assert exports[0]["resource_attributes"][0]["project_type"] == "observe" + for secret in ("test-vendor-key", "synthetic-fi-key", "synthetic-fi-secret"): + assert secret not in json.dumps(spans) + return spans[0], _flatten_attributes(spans[0]["attributes"]) + + +@pytest.mark.parametrize("base_url", BASE_URLS, ids=["model-api", "dedicated", "regional"]) +@pytest.mark.parametrize("use_async", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("resource,outcome", [ + ("chat", "success"), ("chat", "stream"), ("chat", "error"), + ("responses", "success"), ("responses", "stream"), ("responses", "error"), + ("embeddings", "success"), ("embeddings", "error"), +]) +@pytest.mark.parametrize("hidden", [False, True], ids=["content", "masked"]) +def test_baseten_wire_contract(receiver, wire_tracer, base_url, use_async, resource, outcome, hidden): + provider = wire_tracer(hidden) + seen = [] + route = "chat/completions" if resource == "chat" else resource + + def handler(request): + seen.append(request) + assert request.url.path == urlsplit(base_url).path + "/" + route + assert request.headers["authorization"] == "Bearer test-vendor-key" + if outcome == "error": + return unauthorized(request) + if outcome == "stream": + return chat_stream() if resource == "chat" else responses_stream() + body = {"chat": chat_completion, "responses": responses_body, + "embeddings": embeddings_body}[resource]() + return httpx.Response(200, json=body) + + kwargs = {"model": REQUEST_MODEL} + if resource == "chat": + kwargs.update(messages=[{"role": "user", "content": PROMPT}], stream=outcome == "stream") + else: + kwargs["input"] = PROMPT + if resource == "responses": + kwargs["stream"] = outcome == "stream" + else: + kwargs["encoding_format"] = "float" + + def create(sdk): + endpoint = sdk.chat.completions if resource == "chat" else getattr(sdk, resource) + return endpoint.create(**kwargs) + + async def run_async(): + async with async_client(handler, base_url=base_url) as sdk: + response = await create(sdk) + if outcome == "stream": + async for _ in response: + pass + + def run_sync(): + with client(handler, base_url=base_url) as sdk: + response = create(sdk) + if outcome == "stream": + list(response) + + if outcome == "error": + with pytest.raises(openai.AuthenticationError): + asyncio.run(run_async()) if use_async else run_sync() + else: + asyncio.run(run_async()) if use_async else run_sync() + assert len(seen) == 1 + span, attrs = only_span(receiver, provider) + assert attrs["gen_ai.provider.name"] == "baseten" + assert attrs["gen_ai.request.model"] == ( + REQUEST_MODEL if outcome == "error" or (resource == "chat" and outcome == "stream") + else RESPONSE_MODEL + ) + assert span["status"]["code"] == ("STATUS_CODE_ERROR" if outcome == "error" else "STATUS_CODE_OK") + if resource == "chat" and outcome == "success": + assert int(attrs["gen_ai.usage.input_tokens"]) == 5 + assert int(attrs["gen_ai.usage.output_tokens"]) == 2 + if hidden: + assert PROMPT not in json.dumps(span) + assert ANSWER not in json.dumps(span) + + +@pytest.mark.parametrize("base_url,expected", [ + ("https://MODEL-Abc123.API.BASETEN.CO.:443/v1", "baseten"), + ("https://model-abc.api.baseten.co.customer.example/v1", "openai"), + ("https://proxy.model-abc.api.baseten.co/v1", "openai"), + ("https://model-abc.api.baseten.co@customer.example/v1", "openai"), + ("https://customer.example/model-abc.api.baseten.co/v1", "openai"), +]) +def test_baseten_url_controls(receiver, wire_tracer, base_url, expected): + provider = wire_tracer() + with client(lambda _: httpx.Response(200, json=chat_completion()), base_url=base_url) as sdk: + sdk.chat.completions.create(model=REQUEST_MODEL, messages=[{"role": "user", "content": PROMPT}]) + _, attrs = only_span(receiver, provider) + assert attrs["gen_ai.provider.name"] == expected + + +def test_baseten_client_endpoint_changes(receiver, wire_tracer): + provider = wire_tracer() + with client(lambda _: httpx.Response(200, json=chat_completion())) as sdk: + for base_url, expected in [ + *[(url, "baseten") for url in BASE_URLS], + ("https://customer.example/v1", "openai"), + ]: + sdk.base_url = base_url + sdk.chat.completions.create( + model=REQUEST_MODEL, messages=[{"role": "user", "content": PROMPT}], + ) + _, attrs = only_span(receiver, provider) + assert attrs["gen_ai.provider.name"] == expected + receiver.clear() diff --git a/python/frameworks/openai/traceai_openai/_providers.py b/python/frameworks/openai/traceai_openai/_providers.py index 2c426e46f..22330501b 100644 --- a/python/frameworks/openai/traceai_openai/_providers.py +++ b/python/frameworks/openai/traceai_openai/_providers.py @@ -6,6 +6,7 @@ instrumentor's historical label, "openai", rather than a guess. """ +import re from typing import Optional from fi_instrumentation.fi_types import FiLLMProviderValues @@ -58,6 +59,10 @@ def provider_from_host(host: Optional[str]) -> str: host = host.lower().rstrip(".") if host in _HOSTS: return _HOSTS[host] + # Dedicated and regional Baseten endpoints use a single model-* DNS label. + # Bound it to 63 characters; do not label arbitrary subdomains or look-alikes. + if re.fullmatch(r"model-[a-z0-9](?:[a-z0-9-]{0,55}[a-z0-9])?\.api\.baseten\.co", host): + return "baseten" for suffix, provider in _SUFFIXES: if host.endswith(suffix): return provider From 7b29e911e5a43315dde775c2801d52a954908063 Mon Sep 17 00:00:00 2001 From: Nikhil Pareek Date: Mon, 2 Nov 2026 09:23:32 -0800 Subject: [PATCH 2/7] fix(openai): finalize spans on explicit close and cancellation Preserve SDK cleanup and cancellation propagation while ending each span once. Cover native floor/latest transports, close failures, error privacy, and repeated cleanup in Baseten CI contracts. Co-Authored-By: Claude Code --- .../workflows/baseten-python-contracts.yml | 2 + .../openai/tests/test_stream_lifecycle.py | 258 ++++++++++++++++++ .../openai/traceai_openai/_request.py | 3 +- .../openai/traceai_openai/_stream.py | 33 ++- .../openai/traceai_openai/_with_span.py | 2 +- 5 files changed, 287 insertions(+), 11 deletions(-) create mode 100644 python/frameworks/openai/tests/test_stream_lifecycle.py diff --git a/.github/workflows/baseten-python-contracts.yml b/.github/workflows/baseten-python-contracts.yml index d9a2a411d..15cc6dd54 100644 --- a/.github/workflows/baseten-python-contracts.yml +++ b/.github/workflows/baseten-python-contracts.yml @@ -53,6 +53,8 @@ jobs: env PATH="$PATH" PYTHONDONTWRITEBYTECODE=1 \ python -m pytest python/frameworks/openai/tests/test_baseten_provider.py \ python/frameworks/openai/tests/test_openai_compatible_spans.py \ + python/frameworks/openai/tests/test_stream_lifecycle.py \ + python/frameworks/openai/tests/test_stream_context_manager.py \ --noconftest -p no:cacheprovider -o addopts='' -q sudo unshare --net -- bash -c 'ip link set lo up; exec "$@"' bash \ env PATH="$PATH" PYTHONDONTWRITEBYTECODE=1 \ diff --git a/python/frameworks/openai/tests/test_stream_lifecycle.py b/python/frameworks/openai/tests/test_stream_lifecycle.py new file mode 100644 index 000000000..06daa6160 --- /dev/null +++ b/python/frameworks/openai/tests/test_stream_lifecycle.py @@ -0,0 +1,258 @@ +"""Explicit close and task cancellation through the SDK's native HTTP transport.""" + +import asyncio +import json +import os + +import openai +import pytest +from fi_instrumentation import TraceConfig +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 + +from traceai_openai import OpenAIInstrumentor + +# OpenAI 3 uses httpx2; do not exercise it through the legacy httpx adapter. +if int(openai.__version__.split(".")[0]) >= 3: + import httpx2 as http +else: + import httpx as http + +ARGS = {"model": "requested-model", "messages": [{"role": "user", "content": "hello"}]} +CHUNK = { + "id": "test", "object": "chat.completion.chunk", "created": 1, "model": "served-model", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "partial-answer"}, + "finish_reason": None}], +} +CHUNK_BYTES = ("data: " + json.dumps(CHUNK) + "\n\n").encode() + + +@pytest.fixture(params=[ + "https://model-synthetic.api.baseten.co/environments/production/sync/v1", + "https://model-synthetic-us-west.api.baseten.co/sync/v1", +]) +def base_url(request): + return request.param + + +@pytest.fixture +def tracing(monkeypatch, request): + for name in tuple(os.environ): + if name.startswith(("FI_", "OTEL_", "OPENAI_")): + monkeypatch.delenv(name) + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + instrumentor = OpenAIInstrumentor() + instrumentor.instrument( + tracer_provider=provider, + config=TraceConfig(hide_error_details=getattr(request, "param", False)), + ) + yield provider, exporter + instrumentor.uninstrument() + provider.shutdown() + + +def client(handler, base_url, use_async=False): + sdk_type = openai.AsyncOpenAI if use_async else openai.OpenAI + http_type = http.AsyncClient if use_async else http.Client + return sdk_type( + api_key="synthetic-key", base_url=base_url, max_retries=0, + http_client=http_type(transport=http.MockTransport(handler), trust_env=False), + ) + + +def streaming(_): + return http.Response( + 200, headers={"content-type": "text/event-stream"}, + content=CHUNK_BYTES + b"data: [DONE]\n\n", + ) + + +def only_span(tracing, status): + provider, exporter = tracing + assert provider.force_flush() + spans = exporter.get_finished_spans() + assert len(spans) == 1 + span = spans[0] + assert span.status.status_code == status + assert span.attributes["gen_ai.provider.name"] == "baseten" + return span + + +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.parametrize("consumption", ["empty", "partial", "exhausted"]) +def test_explicit_close_finishes_once(tracing, base_url, use_async, consumption): + if use_async: + async def run(): + async with client(streaming, base_url, True) as sdk: + stream = await sdk.chat.completions.create(**ARGS, stream=True) + if consumption == "partial": + await stream.__anext__() + elif consumption == "exhausted": + async for _ in stream: + pass + assert await stream.close() is None + assert stream.response.is_closed + only_span(tracing, StatusCode.OK) + # Repeated close and context-manager cleanup must not end twice. + await stream.close() + await stream.__aexit__(None, None, None) + asyncio.run(run()) + else: + with client(streaming, base_url) as sdk: + stream = sdk.chat.completions.create(**ARGS, stream=True) + if consumption == "partial": + next(stream) + elif consumption == "exhausted": + list(stream) + assert stream.close() is None + assert stream.response.is_closed + only_span(tracing, StatusCode.OK) + stream.close() + stream.__exit__(None, None, None) + span = only_span(tracing, StatusCode.OK) + if consumption != "empty": + assert span.attributes["output.value"] == "partial-answer" + assert span.attributes["gen_ai.request.model"] == ARGS["model"] + + +@pytest.mark.parametrize("use_async", [False, True]) +def test_close_failure_preserves_exception_and_finishes(tracing, base_url, use_async, monkeypatch): + error = RuntimeError("synthetic close failure") + calls = [] + + if use_async: + async def fail_close(): + calls.append(True) + raise error + + async def run(): + async with client(streaming, base_url, True) as sdk: + stream = await sdk.chat.completions.create(**ARGS, stream=True) + await stream.__anext__() + with monkeypatch.context() as patch: + patch.setattr(stream.response, "aclose", fail_close) + with pytest.raises(RuntimeError) as caught: + await stream.close() + assert caught.value is error + await stream.close() + asyncio.run(run()) + else: + def fail_close(): + calls.append(True) + raise error + + with client(streaming, base_url) as sdk: + stream = sdk.chat.completions.create(**ARGS, stream=True) + next(stream) + with monkeypatch.context() as patch: + patch.setattr(stream.response, "close", fail_close) + with pytest.raises(RuntimeError) as caught: + stream.close() + assert caught.value is error + stream.close() + assert calls == [True] + span = only_span(tracing, StatusCode.ERROR) + assert sum(event.name == "exception" for event in span.events) == 1 + + +@pytest.mark.parametrize("tracing,hide_errors", [(False, False), (True, True)], indirect=["tracing"]) +def test_cancel_pending_request_finishes(tracing, base_url, hide_errors): + async def run(): + reached_transport = asyncio.Event() + + async def handler(_): + reached_transport.set() + await asyncio.Event().wait() + + async with client(handler, base_url, True) as sdk: + task = asyncio.create_task(sdk.chat.completions.create(**ARGS)) + await asyncio.wait_for(reached_transport.wait(), 2) + task.cancel("synthetic request cancellation") + with pytest.raises(asyncio.CancelledError, match="synthetic request cancellation"): + await task + # Finalization must not depend on subsequent client cleanup. + only_span(tracing, StatusCode.ERROR) + asyncio.run(run()) + span = only_span(tracing, StatusCode.ERROR) + if hide_errors: + assert "synthetic request cancellation" not in span.to_json() + else: + assert any(event.attributes["exception.type"] == "asyncio.exceptions.CancelledError" + for event in span.events if event.name == "exception") + + +@pytest.mark.parametrize("tracing,hide_errors", [(False, False), (True, True)], indirect=["tracing"]) +def test_cancel_stream_iteration_finishes_before_close(tracing, base_url, hide_errors): + async def run(): + waiting_for_chunk = asyncio.Event() + closed = [] + + class SlowBody(http.AsyncByteStream): + async def __aiter__(self): + yield CHUNK_BYTES + waiting_for_chunk.set() + await asyncio.Event().wait() + + async def aclose(self): + closed.append(True) + + def handler(_): + return http.Response( + 200, headers={"content-type": "text/event-stream"}, stream=SlowBody(), + ) + + async with client(handler, base_url, True) as sdk: + stream = await sdk.chat.completions.create(**ARGS, stream=True) + await stream.__anext__() + task = asyncio.create_task(stream.__anext__()) + await asyncio.wait_for(waiting_for_chunk.wait(), 2) + task.cancel("synthetic stream cancellation") + with pytest.raises(asyncio.CancelledError, match="synthetic stream cancellation"): + await task + only_span(tracing, StatusCode.ERROR) + await stream.close() + await stream.close() + await stream.__aexit__(None, None, None) + assert stream.response.is_closed + assert closed == [True] + asyncio.run(run()) + span = only_span(tracing, StatusCode.ERROR) + assert span.attributes["output.value"] == "partial-answer" + if hide_errors: + assert "synthetic stream cancellation" not in span.to_json() + else: + assert sum(event.name == "exception" for event in span.events) == 1 + + +def test_cancel_async_close_finishes(tracing, base_url): + async def run(): + closing = asyncio.Event() + + class SlowCloseBody(http.AsyncByteStream): + async def __aiter__(self): + yield CHUNK_BYTES + + async def aclose(self): + closing.set() + await asyncio.Event().wait() + + def handler(_): + return http.Response( + 200, headers={"content-type": "text/event-stream"}, stream=SlowCloseBody(), + ) + + async with client(handler, base_url, True) as sdk: + stream = await sdk.chat.completions.create(**ARGS, stream=True) + await stream.__anext__() + task = asyncio.create_task(stream.close()) + await asyncio.wait_for(closing.wait(), 2) + task.cancel("synthetic close cancellation") + with pytest.raises(asyncio.CancelledError, match="synthetic close cancellation"): + await task + only_span(tracing, StatusCode.ERROR) + asyncio.run(run()) + only_span(tracing, StatusCode.ERROR) diff --git a/python/frameworks/openai/traceai_openai/_request.py b/python/frameworks/openai/traceai_openai/_request.py index c6e8cf426..19f8ff394 100644 --- a/python/frameworks/openai/traceai_openai/_request.py +++ b/python/frameworks/openai/traceai_openai/_request.py @@ -406,7 +406,8 @@ async def __call__( add_io_to_span_attributes( with_span, None, response, is_streaming=self._is_streaming(response) ) - except Exception as exception: + except BaseException as exception: + # end_on_exit=False: cancellation also needs explicit finalization. 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 a3a60042e..efab5c65e 100644 --- a/python/frameworks/openai/traceai_openai/_stream.py +++ b/python/frameworks/openai/traceai_openai/_stream.py @@ -26,7 +26,7 @@ def get_extra_attributes(self) -> Iterator[Tuple[str, AttributeValue]]: ... def _finish_on_exit(stream: Any, error: Optional[BaseException]) -> None: - """Finish even on an early context exit, without ending an exhausted span twice.""" + """Finish on close, context exit or error without ending a span twice.""" if stream._self_is_finished: return if error is None: @@ -74,6 +74,16 @@ def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> Any: finally: _finish_on_exit(self, exc) + def close(self) -> Any: + error: Optional[BaseException] = None + try: + return self.__wrapped__.close() + except BaseException as close_error: + error = close_error + raise + finally: + _finish_on_exit(self, error) + def __iter__(self) -> Iterator[Any]: return self @@ -164,6 +174,16 @@ async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> Any: finally: _finish_on_exit(self, exc) + async def close(self) -> Any: + error: Optional[BaseException] = None + try: + return await self.__wrapped__.close() + except BaseException as close_error: + error = close_error + raise + finally: + _finish_on_exit(self, error) + def __aiter__(self): return self @@ -175,14 +195,9 @@ async def __anext__(self): status = trace_api.Status(status_code=trace_api.StatusCode.OK) self._finish_tracing(status=status) raise - except Exception as exception: - if not self._self_is_finished: - status = trace_api.Status( - status_code=trace_api.StatusCode.ERROR, - description=f"{type(exception).__name__}: {exception}", - ) - self._self_with_span.record_exception(exception) - self._finish_tracing(status=status) + except BaseException as exception: + # Cancellation is a BaseException; finalize without swallowing it. + _finish_on_exit(self, exception) raise else: self._process_chunk(chunk) diff --git a/python/frameworks/openai/traceai_openai/_with_span.py b/python/frameworks/openai/traceai_openai/_with_span.py index 5f5419ad0..923a31af8 100644 --- a/python/frameworks/openai/traceai_openai/_with_span.py +++ b/python/frameworks/openai/traceai_openai/_with_span.py @@ -49,7 +49,7 @@ 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: From ab83f8b2faa6e945f0e866a98e12391740f8b4e0 Mon Sep 17 00:00:00 2001 From: Nikhil Pareek Date: Mon, 2 Nov 2026 09:33:30 -0800 Subject: [PATCH 3/7] test(openai): support legacy core and Python cancellation semantics Keep all lifecycle assertions on Python 3.10 while matching cancellation messages only where Task propagates them. Only error-masking variants skip on the Poetry-locked core that predates that optional feature. Co-Authored-By: Claude Code --- .../openai/tests/test_stream_lifecycle.py | 28 ++++++++++++++----- 1 file changed, 21 insertions(+), 7 deletions(-) diff --git a/python/frameworks/openai/tests/test_stream_lifecycle.py b/python/frameworks/openai/tests/test_stream_lifecycle.py index 06daa6160..41b52f7f9 100644 --- a/python/frameworks/openai/tests/test_stream_lifecycle.py +++ b/python/frameworks/openai/tests/test_stream_lifecycle.py @@ -1,8 +1,10 @@ """Explicit close and task cancellation through the SDK's native HTTP transport.""" import asyncio +import inspect import json import os +import sys import openai import pytest @@ -42,14 +44,16 @@ def tracing(monkeypatch, request): for name in tuple(os.environ): if name.startswith(("FI_", "OTEL_", "OPENAI_")): monkeypatch.delenv(name) + config = TraceConfig() + if getattr(request, "param", False): + if "hide_error_details" not in inspect.signature(TraceConfig).parameters: + pytest.skip("The locked legacy core predates error-detail masking") + config = TraceConfig(hide_error_details=True) exporter = InMemorySpanExporter() provider = TracerProvider() provider.add_span_processor(SimpleSpanProcessor(exporter)) instrumentor = OpenAIInstrumentor() - instrumentor.instrument( - tracer_provider=provider, - config=TraceConfig(hide_error_details=getattr(request, "param", False)), - ) + instrumentor.instrument(tracer_provider=provider, config=config) yield provider, exporter instrumentor.uninstrument() provider.shutdown() @@ -71,6 +75,13 @@ def streaming(_): ) +def cancelled(message): + # Task.cancel(msg) propagates msg to the awaiter only on Python 3.11+. + return pytest.raises( + asyncio.CancelledError, match=message if sys.version_info >= (3, 11) else None, + ) + + def only_span(tracing, status): provider, exporter = tracing assert provider.force_flush() @@ -172,8 +183,9 @@ async def handler(_): task = asyncio.create_task(sdk.chat.completions.create(**ARGS)) await asyncio.wait_for(reached_transport.wait(), 2) task.cancel("synthetic request cancellation") - with pytest.raises(asyncio.CancelledError, match="synthetic request cancellation"): + with cancelled("synthetic request cancellation"): await task + assert task.cancelled() # Finalization must not depend on subsequent client cleanup. only_span(tracing, StatusCode.ERROR) asyncio.run(run()) @@ -211,8 +223,9 @@ def handler(_): task = asyncio.create_task(stream.__anext__()) await asyncio.wait_for(waiting_for_chunk.wait(), 2) task.cancel("synthetic stream cancellation") - with pytest.raises(asyncio.CancelledError, match="synthetic stream cancellation"): + with cancelled("synthetic stream cancellation"): await task + assert task.cancelled() only_span(tracing, StatusCode.ERROR) await stream.close() await stream.close() @@ -251,8 +264,9 @@ def handler(_): task = asyncio.create_task(stream.close()) await asyncio.wait_for(closing.wait(), 2) task.cancel("synthetic close cancellation") - with pytest.raises(asyncio.CancelledError, match="synthetic close cancellation"): + with cancelled("synthetic close cancellation"): await task + assert task.cancelled() only_span(tracing, StatusCode.ERROR) asyncio.run(run()) only_span(tracing, StatusCode.ERROR) From 941d1de433bb4e174599bdf6c4986b542a36b366 Mon Sep 17 00:00:00 2001 From: Nikhil Pareek Date: Mon, 2 Nov 2026 09:59:17 -0800 Subject: [PATCH 4/7] fix(openai): finalize async stream aclose alias Route the SDK alias through instrumented close so early cleanup, close errors and cancellation finish the span once. Cover direct aclose and contextlib.aclosing on dedicated and regional Baseten streams. Co-Authored-By: Claude Code --- .../openai/tests/test_stream_lifecycle.py | 65 +++++++++++++++++-- .../openai/traceai_openai/_stream.py | 4 ++ 2 files changed, 64 insertions(+), 5 deletions(-) diff --git a/python/frameworks/openai/tests/test_stream_lifecycle.py b/python/frameworks/openai/tests/test_stream_lifecycle.py index 41b52f7f9..480c0bfea 100644 --- a/python/frameworks/openai/tests/test_stream_lifecycle.py +++ b/python/frameworks/openai/tests/test_stream_lifecycle.py @@ -1,6 +1,7 @@ """Explicit close and task cancellation through the SDK's native HTTP transport.""" import asyncio +from contextlib import aclosing import inspect import json import os @@ -130,8 +131,54 @@ async def run(): assert span.attributes["gen_ai.request.model"] == ARGS["model"] -@pytest.mark.parametrize("use_async", [False, True]) -def test_close_failure_preserves_exception_and_finishes(tracing, base_url, use_async, monkeypatch): +@pytest.mark.skipif( + not hasattr(openai.AsyncStream, "aclose"), + reason="This SDK version does not expose AsyncStream.aclose", +) +@pytest.mark.parametrize("consumption", ["empty", "partial", "exhausted"]) +@pytest.mark.parametrize("cleanup", ["aclose", "stdlib_aclosing"]) +def test_async_close_alias_finishes_once(tracing, base_url, consumption, cleanup): + async def run(): + async with client(streaming, base_url, True) as sdk: + stream = await sdk.chat.completions.create(**ARGS, stream=True) + + async def consume(): + if consumption == "partial": + await stream.__anext__() + elif consumption == "exhausted": + async for _ in stream: + pass + + if cleanup == "stdlib_aclosing": + async with aclosing(stream): + await consume() + else: + await consume() + assert await stream.aclose() is None + assert stream.response.is_closed + only_span(tracing, StatusCode.OK) + # Interleaved aliases and context cleanup must not export duplicates. + await stream.aclose() + await stream.close() + await stream.__aexit__(None, None, None) + + asyncio.run(run()) + span = only_span(tracing, StatusCode.OK) + if consumption != "empty": + assert span.attributes["output.value"] == "partial-answer" + assert span.attributes["gen_ai.request.model"] == ARGS["model"] + + +@pytest.mark.parametrize("use_async,close_method", [ + (False, "close"), (True, "close"), + pytest.param(True, "aclose", marks=pytest.mark.skipif( + not hasattr(openai.AsyncStream, "aclose"), + reason="This SDK version does not expose AsyncStream.aclose", + )), +]) +def test_close_failure_preserves_exception_and_finishes( + tracing, base_url, use_async, close_method, monkeypatch, +): error = RuntimeError("synthetic close failure") calls = [] @@ -147,8 +194,9 @@ async def run(): with monkeypatch.context() as patch: patch.setattr(stream.response, "aclose", fail_close) with pytest.raises(RuntimeError) as caught: - await stream.close() + await getattr(stream, close_method)() assert caught.value is error + only_span(tracing, StatusCode.ERROR) await stream.close() asyncio.run(run()) else: @@ -241,7 +289,14 @@ def handler(_): assert sum(event.name == "exception" for event in span.events) == 1 -def test_cancel_async_close_finishes(tracing, base_url): +@pytest.mark.parametrize("close_method", [ + "close", + pytest.param("aclose", marks=pytest.mark.skipif( + not hasattr(openai.AsyncStream, "aclose"), + reason="This SDK version does not expose AsyncStream.aclose", + )), +]) +def test_cancel_async_close_finishes(tracing, base_url, close_method): async def run(): closing = asyncio.Event() @@ -261,7 +316,7 @@ def handler(_): async with client(handler, base_url, True) as sdk: stream = await sdk.chat.completions.create(**ARGS, stream=True) await stream.__anext__() - task = asyncio.create_task(stream.close()) + task = asyncio.create_task(getattr(stream, close_method)()) await asyncio.wait_for(closing.wait(), 2) task.cancel("synthetic close cancellation") with cancelled("synthetic close cancellation"): diff --git a/python/frameworks/openai/traceai_openai/_stream.py b/python/frameworks/openai/traceai_openai/_stream.py index efab5c65e..e8ee2cdee 100644 --- a/python/frameworks/openai/traceai_openai/_stream.py +++ b/python/frameworks/openai/traceai_openai/_stream.py @@ -184,6 +184,10 @@ async def close(self) -> Any: finally: _finish_on_exit(self, error) + async def aclose(self) -> Any: + # The SDK alias must use our close path, not bypass span finalization. + return await self.close() + def __aiter__(self): return self From 600311db8e34194cd0aa2e5bee7810fce0c3211d Mon Sep 17 00:00:00 2001 From: Nikhil Pareek Date: Mon, 2 Nov 2026 13:14:29 -0800 Subject: [PATCH 5/7] fix(openai): finish high-level SDK stream spans on response close Co-Authored-By: Claude Code --- .../workflows/baseten-python-contracts.yml | 2 + .../tests/test_high_level_stream_error.py | 39 +++++++++++ .../tests/test_high_level_stream_lifecycle.py | 67 +++++++++++++++++++ .../openai/traceai_openai/_stream.py | 67 ++++++++++++++++++- 4 files changed, 173 insertions(+), 2 deletions(-) create mode 100644 python/frameworks/openai/tests/test_high_level_stream_error.py create mode 100644 python/frameworks/openai/tests/test_high_level_stream_lifecycle.py diff --git a/.github/workflows/baseten-python-contracts.yml b/.github/workflows/baseten-python-contracts.yml index 15cc6dd54..2ef37d5de 100644 --- a/.github/workflows/baseten-python-contracts.yml +++ b/.github/workflows/baseten-python-contracts.yml @@ -55,6 +55,8 @@ jobs: python/frameworks/openai/tests/test_openai_compatible_spans.py \ python/frameworks/openai/tests/test_stream_lifecycle.py \ python/frameworks/openai/tests/test_stream_context_manager.py \ + python/frameworks/openai/tests/test_high_level_stream_lifecycle.py \ + python/frameworks/openai/tests/test_high_level_stream_error.py \ --noconftest -p no:cacheprovider -o addopts='' -q sudo unshare --net -- bash -c 'ip link set lo up; exec "$@"' bash \ env PATH="$PATH" PYTHONDONTWRITEBYTECODE=1 \ diff --git a/python/frameworks/openai/tests/test_high_level_stream_error.py b/python/frameworks/openai/tests/test_high_level_stream_error.py new file mode 100644 index 000000000..260bbfdfc --- /dev/null +++ b/python/frameworks/openai/tests/test_high_level_stream_error.py @@ -0,0 +1,39 @@ +"""The high-level helper must retain user exceptions on unfinished streams.""" +import asyncio + +import pytest +from opentelemetry.trace import StatusCode + +from .test_high_level_stream_lifecycle import _response_events, _stream_context +from .test_stream_lifecycle import base_url, client, streaming, tracing + + +@pytest.mark.parametrize("resource", ["chat", "responses"]) +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.parametrize("tracing,hide_errors", [(False, False), (True, True)], indirect=["tracing"]) +def test_high_level_context_exception_finishes_once(tracing, base_url, resource, use_async, hide_errors): + provider, exporter = tracing + handler = streaming if resource == "chat" else _response_events + if use_async: + async def run(): + async with client(handler, base_url, True) as sdk: + with pytest.raises(ValueError, match="synthetic consumer error"): + async with _stream_context(sdk, resource) as stream: + await stream.__anext__() + raise ValueError("synthetic consumer error") + assert stream._response.is_closed + asyncio.run(run()) + else: + with client(handler, base_url) as sdk: + with pytest.raises(ValueError, match="synthetic consumer error"): + with _stream_context(sdk, resource) as stream: + next(stream) + raise ValueError("synthetic consumer error") + assert stream._response.is_closed + assert provider.force_flush() + spans = exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].status.status_code == StatusCode.ERROR + assert spans[0].attributes["gen_ai.provider.name"] == "baseten" + if hide_errors: + assert "synthetic consumer error" not in spans[0].to_json() diff --git a/python/frameworks/openai/tests/test_high_level_stream_lifecycle.py b/python/frameworks/openai/tests/test_high_level_stream_lifecycle.py new file mode 100644 index 000000000..fbc047b1a --- /dev/null +++ b/python/frameworks/openai/tests/test_high_level_stream_lifecycle.py @@ -0,0 +1,67 @@ +"""High-level SDK stream helpers must finish spans when their context closes early.""" +import asyncio +import json + +import pytest +from opentelemetry.trace import StatusCode + +from .test_stream_lifecycle import ARGS, base_url, client, http, streaming, tracing + + +def _response_events(_): + response = { + "id": "resp_synthetic", "object": "response", "created_at": 1, + "model": "served-model", "status": "completed", "parallel_tool_calls": True, + "tool_choice": "auto", "tools": [], "output": [], + "usage": {"input_tokens": 5, "output_tokens": 0, "total_tokens": 5}, + } + events = [ + {"type": "response.created", "sequence_number": 0, + "response": {**response, "status": "in_progress"}}, + {"type": "response.completed", "sequence_number": 1, "response": response}, + ] + content = "".join( + f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in events + ).encode() + return http.Response(200, headers={"content-type": "text/event-stream"}, content=content) + + +def _stream_context(sdk, resource): + if resource == "responses": + return sdk.responses.stream(model="requested-model", input="synthetic prompt") + completions = sdk.chat.completions + if not hasattr(completions, "stream"): + completions = sdk.beta.chat.completions + return completions.stream(**ARGS) + + +@pytest.mark.parametrize("resource", ["chat", "responses"]) +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.parametrize("consumption", ["empty", "partial", "exhausted"]) +def test_high_level_context_finalizes_on_early_exit(tracing, base_url, resource, use_async, consumption): + provider, exporter = tracing + handler = streaming if resource == "chat" else _response_events + if use_async: + async def run(): + async with client(handler, base_url, True) as sdk: + async with _stream_context(sdk, resource) as stream: + if consumption == "partial": + await stream.__anext__() + elif consumption == "exhausted": + async for _ in stream: + pass + assert stream._response.is_closed + asyncio.run(run()) + else: + with client(handler, base_url) as sdk: + with _stream_context(sdk, resource) as stream: + if consumption == "partial": + next(stream) + elif consumption == "exhausted": + list(stream) + assert stream._response.is_closed + assert provider.force_flush() + spans = exporter.get_finished_spans() + assert len(spans) == 1, "SDK closed the HTTP response without ending its span" + assert spans[0].status.status_code == StatusCode.OK + assert spans[0].attributes["gen_ai.provider.name"] == "baseten" diff --git a/python/frameworks/openai/traceai_openai/_stream.py b/python/frameworks/openai/traceai_openai/_stream.py index e8ee2cdee..db5cfcfd2 100644 --- a/python/frameworks/openai/traceai_openai/_stream.py +++ b/python/frameworks/openai/traceai_openai/_stream.py @@ -1,4 +1,7 @@ +import asyncio import logging +import sys +import weakref from typing import Any, AsyncIterator, Generator, Iterator, Optional, Protocol, Tuple from openai.types.chat import ChatCompletionChunk @@ -46,6 +49,7 @@ class _Stream(ObjectProxy): # type: ignore "_self_iteration_count", "_self_is_finished", "_self_response_accumulator", + "_self_defer_close", ) def __init__( @@ -59,6 +63,29 @@ def __init__( self._self_iteration_count = 0 self._self_is_finished = with_span.is_finished self._self_response_accumulator = response_accumulator + self._self_defer_close = False + # The SDK's high-level .stream() helpers close the HTTP response + # directly, bypassing raw_stream.close() on early context exit. + response = stream.response + original_close = response.close + stream_ref = weakref.ref(self) + + def close_response() -> Any: + # A high-level context manager closes the response while handling + # the user's exception; preserve that failure on the span. + error: Optional[BaseException] = sys.exc_info()[1] + try: + return original_close() + except BaseException as close_error: + error = close_error + raise + finally: + response.close = original_close + traced = stream_ref() + if traced is not None and not traced._self_defer_close: + _finish_on_exit(traced, error) + + response.close = close_response def __enter__(self) -> "_Stream": # ObjectProxy otherwise returns the SDK stream, bypassing our iterator. @@ -66,22 +93,26 @@ def __enter__(self) -> "_Stream": return self def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> Any: + self._self_defer_close = True try: return self.__wrapped__.__exit__(exc_type, exc, tb) except BaseException as close_error: exc = close_error raise finally: + self._self_defer_close = False _finish_on_exit(self, exc) def close(self) -> Any: error: Optional[BaseException] = None + self._self_defer_close = True try: return self.__wrapped__.close() except BaseException as close_error: error = close_error raise finally: + self._self_defer_close = False _finish_on_exit(self, error) def __iter__(self) -> Iterator[Any]: @@ -89,7 +120,11 @@ def __iter__(self) -> Iterator[Any]: def __next__(self) -> Any: try: - chunk = self.__wrapped__.__next__() + self._self_defer_close = True + try: + chunk = self.__wrapped__.__next__() + finally: + self._self_defer_close = False except StopIteration: if not self._self_is_finished: status = trace_api.Status(status_code=trace_api.StatusCode.OK) @@ -147,6 +182,7 @@ class _AsyncStream(ObjectProxy): "_self_iteration_count", "_self_is_finished", "_self_response_accumulator", + "_self_defer_close", ) def __init__( @@ -160,28 +196,51 @@ def __init__( self._self_iteration_count = 0 self._self_is_finished = with_span.is_finished self._self_response_accumulator = response_accumulator + self._self_defer_close = None + response = stream.response + original_close = response.aclose + stream_ref = weakref.ref(self) + + async def close_response() -> Any: + error: Optional[BaseException] = sys.exc_info()[1] + try: + return await original_close() + except BaseException as close_error: + error = close_error + raise + finally: + response.aclose = original_close + traced = stream_ref() + if traced is not None and traced._self_defer_close is not asyncio.current_task(): + _finish_on_exit(traced, error) + + response.aclose = close_response async def __aenter__(self) -> "_AsyncStream": await self.__wrapped__.__aenter__() return self async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> Any: + self._self_defer_close = asyncio.current_task() try: return await self.__wrapped__.__aexit__(exc_type, exc, tb) except BaseException as close_error: exc = close_error raise finally: + self._self_defer_close = None _finish_on_exit(self, exc) async def close(self) -> Any: error: Optional[BaseException] = None + self._self_defer_close = asyncio.current_task() try: return await self.__wrapped__.close() except BaseException as close_error: error = close_error raise finally: + self._self_defer_close = None _finish_on_exit(self, error) async def aclose(self) -> Any: @@ -193,7 +252,11 @@ def __aiter__(self): async def __anext__(self): try: - chunk = await self.__wrapped__.__anext__() + self._self_defer_close = asyncio.current_task() + try: + chunk = await self.__wrapped__.__anext__() + finally: + self._self_defer_close = None except StopAsyncIteration: if not self._self_is_finished: status = trace_api.Status(status_code=trace_api.StatusCode.OK) From 596221ffdbbd35a3cb3ca71fb5291ff849898d00 Mon Sep 17 00:00:00 2001 From: Nikhil Pareek Date: Mon, 2 Nov 2026 13:32:33 -0800 Subject: [PATCH 6/7] fix(openai): finalize synchronous streams on interruption Co-Authored-By: Claude Code --- .../openai/tests/test_stream_lifecycle.py | 33 +++++++++++++++++++ .../openai/traceai_openai/_stream.py | 11 ++----- 2 files changed, 36 insertions(+), 8 deletions(-) diff --git a/python/frameworks/openai/tests/test_stream_lifecycle.py b/python/frameworks/openai/tests/test_stream_lifecycle.py index 480c0bfea..b1ed19ed5 100644 --- a/python/frameworks/openai/tests/test_stream_lifecycle.py +++ b/python/frameworks/openai/tests/test_stream_lifecycle.py @@ -218,6 +218,39 @@ def fail_close(): assert sum(event.name == "exception" for event in span.events) == 1 +@pytest.mark.parametrize("tracing,hide_errors", [(False, False), (True, True)], indirect=["tracing"]) +def test_sync_stream_interrupt_finishes_before_close(tracing, base_url, hide_errors): + class InterruptBody(http.SyncByteStream): + def __iter__(self): + yield CHUNK_BYTES + raise KeyboardInterrupt("synthetic stream interruption") + + def close(self): + pass + + def handler(_): + return http.Response( + 200, headers={"content-type": "text/event-stream"}, stream=InterruptBody(), + ) + + with client(handler, base_url) as sdk: + stream = sdk.chat.completions.create(**ARGS, stream=True) + assert next(stream).choices[0].delta.content == "partial-answer" + with pytest.raises(KeyboardInterrupt, match="synthetic stream interruption"): + next(stream) + span = only_span(tracing, StatusCode.ERROR) + stream.close() + stream.close() + assert stream.response.is_closed + + assert only_span(tracing, StatusCode.ERROR) is span + assert span.attributes["output.value"] == "partial-answer" + if hide_errors: + assert "synthetic stream interruption" not in span.to_json() + else: + assert sum(event.name == "exception" for event in span.events) == 1 + + @pytest.mark.parametrize("tracing,hide_errors", [(False, False), (True, True)], indirect=["tracing"]) def test_cancel_pending_request_finishes(tracing, base_url, hide_errors): async def run(): diff --git a/python/frameworks/openai/traceai_openai/_stream.py b/python/frameworks/openai/traceai_openai/_stream.py index db5cfcfd2..956d4d11f 100644 --- a/python/frameworks/openai/traceai_openai/_stream.py +++ b/python/frameworks/openai/traceai_openai/_stream.py @@ -130,14 +130,9 @@ def __next__(self) -> Any: status = trace_api.Status(status_code=trace_api.StatusCode.OK) self._finish_tracing(status=status) raise - except Exception as exception: - if not self._self_is_finished: - status = trace_api.Status( - status_code=trace_api.StatusCode.ERROR, - description=f"{type(exception).__name__}: {exception}", - ) - self._self_with_span.record_exception(exception) - self._finish_tracing(status=status) + except BaseException as exception: + # Interruptions also terminate the stream; preserve the error span. + _finish_on_exit(self, exception) raise else: self._process_chunk(chunk) From 8da1c2387edcbff9a46f516b67d132ccf005bac5 Mon Sep 17 00:00:00 2001 From: Nikhil Pareek Date: Mon, 2 Nov 2026 13:51:25 -0800 Subject: [PATCH 7/7] fix(openai): finalize synchronous requests on interruption Co-Authored-By: Claude Code --- .../openai/tests/test_request_interrupt.py | 83 +++++++++++++++++++ .../openai/traceai_openai/_request.py | 3 +- 2 files changed, 85 insertions(+), 1 deletion(-) create mode 100644 python/frameworks/openai/tests/test_request_interrupt.py diff --git a/python/frameworks/openai/tests/test_request_interrupt.py b/python/frameworks/openai/tests/test_request_interrupt.py new file mode 100644 index 000000000..a944ec05c --- /dev/null +++ b/python/frameworks/openai/tests/test_request_interrupt.py @@ -0,0 +1,83 @@ +"""Native HTTP request interruption must finish the span before propagating.""" + +import os + +import openai +import pytest +from fi_instrumentation import TraceConfig +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 +from traceai_openai import OpenAIInstrumentor + +if int(openai.__version__.split(".")[0]) >= 3: + import httpx2 as http +else: + import httpx as http + + +@pytest.mark.parametrize("base_url", [ + "https://model-synthetic.api.baseten.co/sync/v1", + "https://model-synthetic-us-west.api.baseten.co/sync/v1", +]) +@pytest.mark.parametrize("hide_errors", [False, True]) +@pytest.mark.parametrize("interrupt", [False, True], ids=["success-control", "interrupted"]) +def test_sync_request_interrupt_finishes_span(base_url, hide_errors, interrupt, monkeypatch): + for name in tuple(os.environ): + if name.startswith(("FI_", "OTEL_", "OPENAI_")): + monkeypatch.delenv(name) + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + instrumentor = OpenAIInstrumentor() + instrumentor.instrument( + tracer_provider=provider, config=TraceConfig(hide_error_details=hide_errors), + ) + error = KeyboardInterrupt("synthetic request interruption") + + def handle(_): + if interrupt: + raise error + return http.Response(200, json={ + "id": "synthetic", "object": "chat.completion", "created": 1, + "model": "served-model", "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop", + }], + }) + + try: + with openai.OpenAI( + api_key="synthetic", base_url=base_url, max_retries=0, + http_client=http.Client(transport=http.MockTransport(handle), trust_env=False), + ) as client: + if interrupt: + with pytest.raises(KeyboardInterrupt) as caught: + client.chat.completions.create( + model="requested-model", messages=[{"role": "user", "content": "hi"}], + ) + assert caught.value is error + else: + client.chat.completions.create( + model="requested-model", messages=[{"role": "user", "content": "hi"}], + ) + # Must finish before client close; force_flush cannot finish an active span. + assert provider.force_flush() + spans = exporter.get_finished_spans() + assert len(spans) == 1 + span = spans[0] + assert span.attributes["gen_ai.provider.name"] == "baseten" + assert span.status.status_code == ( + StatusCode.ERROR if interrupt else StatusCode.OK + ) + if interrupt: + if hide_errors: + assert "synthetic request interruption" not in span.to_json() + else: + assert sum(event.name == "exception" for event in span.events) == 1 + assert exporter.get_finished_spans() == spans + finally: + instrumentor.uninstrument() + provider.shutdown() diff --git a/python/frameworks/openai/traceai_openai/_request.py b/python/frameworks/openai/traceai_openai/_request.py index 19f8ff394..e9a60050d 100644 --- a/python/frameworks/openai/traceai_openai/_request.py +++ b/python/frameworks/openai/traceai_openai/_request.py @@ -340,7 +340,8 @@ def __call__( cast_to=cast_to, request_parameters=request_parameters, ) - except Exception as e: + except BaseException as e: + # end_on_exit=False: an interrupted request also needs finalization. status = trace_api.Status( status_code=trace_api.StatusCode.ERROR, description=f"{type(e).__name__}: {e}",