From 336b72f785ab2170ac7059a4fd465cd8093654b4 Mon Sep 17 00:00:00 2001 From: Robin Date: Wed, 16 Sep 2026 12:22:40 -0700 Subject: [PATCH] root latency ends at the last chunk, tracing failures never escape into the caller's loop, parent span captured at call time, private attributes never delegate, KeyboardInterrupt keeps partial output, time-to-first-token omitted when no chunk arrived --- README.md | 6 ++ TRACING.md | 2 +- agentx/integrations/_traced_call.py | 69 +++++++++++++-- agentx/integrations/anthropic.py | 74 +++++++++++----- agentx/integrations/nvidia_nim.py | 12 ++- agentx/integrations/openai.py | 25 +++++- agentx/monitor/__init__.py | 6 +- agentx/monitor/alert_rules.py | 38 ++++++--- agentx/tracing/tracer.py | 6 +- tests/test_alert_rules.py | 12 +++ tests/test_integrations.py | 125 +++++++++++++++++++++++++++- 11 files changed, 322 insertions(+), 53 deletions(-) diff --git a/README.md b/README.md index 699e1a5..e3de9d5 100644 --- a/README.md +++ b/README.md @@ -193,6 +193,12 @@ extra: | LlamaIndex | `pip install "agentx-python[llamaindex]"` | `AgentXLlamaIndexHandler` | | AutoGen | `pip install "agentx-python[autogen]"` | `AgentXAutoGenObserver` | +Raw-client patches (`patch_openai_client`, `patch_nim_client`, `patch_anthropic_client`) trace +streaming calls too: the returned stream is a transparent proxy that assembles the reply from the +chunks you consume, with latency measured to the last chunk and the time to first token in the +trace metadata. OpenAI-compatible endpoints only send token usage on streams when you pass +`stream_options={"include_usage": True}`. + > **Warning: pick ONE instrumentation layer per LLM call.** Do not combine > `AgentXCallbackHandler` (or any framework integration) with a patched provider client > (`patch_openai_client`, `patch_anthropic_client`, `patch_genai_client`) on the same code diff --git a/TRACING.md b/TRACING.md index 36f7bee..a80af33 100644 --- a/TRACING.md +++ b/TRACING.md @@ -191,7 +191,7 @@ works. The label resolves in priority order: | `AgentXCallbackHandler` (LangChain/LangGraph) | `langchain` | | `AgentXCrewObserver` | `crewai` | | `AgentXTracingProcessor` (OpenAI Agents SDK) | `openai-agents` | - | `patch_openai_client` | `openai` | + | `patch_openai_client` (streaming included) | `openai` | | `patch_nim_client` (NVIDIA NIM) | `nvidia-nim` | | `patch_anthropic_client` | `anthropic` | | `patch_genai_client` | `google-genai` | diff --git a/agentx/integrations/_traced_call.py b/agentx/integrations/_traced_call.py index 11c704b..18c997d 100644 --- a/agentx/integrations/_traced_call.py +++ b/agentx/integrations/_traced_call.py @@ -14,9 +14,16 @@ import asyncio import inspect import json +import logging +import threading import time from typing import Any, Callable, Dict, Optional +logger = logging.getLogger(__name__) + +# Sentinel for finish_llm_call's `active_span`: "not passed" is distinct from "passed None". +_UNSET: Any = object() + from agentx.tracing.tracer import Tracer, _safe_serialize @@ -109,6 +116,7 @@ def finish_llm_call( cache_write_tokens: Optional[int] = None, tool_definitions: Optional[list] = None, call_metadata: Optional[Dict[str, Any]] = None, + active_span: Any = _UNSET, ) -> None: """ Close out one raw-client LLM call - shared by the ``on_finish``/exit @@ -129,7 +137,11 @@ def finish_llm_call( if tool_definitions: metadata = {**(metadata or {}), "tools": tool_definitions} - active_span = tracer.current_span + # The parent is the span that was active when the CALL was made. Streaming patches pass it + # explicitly: a stream finalizes later (exhaustion, close, or garbage collection), by which + # time a different span may be active, and the call must not be grafted onto it. + if active_span is _UNSET: + active_span = tracer.current_span if active_span is not None: # The definitions describe the whole call's toolbox - attach them to the enclosing # span's metadata (first capture wins) so the ROOT trace carries them for the @@ -168,6 +180,9 @@ def finish_llm_call( ) span.__enter__() span._start = start_t + # The call ended at end_t (a stream's last chunk), not at whatever later moment this + # runs - __exit__ honors the override instead of measuring to time.time(). + span._end_override = end_t span.input = input_repr span.output = output if error: @@ -217,7 +232,13 @@ class TracedStream: Attribute access falls through to the wrapped stream (``.response``, provider helpers), and ``__iter__``/``__aiter__`` return ``self`` so early - ``break`` leaves no half-driven generator behind. + ``break`` leaves no half-driven generator behind. It is a proxy, not a + subclass: ``isinstance(stream, openai.Stream)`` is False and ``repr()`` + shows the proxy - branch on ``stream=True`` in your own code, not on type. + + Tracing never breaks the caller: a failure while building or sending the + trace is logged and swallowed, and the stream's own iteration/close + semantics are untouched. """ def __init__( @@ -230,6 +251,8 @@ def __init__( self._accumulator = accumulator self._on_finish = on_finish self._done = False + # A watchdog close() racing the reader's StopIteration must not finalize twice. + self._done_lock = threading.Lock() self._first_chunk_t: Optional[float] = None self._last_chunk_t: Optional[float] = None self._sync_iter: Any = None @@ -250,9 +273,10 @@ def _observe(self, chunk: Any) -> None: pass def _finish(self, error: Optional[str]) -> None: - if self._done: - return - self._done = True + with self._done_lock: + if self._done: + return + self._done = True try: result = self._accumulator.result() except Exception: @@ -264,7 +288,12 @@ def _finish(self, error: Optional[str]) -> None: # The response "ended" at its last chunk, not at whatever later moment the caller closed # or dropped the stream - that is the latency the user experienced. result["end_t"] = self._last_chunk_t if self._last_chunk_t is not None else time.time() - self._on_finish(result, error) + try: + self._on_finish(result, error) + except Exception: + # Building or sending the trace failed. The caller's stream ended normally and must + # see it end normally - tracing is never allowed to raise into inference code. + logger.debug("Streamed call could not be traced", exc_info=True) @property def first_chunk_at(self) -> Optional[float]: @@ -283,9 +312,14 @@ def __next__(self) -> Any: except StopIteration: self._finish(None) raise - except BaseException as exc: + except Exception as exc: self._finish(str(exc)) raise + except BaseException: + # KeyboardInterrupt / GeneratorExit: a cancellation, not the provider failing - + # record what streamed so far without inventing an error message. + self._finish(None) + raise self._observe(chunk) return chunk @@ -302,9 +336,12 @@ async def __anext__(self) -> Any: except StopAsyncIteration: self._finish(None) raise - except BaseException as exc: + except Exception as exc: self._finish(str(exc)) raise + except BaseException: + self._finish(None) + raise self._observe(chunk) return chunk @@ -338,7 +375,17 @@ def close(self) -> None: close = getattr(self._stream, "close", None) try: if close is not None: - close() + result = close() + if inspect.isawaitable(result): + # openai's AsyncStream spells its close `async def close()`. A sync close() + # on it (an easy slip inside async code) would drop the coroutine and leak + # the connection; run it on the loop when there is one, else at least don't + # leave an un-awaited coroutine behind. + try: + asyncio.get_running_loop().create_task(result) + except RuntimeError: + result.close() # type: ignore[union-attr] + logger.warning("close() called on an async stream outside an event loop - use aclose()") finally: self._finish(None) @@ -355,6 +402,10 @@ async def aclose(self) -> None: self._finish(None) def __getattr__(self, item: str) -> Any: + # Only public attributes delegate. Private names must resolve on the proxy itself, or a + # half-constructed instance (no _stream yet) would recurse forever looking for it. + if item.startswith("_"): + raise AttributeError(item) return getattr(self._stream, item) def __del__(self) -> None: diff --git a/agentx/integrations/anthropic.py b/agentx/integrations/anthropic.py index fd3e02b..61e8f71 100644 --- a/agentx/integrations/anthropic.py +++ b/agentx/integrations/anthropic.py @@ -232,13 +232,21 @@ def patched_create(*args, **kwargs): input_repr = _safe_serialize(input_messages) if kwargs.get("stream"): + # Parent fixed at call time - see openai.py's patched_create for why. + parent = tracer.current_span + def on_stream_finish(collected: Dict[str, Any], error: Optional[str]) -> None: + ttft = collected.get("time_to_first_token_ms") + call_metadata: Dict[str, Any] = {"streaming": True} + if ttft is not None: + call_metadata["timeToFirstTokenMs"] = ttft finish_llm_call( tracer, name=name, framework="anthropic", metadata=metadata, - call_metadata={"streaming": True, "timeToFirstTokenMs": collected.get("time_to_first_token_ms")}, + call_metadata=call_metadata, + active_span=parent, session_id=session_id, start_t=start_t, end_t=collected.get("end_t") or time.time(), @@ -318,8 +326,9 @@ def patched_stream(*args, **kwargs): # only shows up in whether `with`/`async with` and # `get_final_message()` are used, handled inside `_TracedStream`. start_t = time.time() + parent = tracer.current_span ctx = original_stream(*args, **kwargs) - input_repr = _safe_serialize(_prepend_system(kwargs.get("messages"), kwargs.get("system"))) + input_repr = _safe_serialize(_prepend_system(kwargs.get("messages") or (args[0] if args else None), kwargs.get("system"))) model = kwargs.get("model") tool_definitions = capture_tool_definitions(kwargs.get("tools")) @@ -365,6 +374,29 @@ class _TracedStream: """ _inner: Any = None + _sent: bool = False + + # What streamed so far, WITHOUT draining the rest of the response: the SDK's + # get_final_message() calls until_done(), which would turn an early `break` into a + # blocking read of every remaining token. The snapshot is the final message once the + # stream was consumed, and honestly partial when the caller stopped early. + def _snapshot(self_inner): + inner = self_inner._inner + if inner is None: + return None + try: + return getattr(inner, "current_message_snapshot", None) + except Exception: + return None + + def _send_once(self_inner, end_t: float, error: Optional[str], snapshot: Any) -> None: + if self_inner._sent: + return + self_inner._sent = True + try: + build_and_send(end_t, error, snapshot) + except Exception: + pass # tracing never raises into the caller def __enter__(self_inner): self_inner._inner = ctx.__enter__() @@ -373,15 +405,12 @@ def __enter__(self_inner): def __exit__(self_inner, exc_type, exc_val, tb): end_t = time.time() error = str(exc_val) if exc_val else None - final_message = None - if error is None and self_inner._inner is not None: - try: - final_message = self_inner._inner.get_final_message() - except Exception: - pass - result = ctx.__exit__(exc_type, exc_val, tb) - build_and_send(end_t, error, final_message) - return result + # Snapshot BEFORE the manager closes the stream (the earlier bug read it after). + snapshot = self_inner._snapshot() + try: + return ctx.__exit__(exc_type, exc_val, tb) + finally: + self_inner._send_once(end_t, error, snapshot) async def __aenter__(self_inner): self_inner._inner = await ctx.__aenter__() @@ -390,16 +419,19 @@ async def __aenter__(self_inner): async def __aexit__(self_inner, exc_type, exc_val, tb): end_t = time.time() error = str(exc_val) if exc_val else None - final_message = None - if error is None and self_inner._inner is not None: - try: - raw = self_inner._inner.get_final_message() - final_message = await raw if inspect.isawaitable(raw) else raw - except Exception: - pass - result = await ctx.__aexit__(exc_type, exc_val, tb) - build_and_send(end_t, error, final_message) - return result + snapshot = self_inner._snapshot() + try: + return await ctx.__aexit__(exc_type, exc_val, tb) + finally: + self_inner._send_once(end_t, error, snapshot) + + def __del__(self_inner): + # A helper stream that was entered but never exited still records what it saw. + try: + if self_inner._inner is not None: + self_inner._send_once(time.time(), None, self_inner._snapshot()) + except Exception: + pass def __iter__(self_inner): return iter(ctx) diff --git a/agentx/integrations/nvidia_nim.py b/agentx/integrations/nvidia_nim.py index 1ac71a3..4120c7c 100644 --- a/agentx/integrations/nvidia_nim.py +++ b/agentx/integrations/nvidia_nim.py @@ -27,8 +27,11 @@ usage comes straight off the response's OpenAI-shaped ``usage`` block; NIM reports no prompt-cache fields, so cache token counts stay unset. -Streaming calls (``stream=True``) are passed through untouched and are not -currently traced - same posture as ``patch_openai_client``, see its docstring. +Streaming calls (``stream=True``) are traced too, exactly as +``patch_openai_client`` traces them: the stream is wrapped in a transparent +proxy that assembles the reply from the consumed chunks (token usage when the +endpoint sends it on the final chunk, e.g. with +``stream_options={"include_usage": True}``). Requires: ``pip install "agentx-python[nvidia-nim]"`` (installs the ``openai`` client package; there is no separate NIM SDK dependency). @@ -56,8 +59,9 @@ def patch_nim_client( call with ``framework="nvidia-nim"``. The original method is still called and its return value passed through - unchanged. Sync and async clients both work; ``stream=True`` calls pass - through untraced. Patching is idempotent - and because it shares the guard + unchanged. Sync and async clients both work; ``stream=True`` calls are + traced through the same stream proxy as ``patch_openai_client``. Patching + is idempotent - and because it shares the guard with ``patch_openai_client``, whichever of the two patched a given client first wins (patch each client with the integration that matches where its ``base_url`` actually points). diff --git a/agentx/integrations/openai.py b/agentx/integrations/openai.py index 3f6ac0f..d884c98 100644 --- a/agentx/integrations/openai.py +++ b/agentx/integrations/openai.py @@ -30,7 +30,14 @@ import time from typing import Any, Dict, Optional, Tuple +import logging + from agentx.tracing.tracer import Tracer, _safe_serialize + +logger = logging.getLogger(__name__) +# Warn once per process, not per call: a streamed OpenAI call carries no usage unless the caller +# asked for it, and a silent zero would under-report every streaming app's spend. +_warned_stream_usage = False from agentx.integrations._traced_call import ( StreamAccumulator, capture_tool_definitions, @@ -203,14 +210,30 @@ def patched_create(*args, **kwargs): tool_definitions = capture_tool_definitions(kwargs.get("tools")) if kwargs.get("stream"): + # The parent is fixed at call time: the stream finalizes later, possibly inside an + # unrelated span (or none), and must not attach to whatever is active then. + parent = tracer.current_span + def on_stream_finish(collected: Dict[str, Any], error: Optional[str]) -> None: + global _warned_stream_usage ttft = collected.get("time_to_first_token_ms") + call_metadata: Dict[str, Any] = {"streaming": True} + if ttft is not None: + call_metadata["timeToFirstTokenMs"] = ttft + if error is None and collected.get("input_tokens") is None and not _warned_stream_usage: + _warned_stream_usage = True + logger.warning( + "Streamed %s call carried no token usage - pass stream_options={\"include_usage\": True} " + "so traces (and cost) reflect streamed traffic.", + framework, + ) finish_llm_call( tracer, name=name, framework=framework, metadata=metadata, - call_metadata={"streaming": True, "timeToFirstTokenMs": ttft}, + call_metadata=call_metadata, + active_span=parent, session_id=session_id, start_t=start_t, end_t=collected.get("end_t") or time.time(), diff --git a/agentx/monitor/__init__.py b/agentx/monitor/__init__.py index 03c4bc0..8fb2f28 100644 --- a/agentx/monitor/__init__.py +++ b/agentx/monitor/__init__.py @@ -19,14 +19,14 @@ from agentx.monitor.signals import MonitorSignalClient __all__ = [ - "AlertEvent", - "AlertRule", - "AlertRulesClient", "AgentXImprovementGroupsError", "AgentXJudgeScorersError", "AgentXMonitorError", "AgentXScorerGroupsError", "AgentXScorersError", + "AlertEvent", + "AlertRule", + "AlertRulesClient", "ImprovementGroupsClient", "JudgeScorer", "JudgeScorerBuilder", diff --git a/agentx/monitor/alert_rules.py b/agentx/monitor/alert_rules.py index d2d3c82..d1c4edd 100644 --- a/agentx/monitor/alert_rules.py +++ b/agentx/monitor/alert_rules.py @@ -17,6 +17,22 @@ } +def _validate( + metric: Optional[str] = None, operator: Optional[str] = None, channels: Optional[List[Dict[str, str]]] = None +) -> None: + """Local checks for the fields the engine would otherwise 400 on - ``None`` means "not given" + (an update that leaves the field alone).""" + if metric is not None and metric not in ALERT_METRICS: + raise ValueError(f"metric must be one of {ALERT_METRICS}, got {metric!r}") + if operator is not None and operator not in ("gt", "lt"): + raise ValueError(f"operator must be 'gt' or 'lt', got {operator!r}") + if channels is not None: + for channel in channels: + kind = channel.get("kind") if isinstance(channel, dict) else None + if kind not in ALERT_CHANNEL_KINDS: + raise ValueError(f"channel kind must be one of {ALERT_CHANNEL_KINDS}, got {kind!r}") + + class AlertRule(dict): """Wire object for one KPI alert rule (dict subclass so unknown fields round-trip).""" @@ -48,12 +64,12 @@ class AlertEvent(dict): @property def kind(self) -> str: - return str(self.get("kind")) + return str(self.get("kind") or "") @property def delivered(self) -> bool: deliveries = self.get("deliveries") or [] - return bool(deliveries) and all(bool(d.get("ok")) for d in deliveries) + return bool(deliveries) and all(isinstance(d, dict) and bool(d.get("ok")) for d in deliveries) def slack(url: str) -> Dict[str, str]: @@ -138,13 +154,7 @@ def create( """Create a rule. ``operator`` is ``"gt"`` (above) or ``"lt"`` (below); rates are fractions (``0.10`` = 10%), latency is milliseconds, cost is USD. ``channels`` takes the dicts the module-level helpers build (``slack(url)``, ``pagerduty(key)``, ...).""" - if metric not in ALERT_METRICS: - raise ValueError(f"metric must be one of {ALERT_METRICS}, got {metric!r}") - if operator not in ("gt", "lt"): - raise ValueError(f"operator must be 'gt' or 'lt', got {operator!r}") - for channel in channels: - if channel.get("kind") not in ALERT_CHANNEL_KINDS: - raise ValueError(f"channel kind must be one of {ALERT_CHANNEL_KINDS}, got {channel.get('kind')!r}") + _validate(metric=metric, operator=operator, channels=channels) payload: Dict[str, Any] = { "name": name, "metric": metric, @@ -165,7 +175,10 @@ def create( def update(self, rule_id: str, **fields: Any) -> AlertRule: """Sparse update; snake_case kwargs are mapped to the wire. Changing the metric, - operator, threshold, window, or agent resets the rule's firing state.""" + operator, threshold, window, or agent resets the rule's firing state, and a rule that + was firing sends its channels a final ``resolved`` notification first. The same local + checks as ``create`` apply to whichever of ``metric``, ``operator``, ``channels`` are + given.""" payload: Dict[str, Any] = {} for key, value in fields.items(): wire_key = _ALIASES.get(key, key) @@ -175,12 +188,13 @@ def update(self, rule_id: str, **fields: Any) -> AlertRule: "silently ignore this (see AlertRulesClient.create for the field names)." ) payload[wire_key] = value + _validate(metric=payload.get("metric"), operator=payload.get("operator"), channels=payload.get("channels")) data = self._request("PUT", f"/agent-monitoring/alert-rules/{rule_id}", json=payload) return AlertRule(data.get("rule", data)) def delete(self, rule_id: str) -> None: - """Deletes the rule and its history. A PagerDuty incident the rule opened is not - resolved by this - close it in PagerDuty.""" + """Deletes the rule and its history. A rule that is firing sends its channels a final + ``resolved`` notification (closing the PagerDuty incident it opened) before it goes.""" self._request("DELETE", f"/agent-monitoring/alert-rules/{rule_id}", retry=False) def events(self, rule_id: str, limit: int = 50) -> List[AlertEvent]: diff --git a/agentx/tracing/tracer.py b/agentx/tracing/tracer.py index a5a1d32..03beeda 100644 --- a/agentx/tracing/tracer.py +++ b/agentx/tracing/tracer.py @@ -165,6 +165,9 @@ def __init__( self.tool_calls: list = [] self._start: Optional[float] = None + # Set by callers that know when the work actually ended (a streamed LLM call's last + # chunk) so __exit__ does not measure to "now" - see finish_llm_call's root path. + self._end_override: Optional[float] = None self._error: Optional[str] = None self._captured_model: Optional[str] = None @@ -212,7 +215,8 @@ def __enter__(self) -> "_TraceSpan": def __exit__(self, exc_type, exc_val, tb): self._tracer._pop_active_span(self) - latency_ms = int((time.time() - self._start) * 1000) if self._start else None + ended_at = self._end_override if self._end_override is not None else time.time() + latency_ms = int((ended_at - self._start) * 1000) if self._start else None if exc_val is not None and self._error is None: self._error = str(exc_val) diff --git a/tests/test_alert_rules.py b/tests/test_alert_rules.py index 51c2fef..0bf1419 100644 --- a/tests/test_alert_rules.py +++ b/tests/test_alert_rules.py @@ -89,6 +89,18 @@ def test_update_maps_snake_case_and_refuses_unknown_keys(): AlertRulesClient(fake).update("a1", sample_rate=0.5) # type: ignore[arg-type] +def test_update_applies_the_same_local_checks_as_create(): + fake = FakeMonitorClient([]) + client = AlertRulesClient(fake) # type: ignore[arg-type] + with pytest.raises(ValueError, match="metric"): + client.update("a1", metric="vibes") + with pytest.raises(ValueError, match="operator"): + client.update("a1", operator="ge") + with pytest.raises(ValueError, match="channel kind"): + client.update("a1", channels=[{"kind": "sms", "target": "1"}]) + assert fake.calls == [] + + def test_events_test_preview_and_sweep_paths(): fake = FakeMonitorClient( [ diff --git a/tests/test_integrations.py b/tests/test_integrations.py index 4fed5aa..865e888 100644 --- a/tests/test_integrations.py +++ b/tests/test_integrations.py @@ -278,10 +278,17 @@ def __init__(self): self.closed = False self.text_stream = iter(["streamed ", "reply"]) - def get_final_message(self): + @property + def current_message_snapshot(self): assert not self.closed, "must be read before the manager closes the stream" return final + def get_final_message(self): + # The SDK's get_final_message() drains the rest of the response (until_done()): + # an early `break` would block until the model finished. The wrapper must never + # call it. + raise AssertionError("get_final_message() drains the stream - use the snapshot") + class FakeManager: def __init__(self): self.stream = FakeMessageStream() @@ -314,6 +321,122 @@ def stream(self, **kwargs): assert kwargs["output_tokens"] == 4 +def test_stream_root_trace_latency_ends_at_the_last_chunk_not_at_finalization(): + # The proxy hands finish_llm_call the last-chunk time; the root span must honor it instead + # of measuring to whatever later moment the caller dropped the stream. + from agentx.integrations.openai import patch_openai_client + + stream = _FakeStream(_stream_chunks()) + + class FakeCompletions: + def create(self, **kwargs): + return stream + + client = _fake_openai_client(FakeCompletions()) + tracer = make_tracer() + patch_openai_client(client, tracer, name="gpt-agent") + result = client.chat.completions.create(model="gpt-4o-mini", messages=[], stream=True) + for _ in result: + pass + time.sleep(0.15) # the caller holds the exhausted stream a while before closing it + result.close() + _, kwargs = tracer._send.call_args + assert kwargs["latency_ms"] < 100 + + +def test_stream_tracing_failure_never_escapes_into_the_callers_loop(): + from agentx.integrations.openai import patch_openai_client + + class FakeCompletions: + def create(self, **kwargs): + return _FakeStream(_stream_chunks()) + + client = _fake_openai_client(FakeCompletions()) + tracer = make_tracer() + tracer._send = MagicMock(side_effect=RuntimeError("ingest exploded")) + patch_openai_client(client, tracer, name="gpt-agent") + chunks = list(client.chat.completions.create(model="gpt-4o-mini", messages=[], stream=True)) + assert len(chunks) == 4 # the loop ended normally despite the tracer raising + + +def test_stream_parent_is_the_span_active_at_call_time_not_at_finalization(): + from agentx.integrations.openai import patch_openai_client + + class FakeCompletions: + def create(self, **kwargs): + return _FakeStream(_stream_chunks()) + + client = _fake_openai_client(FakeCompletions()) + tracer = make_tracer() + tracer._dispatch = MagicMock(return_value=None) + patch_openai_client(client, tracer, name="gpt-agent") + + # Created with NO active span, consumed, then finalized while an unrelated span is active. + stream = client.chat.completions.create(model="gpt-4o-mini", messages=[], stream=True) + for _ in stream: + pass + with tracer.trace("unrelated-task") as unrelated: + stream.close() + assert unrelated._child_span_count == 0, "the stream must not graft onto the unrelated span" + # It became its own root trace instead (two _send calls: the stream's root + unrelated-task). + assert tracer._send.call_count == 2 + names = [c.kwargs["name"] for c in tracer._send.call_args_list] + assert "gpt-agent" in names and "unrelated-task" in names + + +def test_stream_proxy_private_attributes_never_delegate(): + from agentx.integrations._traced_call import TracedStream + + proxy = TracedStream.__new__(TracedStream) # half-constructed: no _stream yet + with pytest.raises(AttributeError): + _ = proxy._done + del proxy # __del__ on the half-built object must not recurse or raise + + +def test_stream_keyboard_interrupt_records_partial_output_without_an_error(): + from agentx.integrations.openai import patch_openai_client + + class InterruptingStream(_FakeStream): + def __iter__(self): + yield _chunk(content="Hel") + raise KeyboardInterrupt() + + class FakeCompletions: + def create(self, **kwargs): + return InterruptingStream([]) + + client = _fake_openai_client(FakeCompletions()) + tracer = make_tracer() + patch_openai_client(client, tracer, name="gpt-agent") + with pytest.raises(KeyboardInterrupt): + list(client.chat.completions.create(model="gpt-4o-mini", messages=[], stream=True)) + _, kwargs = tracer._send.call_args + assert kwargs["output"] == "Hel" + assert kwargs["error"] is None + + +def test_stream_without_a_first_chunk_omits_time_to_first_token(): + from agentx.integrations.openai import patch_openai_client + + class DeadStream(_FakeStream): + def __iter__(self): + raise RuntimeError("connection refused before the first chunk") + yield # noqa: unreachable - makes this a generator like the real Stream + + class FakeCompletions: + def create(self, **kwargs): + return DeadStream([]) + + client = _fake_openai_client(FakeCompletions()) + tracer = make_tracer() + patch_openai_client(client, tracer, name="gpt-agent") + with pytest.raises(RuntimeError): + list(client.chat.completions.create(model="gpt-4o-mini", messages=[], stream=True)) + _, kwargs = tracer._send.call_args + assert kwargs["metadata"]["streaming"] is True + assert "timeToFirstTokenMs" not in kwargs["metadata"] + + def test_abandoned_stream_still_records_what_it_saw_when_collected(): import gc from agentx.integrations.openai import patch_openai_client