Initial import

mcp_python/__init__.py

@@ -0,0 +1,104 @@
+from .client.session import ClientSession
+from .client.stdio import StdioServerParameters, stdio_client
+from .server.session import ServerSession
+from .server.stdio import stdio_server
+from .shared.exceptions import McpError
+from .types import (
+    CallToolRequest,
+    ClientCapabilities,
+    ClientNotification,
+    ClientRequest,
+    ClientResult,
+    CompleteRequest,
+    CreateMessageRequest,
+    CreateMessageResult,
+    ErrorData,
+    GetPromptRequest,
+    GetPromptResult,
+    Implementation,
+    IncludeContext,
+    InitializedNotification,
+    InitializeRequest,
+    InitializeResult,
+    JSONRPCError,
+    JSONRPCRequest,
+    JSONRPCResponse,
+    ListPromptsRequest,
+    ListPromptsResult,
+    ListResourcesRequest,
+    ListResourcesResult,
+    ListToolsResult,
+    LoggingLevel,
+    LoggingMessageNotification,
+    Notification,
+    PingRequest,
+    ProgressNotification,
+    ReadResourceRequest,
+    ReadResourceResult,
+    Resource,
+    ResourceUpdatedNotification,
+    Role as SamplingRole,
+    SamplingMessage,
+    ServerCapabilities,
+    ServerNotification,
+    ServerRequest,
+    ServerResult,
+    SetLevelRequest,
+    StopReason,
+    SubscribeRequest,
+    Tool,
+    UnsubscribeRequest,
+)
+
+__all__ = [
+    "CallToolRequest",
+    "ClientCapabilities",
+    "ClientNotification",
+    "ClientRequest",
+    "ClientResult",
+    "ClientSession",
+    "CreateMessageRequest",
+    "CreateMessageResult",
+    "ErrorData",
+    "GetPromptRequest",
+    "GetPromptResult",
+    "Implementation",
+    "IncludeContext",
+    "InitializeRequest",
+    "InitializeResult",
+    "InitializedNotification",
+    "JSONRPCError",
+    "JSONRPCRequest",
+    "ListPromptsRequest",
+    "ListPromptsResult",
+    "ListResourcesRequest",
+    "ListResourcesResult",
+    "ListToolsResult",
+    "LoggingLevel",
+    "LoggingMessageNotification",
+    "McpError",
+    "Notification",
+    "PingRequest",
+    "ProgressNotification",
+    "ReadResourceRequest",
+    "ReadResourceResult",
+    "ResourceUpdatedNotification",
+    "Resource",
+    "SamplingMessage",
+    "SamplingRole",
+    "ServerCapabilities",
+    "ServerNotification",
+    "ServerRequest",
+    "ServerResult",
+    "ServerSession",
+    "SetLevelRequest",
+    "StdioServerParameters",
+    "StopReason",
+    "SubscribeRequest",
+    "Tool",
+    "UnsubscribeRequest",
+    "stdio_client",
+    "stdio_server",
+    "CompleteRequest",
+    "JSONRPCResponse",
+]

mcp_python/client/__main__.py

@@ -0,0 +1,76 @@
+import logging
+import sys
+from functools import partial
+from urllib.parse import urlparse
+
+import anyio
+import click
+
+from mcp_python.client.session import ClientSession
+from mcp_python.client.sse import sse_client
+from mcp_python.client.stdio import StdioServerParameters, stdio_client
+
+if not sys.warnoptions:
+    import warnings
+
+    warnings.simplefilter("ignore")
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger("client")
+
+
+async def receive_loop(session: ClientSession):
+    logger.info("Starting receive loop")
+    async for message in session.incoming_messages:
+        if isinstance(message, Exception):
+            logger.error("Error: %s", message)
+            continue
+
+        logger.info("Received message from server: %s", message)
+
+
+async def run_session(read_stream, write_stream):
+    async with (
+        ClientSession(read_stream, write_stream) as session,
+        anyio.create_task_group() as tg,
+    ):
+        tg.start_soon(receive_loop, session)
+
+        logger.info("Initializing session")
+        await session.initialize()
+        logger.info("Initialized")
+
+
+async def main(command_or_url: str, args: list[str], env: list[tuple[str, str]]):
+    env_dict = dict(env)
+
+    if urlparse(command_or_url).scheme in ("http", "https"):
+        # Use SSE client for HTTP(S) URLs
+        async with sse_client(command_or_url) as streams:
+            await run_session(*streams)
+    else:
+        # Use stdio client for commands
+        server_parameters = StdioServerParameters(
+            command=command_or_url, args=args, env=env_dict
+        )
+        async with stdio_client(server_parameters) as streams:
+            await run_session(*streams)
+
+
+@click.command()
+@click.argument("command_or_url")
+@click.argument("args", nargs=-1)
+@click.option(
+    "--env",
+    "-e",
+    multiple=True,
+    nargs=2,
+    metavar="KEY VALUE",
+    help="Environment variables to set. Can be used multiple times.",
+)
+def cli(*args, **kwargs):
+    anyio.run(partial(main, *args, **kwargs), backend="trio")
+
+
+if __name__ == "__main__":
+    cli()

mcp_python/client/session.py

@@ -0,0 +1,211 @@
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+from pydantic import AnyUrl
+
+from mcp_python.shared.session import BaseSession
+from mcp_python.shared.version import SUPPORTED_PROTOCOL_VERSION
+from mcp_python.types import (
+    CallToolResult,
+    ClientCapabilities,
+    ClientNotification,
+    ClientRequest,
+    ClientResult,
+    EmptyResult,
+    Implementation,
+    InitializedNotification,
+    InitializeResult,
+    JSONRPCMessage,
+    ListResourcesResult,
+    LoggingLevel,
+    ReadResourceResult,
+    ServerNotification,
+    ServerRequest,
+)
+
+
+class ClientSession(
+    BaseSession[
+        ClientRequest,
+        ClientNotification,
+        ClientResult,
+        ServerRequest,
+        ServerNotification,
+    ]
+):
+    def __init__(
+        self,
+        read_stream: MemoryObjectReceiveStream[JSONRPCMessage | Exception],
+        write_stream: MemoryObjectSendStream[JSONRPCMessage],
+    ) -> None:
+        super().__init__(read_stream, write_stream, ServerRequest, ServerNotification)
+
+    async def initialize(self) -> InitializeResult:
+        from mcp_python.types import (
+            InitializeRequest,
+            InitializeRequestParams,
+        )
+
+        result = await self.send_request(
+            ClientRequest(
+                InitializeRequest(
+                    method="initialize",
+                    params=InitializeRequestParams(
+                        protocolVersion=SUPPORTED_PROTOCOL_VERSION,
+                        capabilities=ClientCapabilities(
+                            sampling=None, experimental=None
+                        ),
+                        clientInfo=Implementation(name="mcp_python", version="0.1.0"),
+                    ),
+                )
+            ),
+            InitializeResult,
+        )
+
+        if result.protocolVersion != SUPPORTED_PROTOCOL_VERSION:
+            raise RuntimeError(
+                f"Unsupported protocol version from the server: {result.protocolVersion}"
+            )
+
+        await self.send_notification(
+            ClientNotification(
+                InitializedNotification(method="notifications/initialized")
+            )
+        )
+
+        return result
+
+    async def send_ping(self) -> EmptyResult:
+        """Send a ping request."""
+        from mcp_python.types import PingRequest
+
+        return await self.send_request(
+            ClientRequest(
+                PingRequest(
+                    method="ping",
+                )
+            ),
+            EmptyResult,
+        )
+
+    async def send_progress_notification(
+        self, progress_token: str | int, progress: float, total: float | None = None
+    ) -> None:
+        """Send a progress notification."""
+        from mcp_python.types import (
+            ProgressNotification,
+            ProgressNotificationParams,
+        )
+
+        await self.send_notification(
+            ClientNotification(
+                ProgressNotification(
+                    method="notifications/progress",
+                    params=ProgressNotificationParams(
+                        progressToken=progress_token,
+                        progress=progress,
+                        total=total,
+                    ),
+                ),
+            )
+        )
+
+    async def set_logging_level(self, level: LoggingLevel) -> EmptyResult:
+        """Send a logging/setLevel request."""
+        from mcp_python.types import (
+            SetLevelRequest,
+            SetLevelRequestParams,
+        )
+
+        return await self.send_request(
+            ClientRequest(
+                SetLevelRequest(
+                    method="logging/setLevel",
+                    params=SetLevelRequestParams(level=level),
+                )
+            ),
+            EmptyResult,
+        )
+
+    async def list_resources(self) -> ListResourcesResult:
+        """Send a resources/list request."""
+        from mcp_python.types import (
+            ListResourcesRequest,
+        )
+
+        return await self.send_request(
+            ClientRequest(
+                ListResourcesRequest(
+                    method="resources/list",
+                )
+            ),
+            ListResourcesResult,
+        )
+
+    async def read_resource(self, uri: AnyUrl) -> ReadResourceResult:
+        """Send a resources/read request."""
+        from mcp_python.types import (
+            ReadResourceRequest,
+            ReadResourceRequestParams,
+        )
+
+        return await self.send_request(
+            ClientRequest(
+                ReadResourceRequest(
+                    method="resources/read",
+                    params=ReadResourceRequestParams(uri=uri),
+                )
+            ),
+            ReadResourceResult,
+        )
+
+    async def subscribe_resource(self, uri: AnyUrl) -> EmptyResult:
+        """Send a resources/subscribe request."""
+        from mcp_python.types import (
+            SubscribeRequest,
+            SubscribeRequestParams,
+        )
+
+        return await self.send_request(
+            ClientRequest(
+                SubscribeRequest(
+                    method="resources/subscribe",
+                    params=SubscribeRequestParams(uri=uri),
+                )
+            ),
+            EmptyResult,
+        )
+
+    async def unsubscribe_resource(self, uri: AnyUrl) -> EmptyResult:
+        """Send a resources/unsubscribe request."""
+        from mcp_python.types import (
+            UnsubscribeRequest,
+            UnsubscribeRequestParams,
+        )
+
+        return await self.send_request(
+            ClientRequest(
+                UnsubscribeRequest(
+                    method="resources/unsubscribe",
+                    params=UnsubscribeRequestParams(uri=uri),
+                )
+            ),
+            EmptyResult,
+        )
+
+    async def call_tool(
+        self, name: str, arguments: dict | None = None
+    ) -> CallToolResult:
+        """Send a tools/call request."""
+        from mcp_python.types import (
+            CallToolRequest,
+            CallToolRequestParams,
+        )
+
+        return await self.send_request(
+            ClientRequest(
+                CallToolRequest(
+                    method="tools/call",
+                    params=CallToolRequestParams(name=name, arguments=arguments),
+                )
+            ),
+            CallToolResult,
+        )

mcp_python/client/sse.py

@@ -0,0 +1,129 @@
+import logging
+from contextlib import asynccontextmanager
+from urllib.parse import urljoin, urlparse
+
+import anyio
+import httpx
+from anyio.abc import TaskStatus
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+from httpx_sse import aconnect_sse
+
+from mcp_python.types import JSONRPCMessage
+
+logger = logging.getLogger(__name__)
+
+
+def remove_request_params(url: str) -> str:
+    return urljoin(url, urlparse(url).path)
+
+
+@asynccontextmanager
+async def sse_client(url: str, timeout: float = 5, sse_read_timeout: float = 60 * 5):
+    """
+    Client transport for SSE.
+
+    `sse_read_timeout` determines how long (in seconds) the client will wait for a new event before disconnecting. All other HTTP operations are controlled by `timeout`.
+    """
+    read_stream: MemoryObjectReceiveStream[JSONRPCMessage | Exception]
+    read_stream_writer: MemoryObjectSendStream[JSONRPCMessage | Exception]
+
+    write_stream: MemoryObjectSendStream[JSONRPCMessage]
+    write_stream_reader: MemoryObjectReceiveStream[JSONRPCMessage]
+
+    read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
+    write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
+
+    async with anyio.create_task_group() as tg:
+        try:
+            logger.info(f"Connecting to SSE endpoint: {remove_request_params(url)}")
+            async with httpx.AsyncClient() as client:
+                async with aconnect_sse(
+                    client,
+                    "GET",
+                    url,
+                    timeout=httpx.Timeout(timeout, read=sse_read_timeout),
+                ) as event_source:
+                    event_source.response.raise_for_status()
+                    logger.debug("SSE connection established")
+
+                    async def sse_reader(
+                        task_status: TaskStatus[str] = anyio.TASK_STATUS_IGNORED,
+                    ):
+                        try:
+                            async for sse in event_source.aiter_sse():
+                                logger.debug(f"Received SSE event: {sse.event}")
+                                match sse.event:
+                                    case "endpoint":
+                                        endpoint_url = urljoin(url, sse.data)
+                                        logger.info(
+                                            f"Received endpoint URL: {endpoint_url}"
+                                        )
+
+                                        url_parsed = urlparse(url)
+                                        endpoint_parsed = urlparse(endpoint_url)
+                                        if (
+                                            url_parsed.netloc != endpoint_parsed.netloc
+                                            or url_parsed.scheme
+                                            != endpoint_parsed.scheme
+                                        ):
+                                            error_msg = f"Endpoint origin does not match connection origin: {endpoint_url}"
+                                            logger.error(error_msg)
+                                            raise ValueError(error_msg)
+
+                                        task_status.started(endpoint_url)
+
+                                    case "message":
+                                        try:
+                                            message = (
+                                                JSONRPCMessage.model_validate_json(
+                                                    sse.data
+                                                )
+                                            )
+                                            logger.debug(
+                                                f"Received server message: {message}"
+                                            )
+                                        except Exception as exc:
+                                            logger.error(
+                                                f"Error parsing server message: {exc}"
+                                            )
+                                            await read_stream_writer.send(exc)
+                                            continue
+
+                                        await read_stream_writer.send(message)
+                        except Exception as exc:
+                            logger.error(f"Error in sse_reader: {exc}")
+                            await read_stream_writer.send(exc)
+                        finally:
+                            await read_stream_writer.aclose()
+
+                    async def post_writer(endpoint_url: str):
+                        try:
+                            async with write_stream_reader:
+                                async for message in write_stream_reader:
+                                    logger.debug(f"Sending client message: {message}")
+                                    response = await client.post(
+                                        endpoint_url,
+                                        json=message.model_dump(by_alias=True, mode="json"),
+                                    )
+                                    response.raise_for_status()
+                                    logger.debug(
+                                        f"Client message sent successfully: {response.status_code}"
+                                    )
+                        except Exception as exc:
+                            logger.error(f"Error in post_writer: {exc}")
+                        finally:
+                            await write_stream.aclose()
+
+                    endpoint_url = await tg.start(sse_reader)
+                    logger.info(
+                        f"Starting post writer with endpoint URL: {endpoint_url}"
+                    )
+                    tg.start_soon(post_writer, endpoint_url)
+
+                    try:
+                        yield read_stream, write_stream
+                    finally:
+                        tg.cancel_scope.cancel()
+        finally:
+            await read_stream_writer.aclose()
+            await write_stream.aclose()

mcp_python/client/stdio.py

@@ -0,0 +1,84 @@
+import sys
+from contextlib import asynccontextmanager
+
+import anyio
+import anyio.lowlevel
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+from anyio.streams.text import TextReceiveStream
+from pydantic import BaseModel, Field
+
+from mcp_python.types import JSONRPCMessage
+
+
+class StdioServerParameters(BaseModel):
+    command: str
+    """The executable to run to start the server."""
+
+    args: list[str] = Field(default_factory=list)
+    """Command line arguments to pass to the executable."""
+
+    env: dict[str, str] = Field(default_factory=dict)
+    """
+    The environment to use when spawning the process.
+
+    The environment is NOT inherited from the parent process by default.
+    """
+
+
+@asynccontextmanager
+async def stdio_client(server: StdioServerParameters):
+    """
+    Client transport for stdio: this will connect to a server by spawning a process and communicating with it over stdin/stdout.
+    """
+    read_stream: MemoryObjectReceiveStream[JSONRPCMessage | Exception]
+    read_stream_writer: MemoryObjectSendStream[JSONRPCMessage | Exception]
+
+    write_stream: MemoryObjectSendStream[JSONRPCMessage]
+    write_stream_reader: MemoryObjectReceiveStream[JSONRPCMessage]
+
+    read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
+    write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
+
+    process = await anyio.open_process(
+        [server.command, *server.args], env=server.env, stderr=sys.stderr
+    )
+
+    async def stdout_reader():
+        assert process.stdout, "Opened process is missing stdout"
+
+        try:
+            async with read_stream_writer:
+                buffer = ""
+                async for chunk in TextReceiveStream(process.stdout):
+                    lines = (buffer + chunk).split("\n")
+                    buffer = lines.pop()
+
+                    for line in lines:
+                        try:
+                            message = JSONRPCMessage.model_validate_json(line)
+                        except Exception as exc:
+                            await read_stream_writer.send(exc)
+                            continue
+
+                        await read_stream_writer.send(message)
+        except anyio.ClosedResourceError:
+            await anyio.lowlevel.checkpoint()
+
+    async def stdin_writer():
+        assert process.stdin, "Opened process is missing stdin"
+
+        try:
+            async with write_stream_reader:
+                async for message in write_stream_reader:
+                    json = message.model_dump_json(by_alias=True)
+                    await process.stdin.send((json + "\n").encode())
+        except anyio.ClosedResourceError:
+            await anyio.lowlevel.checkpoint()
+
+    async with (
+        anyio.create_task_group() as tg,
+        process,
+    ):
+        tg.start_soon(stdout_reader)
+        tg.start_soon(stdin_writer)
+        yield read_stream, write_stream

mcp_python/server/__init__.py

@@ -0,0 +1,347 @@
+import contextvars
+import logging
+import warnings
+from collections.abc import Awaitable, Callable
+from typing import Any
+
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+from pydantic import AnyUrl
+
+from mcp_python.server import types
+from mcp_python.server.session import ServerSession
+from mcp_python.server.stdio import stdio_server as stdio_server
+from mcp_python.shared.context import RequestContext
+from mcp_python.shared.session import RequestResponder
+from mcp_python.types import (
+    METHOD_NOT_FOUND,
+    CallToolRequest,
+    ClientNotification,
+    ClientRequest,
+    CompleteRequest,
+    ErrorData,
+    JSONRPCMessage,
+    ListResourcesRequest,
+    ListResourcesResult,
+    LoggingLevel,
+    ProgressNotification,
+    Prompt,
+    PromptReference,
+    ReadResourceRequest,
+    ReadResourceResult,
+    Resource,
+    ResourceReference,
+    ServerResult,
+    SetLevelRequest,
+    SubscribeRequest,
+    UnsubscribeRequest,
+)
+
+logger = logging.getLogger(__name__)
+
+
+request_ctx: contextvars.ContextVar[RequestContext] = contextvars.ContextVar(
+    "request_ctx"
+)
+
+
+class Server:
+    def __init__(self, name: str):
+        self.name = name
+        self.request_handlers: dict[type, Callable[..., Awaitable[ServerResult]]] = {}
+        self.notification_handlers: dict[type, Callable[..., Awaitable[None]]] = {}
+        logger.info(f"Initializing server '{name}'")
+
+    @property
+    def request_context(self) -> RequestContext:
+        """If called outside of a request context, this will raise a LookupError."""
+        return request_ctx.get()
+
+    def list_prompts(self):
+        from mcp_python.types import ListPromptsRequest, ListPromptsResult
+
+        def decorator(func: Callable[[], Awaitable[list[Prompt]]]):
+            logger.debug(f"Registering handler for PromptListRequest")
+
+            async def handler(_: Any):
+                prompts = await func()
+                return ServerResult(ListPromptsResult(prompts=prompts))
+
+            self.request_handlers[ListPromptsRequest] = handler
+            return func
+
+        return decorator
+
+    def get_prompt(self):
+        from mcp_python.types import (
+            GetPromptRequest,
+            GetPromptResult,
+            ImageContent,
+            Role as Role,
+            SamplingMessage,
+            TextContent,
+        )
+
+        def decorator(
+            func: Callable[
+                [str, dict[str, str] | None], Awaitable[types.PromptResponse]
+            ],
+        ):
+            logger.debug(f"Registering handler for GetPromptRequest")
+
+            async def handler(req: GetPromptRequest):
+                prompt_get = await func(req.params.name, req.params.arguments)
+                messages = []
+                for message in prompt_get.messages:
+                    match message.content:
+                        case str() as text_content:
+                            content = TextContent(type="text", text=text_content)
+                        case types.ImageContent() as img_content:
+                            content = ImageContent(
+                                type="image",
+                                data=img_content.data,
+                                mimeType=img_content.mime_type,
+                            )
+                        case _:
+                            raise ValueError(
+                                f"Unexpected content type: {type(message.content)}"
+                            )
+
+                    sampling_message = SamplingMessage(
+                        role=message.role, content=content
+                    )
+                    messages.append(sampling_message)
+
+                return ServerResult(
+                    GetPromptResult(description=prompt_get.desc, messages=messages)
+                )
+
+            self.request_handlers[GetPromptRequest] = handler
+            return func
+
+        return decorator
+
+    def list_resources(self):
+        def decorator(func: Callable[[], Awaitable[list[Resource]]]):
+            logger.debug(f"Registering handler for ListResourcesRequest")
+
+            async def handler(_: Any):
+                resources = await func()
+                return ServerResult(
+                    ListResourcesResult(resources=resources, resourceTemplates=None)
+                )
+
+            self.request_handlers[ListResourcesRequest] = handler
+            return func
+
+        return decorator
+
+    def read_resource(self):
+        from mcp_python.types import (
+            BlobResourceContents,
+            TextResourceContents,
+        )
+
+        def decorator(func: Callable[[AnyUrl], Awaitable[str | bytes]]):
+            logger.debug(f"Registering handler for ReadResourceRequest")
+
+            async def handler(req: ReadResourceRequest):
+                result = await func(req.params.uri)
+                match result:
+                    case str(s):
+                        content = TextResourceContents(
+                            uri=req.params.uri,
+                            text=s,
+                            mimeType="text/plain",
+                        )
+                    case bytes(b):
+                        import base64
+
+                        content = BlobResourceContents(
+                            uri=req.params.uri,
+                            blob=base64.urlsafe_b64encode(b).decode(),
+                            mimeType="application/octet-stream",
+                        )
+
+                return ServerResult(
+                    ReadResourceResult(
+                        contents=[content],
+                    )
+                )
+
+            self.request_handlers[ReadResourceRequest] = handler
+            return func
+
+        return decorator
+
+    def set_logging_level(self):
+        from mcp_python.types import EmptyResult
+
+        def decorator(func: Callable[[LoggingLevel], Awaitable[None]]):
+            logger.debug(f"Registering handler for SetLevelRequest")
+
+            async def handler(req: SetLevelRequest):
+                await func(req.params.level)
+                return ServerResult(EmptyResult())
+
+            self.request_handlers[SetLevelRequest] = handler
+            return func
+
+        return decorator
+
+    def subscribe_resource(self):
+        from mcp_python.types import EmptyResult
+
+        def decorator(func: Callable[[AnyUrl], Awaitable[None]]):
+            logger.debug(f"Registering handler for SubscribeRequest")
+
+            async def handler(req: SubscribeRequest):
+                await func(req.params.uri)
+                return ServerResult(EmptyResult())
+
+            self.request_handlers[SubscribeRequest] = handler
+            return func
+
+        return decorator
+
+    def unsubscribe_resource(self):
+        from mcp_python.types import EmptyResult
+
+        def decorator(func: Callable[[AnyUrl], Awaitable[None]]):
+            logger.debug(f"Registering handler for UnsubscribeRequest")
+
+            async def handler(req: UnsubscribeRequest):
+                await func(req.params.uri)
+                return ServerResult(EmptyResult())
+
+            self.request_handlers[UnsubscribeRequest] = handler
+            return func
+
+        return decorator
+
+    def call_tool(self):
+        from mcp_python.types import CallToolResult
+
+        def decorator(func: Callable[..., Awaitable[Any]]):
+            logger.debug(f"Registering handler for CallToolRequest")
+
+            async def handler(req: CallToolRequest):
+                result = await func(req.params.name, **(req.params.arguments or {}))
+                return ServerResult(CallToolResult(toolResult=result))
+
+            self.request_handlers[CallToolRequest] = handler
+            return func
+
+        return decorator
+
+    def progress_notification(self):
+        def decorator(
+            func: Callable[[str | int, float, float | None], Awaitable[None]],
+        ):
+            logger.debug(f"Registering handler for ProgressNotification")
+
+            async def handler(req: ProgressNotification):
+                await func(
+                    req.params.progressToken, req.params.progress, req.params.total
+                )
+
+            self.notification_handlers[ProgressNotification] = handler
+            return func
+
+        return decorator
+
+    def completion(self):
+        """Provides completions for prompts and resource templates"""
+        from mcp_python.types import CompleteResult, Completion, CompletionArgument
+
+        def decorator(
+            func: Callable[
+                [PromptReference | ResourceReference, CompletionArgument],
+                Awaitable[Completion | None],
+            ],
+        ):
+            logger.debug(f"Registering handler for CompleteRequest")
+
+            async def handler(req: CompleteRequest):
+                completion = await func(req.params.ref, req.params.argument)
+                return ServerResult(
+                    CompleteResult(
+                        completion=completion
+                        if completion is not None
+                        else Completion(values=[], total=None, hasMore=None),
+                    )
+                )
+
+            self.request_handlers[CompleteRequest] = handler
+            return func
+
+        return decorator
+
+    async def run(
+        self,
+        read_stream: MemoryObjectReceiveStream[JSONRPCMessage | Exception],
+        write_stream: MemoryObjectSendStream[JSONRPCMessage],
+    ):
+        with warnings.catch_warnings(record=True) as w:
+            async with ServerSession(read_stream, write_stream) as session:
+                async for message in session.incoming_messages:
+                    logger.debug(f"Received message: {message}")
+
+                    match message:
+                        case RequestResponder(request=ClientRequest(root=req)):
+                            logger.info(
+                                f"Processing request of type {type(req).__name__}"
+                            )
+                            if type(req) in self.request_handlers:
+                                handler = self.request_handlers[type(req)]
+                                logger.debug(
+                                    f"Dispatching request of type {type(req).__name__}"
+                                )
+
+                                try:
+                                    # Set our global state that can be retrieved via
+                                    # app.get_request_context()
+                                    token = request_ctx.set(
+                                        RequestContext(
+                                            message.request_id,
+                                            message.request_meta,
+                                            session,
+                                        )
+                                    )
+                                    response = await handler(req)
+                                    # Reset the global state after we are done
+                                    request_ctx.reset(token)
+                                except Exception as err:
+                                    response = ErrorData(
+                                        code=0, message=str(err), data=None
+                                    )
+
+                                await message.respond(response)
+                            else:
+                                await message.respond(
+                                    ErrorData(
+                                        code=METHOD_NOT_FOUND,
+                                        message="Method not found",
+                                    )
+                                )
+
+                            logger.debug("Response sent")
+                        case ClientNotification(root=notify):
+                            if type(notify) in self.notification_handlers:
+                                assert type(notify) in self.notification_handlers
+
+                                handler = self.notification_handlers[type(notify)]
+                                logger.debug(
+                                    f"Dispatching notification of type {type(notify).__name__}"
+                                )
+
+                                try:
+                                    await handler(notify)
+                                except Exception as err:
+                                    logger.error(
+                                        f"Uncaught exception in notification handler: {err}"
+                                    )
+
+                    for warning in w:
+                        logger.info(
+                            f"Warning: {warning.category.__name__}: {warning.message}"
+                        )

mcp_python/server/__main__.py

@@ -0,0 +1,35 @@
+import logging
+import sys
+
+import anyio
+
+from mcp_python.server.session import ServerSession
+from mcp_python.server.stdio import stdio_server
+
+if not sys.warnoptions:
+    import warnings
+
+    warnings.simplefilter("ignore")
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger("server")
+
+
+async def receive_loop(session: ServerSession):
+    logger.info("Starting receive loop")
+    async for message in session.incoming_messages:
+        if isinstance(message, Exception):
+            logger.error("Error: %s", message)
+            continue
+
+        logger.info("Received message from client: %s", message)
+
+
+async def main():
+    async with stdio_server() as (read_stream, write_stream):
+        async with ServerSession(read_stream, write_stream) as session, write_stream:
+            await receive_loop(session)
+
+
+if __name__ == "__main__":
+    anyio.run(main, backend="trio")

mcp_python/server/session.py

@@ -0,0 +1,203 @@
+from enum import Enum
+from typing import Any
+
+import anyio
+import anyio.lowlevel
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+from pydantic import AnyUrl
+
+from mcp_python.shared.session import (
+    BaseSession,
+    RequestResponder,
+)
+from mcp_python.shared.version import SUPPORTED_PROTOCOL_VERSION
+from mcp_python.types import (
+    ClientNotification,
+    ClientRequest,
+    CreateMessageResult,
+    EmptyResult,
+    Implementation,
+    IncludeContext,
+    InitializedNotification,
+    InitializeRequest,
+    InitializeResult,
+    JSONRPCMessage,
+    LoggingLevel,
+    SamplingMessage,
+    ServerCapabilities,
+    ServerNotification,
+    ServerRequest,
+    ServerResult,
+)
+
+
+class InitializationState(Enum):
+    NotInitialized = 1
+    Initializing = 2
+    Initialized = 3
+
+
+class ServerSession(
+    BaseSession[
+        ServerRequest,
+        ServerNotification,
+        ServerResult,
+        ClientRequest,
+        ClientNotification,
+    ]
+):
+    _initialized: InitializationState = InitializationState.NotInitialized
+
+    def __init__(
+        self,
+        read_stream: MemoryObjectReceiveStream[JSONRPCMessage | Exception],
+        write_stream: MemoryObjectSendStream[JSONRPCMessage],
+    ) -> None:
+        super().__init__(read_stream, write_stream, ClientRequest, ClientNotification)
+        self._initialization_state = InitializationState.NotInitialized
+
+    async def _received_request(
+        self, responder: RequestResponder[ClientRequest, ServerResult]
+    ):
+        match responder.request.root:
+            case InitializeRequest():
+                self._initialization_state = InitializationState.Initializing
+                await responder.respond(
+                    ServerResult(
+                        InitializeResult(
+                            protocolVersion=SUPPORTED_PROTOCOL_VERSION,
+                            capabilities=ServerCapabilities(
+                                logging=None,
+                                resources=None,
+                                tools=None,
+                                experimental=None,
+                                prompts={},
+                            ),
+                            serverInfo=Implementation(
+                                name="mcp_python", version="0.1.0"
+                            ),
+                        )
+                    )
+                )
+            case _:
+                if self._initialization_state != InitializationState.Initialized:
+                    raise RuntimeError(
+                        "Received request before initialization was complete"
+                    )
+
+    async def _received_notification(self, notification: ClientNotification) -> None:
+        # Need this to avoid ASYNC910
+        await anyio.lowlevel.checkpoint()
+        match notification.root:
+            case InitializedNotification():
+                self._initialization_state = InitializationState.Initialized
+            case _:
+                if self._initialization_state != InitializationState.Initialized:
+                    raise RuntimeError(
+                        "Received notification before initialization was complete"
+                    )
+
+    async def send_log_message(
+        self, level: LoggingLevel, data: Any, logger: str | None = None
+    ) -> None:
+        """Send a log message notification."""
+        from mcp_python.types import (
+            LoggingMessageNotification,
+            LoggingMessageNotificationParams,
+        )
+
+        await self.send_notification(
+            ServerNotification(
+                LoggingMessageNotification(
+                    method="notifications/message",
+                    params=LoggingMessageNotificationParams(
+                        level=level,
+                        data=data,
+                        logger=logger,
+                    ),
+                )
+            )
+        )
+
+    async def send_resource_updated(self, uri: AnyUrl) -> None:
+        """Send a resource updated notification."""
+        from mcp_python.types import (
+            ResourceUpdatedNotification,
+            ResourceUpdatedNotificationParams,
+        )
+
+        await self.send_notification(
+            ServerNotification(
+                ResourceUpdatedNotification(
+                    method="notifications/resources/updated",
+                    params=ResourceUpdatedNotificationParams(uri=uri),
+                )
+            )
+        )
+
+    async def request_create_message(
+        self,
+        messages: list[SamplingMessage],
+        *,
+        max_tokens: int,
+        system_prompt: str | None = None,
+        include_context: IncludeContext | None = None,
+        temperature: float | None = None,
+        stop_sequences: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
+    ) -> CreateMessageResult:
+        """Send a sampling/create_message request."""
+        from mcp_python.types import (
+            CreateMessageRequest,
+            CreateMessageRequestParams,
+        )
+
+        return await self.send_request(
+            ServerRequest(
+                CreateMessageRequest(
+                    method="sampling/createMessage",
+                    params=CreateMessageRequestParams(
+                        messages=messages,
+                        systemPrompt=system_prompt,
+                        includeContext=include_context,
+                        temperature=temperature,
+                        maxTokens=max_tokens,
+                        stopSequences=stop_sequences,
+                        metadata=metadata,
+                    ),
+                )
+            ),
+            CreateMessageResult,
+        )
+
+    async def send_ping(self) -> EmptyResult:
+        """Send a ping request."""
+        from mcp_python.types import PingRequest
+
+        return await self.send_request(
+            ServerRequest(
+                PingRequest(
+                    method="ping",
+                )
+            ),
+            EmptyResult,
+        )
+
+    async def send_progress_notification(
+        self, progress_token: str | int, progress: float, total: float | None = None
+    ) -> None:
+        """Send a progress notification."""
+        from mcp_python.types import ProgressNotification, ProgressNotificationParams
+
+        await self.send_notification(
+            ServerNotification(
+                ProgressNotification(
+                    method="notifications/progress",
+                    params=ProgressNotificationParams(
+                        progressToken=progress_token,
+                        progress=progress,
+                        total=total,
+                    ),
+                )
+            )
+        )

mcp_python/server/sse.py

@@ -0,0 +1,133 @@
+import logging
+from contextlib import asynccontextmanager
+from typing import Any
+from urllib.parse import quote
+from uuid import UUID, uuid4
+
+import anyio
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+from pydantic import ValidationError
+from sse_starlette import EventSourceResponse
+from starlette.requests import Request
+from starlette.responses import Response
+from starlette.types import Receive, Scope, Send
+
+from mcp_python.types import JSONRPCMessage
+
+logger = logging.getLogger(__name__)
+
+
+class SseServerTransport:
+    """
+    SSE server transport for MCP. This class provides _two_ ASGI applications, suitable to be used with a framework like Starlette and a server like Hypercorn:
+
+        1. connect_sse() is an ASGI application which receives incoming GET requests, and sets up a new SSE stream to send server messages to the client.
+        2. handle_post_message() is an ASGI application which receives incoming POST requests, which should contain client messages that link to a previously-established SSE session.
+    """
+
+    _endpoint: str
+    _read_stream_writers: dict[UUID, MemoryObjectSendStream[JSONRPCMessage | Exception]]
+
+    def __init__(self, endpoint: str) -> None:
+        """
+        Creates a new SSE server transport, which will direct the client to POST messages to the relative or absolute URL given.
+        """
+
+        super().__init__()
+        self._endpoint = endpoint
+        self._read_stream_writers = {}
+        logger.debug(f"SseServerTransport initialized with endpoint: {endpoint}")
+
+    @asynccontextmanager
+    async def connect_sse(self, scope: Scope, receive: Receive, send: Send):
+        if scope["type"] != "http":
+            logger.error("connect_sse received non-HTTP request")
+            raise ValueError("connect_sse can only handle HTTP requests")
+
+        logger.debug("Setting up SSE connection")
+        read_stream: MemoryObjectReceiveStream[JSONRPCMessage | Exception]
+        read_stream_writer: MemoryObjectSendStream[JSONRPCMessage | Exception]
+
+        write_stream: MemoryObjectSendStream[JSONRPCMessage]
+        write_stream_reader: MemoryObjectReceiveStream[JSONRPCMessage]
+
+        read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
+        write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
+
+        session_id = uuid4()
+        session_uri = f"{quote(self._endpoint)}?session_id={session_id.hex}"
+        self._read_stream_writers[session_id] = read_stream_writer
+        logger.debug(f"Created new session with ID: {session_id}")
+
+        sse_stream_writer, sse_stream_reader = anyio.create_memory_object_stream(
+            0, dict[str, Any]
+        )
+
+        async def sse_writer():
+            logger.debug("Starting SSE writer")
+            async with sse_stream_writer, write_stream_reader:
+                await sse_stream_writer.send({"event": "endpoint", "data": session_uri})
+                logger.debug(f"Sent endpoint event: {session_uri}")
+
+                async for message in write_stream_reader:
+                    logger.debug(f"Sending message via SSE: {message}")
+                    await sse_stream_writer.send(
+                        {
+                            "event": "message",
+                            "data": message.model_dump_json(by_alias=True),
+                        }
+                    )
+
+        async with anyio.create_task_group() as tg:
+            response = EventSourceResponse(
+                content=sse_stream_reader, data_sender_callable=sse_writer
+            )
+            logger.debug("Starting SSE response task")
+            tg.start_soon(response, scope, receive, send)
+
+            logger.debug("Yielding read and write streams")
+            yield (read_stream, write_stream)
+
+    async def handle_post_message(
+        self, scope: Scope, receive: Receive, send: Send
+    ) -> None:
+        logger.debug("Handling POST message")
+        request = Request(scope, receive)
+
+        session_id_param = request.query_params.get("session_id")
+        if session_id_param is None:
+            logger.warning("Received request without session_id")
+            response = Response("session_id is required", status_code=400)
+            return await response(scope, receive, send)
+
+        try:
+            session_id = UUID(hex=session_id_param)
+            logger.debug(f"Parsed session ID: {session_id}")
+        except ValueError:
+            logger.warning(f"Received invalid session ID: {session_id_param}")
+            response = Response("Invalid session ID", status_code=400)
+            return await response(scope, receive, send)
+
+        writer = self._read_stream_writers.get(session_id)
+        if not writer:
+            logger.warning(f"Could not find session for ID: {session_id}")
+            response = Response("Could not find session", status_code=404)
+            return await response(scope, receive, send)
+
+        json = await request.json()
+        logger.debug(f"Received JSON: {json}")
+
+        try:
+            message = JSONRPCMessage.model_validate(json)
+            logger.debug(f"Validated client message: {message}")
+        except ValidationError as err:
+            logger.error(f"Failed to parse message: {err}")
+            response = Response("Could not parse message", status_code=400)
+            await response(scope, receive, send)
+            await writer.send(err)
+            return
+
+        logger.debug(f"Sending message to writer: {message}")
+        response = Response("Accepted", status_code=202)
+        await response(scope, receive, send)
+        await writer.send(message)

mcp_python/server/stdio.py

@@ -0,0 +1,60 @@
+import sys
+from contextlib import asynccontextmanager
+
+import anyio
+import anyio.lowlevel
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+
+from mcp_python.types import JSONRPCMessage
+
+
+@asynccontextmanager
+async def stdio_server(
+    stdin: anyio.AsyncFile | None = None, stdout: anyio.AsyncFile | None = None
+):
+    """
+    Server transport for stdio: this communicates with an MCP client by reading from the current process' stdin and writing to stdout.
+    """
+    # Purposely not using context managers for these, as we don't want to close standard process handles.
+    if not stdin:
+        stdin = anyio.wrap_file(sys.stdin)
+    if not stdout:
+        stdout = anyio.wrap_file(sys.stdout)
+
+    read_stream: MemoryObjectReceiveStream[JSONRPCMessage | Exception]
+    read_stream_writer: MemoryObjectSendStream[JSONRPCMessage | Exception]
+
+    write_stream: MemoryObjectSendStream[JSONRPCMessage]
+    write_stream_reader: MemoryObjectReceiveStream[JSONRPCMessage]
+
+    read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
+    write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
+
+    async def stdin_reader():
+        try:
+            async with read_stream_writer:
+                async for line in stdin:
+                    try:
+                        message = JSONRPCMessage.model_validate_json(line)
+                    except Exception as exc:
+                        await read_stream_writer.send(exc)
+                        continue
+
+                    await read_stream_writer.send(message)
+        except anyio.ClosedResourceError:
+            await anyio.lowlevel.checkpoint()
+
+    async def stdout_writer():
+        try:
+            async with write_stream_reader:
+                async for message in write_stream_reader:
+                    json = message.model_dump_json(by_alias=True)
+                    await stdout.write(json + "\n")
+                    await stdout.flush()
+        except anyio.ClosedResourceError:
+            await anyio.lowlevel.checkpoint()
+
+    async with anyio.create_task_group() as tg:
+        tg.start_soon(stdin_reader)
+        tg.start_soon(stdout_writer)
+        yield read_stream, write_stream

mcp_python/server/types.py

@@ -0,0 +1,27 @@
+"""
+This module provides simpler types to use with the server for managing prompts.
+"""
+
+from dataclasses import dataclass
+from typing import Literal
+
+from mcp_python.types import Role
+
+
+@dataclass
+class ImageContent:
+    type: Literal["image"]
+    data: str
+    mime_type: str
+
+
+@dataclass
+class Message:
+    role: Role
+    content: str | ImageContent
+
+
+@dataclass
+class PromptResponse:
+    messages: list[Message]
+    desc: str | None = None

mcp_python/server/websocket.py

@@ -0,0 +1,58 @@
+import logging
+from contextlib import asynccontextmanager
+
+import anyio
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+from starlette.types import Receive, Scope, Send
+from starlette.websockets import WebSocket
+
+from mcp_python.types import JSONRPCMessage
+
+logger = logging.getLogger(__name__)
+
+
+@asynccontextmanager
+async def websocket_server(scope: Scope, receive: Receive, send: Send):
+    """
+    WebSocket server transport for MCP. This is an ASGI application, suitable to be used with a framework like Starlette and a server like Hypercorn.
+    """
+
+    websocket = WebSocket(scope, receive, send)
+    await websocket.accept(subprotocol="mcp")
+
+    read_stream: MemoryObjectReceiveStream[JSONRPCMessage | Exception]
+    read_stream_writer: MemoryObjectSendStream[JSONRPCMessage | Exception]
+
+    write_stream: MemoryObjectSendStream[JSONRPCMessage]
+    write_stream_reader: MemoryObjectReceiveStream[JSONRPCMessage]
+
+    read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
+    write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
+
+    async def ws_reader():
+        try:
+            async with read_stream_writer:
+                async for message in websocket.iter_json():
+                    try:
+                        client_message = JSONRPCMessage.model_validate(message)
+                    except Exception as exc:
+                        await read_stream_writer.send(exc)
+                        continue
+
+                    await read_stream_writer.send(client_message)
+        except anyio.ClosedResourceError:
+            await websocket.close()
+
+    async def ws_writer():
+        try:
+            async with write_stream_reader:
+                async for message in write_stream_reader:
+                    obj = message.model_dump(by_alias=True, mode="json")
+                    await websocket.send_json(obj)
+        except anyio.ClosedResourceError:
+            await websocket.close()
+
+    async with anyio.create_task_group() as tg:
+        tg.start_soon(ws_reader)
+        tg.start_soon(ws_writer)
+        yield (read_stream, write_stream)

mcp_python/shared/context.py

@@ -0,0 +1,14 @@
+from dataclasses import dataclass
+from typing import Generic, TypeVar
+
+from mcp_python.shared.session import BaseSession
+from mcp_python.types import RequestId, RequestParams
+
+SessionT = TypeVar("SessionT", bound=BaseSession)
+
+
+@dataclass
+class RequestContext(Generic[SessionT]):
+    request_id: RequestId
+    meta: RequestParams.Meta | None
+    session: SessionT

mcp_python/shared/exceptions.py

@@ -0,0 +1,9 @@
+from mcp_python.types import ErrorData
+
+
+class McpError(Exception):
+    """
+    Exception type raised when an error arrives over an MCP connection.
+    """
+
+    error: ErrorData

mcp_python/shared/progress.py

@@ -0,0 +1,40 @@
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+
+from pydantic import BaseModel
+
+from mcp_python.shared.context import RequestContext
+from mcp_python.shared.session import BaseSession
+from mcp_python.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_python/shared/session.py

@@ -0,0 +1,244 @@
+from contextlib import AbstractAsyncContextManager
+from typing import Generic, TypeVar
+
+import anyio
+import anyio.lowlevel
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+from pydantic import BaseModel
+
+from mcp_python.shared.exceptions import McpError
+from mcp_python.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],
+    ) -> 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._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")
+        )
+
+        # TODO: Support progress callbacks
+
+        await self._write_stream.send(JSONRPCMessage(jsonrpc_request))
+
+        response_or_error = await response_stream_reader.receive()
+        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")
+        )
+
+        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"),
+            )
+            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")
+                    )
+                    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")
+                    )
+
+                    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(
+                                f"Received response with an unknown 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_python/shared/version.py

@@ -0,0 +1 @@
+SUPPORTED_PROTOCOL_VERSION = 1