wehub-resource-sync
49b9bb6724
Deploy Docs / deploy-docs (push) Failing after 1s
Conformance Tests / client-conformance (push) Failing after 3s
Conformance Tests / server-conformance (push) Failing after 1s
GitHub Actions Security Analysis / zizmor (push) Failing after 1s
CI / checks (push) Failing after 59m20s
CI / all-green (push) Waiting to run
109 lines
4.2 KiB
Python
109 lines
4.2 KiB
Python
"""Shared helpers for the interaction suite.
|
|
|
|
Keep this module small: it exists only for (a) types that every test would otherwise have to
|
|
assemble from the SDK's internals to annotate a client callback, and (b) the recording transport
|
|
used by the wire-level tests. Server fixtures and assertion helpers belong in the test that uses
|
|
them.
|
|
"""
|
|
|
|
from types import TracebackType
|
|
|
|
import anyio
|
|
from mcp_types import ClientResult, ServerNotification, ServerRequest
|
|
from typing_extensions import Self
|
|
|
|
from mcp.client._transport import ReadStream, Transport, TransportStreams, WriteStream
|
|
from mcp.shared.message import SessionMessage
|
|
from mcp.shared.session import RequestResponder
|
|
|
|
# TODO: this union is the parameter type of every client message handler (MessageHandlerFnT),
|
|
# but the SDK does not export a name for it -- writing a correctly-typed handler requires
|
|
# importing RequestResponder from mcp.shared.session and assembling the union by hand. It
|
|
# should be a named, exported alias next to MessageHandlerFnT (like ClientRequestContext is
|
|
# for the request callbacks), at which point this alias can be deleted.
|
|
IncomingMessage = RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception
|
|
"""Everything a client message handler can receive."""
|
|
|
|
|
|
class _RecordingReadStream:
|
|
"""Delegates to a read stream, appending every received message to a log."""
|
|
|
|
def __init__(self, inner: ReadStream[SessionMessage | Exception], log: list[SessionMessage | Exception]) -> None:
|
|
self._inner = inner
|
|
self._log = log
|
|
|
|
async def receive(self) -> SessionMessage | Exception:
|
|
item = await self._inner.receive()
|
|
self._log.append(item)
|
|
return item
|
|
|
|
async def aclose(self) -> None:
|
|
await self._inner.aclose()
|
|
|
|
def __aiter__(self) -> Self:
|
|
return self
|
|
|
|
async def __anext__(self) -> SessionMessage | Exception:
|
|
try:
|
|
return await self.receive()
|
|
except anyio.EndOfStream:
|
|
raise StopAsyncIteration from None
|
|
|
|
async def __aenter__(self) -> Self:
|
|
return self
|
|
|
|
async def __aexit__(
|
|
self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None
|
|
) -> bool | None:
|
|
await self.aclose()
|
|
return None
|
|
|
|
|
|
class _RecordingWriteStream:
|
|
"""Delegates to a write stream, appending every sent message to a log."""
|
|
|
|
def __init__(self, inner: WriteStream[SessionMessage], log: list[SessionMessage]) -> None:
|
|
self._inner = inner
|
|
self._log = log
|
|
|
|
async def send(self, item: SessionMessage, /) -> None:
|
|
# Record only after the inner send returns: a failed or cancelled send never reached the transport.
|
|
await self._inner.send(item)
|
|
self._log.append(item)
|
|
|
|
async def aclose(self) -> None:
|
|
await self._inner.aclose()
|
|
|
|
async def __aenter__(self) -> Self:
|
|
return self
|
|
|
|
async def __aexit__(
|
|
self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None
|
|
) -> bool | None:
|
|
await self.aclose()
|
|
return None
|
|
|
|
|
|
class RecordingTransport:
|
|
"""Wraps a Transport and records every message crossing the client's transport boundary.
|
|
|
|
`sent` holds everything the client wrote towards the server; `received` holds everything the
|
|
server delivered to the client. The recording sits at the transport seam -- the exact payloads
|
|
a real transport would serialise -- and never touches the session, so wire-level assertions
|
|
written against it survive changes to the receive path.
|
|
"""
|
|
|
|
def __init__(self, inner: Transport) -> None:
|
|
self.inner = inner
|
|
self.sent: list[SessionMessage] = []
|
|
self.received: list[SessionMessage | Exception] = []
|
|
|
|
async def __aenter__(self) -> TransportStreams:
|
|
read_stream, write_stream = await self.inner.__aenter__()
|
|
return _RecordingReadStream(read_stream, self.received), _RecordingWriteStream(write_stream, self.sent)
|
|
|
|
async def __aexit__(
|
|
self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None
|
|
) -> bool | None:
|
|
return await self.inner.__aexit__(exc_type, exc_val, exc_tb)
|