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

agent_trace/interceptor/stdio_hook.py

@@ -0,0 +1,245 @@
+"""
+MCP stdio-transport capture layer.
+
+``agent-trace``'s only capture layer historically has been HTTP-level
+(``httpx``/``requests`` — see ``httpx_hook.py``/``requests_patch.py``). MCP's
+``stdio`` transport communicates over subprocess stdin/stdout pipes instead —
+zero HTTP traffic touches the wire, so an MCP-stdio-related failure (e.g. a
+startup/tool-loading crash) is completely invisible to those interceptors,
+regardless of which framework wraps the MCP client.
+
+``recording_stdio_client`` wraps ``mcp.client.stdio.stdio_client`` the same
+way ``RecordingTransport`` wraps a real ``httpx`` transport: the real
+subprocess/streams are used unmodified, every JSON-RPC frame that flows
+through is persisted to a :class:`~agent_trace._replay.fixture.Fixture`
+*before* being handed back to the caller, and the caller (typically an MCP
+``ClientSession``) is none the wiser.
+
+Usage::
+
+    from mcp import ClientSession, StdioServerParameters
+    from agent_trace.interceptor.stdio_hook import recording_stdio_client
+
+    server = StdioServerParameters(command="my-mcp-server", args=[])
+    async with recording_stdio_client(server, fixture) as (read, write):
+        async with ClientSession(read, write) as session:
+            await session.initialize()
+            tools = await session.list_tools()
+"""
+
+from __future__ import annotations
+
+import logging
+import sys
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from typing import TYPE_CHECKING, Any, TextIO
+
+if TYPE_CHECKING:
+    from agent_trace._replay.fixture import Fixture
+
+__all__ = ["recording_stdio_client"]
+
+logger = logging.getLogger(__name__)
+
+_INSTALL_HINT = (
+    "MCP stdio recording requires the mcp package.\n"
+    "Install it with:\n\n"
+    "    pip install mcp\n"
+)
+
+
+def _require_mcp() -> Any:
+    """Lazy import guard — raises a clear error if mcp is absent."""
+    try:
+        import mcp
+
+        return mcp
+    except ImportError as exc:
+        raise ImportError(_INSTALL_HINT) from exc
+
+
+def _classify_frame(root: Any) -> tuple[str, str | None, str | None]:
+    """Classify a ``JSONRPCMessage.root`` union member.
+
+    Returns ``(frame_type, rpc_id, method)`` where ``frame_type`` is one of
+    ``"request"``, ``"notification"``, ``"response"``, ``"error"``, or
+    ``"unknown"``.
+
+    Distinguishing the four ``mcp.types`` JSON-RPC frame shapes:
+
+    - ``JSONRPCRequest``:      has ``method`` *and* ``id``.
+    - ``JSONRPCNotification``: has ``method``, no ``id``.
+    - ``JSONRPCResponse``:     has ``id`` and ``result``, no ``method``.
+    - ``JSONRPCError``:        has ``id`` and ``error``, no ``method``.
+    """
+    rpc_id = str(root.id) if getattr(root, "id", None) is not None else None
+    method = getattr(root, "method", None)
+    if getattr(root, "error", None) is not None:
+        return "error", rpc_id, method
+    if hasattr(root, "result"):
+        return "response", rpc_id, method
+    if method is not None and rpc_id is not None:
+        return "request", rpc_id, method
+    if method is not None:
+        return "notification", rpc_id, method
+    return "unknown", rpc_id, method  # pragma: no cover — defensive fallback
+
+
+class _RecordingReceiveStream:
+    """Tees every ``SessionMessage`` read from *inner* into a Fixture.
+
+    Duck-types anyio's ``ObjectReceiveStream`` interface (``receive``,
+    ``aclose``, ``__aiter__``/``__anext__``, async-context-manager) so it can
+    be handed directly to ``ClientSession`` in place of the real stream.
+    """
+
+    def __init__(
+        self,
+        inner: Any,
+        fixture: Fixture,
+        server_command: str,
+    ) -> None:
+        self._inner = inner
+        self._fixture = fixture
+        self._server_command = server_command
+
+    def __aiter__(self) -> _RecordingReceiveStream:
+        return self
+
+    async def __anext__(self) -> Any:
+        import anyio
+
+        try:
+            return await self.receive()
+        except anyio.EndOfStream:
+            raise StopAsyncIteration from None
+
+    async def receive(self) -> Any:
+        item = await self._inner.receive()
+        self._record(item)
+        return item
+
+    def _record(self, item: Any) -> None:
+        from mcp.shared.message import SessionMessage
+
+        # read_stream carries `SessionMessage | Exception` — malformed frames
+        # the real client failed to parse arrive as bare exceptions; nothing
+        # to record on the wire for those.
+        if not isinstance(item, SessionMessage):
+            return
+        try:
+            root = item.message.root
+            frame_type, rpc_id, method = _classify_frame(root)
+            payload = item.message.model_dump_json(by_alias=True, exclude_none=True)
+            self._fixture.record_mcp_frame(
+                server_command=self._server_command,
+                direction="from_server",
+                frame_type=frame_type,
+                rpc_id=rpc_id,
+                method=method,
+                payload=payload,
+            )
+        except Exception:
+            logger.debug(
+                "agent-trace: failed to record inbound MCP frame", exc_info=True
+            )
+
+    async def aclose(self) -> None:
+        await self._inner.aclose()
+
+    async def __aenter__(self) -> _RecordingReceiveStream:
+        return self
+
+    async def __aexit__(self, *exc_info: Any) -> None:
+        await self.aclose()
+
+
+class _RecordingSendStream:
+    """Tees every ``SessionMessage`` sent through *inner* into a Fixture."""
+
+    def __init__(
+        self,
+        inner: Any,
+        fixture: Fixture,
+        server_command: str,
+    ) -> None:
+        self._inner = inner
+        self._fixture = fixture
+        self._server_command = server_command
+
+    async def send(self, item: Any) -> None:
+        self._record(item)
+        await self._inner.send(item)
+
+    def _record(self, item: Any) -> None:
+        from mcp.shared.message import SessionMessage
+
+        if not isinstance(item, SessionMessage):
+            return  # pragma: no cover — write_stream is SessionMessage-only
+        try:
+            root = item.message.root
+            frame_type, rpc_id, method = _classify_frame(root)
+            payload = item.message.model_dump_json(by_alias=True, exclude_none=True)
+            self._fixture.record_mcp_frame(
+                server_command=self._server_command,
+                direction="to_server",
+                frame_type=frame_type,
+                rpc_id=rpc_id,
+                method=method,
+                payload=payload,
+            )
+        except Exception:
+            logger.debug(
+                "agent-trace: failed to record outbound MCP frame", exc_info=True
+            )
+
+    async def aclose(self) -> None:
+        await self._inner.aclose()
+
+    async def __aenter__(self) -> _RecordingSendStream:
+        return self
+
+    async def __aexit__(self, *exc_info: Any) -> None:
+        await self.aclose()
+
+
+@asynccontextmanager
+async def recording_stdio_client(
+    server: Any,
+    fixture: Fixture,
+    errlog: TextIO | None = None,
+) -> AsyncIterator[tuple[_RecordingReceiveStream, _RecordingSendStream]]:
+    """Wrap ``mcp.client.stdio.stdio_client`` to record every JSON-RPC frame.
+
+    Parameters
+    ----------
+    server:
+        An ``mcp.StdioServerParameters`` instance describing the subprocess
+        to launch (``command``, ``args``, ``env``, ``cwd``).
+    fixture:
+        Open :class:`~agent_trace._replay.fixture.Fixture` to record into.
+    errlog:
+        Passed through to ``stdio_client`` for the child process's stderr
+        (defaults to ``sys.stderr``, matching the real client's default).
+
+    Yields
+    ------
+    tuple[_RecordingReceiveStream, _RecordingSendStream]
+        Drop-in replacements for the ``(read_stream, write_stream)`` pair
+        ``stdio_client`` normally yields — pass these straight into
+        ``mcp.ClientSession``.
+    """
+    _require_mcp()
+    from mcp.client.stdio import stdio_client
+
+    command = f"{server.command} {' '.join(server.args)}".strip()
+    effective_errlog: TextIO = errlog if errlog is not None else sys.stderr
+
+    async with stdio_client(server, errlog=effective_errlog) as (
+        read_stream,
+        write_stream,
+    ):
+        recording_read = _RecordingReceiveStream(read_stream, fixture, command)
+        recording_write = _RecordingSendStream(write_stream, fixture, command)
+        yield recording_read, recording_write

agent_trace/interceptor/warnings_hook.py

@@ -0,0 +1,132 @@
+"""
+Python `warnings`-capture layer for library-raised runtime warnings.
+
+Some real, reproducible failures never raise an exception and never touch
+the network at all — e.g. langgraph#5628's ``RuntimeWarning: Failed to trim
+messages to fit within max_tokens limit before summarization``, raised
+entirely inside local token-counting logic before any HTTP call is
+attempted. Nothing in agent-trace previously captured arbitrary
+library-raised ``warnings.warn()`` calls (the only prior ``warnings.warn()``
+call sites in the codebase were agent-trace's own internal warnings in
+``httpx_hook.py``/``requests_patch.py``) — this is a related but distinct
+gap from ``agent_trace.interceptor.logging_hook`` (that one targets
+Python's ``logging`` module for a different subsystem's warning path —
+LangGraph's pregel-scheduler log line for langgraph#5464; this one targets
+Python's ``warnings`` module, requiring a ``warnings.catch_warnings()``/
+``showwarning`` override instead of a ``logging.Handler``).
+
+Usage::
+
+    from agent_trace import tracer
+    from agent_trace.interceptor.warnings_hook import capture_warnings
+
+    with tracer.start_trace("my_graph") as trace:
+        with capture_warnings(tracer):
+            graph.invoke(state, config={"callbacks": [...]})
+
+Every ``UserWarning``/``RuntimeWarning`` (or whatever *categories* is
+passed) raised anywhere during the ``with`` block is persisted as a
+``runtime_warning`` event on a dedicated ``warnings:capture`` span —
+independent of, and in addition to, whatever HTTP/callback capture is also
+active.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import logging
+import warnings
+from collections.abc import Generator
+from typing import TYPE_CHECKING, Any
+
+from agent_trace.core.span import Span, SpanStatus
+
+if TYPE_CHECKING:
+    from agent_trace import Tracer
+
+__all__ = ["capture_warnings"]
+
+logger = logging.getLogger(__name__)
+
+# Bounds — this is diagnostic capture, not a full warnings mirror. A
+# pathological hot loop emitting thousands of warnings must not grow a
+# single span's event list unboundedly; warnings.captured_count keeps
+# counting past the cap.
+_MAX_ATTR_LEN = 4_000
+_MAX_WARNING_EVENTS = 200
+
+_DEFAULT_CATEGORIES: tuple[type[Warning], ...] = (UserWarning, RuntimeWarning)
+
+
+@contextlib.contextmanager
+def capture_warnings(
+    tracer: Tracer,
+    *,
+    categories: tuple[type[Warning], ...] = _DEFAULT_CATEGORIES,
+    span_name: str = "warnings:capture",
+) -> Generator[Span, None, None]:
+    """Install a ``warnings.catch_warnings()`` context with a custom
+    ``showwarning`` override for the lifetime of the ``with`` block,
+    persisting every caught warning matching *categories* onto a dedicated
+    ``warnings:capture`` span (opened via *tracer*).
+
+    Uses ``warnings.simplefilter("always")`` inside the ``catch_warnings()``
+    context so a warning that Python would otherwise only show once per
+    location (the default ``"default"`` filter action) is still captured on
+    every occurrence, without changing any filter state outside this
+    ``with`` block — ``catch_warnings()`` saves and restores
+    ``warnings.filters``/``showwarning`` on exit regardless of how the block
+    exits.
+
+    The span closes ``OK`` when the block exits normally (warning 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)
+    count = 0
+
+    def _showwarning(
+        message: Warning | str,
+        category: type[Warning],
+        filename: str,
+        lineno: int,
+        file: Any = None,
+        line: str | None = None,
+    ) -> None:
+        nonlocal count
+        if not (isinstance(category, type) and issubclass(category, categories)):
+            return
+        count += 1
+        if count > _MAX_WARNING_EVENTS:
+            return
+        try:
+            span.add_event(
+                "runtime_warning",
+                attributes={
+                    "warning.category": category.__name__,
+                    "warning.message": str(message)[:_MAX_ATTR_LEN],
+                    "warning.filename": str(filename),
+                    "warning.lineno": int(lineno),
+                },
+            )
+        except Exception:
+            logger.debug(
+                "agent-trace: failed to record captured warning event",
+                exc_info=True,
+            )
+
+    try:
+        with warnings.catch_warnings():
+            warnings.simplefilter("always")
+            warnings.showwarning = _showwarning
+            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("warnings.captured_count", count)
+        if span.end_time is None:
+            span.end(SpanStatus.OK)

agent_trace/interceptor/websocket_hook.py

@@ -0,0 +1,323 @@
+"""
+WebSocket connection wrappers for recording and replaying Realtime-API-style
+duplex sessions (e.g. the OpenAI Agents SDK's Realtime API, which opens a
+persistent ``websockets.asyncio.client.ClientConnection`` and exchanges many
+JSON events over it — tool calls, handoffs, audio deltas — rather than the
+discrete request/response model ``RecordingTransport``/``ReplayTransport``
+(``httpx_hook.py``) assume).
+
+RecordingWebSocketConnection wraps a real connection object (anything
+exposing the ``websockets`` client protocol: async ``send``, async ``recv``,
+async iteration, and ``close``): frames flow through unmodified in both
+directions while each one is teed into the fixture store as it passes.
+
+ReplayWebSocketConnection never touches the network. It serves previously
+recorded inbound ("recv") frames from the fixture, in the order they were
+captured, so a Realtime session can be replayed offline at zero API cost.
+Frames sent by the caller during replay are accepted but not forwarded
+anywhere — there is no live peer to send them to, matching how
+``ReplayTransport`` ignores request bodies and only replays responses keyed
+by (method, url).
+
+RecordingConnect is a drop-in stand-in for ``websockets.connect(...)`` (the
+callable this module patches) that supports both calling conventions the
+``websockets`` library provides::
+
+    ws = await websockets.connect(uri, **kwargs)
+    async with websockets.connect(uri, **kwargs) as ws:
+
+and returns a RecordingWebSocketConnection either way.
+"""
+
+from __future__ import annotations
+
+import logging
+import uuid
+import warnings
+from collections.abc import AsyncIterator
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from agent_trace._replay.fixture import Fixture
+
+from agent_trace.core.exceptions import NetworkGuardError, guard_active
+
+__all__ = [
+    "NetworkGuardError",
+    "RecordingConnect",
+    "RecordingWebSocketConnection",
+    "ReplayWebSocketConnection",
+    "WebSocketReplayExhaustedError",
+]
+
+logger = logging.getLogger(__name__)
+
+DIRECTION_SEND = "send"
+DIRECTION_RECV = "recv"
+
+
+class WebSocketReplayExhaustedError(EOFError):
+    """Raised by :meth:`ReplayWebSocketConnection.recv` when no more recorded
+    inbound frames remain and ``AGENT_TRACE_NETWORK_GUARD`` is not set.
+
+    Unlike ``ReplayTransport``'s HTTP fallback (which can fall through to a
+    real request), there is no live peer to fall back to for a duplex
+    WebSocket session during replay, so this is always the terminal signal —
+    the caller should treat it the way it would treat the underlying library's
+    own "connection closed" exception.
+    """
+
+
+def _decode_for_recording(message: Any) -> tuple[str, str]:
+    """Return (payload_as_text, frame_type) for a frame about to be recorded."""
+    if isinstance(message, (bytes, bytearray)):
+        return message.decode("utf-8", errors="replace"), "binary"
+    return str(message), "text"
+
+
+def _encode_from_frame(frame: dict[str, Any]) -> Any:
+    """Inverse of _decode_for_recording — reconstruct the original message shape."""
+    if frame["frame_type"] == "binary":
+        return frame["payload"].encode("utf-8")
+    return frame["payload"]
+
+
+class RecordingWebSocketConnection:
+    """Wraps a real WebSocket client connection, recording every frame.
+
+    Duck-types the subset of ``websockets.asyncio.client.ClientConnection``
+    that agent-trace's supported callers use — ``send``, ``recv``, async
+    iteration (``async for message in ws``), and ``close`` — and forwards
+    every other attribute access to the wrapped connection (e.g.
+    ``close_code``, ``state``, ``request``, ``response``), so instrumented
+    code can use this exactly like the real connection.
+
+    Parameters
+    ----------
+    inner:
+        The real, already-connected WebSocket client connection.
+    fixture:
+        Open Fixture instance where frames will be written.
+    url:
+        The URL the connection was opened against — stored alongside every
+        frame so a fixture holding multiple connections can be filtered.
+    connection_id:
+        Identifier grouping frames from this connection in the fixture.
+        Auto-generated if not supplied.
+    """
+
+    def __init__(
+        self,
+        inner: Any,
+        fixture: Fixture,
+        url: str,
+        connection_id: str | None = None,
+    ) -> None:
+        self._inner = inner
+        self._fixture = fixture
+        self._url = url
+        self.connection_id: str = connection_id or uuid.uuid4().hex[:16]
+
+    async def send(self, message: Any, *args: Any, **kwargs: Any) -> None:
+        """Send *message*, recording it as an outbound ("send") frame first.
+
+        Recording before the awaited send (rather than after) means a frame
+        that raises mid-send (e.g. connection dropped) is still captured —
+        matching the intent, if not the exact ordering, of
+        ``RecordingTransport`` which always records a completed exchange.
+        """
+        payload, frame_type = _decode_for_recording(message)
+        self._fixture.record_ws_frame(
+            connection_id=self.connection_id,
+            url=self._url,
+            direction=DIRECTION_SEND,
+            payload=payload,
+            frame_type=frame_type,
+        )
+        await self._inner.send(message, *args, **kwargs)
+
+    async def recv(self, *args: Any, **kwargs: Any) -> Any:
+        """Receive the next message, recording it as an inbound ("recv") frame."""
+        message = await self._inner.recv(*args, **kwargs)
+        payload, frame_type = _decode_for_recording(message)
+        self._fixture.record_ws_frame(
+            connection_id=self.connection_id,
+            url=self._url,
+            direction=DIRECTION_RECV,
+            payload=payload,
+            frame_type=frame_type,
+        )
+        return message
+
+    def __aiter__(self) -> AsyncIterator[Any]:
+        return self._iter_messages()
+
+    async def _iter_messages(self) -> AsyncIterator[Any]:
+        async for message in self._inner:
+            payload, frame_type = _decode_for_recording(message)
+            self._fixture.record_ws_frame(
+                connection_id=self.connection_id,
+                url=self._url,
+                direction=DIRECTION_RECV,
+                payload=payload,
+                frame_type=frame_type,
+            )
+            yield message
+
+    async def close(self, *args: Any, **kwargs: Any) -> None:
+        await self._inner.close(*args, **kwargs)
+
+    async def __aenter__(self) -> RecordingWebSocketConnection:
+        return self
+
+    async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None:
+        await self.close()
+
+    def __getattr__(self, name: str) -> Any:
+        # Passthrough for anything not explicitly wrapped above (close_code,
+        # state, request, response, ping, pong, ...).
+        return getattr(self._inner, name)
+
+
+class ReplayWebSocketConnection:
+    """Serves recorded WS frames from a Fixture without any network I/O.
+
+    Parameters
+    ----------
+    fixture:
+        Open Fixture instance to serve recorded frames from.
+    url:
+        The URL the connection would have been opened against — used only
+        for NetworkGuardError messages.
+    connection_id:
+        Identifier selecting which connection's frames to serve. Must match
+        the ``connection_id`` used when the session was recorded.
+    clock:
+        Optional FixtureClock to advance with each frame's recorded_at
+        timestamp, reproducing original timing during replay.
+    """
+
+    def __init__(
+        self,
+        fixture: Fixture,
+        url: str,
+        connection_id: str,
+        clock: Any | None = None,
+    ) -> None:
+        self._fixture = fixture
+        self._url = url
+        self.connection_id = connection_id
+        self._clock = clock
+        self.close_code: int | None = None
+
+    async def send(self, message: Any, *args: Any, **kwargs: Any) -> None:
+        """Accept a send during replay without forwarding it anywhere.
+
+        There is no live peer in replay mode — this mirrors how
+        ``ReplayTransport`` ignores the outbound request body and only ever
+        replays canned responses.
+        """
+        return None
+
+    async def recv(self, *args: Any, **kwargs: Any) -> Any:
+        """Return the next recorded inbound frame, or raise/warn if exhausted."""
+        frame = self._fixture.next_ws_frame(self.connection_id, DIRECTION_RECV)
+        if frame is None:
+            return self._on_exhausted()
+        return self._deliver(frame)
+
+    def __aiter__(self) -> AsyncIterator[Any]:
+        return self._iter_messages()
+
+    async def _iter_messages(self) -> AsyncIterator[Any]:
+        while True:
+            frame = self._fixture.next_ws_frame(self.connection_id, DIRECTION_RECV)
+            if frame is None:
+                # Mirrors a real connection closing: the async-for loop ends
+                # normally, exactly like `_listen_for_messages`'s
+                # `async for message in self._websocket:` does on
+                # ConnectionClosedOK.
+                return
+            yield self._deliver(frame)
+
+    async def close(self, code: int = 1000, reason: str = "") -> None:
+        self.close_code = code
+
+    async def __aenter__(self) -> ReplayWebSocketConnection:
+        return self
+
+    async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None:
+        await self.close()
+
+    def _deliver(self, frame: dict[str, Any]) -> Any:
+        if self._clock is not None:
+            self._clock.advance(float(frame["recorded_at"]))
+        return _encode_from_frame(frame)
+
+    def _on_exhausted(self) -> Any:
+        if guard_active():
+            raise NetworkGuardError(
+                f"No recorded inbound WS frame for connection "
+                f"{self.connection_id!r} ({self._url}) and "
+                "AGENT_TRACE_NETWORK_GUARD=1 is set.  Run in recording mode "
+                "first to capture this session."
+            )
+        warnings.warn(
+            f"agent-trace: no more recorded WS frames for connection "
+            f"{self.connection_id!r} ({self._url}); there is no live "
+            "connection to fall back to during replay.  Set "
+            "AGENT_TRACE_NETWORK_GUARD=1 to make this an error.",
+            stacklevel=2,
+        )
+        raise WebSocketReplayExhaustedError(
+            f"agent-trace: WS fixture exhausted for connection "
+            f"{self.connection_id!r} ({self._url})"
+        )
+
+
+class RecordingConnect:
+    """Drop-in stand-in for ``websockets.connect(...)`` that records frames.
+
+    Constructed with the *real* ``connect`` callable/class (so the patch
+    installer keeps full control of what "real" means, exactly like
+    ``RecordingTransport`` accepts an ``inner`` transport) plus the fixture
+    to record into.  Supports both calling conventions ``websockets.connect``
+    provides:
+
+        ws = await websockets.connect(uri, **kwargs)
+        async with websockets.connect(uri, **kwargs) as ws:
+
+    and returns a :class:`RecordingWebSocketConnection` either way.
+    """
+
+    def __init__(
+        self,
+        real_connect: Any,
+        fixture: Fixture,
+        uri: str,
+        *args: Any,
+        connection_id: str | None = None,
+        **kwargs: Any,
+    ) -> None:
+        self._real = real_connect(uri, *args, **kwargs)
+        self._fixture = fixture
+        self._uri = uri
+        self._connection_id = connection_id
+        self._wrapped: RecordingWebSocketConnection | None = None
+
+    def __await__(self) -> Any:
+        return self._connect().__await__()
+
+    async def _connect(self) -> RecordingWebSocketConnection:
+        inner = await self._real
+        self._wrapped = RecordingWebSocketConnection(
+            inner, self._fixture, self._uri, connection_id=self._connection_id
+        )
+        return self._wrapped
+
+    async def __aenter__(self) -> RecordingWebSocketConnection:
+        return await self._connect()
+
+    async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None:
+        if self._wrapped is not None:
+            await self._wrapped.close()