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

agent_trace/integrations/streaming.py

@@ -0,0 +1,218 @@
+"""
+Outbound event-stream/SSE (or websocket) delivery instrumentation.
+
+agent-trace's LangGraph integration only hooks ``langchain_core`` callbacks
+that fire *in-process* during graph execution — the same layer LangSmith
+already instruments. Failures where the backend graph completes
+successfully but the frontend never receives the terminal event (e.g.
+dropped/stalled SSE connections, a silent exception raised inside a
+streaming response writer, a websocket send that raises after the client
+disconnected) happen entirely *downstream* of that layer — langgraph#6202's
+failure class. There was previously zero code anywhere in agent-trace that
+touched this transport boundary (confirmed via repo-wide grep for
+``stream|sse|websocket|asgi|fastapi`` prior to this module).
+
+This module instruments the layer that actually carries chunks out to a
+browser client, independent of the in-process LangGraph/LangChain callback
+layer — a thin, framework-agnostic wrapper around whatever iterable/
+async-iterable/send-callable a server framework hands its transport (a
+Starlette/FastAPI ``StreamingResponse`` body iterator, the generator behind
+LangGraph's own ``astream_events``, or a websocket's ``.send()`` method).
+No hard dependency on any specific web framework — it only needs a
+generator/async-generator or an async callable.
+
+Usage (SSE via a Starlette/FastAPI ``StreamingResponse``)::
+
+    from agent_trace.integrations.streaming import traced_sse_astream
+
+    async def event_source():
+        async for chunk in graph.astream_events(state, version="v2"):
+            yield f"data: {json.dumps(chunk)}\\n\\n"
+
+    return StreamingResponse(
+        traced_sse_astream(tracer, event_source()),
+        media_type="text/event-stream",
+    )
+
+Usage (websocket)::
+
+    from agent_trace.integrations.streaming import traced_websocket_send
+
+    send = traced_websocket_send(tracer, websocket.send_text)
+    await send(json.dumps(chunk))
+"""
+
+from __future__ import annotations
+
+from collections.abc import (
+    AsyncGenerator,
+    AsyncIterable,
+    Awaitable,
+    Callable,
+    Generator,
+    Iterable,
+)
+from typing import TYPE_CHECKING, Any, TypeVar
+
+from agent_trace.core.span import Span, SpanStatus
+
+if TYPE_CHECKING:
+    from agent_trace import Tracer
+
+__all__ = [
+    "traced_sse_astream",
+    "traced_sse_stream",
+    "traced_websocket_send",
+]
+
+# Bounds — this records enough of each outbound chunk to diagnose a
+# dropped/stalled transport, not a full-fidelity mirror of every byte sent.
+_MAX_CHUNK_PREVIEW_LEN = 500
+_MAX_TRANSPORT_EVENTS = 500
+
+T = TypeVar("T")
+
+
+def _stringify_chunk(item: Any) -> str:
+    """Best-effort, bounded preview of an outbound chunk for a SpanEvent
+    attribute. Never raises — falls back to a placeholder."""
+    try:
+        if isinstance(item, bytes):
+            text = item.decode("utf-8", errors="replace")
+        else:
+            text = str(item)
+    except Exception:
+        return "<unrepresentable-chunk>"
+    if len(text) > _MAX_CHUNK_PREVIEW_LEN:
+        text = text[:_MAX_CHUNK_PREVIEW_LEN] + "...<truncated>"
+    return text
+
+
+def _record_transport_chunk_sent(span: Span, index: int, item: Any) -> None:
+    """Append one bounded transport_chunk_sent SpanEvent, capped at
+    _MAX_TRANSPORT_EVENTS so a very long-lived stream can't grow a single
+    span's event list without limit. transport.chunk_count (set at span
+    close) still reflects every chunk that was actually sent."""
+    if index >= _MAX_TRANSPORT_EVENTS:
+        return
+    span.add_event(
+        "transport_chunk_sent",
+        attributes={
+            "transport.index": index,
+            "transport.chunk_preview": _stringify_chunk(item),
+        },
+    )
+
+
+def traced_sse_stream(
+    tracer: Tracer,
+    source: Iterable[T],
+    *,
+    span_name: str = "transport:sse",
+) -> Generator[T, None, None]:
+    """Wrap a synchronous outbound SSE/streaming-response iterable (e.g. the
+    body iterator handed to a WSGI/Starlette ``StreamingResponse``) so each
+    chunk's actual send moment — and a bounded preview of its content — lands
+    on the trace timeline, independent of the in-process LangGraph callback
+    layer.
+
+    Opens one ``transport:sse`` span (customizable via *span_name*) that
+    stays open for the lifetime of iteration, closing ``OK`` once the
+    source is exhausted (every chunk successfully handed off to the
+    transport), ``ERROR`` if the source itself raises (e.g. a broken pipe
+    surfaced back into application code) with the exception recorded onto
+    the span then re-raised unchanged, or ``CANCELLED`` if the *caller*
+    (the ASGI/WSGI server) stops iterating early — the exact shape of a
+    dropped/stalled client connection — detected via ``GeneratorExit`` when
+    this generator is closed/garbage-collected before exhaustion.
+    """
+    span = tracer.start_span(span_name)
+    index = 0
+    status = SpanStatus.OK
+    try:
+        for item in source:
+            _record_transport_chunk_sent(span, index, item)
+            index += 1
+            yield item
+    except GeneratorExit:
+        status = SpanStatus.CANCELLED
+        raise
+    except Exception as exc:
+        status = SpanStatus.ERROR
+        span.record_exception(exc)
+        raise
+    finally:
+        span.set_attribute("transport.chunk_count", index)
+        if span.end_time is None:
+            span.end(status)
+
+
+async def traced_sse_astream(
+    tracer: Tracer,
+    source: AsyncIterable[T],
+    *,
+    span_name: str = "transport:sse",
+) -> AsyncGenerator[T, None]:
+    """Async equivalent of :func:`traced_sse_stream` — wraps an async
+    outbound SSE generator (e.g. the body passed to FastAPI's
+    ``StreamingResponse``, or the generator behind ``graph.astream_events``)
+    the same way."""
+    span = tracer.start_span(span_name)
+    index = 0
+    status = SpanStatus.OK
+    try:
+        async for item in source:
+            _record_transport_chunk_sent(span, index, item)
+            index += 1
+            yield item
+    except GeneratorExit:
+        status = SpanStatus.CANCELLED
+        raise
+    except Exception as exc:
+        status = SpanStatus.ERROR
+        span.record_exception(exc)
+        raise
+    finally:
+        span.set_attribute("transport.chunk_count", index)
+        if span.end_time is None:
+            span.end(status)
+
+
+def traced_websocket_send(
+    tracer: Tracer,
+    send_fn: Callable[[T], Awaitable[Any]],
+    *,
+    span_name: str = "transport:websocket",
+) -> Callable[[T], Awaitable[Any]]:
+    """Wrap an async websocket ``send``-shaped callable (e.g.
+    ``websocket.send_text``/``send_json``/``send_bytes``) so every outbound
+    send is recorded with a timestamp and success/failure status,
+    independent of the in-process LangGraph callback layer.
+
+    Returns a new async callable with the same signature as *send_fn*.
+    Opens one dedicated span per call (rather than one long-lived span for
+    the whole connection, since a websocket connection's lifetime is not
+    naturally scoped the way a single HTTP streaming response's is) —
+    closing ``OK`` immediately after a successful send, or ``ERROR`` (with
+    the exception recorded, then re-raised unchanged) if the send itself
+    raises — e.g. because the client already disconnected.
+    """
+
+    async def _wrapped(item: T) -> Any:
+        span = tracer.start_span(span_name)
+        # _stringify_chunk() is itself defensive (never raises) and
+        # set_attribute() on a plain string cannot raise either.
+        span.set_attribute("transport.chunk_preview", _stringify_chunk(item))
+        try:
+            result = await send_fn(item)
+        except Exception as exc:
+            span.record_exception(exc)
+            if span.end_time is None:
+                span.end(SpanStatus.ERROR)
+            raise
+        else:
+            if span.end_time is None:
+                span.end(SpanStatus.OK)
+            return result
+
+    return _wrapped

agent_trace/interceptor/__init__.py

@@ -0,0 +1,64 @@
+"""
+agent-trace transport interceptors.
+
+Each submodule patches one HTTP/RPC/WS client library (httpx, requests,
+aiohttp, botocore/boto3, grpc, websockets, MCP stdio, ...) to record real
+traffic into a ``Fixture`` and/or replay it offline — see
+``httpx_hook.py``'s module docstring for the canonical Recording*/Replay*
+Transport shape every other interceptor here mirrors.
+
+Submodules are exported lazily here (PEP 562 module ``__getattr__``) so that
+``import agent_trace.interceptor`` never fails regardless of which optional
+client libraries happen to be installed. Most submodules (httpx_hook.py,
+aiohttp_hook.py, botocore_hook.py, websocket_hook.py, stdio_hook.py) already
+defer their own ``import <library>`` to call time and are always safely
+importable; ``grpc_hook.py`` and ``requests_patch.py`` import their target
+library eagerly (``grpc``/``requests``) since ``Tracer._patch_grpc`` /
+``Tracer._patch_requests`` in ``agent_trace/__init__.py`` already wrap the
+*submodule* import itself in ``try/except ImportError`` — lazy resolution
+here means accessing e.g. ``agent_trace.interceptor.grpc_hook`` only fails
+at that point (with the normal ``ImportError``), not merely by importing
+this package.
+"""
+
+from __future__ import annotations
+
+import importlib
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from agent_trace.interceptor import (
+        aiohttp_hook,
+        botocore_hook,
+        grpc_hook,
+        httpx_hook,
+        logging_hook,
+        requests_patch,
+        sse,
+        stdio_hook,
+        warnings_hook,
+        websocket_hook,
+    )
+
+__all__ = [
+    "aiohttp_hook",
+    "botocore_hook",
+    "grpc_hook",
+    "httpx_hook",
+    "logging_hook",
+    "requests_patch",
+    "sse",
+    "stdio_hook",
+    "warnings_hook",
+    "websocket_hook",
+]
+
+
+def __getattr__(name: str) -> Any:
+    if name in __all__:
+        return importlib.import_module(f"agent_trace.interceptor.{name}")
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+def __dir__() -> list[str]:
+    return sorted(__all__)

agent_trace/interceptor/aiohttp_hook.py

@@ -0,0 +1,152 @@
+"""
+aiohttp interceptor for recording HTTP exchanges made via aiohttp.ClientSession.
+
+Why a third interceptor (alongside httpx_hook.py / requests_patch.py)?
+LiteLLM — the standard route agent frameworks (crewAI's fallback path,
+LangChain's non-native providers, etc.) use to reach non-OpenAI/Anthropic
+backends such as Gemini/Vertex AI — defaults to routing outbound async calls
+through an aiohttp-based transport (`_should_use_aiohttp_transport()` returns
+True unless explicitly disabled). That traffic never touches
+`httpx.Client`/`httpx.AsyncClient` or `requests.Session` at all, so
+agent-trace's two existing interceptors silently miss it with no error or
+warning.
+
+`aiohttp.ClientSession` has no pluggable "transport" object the way httpx
+does, so recording is implemented by wrapping `ClientSession._request` — the
+single coroutine every request-verb helper (`get`/`post`/`request`/...)
+funnels through, and the same method the `aioresponses` mocking library
+itself patches for exactly this reason (see `patch("aiohttp.client.
+ClientSession._request", ...)` in aioresponses' own source).
+
+Deliberately NOT a `ClientSession` subclass: the installed aiohttp (3.12+)
+raises `DeprecationWarning: Inheritance class ... from ClientSession is
+discouraged` from `ClientSession.__init_subclass__` for any subclass,
+confirmed by direct introspection of the installed package. `make_recording_
+request` instead builds a plain replacement function for `ClientSession.
+_request` that `Tracer._patch_aiohttp` assigns onto the class directly —
+the same class-level monkey-patch pattern already used for `httpx.Client.
+__init__`/`httpx.AsyncClient.__init__` and `requests.Session.get_adapter`.
+
+The replacement wraps the real (or previously patched) `_request()` call,
+eagerly reads the response body (aiohttp caches it on
+`ClientResponse._body`, so callers can still call `.read()`/`.text()`/
+`.json()` afterwards without re-consuming the stream), and persists the
+exchange to the fixture before returning the original response object
+unmodified.
+
+Only a recording path is provided here (no replay counterpart): this
+interceptor's job is to stop agent-trace from *silently missing* aiohttp
+traffic during recording. Offline replay of aiohttp-captured fixtures is a
+separate, larger piece of work (constructing a faithful fake
+`aiohttp.ClientResponse` without a live connection) and is not required to
+close the gap this interceptor targets.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable
+
+    import aiohttp
+    from yarl import URL
+
+    from agent_trace._replay.fixture import Fixture
+
+__all__ = [
+    "make_recording_request",
+]
+
+logger = logging.getLogger(__name__)
+
+
+def _serialize_request_body(kwargs: dict[str, Any]) -> str:
+    """Best-effort reconstruction of the body bytes aiohttp will send.
+
+    Mirrors RecordingAdapter's (requests_patch.py) bytes/str/else fallback:
+    `json=` is serialized the same way `json.dumps` would represent it (the
+    exact separators aiohttp's internal `json_serialize` uses can differ,
+    but the semantic content — e.g. the model/messages payload a developer
+    needs to root-cause a bug — is preserved); `data=` is decoded when it's
+    already bytes/str (the common case for SDKs like LiteLLM's aiohttp
+    transport, which hand aiohttp pre-serialized JSON bytes); anything else
+    (FormData, an IO stream, a dict passed as `data=`) falls back to `str()`
+    rather than raising, since raw evidence beats no evidence.
+    """
+    body = kwargs.get("json")
+    if body is not None:
+        try:
+            return json.dumps(body)
+        except (TypeError, ValueError):
+            return str(body)
+
+    data = kwargs.get("data")
+    if data is None:
+        return ""
+    if isinstance(data, bytes):
+        return data.decode("utf-8", errors="replace")
+    if isinstance(data, str):
+        return data
+    return str(data)
+
+
+async def _record_exchange(
+    fixture: Fixture,
+    method: str,
+    str_or_url: str | URL,
+    kwargs: dict[str, Any],
+    response: aiohttp.ClientResponse,
+) -> None:
+    """Read *response*'s body and persist the exchange to *fixture*.
+
+    `response.read()` caches the body on `ClientResponse._body`, so calling
+    it here does not prevent the caller from reading the response again via
+    `.read()` / `.text()` / `.json()`.
+    """
+    await response.read()
+
+    req_headers = dict(kwargs.get("headers") or {})
+    req_body = _serialize_request_body(kwargs)
+
+    resp_headers = dict(response.headers)
+    resp_body = await response.text()
+
+    fixture.record_exchange(
+        url=str(str_or_url),
+        method=str(method).upper(),
+        request_headers=req_headers,
+        request_body=req_body,
+        response_status=response.status,
+        response_headers=resp_headers,
+        response_body=resp_body,
+    )
+
+
+def make_recording_request(
+    fixture: Fixture,
+    original_request: Callable[..., Awaitable[aiohttp.ClientResponse]],
+) -> Callable[..., Awaitable[aiohttp.ClientResponse]]:
+    """Build a replacement for `aiohttp.ClientSession._request` that records.
+
+    *original_request* is the real (or previously-patched) `_request`
+    method, called through unconditionally so behaviour — redirects, auth,
+    proxies, etc. — is unchanged; only the body-eager-read + fixture-record
+    step is added around it.  Used by `Tracer._patch_aiohttp` to patch
+    `aiohttp.ClientSession` at the class level so plain `aiohttp.ClientSession()`
+    call sites inside third-party SDKs are captured with zero code changes.
+    """
+
+    async def _patched_request(
+        session_self: aiohttp.ClientSession,
+        method: str,
+        str_or_url: str | URL,
+        **kwargs: Any,
+    ) -> aiohttp.ClientResponse:
+        response = await original_request(session_self, method, str_or_url, **kwargs)
+        await _record_exchange(fixture, method, str_or_url, kwargs, response)
+        return response
+
+    return _patched_request

agent_trace/interceptor/botocore_hook.py

@@ -0,0 +1,223 @@
+"""
+botocore session wrappers for recording and replaying AWS SDK (boto3) HTTP
+exchanges — the botocore equivalent of httpx_hook.py / requests_patch.py.
+
+Every boto3 service client (bedrock-runtime, sagemaker-runtime, s3, ...)
+ultimately routes its outbound HTTP request through a single low-level
+object: ``botocore.httpsession.URLLib3Session``, whose ``send(request)``
+method takes an ``AWSPreparedRequest`` (``.method``/``.url``/``.headers``/
+``.body``) and returns an ``AWSResponse`` (``.status_code``/``.headers``/
+``.content``/``.text``) — the same request-in/response-out shape as
+``requests.HTTPAdapter.send`` and ``httpx.BaseTransport.handle_request``.
+
+RecordingSession wraps a real session-like object (anything exposing
+``.send(request)``): it lets the request go through, reads the full
+response body, and saves the exchange to the fixture before returning the
+response to the caller.
+
+ReplaySession never touches the network: it looks up the next recorded
+response for the requested (method, url) from the fixture, mirroring
+ReplayTransport/ReplayAdapter.
+
+Streaming operations (e.g. Bedrock's ``InvokeModel``, whose response body
+is delivered via ``StreamingBody``, or ``InvokeModelWithResponseStream``/
+``ConverseStream``, whose response is an AWS event stream) set
+``request.stream_output = True``.  For those, botocore does *not* eagerly
+buffer the response at the urllib3 layer, so ``response.content`` is only
+populated the first time something reads it.  RecordingSession still drains
+and records these — mirroring httpx_hook.py's ``response.read()``/
+``response.aread()`` eager-buffering of every exchange, including SSE — but
+afterwards replaces ``response.raw`` with a fresh in-memory buffer over the
+same bytes so downstream body consumption (``StreamingBody.read()``, event
+stream iteration) keeps working transparently.  True non-buffering
+pass-through capture for streaming responses is a known follow-up, not
+yet implemented, exactly as it is for the httpx/SSE case.
+"""
+
+from __future__ import annotations
+
+import io
+import logging
+import warnings
+from typing import TYPE_CHECKING, Any, Protocol
+
+if TYPE_CHECKING:
+    from agent_trace._replay.fixture import Fixture
+
+from agent_trace.core.exceptions import NetworkGuardError, guard_active
+
+__all__ = [
+    "NetworkGuardError",
+    "RecordingSession",
+    "ReplaySession",
+]
+
+logger = logging.getLogger(__name__)
+
+
+class _SendableSession(Protocol):
+    """Structural type for anything exposing botocore's session.send(request)."""
+
+    def send(self, request: Any) -> Any: ...
+
+
+def _stringify_headers(headers: Any) -> dict[str, str]:
+    """Coerce a botocore/urllib3 header mapping to a plain ``dict[str, str]``.
+
+    botocore request headers are commonly a mix of ``bytes`` and ``str``
+    values (e.g. ``Content-Type``/``Authorization`` are bytes, while
+    ``Content-Length`` is str) — confirmed via direct inspection of a real
+    signed ``AWSPreparedRequest``.  ``Fixture.record_exchange`` JSON-encodes
+    the header dict, which fails on bytes values, so normalise everything to
+    str here.
+    """
+    result: dict[str, str] = {}
+    for raw_key, raw_value in dict(headers).items():
+        str_key = (
+            raw_key.decode("utf-8", errors="replace")
+            if isinstance(raw_key, bytes)
+            else str(raw_key)
+        )
+        str_value = (
+            raw_value.decode("utf-8", errors="replace")
+            if isinstance(raw_value, bytes)
+            else str(raw_value)
+        )
+        result[str_key] = str_value
+    return result
+
+
+def _stringify_body(body: Any) -> str:
+    """Coerce a botocore request body (bytes | str | file-like | None) to str."""
+    if isinstance(body, bytes):
+        return body.decode("utf-8", errors="replace")
+    if isinstance(body, str):
+        return body
+    return ""
+
+
+class RecordingSession:
+    """botocore session wrapper that records every exchange to a Fixture.
+
+    Parameters
+    ----------
+    fixture:
+        Open Fixture instance where exchanges will be written.
+    inner:
+        The real session (or session-like object exposing ``.send(request)``)
+        to use for outbound requests.  Defaults to a fresh
+        ``botocore.httpsession.URLLib3Session()`` when None.
+    """
+
+    def __init__(
+        self,
+        fixture: Fixture,
+        inner: _SendableSession | None = None,
+    ) -> None:
+        self._fixture = fixture
+        if inner is not None:
+            self._inner: _SendableSession = inner
+        else:
+            import botocore.httpsession
+
+            self._inner = botocore.httpsession.URLLib3Session()
+
+    def send(self, request: Any) -> Any:
+        """Forward the request, record the exchange, return the response."""
+        response = self._inner.send(request)
+
+        # Force the response body to be fully read so we can persist it.
+        # For non-streaming operations botocore already preloaded this at
+        # the urllib3 layer (see URLLib3Session.send), so this is a no-op
+        # cache hit; for streaming operations (stream_output=True) this is
+        # the first read and drains the live connection.
+        content: bytes = response.content
+
+        if request.stream_output:
+            # StreamingBody / EventStream wrap `response.raw` directly
+            # rather than `.content` (see
+            # botocore.endpoint.convert_to_response_dict).  We just drained
+            # the real socket above to record it, so hand the caller a
+            # fresh in-memory buffer over the same bytes — BytesIO supports
+            # the same incremental .read()/.readinto() interface — instead
+            # of the now-exhausted urllib3 stream.
+            response.raw = io.BytesIO(content)
+
+        self._fixture.record_exchange(
+            url=str(request.url),
+            method=str(request.method),
+            request_headers=_stringify_headers(request.headers),
+            request_body=_stringify_body(request.body),
+            response_status=int(response.status_code),
+            response_headers=_stringify_headers(response.headers),
+            response_body=response.text,
+        )
+
+        return response
+
+    def close(self) -> None:
+        close = getattr(self._inner, "close", None)
+        if close is not None:
+            close()
+
+
+class ReplaySession:
+    """botocore session wrapper that serves responses from a Fixture without
+    network I/O.
+    """
+
+    def __init__(
+        self,
+        fixture: Fixture,
+        clock: Any | None = None,
+    ) -> None:
+        self._fixture = fixture
+        self._clock = clock
+
+    def send(self, request: Any) -> Any:
+        """Return the next recorded response for *(method, url)*."""
+        import botocore.awsrequest
+
+        url = str(request.url)
+        method = str(request.method)
+        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,
+            )
+            import botocore.httpsession
+
+            fallback = botocore.httpsession.URLLib3Session()
+            try:
+                return fallback.send(request)
+            finally:
+                fallback.close()
+
+        if self._clock is not None:
+            self._clock.advance(float(exchange["recorded_at"]))
+
+        content = exchange["response_body"].encode("utf-8")
+        response = botocore.awsrequest.AWSResponse(
+            url,
+            int(exchange["response_status"]),
+            exchange["response_headers"],
+            io.BytesIO(content),
+        )
+        # AWSResponse.content lazily reads `.raw`; pre-seed the cache so
+        # repeated access (ours and the caller's) doesn't require `.raw` to
+        # still be positioned at the start.
+        response._content = content
+        return response
+
+    def close(self) -> None:
+        pass  # No resources to release; fixture lifecycle is managed externally.