wcgw 5.0.2__py3-none-any.whl → 5.1.0__py3-none-any.whl

mcp_wcgw/shared/memory.py

@@ -1,87 +0,0 @@
-"""
-In-memory transports
-"""
-
-from contextlib import asynccontextmanager
-from datetime import timedelta
-from typing import AsyncGenerator
-
-import anyio
-from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
-
-from mcp_wcgw.client.session import ClientSession
-from mcp_wcgw.server import Server
-from mcp_wcgw.types import JSONRPCMessage
-
-MessageStream = tuple[
-    MemoryObjectReceiveStream[JSONRPCMessage | Exception],
-    MemoryObjectSendStream[JSONRPCMessage],
-]
-
-
-@asynccontextmanager
-async def create_client_server_memory_streams() -> (
-    AsyncGenerator[tuple[MessageStream, MessageStream], None]
-):
-    """
-    Creates a pair of bidirectional memory streams for client-server communication.
-
-    Returns:
-        A tuple of (client_streams, server_streams) where each is a tuple of
-        (read_stream, write_stream)
-    """
-    # Create streams for both directions
-    server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[
-        JSONRPCMessage | Exception
-    ](1)
-    client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[
-        JSONRPCMessage | Exception
-    ](1)
-
-    client_streams = (server_to_client_receive, client_to_server_send)
-    server_streams = (client_to_server_receive, server_to_client_send)
-
-    async with (
-        server_to_client_receive,
-        client_to_server_send,
-        client_to_server_receive,
-        server_to_client_send,
-    ):
-        yield client_streams, server_streams
-
-
-@asynccontextmanager
-async def create_connected_server_and_client_session(
-    server: Server,
-    read_timeout_seconds: timedelta | None = None,
-    raise_exceptions: bool = False,
-) -> AsyncGenerator[ClientSession, None]:
-    """Creates a ClientSession that is connected to a running MCP server."""
-    async with create_client_server_memory_streams() as (
-        client_streams,
-        server_streams,
-    ):
-        client_read, client_write = client_streams
-        server_read, server_write = server_streams
-
-        # Create a cancel scope for the server task
-        async with anyio.create_task_group() as tg:
-            tg.start_soon(
-                lambda: server.run(
-                    server_read,
-                    server_write,
-                    server.create_initialization_options(),
-                    raise_exceptions=raise_exceptions,
-                )
-            )
-
-            try:
-                async with ClientSession(
-                    read_stream=client_read,
-                    write_stream=client_write,
-                    read_timeout_seconds=read_timeout_seconds,
-                ) as client_session:
-                    await client_session.initialize()
-                    yield client_session
-            finally:
-                tg.cancel_scope.cancel()

mcp_wcgw/shared/progress.py

@@ -1,40 +0,0 @@
-from contextlib import contextmanager
-from dataclasses import dataclass, field
-
-from pydantic import BaseModel
-
-from mcp_wcgw.shared.context import RequestContext
-from mcp_wcgw.shared.session import BaseSession
-from mcp_wcgw.types import ProgressToken
-
-
-class Progress(BaseModel):
-    progress: float
-    total: float | None
-
-
-@dataclass
-class ProgressContext:
-    session: BaseSession
-    progress_token: ProgressToken
-    total: float | None
-    current: float = field(default=0.0, init=False)
-
-    async def progress(self, amount: float) -> None:
-        self.current += amount
-
-        await self.session.send_progress_notification(
-            self.progress_token, self.current, total=self.total
-        )
-
-
-@contextmanager
-def progress(ctx: RequestContext, total: float | None = None):
-    if ctx.meta is None or ctx.meta.progressToken is None:
-        raise ValueError("No progress token provided")
-
-    progress_ctx = ProgressContext(ctx.session, ctx.meta.progressToken, total)
-    try:
-        yield progress_ctx
-    finally:
-        pass

mcp_wcgw/shared/session.py

@@ -1,288 +0,0 @@
-from contextlib import AbstractAsyncContextManager
-from datetime import timedelta
-from typing import Generic, TypeVar
-
-import anyio
-import anyio.lowlevel
-import httpx
-from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
-from pydantic import BaseModel
-
-from mcp_wcgw.shared.exceptions import McpError
-from mcp_wcgw.types import (
-    ClientNotification,
-    ClientRequest,
-    ClientResult,
-    ErrorData,
-    JSONRPCError,
-    JSONRPCMessage,
-    JSONRPCNotification,
-    JSONRPCRequest,
-    JSONRPCResponse,
-    RequestParams,
-    ServerNotification,
-    ServerRequest,
-    ServerResult,
-)
-
-SendRequestT = TypeVar("SendRequestT", ClientRequest, ServerRequest)
-SendResultT = TypeVar("SendResultT", ClientResult, ServerResult)
-SendNotificationT = TypeVar("SendNotificationT", ClientNotification, ServerNotification)
-ReceiveRequestT = TypeVar("ReceiveRequestT", ClientRequest, ServerRequest)
-ReceiveResultT = TypeVar("ReceiveResultT", bound=BaseModel)
-ReceiveNotificationT = TypeVar(
-    "ReceiveNotificationT", ClientNotification, ServerNotification
-)
-
-RequestId = str | int
-
-
-class RequestResponder(Generic[ReceiveRequestT, SendResultT]):
-    def __init__(
-        self,
-        request_id: RequestId,
-        request_meta: RequestParams.Meta | None,
-        request: ReceiveRequestT,
-        session: "BaseSession",
-    ) -> None:
-        self.request_id = request_id
-        self.request_meta = request_meta
-        self.request = request
-        self._session = session
-        self._responded = False
-
-    async def respond(self, response: SendResultT | ErrorData) -> None:
-        assert not self._responded, "Request already responded to"
-        self._responded = True
-
-        await self._session._send_response(
-            request_id=self.request_id, response=response
-        )
-
-
-class BaseSession(
-    AbstractAsyncContextManager,
-    Generic[
-        SendRequestT,
-        SendNotificationT,
-        SendResultT,
-        ReceiveRequestT,
-        ReceiveNotificationT,
-    ],
-):
-    """
-    Implements an MCP "session" on top of read/write streams, including features
-    like request/response linking, notifications, and progress.
-
-    This class is an async context manager that automatically starts processing
-    messages when entered.
-    """
-
-    _response_streams: dict[
-        RequestId, MemoryObjectSendStream[JSONRPCResponse | JSONRPCError]
-    ]
-    _request_id: int
-
-    def __init__(
-        self,
-        read_stream: MemoryObjectReceiveStream[JSONRPCMessage | Exception],
-        write_stream: MemoryObjectSendStream[JSONRPCMessage],
-        receive_request_type: type[ReceiveRequestT],
-        receive_notification_type: type[ReceiveNotificationT],
-        # If none, reading will never time out
-        read_timeout_seconds: timedelta | None = None,
-    ) -> None:
-        self._read_stream = read_stream
-        self._write_stream = write_stream
-        self._response_streams = {}
-        self._request_id = 0
-        self._receive_request_type = receive_request_type
-        self._receive_notification_type = receive_notification_type
-        self._read_timeout_seconds = read_timeout_seconds
-
-        self._incoming_message_stream_writer, self._incoming_message_stream_reader = (
-            anyio.create_memory_object_stream[
-                RequestResponder[ReceiveRequestT, SendResultT]
-                | ReceiveNotificationT
-                | Exception
-            ]()
-        )
-
-    async def __aenter__(self):
-        self._task_group = anyio.create_task_group()
-        await self._task_group.__aenter__()
-        self._task_group.start_soon(self._receive_loop)
-        return self
-
-    async def __aexit__(self, exc_type, exc_val, exc_tb):
-        # Using BaseSession as a context manager should not block on exit (this
-        # would be very surprising behavior), so make sure to cancel the tasks
-        # in the task group.
-        self._task_group.cancel_scope.cancel()
-        return await self._task_group.__aexit__(exc_type, exc_val, exc_tb)
-
-    async def send_request(
-        self,
-        request: SendRequestT,
-        result_type: type[ReceiveResultT],
-    ) -> ReceiveResultT:
-        """
-        Sends a request and wait for a response. Raises an McpError if the
-        response contains an error.
-
-        Do not use this method to emit notifications! Use send_notification()
-        instead.
-        """
-
-        request_id = self._request_id
-        self._request_id = request_id + 1
-
-        response_stream, response_stream_reader = anyio.create_memory_object_stream[
-            JSONRPCResponse | JSONRPCError
-        ](1)
-        self._response_streams[request_id] = response_stream
-
-        jsonrpc_request = JSONRPCRequest(
-            jsonrpc="2.0",
-            id=request_id,
-            **request.model_dump(by_alias=True, mode="json", exclude_none=True),
-        )
-
-        # TODO: Support progress callbacks
-
-        await self._write_stream.send(JSONRPCMessage(jsonrpc_request))
-
-        try:
-            with anyio.fail_after(
-                None
-                if self._read_timeout_seconds is None
-                else self._read_timeout_seconds.total_seconds()
-            ):
-                response_or_error = await response_stream_reader.receive()
-        except TimeoutError:
-            raise McpError(
-                ErrorData(
-                    code=httpx.codes.REQUEST_TIMEOUT,
-                    message=(
-                        f"Timed out while waiting for response to "
-                        f"{request.__class__.__name__}. Waited "
-                        f"{self._read_timeout_seconds} seconds."
-                    ),
-                )
-            )
-
-        if isinstance(response_or_error, JSONRPCError):
-            raise McpError(response_or_error.error)
-        else:
-            return result_type.model_validate(response_or_error.result)
-
-    async def send_notification(self, notification: SendNotificationT) -> None:
-        """
-        Emits a notification, which is a one-way message that does not expect
-        a response.
-        """
-        jsonrpc_notification = JSONRPCNotification(
-            jsonrpc="2.0",
-            **notification.model_dump(by_alias=True, mode="json", exclude_none=True),
-        )
-
-        await self._write_stream.send(JSONRPCMessage(jsonrpc_notification))
-
-    async def _send_response(
-        self, request_id: RequestId, response: SendResultT | ErrorData
-    ) -> None:
-        if isinstance(response, ErrorData):
-            jsonrpc_error = JSONRPCError(jsonrpc="2.0", id=request_id, error=response)
-            await self._write_stream.send(JSONRPCMessage(jsonrpc_error))
-        else:
-            jsonrpc_response = JSONRPCResponse(
-                jsonrpc="2.0",
-                id=request_id,
-                result=response.model_dump(
-                    by_alias=True, mode="json", exclude_none=True
-                ),
-            )
-            await self._write_stream.send(JSONRPCMessage(jsonrpc_response))
-
-    async def _receive_loop(self) -> None:
-        async with (
-            self._read_stream,
-            self._write_stream,
-            self._incoming_message_stream_writer,
-        ):
-            async for message in self._read_stream:
-                if isinstance(message, Exception):
-                    await self._incoming_message_stream_writer.send(message)
-                elif isinstance(message.root, JSONRPCRequest):
-                    validated_request = self._receive_request_type.model_validate(
-                        message.root.model_dump(
-                            by_alias=True, mode="json", exclude_none=True
-                        )
-                    )
-                    responder = RequestResponder(
-                        request_id=message.root.id,
-                        request_meta=validated_request.root.params._meta
-                        if validated_request.root.params
-                        else None,
-                        request=validated_request,
-                        session=self,
-                    )
-
-                    await self._received_request(responder)
-                    if not responder._responded:
-                        await self._incoming_message_stream_writer.send(responder)
-                elif isinstance(message.root, JSONRPCNotification):
-                    notification = self._receive_notification_type.model_validate(
-                        message.root.model_dump(
-                            by_alias=True, mode="json", exclude_none=True
-                        )
-                    )
-
-                    await self._received_notification(notification)
-                    await self._incoming_message_stream_writer.send(notification)
-                else:  # Response or error
-                    stream = self._response_streams.pop(message.root.id, None)
-                    if stream:
-                        await stream.send(message.root)
-                    else:
-                        await self._incoming_message_stream_writer.send(
-                            RuntimeError(
-                                "Received response with an unknown "
-                                f"request ID: {message}"
-                            )
-                        )
-
-    async def _received_request(
-        self, responder: RequestResponder[ReceiveRequestT, SendResultT]
-    ) -> None:
-        """
-        Can be overridden by subclasses to handle a request without needing to
-        listen on the message stream.
-
-        If the request is responded to within this method, it will not be
-        forwarded on to the message stream.
-        """
-
-    async def _received_notification(self, notification: ReceiveNotificationT) -> None:
-        """
-        Can be overridden by subclasses to handle a notification without needing
-        to listen on the message stream.
-        """
-
-    async def send_progress_notification(
-        self, progress_token: str | int, progress: float, total: float | None = None
-    ) -> None:
-        """
-        Sends a progress notification for a request that is currently being
-        processed.
-        """
-
-    @property
-    def incoming_messages(
-        self,
-    ) -> MemoryObjectReceiveStream[
-        RequestResponder[ReceiveRequestT, SendResultT]
-        | ReceiveNotificationT
-        | Exception
-    ]:
-        return self._incoming_message_stream_reader

mcp_wcgw/shared/version.py

@@ -1,3 +0,0 @@
-from mcp_wcgw.types import LATEST_PROTOCOL_VERSION
-
-SUPPORTED_PROTOCOL_VERSIONS = [1, LATEST_PROTOCOL_VERSION]