root latency ends at the last chunk, tracing failures never escape in…

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

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` |

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:

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)

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).

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(),

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",