agent-observability-trace-cli 0.1.2__py3-none-any.whl

agent_trace/interceptor/logging_hook.py

@@ -0,0 +1,137 @@
+"""
+Logging-capture layer for non-HTTP, non-callback runtime warnings.
+
+Some real, reproducible agent failures never raise an exception and never
+touch the network — e.g. LangGraph's pregel scheduler logging ``"wrote to
+unknown channel X, ignoring it"`` when ``push_ui_message`` targets an
+unregistered state channel (langgraph#5464). Neither the HTTP interceptor
+(``agent_trace.interceptor.httpx_hook``/``requests_patch``) nor
+``LangGraphTracer``'s callback handlers
+(``agent_trace.integrations.langgraph``) can ever see this class of
+failure — nothing in agent-trace previously attached to Python's own
+``logging`` module at all.
+
+Usage::
+
+    from agent_trace import tracer
+    from agent_trace.interceptor.logging_hook import capture_logging
+
+    with tracer.start_trace("my_graph") as trace:
+        with capture_logging(tracer, logger_names=["langgraph"]):
+            graph.invoke(state, config={"callbacks": [...]})
+
+Every ``WARNING``-or-above record emitted by any of the named loggers while
+the ``with`` block is active is persisted as a ``runtime_log`` event on a
+dedicated ``logging:capture`` span — independent of, and in addition to,
+whatever HTTP/callback capture is also active.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import logging
+from collections.abc import Generator, Sequence
+from typing import TYPE_CHECKING
+
+from agent_trace.core.span import Span, SpanStatus
+
+if TYPE_CHECKING:
+    from agent_trace import Tracer
+
+__all__ = ["capture_logging"]
+
+logger = logging.getLogger(__name__)
+
+# Bounds — this is diagnostic capture, not a full log mirror. A pathological
+# hot loop logging thousands of warnings must not grow a single span's event
+# list unboundedly; logging.captured_count keeps counting past the cap.
+_MAX_ATTR_LEN = 4_000
+_MAX_LOG_EVENTS = 200
+
+# Loggers known to emit real, reproducible non-exception runtime warnings
+# relevant to agent failures — e.g. langgraph#5464's pregel-scheduler
+# "wrote to unknown channel" line. Callers can pass their own logger_names
+# to widen or narrow this.
+_DEFAULT_LOGGER_NAMES: tuple[str, ...] = ("langgraph", "langchain_core")
+
+
+class _SpanLoggingHandler(logging.Handler):
+    """logging.Handler that persists each record as a SpanEvent on *span*.
+
+    Formatting/attribute-extraction failures on a single record must never
+    break the caller's real logging pipeline — every step is best-effort.
+    """
+
+    def __init__(self, span: Span, level: int = logging.WARNING) -> None:
+        super().__init__(level=level)
+        self._span = span
+        self.captured_count = 0
+
+    def emit(self, record: logging.LogRecord) -> None:
+        self.captured_count += 1
+        if self.captured_count > _MAX_LOG_EVENTS:
+            return
+        try:
+            message = self.format(record)
+        except Exception:
+            message = record.getMessage()
+        try:
+            self._span.add_event(
+                "runtime_log",
+                attributes={
+                    "log.level": record.levelname,
+                    "log.logger": record.name,
+                    "log.message": str(message)[:_MAX_ATTR_LEN],
+                    "log.filename": str(record.pathname),
+                    "log.lineno": int(record.lineno),
+                },
+            )
+        except Exception:
+            logger.debug(
+                "agent-trace: failed to record captured log event",
+                exc_info=True,
+            )
+
+
+@contextlib.contextmanager
+def capture_logging(
+    tracer: Tracer,
+    *,
+    logger_names: Sequence[str] = _DEFAULT_LOGGER_NAMES,
+    level: int = logging.WARNING,
+    span_name: str = "logging:capture",
+) -> Generator[Span, None, None]:
+    """Attach a logging.Handler to *logger_names* for the lifetime of the
+    ``with`` block, persisting each captured record onto a dedicated
+    ``logging:capture`` span (opened via *tracer*).
+
+    *logger_names* defaults to LangGraph/LangChain's own logger namespaces
+    — the loggers a pregel-scheduler warning like langgraph#5464's actually
+    goes through. Pass an explicit list to widen/narrow the set.
+
+    The span closes ``OK`` when the block exits normally (log capture is a
+    diagnostic side-channel, not a success/failure signal in itself) or
+    ``ERROR`` if the block itself raises (the exception is recorded onto the
+    span, then re-raised unchanged).
+    """
+    span = tracer.start_span(span_name)
+    handler = _SpanLoggingHandler(span, level=level)
+    attached: list[logging.Logger] = []
+    try:
+        for name in logger_names:
+            target = logging.getLogger(name)
+            target.addHandler(handler)
+            attached.append(target)
+        yield span
+    except Exception as exc:
+        span.record_exception(exc)
+        if span.end_time is None:
+            span.end(SpanStatus.ERROR)
+        raise
+    else:
+        span.set_attribute("logging.captured_count", handler.captured_count)
+        if span.end_time is None:
+            span.end(SpanStatus.OK)
+    finally:
+        for target in attached:
+            target.removeHandler(handler)

agent_trace/interceptor/requests_patch.py

@@ -0,0 +1,184 @@
+"""
+requests adapters for recording and replaying HTTP exchanges.
+
+These are the requests-library equivalents of httpx_hook.py's transports.
+RecordingAdapter wraps the real HTTPAdapter and saves each exchange.
+ReplayAdapter serves responses from the fixture without touching the network.
+
+Why two separate modules (httpx_hook + requests_patch)?
+Many AI SDKs use httpx (e.g. the Anthropic Python SDK); others use requests
+(e.g. OpenAI's legacy client).  Supporting both lets agent-trace intercept
+the full HTTP layer regardless of which SDK the user chooses.
+"""
+
+from __future__ import annotations
+
+import io
+import logging
+import time
+import warnings
+from typing import TYPE_CHECKING, Any
+
+from requests import PreparedRequest, Response
+from requests.adapters import BaseAdapter, HTTPAdapter
+
+if TYPE_CHECKING:
+    from agent_trace._replay.fixture import Fixture
+
+from agent_trace.core.exceptions import NetworkGuardError, guard_active
+
+__all__ = [
+    "NetworkGuardError",
+    "RecordingAdapter",
+    "ReplayAdapter",
+]
+
+logger = logging.getLogger(__name__)
+
+
+class RecordingAdapter(HTTPAdapter):
+    """requests HTTPAdapter that persists each exchange to a Fixture.
+
+    Mount this adapter on a requests.Session before making calls:
+
+        session = requests.Session()
+        session.mount("https://", RecordingAdapter(fixture))
+
+    Pass *inner* to wrap an existing adapter (preserves connection pooling
+    and custom adapters) instead of creating a fresh HTTPAdapter pool.
+    """
+
+    def __init__(
+        self,
+        fixture: Fixture,
+        inner: BaseAdapter | None = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self._fixture = fixture
+        self._inner = inner
+
+    def send(
+        self,
+        request: PreparedRequest,
+        *args: Any,
+        **kwargs: Any,
+    ) -> Response:
+        """Send the request, record the exchange, return the response.
+
+        A pre-response exception (connection refused, DNS failure, TLS
+        failure, ...) is recorded too — as a failed-before-response exchange
+        (error_type/error_message, no response_status) — instead of being
+        silently lost, then re-raised unchanged so the caller sees the exact
+        same failure it would without recording active.
+        """
+        url = str(request.url or "")
+        method = str(request.method or "GET").upper()
+        req_headers = dict(request.headers or {})
+        # PreparedRequest.body can be bytes, str, or None.
+        body = request.body
+        if isinstance(body, bytes):
+            req_body = body.decode("utf-8", errors="replace")
+        elif isinstance(body, str):
+            req_body = body
+        else:
+            req_body = ""
+
+        start = time.monotonic()
+        try:
+            if self._inner is not None:
+                response: Response = self._inner.send(request, *args, **kwargs)
+            else:
+                response = super().send(request, *args, **kwargs)
+        except Exception as exc:
+            duration_ms = (time.monotonic() - start) * 1000
+            self._fixture.record_exchange(
+                url=url,
+                method=method,
+                request_headers=req_headers,
+                request_body=req_body,
+                duration_ms=duration_ms,
+                error_type=type(exc).__qualname__,
+                error_message=str(exc),
+            )
+            raise
+        duration_ms = (time.monotonic() - start) * 1000
+
+        resp_headers = dict(response.headers)
+        resp_body = response.text  # reads and caches the body
+
+        self._fixture.record_exchange(
+            url=url,
+            method=method,
+            request_headers=req_headers,
+            request_body=req_body,
+            response_status=response.status_code,
+            response_headers=resp_headers,
+            response_body=resp_body,
+            duration_ms=duration_ms,
+        )
+
+        return response
+
+
+class ReplayAdapter(BaseAdapter):
+    """requests adapter that serves responses from a Fixture without network I/O.
+
+    Mount this adapter on a requests.Session to intercept all outbound calls:
+
+        session = requests.Session()
+        session.mount("https://", ReplayAdapter(fixture))
+        session.mount("http://", ReplayAdapter(fixture))
+    """
+
+    def __init__(self, fixture: Fixture) -> None:
+        super().__init__()
+        self._fixture = fixture
+
+    def close(self) -> None:
+        pass
+
+    def send(
+        self,
+        request: PreparedRequest,
+        *args: Any,
+        **kwargs: Any,
+    ) -> Response:
+        """Return the next recorded response for this request without I/O."""
+        url = str(request.url or "")
+        method = str(request.method or "GET").upper()
+
+        exchange: dict[str, Any] | None = self._fixture.next_exchange(url, method)
+
+        if exchange is None:
+            if guard_active():
+                raise NetworkGuardError(
+                    f"No recorded exchange for {method} {url} and "
+                    "AGENT_TRACE_NETWORK_GUARD=1 is set.  "
+                    "Run in recording mode first to capture this request."
+                )
+            warnings.warn(
+                f"agent-trace: no fixture entry for {method} {url}; "
+                "falling through to live network.  Set AGENT_TRACE_NETWORK_GUARD=1 "
+                "to make this an error.",
+                stacklevel=2,
+            )
+            fallback = HTTPAdapter()
+            try:
+                return fallback.send(request, *args, **kwargs)
+            finally:
+                fallback.close()
+
+        response = Response()
+        response.status_code = int(exchange["response_status"])
+        response.headers.update(exchange["response_headers"])
+        # _content must be set as bytes so that response.text works correctly.
+        content = exchange["response_body"].encode("utf-8")
+        response._content = content
+        response.encoding = "utf-8"
+        response.url = url
+        response.request = request
+        # Wrap bytes in a BytesIO so raw.read() works if callers check it.
+        response.raw = io.BytesIO(content)
+
+        return response

agent_trace/interceptor/sse.py

@@ -0,0 +1,176 @@
+"""
+SSE-aware parsing of streamed HTTP response bodies stored in a Fixture.
+
+RecordingTransport/AsyncRecordingTransport already capture the full raw
+Server-Sent-Events text off the wire — framework-agnostic, regardless of
+provider (OpenAI/Azure/Groq-style ``data: {...}`` chunks, Anthropic's
+``event: ...`` + ``data: ...`` pairs, ...). What's missing is turning that
+opaque text into diffable, addressable per-chunk data: today a developer has
+to hand-parse ``response_body`` themselves to see what each streamed delta
+actually contained.
+
+This module is deliberately read-only and operates purely on an already
+-recorded ``response_body`` string (or a fixture exchange dict) — it does not
+change how RecordingTransport captures bytes off the wire. Use it against
+``Fixture.all_exchanges()``/``Fixture.next_exchange()`` output, e.g.:
+
+    from agent_trace.interceptor.sse import (
+        is_sse_exchange,
+        parse_sse_events,
+        reconstruct_streamed_message,
+    )
+
+    for exchange in fixture.all_exchanges():
+        if is_sse_exchange(exchange):
+            events = parse_sse_events(exchange["response_body"])
+            merged = reconstruct_streamed_message(events)
+"""
+
+from __future__ import annotations
+
+import json
+from typing import Any
+
+__all__ = [
+    "is_sse_exchange",
+    "parse_sse_events",
+    "reconstruct_streamed_message",
+]
+
+# Sentinel value some providers (OpenAI-style) send as the final SSE data
+# line to signal "no more chunks" — not JSON, must be special-cased rather
+# than fed to json.loads.
+_DONE_SENTINEL = "[DONE]"
+
+
+def is_sse_exchange(exchange: dict[str, Any]) -> bool:
+    """True if *exchange* (as returned by Fixture.all_exchanges()/
+    next_exchange()) looks like a captured Server-Sent-Events response.
+
+    Checks the recorded ``Content-Type`` response header first (the
+    authoritative signal); falls back to sniffing the body for an SSE
+    ``data: `` line when headers are missing/lowercased differently than
+    expected, since some proxies/test fixtures don't preserve casing.
+    """
+    headers = exchange.get("response_headers") or {}
+    for key, value in headers.items():
+        if str(key).lower() != "content-type":
+            continue
+        if "text/event-stream" in str(value).lower():
+            return True
+    body = exchange.get("response_body") or ""
+    return _looks_like_sse_body(body)
+
+
+def _looks_like_sse_body(body: str) -> bool:
+    """Best-effort sniff: at least one line starts with the SSE ``data: ``
+    field prefix. Deliberately conservative — a false negative just means a
+    caller has to fall back to treating the body as an opaque string, which
+    is the pre-existing behavior anyway."""
+    for line in body.splitlines():
+        if line.startswith("data:"):
+            return True
+    return False
+
+
+def parse_sse_events(body: str) -> list[dict[str, Any] | str]:
+    """Split a raw SSE response body on ``data: `` boundaries and parse each
+    event's payload.
+
+    Returns an ordered list — one entry per ``data:`` line found, in the
+    order they appear in *body* (the order they were streamed). Each entry
+    is either:
+
+    - a parsed JSON object, for the common ``data: {...}`` shape every major
+      provider's streaming chat-completions endpoint uses, or
+    - the raw string payload, when the line isn't valid JSON (e.g. the
+      ``[DONE]`` sentinel, or a provider using a non-JSON SSE payload).
+
+    A blank body, or a body with no ``data:`` lines at all, returns ``[]`` —
+    callers should treat that as "not an SSE payload" rather than an error.
+    """
+    events: list[dict[str, Any] | str] = []
+    for line in body.splitlines():
+        if not line.startswith("data:"):
+            continue
+        payload = line[len("data:") :].strip()
+        if not payload or payload == _DONE_SENTINEL:
+            continue
+        try:
+            events.append(json.loads(payload))
+        except (json.JSONDecodeError, TypeError):
+            events.append(payload)
+    return events
+
+
+def _delta_from_event(event: dict[str, Any] | str) -> dict[str, Any] | None:
+    """Extract the OpenAI-style ``choices[0].delta`` dict from one parsed SSE
+    event, or None if this event isn't that shape (a non-dict event, a
+    provider using a different streaming schema, etc.)."""
+    if not isinstance(event, dict):
+        return None
+    choices = event.get("choices")
+    if not choices or not isinstance(choices, list):
+        return None
+    first = choices[0]
+    if not isinstance(first, dict):
+        return None
+    delta = first.get("delta")
+    return delta if isinstance(delta, dict) else None
+
+
+def reconstruct_streamed_message(
+    events: list[dict[str, Any] | str],
+) -> dict[str, Any]:
+    """Merge an ordered list of parsed SSE events (from parse_sse_events())
+    into one addressable, reassembled message.
+
+    Concatenates every ``delta.content`` fragment into ``content``, and
+    merges every ``delta.tool_calls[]`` fragment — keyed by its ``index``
+    field, the same index OpenAI-style providers use to identify which
+    parallel tool call a given argument fragment belongs to — into
+    ``tool_calls``, concatenating each tool call's ``function.arguments``
+    string across chunks in the order they arrived.
+
+    Returns ``{"content": str, "tool_calls": {index: {...}}}``. Events that
+    don't match the OpenAI-style ``choices[0].delta`` shape are ignored
+    (not raised on) — this reconstruction targets that one well-known wire
+    shape; a caller working with a different provider's streaming schema
+    should read ``events`` directly instead.
+    """
+    content_parts: list[str] = []
+    tool_calls: dict[int, dict[str, Any]] = {}
+
+    for event in events:
+        delta = _delta_from_event(event)
+        if delta is None:
+            continue
+
+        content_piece = delta.get("content")
+        if isinstance(content_piece, str):
+            content_parts.append(content_piece)
+
+        for tc_fragment in delta.get("tool_calls") or []:
+            if not isinstance(tc_fragment, dict):
+                continue
+            index = tc_fragment.get("index", 0)
+            try:
+                index = int(index)
+            except (TypeError, ValueError):
+                index = 0
+            entry = tool_calls.setdefault(
+                index,
+                {"id": None, "function": {"name": "", "arguments": ""}},
+            )
+            if tc_fragment.get("id"):
+                entry["id"] = tc_fragment["id"]
+            fn_fragment = tc_fragment.get("function") or {}
+            if fn_fragment.get("name"):
+                entry["function"]["name"] += fn_fragment["name"]
+            if fn_fragment.get("arguments"):
+                entry["function"]["arguments"] += fn_fragment["arguments"]
+
+    return {
+        "content": "".join(content_parts),
+        "tool_calls": tool_calls,
+    }