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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: telemetry-dev-openrouter
3
- Version: 0.2.1
3
+ Version: 0.2.3
4
4
  Summary: OpenRouter integration for telemetry.dev Python SDK
5
5
  Keywords: telemetry,opentelemetry,openrouter,llm,genai,tracing
6
6
  Author: telemetry.dev
@@ -134,9 +134,12 @@ still execute and return their normal SDK values without exporting telemetry.
134
134
  ## Limitations
135
135
 
136
136
  - Unconsumed streams end their spans only when the stream is exhausted, errors, or is closed.
137
- - Streaming output is retained up to 64 KiB and 1,024 items; terminal metadata and usage can still
138
- be recorded after the bound is reached. The configured `max_attribute_length` does not shrink this
139
- retention bound; the core applies it to the exported attribute after the mask runs.
137
+ - Streaming output is kept up to 16 MiB, and all streams in flight share the `max_stream_bytes`
138
+ bound of the core (64 MiB by default); terminal metadata and usage can still be recorded after
139
+ the bound is reached. The configured `max_attribute_length` does not shrink this
140
+ retention bound; the core applies it to the exported attribute after the mask runs. These bounds
141
+ need a `telemetry-dev` core with `max_stream_bytes`; with an older core, each stream keeps up to
142
+ 64 KiB and 1,024 items.
140
143
  - A chat stream that is cut at the bound, closed early, or fails, and a Responses stream that is cut
141
144
  at the bound, ends without a terminal response, or ends with an `error` event set
142
145
  `telemetry.dev.capture.truncated = true`. With a mask configured, the incomplete reconstructed
@@ -111,9 +111,12 @@ still execute and return their normal SDK values without exporting telemetry.
111
111
  ## Limitations
112
112
 
113
113
  - Unconsumed streams end their spans only when the stream is exhausted, errors, or is closed.
114
- - Streaming output is retained up to 64 KiB and 1,024 items; terminal metadata and usage can still
115
- be recorded after the bound is reached. The configured `max_attribute_length` does not shrink this
116
- retention bound; the core applies it to the exported attribute after the mask runs.
114
+ - Streaming output is kept up to 16 MiB, and all streams in flight share the `max_stream_bytes`
115
+ bound of the core (64 MiB by default); terminal metadata and usage can still be recorded after
116
+ the bound is reached. The configured `max_attribute_length` does not shrink this
117
+ retention bound; the core applies it to the exported attribute after the mask runs. These bounds
118
+ need a `telemetry-dev` core with `max_stream_bytes`; with an older core, each stream keeps up to
119
+ 64 KiB and 1,024 items.
117
120
  - A chat stream that is cut at the bound, closed early, or fails, and a Responses stream that is cut
118
121
  at the bound, ends without a terminal response, or ends with an `error` event set
119
122
  `telemetry.dev.capture.truncated = true`. With a mask configured, the incomplete reconstructed
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "telemetry-dev-openrouter"
3
- version = "0.2.1"
3
+ version = "0.2.3"
4
4
  description = "OpenRouter integration for telemetry.dev Python SDK"
5
5
  readme = "README.md"
6
6
  license = "MIT"
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "telemetry-dev-openrouter"
3
- version = "0.2.1"
3
+ version = "0.2.3"
4
4
  description = "OpenRouter integration for telemetry.dev Python SDK"
5
5
  readme = "README.md"
6
6
  license = "MIT"
@@ -1,10 +1,13 @@
1
1
  from __future__ import annotations
2
2
 
3
3
  import asyncio
4
+ import os
4
5
  import threading
5
6
  import time
7
+ import weakref
6
8
  from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Mapping, Sequence
7
- from functools import wraps
9
+ from functools import partial, wraps
10
+ from io import StringIO
8
11
  from typing import Any, TypeVar, cast
9
12
 
10
13
  import telemetry_dev
@@ -12,7 +15,7 @@ from openrouter.chat import Chat
12
15
  from openrouter.embeddings import Embeddings
13
16
  from openrouter.responses import Responses
14
17
 
15
- __version__ = "0.2.1"
18
+ __version__ = "0.2.3"
16
19
 
17
20
  RequestMapper = Callable[[Mapping[str, Any]], tuple[str, dict[str, Any]]]
18
21
  ResponseMapper = Callable[[Any], dict[str, Any]]
@@ -21,10 +24,8 @@ _PROVIDER = "openrouter"
21
24
  _WRAPPED_ATTR = "_telemetry_dev_openrouter_wrapped"
22
25
  _ORIGINAL_ATTR = "_telemetry_dev_openrouter_original"
23
26
  _CAPTURE_TRUNCATED_ATTRIBUTE = "telemetry.dev.capture.truncated"
24
- # Retention uses fixed bounds, the same as the TypeScript integration; the core applies
25
- # max_attribute_length after the mask.
26
- _STREAM_CAPTURE_MAX_BYTES = 64 * 1024
27
- _STREAM_CAPTURE_MAX_ITEMS = 1024
27
+ # The core applies max_attribute_length after the mask. More choices than this are not kept.
28
+ _MAX_STREAM_CHOICES = 1024
28
29
  _FINISH_REASON_MAX_LENGTH = 128
29
30
  _TERMINAL_RESPONSE_EVENTS = frozenset(
30
31
  {"response.completed", "response.failed", "response.incomplete"}
@@ -33,6 +34,22 @@ _OMIT = object()
33
34
  _ORIGINALS: list[tuple[type[Any], str, Any]] = []
34
35
  _installed = False
35
36
  _install_lock = threading.Lock()
37
+
38
+
39
+ def _reset_install_lock_in_child() -> None:
40
+ """A thread of the parent can hold the lock at a fork, and only this thread exists here."""
41
+ global _install_lock
42
+ _install_lock = threading.Lock()
43
+
44
+
45
+ # A hook is enough: a fork that runs no hooks, as uWSGI's fork of its workers by default,
46
+ # starts from its master, which serves no requests. os.register_at_fork cannot remove a hook,
47
+ # and importlib.reload runs this module again with the same globals, so the hook of the first
48
+ # run stays. Windows has no fork.
49
+ if hasattr(os, "register_at_fork") and "_install_fork_hook" not in globals():
50
+ _install_fork_hook = True
51
+ os.register_at_fork(after_in_child=_reset_install_lock_in_child)
52
+
36
53
  _T = TypeVar("_T")
37
54
 
38
55
 
@@ -174,6 +191,12 @@ def _cost_usd(raw_usage: Any) -> float | None:
174
191
  return cost if cost is not None else upstream_cost
175
192
 
176
193
 
194
+ def _prompt_attributes(messages: object) -> dict[str, Any]:
195
+ # Older cores have no prompt_attributes_for and record a prompt only inside with_prompt.
196
+ read = getattr(telemetry_dev, "prompt_attributes_for", None)
197
+ return cast("dict[str, Any]", read(messages)) if callable(read) else {}
198
+
199
+
177
200
  def _chat_request(params: Mapping[str, Any]) -> tuple[str, dict[str, Any]]:
178
201
  model = _string(params.get("model"))
179
202
  max_completion_tokens = _number(params.get("max_completion_tokens"))
@@ -193,6 +216,7 @@ def _chat_request(params: Mapping[str, Any]) -> tuple[str, dict[str, Any]]:
193
216
  "seed": _number(params.get("seed")),
194
217
  "frequency_penalty": _number(params.get("frequency_penalty")),
195
218
  "presence_penalty": _number(params.get("presence_penalty")),
219
+ "attributes": _prompt_attributes(params.get("messages")) or None,
196
220
  },
197
221
  )
198
222
 
@@ -260,6 +284,7 @@ def _responses_request(params: Mapping[str, Any]) -> tuple[str, dict[str, Any]]:
260
284
  "max_tokens": _number(params.get("max_output_tokens")),
261
285
  "frequency_penalty": _number(params.get("frequency_penalty")),
262
286
  "presence_penalty": _number(params.get("presence_penalty")),
287
+ "attributes": _prompt_attributes(params.get("input")) or None,
263
288
  },
264
289
  )
265
290
 
@@ -382,22 +407,107 @@ def _end_once(handle: telemetry_dev.SpanHandle) -> Callable[..., None]:
382
407
  return end
383
408
 
384
409
 
410
+ # On Python 3.10 and 3.11, StringIO keeps up to 100,000 written strings in a list before it
411
+ # joins them, about 60 bytes each that the stream bound does not charge: 15 bytes a character
412
+ # for 4-character deltas. Deltas are joined 256 at a time, so that those waiting hold under
413
+ # about 22 KB, onto a string that is written once it has 4,096 characters, so that what is
414
+ # written costs about 1.5% more than its text. Python 3.12 dropped the list.
415
+ _DELTA_BATCH = 256
416
+ _WRITE_CHARS = 4096
417
+
418
+
419
+ class _DeltaBuffer:
420
+ # A write takes several steps, unlike one StringIO call, so without the lock a read on
421
+ # another thread, such as one that closes the stream, could copy a batch twice or drop it.
422
+ # A fork copies the lock held if another thread was writing, and that thread does not run
423
+ # in the child to release it. Every fork, also one that runs no at-fork hooks, changes the
424
+ # process ID, so a child takes a new lock. The check is inline because it runs per delta.
425
+ # A signal handler runs on the thread it interrupts, so one that reads the stream, such as
426
+ # one that closes it, can read inside a write. So the lock is reentrant, and each step takes
427
+ # the deltas and the tail out of the buffer before it joins or writes them, so that a read
428
+ # inside the step does not write them as well; that read returns a prefix, which can leave
429
+ # out the deltas being joined. A handler that writes inside a step is not supported. An
430
+ # exception inside a step, such as KeyboardInterrupt from a handler, puts back the text the
431
+ # step took unless the StringIO has it, so that a later read still has every delta.
432
+ __slots__ = ("_buffer", "_lock", "_pending", "_pid", "_tail")
433
+
434
+ def __init__(self) -> None:
435
+ self._buffer = StringIO()
436
+ self._lock = threading.RLock()
437
+ self._pid = os.getpid()
438
+ self._pending: list[str] = []
439
+ self._tail = ""
440
+
441
+ def write(self, text: str) -> None:
442
+ pid = os.getpid()
443
+ if self._pid != pid:
444
+ self._lock = threading.RLock()
445
+ self._pid = pid
446
+ with self._lock:
447
+ self._pending.append(text)
448
+ if len(self._pending) >= _DELTA_BATCH:
449
+ self._join(_WRITE_CHARS)
450
+
451
+ def getvalue(self) -> str:
452
+ pid = os.getpid()
453
+ if self._pid != pid:
454
+ self._lock = threading.RLock()
455
+ self._pid = pid
456
+ with self._lock:
457
+ self._join(1)
458
+ return self._buffer.getvalue()
459
+
460
+ def _join(self, write_at: int) -> None:
461
+ if not self._pending and len(self._tail) < write_at:
462
+ return
463
+ written = self._buffer.tell()
464
+ pending: list[str] = []
465
+ tail = ""
466
+ try:
467
+ pending, self._pending = self._pending, []
468
+ tail, self._tail = self._tail, ""
469
+ joined = "".join([tail, *pending])
470
+ if len(joined) >= write_at:
471
+ self._buffer.write(joined)
472
+ else:
473
+ self._tail = joined
474
+ except BaseException:
475
+ if self._buffer.tell() == written:
476
+ self._pending[:0] = pending
477
+ self._tail = tail + self._tail
478
+ raise
479
+
480
+
385
481
  class _ChatChoice:
386
482
  def __init__(self) -> None:
387
483
  self.role: str | None = None
388
- self.content = ""
389
- self.reasoning = ""
484
+ # Deltas are written to buffers and joined when output is read: repeated concatenation
485
+ # copies the whole text for each delta.
486
+ self.content_fragments = _DeltaBuffer()
487
+ self.reasoning_fragments = _DeltaBuffer()
390
488
  self.reasoning_details: list[Any] = []
391
- self.refusal = ""
489
+ self.refusal_fragments = _DeltaBuffer()
392
490
  self.tool_calls: dict[int, dict[str, Any]] = {}
393
491
  self.finish_reason: str | None = None
394
492
 
493
+ @property
494
+ def content(self) -> str:
495
+ return self.content_fragments.getvalue()
496
+
497
+ @property
498
+ def reasoning(self) -> str:
499
+ return self.reasoning_fragments.getvalue()
500
+
501
+ @property
502
+ def refusal(self) -> str:
503
+ return self.refusal_fragments.getvalue()
504
+
395
505
 
396
506
  def _choice_state(
397
507
  states: dict[int, _ChatChoice], index: int, budget: telemetry_dev.CaptureBudget
398
508
  ) -> _ChatChoice | None:
399
509
  if index not in states:
400
- if len(states) >= budget.max_items:
510
+ if len(states) >= _MAX_STREAM_CHOICES:
401
511
  budget.truncated = True
402
512
  return None
403
513
  states[index] = _ChatChoice()
@@ -422,7 +532,10 @@ def _merge_tool_call(state: _ChatChoice, delta: Any) -> None:
422
532
  if name is not None:
423
533
  current_function["name"] = name
424
534
  if isinstance(arguments, str):
425
- current_function["arguments"] = f"{current_function.get('arguments', '')}{arguments}"
535
+ fragments = current_function.get("arguments")
536
+ if not isinstance(fragments, _DeltaBuffer):
537
+ fragments = current_function["arguments"] = _DeltaBuffer()
538
+ fragments.write(arguments)
426
539
  current["function"] = current_function
427
540
  state.tool_calls[tool_index] = current
428
541
 
@@ -443,15 +556,15 @@ def _record_chat_chunk(
443
556
  if isinstance(role, str):
444
557
  state.role = role
445
558
  if isinstance(content, str) and budget.accept(content):
446
- state.content += content
559
+ state.content_fragments.write(content)
447
560
  if isinstance(reasoning, str) and budget.accept(reasoning):
448
- state.reasoning += reasoning
561
+ state.reasoning_fragments.write(reasoning)
449
562
  for detail in _sequence_items(_field(delta, "reasoning_details")):
450
563
  native_detail = _native(detail)
451
564
  if native_detail is not None and budget.accept(native_detail):
452
565
  state.reasoning_details.append(native_detail)
453
566
  if isinstance(refusal, str) and budget.accept(refusal):
454
- state.refusal += refusal
567
+ state.refusal_fragments.write(refusal)
455
568
  for tool_call in _sequence_items(_field(delta, "tool_calls")):
456
569
  if budget.accept(tool_call):
457
570
  _merge_tool_call(state, tool_call)
@@ -520,6 +633,16 @@ def _response_event_has_output(event: Any) -> bool:
520
633
  )
521
634
 
522
635
 
636
+ def _joined_tool_call(call: dict[str, Any]) -> dict[str, Any]:
637
+ function = call.get("function")
638
+ if not isinstance(function, dict):
639
+ return call
640
+ arguments = cast(dict[str, Any], function).get("arguments")
641
+ if not isinstance(arguments, _DeltaBuffer):
642
+ return call
643
+ return {**call, "function": {**function, "arguments": arguments.getvalue()}}
644
+
645
+
523
646
  def _chat_output(states: Mapping[int, _ChatChoice]) -> list[dict[str, Any]]:
524
647
  output: list[dict[str, Any]] = []
525
648
  for _, state in sorted(states.items()):
@@ -535,13 +658,57 @@ def _chat_output(states: Mapping[int, _ChatChoice]) -> list[dict[str, Any]]:
535
658
  if state.refusal:
536
659
  message["refusal"] = state.refusal
537
660
  if state.tool_calls:
538
- message["tool_calls"] = [call for _, call in sorted(state.tool_calls.items())]
661
+ message["tool_calls"] = [
662
+ _joined_tool_call(call) for _, call in sorted(state.tool_calls.items())
663
+ ]
539
664
  output.append(message)
540
665
  return output
541
666
 
542
667
 
543
668
  def _stream_capture_budget() -> telemetry_dev.CaptureBudget:
544
- return telemetry_dev.CaptureBudget(_STREAM_CAPTURE_MAX_BYTES, _STREAM_CAPTURE_MAX_ITEMS)
669
+ # Older cores have no for_stream and bound each budget at 64 KiB on its own.
670
+ for_stream = getattr(telemetry_dev.CaptureBudget, "for_stream", None)
671
+ budget = for_stream() if callable(for_stream) else telemetry_dev.CaptureBudget()
672
+ return cast(telemetry_dev.CaptureBudget, budget)
673
+
674
+
675
+ def _release_budget(budget: telemetry_dev.CaptureBudget) -> None:
676
+ release = getattr(budget, "release", None)
677
+ if callable(release):
678
+ release()
679
+
680
+
681
+ def _measuring_budget(stream: telemetry_dev.CaptureBudget) -> telemetry_dev.CaptureBudget:
682
+ """An empty budget that measures a value replacing everything ``stream`` holds, within the
683
+ bound ``stream`` has now, without holding bytes itself."""
684
+ return telemetry_dev.CaptureBudget(stream.max_bytes, stream.max_items)
685
+
686
+
687
+ def _adopt_budget(
688
+ stream: telemetry_dev.CaptureBudget, measured: telemetry_dev.CaptureBudget
689
+ ) -> bool:
690
+ """Takes the counts of a budget that measured a replacement into the stream's own budget.
691
+ Returns False, and takes nothing, when another stream took the room after the measure."""
692
+ # Older cores have no reserve_bytes; there the measure against max_bytes stands.
693
+ reserve = getattr(stream, "reserve_bytes", None)
694
+ if callable(reserve):
695
+ if not reserve(measured.bytes_used):
696
+ return False
697
+ else:
698
+ stream.bytes_used = measured.bytes_used
699
+ stream.items_used = measured.items_used
700
+ stream.truncated = measured.truncated
701
+ return True
702
+
703
+
704
+ def _ending_stream(end: Callable[..., None], release: Callable[[], None]) -> Callable[..., None]:
705
+ def end_stream(**fields: Any) -> None:
706
+ try:
707
+ end(**fields)
708
+ finally:
709
+ release()
710
+
711
+ return end_stream
545
712
 
546
713
 
547
714
  def _capture_masked(handle: telemetry_dev.SpanHandle) -> bool:
@@ -611,6 +778,15 @@ def _hook_response_close(inner: Any, finish: Callable[[], None]) -> None:
611
778
  response = getattr(inner, "response", None)
612
779
  if response is None:
613
780
  return
781
+ # ``finish`` is a method of the stream wrapper, which holds ``response``: held weakly, the hook
782
+ # adds no reference cycle, so a dropped wrapper and its stream budget are freed at once.
783
+ finish_method = weakref.WeakMethod(finish)
784
+
785
+ def finish_if_alive() -> None:
786
+ method = finish_method()
787
+ if method is not None:
788
+ method()
789
+
614
790
  close = getattr(response, "close", None)
615
791
  if callable(close):
616
792
 
@@ -618,7 +794,7 @@ def _hook_response_close(inner: Any, finish: Callable[[], None]) -> None:
618
794
  try:
619
795
  return close(*args, **kwargs)
620
796
  finally:
621
- finish()
797
+ finish_if_alive()
622
798
 
623
799
  response.close = _close_hook
624
800
  aclose = getattr(response, "aclose", None)
@@ -629,7 +805,7 @@ def _hook_response_close(inner: Any, finish: Callable[[], None]) -> None:
629
805
  try:
630
806
  return await aclose_fn(*args, **kwargs)
631
807
  finally:
632
- finish()
808
+ finish_if_alive()
633
809
 
634
810
  response.aclose = _aclose_hook
635
811
 
@@ -657,7 +833,6 @@ async def _close_async_stream(inner: Any) -> None:
657
833
  class _InstrumentedStream:
658
834
  def __init__(self, inner: Any, handle: telemetry_dev.SpanHandle, started_at: float) -> None:
659
835
  self._inner = inner
660
- self._end = _end_once(handle)
661
836
  self._handle = handle
662
837
  self._started_at = started_at
663
838
  self._states: dict[int, _ChatChoice] = {}
@@ -669,6 +844,9 @@ class _InstrumentedStream:
669
844
  self._completed = False
670
845
  self._mask_output_when_incomplete = _capture_masked(handle)
671
846
  self._budget = _stream_capture_budget()
847
+ # Bound to the budget, not self, so the callback adds no reference cycle.
848
+ self._end_stream = _ending_stream(_end_once(handle), partial(_release_budget, self._budget))
849
+ self._ended = False
672
850
  _hook_response_close(inner, self._on_response_close)
673
851
 
674
852
  def _iterate(self) -> Iterator[Any]:
@@ -739,7 +917,15 @@ class _InstrumentedStream:
739
917
  def __getattr__(self, name: str) -> Any:
740
918
  return getattr(self._inner, name)
741
919
 
920
+ def _end(self, **fields: Any) -> None:
921
+ # The span has ended and the stream budget no longer counts bytes, so later events are
922
+ # not kept.
923
+ self._ended = True
924
+ self._end_stream(**fields)
925
+
742
926
  def _record(self, chunk: Any, received_at: float) -> bool:
927
+ if self._ended:
928
+ return True
743
929
  update = _record_chat_chunk(chunk, self._states, self._budget)
744
930
  if _chat_chunk_has_output(chunk):
745
931
  record_output_chunk = getattr(self._handle, "record_output_chunk", None)
@@ -762,7 +948,6 @@ class _InstrumentedStream:
762
948
  class _InstrumentedAsyncStream:
763
949
  def __init__(self, inner: Any, handle: telemetry_dev.SpanHandle, started_at: float) -> None:
764
950
  self._inner = inner
765
- self._end = _end_once(handle)
766
951
  self._handle = handle
767
952
  self._started_at = started_at
768
953
  self._states: dict[int, _ChatChoice] = {}
@@ -774,6 +959,9 @@ class _InstrumentedAsyncStream:
774
959
  self._completed = False
775
960
  self._mask_output_when_incomplete = _capture_masked(handle)
776
961
  self._budget = _stream_capture_budget()
962
+ # Bound to the budget, not self, so the callback adds no reference cycle.
963
+ self._end_stream = _ending_stream(_end_once(handle), partial(_release_budget, self._budget))
964
+ self._ended = False
777
965
  _hook_response_close(inner, self._on_response_close)
778
966
 
779
967
  async def _aiterate(self) -> AsyncIterator[Any]:
@@ -847,7 +1035,15 @@ class _InstrumentedAsyncStream:
847
1035
  def __getattr__(self, name: str) -> Any:
848
1036
  return getattr(self._inner, name)
849
1037
 
1038
+ def _end(self, **fields: Any) -> None:
1039
+ # The span has ended and the stream budget no longer counts bytes, so later events are
1040
+ # not kept.
1041
+ self._ended = True
1042
+ self._end_stream(**fields)
1043
+
850
1044
  def _record(self, chunk: Any, received_at: float) -> bool:
1045
+ if self._ended:
1046
+ return True
851
1047
  update = _record_chat_chunk(chunk, self._states, self._budget)
852
1048
  if _chat_chunk_has_output(chunk):
853
1049
  record_output_chunk = getattr(self._handle, "record_output_chunk", None)
@@ -867,6 +1063,37 @@ class _InstrumentedAsyncStream:
867
1063
  return _chunk_error(chunk) is None
868
1064
 
869
1065
 
1066
+ def _write_fragment(target: dict[str, Any], key: str, value: str) -> None:
1067
+ """Appends a delta to ``target[key]`` in a buffer, since repeated concatenation copies the
1068
+ whole text for each delta. _joined_item reads the text back."""
1069
+ fragments = target.get(key)
1070
+ if not isinstance(fragments, _DeltaBuffer):
1071
+ fragments = _DeltaBuffer()
1072
+ fragments.write(f"{target.get(key, '')}")
1073
+ target[key] = fragments
1074
+ fragments.write(value)
1075
+
1076
+
1077
+ def _joined_value(value: Any) -> Any:
1078
+ return value.getvalue() if isinstance(value, _DeltaBuffer) else value
1079
+
1080
+
1081
+ def _joined_item(item: Any) -> Any:
1082
+ if not isinstance(item, dict):
1083
+ return item
1084
+ joined: dict[str, Any] = {}
1085
+ for key, value in cast(dict[str, Any], item).items():
1086
+ if key in ("content", "summary") and isinstance(value, list):
1087
+ value = [
1088
+ {name: _joined_value(field) for name, field in cast(dict[str, Any], part).items()}
1089
+ if isinstance(part, dict)
1090
+ else part
1091
+ for part in cast(list[Any], value)
1092
+ ]
1093
+ joined[key] = _joined_value(value)
1094
+ return joined
1095
+
1096
+
870
1097
  class _ResponsesStreamState:
871
1098
  def __init__(self) -> None:
872
1099
  self.partial: dict[str, Any] = {}
@@ -895,7 +1122,10 @@ class _ResponsesStreamState:
895
1122
 
896
1123
  def _output_with_item(self, output_index: int, item: dict[str, Any]) -> list[dict[str, Any]]:
897
1124
  items = {**self._items, output_index: item}
898
- return [*[value for _, value in sorted(items.items())], *self._synthetic_events]
1125
+ return [
1126
+ *[_joined_item(value) for _, value in sorted(items.items())],
1127
+ *self._synthetic_events,
1128
+ ]
899
1129
 
900
1130
  @staticmethod
901
1131
  def _can_accept(budget: telemetry_dev.CaptureBudget, value: Any) -> bool:
@@ -906,14 +1136,19 @@ class _ResponsesStreamState:
906
1136
  return candidate.accept(value)
907
1137
 
908
1138
  def _replace_output(self, output: Any) -> bool:
909
- replacement = _stream_capture_budget()
1139
+ replacement = _measuring_budget(self._budget)
910
1140
  if not replacement.accept(output):
911
1141
  self._budget.truncated = True
912
1142
  return False
913
1143
  replacement.truncated = self._provider_events_truncated
914
- self._budget = replacement
1144
+ if not _adopt_budget(self._budget, replacement):
1145
+ self._budget.truncated = True
1146
+ return False
915
1147
  return True
916
1148
 
1149
+ def release(self) -> None:
1150
+ _release_budget(self._budget)
1151
+
917
1152
  def _sync_output(self) -> None:
918
1153
  self.retained_output = self._output()
919
1154
  self.partial["output"] = self.retained_output
@@ -999,7 +1234,7 @@ class _ResponsesStreamState:
999
1234
  if event_type in {"response.output_item.added", "response.output_item.done"}:
1000
1235
  raw_item = _field(event, "item")
1001
1236
  replacement = event_type.endswith(".done")
1002
- preflight = _stream_capture_budget() if replacement else self._budget
1237
+ preflight = _measuring_budget(self._budget) if replacement else self._budget
1003
1238
  if not self._can_accept(preflight, raw_item):
1004
1239
  self._budget.truncated = True
1005
1240
  return False
@@ -1017,7 +1252,7 @@ class _ResponsesStreamState:
1017
1252
  if event_type in {"response.content_part.added", "response.content_part.done"}:
1018
1253
  event_part = _field(event, "part")
1019
1254
  replacement = event_type.endswith(".done")
1020
- preflight = _stream_capture_budget() if replacement else self._budget
1255
+ preflight = _measuring_budget(self._budget) if replacement else self._budget
1021
1256
  if not self._can_accept(preflight, event_part):
1022
1257
  self._budget.truncated = True
1023
1258
  return False
@@ -1119,7 +1354,7 @@ class _ResponsesStreamState:
1119
1354
  else:
1120
1355
  if not self._budget.accept(value):
1121
1356
  return False
1122
- item[field] = f"{item.get(field, '')}{value}"
1357
+ _write_fragment(item, field, value)
1123
1358
  name = _field(event, "name")
1124
1359
  if name is not None:
1125
1360
  item["name"] = name
@@ -1130,7 +1365,7 @@ class _ResponsesStreamState:
1130
1365
  }:
1131
1366
  event_part = _field(event, "part")
1132
1367
  replacement = event_type.endswith(".done")
1133
- preflight = _stream_capture_budget() if replacement else self._budget
1368
+ preflight = _measuring_budget(self._budget) if replacement else self._budget
1134
1369
  if not self._can_accept(preflight, event_part):
1135
1370
  self._budget.truncated = True
1136
1371
  return False
@@ -1245,7 +1480,7 @@ class _ResponsesStreamState:
1245
1480
  self._summary[summary_key] = part
1246
1481
  self._sync_item(output_index, item)
1247
1482
  else:
1248
- summary["text"] = f"{summary.get('text', '')}{value}"
1483
+ _write_fragment(summary, "text", value)
1249
1484
  return True
1250
1485
  content_index = self._event_index(event, "content_index", "contentIndex")
1251
1486
  content_key = (output_index, content_index)
@@ -1284,9 +1519,9 @@ class _ResponsesStreamState:
1284
1519
  self._content[content_key] = replacement
1285
1520
  self._sync_item(output_index, item)
1286
1521
  elif refusal:
1287
- part["refusal"] = f"{part.get('refusal', '')}{value}"
1522
+ _write_fragment(part, "refusal", value)
1288
1523
  else:
1289
- part["text"] = f"{part.get('text', '')}{value}"
1524
+ _write_fragment(part, "text", value)
1290
1525
  return True
1291
1526
 
1292
1527
  def record(self, event: Any) -> None:
@@ -1294,7 +1529,7 @@ class _ResponsesStreamState:
1294
1529
  response = _field(event, "response")
1295
1530
  if response is not None:
1296
1531
  fields = _responses_response(response, include_error=False, include_output=False)
1297
- response_budget = _stream_capture_budget()
1532
+ response_budget = _measuring_budget(self._budget)
1298
1533
  raw_output = _field(response, "output")
1299
1534
  if raw_output is not None:
1300
1535
  if not self._can_accept(response_budget, raw_output):
@@ -1306,7 +1541,11 @@ class _ResponsesStreamState:
1306
1541
  self._budget.truncated = True
1307
1542
  else:
1308
1543
  response_artifact = [*response_items, *self._synthetic_events]
1544
+ adopted = False
1309
1545
  if response_budget.accept(response_artifact):
1546
+ response_budget.truncated = self._provider_events_truncated
1547
+ adopted = _adopt_budget(self._budget, response_budget)
1548
+ if adopted:
1310
1549
  self._items.clear()
1311
1550
  self._content.clear()
1312
1551
  self._summary.clear()
@@ -1314,8 +1553,6 @@ class _ResponsesStreamState:
1314
1553
  if isinstance(item, dict):
1315
1554
  self._hydrate_item(index, cast(dict[str, Any], item))
1316
1555
  self.retained_output = response_artifact
1317
- response_budget.truncated = self._provider_events_truncated
1318
- self._budget = response_budget
1319
1556
  else:
1320
1557
  self._budget.truncated = True
1321
1558
  self.partial = fields
@@ -1329,6 +1566,9 @@ class _ResponsesStreamState:
1329
1566
 
1330
1567
  def fields(self, mask_output: bool, *, failed: bool = False) -> dict[str, Any]:
1331
1568
  fields = dict(self.partial)
1569
+ output = fields.get("output")
1570
+ if isinstance(output, list):
1571
+ fields["output"] = [_joined_item(item) for item in cast(list[Any], output)]
1332
1572
  if failed or self._budget.truncated or not self.saw_terminal_snapshot:
1333
1573
  _mark_capture_incomplete(fields, mask_output)
1334
1574
  return fields
@@ -1338,10 +1578,12 @@ class _InstrumentedResponsesStream:
1338
1578
  def __init__(self, inner: Any, handle: telemetry_dev.SpanHandle, started_at: float) -> None:
1339
1579
  self._inner = inner
1340
1580
  self._handle = handle
1341
- self._end = _end_once(handle)
1342
1581
  self._started_at = started_at
1343
1582
  self._saw_first = False
1344
1583
  self._state = _ResponsesStreamState()
1584
+ # Bound to the state, not self, so the callback adds no reference cycle.
1585
+ self._end_stream = _ending_stream(_end_once(handle), self._state.release)
1586
+ self._ended = False
1345
1587
  self._mask_output_when_incomplete = _capture_masked(handle)
1346
1588
  self._consume: Iterator[Any] | None = None
1347
1589
  self._in_next = False
@@ -1409,7 +1651,15 @@ class _InstrumentedResponsesStream:
1409
1651
  def __getattr__(self, name: str) -> Any:
1410
1652
  return getattr(self._inner, name)
1411
1653
 
1654
+ def _end(self, **fields: Any) -> None:
1655
+ # The span has ended and the stream budget no longer counts bytes, so later events are
1656
+ # not kept.
1657
+ self._ended = True
1658
+ self._end_stream(**fields)
1659
+
1412
1660
  def _record(self, event: Any, received_at: float) -> None:
1661
+ if self._ended:
1662
+ return
1413
1663
  if _response_event_has_output(event):
1414
1664
  record_output_chunk = getattr(self._handle, "record_output_chunk", None)
1415
1665
  if callable(record_output_chunk):
@@ -1436,10 +1686,12 @@ class _InstrumentedAsyncResponsesStream:
1436
1686
  def __init__(self, inner: Any, handle: telemetry_dev.SpanHandle, started_at: float) -> None:
1437
1687
  self._inner = inner
1438
1688
  self._handle = handle
1439
- self._end = _end_once(handle)
1440
1689
  self._started_at = started_at
1441
1690
  self._saw_first = False
1442
1691
  self._state = _ResponsesStreamState()
1692
+ # Bound to the state, not self, so the callback adds no reference cycle.
1693
+ self._end_stream = _ending_stream(_end_once(handle), self._state.release)
1694
+ self._ended = False
1443
1695
  self._mask_output_when_incomplete = _capture_masked(handle)
1444
1696
  self._consume: AsyncIterator[Any] | None = None
1445
1697
  self._in_next = False
@@ -1510,7 +1762,15 @@ class _InstrumentedAsyncResponsesStream:
1510
1762
  def __getattr__(self, name: str) -> Any:
1511
1763
  return getattr(self._inner, name)
1512
1764
 
1765
+ def _end(self, **fields: Any) -> None:
1766
+ # The span has ended and the stream budget no longer counts bytes, so later events are
1767
+ # not kept.
1768
+ self._ended = True
1769
+ self._end_stream(**fields)
1770
+
1513
1771
  def _record(self, event: Any, received_at: float) -> None:
1772
+ if self._ended:
1773
+ return
1514
1774
  if _response_event_has_output(event):
1515
1775
  record_output_chunk = getattr(self._handle, "record_output_chunk", None)
1516
1776
  if callable(record_output_chunk):