telemetry-dev-openai 0.2.1__tar.gz → 0.2.3__tar.gz

{telemetry_dev_openai-0.2.1 → telemetry_dev_openai-0.2.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: telemetry-dev-openai
-Version: 0.2.1
+Version: 0.2.3
 Summary: OpenAI integration for telemetry.dev Python SDK
 Keywords: telemetry,opentelemetry,openai,llm,genai,tracing
 Author: telemetry.dev
@@ -101,9 +101,9 @@ The integration maps native OpenAI request/response shapes directly into telemet
 
 Chat completion streams are traced. Requests are sent unchanged by default, so token usage is only captured when the caller sets `stream_options={"include_usage": True}` themselves. Pass `inject_stream_usage=True` to `wrap_openai` or `instrument_openai` to inject it automatically; the synthetic usage-only chunk is then hidden from the caller. Injection is opt-in because some providers reject `stream_options` — for example Azure OpenAI "on your data" (`data_sources`) returns 400 for it while plain `stream=True` works.
 
-Streamed chat output and finish reasons use separate budgets of up to 48 KiB and 1,000 items; the configured `max_attribute_length` does not shrink these retention bounds; the core applies it to the exported attribute after the mask runs, with its `...[truncated]` marker, without setting `telemetry.dev.capture.truncated`. When present, `gen_ai.output.messages` contains bounded partial output. If output and/or finish-reason capture is incomplete, including when a choice ends without `finish_reason`, `telemetry.dev.capture.truncated` is `true`. With a mask configured, incomplete output is omitted because the mask cannot inspect the complete value. When output capture is disabled or omitted, response IDs and models, finish reasons, usage, timing, and errors may still be recorded.
+Streamed chat output is kept up to 16 MiB, and all streams in flight share the `max_stream_bytes` bound of the core (64 MiB by default). These bounds need a `telemetry-dev` core with `max_stream_bytes`; with an older core, each stream keeps up to 64 KiB and 1,024 items. Each streamed finish reason is cut to its first 128 characters, without setting `telemetry.dev.capture.truncated`, and the finish reasons of a stream have their own budget, which fits at least 1,000 of them; the configured `max_attribute_length` does not shrink these retention bounds; the core applies it to the exported attribute after the mask runs, with its `...[truncated]` marker, without setting `telemetry.dev.capture.truncated`. A chat stream inspects at most 1,000 choices and 1,000 tool calls in each chunk; past these caps the rest is not kept and `telemetry.dev.capture.truncated` is set. When present, `gen_ai.output.messages` contains bounded partial output. If output and/or finish-reason capture is incomplete, including when a choice ends without `finish_reason`, `telemetry.dev.capture.truncated` is `true`. With a mask configured, incomplete output is omitted because the mask cannot inspect the complete value. When output capture is disabled or omitted, response IDs and models, finish reasons, usage, timing, and errors may still be recorded.
 
-Capture flags do not gate stop sequences or caller-supplied metadata or raw attributes; redact those separately when needed. Exception messages and stack traces are recorded only when both `capture_input` and `capture_output` are on, and pass through `mask`. Older `telemetry-dev` versions record them regardless of capture flags.
+Capture flags do not gate stop sequences or caller-supplied metadata or raw attributes; redact those separately when needed. Exception messages and stack traces are recorded only when both `capture_input` and `capture_output` are on, and pass through `mask`.
 
 ```py
 client = wrap_openai(OpenAI(), inject_stream_usage=True)
@@ -111,7 +111,7 @@ client = wrap_openai(OpenAI(), inject_stream_usage=True)
 
 Responses API streams are traced through `responses.create(stream=True)` and `responses.stream(response_id=...)`; terminal `response.completed`, `response.failed`, and `response.incomplete` events close the span.
 
-Responses and transcription streams follow the same rule. Responses snapshots are retained within 48 KiB and the core default of 1,024 items. A Responses stream that ends without a terminal snapshot, ends with an `error` event, fails to map an event, or has a terminal snapshot over the limit sets `telemetry.dev.capture.truncated`; it keeps the last snapshot that fit, unless a mask is configured, in which case output is omitted. Transcription text has its own retention limit of up to 64 KiB with no item limit. A complete non-streamed transcript is passed to the core whole. A transcription stream that ends before `transcript.text.done` (an error, early close, or end of stream) or whose text exceeds the limit is flagged the same way and keeps its bounded partial text only when no mask is configured.
+Responses and transcription streams follow the same rule. Responses snapshots and streamed transcription text are kept within the same stream bounds. A Responses stream that ends without a terminal snapshot, ends with an `error` event, fails to map an event, or has a terminal snapshot over the limit sets `telemetry.dev.capture.truncated`; it keeps the last snapshot that fit, unless a mask is configured, in which case output is omitted. A complete non-streamed transcript is passed to the core whole. A transcription stream that ends before `transcript.text.done` (an error, early close, or end of stream) or whose text exceeds the limit is flagged the same way and keeps its bounded partial text only when no mask is configured.
 
 ## Embeddings
 

{telemetry_dev_openai-0.2.1 → telemetry_dev_openai-0.2.3}/README.md

@@ -78,9 +78,9 @@ The integration maps native OpenAI request/response shapes directly into telemet
 
 Chat completion streams are traced. Requests are sent unchanged by default, so token usage is only captured when the caller sets `stream_options={"include_usage": True}` themselves. Pass `inject_stream_usage=True` to `wrap_openai` or `instrument_openai` to inject it automatically; the synthetic usage-only chunk is then hidden from the caller. Injection is opt-in because some providers reject `stream_options` — for example Azure OpenAI "on your data" (`data_sources`) returns 400 for it while plain `stream=True` works.
 
-Streamed chat output and finish reasons use separate budgets of up to 48 KiB and 1,000 items; the configured `max_attribute_length` does not shrink these retention bounds; the core applies it to the exported attribute after the mask runs, with its `...[truncated]` marker, without setting `telemetry.dev.capture.truncated`. When present, `gen_ai.output.messages` contains bounded partial output. If output and/or finish-reason capture is incomplete, including when a choice ends without `finish_reason`, `telemetry.dev.capture.truncated` is `true`. With a mask configured, incomplete output is omitted because the mask cannot inspect the complete value. When output capture is disabled or omitted, response IDs and models, finish reasons, usage, timing, and errors may still be recorded.
+Streamed chat output is kept up to 16 MiB, and all streams in flight share the `max_stream_bytes` bound of the core (64 MiB by default). These bounds need a `telemetry-dev` core with `max_stream_bytes`; with an older core, each stream keeps up to 64 KiB and 1,024 items. Each streamed finish reason is cut to its first 128 characters, without setting `telemetry.dev.capture.truncated`, and the finish reasons of a stream have their own budget, which fits at least 1,000 of them; the configured `max_attribute_length` does not shrink these retention bounds; the core applies it to the exported attribute after the mask runs, with its `...[truncated]` marker, without setting `telemetry.dev.capture.truncated`. A chat stream inspects at most 1,000 choices and 1,000 tool calls in each chunk; past these caps the rest is not kept and `telemetry.dev.capture.truncated` is set. When present, `gen_ai.output.messages` contains bounded partial output. If output and/or finish-reason capture is incomplete, including when a choice ends without `finish_reason`, `telemetry.dev.capture.truncated` is `true`. With a mask configured, incomplete output is omitted because the mask cannot inspect the complete value. When output capture is disabled or omitted, response IDs and models, finish reasons, usage, timing, and errors may still be recorded.
 
-Capture flags do not gate stop sequences or caller-supplied metadata or raw attributes; redact those separately when needed. Exception messages and stack traces are recorded only when both `capture_input` and `capture_output` are on, and pass through `mask`. Older `telemetry-dev` versions record them regardless of capture flags.
+Capture flags do not gate stop sequences or caller-supplied metadata or raw attributes; redact those separately when needed. Exception messages and stack traces are recorded only when both `capture_input` and `capture_output` are on, and pass through `mask`.
 
 ```py
 client = wrap_openai(OpenAI(), inject_stream_usage=True)
@@ -88,7 +88,7 @@ client = wrap_openai(OpenAI(), inject_stream_usage=True)
 
 Responses API streams are traced through `responses.create(stream=True)` and `responses.stream(response_id=...)`; terminal `response.completed`, `response.failed`, and `response.incomplete` events close the span.
 
-Responses and transcription streams follow the same rule. Responses snapshots are retained within 48 KiB and the core default of 1,024 items. A Responses stream that ends without a terminal snapshot, ends with an `error` event, fails to map an event, or has a terminal snapshot over the limit sets `telemetry.dev.capture.truncated`; it keeps the last snapshot that fit, unless a mask is configured, in which case output is omitted. Transcription text has its own retention limit of up to 64 KiB with no item limit. A complete non-streamed transcript is passed to the core whole. A transcription stream that ends before `transcript.text.done` (an error, early close, or end of stream) or whose text exceeds the limit is flagged the same way and keeps its bounded partial text only when no mask is configured.
+Responses and transcription streams follow the same rule. Responses snapshots and streamed transcription text are kept within the same stream bounds. A Responses stream that ends without a terminal snapshot, ends with an `error` event, fails to map an event, or has a terminal snapshot over the limit sets `telemetry.dev.capture.truncated`; it keeps the last snapshot that fit, unless a mask is configured, in which case output is omitted. A complete non-streamed transcript is passed to the core whole. A transcription stream that ends before `transcript.text.done` (an error, early close, or end of stream) or whose text exceeds the limit is flagged the same way and keeps its bounded partial text only when no mask is configured.
 
 ## Embeddings
 

{telemetry_dev_openai-0.2.1 → telemetry_dev_openai-0.2.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "telemetry-dev-openai"
-version = "0.2.1"
+version = "0.2.3"
 description = "OpenAI integration for telemetry.dev Python SDK"
 readme = "README.md"
 license = "MIT"

{telemetry_dev_openai-0.2.1 → telemetry_dev_openai-0.2.3}/pyproject.toml.orig

@@ -1,6 +1,6 @@
 [project]
 name = "telemetry-dev-openai"
-version = "0.2.1"
+version = "0.2.3"
 description = "OpenAI integration for telemetry.dev Python SDK"
 readme = "README.md"
 license = "MIT"

{telemetry_dev_openai-0.2.1 → telemetry_dev_openai-0.2.3}/src/telemetry_dev_openai/__init__.py

@@ -1,10 +1,13 @@
 from __future__ import annotations
 
+import os
+import re
 import threading
 import time
 import types
+import weakref
 from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Mapping, Sequence
-from functools import wraps
+from functools import partial, wraps
 from io import StringIO
 from typing import Any, Literal, TypeVar, cast
 from urllib.parse import urlsplit
@@ -20,7 +23,7 @@ from openai.resources.embeddings import AsyncEmbeddings, Embeddings
 from openai.resources.images import AsyncImages, Images
 from openai.resources.responses.responses import AsyncResponses, Responses
 
-__version__ = "0.2.1"
+__version__ = "0.2.3"
 
 ProviderResolver = Callable[[object | None], str]
 RequestMapper = Callable[[Mapping[str, Any]], tuple[str, dict[str, Any]]]
@@ -31,11 +34,30 @@ _ORIGINAL_ATTR = "_telemetry_dev_openai_original"
 _ORIGINALS: list[tuple[type[Any], str, Any, Any]] = []
 _installed = False
 _install_lock = threading.Lock()
+
+
+def _reset_install_lock_in_child() -> None:
+    """A thread of the parent can hold the lock at a fork, and only this thread exists here."""
+    global _install_lock
+    _install_lock = threading.Lock()
+
+
+# A hook is enough: a fork that runs no hooks, as uWSGI's fork of its workers by default,
+# starts from its master, which serves no requests. os.register_at_fork cannot remove a hook,
+# and importlib.reload runs this module again with the same globals, so the hook of the first
+# run stays. Windows has no fork.
+if hasattr(os, "register_at_fork") and "_install_fork_hook" not in globals():
+    _install_fork_hook = True
+    os.register_at_fork(after_in_child=_reset_install_lock_in_child)
+
 _T = TypeVar("_T")
 _OMIT = object()
-_CHAT_STREAM_CAPTURE_MAX_BYTES = 48 * 1024
-_TRANSCRIPT_CAPTURE_MAX_BYTES = 64 * 1024
+# More choices, tool calls, or finish reasons than this are not kept.
 _CHAT_STREAM_CAPTURE_MAX_ITEMS = 1000
+# Each streamed finish reason is cut to this many characters, as in the other integrations.
+_FINISH_REASON_MAX_LENGTH = 128
+_TEXT_CHUNK_CHARS = 4096
+_SURROGATES = re.compile("[\ud800-\udfff]")
 
 
 class _CaptureLimit(Exception):
@@ -65,6 +87,12 @@ def _structure_cost(shape: object) -> tuple[int, int]:
 
 
 _FINISH_REASON_BYTES, _FINISH_REASON_ITEMS = _structure_cost({"finish_reason": "", "index": 0})
+# Finish reasons stay outside the stream budget so a terminal one is still exported. Their own
+# budget fits _CHAT_STREAM_CAPTURE_MAX_ITEMS of the largest records a cut reason can take.
+_FINISH_REASON_BUDGET_BYTES = _CHAT_STREAM_CAPTURE_MAX_ITEMS * (
+    _FINISH_REASON_BYTES
+    + telemetry_dev.CaptureBudget.json_string_bytes("\x00" * _FINISH_REASON_MAX_LENGTH)
+)
 _TOOL_CALLS_FIELD_BYTES = _field_bytes("tool_calls")
 _CONTENT_FIELD_BYTES = _field_bytes("content")
 _FUNCTION_CALL_FIELD_BYTES = _field_bytes("function_call")
@@ -147,6 +175,7 @@ def _bounded_responses_capture(
     *,
     capture_enabled: bool | None = None,
     max_bytes: int | None = None,
+    budget: telemetry_dev.CaptureBudget | None = None,
     report_error: Callable[[BaseException], None] | None = None,
 ) -> tuple[Any | None, bool]:
     client = telemetry_dev.get_client()
@@ -156,11 +185,12 @@ def _bounded_responses_capture(
         capture_enabled = client.capture_input if capture == "input" else client.capture_output
     if not capture_enabled:
         return None, False
-    budget = (
-        telemetry_dev.CaptureBudget(max_bytes=max_bytes)
-        if max_bytes is not None
-        else telemetry_dev.CaptureBudget()
-    )
+    if budget is None:
+        budget = (
+            telemetry_dev.CaptureBudget(max_bytes=max_bytes)
+            if max_bytes is not None
+            else telemetry_dev.CaptureBudget()
+        )
     ancestors: set[int] = set()
 
     def reserve(byte_count: int) -> None:
@@ -169,7 +199,8 @@ def _bounded_responses_capture(
             or budget.items_used >= budget.max_items
         ):
             raise _CaptureLimit
-        budget.bytes_used += byte_count
+        if not _reserve_bytes(budget, budget.bytes_used + byte_count):
+            raise _CaptureLimit
         budget.items_used += 1
 
     def convert(item: Any, parent_type: str | None, depth: int) -> Any:
@@ -313,45 +344,103 @@ def _sent_params(params: Mapping[str, Any]) -> dict[str, Any]:
     }
 
 
-def _transcript_capture_budget() -> telemetry_dev.CaptureBudget:
-    return telemetry_dev.CaptureBudget(max_bytes=_TRANSCRIPT_CAPTURE_MAX_BYTES)
+def _stream_budget() -> telemetry_dev.CaptureBudget:
+    # Older cores have no for_stream and bound each budget at 64 KiB on its own.
+    for_stream = getattr(telemetry_dev.CaptureBudget, "for_stream", None)
+    budget = for_stream() if callable(for_stream) else telemetry_dev.CaptureBudget()
+    return cast(telemetry_dev.CaptureBudget, budget)
+
+
+def _release_budget(budget: telemetry_dev.CaptureBudget) -> None:
+    release = getattr(budget, "release", None)
+    if callable(release):
+        release()
+
+
+def _reserve_bytes(budget: telemetry_dev.CaptureBudget, total: int) -> bool:
+    # Older cores have no reserve_bytes; there the caller's check against max_bytes stands.
+    reserve = getattr(budget, "reserve_bytes", None)
+    if callable(reserve):
+        return bool(reserve(total))
+    budget.bytes_used = total
+    return True
+
+
+def _ending_stream(end: Callable[..., None], release: Callable[[], None]) -> Callable[..., None]:
+    def end_stream(**fields: Any) -> None:
+        try:
+            end(**fields)
+        finally:
+            release()
+
+    return end_stream
+
+
+def _character_bytes(character: str) -> int:
+    code = ord(character)
+    if code <= 0x7F:
+        return 1
+    if code <= 0x7FF:
+        return 2
+    if 0xD800 <= code <= 0xDFFF:
+        return 6
+    if code <= 0xFFFF:
+        return 3
+    return 4
+
+
+def _text_bytes(value: str) -> int:
+    """The bytes _character_bytes charges for ``value``: its UTF-8 width, and 6 per surrogate."""
+    return len(value.encode("utf-8", "surrogatepass")) + 3 * len(_SURROGATES.findall(value))
 
 
 def _capture_text(value: str, budget: telemetry_dev.CaptureBudget) -> str:
     if not value or budget.truncated:
         return ""
-    if budget.remaining_bytes <= 0:
+    # Read once: the remaining bytes of a stream budget take a lock.
+    remaining = budget.remaining_bytes
+    if remaining <= 0:
         budget.truncated = True
         return ""
     byte_count = 0
-    retained_characters = 0
-    for index, character in enumerate(value):
-        code = ord(character)
-        if code <= 0x7F:
-            character_bytes = 1
-        elif code <= 0x7FF:
-            character_bytes = 2
-        elif 0xD800 <= code <= 0xDFFF:
-            character_bytes = 6
-        elif code <= 0xFFFF:
-            character_bytes = 3
-        else:
-            character_bytes = 4
-        if byte_count + character_bytes > budget.remaining_bytes:
-            budget.truncated = True
-            break
-        byte_count += character_bytes
-        retained_characters = index + 1
-    budget.bytes_used += byte_count
+    retained_characters = len(value)
+    # Whole chunks are measured at once; only the chunk that does not fit is walked by character.
+    # Every character costs at least one byte, so the scan stops within remaining + one chunk.
+    for start in range(0, len(value), _TEXT_CHUNK_CHARS):
+        chunk = value[start : start + _TEXT_CHUNK_CHARS]
+        chunk_bytes = _text_bytes(chunk)
+        if byte_count + chunk_bytes <= remaining:
+            byte_count += chunk_bytes
+            continue
+        retained_characters = start
+        for character in chunk:
+            character_bytes = _character_bytes(character)
+            if byte_count + character_bytes > remaining:
+                break
+            byte_count += character_bytes
+            retained_characters += 1
+        break
+    if not _reserve_bytes(budget, budget.bytes_used + byte_count):
+        budget.truncated = True
+        return ""
     if retained_characters < len(value):
         budget.truncated = True
         return value[:retained_characters]
     return value
 
 
+def _prompt_attributes(messages: object) -> dict[str, Any]:
+    # Older cores have no prompt_attributes_for and record a prompt only inside with_prompt.
+    read = getattr(telemetry_dev, "prompt_attributes_for", None)
+    return cast("dict[str, Any]", read(messages)) if callable(read) else {}
+
+
 def _chat_request(params: Mapping[str, Any]) -> tuple[str, dict[str, Any]]:
     model = _string(params.get("model"))
+    attributes = _prompt_attributes(params.get("messages"))
     captured, truncated = _bounded_responses_capture(params.get("messages"), "input")
+    if truncated:
+        attributes["telemetry.dev.capture.truncated"] = True
     return (
         f"chat {model or 'unknown'}",
         {
@@ -366,7 +455,7 @@ def _chat_request(params: Mapping[str, Any]) -> tuple[str, dict[str, Any]]:
             "seed": _number(params.get("seed")),
             "frequency_penalty": _number(params.get("frequency_penalty")),
             "presence_penalty": _number(params.get("presence_penalty")),
-            "attributes": {"telemetry.dev.capture.truncated": True} if truncated else None,
+            "attributes": attributes or None,
         },
     )
 
@@ -432,7 +521,10 @@ def _chat_response(response: Any) -> dict[str, Any]:
 
 def _responses_request(params: Mapping[str, Any]) -> tuple[str, dict[str, Any]]:
     model = _string(params.get("model"))
+    attributes = _prompt_attributes(params.get("input"))
     captured, truncated = _bounded_responses_capture(params.get("input"), "input")
+    if truncated:
+        attributes["telemetry.dev.capture.truncated"] = True
     return (
         f"chat {model or 'unknown'}",
         {
@@ -443,7 +535,7 @@ def _responses_request(params: Mapping[str, Any]) -> tuple[str, dict[str, Any]]:
             "temperature": _number(params.get("temperature")),
             "top_p": _number(params.get("top_p")),
             "max_tokens": _number(params.get("max_output_tokens")),
-            "attributes": {"telemetry.dev.capture.truncated": True} if truncated else None,
+            "attributes": attributes or None,
         },
     )
 
@@ -680,12 +772,92 @@ def _end_once(handle: telemetry_dev.SpanHandle) -> Callable[..., None]:
     return end
 
 
+# On Python 3.10 and 3.11, StringIO keeps up to 100,000 written strings in a list before it
+# joins them, about 60 bytes each that the stream bound does not charge: 15 bytes a character
+# for 4-character deltas. Deltas are joined 256 at a time, so that those waiting hold under
+# about 22 KB, onto a string that is written once it has 4,096 characters, so that what is
+# written costs about 1.5% more than its text. Python 3.12 dropped the list.
+_DELTA_BATCH = 256
+_WRITE_CHARS = 4096
+
+
+class _DeltaBuffer:
+    # A write takes several steps, unlike one StringIO call, so without the lock a read on
+    # another thread, such as one that closes the stream, could copy a batch twice or drop it.
+    # A fork copies the lock held if another thread was writing, and that thread does not run
+    # in the child to release it. Every fork, also one that runs no at-fork hooks, changes the
+    # process ID, so a child takes a new lock. The check is inline because it runs per delta.
+    # A signal handler runs on the thread it interrupts, so one that reads the stream, such as
+    # one that closes it, can read inside a write. So the lock is reentrant, and each step takes
+    # the deltas and the tail out of the buffer before it joins or writes them, so that a read
+    # inside the step does not write them as well; that read returns a prefix, which can leave
+    # out the deltas being joined. A handler that writes inside a step is not supported. An
+    # exception inside a step, such as KeyboardInterrupt from a handler, puts back the text the
+    # step took unless the StringIO has it, so that a later read still has every delta.
+    __slots__ = ("_buffer", "_length", "_lock", "_pending", "_pid", "_tail")
+
+    def __init__(self) -> None:
+        self._buffer = StringIO()
+        self._length = 0
+        self._lock = threading.RLock()
+        self._pid = os.getpid()
+        self._pending: list[str] = []
+        self._tail = ""
+
+    def write(self, text: str) -> None:
+        pid = os.getpid()
+        if self._pid != pid:
+            self._lock = threading.RLock()
+            self._pid = pid
+        with self._lock:
+            # Counted first: a signal handler runs only once a call returns, so none runs
+            # between the count and the append, and one that raises after the append, or in
+            # a join, finds the delta counted.
+            self._length += len(text)
+            self._pending.append(text)
+            if len(self._pending) >= _DELTA_BATCH:
+                self._join(_WRITE_CHARS)
+
+    def tell(self) -> int:
+        # Read on every content delta, so it counts the pending deltas instead of writing them.
+        return self._length
+
+    def getvalue(self) -> str:
+        pid = os.getpid()
+        if self._pid != pid:
+            self._lock = threading.RLock()
+            self._pid = pid
+        with self._lock:
+            self._join(1)
+            return self._buffer.getvalue()
+
+    def _join(self, write_at: int) -> None:
+        if not self._pending and len(self._tail) < write_at:
+            return
+        written = self._buffer.tell()
+        pending: list[str] = []
+        tail = ""
+        try:
+            pending, self._pending = self._pending, []
+            tail, self._tail = self._tail, ""
+            joined = "".join([tail, *pending])
+            if len(joined) >= write_at:
+                self._buffer.write(joined)
+            else:
+                self._tail = joined
+        except BaseException:
+            if self._buffer.tell() == written:
+                self._pending[:0] = pending
+                self._tail = tail + self._tail
+            raise
+
+
 class _ChatChoice:
     def __init__(self) -> None:
         self.role: str | None = None
         self.role_resolved = True
-        self.content_fragments = StringIO()
-        self.refusal_fragments = StringIO()
+        self.content_fragments = _DeltaBuffer()
+        self.refusal_fragments = _DeltaBuffer()
         self.function_call: dict[str, Any] | None = None
         self.unresolved_function_scalars: set[str] = set()
         self.tool_calls: dict[int, dict[str, Any]] = {}
@@ -725,10 +897,7 @@ class _ChatToolCallDelta:
 
 
 def _chat_capture_budget(*, reserve_output_list: bool = False) -> telemetry_dev.CaptureBudget:
-    budget = telemetry_dev.CaptureBudget(
-        _CHAT_STREAM_CAPTURE_MAX_BYTES,
-        _CHAT_STREAM_CAPTURE_MAX_ITEMS,
-    )
+    budget = _stream_budget()
     if reserve_output_list:
         budget.accept([])
     return budget
@@ -756,7 +925,11 @@ def _replace_finish_reason(
         reservations.pop(index, None)
         return False
     reservation = (_FINISH_REASON_BYTES + value_bytes, _FINISH_REASON_ITEMS)
-    budget.bytes_used = base_bytes + reservation[0]
+    if not _reserve_bytes(budget, base_bytes + reservation[0]):
+        budget.bytes_used = base_bytes
+        budget.items_used = base_items
+        reservations.pop(index, None)
+        return False
     budget.items_used = base_items + reservation[1]
     reservations[index] = reservation
     return True
@@ -777,7 +950,10 @@ def _reserve_chat_budget(
         if not recoverable:
             budget.truncated = True
         return False
-    budget.bytes_used += byte_count
+    if not _reserve_bytes(budget, budget.bytes_used + byte_count):
+        if not recoverable:
+            budget.truncated = True
+        return False
     budget.items_used += item_count
     return True
 
@@ -825,8 +1001,7 @@ def _replace_chat_scalar(
     value_bytes = telemetry_dev.CaptureBudget.json_string_bytes(value, available_bytes)
     if value_bytes > available_bytes:
         return False
-    budget.bytes_used = base_bytes + _BASE_ITEM_BYTES + value_bytes
-    return True
+    return _reserve_bytes(budget, base_bytes + _BASE_ITEM_BYTES + value_bytes)
 
 
 def _release_chat_field(field_name: str, value: str, budget: telemetry_dev.CaptureBudget) -> None:
@@ -899,8 +1074,8 @@ def _capture_tool_call_payload(
             unresolved.add(unresolved_name)
     if not budget.truncated and value is not None:
         fragments = payload.get(value_key)
-        had_value = isinstance(fragments, StringIO)
-        value_fragments = fragments if had_value else StringIO()
+        had_value = isinstance(fragments, _DeltaBuffer)
+        value_fragments = fragments if had_value else _DeltaBuffer()
         if value or not had_value:
             if _capture_chat_string(
                 value,
@@ -1053,8 +1228,8 @@ def _capture_function_call_delta(
             state.unresolved_function_scalars.add("name")
     if not budget.truncated and arguments is not None:
         fragments = current.get("arguments")
-        had_arguments = isinstance(fragments, StringIO)
-        argument_fragments = fragments if had_arguments else StringIO()
+        had_arguments = isinstance(fragments, _DeltaBuffer)
+        argument_fragments = fragments if had_arguments else _DeltaBuffer()
         if arguments or not had_arguments:
             if _capture_chat_string(
                 arguments,
@@ -1127,6 +1302,7 @@ def _record_chat_chunk(
             )
         )
         if finish_reason is not None:
+            finish_reason = finish_reason[:_FINISH_REASON_MAX_LENGTH]
             if unterminated_choices is not None:
                 unterminated_choices.discard(choice_index)
             if _replace_finish_reason(
@@ -1291,7 +1467,7 @@ def _chat_message(state: _ChatChoice) -> dict[str, Any]:
     if state.function_call is not None:
         captured_function_call = dict(state.function_call)
         arguments = captured_function_call.get("arguments")
-        if isinstance(arguments, StringIO):
+        if isinstance(arguments, _DeltaBuffer):
             captured_function_call["arguments"] = arguments.getvalue()
         if "name" in state.unresolved_function_scalars:
             captured_function_call.pop("name", None)
@@ -1310,7 +1486,7 @@ def _chat_message(state: _ChatChoice) -> dict[str, Any]:
                 if isinstance(raw_payload, Mapping):
                     captured_payload = dict(cast(Mapping[str, Any], raw_payload))
                     fragments = captured_payload.get(value_key)
-                    if isinstance(fragments, StringIO):
+                    if isinstance(fragments, _DeltaBuffer):
                         captured_payload[value_key] = fragments.getvalue()
                 else:
                     captured_payload = {}
@@ -1431,6 +1607,15 @@ def _hook_response_close(inner: Any, finish: Callable[[], None]) -> None:
     response = getattr(inner, "response", None)
     if response is None:
         return
+    # ``finish`` is a method of the stream wrapper, which holds ``response``: held weakly, the hook
+    # adds no reference cycle, so a dropped wrapper and its stream budget are freed at once.
+    finish_method = weakref.WeakMethod(finish)
+
+    def finish_if_alive() -> None:
+        method = finish_method()
+        if method is not None:
+            method()
+
     close = getattr(response, "close", None)
     if callable(close):
 
@@ -1438,7 +1623,7 @@ def _hook_response_close(inner: Any, finish: Callable[[], None]) -> None:
             try:
                 return close(*args, **kwargs)
             finally:
-                finish()
+                finish_if_alive()
 
         response.close = _close_hook
     aclose = getattr(response, "aclose", None)
@@ -1449,7 +1634,7 @@ def _hook_response_close(inner: Any, finish: Callable[[], None]) -> None:
             try:
                 return await aclose_fn(*args, **kwargs)
             finally:
-                finish()
+                finish_if_alive()
 
         response.aclose = _aclose_hook
 
@@ -1504,7 +1689,6 @@ class _InstrumentedStream:
         started_at: float,
     ) -> None:
         self._inner = inner
-        self._end = _end_once(handle)
         self._handle = handle
         self._injected_usage = injected_usage
         self._started_at = started_at
@@ -1524,7 +1708,9 @@ class _InstrumentedStream:
         self._consume: Iterator[Any] | None = None
         self._in_next = False
         self._budget = _chat_capture_budget(reserve_output_list=self._capture_output)
-        self._finish_reason_budget = _chat_capture_budget()
+        # Bound to the budget, not self, so the callback adds no reference cycle.
+        self._end = _ending_stream(_end_once(handle), partial(_release_budget, self._budget))
+        self._finish_reason_budget = telemetry_dev.CaptureBudget(_FINISH_REASON_BUDGET_BYTES)
         _hook_response_close(inner, self._on_response_close)
 
     def _iterate(self) -> Iterator[Any]:
@@ -1667,7 +1853,6 @@ class _InstrumentedAsyncStream:
         started_at: float,
     ) -> None:
         self._inner = inner
-        self._end = _end_once(handle)
         self._handle = handle
         self._injected_usage = injected_usage
         self._started_at = started_at
@@ -1687,7 +1872,9 @@ class _InstrumentedAsyncStream:
         self._consume: AsyncIterator[Any] | None = None
         self._in_next = False
         self._budget = _chat_capture_budget(reserve_output_list=self._capture_output)
-        self._finish_reason_budget = _chat_capture_budget()
+        # Bound to the budget, not self, so the callback adds no reference cycle.
+        self._end = _ending_stream(_end_once(handle), partial(_release_budget, self._budget))
+        self._finish_reason_budget = telemetry_dev.CaptureBudget(_FINISH_REASON_BUDGET_BYTES)
         _hook_response_close(inner, self._on_response_close)
 
     async def _aiterate(self) -> AsyncIterator[Any]:
@@ -1839,7 +2026,12 @@ class _InstrumentedResponsesStream:
     def __init__(self, inner: Any, handle: telemetry_dev.SpanHandle, started_at: float) -> None:
         self._inner = inner
         self._handle = handle
-        self._end = _end_once(handle)
+        self._output_budget = _stream_budget()
+        # Bound to the budget, not self, so the callback adds no reference cycle.
+        self._end_stream = _ending_stream(
+            _end_once(handle), partial(_release_budget, self._output_budget)
+        )
+        self._ended = False
         (
             self._capture_output,
             self._mask_output_when_incomplete,
@@ -1927,7 +2119,15 @@ class _InstrumentedResponsesStream:
     def __getattr__(self, name: str) -> Any:
         return getattr(self._inner, name)
 
+    def _end(self, **fields: Any) -> None:
+        # The span has ended and the stream budget no longer counts bytes, so later events are
+        # not kept.
+        self._ended = True
+        self._end_stream(**fields)
+
     def _record(self, event: Any, received_at: float) -> None:
+        if self._ended:
+            return
         if _response_event_has_output(event):
             self._handle.record_output_chunk(received_at * 1000)
         if not self._saw_first:
@@ -1945,13 +2145,24 @@ class _InstrumentedResponsesStream:
             output = None
             truncated = False
             if raw_output is not None:
+                # The snapshot replaces the output kept so far.
+                measured = telemetry_dev.CaptureBudget(self._output_budget.max_bytes)
                 output, truncated = _bounded_responses_capture(
                     raw_output,
                     capture_enabled=self._capture_output,
-                    max_bytes=_CHAT_STREAM_CAPTURE_MAX_BYTES,
+                    budget=measured,
                     report_error=self._report_error,
                 )
-                if not truncated and output:
+                # A snapshot that fits replaces the output kept so far. Until one fits, the
+                # stream keeps the cut one, so it holds that one's bytes.
+                keeps = output and (
+                    not truncated
+                    or (self._retained_output is None and not self._mask_output_when_incomplete)
+                )
+                if keeps and not _reserve_bytes(self._output_budget, measured.bytes_used):
+                    # Another stream took the room after this one measured.
+                    output, truncated = None, True
+                elif keeps and not truncated:
                     self._retained_output = output
             self._partial = fields
             if self._retained_output is not None and not (
@@ -1985,7 +2196,12 @@ class _InstrumentedAsyncResponsesStream:
     def __init__(self, inner: Any, handle: telemetry_dev.SpanHandle, started_at: float) -> None:
         self._inner = inner
         self._handle = handle
-        self._end = _end_once(handle)
+        self._output_budget = _stream_budget()
+        # Bound to the budget, not self, so the callback adds no reference cycle.
+        self._end_stream = _ending_stream(
+            _end_once(handle), partial(_release_budget, self._output_budget)
+        )
+        self._ended = False
         (
             self._capture_output,
             self._mask_output_when_incomplete,
@@ -2078,7 +2294,15 @@ class _InstrumentedAsyncResponsesStream:
     def __getattr__(self, name: str) -> Any:
         return getattr(self._inner, name)
 
+    def _end(self, **fields: Any) -> None:
+        # The span has ended and the stream budget no longer counts bytes, so later events are
+        # not kept.
+        self._ended = True
+        self._end_stream(**fields)
+
     def _record(self, event: Any, received_at: float) -> None:
+        if self._ended:
+            return
         if _response_event_has_output(event):
             self._handle.record_output_chunk(received_at * 1000)
         if not self._saw_first:
@@ -2096,13 +2320,24 @@ class _InstrumentedAsyncResponsesStream:
             output = None
             truncated = False
             if raw_output is not None:
+                # The snapshot replaces the output kept so far.
+                measured = telemetry_dev.CaptureBudget(self._output_budget.max_bytes)
                 output, truncated = _bounded_responses_capture(
                     raw_output,
                     capture_enabled=self._capture_output,
-                    max_bytes=_CHAT_STREAM_CAPTURE_MAX_BYTES,
+                    budget=measured,
                     report_error=self._report_error,
                 )
-                if not truncated and output:
+                # A snapshot that fits replaces the output kept so far. Until one fits, the
+                # stream keeps the cut one, so it holds that one's bytes.
+                keeps = output and (
+                    not truncated
+                    or (self._retained_output is None and not self._mask_output_when_incomplete)
+                )
+                if keeps and not _reserve_bytes(self._output_budget, measured.bytes_used):
+                    # Another stream took the room after this one measured.
+                    output, truncated = None, True
+                elif keeps and not truncated:
                     self._retained_output = output
             self._partial = fields
             if self._retained_output is not None and not (
@@ -2136,6 +2371,7 @@ def _media_stream_event_fields(
     event: Any,
     response_mapper: ResponseMapper,
     capture_output: bool,
+    budget: telemetry_dev.CaptureBudget,
 ) -> dict[str, Any]:
     if response_mapper is not _text_media_response:
         return response_mapper(event)
@@ -2143,7 +2379,9 @@ def _media_stream_event_fields(
     if capture_output and _string(_field(event, "type")) == "transcript.text.done":
         text = _string(_field(event, "text"))
         if text is not None:
-            budget = _transcript_capture_budget()
+            # The terminal text replaces the deltas, within the room the stream has.
+            budget.bytes_used = 0
+            budget.truncated = False
             fields["output"] = _capture_text(text, budget)
             if budget.truncated:
                 fields["attributes"] = {"telemetry.dev.capture.truncated": True}
@@ -2153,7 +2391,7 @@ def _media_stream_event_fields(
 def _media_stream_end_fields(
     fields: Mapping[str, Any],
     response_mapper: ResponseMapper,
-    text_parts: Sequence[str],
+    text_fragments: _DeltaBuffer,
     budget: telemetry_dev.CaptureBudget,
     mask_output_when_incomplete: bool,
     stream_incomplete: bool,
@@ -2168,8 +2406,9 @@ def _media_stream_end_fields(
         )
         if "output" not in result:
             capture_incomplete = capture_incomplete or budget.truncated
-            if text_parts and not (mask_output_when_incomplete and capture_incomplete):
-                result["output"] = "".join(text_parts)
+            text = text_fragments.getvalue()
+            if text and not (mask_output_when_incomplete and capture_incomplete):
+                result["output"] = text
         else:
             capture_incomplete = capture_incomplete or terminal_truncated
         if mask_output_when_incomplete and capture_incomplete:
@@ -2228,15 +2467,17 @@ class _InstrumentedMediaStream:
         self, inner: Any, handle: telemetry_dev.SpanHandle, response_mapper: ResponseMapper
     ) -> None:
         self._inner = inner
-        self._end = _end_once(handle)
         self._fields: dict[str, Any] = {}
         self._response_mapper = response_mapper
         (
             self._capture_output,
             self._mask_output_when_incomplete,
         ) = _chat_span_capture_policy(handle)
-        self._budget = _transcript_capture_budget()
-        self._text_parts: list[str] = []
+        self._budget = _stream_budget()
+        # Bound to the budget, not self, so the callback adds no reference cycle.
+        self._end_stream = _ending_stream(_end_once(handle), partial(_release_budget, self._budget))
+        self._ended = False
+        self._text_fragments = _DeltaBuffer()
         self._saw_delta = False
         self._saw_terminal = False
         self._report_error = _chat_span_reporter(handle)
@@ -2292,7 +2533,15 @@ class _InstrumentedMediaStream:
             if failure is not None:
                 self._end(**self._finish_fields(), error=failure)
 
+    def _end(self, **fields: Any) -> None:
+        # The span has ended and the stream budget no longer counts bytes, so later events are
+        # not kept.
+        self._ended = True
+        self._end_stream(**fields)
+
     def _record(self, event: Any) -> None:
+        if self._ended:
+            return
         event_type = _string(_field(event, "type")) or ""
         if event_type in _MEDIA_TERMINAL_EVENTS:
             self._saw_terminal = True
@@ -2300,6 +2549,7 @@ class _InstrumentedMediaStream:
             event,
             self._response_mapper,
             self._capture_output,
+            self._budget,
         )
         if "output" in mapped:
             self._fields = _clean_fields(mapped)
@@ -2315,7 +2565,7 @@ class _InstrumentedMediaStream:
             if text:
                 captured = _capture_text(text, self._budget)
                 if captured:
-                    self._text_parts.append(captured)
+                    self._text_fragments.write(captured)
         if event_type in _MEDIA_TERMINAL_EVENTS:
             error = _media_stream_error(event, event_type)
             self._end(**self._finish_fields(), **({"error": error} if error is not None else {}))
@@ -2324,7 +2574,7 @@ class _InstrumentedMediaStream:
         return _media_stream_end_fields(
             self._fields,
             self._response_mapper,
-            self._text_parts,
+            self._text_fragments,
             self._budget,
             self._mask_output_when_incomplete,
             not self._saw_terminal or self._mapping_failed,
@@ -2345,15 +2595,17 @@ class _InstrumentedAsyncMediaStream:
         self, inner: Any, handle: telemetry_dev.SpanHandle, response_mapper: ResponseMapper
     ) -> None:
         self._inner = inner
-        self._end = _end_once(handle)
         self._fields: dict[str, Any] = {}
         self._response_mapper = response_mapper
         (
             self._capture_output,
             self._mask_output_when_incomplete,
         ) = _chat_span_capture_policy(handle)
-        self._budget = _transcript_capture_budget()
-        self._text_parts: list[str] = []
+        self._budget = _stream_budget()
+        # Bound to the budget, not self, so the callback adds no reference cycle.
+        self._end_stream = _ending_stream(_end_once(handle), partial(_release_budget, self._budget))
+        self._ended = False
+        self._text_fragments = _DeltaBuffer()
         self._saw_delta = False
         self._saw_terminal = False
         self._report_error = _chat_span_reporter(handle)
@@ -2416,7 +2668,15 @@ class _InstrumentedAsyncMediaStream:
             if failure is not None:
                 self._end(**self._finish_fields(), error=failure)
 
+    def _end(self, **fields: Any) -> None:
+        # The span has ended and the stream budget no longer counts bytes, so later events are
+        # not kept.
+        self._ended = True
+        self._end_stream(**fields)
+
     def _record(self, event: Any) -> None:
+        if self._ended:
+            return
         event_type = _string(_field(event, "type")) or ""
         if event_type in _MEDIA_TERMINAL_EVENTS:
             self._saw_terminal = True
@@ -2424,6 +2684,7 @@ class _InstrumentedAsyncMediaStream:
             event,
             self._response_mapper,
             self._capture_output,
+            self._budget,
         )
         if "output" in mapped:
             self._fields = _clean_fields(mapped)
@@ -2439,7 +2700,7 @@ class _InstrumentedAsyncMediaStream:
             if text:
                 captured = _capture_text(text, self._budget)
                 if captured:
-                    self._text_parts.append(captured)
+                    self._text_fragments.write(captured)
         if event_type in _MEDIA_TERMINAL_EVENTS:
             error = _media_stream_error(event, event_type)
             self._end(**self._finish_fields(), **({"error": error} if error is not None else {}))
@@ -2448,7 +2709,7 @@ class _InstrumentedAsyncMediaStream:
         return _media_stream_end_fields(
             self._fields,
             self._response_mapper,
-            self._text_parts,
+            self._text_fragments,
             self._budget,
             self._mask_output_when_incomplete,
             not self._saw_terminal or self._mapping_failed,